Repository navigation
fix(examples): anchor logit_trace chunk 0 with the BOS token - #2320
Merged
Merged
Conversation
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
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.
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.
Summary
examples/logit_traceleft chunk 0 without the BOS anchor (no BOS prepend, and no context pass underPREFILL), so on a model trained behind a document separator the first window collapsed and pooled perplexity or disagreement figures read like a defect.PREFILL, chunk 0 gets the BOS prepended like every other chunk (target_offset = bos_prefix.len()).PREFILL, a chunk with no history gets a context pass of the BOS alone, so the traced forward keeps the chunk's width andtarget_offset0. An emptybos_prefixskips it, so the data rows are unchanged for tokenizers without a BOS.# bos_anchormetadata line records the rule.compare_logit_traces.pyreads only themodelkey, so parsing is unaffected.scripts/laguna_oracle_trace.pyplan_chunksand 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 indocs/supported-models.mdcarries the re-measured figures.Validation (GB10, real checkpoints)
Checkpoint
poolside/Laguna-XS-2.1-NVFP4(downloaded withhf downloadto~/models/mlx/laguna-xs-2.1-nvfp4), corpusdocs/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).compare_logit_traces.py oracle mlxcelafter 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.PREFILLchunk 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(emptybos_prefix):# bos_anchor none, and the data rows are byte-identical before and after at256 2 8 0and8 4 8 64.cargo clippy --profile test-fast --features cuda --example logit_trace -- -D warningsandcargo fmt --all -- --checkare clean. Thescripts/cicontract checks pass.Closes #1785