[WIP][feat] prompt-native generative recommendation - #625
Draft
WhiteSwan1 wants to merge 88 commits into
Draft
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>
_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>
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>
Adds the config surface for prompt-native generative recommendation and the compiler that turns it into the artifacts the data layer, model and serving read: a resolved SidSpace, a PromptPlan walk order, a ModulePlan projection topology, and an extended tokenizer. The compiler resolves no physical dimension. Slot fill mode is derived, not configured: a lone sequence member declaring no embedding renders INLINE, and everything else is PROJECTED and reaches the LM hidden size through the slot's projection, which the model sizes from group_total_dim at __init__. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Walks a compiled PromptPlan to build the packed token stream: static runs, the base_vocab shift on INLINE SID slots, sentinels plus recorded hole positions on PROJECTED ones, and labels that cover the response span only. Band validation lives here rather than in a feature because the assembler owns the codebook, and running it in the worker fails the offending sample instead of letting one rank raise and hang its peers on the collective. Its error names offset_codebook, since a raw or origin column is well formed and would otherwise train silently on pre-relocation SIDs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Reconciles a prompt slot's group width with the LM hidden size: an optional body, then a bare Linear that is structural rather than configurable. Every Perceptron applies its activation and MLP.output_dim() raises on an empty stack, so ending on the MLP would apply an activation to the LM input space and an empty body would not degrade to a linear. The module takes in_dim from the caller. Nothing here or in the compiled plan stores a dimension: the model resolves both ends at __init__, from group_total_dim and the backbone config. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The model reads target_vocab, the module plan and the decode bands off a CompiledPrompt and never reads the prompt's structure. Its config carries only what belongs to the LM: the template, slots, SID space and tokenizer are prompt_config's, so hf_model_id now names the weights alone. The forward gathers the assembled ids through the LM's own input embedding, then index_copy overwrites the projected positions -- out of place, since that gather carries grad. Padding is confined to one adapter at the LM boundary, which takes the collator's max_seqlen rather than lengths.max(): deriving the width here would sync the device to the host on every step. ParamDtype is nested in PromptModelConfig because proto2 enum values are siblings of their enclosing scope, and a top-level one collides. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Wires band-restricted beam decode onto the compiled SidSpace: predict returns the loss while training and the decoded local codes at inference, undoing both the base_vocab and the per-level shift. dynamic_beam_search now prefills from embeddings rather than ids. A projected slot has no vocabulary id, so an id-only prefill cannot express a prompt that contains one. Decode steps still pass ids, because a generated token is always a real vocabulary row. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Every entry point now compiles prompt_config once, right after the features exist, and threads the result to both consumers: the dataloader, whose workers assemble each batch's token stream into additional_infos, and the model, which reads target_vocab and the decode bands off it. Assembly runs after the data parser, not inside any feature's _parse, so the walk stays pure-integer and FG-free. max_seqlen is computed there, on the host, because the model must not derive a shape on the device. Passing prompt through _create_model is conditional so models that take no such kwarg are untouched. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Exercises the real path -- compile_prompt, assemble_into, _create_model, forward and backward -- on a two-layer Qwen saved locally, so the test needs no download. It pins the properties the unit tests cannot see from one layer: that the model resizes to the compiled target_vocab, that loss reaches the backbone embedding, that supervision covers the answer alone, and that raw codes are rejected by the assembler before the model is ever built. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Deletes SidFeature, BaseGenerativeModel, GenerativeQwen, their proto messages, tests and mock config. The prompt-native stack replaces all of it: the per-level offset now arrives in the data, so no feature type is needed to apply it, and the template, SID space and tokenizer live in prompt_config rather than in the model. Nothing is kept for compatibility. The design this implements is design_v2.md. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds a BaseModel.save_assets hook, symmetric with init_from_pretrained, called for every model.ckpt-N/. The prompt-native model writes sid_space.json, prompt_plan.json, the hashes and the extended tokenizer, so a checkpoint describes its own vocabulary rather than relying on config supplied out of band -- which is where offline/online skew comes from. Restore compares the two hashes. A vocab_hash mismatch raises: the SID space or tokenizer changed, so the decode bands no longer address the rows these weights learned, and the run would emit plausible output instead of failing. A plan_hash mismatch only reshapes the prompt, so it warns. ModulePlan is not persisted. It is model-only and rebuilt from config at every __init__, so writing it would create a second source of truth. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Unit tests passed while the pipeline could not run. Running it found: train_and_evaluate never passed prompt= to _create_model, so the model raised at construction. The earlier wiring matched other call sites by shape and missed this one. assemble_into read a dense sequence feature's values as flat, but the parser emits (total, value_dim); the (n, 1) slices made a ragged list. It now flattens, with a test that feeds the parser's real shape. The model implemented none of init_loss, loss, init_metric, update_metric or update_train_metric, which TrainWrapper and the eval loop all call. write_hf_assets required an hf_tokenizer. A prompt-native model owns none: its extended vocabulary is a separately versioned artifact that save_assets writes to prompt/tokenizer, so the tokenizer step is now optional. Verified end to end: three steps train, eval reports ce_loss, the checkpoint carries prompt/ assets, resume works, and a changed codebook is refused at restore. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A PROJECTED slot reaches the model through EmbeddingGroup, but nothing built one: the compiler produced a ModulePlan without the FeatureGroupConfig it implies, so the model raised on embedding_group at construction. Step 12 of the compile algorithm was specified and never implemented; no unit test caught it because every earlier test used an INLINE-only prompt. The compiler now derives one group per projected slot and the model builds its EmbeddingGroup from them. Derived rather than declared, because a prompt group is never shared with a model tower and most of FeatureGroupConfig is meaningless here. Verified by training both shapes end to end. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two gaps, both found by running tzrec.predict rather than by a unit test. predict_checkpoint never passed prompt= to _create_model, so the model raised at construction. That is the third call site with its own shape; all four are now explicit. PredictPipelineSparseDist FX-traces the model, and beam decode reads host ints and branches on them, so tracing died on a Proxy. The decode loop is now a single torch.fx.wrap leaf. One leaf rather than a wrapped inner helper, because every host read inside the loop is untraceable and wrapping one only moves the failure to the next. Verified on both checkpoint shapes: 64 rows in, 64 out, generated_sids of shape (num_return, num_levels), every code inside [0, codebook) -- so the bands held and detokenize inverted both shifts. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The HF branch returns before a model exists, so save_assets never ran and the exported directory carried weights with no vocabulary: serving would have had to reach back into model_dir for the SID space and tokenizer. The branch now copies the checkpoint's prompt/ forward, so one directory is the whole contract. TorchScript export of a prompt-native model died on a missing prompt_input_ids. That is architectural, not a bug: the model's input is an assembled token stream the dataloader builds, and export has no dataloader. The design exports the LM as a HuggingFace directory and reserves TorchScript for the prompt front end, so the format is now refused up front with a message that says which knob to set. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Covers what an operator has to get right: reading offset_codebook rather than codebook or origin_codebook, the config surface, how a slot's fill mode is derived rather than configured, the artifacts a checkpoint carries, and why export is HuggingFace-only. The troubleshooting section is keyed on the exact error strings the code raises, including the two that are fatal by design -- a wrong SID column and a vocabulary that no longer matches the checkpoint. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…c_qwen # Conflicts: # tzrec/main_test.py # tzrec/version.py
save_assets ran on every rank, so a multi-rank save had all of them doing json.dump and copytree to the same paths at once. The files survived a 2-rank local run, but concurrent writes to one path can interleave into a truncated artifact, and a checkpoint whose sid_space.json is truncated is one that restore cannot validate. Guarded in save_prompt_assets, symmetric with write_hf_assets. copy_prompt_assets needs none: the export path already calls it inside is_rank_zero. Verified at --nproc-per-node=2 for both the INLINE-only and projected configs, train and predict. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
suffix_keep was always None, so the model scored logits over every position instead of the answer. It came out None because the answer slot's width was derived from sequence_length, which the answer feature does not declare and should not have to: the answer is exactly one SID item, so its width is the codebook depth. At the toy vocab the smoke runs use, the difference is invisible. At a real vocabulary it is not: full-length logits are batch x length x vocab in fp32, which is terabytes where the window is megabytes, so the first training step would have failed to allocate. The width now comes from the codebook, and an unbounded response is rejected at compile rather than silently falling back to scoring everything. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_create_model built a kwargs dict to avoid handing prompt= to models that do not declare it. BaseModule already absorbs unknown kwargs and forwards none to nn.Module, so a None prompt reaching an unrelated model was never a problem and the conditional guarded nothing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Everything independent of how a family runs its transformer moves to the base: building the LM empty, resizing to the compiled vocabulary, wiring slot projections, the SID coordinate conversions, and the loss, metric and checkpoint hooks. The subclass keeps predict, the teacher-forced loss, the decode loop and the parameters only those need -- ignore_index, generated_sids_key and the beam schedule. Those differ irreducibly rather than by configuration: a decoder-only family reaches past lm(...) into body and head so logits cover a suffix window, while an encoder-decoder would pass labels and get a loss back. Making them abstract states that, where a parameterized hook would imply the difference is one of degree. Llama and Mistral share Qwen's attribute layout exactly, so they need no subclass of their own; the split earns its keep only when a family with a different forward arrives. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collapses the two identical wrapper-unwind walks into one checkpoint_util.unwrap_to(model, attr); hf_export_util had a byte-for-byte copy differing only in the marker attribute. Stops round-tripping constants through the device. The SID band edges were built as tensors and immediately read back with int(), which is 2*num_levels host syncs per decode to recover values the caller already had; they now stay Python ints. The level offsets decode subtracts every step become a buffer, and the assembler's per-level bounds are hoisted out of the per-row path. Drops what nothing reads: PromptPlan.slot_index, Static.owner_slot_id (always None, so its documented drop behaviour never existed), PromptProjection.in_dim and the one-line _slot_in_dim wrapper. Inference no longer builds and scatters a label tensor it discards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two independent defects. _slot_width read sequence_length off the feature's own config, but a member of a SequenceFeature group never sets that field -- the cap is resolved by BaseFeature.sequence_length from the group and passed in by create_features. A grouped feature therefore compiled as UNBOUNDED and the resulting error told the user to set a field they had already set. Reading the resolved property fixes both. The teacher-forced forward reached _unpack, which converts the collator's max_seqlen to a host int. TrainPipelineSparseDist symbolically traces the model whenever a sharded module exists, and int() on a Proxy raises. _fx_wrapped_loss makes that path one FX leaf, the same treatment decode already had; the embedding lookup stays outside the leaf so the pipeline can still see the sharded module and prefetch it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_unpack right-padded: mask = columns < lengths puts real tokens at the left of each row. Both consumers index from the right. _forward_loss scores slice(-suffix_keep, None) and dynamic_beam_search prefills from last_hidden_state[:, -1, :], so for any row shorter than the collator's max_seqlen the loss window fell on padding -- whose labels are ignore_index -- and decode prefilled from a masked-out position. Rows below the batch width were therefore trained on a fraction of their SID levels, or none, and scored from a garbage hidden state. dynamic_beam_search already documented its input as left-padded, and its decode loop only works that way: it appends mask ones on the right for generated tokens and derives positions from the mask's running count. Flipping the mask to pad on the left satisfies that contract and leaves the row-major scatter, which fills in packed order either way, intact. The three existing _unpack tests asserted the right-padded layout, so they encoded the bug; they now pin the left-padded one. Added tests that every row ends on its own final token and that a short row contributes the same number of supervised positions as a long one -- all five fail against the previous mask. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The base class had no test file. It was exercised only incidentally, by the integration test constructing a Qwen model, which reaches its happy path and nothing else: both __init__ guards, the shared-projection width check, the projected scatter in _prompt_embeds, the loss and metric hooks, save_assets and init_from_pretrained had no coverage at all. DetokenizeTest was worse than absent. It lived in the Qwen test file though _detokenize belongs to the base, and it never called the method -- it recomputed tokens - base_vocab - offsets inline and asserted on its own arithmetic, so it would have passed with _detokenize deleted. It is replaced by tests that call the method on a constructed model. Every test here was checked by mutating the behavior it claims to cover and confirming it fails. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Ruff's isort rule orders torchrec before transformers and wraps the assembler import past 88 columns; the committed file had neither, so Code Style CI failed on a hook that fixes files in place. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
assemble_into constructed a PromptAssembler on every batch, though the plan it walks and the band edges it validates against are fixed for the run. The per-batch reshaping moves into PromptAssembler.assemble_batch, so BaseDataset can hold one instance built in __init__; assemble_into stays as a one-shot wrapper for callers that assemble a single batch. A training run now constructs two assemblers, one per dataloader, instead of one per step. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replaces the SID-feature-based generative stack with a prompt-native one: a new
prompt_configowns the template, slots, SID space and tokenizer, andmodel_configkeeps only what belongs to the LM.Why
The per-level SID offset used to be applied twice — once in the SID tool's
sid_offset_codes, once again inSidFeature._parse— with no hash covering thepair. The collision tool now emits an
offset_codebookcolumn, so the offsetarrives already applied and that second implementation can go.
What remains inside tzrec is a uniform
+ base_vocab, which is notposition-dependent and therefore needs no bespoke feature type.
SidFeatureisdeleted; an INLINE slot is an ordinary
sequence_raw_featureand a PROJECTED oneis a
sequence_id_featurewith an injectednum_buckets.The offset form is also self-validating: a raw code
100is legal under codebook4096 and 8192 alike, so a tool/config mismatch passes a range check, whereas
offset codes carry
cumsum(codebook[:l])inside the value and throw every levelWhat is here
tzrec/protos/prompt.protoPromptConfig,PromptSlot,PromptProjection,SidSpacetzrec/prompt/compile.pySidSpace, aPromptPlan, aModulePlanand an extended tokenizertzrec/prompt/assembler.pytzrec/prompt/persist.pytzrec/models/prompt_generative_qwen.pyinputs_embedsforward,index_copyscatter, band-constrained decodetzrec/modules/prompt_projection.pydocs/source/models/prompt_generative_qwen.mdRemoved:
sid_feature.py,generative_model.py,generative_qwen.py,generative_model.protoand their tests. Nothing is kept for compatibility.Two properties worth reviewing specifically:
in_dim/out_dim;the model resolves both at
__init__fromgroup_total_dimand the backboneconfig, so one compiled plan is valid across model sizes.
max_seqlenis computed by thecollator on the host and carried in
additional_infos; the model never callslengths.max().Test Plan
Unit: 55 tests across compile, assemble, projection, persistence and the model.
pre-commit run -aandpython scripts/pyre_check.pyboth clean.End to end, on a two-layer Qwen saved locally (no download), both an INLINE-only
prompt and one with an added PROJECTED slot:
tzrec.train_evalce_loss 4.79againstln(128) = 4.85— near-random from a fresh init, as expectedce_losscomputed and written totrain_eval_result_v2.txtmodel.ckpt-N/prompt/carriessid_space.json,prompt_plan.json,prompt_hashes.json,tokenizer/--continue_train8,8,8is refused at restoretzrec.predictgenerated_sidsof shape(num_return, num_levels); every code inside[0, codebook), so the bands held and detokenize inverted both shiftstzrec.export(HF)AutoModelForCausalLM.from_pretrained; export dir carries the prompt contractEight defects were found by running the pipeline rather than by unit tests —
four missed
_create_modelcall sites, a(total, value_dim)shape theassembler mis-read, the absent loss/metric hooks, an FX
Proxyin decode, andthe un-derived feature groups a projected slot needs. All are fixed and covered.
Not done
--nproc-per-node=1, sosharding, DMP and the planner have not been exercised. The projected config is
the one with a sparse table to shard.
cache_ids(§11 of the design). Withoutcache_ids, a stock radix prefix cache keyed on token ids will reuse KV acrossusers at a sentinel position — a correctness bug, so it must land before any
projected slot is served.
export_format: HF;the model's input is a stream the dataloader assembles, which an export-time
dummy batch cannot supply.
Open design question
design_v2.md§8.1 has an "INLINE text" branch, butTokenizeFeatureextendsIdFeatureand is therefore always sparse, hence always PROJECTED. No existingtype can be an INLINE text slot. I implemented INLINE as unconditionally meaning
SID. Either the branch goes, or
TokenizeFeaturegains a no-embedding mode.