The Vocab Mapping Fix: Diagnosing and Remedying a Critical Mismatch in EAGLE-3 Fine-Tuning

Introduction

In the high-stakes world of speculative decoding for large language models, a single wrong assumption can derail hours of training and computation. Message 4991 captures a pivotal moment in a complex debugging session: the assistant reports the first successful results after identifying and fixing a fundamental data alignment issue that had rendered an entire fine-tuning run effectively random. What makes this message significant is not just the reported improvement—a loss drop from ~18-20 to ~8-9, and accuracy rising from ~2-7% to ~12-24%—but the diagnostic journey it culminates in, and the clear thinking process it reveals about how neural network components interact across model variants.

The Context: A Fine-Tuning Effort Gone Wrong

The broader session was an ambitious attempt to improve Kimi-K2.5 inference throughput using EAGLE-3 speculative decoding. The assistant had already trained a from-scratch EAGLE-3 drafter that achieved 74.7% validation accuracy, but the user wanted to explore whether fine-tuning an existing drafter—specifically the AQ-MedAI K2 EAGLE-3 drafter—could provide a better starting point. The AQ-MedAI drafter had been designed for the Kimi-K2 model (the predecessor to K2.5), and a quick probe (Phase 0) showed it achieved an accept_len of ~1.5 when deployed directly with K2.5, suggesting the representations were "correlated but not aligned" and that fine-tuning could bridge the gap.

Phase 1 began with a fine-tuning run using the existing training pipeline, loading AQ-MedAI's weights as a starting point with a conservative learning rate of 5e-5. The initial results were alarming: loss values of ~18-20 and accuracy of ~2-7%—essentially random performance. This was deeply puzzling because the same drafter had demonstrated meaningful predictive power (accept_len ~1.5) when deployed in SGLang. The user's sharp question in message 4975—"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?"—triggered a systematic investigation.

The Diagnostic Breakthrough

The assistant's response to the user's question reveals a methodical debugging approach. Rather than guessing, the assistant enumerated three possible causes:

  1. Vocab mapping mismatch: The training data's target labels might be computed using a different draft-to-target token mapping than what AQ-MedAI's lm_head was trained on.
  2. Hidden state format differences: The training data pipeline might feed hidden states differently than SGLang does at inference.
  3. Verifier lm_head interaction: The loss computation uses the verifier's lm_head with a specific mask that might not align with AQ-MedAI's trained weights. The assistant immediately identified the first cause as the most likely and wrote a diagnostic script to compare the vocab mappings. The results were stark: only 252 out of 32,000 draft vocab positions matched between AQ-MedAI's mapping and K2.5's mapping. The first ~50 positions overlapped (likely special tokens and byte-level tokens), but after that the orderings diverged completely. While 25,750 of AQ-MedAI's draft tokens existed somewhere in K2.5's draft vocab, they were at different positions—meaning that when the fine-tuning pipeline computed the loss, it was comparing the lm_head's output at position 500 against a label that meant something entirely different from what the lm_head had been trained to predict at that position. This is the kind of subtle bug that can waste enormous amounts of compute time. The training loop was not broken—it was faithfully computing cross-entropy loss—but the labels were scrambled relative to the model's output space. The model was predicting reasonable tokens for its own vocab ordering, but the loss function was comparing those predictions against targets that belonged to a different ordering. The result was a loss value indistinguishable from random chance.

The Fix: Remapping Through Permutation

The assistant considered two options:

The Results: Message 4991

Message 4991 reports the outcome of that fix. The assistant writes:

Massive improvement! Initial metrics: - loss_0: ~8.3-9.3 (was ~18-20 before the fix) - cond_acc_0: ~12-24% (was ~2-7% before)

>

This confirms the vocab mapping was the root cause. The loss is still high because the K2 drafter needs to adapt to K2.5 hidden states, but it's starting from a meaningful place now instead of random.

The message then includes a bash command that waits 120 seconds and checks the latest metrics, showing loss_0 at 8.88 and accuracy at 16.27% as the learning rate warms up.

Several aspects of this message reveal the assistant's thinking:

First, the framing as "Massive improvement!" is telling. The assistant is not just reporting numbers—it's interpreting them in context. A drop from ~18-20 to ~8-9 for cross-entropy loss over a 32,000-class vocabulary is indeed dramatic. Random guessing on a uniform distribution over 32,000 classes would yield a loss of ln(32000) ≈ 10.37. So the pre-fix loss of ~18-20 was actually worse than random—the model was confidently wrong because the labels were systematically scrambled. The post-fix loss of ~8-9 is below random chance, meaning the model is now learning meaningful patterns.

Second, the assistant's interpretation of the remaining gap. The loss is still high compared to the from-scratch model's final state (~0.9), but the assistant correctly attributes this to the need for the K2 drafter to adapt to K2.5's hidden state representations. The lm_head remapping fixed the label alignment, but the drafter's internal layers (the midlayer transformer block and fc layer) were trained on K2's hidden states and need to adapt to K2.5's different representation space. This is exactly what fine-tuning is supposed to accomplish.

Third, the follow-up check after 120 seconds. The assistant doesn't just report the initial improvement and move on—it waits to see how the training progresses as the learning rate warms up. The metrics at 21:28:59 show loss_0 at 8.88 (slightly improved from the initial ~9.2) and accuracy at 16.27% (improved from ~12%). The learning rate has risen from 0 to 3.82e-6 (still well below the target 5e-5). This shows the assistant is monitoring the training dynamics to ensure the fix is working correctly and not introducing new issues.

Deeper Analysis: What the Numbers Tell Us

The structure of the EAGLE-3 loss is worth understanding. The metrics show loss_0, loss_1, ..., loss_4 corresponding to five speculative steps (ttt-steps=5). loss_0 is the first draft token prediction, which has the richest conditioning (the full hidden state from the target model). loss_1 through loss_4 are subsequent steps that depend on earlier draft tokens, making them progressively harder. The fact that loss_0 drops from ~18 to ~8-9 while loss_1 only drops from ~19.6 to ~13-14 is expected—the first step benefits most from the vocab remapping because it has the strongest signal from the hidden state input.

The accuracy numbers tell a similar story. cond_acc_0 (conditional accuracy for the first step) jumps from ~2-7% to ~12-24%, while cond_acc_1 goes from 0% to ~0.15-1.5%. The second step remains near-random because it depends on the first draft token, and with only ~16% accuracy on step 0, the conditioning for step 1 is often wrong. This is a known property of autoregressive speculative decoding: errors compound, and the later steps are only as good as the earlier ones.

The Broader Significance

This message represents a turning point in the session. Before this fix, the fine-tuning approach was producing garbage and the entire strategy of leveraging AQ-MedAI's pre-trained weights was in doubt. After the fix, the numbers confirm that the approach is viable—the K2 drafter can be adapted to K2.5, it just needs the correct vocab alignment.

The message also demonstrates a crucial principle in ML engineering: when a model performs well at inference but poorly in training, the data pipeline is almost always the culprit. The assistant's systematic debugging—enumerating hypotheses, testing the most likely first, and applying a targeted fix—is a textbook example of how to resolve such discrepancies. The vocab mapping mismatch is a particularly insidious bug because it doesn't crash anything or produce obvious errors; it just silently produces random loss values that could easily be misinterpreted as "the model isn't learning."

Conclusion

Message 4991 is a moment of validation in a complex debugging journey. It confirms that the diagnostic work was correct, the fix was appropriate, and the fine-tuning approach is now on solid ground. The assistant's clear reporting—framing the improvement, contextualizing the remaining gap, and monitoring the training dynamics—provides the user with exactly the information needed to make decisions about next steps. While the ultimate fine-tuning results would later plateau at ~38% accuracy (leading to the approach being abandoned), this message represents the critical inflection point where the problem shifted from "why is this broken?" to "how well can we make this work?"—a fundamentally more productive and optimistic place to be.