[feat] Add Qwen generative recommendation model - #612
Conversation
- 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>
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>
|
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 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>
|
Thanks for github actions — all seven reproduce. Four fixed in Fixed1. HF export bypasses dense EMA ( 3. 5. Stale FG comment in the proto ( PR description says 1-based — fixed; the code has been 0-based since that switch. Not changing2. 4. Non-integral / non-finite codes — the realistic route to Deferred6/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 DocsAgreed. Landing as a separate PR covering the sample contract, the JAGGED_SEQUENCE group rule, |
Summary
Compatibility and configuration
Testing