The Vocab Remapping That Saved (and Doomed) the K2 Fine-Tuning Experiment

Introduction

In the sprawling, multi-week effort to deploy speculative decoding for the Kimi-K2.5 language model on an 8-GPU Blackwell system, message [msg 4990] represents a quiet but pivotal moment. It is the message where a critical bug is confirmed fixed, the first fruits of a diagnosis are visible in the numbers, and yet—unbeknownst to the assistant at this exact moment—the trajectory toward abandoning the entire approach has already been set. This single message, a brief status update accompanied by a bash command and its output, sits at the intersection of a successful debugging intervention and the early warning signs of a failed strategy.

The Context: A Fine-Tuning Experiment in Jeopardy

To understand why message [msg 4990] was written, we must step back into the broader narrative. The assistant had been attempting to improve Kimi-K2.5 inference throughput using EAGLE-3 speculative decoding. After training a from-scratch EAGLE-3 drafter that achieved 74.7% validation accuracy, the assistant pivoted to a promising shortcut: fine-tuning an existing EAGLE-3 drafter from AQ-MedAI (trained for the Kimi-K2 model) onto K2.5 data. The hypothesis was that the K2 weights would provide a strong initialization, accelerating convergence and potentially yielding a better drafter than training from scratch.

The initial fine-tuning attempt ([msg 4971]) produced catastrophic results: losses of ~18-20 and near-zero accuracy (~2-7%). This was essentially random performance—yet the same AQ-MedAI drafter had achieved an acceptance length of ~1.5 when deployed directly in SGLang with K2.5, proving it was predicting above chance. The contradiction demanded investigation.

The user astutely asked in [msg 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?" This question triggered a systematic debugging effort that revealed the root cause.

The Diagnosis: A Silent Vocab Mapping Catastrophe

The assistant's investigation ([msg 4976] through [msg 4981]) uncovered a subtle but devastating bug: a vocab mapping mismatch between the AQ-MedAI drafter and the K2.5 training pipeline. The EAGLE-3 architecture uses a compressed draft vocabulary—32,000 tokens instead of the full 163,840-token vocabulary of the base model. Each position in the draft vocabulary maps to a specific target token ID via a "draft-to-target" (d2t) mapping. The AQ-MedAI drafter had been trained with its own d2t mapping, while the K2.5 training pipeline used a different d2t mapping derived from the K2.5 model's tokenizer.

The diagnostic script revealed the staggering extent of the mismatch: only 252 out of 32,000 draft vocabulary positions mapped to the same target token between the two mappings. This meant that when the fine-tuning script loaded the AQ-MedAI lm_head weights (the output projection layer), it was feeding the model labels computed with a completely different token ordering. Position 500 in the draft vocabulary meant one thing to AQ-MedAI's lm_head but something entirely different to the K2.5 training labels. The result was effectively random gradients—the model was being penalized for predicting the "right" token in the "wrong" position.

The Fix: Remapping the lm_head

Message [msg 4990] is the first report of the fix being applied and producing results. The assistant had modified the fine-tuning script ([msg 4983]) to remap the AQ-MedAI lm_head weights from their draft vocabulary ordering to the K2.5 ordering. The algorithm was straightforward:

  1. Load AQ-MedAI's d2t mapping to determine which target token each of their 32,000 draft positions maps to.
  2. For each AQ-MedAI draft position, find the corresponding position in the K2.5 draft vocabulary (using the K2.5 t2d mapping in reverse).
  3. Permute the lm_head weight rows accordingly.
  4. For the 6,250 AQ-MedAI draft tokens that don't exist in the K2.5 vocabulary at all, initialize those rows from the verifier's lm_head (the fallback initialization). The message states: "25,750 out of 32,000 lm_head rows were remapped from AQ-MedAI, the other 6,250 keep the verifier initialization." This is a concise status update confirming that the remapping logic executed correctly. The 6,250 rows that couldn't be remapped correspond to target tokens that exist in AQ-MedAI's vocabulary but not in K2.5's—a consequence of the two models having different tokenizers and vocabulary coverage.

The Results: Improvement, But Not Salvation

The bash command in the message polls the training logs and reveals the first metrics after the fix:

loss_0: 9.22 (down from ~18-20)
full_acc_0: 12.1% (up from ~2-7%)
cond_acc_0: 12.1%

This is a dramatic improvement. The loss has been cut in half, and accuracy has jumped from near-random to a meaningful 12%. The fix worked. The vocab remapping was the correct diagnosis, and the intervention was successful.

However, the numbers also tell a more sobering story. The loss of 9.22 is still very high compared to the from-scratch model's final loss of ~0.9. The accuracy of 12.1% is far from the 74.7% the from-scratch model achieved. And critically, the per-step breakdown shows a steep degradation: step 0 (the first draft token prediction) achieves 12.1% accuracy, but step 1 drops to 0.2%, and steps 2-4 achieve 0%. The drafter is only confidently predicting the first draft token; subsequent tokens are essentially random.

This pattern—good first-token accuracy collapsing for subsequent tokens—is a hallmark of the fundamental problem: the AQ-MedAI drafter's internal representations (the hidden state processing in the midlayer transformer block) are tuned to K2's hidden state distribution, not K2.5's. The lm_head remapping fixes the output vocabulary alignment, but it cannot fix the mismatch in the hidden state representations that feed into the draft prediction. The drafter is seeing K2.5 hidden states that look different from what it was trained on, and its single transformer layer cannot compensate.

Assumptions and Their Consequences

Several assumptions underpinned this experiment, and message [msg 4990] reveals their limitations:

Assumption 1: The vocab mapping is the dominant source of error. This turned out to be correct—fixing it dropped the loss by half. But it was not the only source of error. The hidden state distribution mismatch remained.

Assumption 2: Fine-tuning can bridge the representation gap. The assistant chose a low learning rate (5e-5) to "gently adapt, not destroy K2 representations" ([msg 4971]). This assumed that the K2 representations were a good starting point that needed only minor adjustment. The subsequent training would show this assumption was wrong—the fine-tuned model plateaued at ~38% accuracy, far below the from-scratch model's 75%.

Assumption 3: The lm_head remapping is sufficient. The assistant correctly noted that only lm_head is affected by vocab ordering ([msg 4984]), but implicitly assumed that fixing the output layer would allow the rest of the network to adapt during fine-tuning. The per-step accuracy collapse suggests the hidden state mismatch is too severe for a single transformer layer to overcome.

Input Knowledge Required

To fully understand message [msg 4990], one needs knowledge of:

Output Knowledge Created

This message produces several important pieces of knowledge:

  1. Confirmation that vocab remapping works: The technique of permuting lm_head rows based on d2t mapping alignment is validated as an effective fix for cross-model fine-tuning.
  2. Quantified mismatch severity: 25,750/32,000 rows remappable, 6,250 unmappable—this quantifies the vocabulary overlap between K2 and K2.5.
  3. Baseline post-fix metrics: Loss ~9.2, accuracy ~12.1% for step 0, with rapid degradation for subsequent steps. These numbers become the baseline for evaluating whether fine-tuning is worth continuing.
  4. Evidence of hidden state mismatch: The per-step accuracy collapse provides strong evidence that the representation mismatch is deeper than the output layer.

The Thinking Process

The message reveals a clear chain of reasoning:

  1. Verify the fix executed correctly: The assistant first confirms the remapping numbers (25,750 remapped, 6,250 fallback), ensuring the code change produced the expected result.
  2. Poll for results: The sleep 30 command shows the assistant knows the training loop needs time to initialize and produce its first metrics.
  3. Interpret the output: The metrics are presented without commentary, but the assistant's understanding is implicit in the structure—loss_0 through loss_4 correspond to the 5 draft-token prediction steps (TTT steps), and the degradation pattern is immediately visible.
  4. No premature conclusion: The assistant does not declare victory or defeat. The message is a checkpoint—"the fix is working, here are the numbers, let's see where this goes." This measured response reflects the understanding that one data point does not make a trend.

Conclusion

Message [msg 4990] is a snapshot of a scientific process in action: a hypothesis (vocab mismatch causing random loss) is tested, the intervention (lm_head remapping) is confirmed correct, and the results (loss halved, accuracy up to 12%) are recorded for evaluation. It is neither a triumph nor a failure—it is data. The message embodies the disciplined, iterative approach of ML engineering: diagnose, fix, measure, and let the numbers guide the next decision. In the broader arc of the session, this message marks the point where the K2 fine-tuning path, while not yet abandoned, begins to show its fundamental limitations. The fix was necessary, correct, and informative—but ultimately insufficient to salvage the approach.