Repository navigation
fix(examples): anchor logit_trace chunk 0 with the BOS token #1785
Description
Found during PR #1779 (the Laguna reference comparison).
Problem
examples/logit_trace.rs encodes the corpus without special tokens and prepends the tokenizer's BOS to every chunk except chunk 0; with PREFILL > 0, chunk 0 also gets no context pass. Chunk 0 therefore starts with no anchor at every width. On a model trained behind a document separator this collapses the first chunk: on Laguna XS 2.1 at width 256, chunk 0 reads perplexity 955 in the float32 reference and 999 in mlxcel, against about 17.8 in a BOS-anchored chunk, and it held 42 of the 60 top-1 disagreements (PR #1779). Both arms mirror the gap, so A/B comparisons stay valid, but a pooled perplexity or disagreement figure reads like a defect.
examples/perplexity.rs has the same c == 0 skip (:87) but encodes with encode(&text, true) (:50-59), so its chunk 0 already starts at the tokenizer's BOS. It needs no change.
Evidence
examples/logit_trace.rs:191-192(encode(text.as_str(), false)),:200-207(the "Same BOS anchoring asexamples/perplexity" comment andbos_prefix),:223-233(if prefill > 0 { ... } else if c == 0 || bos_prefix.is_empty()),:245-256(context pass only whenctx_from < start, never true for chunk 0).scripts/laguna_oracle_trace.pymirrors it on purpose: docstring:44-49,plan_chunks:184-205.
Proposed fix
- Without
PREFILL: prependbos_prefixto chunk 0 as for later chunks, withtarget_offset = bos_prefix.len(). - With
PREFILL: give chunk 0 a context pass ofbos_prefixalone, so the traced forward keeps the chunk's width andtarget_offset0. Skip it whenbos_prefixis empty. - Keep
encode(text, false): encoding with special tokens would shift every chunk boundary, not just chunk 0. - Add a
#metadata line recording the anchoring.scripts/compare_logit_traces.pyreads only themodelmetadata key (:43-63,:107-108), so parsing is unaffected. - Update
plan_chunksand its docstring inscripts/laguna_oracle_trace.pyin the same PR. - Note in
docs/benchmarks.md(:407, "Judging a change that moves the numbers") that older traces are not comparable at chunk 0, and refresh the unanchored-window figures in the Laguna entry ofdocs/supported-models.md:55.
Acceptance criteria
- Chunk 0 carries the BOS at every width, with and without
PREFILL, wheneverbos_prefixis non-empty; with an emptybos_prefixthe data rows are unchanged. - Laguna XS 2.1 at width 256: chunk-0 perplexity drops into the BOS-anchored range (about 17.8 in PR test(laguna): add a checkpoint-code reference trace and forward tests #1779).
- mlxcel and
laguna_oracle_trace.pystill agree on every(chunk, pos, target).
Verification
cargo build --release --features metal,accelerate --example logit_trace
./target/release/examples/logit_trace models/mlx/laguna-xs-2.1-nvfp4 CORPUS.txt 256 2 8 0 > mlxcel.tsv
python3 scripts/laguna_oracle_trace.py models/mlx/laguna-xs-2.1-nvfp4 CORPUS.txt 256 2 8 0 > oracle.tsv
python3 scripts/compare_logit_traces.py oracle.tsv mlxcel.tsvRepeat both traces with 8 80 8 512 (width 8 behind 512 tokens) for the prefill path. Clippy and cargo fmt --all -- --check stay clean.
Activity