The Vocab Mapping That Broke Fine-Tuning: A Diagnostic Pivot in EAGLE-3 Drafter Adaptation

When a machine learning experiment produces results that are mathematically impossible given what you already know, the universe is trying to tell you something. In the middle of a sprawling effort to accelerate Kimi-K2.5 inference through speculative decoding, one message — message 4973 — captured exactly this kind of contradiction. The assistant checked on a fine-tuning run and found loss values of ~18-20 and accuracies near zero, even though the exact same model weights had just demonstrated meaningful predictive power (accept_len ~1.5) when deployed in SGLang. This discrepancy couldn't be explained by bad initialization or insufficient training — it pointed to a fundamental structural mismatch in how the training pipeline was feeding data to the model.

The Message: Raw Output from a Troubled Training Run

The message itself is deceptively simple. The assistant writes:

It's loading. Let me wait a bit more for it to finish model setup and start training. [bash] sleep 60 && ssh root@10.1.230.174 'tail -40 /data/eagle3/synth_100k/logs/train_finetune_aqmedai.log' 2>&1 21:18:08 [speculators.metrics] {'train': {'loss_0': 17.80253028869629, 'full_acc_0': 0.06492192298173904, 'cond_acc_0': 0.06492192298173904, 'loss_1': 19.588577270507812, 'full_acc_1': 0.0, 'cond_acc_1': 0.0, 'loss_2': 20.24933433532715, 'full_acc_2': 0.0, 'cond_acc_2': 0.0, 'loss_3': 20.489913940429688, 'full_acc_3': 0.0, 'cond_acc_3': 0.0, 'loss_4': 20.573558807373047, 'full_acc_4': 0.0, 'cond_acc_4': 0.0, 'loss': 98.70391082763672}, 'epoch': 0, 'lr': 2.384737678855326e-07}

The training metrics show five prediction steps (loss_0 through loss_4) and a total loss of 98.7. The conditional accuracy for steps 1 through 4 is exactly 0.0 — the model is not predicting a single token correctly beyond the first step. Step 0 has a conditional accuracy of ~6.5%, which is barely above random chance for a 32,000-class vocabulary. The learning rate is still in its warmup phase at 2.38e-07, so the model hasn't really started learning yet, but these loss values are catastrophically high. To put this in perspective, the from-scratch EAGLE-3 drafter that was trained earlier in this session achieved a final loss around 0.9. A loss of 18-20 means the model is essentially producing uniform random predictions across the vocabulary.

The Contradiction That Drove the Investigation

What makes this message so significant is the context it sits within. Just minutes earlier, in [msg 4966], the assistant had deployed this exact AQ-MedAI K2 drafter in SGLang and measured an accept_len of approximately 1.5 — meaning the drafter was successfully predicting about 1.5 additional tokens per step beyond the guaranteed first token. The benchmark showed 52 tok/s throughput, which, while slower than the from-scratch drafter's 60 tok/s, was still well above the random baseline of accept_len ~1.0. The drafter was clearly doing something useful when running inside SGLang.

But now, in the training pipeline, the same weights were producing random-level loss. This is the kind of contradiction that forces a re-examination of every assumption in the pipeline. The assistant's immediate reaction — documented in the following message [msg 4974] — was to note that the high loss was "expected at the start" due to misalignment, but the user in [msg 4975] immediately pushed back with the right question: "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?"

This question was the catalyst for one of the most important debugging sequences in the entire session.## The Assumptions That Led to the Contradiction

To understand why this message was written, we need to trace the chain of assumptions that led to the fine-tuning attempt. The session had been pursuing a pragmatic strategy: rather than training an EAGLE-3 drafter from scratch (which had already been done and achieved 75% validation accuracy), the assistant and user were exploring whether a pre-existing drafter from AQ-MedAI — trained for the related Kimi-K2 model — could be fine-tuned to work with K2.5 more efficiently.

The assumption chain went like this:

  1. Architectural compatibility: AQ-MedAI's K2 drafter uses the same LlamaForCausalLMEagle3 architecture as the K2.5 drafter, so the model structure should be compatible.
  2. Hidden state compatibility: The K2 and K2.5 models share the same underlying architecture (hidden_size=7168, head_dim=128, etc.), so the hidden states fed to the drafter should have the same dimensionality and roughly similar distribution.
  3. Vocab mapping equivalence: Both drafters use a draft_vocab_size of 32,000 and map to a target vocab of 163,840. The assistant assumed that the draft-to-target token mappings (d2t) were either identical or close enough that fine-tuning could bridge the gap.
  4. Weight loading correctness: The --finetune-from logic in the training script loaded lm_head weights from the pretrained drafter and expected them to be compatible with the training labels computed using the local vocab mapping. Assumption #3 turned out to be catastrophically wrong. The probe in [msg 4980] revealed that only 252 out of 32,000 draft-to-target token positions matched between AQ-MedAI's mapping and the K2.5 mapping. This meant that when the lm_head predicted "token at position 500," AQ-MedAI's mapping said that position corresponded to target token X, but the training pipeline's mapping said position 500 corresponded to target token Y. The labels were scrambled.

The Input Knowledge Required

To fully understand this message, one needs to understand several layers of the EAGLE-3 speculative decoding architecture. EAGLE-3 works by training a lightweight "draft" model that predicts multiple future tokens in parallel, given the hidden states from the main "verifier" model. The draft model has its own vocabulary (draft_vocab_size=32,000) which is a subset of the verifier's full vocabulary (163,840 for Kimi-K2.5). The mapping between draft token IDs and verifier token IDs is stored in two tensors: t2d (target-to-draft, a boolean mask of size 163,840 indicating which target tokens are in the draft vocab) and d2t (draft-to-target, an integer array of size 32,000 giving the target token ID for each draft position).

The critical detail is that d2t is not just a fixed subset — it's an ordered mapping. Draft position 0 maps to target token 0, draft position 1 maps to target token 1, and so on, but only for the first ~50 tokens (special tokens and byte tokens). After that, the mapping depends on which target tokens were selected as the top-32,000 most frequent tokens in the training data. Since K2 and K2.5 have different training distributions, their top-32,000 tokens differ, and the ordering within the shared tokens is completely different.

This is why the lm_head remapping was necessary: the lm_head is a linear layer with 32,000 output logits, and the row of the weight matrix determines which draft token is predicted. If draft position 42 means target token 42 in one mapping but target token 1,337 in another, then the same row of the lm_head predicts completely different things in the two systems.## The Thinking Process Visible in the Message

The message at index 4973 is relatively short — it's essentially a status check with the training log output. But the thinking process is revealed through what the assistant doesn't say. The assistant expected to see reasonable loss values, or at least loss values that were trending in the right direction. Instead, the log shows numbers that are so high they can only be explained by a systematic error.

The assistant's reasoning at this point, visible in the surrounding messages, follows a careful diagnostic path:

  1. Initial interpretation ([msg 4974]): The assistant initially attributes the high loss to expected misalignment, noting "this is expected at the start" and that the learning rate is still warming up. This is a reasonable first hypothesis — fine-tuning from a different model's weights often starts with high loss.
  2. User challenge ([msg 4975]): The user pushes back with the key insight: if the model was predicting acceptably in SGLang, the weights can't be completely random. This forces a deeper investigation.
  3. Hypothesis generation ([msg 4976]): The assistant stops training and generates three possible causes: vocab mapping mismatch, hidden state format mismatch, or verifier_lm_head mask mismatch.
  4. Systematic elimination ([msg 4978]): The assistant writes a diagnostic script to check the vocab mapping, comparing AQ-MedAI's d2t with the local d2t.
  5. Root cause confirmed (<msg id=4980-4981>): The script reveals only 252/32,000 positional matches. The assistant immediately identifies the fix: remap the lm_head rows using a permutation derived from the two d2t mappings. This diagnostic chain is a textbook example of how to debug a machine learning pipeline: when you see behavior that contradicts known results, don't attribute it to randomness or bad luck — look for a structural bug in the data flow.

The Output Knowledge Created

This message and the investigation it triggered produced several important pieces of knowledge:

  1. The vocab mapping is not portable across model versions: Even between closely related models like K2 and K2.5, the draft-to-target token mappings differ in 31,748 out of 32,000 positions. This means that EAGLE-3 drafters cannot be directly transferred between model versions without a remapping step.
  2. The remapping strategy: The solution was to build a permutation matrix that maps AQ-MedAI's lm_head rows to the correct positions in the local vocab ordering. For the 6,250 AQ-MedAI draft tokens that don't exist in the K2.5 draft vocab at all, the assistant planned to initialize those rows from the verifier's lm_head or randomly.
  3. A debugging methodology: The assistant established a pattern of writing standalone diagnostic scripts (like debug_vocab_mismatch.py) that compare tensors directly, rather than trying to trace through the complex training pipeline. This is a valuable technique for any ML engineer working with multi-stage pipelines.
  4. The limits of fine-tuning transfer: Even after fixing the vocab mapping, the fine-tuning would eventually plateau at ~38% accuracy (as documented in later chunks), converging much slower than the from-scratch model. This led to the conclusion that the K2 weights were a poor initialization for K2.5 — the representations had diverged too much between the two model versions.## The Broader Significance This message sits at a critical inflection point in the session. Before it, the team was pursuing a data-centric strategy: more training data, better fine-tuning, leveraging pre-existing weights. After it, the strategy pivoted dramatically. The fine-tuning approach was abandoned, and the focus shifted to system-level optimization of the verify step — the NCCL all-reduce bottleneck that was consuming 97% of the EAGLE-3 cycle time. The pivot was driven by the realization that fine-tuning was not just slow but fundamentally limited by the vocab mapping mismatch. Even after fixing the remapping, the fine-tuned model plateaued at ~38% accuracy, far below the from-scratch model's 75%. The K2 weights were not a useful starting point — they encoded a different distribution of token relationships that was actively misleading for K2.5. This is a common pattern in deep learning: the most sophisticated optimization strategy (fine-tuning a pre-trained model) can fail spectacularly when there's a subtle data mismatch that isn't visible at the surface level. The vocab mapping was invisible to the training loop — it was baked into the lm_head weights as a permutation of the output space. The training pipeline saw "loss = 18" and assumed the model was bad, when in fact the model was fine but the labels were scrambled. The message also illustrates the importance of the human-in-the-loop in AI-assisted coding sessions. The assistant initially dismissed the high loss as "expected," but the user's intervention — asking the right question about the SGLang inference results — triggered a debugging chain that uncovered the root cause. This is a powerful example of how effective human-AI collaboration works: the AI handles the mechanical work of running experiments and checking logs, while the human provides the critical insight that connects disparate observations.

Conclusion

Message 4973 is, on its surface, a routine status check — the assistant waited for a training run to produce its first metrics and reported them. But the numbers in that log line told a story of contradiction: a model that worked in inference but failed in training. The investigation that followed revealed a fundamental structural mismatch in the vocab mapping, leading to a remapping fix, a failed fine-tuning experiment, and ultimately a strategic pivot to system-level optimization. This message captures the moment when a seemingly straightforward fine-tuning attempt collided with the messy reality of model versioning and tokenizer drift — a collision that would reshape the entire direction of the project.