The Pivot Question: How a User's Skepticism Uncovered a Critical Vocab Mapping Bug in EAGLE-3 Fine-Tuning

Message Overview

In the middle of a complex session fine-tuning an AQ-MedAI K2 EAGLE-3 draft model for Kimi-K2.5, the user asked a single, incisive question that redirected the entire debugging effort:

"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 message, indexed as <msg id=4975>, appears at a critical juncture. The assistant had just launched a fine-tuning run of the AQ-MedAI K2 drafter on 37K K2.5 training samples. The initial training metrics showed a loss of ~18-20 and accuracy of 0-7% — essentially random performance. The assistant had dismissed this as "expected at the start" because "the K2 drafter's weights are misaligned with K2.5 hidden states." But the user spotted a logical contradiction: if the same model was achieving an accept_len of ~1.5 in SGLang inference (meaning it was predicting above chance), how could its training loss be essentially random?

This article examines why this question was so important, what assumptions it challenged, the knowledge it required, and the cascade of discoveries it triggered.

Context: The State of Play Before the Question

To understand the significance of <msg id=4975>, we need to trace the narrative leading up to it. The session was part of a larger effort to improve Kimi-K2.5 inference throughput using speculative decoding with EAGLE-3. The team had already trained a from-scratch EAGLE-3 drafter that achieved 74.7% validation accuracy and ~60 tok/s, but this was still below the baseline of 82 tok/s without speculation.

The user had identified a promising shortcut: AQ-MedAI had released a pre-trained EAGLE-3 drafter for Kimi-K2 (the predecessor model). Since K2 and K2.5 share significant architecture, the hypothesis was that fine-tuning the K2 drafter on K2.5 data would converge faster and achieve better performance than training from scratch.

Phase 0 was a direct probe: deploy the AQ-MedAI K2 drafter unchanged with the K2.5 base model in SGLang. The result was an accept_len of ~1.5 and throughput of 52 tok/s — worse than the from-scratch drafter (60 tok/s) but above the break-even threshold of 1.0. The assistant interpreted this as a "positive signal for fine-tuning," concluding that "the K2 drafter just needs adaptation to K2.5's specific hidden state distribution."

Phase 1 began: fine-tune the K2 drafter on the existing 37K K2.5 training samples. The assistant launched training with a conservative learning rate (5e-5) and 5 epochs, using the --finetune-from flag that loads AQ-MedAI weights and remaps internal layer names.

The first training metrics arrived in <msg id=4973>:

loss_0: 17.80, full_acc_0: 0.065
loss_1: 19.59, full_acc_1: 0.0
loss_2: 20.25, full_acc_2: 0.0

These numbers are essentially random — cross-entropy of ~18 over a 32,000 vocab corresponds to uniform random guessing. The assistant's response in <msg id=4974> was:

"Training is running. The losses are very high (~18-20 for step 0 vs ~0.9 for our from-scratch model's final state) because the K2 drafter's weights are misaligned with K2.5 hidden states — this is expected at the start. Accuracy is ~2-7% for step 0."

This is where the assistant made a critical error in reasoning. The explanation seemed plausible on the surface: different models have different hidden state distributions, so a drafter trained on K2's hidden states would initially perform poorly on K2.5's hidden states. But this explanation failed to account for a crucial data point: the same model had just demonstrated non-random performance in SGLang.

The User's Insight

The user's question in <msg id=4975> cuts to the heart of the contradiction. The reasoning can be reconstructed as follows:

  1. Premise 1: The AQ-MedAI K2 drafter, when deployed in SGLang with the K2.5 base model, achieved an accept_len of ~1.5. This means the drafter's predictions were correct more often than random chance — the model was producing useful draft tokens.
  2. Premise 2: The same drafter, when fine-tuned on K2.5 training data, showed a loss of ~18-20 and accuracy of 0-7%. These are random-level metrics.
  3. Contradiction: If the drafter can predict useful tokens during inference, its underlying representations must contain signal. A loss of ~18 implies the model is guessing randomly, which contradicts the inference behavior.
  4. Hypothesis: The problem is not with the model weights themselves, but with how the training pipeline feeds data to the model. The layers and hidden states might be wired incorrectly, or the labels might not correspond to what the model is actually predicting. The user's phrasing — "Are we passing it correct layers/hidden states etc correctly?" — reveals the specific suspicion: the training data pipeline has a mismatch that makes the training task different from the inference task.

Assumptions Challenged

The assistant had made several assumptions that the user's question exposed as potentially flawed:

Assumption 1: High initial loss is expected during domain adaptation. The assistant assumed that the K2→K2.5 gap would naturally produce high loss at the start of fine-tuning. While some elevation is expected, a loss of ~18 (equivalent to random guessing over 32,000 classes) is not "expected" — it indicates the model is receiving no useful signal at all.

Assumption 2: The training data pipeline is correct. The assistant had previously debugged and fixed hidden state format issues in the training pipeline (in earlier segments). Having fixed those, the assumption was that the pipeline was now correct. The user's question forced a re-examination of this assumption.

Assumption 3: The vocab mapping is consistent. The assistant knew that AQ-MedAI's K2 drafter had its own d2t (draft-to-target) and t2d (target-to-draft) mappings, but assumed the fine-tuning code handled this correctly. The --finetune-from logic loaded weights and skipped t2d/d2t tensors, but the critical question was whether the lm_head weights — which predict logits in the draft vocab space — were compatible with K2.5's target labels.

The Investigation and Discovery

The assistant's response in <msg id=4976> shows immediate recognition of the insight:

"Excellent question. The loss is ~18-20 (cross-entropy over 32000 vocab) and accuracy is ~2-7% — that's essentially random. Yet the same AQ-MedAI drafter achieved accept_len ~1.5 when deployed in SGLang with K2.5, which means it WAS predicting above chance. This strongly suggests the training pipeline is feeding the drafter different data than what SGLang provides at inference."

The assistant killed the training run and began a systematic investigation. The first hypothesis was a vocab mapping mismatch: AQ-MedAI's lm_head was trained to output logits for their draft-to-target mapping, but the training labels were computed using K2.5's mapping. If draft token position 42 maps to target token X in AQ-MedAI's mapping but target token Y in K2.5's mapping, then the lm_head is being trained with wrong labels.

A diagnostic script was written and executed (via <msg id=4978> through <msg id=4980>). The results were stark:

The Fix and Its Impact

The fix was to remap the lm_head weights from AQ-MedAI's draft vocab ordering to K2.5's ordering. The assistant modified the --finetune-from logic in 04_train.py to build a permutation matrix: for each of AQ-MedAI's 32,000 draft positions, find the corresponding target token ID, then find that target token's position in K2.5's draft vocab. The lm_head rows were permuted accordingly. For the 6,250 AQ-MedAI draft tokens with no K2.5 equivalent, the corresponding rows were initialized from the verifier's lm_head.

The results of the fix were immediate and dramatic. In <msg id=4990>, the first training metrics after the fix showed:

Knowledge Required to Understand This Message

To fully grasp the user's question, one needs:

  1. Understanding of EAGLE-3 architecture: EAGLE-3 uses a lightweight "draft" model that predicts multiple tokens in parallel using hidden states from the base model. The draft model has its own vocabulary (draft vocab) and a mapping (d2t) that converts draft token IDs to target model token IDs.
  2. Understanding of vocab mappings: In EAGLE-3, the target model's large vocabulary (163,840 tokens for Kimi-K2.5) is compressed to a smaller draft vocabulary (32,000 tokens) by selecting the most frequent tokens. The t2d array maps target IDs to draft positions, and d2t maps draft positions back to target IDs. Different training runs produce different mappings because they select different top-K tokens.
  3. Understanding of training vs inference data flow: During inference, SGLang feeds hidden states from the base model's transformer layers into the drafter, which predicts draft tokens. During training, the pipeline loads pre-extracted hidden states and computes loss against target tokens. If the hidden state format or vocab mapping differs between these two paths, the training signal is corrupted.
  4. Understanding of cross-entropy loss: A loss of ~18 over 32,000 classes is essentially random (ln(32000) ≈ 10.37 for uniform distribution, but with the softmax temperature and label distribution, ~18 indicates no useful signal).

Knowledge Created by This Message

The user's question created several important outputs:

  1. The vocab mapping mismatch diagnosis: The investigation revealed that only 252/32,000 draft positions had consistent mappings between AQ-MedAI and K2.5. This became a documented finding that would inform all future fine-tuning attempts.
  2. The lm_head remapping technique: The fix developed a general technique for transferring EAGLE-3 draft models between different vocab mappings — building a permutation matrix from the source model's d2t to the target model's d2t. This is reusable for any future model adaptation.
  3. A debugging methodology: The question established a principle: when inference behavior contradicts training metrics, suspect the data pipeline, not the model weights. This became a recurring theme in the session.
  4. A cautionary finding about transfer learning: The K2→K2.5 fine-tuning ultimately plateaued at ~38% accuracy (as documented in the chunk summary), converging slower than the from-scratch model. The vocab mapping fix was necessary but not sufficient — the K2 weights were still a poor initialization for K2.5. This finding led the team to abandon the fine-tuning approach and pivot to system-level optimization.

The Thinking Process

The user's question reveals a sophisticated debugging mindset. Rather than accepting the assistant's explanation that "high loss is expected," the user performed a mental consistency check: inference performance (accept_len ~1.5) implies non-random predictions, which should produce non-random loss. The contradiction means something is fundamentally wrong with how training measures loss versus how inference measures acceptance.

The specific phrasing — "Are we passing it correct layers/hidden states etc correctly?" — shows the user was thinking about the data flow. The "etc" is telling: the user wasn't sure exactly what was wrong, but had identified the general category (data pipeline mismatch) and was asking the assistant to investigate systematically.

This is a classic example of the "skeptical reviewer" role in debugging: someone who doesn't accept surface-level explanations and forces a deeper investigation. The assistant's initial explanation was plausible but wrong, and the user's skepticism saved hours of wasted training time.

Conclusion

The user's question in <msg id=4975> is a masterclass in debugging intuition. By spotting a contradiction between inference behavior and training metrics, the user identified a subtle vocab mapping bug that would have rendered the entire fine-tuning effort meaningless. The fix — remapping the lm_head weights — immediately dropped loss from ~18 to ~9 and improved accuracy from ~7% to ~24%.

More importantly, the question established a debugging principle that guided the remainder of the session: when the numbers don't add up, suspect the data pipeline. This principle would prove invaluable in subsequent debugging of hidden state format mismatches, NCCL communication bottlenecks, and other issues that plagued the EAGLE-3 deployment effort.

The message also demonstrates the value of domain expertise in AI debugging. Understanding how EAGLE-3's vocab mappings work, how training data flows through the pipeline, and what "random loss" looks like numerically — all of this knowledge was required to ask the right question at the right time.