Skip to content

fix(examples): anchor logit_trace chunk 0 with the BOS token - #2320

Merged
inureyes merged 2 commits into
mainfrom
fix/issue-1785-logit-trace-bos-anchor
Oct 15, 2026
Merged

inureyes merged 2 commits into
mainfrom
fix/issue-1785-logit-trace-bos-anchor

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

examples/logit_trace left chunk 0 without the BOS anchor (no BOS prepend, and no context pass under PREFILL), so on a model trained behind a document separator the first window collapsed and pooled perplexity or disagreement figures read like a defect.

  • Without PREFILL, chunk 0 gets the BOS prepended like every other chunk (target_offset = bos_prefix.len()).
  • With PREFILL, a chunk with no history gets a context pass of the BOS alone, so the traced forward keeps the chunk's width and target_offset 0. An empty bos_prefix skips it, so the data rows are unchanged for tokenizers without a BOS.
  • The corpus is still encoded without special tokens; adding them would shift every chunk boundary.
  • A # bos_anchor metadata line records the rule. compare_logit_traces.py reads only the model key, so parsing is unaffected.
  • scripts/laguna_oracle_trace.py plan_chunks and its docstring follow the same rule and write the same metadata line.
  • docs/benchmarks.md ("Judging a change that moves the numbers") says older traces are not comparable at chunk 0; the Laguna entry in docs/supported-models.md carries the re-measured figures.

Validation (GB10, real checkpoints)

Checkpoint poolside/Laguna-XS-2.1-NVFP4 (downloaded with hf download to ~/models/mlx/laguna-xs-2.1-nvfp4), corpus docs/architecture.md (4016 tokens). The float32 oracle runs on this aarch64 CUDA host (torch 2.14.1+cu130, transformers 5.19.0 in a scratch venv, --device cuda).

Arm chunk 0 ppl, width 256 chunk 1 ppl
oracle, before 364.3 22.25
oracle, after 19.93 22.25
mlxcel, before 378.2 22.16
mlxcel, after 19.86 22.16
  • compare_logit_traces.py oracle mlxcel after the fix: width 256 (256 2 8 0), 20/512 top-1 disagreements, 0/250 on decided positions; width 8 behind 512 (8 80 8 512), 27/640, 0/317 decided. Both exit 0, so (chunk, pos, target) agree on every row.
  • Only chunk 0 changes: every row of chunks 1+ is byte-identical between the pre-fix and fixed traces, in mlxcel at both shapes and in the oracle at width 256. Under PREFILL chunk 0 (8 rows at the corpus head) goes from perplexity 762279 to 229.3, against 228.2 in the oracle.
  • qwen2.5-0.5b-instruct-4bit (empty bos_prefix): # bos_anchor none, and the data rows are byte-identical before and after at 256 2 8 0 and 8 4 8 64.
  • cargo clippy --profile test-fast --features cuda --example logit_trace -- -D warnings and cargo fmt --all -- --check are clean. The scripts/ci contract checks pass.

Closes #1785

logit_trace encodes the corpus without special tokens and prepended the tokenizer's BOS to every chunk except chunk 0, and with PREFILL chunk 0 got no context pass at all. On a model trained behind a document separator that collapses the first window: Laguna XS 2.1 read perplexity about 950 there in the float32 reference, and pooled figures looked like a defect.

Every chunk now carries the anchor. Without PREFILL the BOS is prepended to chunk 0 as to the others; with PREFILL a chunk without history gets a context pass of the BOS alone, so the traced forward keeps the chunk's width. An empty BOS prefix leaves the data rows unchanged. A `# bos_anchor` metadata line records the rule, scripts/laguna_oracle_trace.py plans its chunks the same way, and docs/benchmarks.md and the Laguna entry in docs/supported-models.md note that older traces differ at chunk 0.

Validated on laguna-xs-2.1-nvfp4 against the float32 oracle on the GB10: chunk 0 at width 256 drops from 364 (reference) and 378 (mlxcel) to 19.9 in both, every later row is byte-identical to the pre-fix trace, and 0 of 250 (width 256) and 0 of 317 (width 8 behind 512) decided positions disagree. qwen2.5-0.5b (no BOS) traces are byte-identical before and after.

Closes #1785
@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:low Low priority area:benchmark Benchmark harness and performance measurement (bench_*.sh, /update-benchmarks) labels Oct 15, 2026
A tokenizer without a BOS wrote `# bos_anchor none chunks all`, which reads as if every chunk were anchored. Both logit_trace and laguna_oracle_trace.py now write `chunks none` in that case.
@inureyes inureyes added status:done Completed and removed status:review Under review labels Oct 15, 2026
@inureyes
inureyes merged commit a9f4ba4 into main Oct 15, 2026
27 checks passed
@inureyes
inureyes deleted the fix/issue-1785-logit-trace-bos-anchor branch October 15, 2026 18:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:benchmark Benchmark harness and performance measurement (bench_*.sh, /update-benchmarks) priority:low Low priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(examples): anchor logit_trace chunk 0 with the BOS token

1 participant