Skip to content

[feat] Add Qwen generative recommendation model - #612

Open
WhiteSwan1 wants to merge 63 commits into
alibaba:masterfrom
WhiteSwan1:support_qwen
Open

[feat] Add Qwen generative recommendation model#612
WhiteSwan1 wants to merge 63 commits into
alibaba:masterfrom
WhiteSwan1:support_qwen

Conversation

@WhiteSwan1

@WhiteSwan1 WhiteSwan1 commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add a reusable GenerativeRecLM base and a Qwen2RecLM implementation for teacher-forced training and beam-search inference.
  • Extend the Qwen tokenizer with per-level SID atoms. Dataset codes stay local 0-based values, while the model derives and applies level_offsets during tokenization and removes them after generation.
  • Add ALGR-style escalating beam search, cold-start Hugging Face weight initialization, DCP checkpoint metadata, and DCP-to-HF export support.
  • Register the model and export configuration in protobuf, add a linear-decay learning-rate scheduler, and cover SID mapping, prompt splicing, generation, and beam behavior with unit tests.

Compatibility and configuration

  • History rows must contain complete items in level order; each code must be in [0, codebook[level]). Data producers must not pre-apply level offsets.
  • The answer is read from data_config.label_field, while history is read from the single JAGGED_SEQUENCE feature group.
  • Hugging Face export is opt-in through export_config.export_format = HF; the existing TorchScript path remains the default.
  • Add transformers 4.51.2 as the pinned runtime backbone dependency.

Testing

  • PYTHONPATH=. python -m unittest tzrec.models.generative_rec_lm_test tzrec.models.qwen2_rec_lm_test (37 tests passed)
  • ruff check on the four modified model and test modules
  • ruff format --check on the same modules
  • git diff --check
  • bash scripts/gen_proto.sh

WhiteSwan1 and others added 30 commits June 8, 2026 09:04
- tzrec/models/generative_rec_lm.py, qwen2_rec_lm.py
- tzrec/protos/models/generative_model.proto + model.proto oneof entry
- tzrec/optim: lr_scheduler additions; optimizer.proto grad-accum/grad-clip
- tzrec/tools/export_genreclm_to_hf.py (DCP -> HF export)

Example scripts and design notes intentionally excluded (to be refactored).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- GenerativeRecLM (base): architecture-agnostic plumbing — vocab extension,
  _tokenize_sids (SID->token-id offset map), _sid_token_rows (jagged read +
  tokenize-once + split, with data-boundary answer-width validation), device
  property, loss/metrics; _build_prompt_tokens / predict are abstract hooks.
- Qwen2RecLM (subclass): ChatML template, causal-LM splice, decoder-only forward.
  _splice_input_ids builds input_ids/mask via pad_sequence and labels in one
  vectorized write (fixed answer width = len(codebook) levels).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- generative_rec_lm_test: registry dispatch, abstract-hook errors, device
  property, _tokenize_sids offset map, _sid_token_rows split/cast/(N,1)-squeeze
  and answer-width validation (ok + violation).
- qwen2_rec_lm_test: splice layout + label masking, left-padding/varied lengths,
  mask keeps trailing eos when pad==eos, _min_first_non_neg_index,
  _build_prompt_tokens buffer registration.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…rence beam search

predict() dispatches on the TER inference flag (BaseModule.is_inference, set by
main.py's set_is_inference before the predict/export wrappers):
- Branch 1 (not is_inference -> train/eval): existing teacher-forced forward +
  suffix-slice + CE loss (moved to _predict_train).
- Branch 2 (is_inference -> inference): _generate beam-searches the SID answer
  from an answer-less prompt (_splice_prompt_ids), emitting num_levels tokens/beam
  and mapping them back to raw SID indices -> {"generated_sids": (B, num_return, L)}.

Beam params (_num_beams/_num_return) default to algr's 50/50 (optional proto
fields). Tests cover prompt splice, is_inference routing, and token->SID map.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…name dispatch

Each LLM family is its own model_config oneof entry whose message-type name
resolves straight to the same-named class (qwen2_rec_lm -> Qwen2RecLM), so
GenerativeRecLM.__new__ and the class_name field are gone.

Proto split:
- GenerativeRecLMConfig = architecture-AGNOSTIC config (codebook, vocab pad,
  feature names, ignore_index, beam params), embedded as `common`.
- The backbone `hf_model_id` is OWNED by the family message (Qwen2RecLM,
  default "Qwen/Qwen2.5-0.5B"), NOT `common` — the registered family IS the
  architecture commitment, so a backbone in the shared block could contradict
  it. (common field 1 reserved.)
- Family-specific chat-template knobs also live on the family message.

Base __init__ reads cfg.common.* (shared) and cfg.hf_model_id (family-owned).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…k VRAM imbalance

One GPU reserved ~25GB more than the rest during 8-GPU training: the CUDA
caching allocator stranded a whole segment generation on whichever rank drew
its shortest batch first (variable seq-len + native allocator never shrinks
reserved). Pre-size the pool up front so it never has to grow mid-run.

- Qwen2RecLM warms the activation pool with a one-shot dummy fwd+bwd at the
  worst-case (batch_size, T_max) on the first training step (earliest the HF
  backbone is on-GPU); T_max = template frame + sequence_length + num_levels.
- Pool length reuses the user_sequence feature's sequence_length
  (GenerativeRecLM._input_sequence_length); _sid_token_rows enforces it with a
  recency-preserving, item-aligned tail clip (keep newest items, drop oldest)
  so it is a guaranteed bound under FG_NONE, which does not truncate.
- Thread data_config.batch_size into _create_model; move num_beams/num_return
  from base to subclass.

Verified on 8xGPU: per-card spread 25GB -> 2.2GB, no OOM.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…=300

Under FG_NONE the reader does not truncate, so the model enforces this length
via the recency-preserving clip in _sid_token_rows and uses it to pre-size the
activation pool. 300 = AL-GR-Tiny's realistic max history (100 items x 3 codes);
the prior loose 1056 (algr max_length) would oversize the warm-up.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…fix export tool schema

Intermediate training checkpoints were only written in TER's DCP format, which
isn't directly from_pretrained-loadable and reduced confidence when validating
results. Now save an HF copy at each periodic checkpoint.

- GenerativeRecLM.export_hf(dir): save self.lm + the rebuilt extended tokenizer
  (base + C0..C{sum-1}) straight from the live model.
- main._train_and_evaluate: after each periodic ckpt_manager.save, rank-0 calls
  _model.export_hf(model_dir/hf_ckpt-{step}) when available (duck-typed; other
  models unaffected).
- export_genreclm_to_hf: fix three new-schema bugs (class_name -> which_msg
  resolution, abstract GenerativeRecLM -> resolved family class + register
  Qwen2RecLM, grl_cfg.codebook -> grl_cfg.common.codebook). The codebook bug
  had exported weights without the tokenizer, breaking predict.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…line tool

The standalone export tool duplicated the backbone+tokenizer save that
GenerativeRecLM.export_hf now owns. Make the tool do only the DCP overlay then
call model.export_hf, so offline and in-training HF exports are one code path
(byte-identical) and there is a single source of truth.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… standalone tool

The standalone tzrec/tools/export_genreclm_to_hf.py re-implemented the model
build + checkpoint restore that TER's export() already does. Delete it and add
a GenerativeRecLM branch directly in main.export: after the existing
checkpoint resolution, overlay the DCP shards and call model.export_hf (the
same single save path used by the in-training checkpoint hook). No separate
tool, no duplicated save; `python -m tzrec.export` now produces the HF dir for
GenerativeRecLM models.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…t-step pad pre-sizing

Import: __init__ builds the empty extended arch (AutoConfig + from_config, no
weight download); new init_from_pretrained() is the sole from_pretrained, invoked
by the pipeline on cold start via a no-op BaseModel hook (no hasattr duck-typing).
Restore/eval/export load weights from DCP, skipping the ~1GB download.

Export: training writes DCP only; each model.ckpt-N/ co-locates HF config+tokenizer
(no weights), gated by export_config.export_format == HF and owned by
CheckpointManager. tzrec.export converts via a standalone dcp_to_hf (recorded
backbone-prefix recorded as data + suffix-match self-heal + strict 1:1 validation,
never a silent partial load). Adds ExportFormat to export.proto; TORCHSCRIPT path
untouched.

Training: replace _warmup_alloc -- a separate unscaled forward+backward that
corrupted the first optimizer step (the lr5e-5 HR 2.6-2.9->1.14-flat regression) --
with first-step left-padding to the worst-case length in _predict_train /
_splice_input_ids / _left_pad. The pad rides the real step (positions masked +
labelled -100, so loss/grad are identical) while pre-sizing the activation pool to
(B, T_max), keeping per-rank reservations uniform.

Tests: rewrite the warmup unit tests to assert the first-step padding contract.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… comments

__init__ was ~105 lines of mixed concerns + heavy narration. Extract three
behavior-preserving private helpers and condense comments to the load-bearing
rationale only:
- _read_common_config(common) -> sid_atoms: proto knobs + codebook guard.
- _build_backbone() -> module: empty bf16 arch (no weight download) + empty-id guard.
- _build_extended_tokenizer(sid_atoms) -> (tokenizer, base): add C0.. atoms,
  resize self.lm, the added==sid_atoms and C0-at-base guards.

__init__ now reads as super().__init__ -> read config -> build backbone -> build
tokenizer -> pad/prompt/debug tail. All instance attributes, guard messages, and
effect order are unchanged; no behavior change. 22 model unit tests pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Brings in upstream's SidRqkmeans (FAISS residual K-Means SID) feature + the
maybe_save / event-time checkpoint refactor + on_train_end hook (3 commits,
ad083a7..3d4d5a8). Conflicts resolved:

- tzrec/protos/model.proto: both sides added a oneof field at 601. Kept
  upstream's `sid_rqkmeans = 601`; moved our `qwen2_rec_lm` to 700 (a clear
  100-block for the generative-LM family) and dropped our `reserved 600` in
  favor of upstream's planned SidRqvae=600. (TER configs are text-format / use
  field names, so the renumber is wire-irrelevant.) gen_proto verified: 27 oneof
  fields, all numbers unique.
- tzrec/main.py: adopted upstream's `ckpt_manager.maybe_save(...)` + `run_eval`
  dispatch, dropping our explicit step/epoch/final save blocks. Our HF-asset
  co-location rides along unchanged: maybe_save -> CheckpointManager.save ->
  the export_format==HF gate.

Auto-merged cleanly: models/model.py (our init_from_pretrained hook + upstream's
on_train_end coexist), utils/checkpoint_util.py (our HF save-gate + upstream's
maybe_save coexist). Verified: hot files compile, protos regenerate, genrec unit
tests pass (22).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…1e-5 collapse

Root cause (overnight investigation): the LM was built in bf16 params with
mixed_precision unset, so the optimizer updated bf16 weights directly with no
fp32 master. Adam's small updates at lr=1e-5 (~1e-5) fall below the bf16 ULP of
the weights and round to zero -> weights freeze -> training collapses (eval ce
hard-plateaus ~5.76, HR ~0). lr=5e-5 masked it (5x larger updates clear the ULP).
ALGR's HF-Trainer bf16:true keeps an fp32 master, so it trains fine at lr1e-5.

Fix: build the LM in fp32 in BOTH paths (`_build_backbone` from_config and
`init_from_pretrained` from_pretrained) so the optimizer keeps fp32 master
weights. Set `mixed_precision:"BF16"` in the run config for bf16 *compute* speed
(autocast) on the fp32 master — the standard AMP pattern, mirroring ALGR.

Proven by a single-variable A/B (only precision changed):
  bf16-params:  lr1e-5 = 0.02 flat (collapse) | lr5e-5 = 2.71 rising
  fp32-master:  lr1e-5 = 1.17 rising (1ep)    | lr5e-5 = 2.68 rising
The fix recovers lr1e-5 ~60x without regressing lr5e-5 (2.68 ~= 2.71).

Tradeoff: fp32 master -> fp32 DCP checkpoint (~2x) + ~64GB/GPU (vs 54). Old bf16
checkpoints still restore (upcast to fp32). Details in
ai_report/MASTER_EXPERIMENT_REPORT.md.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Layer A — _generate now validates each beam against the per-level SID bands;
malformed candidates (early EOS / non-SID / wrong-level atom) collapse to a -1
sentinel that can never match a real item, and the fixed-width canvas removes
the reshape crash when beams stop early. The gate + token<->SID inversion live
in the base GenerativeRecLM (_validate_sid_candidates, _sid_level_bands) as one
source of truth, reused by every family.

Code-review findings:
- alibaba#2 (perf): cache self._suffix_keep; _forward_loss slices a constant suffix
  instead of recomputing it per step -> drops 2 GPU->CPU syncs/step.
- alibaba#5: isolate the rank-0 write_hf_assets call (try/except + log) so an asset
  write error can't abort before the next collective and hang other ranks.
- alibaba#7: make the dense-only contract explicit (embedding_group = None); remove the
  never-called init_input (calling it would wrongly build an unused sharded
  table for the SEQUENCE feature, which flows as raw token ids).
- alibaba#8: drop the dead batch_size plumbing from _create_model + call sites.
- alibaba#9: _resolve_pad_token_id asserts pad/eos present (clear error, not int(None)).
- dtype: single _PARAM_DTYPE source of truth (both builders fp32-master).

Also: streamline comments, pin transformers==4.51.2 (<5.0).
Tests: 25 genrec unit tests (Layer-A valid/malformed/narrow-tail, _suffix_keep
equivalence, pad resolution).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the explicit canvas alloc + min/slice-copy + boolean-index assignment
with F.pad (fixed-width -1 padding) and masked_fill (row invalidation): 6 logic
lines -> 4, no in-place indexing, clearer intent. Behavior is unchanged —
verified value-identical to the previous form on 4 edge cases (valid, narrow
tails w=1/2, all-invalid) + 2000 random fuzz batches.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Condense the multi-line inline rationale comments (CE-suffix width, one-shot
pool pre-sizing, fp32-master build, SID-band gate, beam-order reshape) and the
two longest docstrings (init_from_pretrained, _validate_sid_candidates) to terse
1-3 liners, keeping the load-bearing why (fp32-master underflow, -1 can't match
a real item, suffix-slice OOM, pad==eos mask). Comments only — no behavior
change; ruff clean, 25 genrec tests pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The SID->token offset now has both directions as named helpers next to each
other: _tokenize_sids (sid -> token) and its inverse _detokenize_sids
(token -> sid), used by _validate_sid_candidates instead of the inline
`new_tokens - (base_vocab - 1)`. One owner for the offset constant.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The data reader/DataParser does NOT truncate sequences under FG_NONE (only pyfg
does, in FG_NORMAL/FG_DAG; the native EmbeddingGroup does via to_padded_dense at
forward time). GenerativeRecLM uses FG_NONE and no EmbeddingGroup, so the history
cap is enforced model-side in _sid_token_rows (item-aligned to whole num_levels
items). Correct _read_common_config + _input_sequence_length to say so instead of
"the reader caps every row".

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port ALGR's dynamic_beams schedule (beam width doubles per SID level
50->100->200->400, returns num_beams*2**num_levels candidates) as a
torch-only KV-cached kernel; faithful to ALGR's escalating beam and
exploits the fixed-length, EOS-free SID answer.

- tzrec/models/escalating_beam.py: escalating_beam_search kernel
- qwen2_rec_lm.py: _dynamic_beam_search delegates; _generate dispatch on
  the dynamic_beam flag
- generative_model.proto: dynamic_beam flag
- examples/generative_rec_lm_predict.py: --dynamic_beam / --codebook
- tests: exhaustive==brute-force top-k + validity/left-pad (19 pass)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ding

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…t (JAGGED_SEQUENCE)

Route the genrec-LM's SID retrieval through the framework's standard
init_input/build_input + EmbeddingGroup path (the HSTU idiom) instead of
reaching into batch.sequence_dense_features directly.

- proto: add GenerativeRecLMConfig.history_group_name / label_group_name
  (group-name knobs, defaults "user_seq"/"answer"), keyed by GROUP name like
  HSTU, decoupled from feature names.
- generative_rec_lm: init_input builds the param-free raw-passthrough
  EmbeddingGroup; build_input reads "{group}.sequence"/".sequence_length" and
  tokenizes; _sid_token_rows now takes (values, lengths); group knobs read in
  _read_common_config.
- qwen2_rec_lm: _predict_train / _generate consume build_input.
- example config: one SEQUENCE group -> two single-feature JAGGED_SEQUENCE groups.
- tests updated (+ build_input coverage). Integration-verified that a top-level
  sequence_raw_feature passes raw through a JAGGED_SEQUENCE EmbeddingGroup.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…-configurable

Promote the two GenerativeRecLM class constants to GenerativeRecLMConfig knobs,
keeping their previous values as defaults:
- generated_sids_key (default "generated_sids") -> self._generated_sids_key,
  used by _generate's output dict.
- param_dtype (default "float32") -> self._param_dtype via _DTYPE_BY_NAME
  {float32, bfloat16, float16}; used by _build_backbone / init_from_pretrained.
  An unknown value raises a clear ValueError.
Both read in _read_common_config; tests cover defaults + dtype mapping + validation.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…onstant

The value now lives in the GenerativeRecLMConfig.generated_sids_key proto default
("generated_sids"); _generate already emits self._generated_sids_key.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ne .get()

/simplify cleanup: membership-check + dict-index -> single .get() + None-guard
(same ValueError). No behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…(HSTU-style)

Promote the model's history budget to GenerativeRecLMConfig.max_sequence_length
(field 14), mirroring HSTU's DlrmHSTU.max_seq_len — a model knob distinct from the
user_sequence feature's sequence_length. _max_seq_length now reads
common.max_sequence_length, falling back to the feature's sequence_length when 0
(backward-compatible). It drives the recency-preserving truncation cap
(_sid_token_rows) + the activation-pool pre-size. Example config sets it in common;
tests cover the model-knob and the fallback.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ce_length is the sole source

common.max_sequence_length is now the single source for the model's history
budget; remove the feature-derived fallback (_input_sequence_length) and its
test. 0 = off (no cap / no activation-pool pre-allocation). Proto + example-config
comments updated accordingly.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…s2 example config

max_sequence_length is now a REQUIRED GenerativeRecLMConfig field (like HSTU's
DlrmHSTU.max_seq_len) — every genrec config must set it (0 = explicitly off).
Migrate examples/generative_rec_lm_s2pretrained.config to the new format to keep
it valid + consistent: SEQUENCE group -> two JAGGED_SEQUENCE groups (user_seq /
answer) + max_sequence_length: 1056. Both example configs verified to parse.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…t a feature

Move the answer/target SID stream out of feature_configs into
data_config.label_fields (a list<int64> column -> batch.jagged_labels[label]).
build_input now reads HISTORY from the user_seq JAGGED_SEQUENCE group (EmbeddingGroup)
and the ANSWER from batch.jagged_labels[self._label_name]. Drop the now-unused
label_group_name proto knob (reserved 11). This is semantically correct (the answer
is the target, not an input feature) and lets the label be absent at inference
without the EmbeddingGroup requiring it. Example configs (s1/s2) migrated; tests
updated. (Also wrapped the hf_backbone/hf_tokenizer export-only docstrings.)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
WhiteSwan1 and others added 15 commits July 28, 2026 08:55
GenerativeRecLM/QwenRecLM matched neither base-class pattern in tzrec and used
Rec and LM tokens no other model carries. They now mirror the only comparable
family -- BaseSidModel in sid_model.py with SidRqvae/SidRqkmeans beside it --
as BaseGenerativeModel in generative_model.py with GenerativeQwen beside it,
so the module finally shares a name with the generative_model.proto that
drives it and the family is greppable by one prefix. Oneof tag 700 is
unchanged, but the message rename is a config break: a qwen_rec_lm block must
become generative_qwen.

Also guards two config values that silently did nothing: a max_sequence_length
narrower than one item floored to zero whole items and left the history
uncapped, and a feature named by two feature_groups resolved to whichever came
last instead of being rejected. Both now raise, and the escalating-beam module
is renamed to dynamic_beam to match the dynamic_beam proto field it backs.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… decoder

num_beams plus a dynamic_beam bool described the width schedule indirectly: the
kernel derived num_beams * 2**(j+1) internally, so a policy lived inside a
mechanism and no other shape was expressible. GenerativeModelConfig now carries
repeated beam_widths -- [50,50,50] is a fixed beam, [100,200,400] the escalating
one -- and dynamic_beam_search only caps each entry by what its band and the
surviving prefixes can supply.

That also lets the HF generate branch go from _generate. The two decoders were
never equivalent: HF beam search ranges over the whole vocabulary, so candidates
that are not SID atoms had to be discarded as -1, while the band-restricted
kernel makes every candidate well-formed by construction. Given the same band
mask the two agree exactly across 96 configurations, so no reachable behaviour
is lost, and every recorded experiment already selected the dynamic path.

The tiny-backbone test fixture moves into tzrec/utils/test_util.py so the genrec
test modules share one copy, and the proto's reserved markers are dropped:
tzrec configs are protobuf text format, matched by field name, so a reused tag
cannot mis-parse an existing config.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
predict_checkpoint runs batches through PredictPipelineSparseDist, whose
_rewrite_model FX-traces the model to locate shardable modules. The decode
cannot be traced: it turns a jagged batch into per-row python lists, interleaves
them with the prompt, then runs a beam whose widths depend on the data. So
tzrec.predict died with "Proxy object cannot be iterated" before it reached a
single batch, while train and eval were unaffected -- create_train_pipeline
falls back to the un-traced TrainPipelineBase when a model owns no ShardedModule,
and predict_checkpoint has no equivalent guard.

The whole decode now sits behind one torch.fx.wrap leaf. Wrapping an inner
helper does not work: a leaf returns a single Proxy, so the list structure is
lost and the failure only moves to the caller, through all nine untraceable
sites. With one leaf the trace completes and finds nothing to shard, which is
correct -- a SID feature declares no embedding table. At run time the leaf is an
ordinary call.

_generate now returns the tensor and predict owns the output-key contract, since
a wrapped function has to return a tensor for FX to handle it cleanly.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The H20 lane failed on a ReadTimeoutError from files.pythonhosted.org while
pulling the 10.4 MB transformers wheel: the read stalled for eleven minutes,
pip aborted, and every test that touches `import tzrec` then died on
"No module named 'transformers'" -- nine of them in unrelated HSTU code.

transformers was the only large dependency still coming from public PyPI; every
other third-party wheel already resolves from the project's OSS bucket, which
served 348 MB of fbgemm_gpu_hstu in nineteen seconds during the same failed run.
Fetching it from the same bucket takes the stalling transfer off the
cross-border path (measured 10.4 MB in 0.22 s, sha256 unchanged).

The version pin moves into the filename, matching how faiss and graphlearn are
already referenced.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
data_config.fg_mode is a pipeline-wide switch, so SidFeature refusing anything
but FG_NONE did not merely disable fg for itself -- it blocked every other
feature in the same config from using fg at all. A config that needs FG_NORMAL
for its ordinary features therefore could not carry a SID feature.

SidFeature now emits a passthrough fg config: a plain raw_feature that fg only
splits into per-position values, with _parse folding in the level offsets
afterwards exactly as it does on the FG_NONE path. No fg feature_type can add
level_offsets[i % num_levels] -- fg expressions are not position-aware within a
sequence -- so the arithmetic stays in _parse and fg is used purely to reach the
codes. Verified against real pyfg: FG_NONE on a list column and FG_NORMAL on the
delimited-string column that ODPS and CSV deliver produce identical values and
lengths.

sequence_length is now meaningful, since fg applies it, but it truncates by
value count and keeps the head. A cap that is not a whole number of levels would
hand the model a partial item, so it is rejected with a pointer to
max_sequence_length, which is item-aligned and keeps the recent tail.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
# Conflicts:
#	tzrec/protos/export.proto
#	tzrec/utils/checkpoint_util.py
The URL already carries the version, and no other entry in this file annotates
its pin, so the trailing comment was the only one of its kind.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… hang

_unwrap_hf_model descends .module / .model looking for a model that exposes
hf_backbone, and nothing recorded where it had already been. A cycle in that
chain -- a wrapper whose .model points back at itself or an ancestor -- made the
while-loop spin forever.

It matters because of where the walk runs: CheckpointManager.save() calls
write_hf_assets for EVERY model, not just HF-backed ones, so the hang would land
inside checkpoint save and strand the peers waiting on the all_gather that
follows. Today's four wrappers form a strict chain, so this is latent rather
than live.

A set of visited ids bounds the walk; a repeat means the chain cannot reach an
hf_backbone, which is the same answer as running out of attributes. The
regression test walks a cycle on a worker thread so a reoccurrence fails the
suite instead of hanging it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ring

The genrec files were about 30% comment and docstring. Most of it restated the
code beside it -- shapes already visible in the expression, names already in the
signature, proto semantics repeated in three places.

What stays is the part that costs debugging time to rediscover: why a gap is
encoded as one string (a BPE merge must not span a seam), why fp32 master
weights, why validation runs in the dataloader worker rather than the forward
path, why the decode sits behind an fx leaf, why _fg_op is seeded before the
first raise, and why the wrapper walk is bounded.

Docstrings are ruff-enforced so they are shortened, not removed; several drop to
their summary line.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…l locals

beam_widths and num_return_sequences live on GenerativeModelConfig, so parsing
them belonged in the base rather than the Qwen subclass. _read_beam_config and
DEFAULT_BEAM_WIDTH move to BaseGenerativeModel, called from _read_common_config
once _num_levels is known; the width is read through self so a family can
override it. Its test moves to the base test file for the same reason.

dynamic_beam_search drops the keyword-only marker, and its locals get names that
say what they hold: am -> beam_mask, bsz -> batch_size, h -> outputs,
cur_w -> width, lo_j/hi_j -> band_lo/band_hi, idx -> flat_idx, tok ->
next_token, j -> level.

Also drops the export.proto comments describing ExportFormat, and corrects a
stale one on SidFeature.sequence_length that still claimed fg was forbidden and
the field rejected -- both untrue since the fg passthrough landed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
GenerativeQwenBeamTest claimed to cover the kernel/model seam, but mutation
testing showed it does not: breaking _validate_sid_candidates' sentinel leaves
it passing, because band-restricted output never reaches that path. The two
mutations it does catch -- the band edges and the width cap -- are already
caught by generative_model_test and dynamic_beam_test respectively.

separator and embedding_constraints are unreachable for a SID feature:
_fg_json never emits a separator, and parameter_constraints is only consulted
for a feature with an emb_config, which a SID feature does not have. Removing
either keeps both the unit and integration suites green.

pooling looked equally dead and is not: EmbeddingGroup reads pooling_type for
every member of a group, so removing it fails the integration test with
AttributeError. use_mask and value_dim are likewise read by BaseFeature. All
three stay.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_read_beam_config moved to BaseGenerativeModel, but its test and the shared
fixtures still built GenerativeQwen, so the base test read as though the
subclass owned the method. They now construct BaseGenerativeModel; the only
GenerativeQwen references left are the ones genuinely about the subclass --
registry dispatch, the oneof, and the family proto's hf_model_id default.

Also reaches CHAT_TEMPLATE through self, matching DEFAULT_BEAM_WIDTH. Both are
class constants a family may override and both resolve identically; using two
different spellings for the same intent was the only reason to prefer one.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
DEFAULT_BEAM_WIDTH read as a public constant, but it is neither public nor
constant: nothing outside the class touches it and a family is meant to
override it. _default_beam_width says both.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_default_beam_width is now an instance attribute assigned in __init__, before
_read_common_config parses the beam knobs that consume it.

Note this changes how a family overrides it: a class-level _default_beam_width
on a subclass would be overwritten by this assignment, so an override has to
happen through the constructor or by overriding _read_beam_config. The test
fixtures that build a model with object.__new__ now set it explicitly, since
they bypass __init__.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Both were silently defaultable. beam_widths fell back to a flat width of 50 per
level, so a config that never mentioned the beam still decoded -- at a width
nobody chose, with no way to tell an intended 50 from an unset one.
num_return_sequences carried the same 50 as a proto default even though the
value must relate to the final beam width.

beam_widths now raises when empty, num_return_sequences is proto-required, and
_default_beam_width is gone. codebook was already enforced in code; the proto
and the error message now say it is required rather than merely non-empty.

The mock config gains an explicit [4, 8, 16] schedule, which is what a real
config has to do now too.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@WhiteSwan1 WhiteSwan1 changed the title [WIP] Add Qwen generative recommendation model [feat] Add Qwen generative recommendation model Jul 29, 2026
@WhiteSwan1 WhiteSwan1 added the codex-review Let Codex Review label Jul 29, 2026
@github-actions github-actions Bot removed the codex-review Let Codex Review label Jul 29, 2026
Comment thread tzrec/main.py
Comment thread tzrec/models/generative_model.py
Comment thread tzrec/features/sid_feature.py
Comment thread tzrec/features/sid_feature.py
Comment thread tzrec/protos/feature.proto Outdated
Comment thread tzrec/modules/dynamic_beam.py
Comment thread tzrec/modules/dynamic_beam.py
@github-actions

Copy link
Copy Markdown
Contributor

Static review only; no tests or builds were run.

I left inline comments for the correctness and scalability issues. Two general documentation items also need alignment before merge:

  • The PR description says SID history/labels are 1-based in [1, codebook[level]], but the implementation, proto, and tests enforce 0-based values in [0, codebook[level]). Producers following the current description will shift or reject every boundary value.
  • This adds a public model, feature type, and HF export format without user-facing docs or a runnable example. Please document the data contract, JAGGED_SEQUENCE/prompt requirements, beam/output shape, pretrained loading, and HF export workflow.

WhiteSwan1 and others added 3 commits July 29, 2026 07:07
The HF export branch returns before the model is built and hands the checkpoint
directory to dcp_to_hf, which always reads <checkpoint>/model. Dense EMA lives
in a sibling <checkpoint>/dense_ema that only restore_model overlays, so with
EMA enabled TorchScript export shipped the averaged weights while HF export
silently shipped the raw ones -- same checkpoint, different model, no warning.

The branch now resolves use_dense_ema exactly as the DCP path does and raises
instead of converting. Refusing is preferable to overlaying here: the converter
never builds a model, so it has no place to apply the EMA state without
duplicating restore_model's mapping logic.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
value_dim was configurable but unrepresentable: _parse reshapes by level and
splits by seq_lengths, both of which assume one code per sequence position, so
a wider value lands the level offsets on the wrong components. It is now
rejected at construction rather than described as "stays 1" in a comment.

The SidFeature message comment still claimed FG_NONE-only support, which the fg
passthrough replaced; it now describes what fg actually does.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@WhiteSwan1

Copy link
Copy Markdown
Collaborator Author

Thanks for github actions — all seven reproduce. Four fixed in 07443e3 and 17a4631, two I'd like to leave, one deferred, docs separate.

Fixed

1. HF export bypasses dense EMA (07443e3) — correct, and a regression from the recent upstream merge. dcp_to_hf reads <checkpoint>/model while EMA lives in <checkpoint>/dense_ema, so TorchScript export shipped averaged weights and HF export shipped raw ones. I chose to reject the combination rather than overlay: the converter never builds a model, so applying EMA there would duplicate restore_model's mapping logic. Test added.

3. value_dim unsupported (17a4631) — agreed, _parse assumes one code per position. Rejected at construction now. The field has to stay; BaseFeature reads it unconditionally.

5. Stale FG comment in the proto (17a4631) — correct, it predates the fg passthrough. Fixed.

PR description says 1-based — fixed; the code has been 0-based since that switch.

Not changing

2. num_return_sequences vs. the capped width — real, but the final width only shrinks when beam_widths[-1] > effective[-2] * codebook[-1]. At codebook 512 or 8192 that needs a final width above the codebook itself, which will not fit in memory; even the repo's [4,4,4] mock at [4,8,16] is uncapped. Reaching it takes a codebook around 4 or a decreasing schedule. Intermediate capping is harmless — later levels multiply the width back up — and the failure returns fewer candidates, not wrong ones. The existing check catches the mistake a user actually makes.

4. Non-integral / non-finite codes — the realistic route to NaN is a missing value, and that is already rejected before _parse: pyarrow raises ArrowInvalid: ... 1 nulls, but zero_copy_only was True. So NaN needs a float column with NaN written into it deliberately. The fg point is right that _fg_json declares a float type, but float32 is exact to 2^24 against a largest codebook of 8192 — floats by dtype, never by magnitude.

Deferred

6/7. Beam KV duplication and full-vocab fp32 logits — both accurate. Context the static pass lacked: the band-slice form already cut the sustained score tensor 538 MiB → 25 MiB and level-3 topk 499 ms → 69 ms, and [100, 200, 400] runs complete on H20s. Paged cache indirection and a fused chunked reduction are substantial changes against a kernel verified equivalent to HF beam search across 96 configurations, so I'd rather file them as follow-ups with that profile than take them here.

Docs

Agreed. Landing as a separate PR covering the sample contract, the JAGGED_SEQUENCE group rule, prompt_template, the beam knobs, cold-start loading, and export_format: HF.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant