Skip to content

[WIP][feat] prompt-native generative recommendation - #625

Draft
WhiteSwan1 wants to merge 88 commits into
alibaba:masterfrom
WhiteSwan1:feat/prompt_genrec_qwen
Draft

[WIP][feat] prompt-native generative recommendation#625
WhiteSwan1 wants to merge 88 commits into
alibaba:masterfrom
WhiteSwan1:feat/prompt_genrec_qwen

Conversation

@WhiteSwan1

Copy link
Copy Markdown
Collaborator

WIP — opened for design review, not for merge. Multi-rank is untested and two
items from the design remain unbuilt (see Not done below).

Replaces the SID-feature-based generative stack with a prompt-native one: a new
prompt_config owns the template, slots, SID space and tokenizer, and
model_config keeps 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 in SidFeature._parse — with no hash covering the
pair. The collision tool now emits an offset_codebook column, so the offset
arrives already applied and that second implementation can go.

What remains inside tzrec is a uniform + base_vocab, which is not
position-dependent and therefore needs no bespoke feature type. SidFeature is
deleted; an INLINE slot is an ordinary sequence_raw_feature and a PROJECTED one
is a sequence_id_feature with an injected num_buckets.

The offset form is also self-validating: a raw code 100 is legal under codebook
4096 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 level

= 1 out of band.

What is here

area
tzrec/protos/prompt.proto PromptConfig, PromptSlot, PromptProjection, SidSpace
tzrec/prompt/compile.py compiles the config into a resolved SidSpace, a PromptPlan, a ModulePlan and an extended tokenizer
tzrec/prompt/assembler.py builds the packed varlen token stream in the dataloader worker
tzrec/prompt/persist.py writes the contract beside the weights, checks it on restore
tzrec/models/prompt_generative_qwen.py the model: inputs_embeds forward, index_copy scatter, band-constrained decode
tzrec/modules/prompt_projection.py slot width -> LM hidden size
docs/source/models/prompt_generative_qwen.md operator manual (Chinese)

Removed: sid_feature.py, generative_model.py, generative_qwen.py,
generative_model.proto and their tests. Nothing is kept for compatibility.

Two properties worth reviewing specifically:

  • The compiler resolves no dimension. No artifact stores in_dim/out_dim;
    the model resolves both at __init__ from group_total_dim and the backbone
    config, so one compiled plan is valid across model sizes.
  • No shape is discovered on the device. max_seqlen is computed by the
    collator on the host and carried in additional_infos; the model never calls
    lengths.max().

Test Plan

Unit: 55 tests across compile, assemble, projection, persistence and the model.
pre-commit run -a and python scripts/pyre_check.py both 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:

step result
tzrec.train_eval 3 steps, ce_loss 4.79 against ln(128) = 4.85 — near-random from a fresh init, as expected
eval ce_loss computed and written to train_eval_result_v2.txt
checkpoint model.ckpt-N/prompt/ carries sid_space.json, prompt_plan.json, prompt_hashes.json, tokenizer/
--continue_train restores and continues
hash guard changing the codebook to 8,8,8 is refused at restore
tzrec.predict 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
tzrec.export (HF) loads with AutoModelForCausalLM.from_pretrained; export dir carries the prompt contract

Eight defects were found by running the pipeline rather than by unit tests —
four missed _create_model call sites, a (total, value_dim) shape the
assembler mis-read, the absent loss/metric hooks, an FX Proxy in decode, and
the un-derived feature groups a projected slot needs. All are fixed and covered.

Not done

  • Multi-rank is untested. Everything above is --nproc-per-node=1, so
    sharding, DMP and the planner have not been exercised. The projected config is
    the one with a sparse table to shard.
  • The scripted serving front end and cache_ids (§11 of the design). Without
    cache_ids, a stock radix prefix cache keyed on token ids will reuse KV across
    users at a sentinel position — a correctness bug, so it must land before any
    projected slot is served.
  • TorchScript export is refused with a message pointing at 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, but TokenizeFeature extends
IdFeature and is therefore always sparse, hence always PROJECTED. No existing
type can be an INLINE text slot. I implemented INLINE as unconditionally meaning
SID. Either the branch goes, or TokenizeFeature gains a no-embedding mode.

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 30 commits July 29, 2026 03:56
_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>
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