Remapping the Unmappable: Fixing a Vocab Mismatch in EAGLE-3 Fine-Tuning

In the high-stakes world of speculative decoding for large language models, even a single incorrect assumption can cascade into hours of wasted computation. This message — message 4982 in a long-running opencode session — represents the precise moment when an assistant diagnosed and planned the fix for a subtle but devastating bug: a vocabulary mapping mismatch that had rendered an entire fine-tuning run effectively random. The message is deceptively simple — a four-point plan and a file read — but it encapsulates the critical transition from confusion to clarity, from wasted training to a targeted fix.

The Context: A Promising but Broken Fine-Tuning Run

The session leading up to this message had been an ambitious attempt to improve inference throughput for the Kimi-K2.5 model using EAGLE-3 speculative decoding. The team had already trained a from-scratch EAGLE-3 drafter on 100K samples, achieving 74.7% validation accuracy and a respectable 60 tok/s. But they wanted more. The AQ-MedAI K2 drafter — a publicly available EAGLE-3 model trained for the earlier Kimi-K2 architecture — seemed like a promising starting point for fine-tuning. If its pre-trained representations could be adapted to K2.5, it might outperform the from-scratch model.

Phase 0 had been encouraging: a direct probe of the AQ-MedAI K2 drafter on K2.5 achieved an accept_len of ~1.5 and 52 tok/s. This was worse than the from-scratch drafter (accept_len ~2.0, 60 tok/s) but significantly above random (accept_len = 1.0). The representations were correlated, just not aligned. Fine-tuning should bridge the gap.

Phase 1 launched with high confidence. The assistant killed the inference server, freed the GPUs, copied the training script, and launched a fine-tuning run with a conservative learning rate (5e-5) and 5 epochs. The initial loss, however, was catastrophic: ~18-20 for step 0, compared to ~0.9 for the from-scratch model's final state. Accuracy hovered at 2-7% — essentially random. The user's question in message 4975 cut to the heart of the mystery: "Why is the loss/accuracy basically zero if the model was predicting ok in sglang? Are we passing it correct layers/hidden states etc correctly?"

The Diagnostic Breakthrough

The assistant's investigation in messages 4976-4981 was a masterclass in systematic debugging. It considered three possible causes: (1) a vocab mapping mismatch, (2) a hidden state format discrepancy, or (3) a verifier_lm_head issue. It started with the most likely culprit — the vocab mapping — and wrote a diagnostic script to compare the AQ-MedAI and K2.5 mappings.

The results were stark. Both mappings covered 32,000 target tokens, and 25,750 of those target tokens appeared in both mappings. But the ordering was almost completely different: only 252 out of 32,000 draft-to-target token positions matched between the two mappings. The first ~50 positions happened to align (likely special tokens and byte-level tokens), but after that the orderings diverged entirely.

This explained everything. The AQ-MedAI drafter's lm_head — the final linear layer that projects hidden states to vocabulary logits — had been trained to output logits according to AQ-MedAI's draft-to-target mapping. At position 500, the lm_head predicted "token X" according to AQ-MedAI's mapping. But the training labels were computed using the K2.5 mapping, where position 500 meant "token Y." The loss function was comparing apples to oranges — or more precisely, comparing apples to whatever random fruit happened to be at the same index in a different basket.

The drafter had been predicting above chance in SGLang because SGLang used AQ-MedAI's own vocab mapping for inference. The training pipeline, however, used a different mapping, creating a fatal mismatch between the model's output distribution and the label distribution.## The Message: A Surgical Plan for a Surgical Fix

Message 4982 is the assistant's response after discovering the root cause. It opens with a clear statement of intent: "Now let me update the --finetune-from logic in 04_train.py to remap the lm_head weights." This is not a moment of discovery — that happened in the previous message. This is the moment of action, when understanding crystallizes into a concrete plan.

The four-point approach is worth examining in detail:

Point 1: "Load AQ-MedAI's d2t to get their target token IDs for each draft position." The d2t (draft-to-target) mapping is a tensor that, for each position in the draft vocabulary, gives the offset to the corresponding target token ID. Specifically, target_id = d2t[pos] + pos. This is the key data structure that defines the relationship between the draft model's output space and the base model's vocabulary. Loading AQ-MedAI's d2t allows the assistant to reconstruct the mapping that the AQ-MedAI lm_head was trained against.

Point 2: "For each AQ-MedAI draft position, find the corresponding position in our draft vocab." This is the core remapping logic. The assistant needs to build a permutation: for each row in AQ-MedAI's lm_head (indexed by draft position), find which row in the K2.5 lm_head corresponds to the same target token. If AQ-MedAI's draft position 500 maps to target token 17,423, and K2.5's draft position 742 also maps to target token 17,423, then row 500 of AQ-MedAI's lm_head should become row 742 of the fine-tuned model.

Point 3: "Permute the lm_head rows accordingly." This is the mechanical step — applying the permutation matrix to reorder the lm_head weights. In practice, this means creating a new tensor where new_lm_head[our_draft_id] = old_lm_head[aq_draft_id] for each overlapping target token.

Point 4: "For AQ-MedAI draft tokens that don't exist in our vocab (6,250 tokens), initialize those rows randomly or from the verifier lm_head." This handles the 6,250 AQ-MedAI draft tokens that map to target tokens not covered by the K2.5 draft vocabulary. Those rows have no corresponding weight in the AQ-MedAI checkpoint and must be initialized from scratch. The assistant suggests two options: random initialization (which would start with high loss but learn from scratch) or copying from the verifier's lm_head (which might provide a better starting point since it's already aligned with K2.5's vocabulary).

The Assumptions Embedded in the Plan

This message makes several implicit assumptions that are worth examining. First, it assumes that the lm_head remapping is sufficient — that the hidden state representations learned by the AQ-MedAI drafter's transformer layers are compatible with K2.5's hidden states, and only the output projection needs correction. This is a reasonable assumption given the Phase 0 results (accept_len ~1.5), which demonstrated that the drafter's hidden states were already correlated with K2.5's distribution. The lm_head was the final bottleneck.

Second, it assumes that the remapping is deterministic — that each AQ-MedAI draft position maps to exactly one K2.5 draft position, or to none. The diagnostic showed 25,750 overlapping target tokens and 6,250 missing, which means the mapping is straightforward: 25,750 rows get permuted, and 6,250 rows get new initialization.

Third, it assumes that the remaining weights — the transformer layers (midlayer.* → layers.0.*), embed_tokens, and verifier_lm_head — do not need remapping. The existing --finetune-from logic already handled those by skipping them (they're loaded from the verifier checkpoint, not the drafter). This is correct because those weights are either tied to the base model's vocabulary (embed_tokens, verifier_lm_head) or are architecture weights that don't depend on vocab ordering.

What Knowledge Was Required to Understand This Message

To fully grasp this message, one needs a solid understanding of the EAGLE-3 architecture and its vocab mapping system. EAGLE-3 uses a compressed draft vocabulary: instead of predicting over the full base vocabulary (163,840 tokens for Kimi-K2.5), it predicts over a smaller "draft vocabulary" of the most frequent tokens (32,000 in this case). The t2d (target-to-draft) mapping is a boolean mask indicating which target tokens are in the draft vocabulary, and the d2t (draft-to-target) mapping gives the offset to reconstruct the target token ID from a draft position.

One also needs to understand the fine-tuning pipeline: the --finetune-from flag loads weights from a pretrained drafter checkpoint, but the training labels are computed using the current vocab mapping (K2.5's). If the pretrained checkpoint used a different mapping, the lm_head weights are misaligned with the labels.

The message also assumes familiarity with the codebase structure — specifically, that the remapping logic belongs in 04_train.py around line 320, where the --finetune-from logic already exists.## The Thinking Process: From Symptom to Root Cause

The reasoning visible in this message and its immediate predecessors reveals a sophisticated debugging methodology. The assistant did not jump to conclusions or apply random fixes. Instead, it followed a clear chain of inference:

1. Observe the symptom: Loss ~18-20, accuracy ~2-7% — essentially random performance.
2. Formulate the contradiction: The same drafter achieved accept_len ~1.5 in SGLang, proving it was predicting above chance. Why would training show random performance?
3. Generate hypotheses: Three possible causes — vocab mapping mismatch, hidden state format discrepancy, or verifier_lm_head issue.
4. Test the most likely hypothesis first: Write a diagnostic script comparing the AQ-MedAI and K2.5 vocab mappings.
5. Confirm the root cause: Only 252/32,000 positions match between the two mappings.
6. Design the fix: Remap the lm_head rows using a permutation derived from the target token overlap. This is textbook systematic debugging. The assistant resisted the temptation to tweak hyperparameters or blame the data, instead focusing on the structural mismatch between the pretrained weights and the training pipeline. The key insight — that the lm_head is a vocabulary-ordered projection and must be permuted when the vocab ordering changes — is exactly the kind of subtle bug that can waste days of training time.

The Broader Significance: Transfer Learning Across Vocabularies

This message touches on a fundamental challenge in transfer learning for language models: vocabulary alignment. When fine-tuning a pretrained model, one typically assumes the vocabulary is identical. But EAGLE-3's draft vocabulary is a learned compression of the base vocabulary — different training runs can produce different draft vocabularies even for the same base model. The AQ-MedAI K2 drafter was trained on Kimi-K2, which likely had a different token frequency distribution than Kimi-K2.5, leading to a different draft vocabulary ordering.

This is a general problem that extends beyond EAGLE-3. Any time a model has a learned projection layer tied to a specific vocabulary ordering — whether it's an lm_head, an embedding layer, or a classification head — transferring weights requires careful remapping. The naive approach of loading weights by name (which is what most frameworks do by default) will silently produce wrong results if the vocabularies differ.

The fix planned in this message is elegant because it exploits the structure of the problem: instead of trying to learn the remapping, it computes it directly from the two vocab mappings. The permutation is deterministic, exact for the 25,750 overlapping tokens, and the remaining 6,250 rows are initialized from scratch. This is far more efficient than trying to learn the remapping through gradient descent, which would require the model to unlearn its vocabulary ordering and relearn a new one — a process that could take many epochs or fail entirely.

What Came Next

The story does not end with this fix. After applying the lm_head remapping, the fine-tuning loss dropped to ~9 and accuracy improved to ~24% — a dramatic improvement from random, but still far from the from-scratch model's 75%. The fine-tuning plateaued at ~38% accuracy, converging slower than expected. The assistant eventually concluded that the K2 weights were a poor initialization for K2.5 and abandoned the approach.

But the diagnostic work was not wasted. The vocab remapping fix became a permanent part of the training pipeline, and the deeper lesson — that vocabulary alignment is critical for transfer learning — informed subsequent decisions. The assistant pivoted to system-level optimization of the verify step, addressing the fundamental PCIe communication bottleneck that limited all speculative decoding approaches on this hardware.

Conclusion

Message 4982 captures a pivotal moment in a complex optimization journey. It is the point where confusion gives way to understanding, where a mysterious failure is traced to a precise structural cause, and where a surgical fix is designed to replace a blunt instrument. The message itself is concise — a four-point plan and a file read — but it represents hours of diagnostic work and embodies a debugging methodology worth studying.

The vocab mapping mismatch bug is a cautionary tale about the hidden assumptions in transfer learning. When loading pretrained weights, it is not enough to check that the architecture matches; one must also verify that the vocabulary ordering matches. A mismatch of 252 out of 32,000 positions — less than 1% — was enough to render the entire fine-tuning run effectively random. The fix, once understood, was straightforward: remap the lm_head rows to align the vocabularies. But finding that fix required asking the right question at the right time: "Why is the loss/accuracy basically zero if the model was predicting ok in sglang?"