[feat] fill prompt slots inline from id and tokenize features

[tiankongdeguiji]

A series of commits on the genrec prompt path, merged with master; the automated review's findings are addressed in the last four (inline width in positions and id-space rejection, checkpoint shape check in dcp_to_hf, fixed codes for non-finite floats in hole_keys, doc accuracy) and answered inline.

1. INLINE prompt slots from id and tokenize features

RFC 0001 (§3.1, §3.2, §5.2) designed two INLINE slot kinds the implementation never delivered. The compiler gated INLINE on has_embedding, which is True for every sparse feature, and the assembler applied one global id_shift = base_vocab_size — so the only working INLINE feature was a sequence_raw_feature of flat offset codes.

Now a slot is INLINE when its one sequence member declares no embedding, resolved per feature kind:

• tokenize_feature without embedding_dim — FG's word ids are ids of the prompt tokenizer, so they enter input_ids unshifted (id_shift = 0).
• sequence_id_feature without embedding_dim, including a grouped sequence_feature { id_feature } sub-feature — a multi-value sequence with value_dim = num_levels, one offset SID code (level_offsets[l] + code) per level per item, shifted by base_vocab_size. The compiler rejects any other value_dim.
• sequence_raw_feature with no dense embedding — unchanged legacy path, so existing pipelines keep training on their flat columns.

SlotSeg carries id_shift; the assembler applies it per segment. IdFeature.embedding_dim, TokenizeFeature.embedding_dim and TokenizeFeature.vocab_file become optional. has_embedding is deliberately untouched (RFC §3.2): nothing reaches emb_config for a feature in no group, so the existing assert stays as the loud failure if that invariant is ever broken.

Vocabulary guarantee. A text slot must tokenize with the vocabulary the LM was extended from. load_pipeline_config fills vocab_file from prompt_config.tokenizer_path on every tokenize feature that omits it (top-level, sequence_tokenize_feature, and grouped) — on load rather than in the compiler, because FG_NORMAL builds its pyfg handler inside BaseFeature.__init__, and because that single choke point also covers every tool that calls _create_features. At compile, a user-set vocab_file must be byte-identical (md5) to the prompt tokenizer; a path comparison would not survive export, which rewrites the asset to tok_<md5>.json. A TokenizeFeature that reaches construction with no vocab_file at all (no prompt_config to fill it from) is rejected right there, with a message naming both remedies, rather than being tolerated until FG fails.

What this does not do, by decision: no warning or docs about tokenization seams (a text slot next to template text or another text slot is an inexact splice — libfg has no string-join op, so the exact path is an upstream concat into one column), and no constant prefix/suffix folding (possible via a regex_replace stub under FG_DAG; separate follow-up).

Out of repo and unchanged: the C++ processor's dist-emb path drops ungrouped sparse features (RFC 0002 §5.7); real serving of these slots needs its pass-through list.

2. Compose genrec HF assets at export instead of in checkpoints

Genrec checkpoints carried config.json, generation_config.json and hf_export_meta.json only so dcp_to_hf could read them back at export — but by then the backbone is already live. Export reads the backbone config and generation config off the LM one line before del model.model.lm and hands them to export_hf_assets; dcp_to_hf finds the backbone's checkpoint keys by suffix and requires every match to share one prefix, so no wrapper layout is assumed and a look-alike key elsewhere cannot stand in. Checkpoints keep only weights; an HF-asset failure can no longer abort a checkpoint save. CompiledPrompt now carries the tokenizer compile_prompt extended with the SID atoms, so export writes it via save_tokenizer_dir from the prompt the model was built on instead of compiling a second time, and compile_prompt writes no files at all — nothing at runtime reads the field (the front-end is FX-traced before scripting, and the dataset keeps only the assembler). Export directory contents are unchanged. The deleted test_dist_checkpoint_manager_propagates_hf_asset_failure covered the cross-rank reconcile of a rank-0-only write that no longer exists.

3. hole_keys: dense members could not be scripted

HoleKeyBuilder's dense branch viewed float32 values as int32 to fold their bits. TorchScript only knows aten::view(Tensor, SymInt[] size) — view(torch.int32) compiles to view([3]) (ScalarType int32 is 3) and aten::view.dtype is not scriptable in any spelling — so the exported front-end failed on the first prompt with a raw feature in a projected slot. Eager parity and int-only scripting tests hid it. The fold now uses torch.frexp's exponent and 24-bit mantissa, which identify a float32 exactly and script cleanly; key values for float members change (nothing reads them back — only eager/script parity and cross-device determinism matter), and a float member joins the scripting test.

4. Widen the genrec integration fixture

The mock config is now self-contained and shaped like production prompts: two DEEP slots (profile = 3 scalar ids, two sharing an embedding_name table; ctx = an id plus a raw_feature vector) sharing one projection_name; one grouped sequence_feature whose id, multi-value id (value_dim: 0, mean-pooled) and projected tokenize_feature members feed an mlp-bodied, bias-free projection; the grouped SID history; the inline text slot. The test builds every request column from the data (dense members carry (B, value_dim) float32 and no lengths), asserts HOLE_SLOT_COUNTS and the shared projection module, pins the distributed export's dense_meta.json contract for the user-tiled EBC group and the three EmbeddingCollections, and runs predict_checkpoint to assert decode emits num_return_sequences × num_levels local codes within the codebook.

One FG rule surfaced on the way and is baked into the fixture: a grouped sub-feature's name must differ from its input's base name, because FG prefixes both with the group name and rejects a feature that is its own input.

Test Plan

• tzrec.prompt.compile_test (28), tzrec.prompt.assembler_test (11), tzrec.prompt.hole_keys_test (14), tzrec.utils.config_util_test (11), tzrec.features.tokenize_feature_test (15), tzrec.features.feature_test (46), tzrec.models.genrec_model_test (15), tzrec.models.genrec_causal_lm_model_test (8), tzrec.utils.hf_export_util_test (8), tzrec.utils.checkpoint_util_test (77) — all pass. New cases cover each INLINE kind and its id_shift, the grouped sub-feature, the value_dim and vocabulary rejections, the byte-identical-copy acceptance, scripting with mixed shifts, and vocab_file injection for every config shape.
• tzrec.tests.genrec_integration_test.test_genrec_train_eval_export — the mock config's SID history is now a grouped sequence_feature { id_feature { value_dim: 3 } } fed by list<list<int64>>, plus a {{title}} slot from a tokenize_feature with neither embedding_dim nor vocab_file. Trains, exports, reproduces the front-end's input_ids from raw request tensors, asserts the injected vocab_file in the trained config and the md5-matching FG asset beside fg.json.
• test_genrec_export_distributed_embedding (GPU, A10) — passes; dense_meta.json pins the projected slots' sparse-stage contract.
• tzrec.prompt.hole_keys_test (14, float member now scripted), tzrec.prompt.compile_test (28), tzrec.prompt.assembler_test (11), tzrec.models.genrec_model_test (15), tzrec.models.genrec_causal_lm_model_test (8) pass after the last two commits.
• test_genrec_train_eval_export and test_genrec_export_distributed_embedding (GPU, A10) pass with the widened fixture, including the decode step.
• pre-commit run and pyrefly check clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_011b1CqqpoqwBvBfwxK16UYJ

[github-actions]

Width of an INLINE multi-value SID slot undercounts positions (medium)

_slot_width returns the member's sequence_length, which caps items, but the new INLINE id-feature slot emits value_dim == num_levels codes per item (_inline_counts sums key_lengths), i.e. up to sequence_length × num_levels stream positions. The width unit is positions elsewhere — the response slot gets Width(STATIC, num_levels) (compile.py:418) and _validate phrases the gate as "the prompt can reach {max_total_length} positions".

Failure scenario: max_length: 64 plus an INLINE SID history with sequence_length: 16 and a 3-level codebook — the plan claims a ceiling of 16 for that slot while rows reach 48, so a prompt whose real proven length exceeds max_length passes the compile gate, and the bucket sizing max_total_length feeds to serving (see the warning at compile.py:566-569) is silently short by items × (num_levels − 1). Before this PR the invariant held because only 1-value-per-item raw sequences could be INLINE; the new tests pin the undercount (test_sid_id_feature_without_embedding_is_inline asserts 4 where 12 is reachable).

Suggested fix: for an INLINE slot whose single member declares value_dim > 1, bound by sequence_length × value_dim (or num_levels).

[tiankongdeguiji]

Fixed in 327878c. _slot_width now takes the fill mode: an INLINE slot is bounded by sequence_length × value_dim (every emitted token), and UNBOUNDED when an item may carry any number of values (value_dim: 0, i.e. sequence_tokenize_feature / grouped tokenize). PROJECTED sequence slots keep one position per item. The two tests that pinned the undercount now assert 12 and 48; the fixture's proven maximum goes 26 → 30.

[github-actions]

INLINE id feature may still declare its own id space (medium)

The IdFeature branch validates the shape contract (value_dim == num_levels) but not the value contract: dropping embedding_dim from a bucketized config like sequence_id_feature { num_buckets: 32 ... } makes it INLINE, and its bucket/vocab ids are shifted by base_vocab_size straight into the SID bands — valid-looking token ids, no error anywhere. It is reachable whenever value_dim happens to match num_levels (trivially so for a single-level codebook, since a sequence id feature defaults value_dim to 1), or via vocab_list/vocab_dict remapping.

Since _check_inline_member exists to "reject an INLINE member whose values cannot be LM token ids", consider also rejecting an INLINE id feature that declares any of hash_bucket_size / num_buckets / vocab_list / vocab_dict / zch / dynamicemb — the slot contract is pre-offset SID codes, which leave nothing for a vocabulary to do.

Minor, same function: when the paths differ, _file_md5 (line 159) opens member.vocab_file with a bare open(), so a misconfigured path surfaces as a raw FileNotFoundError traceback (a new compile-time read under FG_NONE, where the vocab was previously never touched) instead of the curated ValueError style used everywhere else in this file.

[tiankongdeguiji]

Both fixed in 327878c. An INLINE id feature that declares hash_bucket_size / num_buckets / zch / dynamicemb / vocab_list / vocab_dict / vocab_file is rejected at compile with a message naming the field and the remedy (drop it, or add embedding_dim to project it). An unreadable vocab_file on an inline text member now raises a ValueError naming the member and path instead of a bare FileNotFoundError.

[github-actions]

config.json can now disagree with the checkpoint's tensor shapes (medium)

The backbone config is read off the live LM — built from the current pipeline config and resized by the current sid_space.target_vocab_size — while the weights come from the checkpoint, mapped by name only. Nothing on this path compares shapes: the front-end restore never holds lm.* keys, and _derive_by_suffix matches names.

Failure scenario: tzrec.export --checkpoint_path <older run> (or any edit to sid_space.codebook / hf_model_name_or_path between train and export). Key names still match within the same architecture family, so the export completes and ships model.safetensors with embed_tokens.weight of shape [V_old, H] beside a config.json declaring vocab_size: V_new. The mismatch surfaces only when the serving engine calls from_pretrained — a late, confusing failure on a directory that looks complete. Before this PR, write_hf_assets wrote the config beside the checkpoint at save time, so the two could not diverge.

Cheap fix in the spirit of "refuse to write a partial model": reader.read_metadata().state_dict_metadata[ck].size already carries the checkpoint shapes — compare against empty.state_dict()[tk].shape (before del empty, line 59) and raise here on mismatch.

[tiankongdeguiji]

Fixed in 1796757. dcp_to_hf now keeps the DCP metadata and compares every mapped tensor's TensorStorageMetadata.size with the meta model's shape before writing; a mismatch raises with the first offenders as (key, config shape, checkpoint shape) and the same "refusing to write a partially-loaded HF model" wording. test_dcp_to_hf_refuses_a_shape_that_drifted grows vocab_size by 64 (identical key names) and expects the refusal.

[github-actions]

The distributed-embedding test lost its only numeric equivalence check (coverage regression)

The previous version of this test was a processor simulator: it looked a request's ids up in the exported sparse_embeddings-*.npz, fed the dense stage's scripted_model.pt, and compared slot_embeds (allclose), hole_keys (equal) and input_ids (equal) against the monolithic export. What remains is file existence plus an exact dense_meta.json name dump — which pins the contract but not the behavior: a wrong input-tiling or a mis-ordered collection feed in the dense stage now passes silently.

This matters more because this PR changed exactly this path's inputs: a dense score member with no sparse-stage entry, the input-tiled profile__ctx__ebc_user group, and multi-value tags__tag_b. I checked the alternatives — export_util_test.py feeds the dense stage random tensors and asserts shapes only, and genrec_model_test.py exercises the eager front-end — so nothing else in the suite numerically verifies the two-stage serving path against the monolithic one. Consider keeping the simulator (the widened fixture's _request(_MEMBERS) already builds every column it needs) even if the dense_meta.json pin stays.

[tiankongdeguiji]

Leaving this as is, deliberately. The simulator was removed at the author's request as too much machinery for what it proved; the dense score member and the frexp fold are exercised end to end by the CPU test's front-end parity check (which drives the scripted HoleKeyBuilder), and the distributed export's contract is pinned by the exact dense_meta.json assertion. Numeric two-stage-vs-monolithic parity is a fair follow-up if the dense stage's input handling changes again.

[github-actions]

The fold is no longer exact/device-independent for non-finite floats (low)

The frexp encoding is injective for every finite float32 (nice), but mantissa.to(torch.int64) for a non-finite mantissa is an out-of-range float→int cast, which is implementation-defined: on x86 CPU +inf, -inf and NaN all collapse to the same INT64_MIN code (so +inf and -inf collide into one hole key), while CUDA's cvt saturates +inf to INT64_MAX and -inf/NaN to INT64_MIN (so the same data folds differently per device). The old view(int32) path was an exact bijection on every device.

Impact is small — dense members carrying infinities are pathological — but it contradicts two guarantees this file states: the module docstring ("a key cannot depend on reduction order, device or batch split") and the new comment ("identify a float32 exactly"). Since a collision here means an engine prefix-cache hit on different inputs, consider either mapping non-finite values to fixed distinct codes before the cast (scriptable via torch.isinf/torch.isnan) or narrowing both claims to finite values. Related test gap: test_the_fold_is_bit_identical_across_devices feeds only int members, so the rewritten float branch is never compared CPU-vs-GPU anywhere (the GPU integration test's hole_keys comparison was dropped in this PR).

[tiankongdeguiji]

Fixed in 8925f5b. +inf, -inf and NaN now take three fixed codes above the finite range (isfinite/isnan/isneginf, all scriptable) instead of going through the out-of-range cast, so the fold is exact for every finite float32 and device-independent for the rest. test_the_fold_is_bit_identical_across_devices gained a DEEP float member carrying all three, compared CPU vs GPU.

[github-actions]

No test can catch a lossy float fold (medium)

Every dense-member assertion in this test is either identical-content stability or a permutation — and the permutation is already separated by the local position term, so it would pass even if the float encoding itself were collapsed. Truncating the mantissa (or dropping it entirely, folding only the exponent) passes this test, the scripting parity test below, and the integration HOLE_KEYS equality, because all three compare the new fold against itself. The int path pins exactly this property — test_different_content_folds_apart ("a changed input must not reuse the cached KV") — but the float path, which this PR rewrote from an exact bit-pattern view to exponent+mantissa, has no equivalent.

Since the encoding's whole justification is that it stays exact, one cheap assertion would pin it, e.g. keys([[1.0]]) != keys([[1.0 + 2**-23]]) (differing only in the last mantissa bit), or N distinct random floats folding to N distinct keys.

[tiankongdeguiji]

Fixed in 8925f5b. test_a_dense_member_folds_its_value now asserts 1.0 ≠ 1.0 + 2**-23 (last mantissa bit), 1.0 ≠ -1.0, and that 0.0, +inf, -inf, NaN fold to four distinct keys — a truncated or exponent-only fold fails all three.

[github-actions]

这两条新 bullet 与代码实际行为有几处偏差（低）：

1. L598 “配置了但内容与之不一致时，prompt 编译报错” 范围偏大：md5 一致性检查只在 _check_inline_member 中执行，而它仅对 INLINE slot 的成员触发（compile.py:406-407）。带 embedding_dim 的 tokenize feature（PROJECTED slot 或普通特征组）即使 vocab_file 内容与 prompt_config.tokenizer_path 不一致也不会报错。建议限定为“作为 inline slot 使用时”。
2. L600 未说明 INLINE 只对序列形态可达：_derive_slot_layout 要求 is_sequence and len(members) == 1，即需要 tokens_as_sequence: true、sequence_tokenize_feature 或 grouped 子特征。标量 tokenize_feature 省略 embedding_dim 后进特征组会命中 embedding_dim ... should be greater than 0 断言，而非文档描述的 inline 行为。
3. L598 “未配置时自动取” 有一个前提：仅当 pipeline 声明了 prompt_config 才会注入（config_util.py:194-195 提前返回）；否则特征构造时直接抛 ValueError（tokenize_feature.py:59-64）。补一句“未配置且无 prompt_config 时报错”会更完整。

[tiankongdeguiji]

三处都已按建议修正（cfb9f8ae）：vocab_file 一致性检查限定为「作为 inline slot 使用时」，并补充「未声明 prompt_config 时必须配置，否则创建特征时报错」；embedding_dim 可省略的说明补充了仅序列形态可 inline（tokens_as_sequence: true、sequence_tokenize_feature 或 sequence_feature 子特征）。IdFeature 一条同时补上了不能再声明 num_buckets/hash_bucket_size/vocab_*/zch/dynamicemb（对应 327878c 的校验）。

[github-actions]

Injection runs before --edit_config_json and is frozen by load→save round trips (low)

load_pipeline_config fills vocab_file at parse time, but every entry point applies edit_config_json after the load (e.g. main.py:781-783 in train, same shape in export/predict). Scenario: --edit_config_json '{"prompt_config.tokenizer_path": "/new/tok.json"}' on a config whose tokenize features omit vocab_file — every feature keeps the old injected path, a field the user never wrote. An inline text slot then dies with the md5 ValueError blaming vocab_file: [old path] (loud but confusing); an embedded tokenize feature silently keeps tokenizing with the old vocabulary. Tools that save the loaded config back (model_dir/pipeline.config, feature_selection, add_feature_info) freeze the injected value permanently, so the documented "omitted → follows prompt_config.tokenizer_path" link stays pinned to the day the tool ran even across an LM upgrade.

_fill_prompt_tokenizer is idempotent and cheap — re-running it after edit_config (or calling it from _create_features's callers post-edit) would close the first half.

[tiankongdeguiji]

Not changing this in this PR, by decision. Injection stays at load time (it has to precede feature construction for FG_NORMAL and cover every tool through load_pipeline_config); editing prompt_config.tokenizer_path afterwards with --edit_config_json is out of scope here. An inline text slot with a stale injected value fails loudly at compile — the md5 check names both paths — and the remedy is to set vocab_file explicitly or edit the config file.

[github-actions]

Code review summary

Reviewed the full change (22 files) across five areas — code quality, performance, test coverage, documentation accuracy, and security/multi-process safety — with each finding verified against the checkout before posting. All five reviews completed.

Overall this is a well-structured change: the per-segment id_shift wiring is TorchScript-clean, the checkpoint→export move of HF assets removes a per-save rank-0 write plus a collective (a net win for the training loop), the deleted unwrap_to/write_hf_assets/HF_EXPORT_META_FILENAME leave zero orphaned references, and the dcp_to_hf suffix+single-prefix derivation fails closed against look-alike keys. The new compile/config tests genuinely pin the behavior they claim (each INLINE kind and its shift, the md5 reject/accept pair, all four vocab-injection shapes).

Posted 7 inline comments. The three I'd weight highest:

1. compile.py _slot_width (medium) — an INLINE multi-value SID slot's width counts items, but the assembler emits items × num_levels positions, so the "proven ceiling" behind the max_length gate and serving bucket sizing is silently short by a factor of num_levels.
2. hf_export_util.py dcp_to_hf (medium) — config.json now comes from the live LM while weights come from the checkpoint mapped by name only; a codebook/backbone edit or an older --checkpoint_path ships an internally inconsistent export that fails late in the engine's from_pretrained. A shape check against the DCP metadata would keep the "refuse to write a partial model" guarantee.
3. genrec_integration_test.py (coverage regression) — the distributed-embedding test dropped its processor simulator, which was the only numeric equivalence check (slot_embeds/hole_keys/input_ids) between the two-stage and monolithic exports, in the same PR that changed that path's inputs.

Also flagged: a validation hole for INLINE id features that still declare num_buckets/vocab (silent id corruption in the SID bands), the frexp fold's non-finite edge (±inf collide on CPU; device-dependent casts contradict the module's device-independence claim) plus the missing discrimination test for float folds, the vocab_file injection running before --edit_config_json, and three small doc-accuracy gaps in feature.md.

🤖 Generated with Claude Code