The Vocab Mapping Mismatch: Debugging Why a Working EAGLE-3 Drafter Produced Random Loss During Fine-Tuning

Introduction

In the course of optimizing Kimi-K2.5 inference throughput using speculative decoding with EAGLE-3, a team of engineers encountered a deeply puzzling contradiction. The AQ-MedAI K2 EAGLE-3 drafter, when deployed directly in SGLang with the K2.5 base model, achieved an acceptance length of approximately 1.5 tokens and a throughput of 52 tok/s—modest but clearly above random chance. Yet when they attempted to fine-tune this same drafter on K2.5 training data, the loss started at ~18-20 (effectively random for a 32,000-vocabulary cross-entropy) and the accuracy hovered at 2-7%, barely distinguishable from a model guessing uniformly. How could the same drafter perform above chance at inference yet appear completely untrained when subjected to fine-tuning?

This article examines a single pivotal message in the conversation ([msg 4978]), where the assistant systematically diagnoses this contradiction. The message represents a critical debugging juncture: the moment when the team realized that the fine-tuning pipeline and the inference pipeline were feeding the drafter fundamentally different data, despite both ostensibly operating on the same model. The assistant's reasoning, the diagnostic script it wrote, and the shell escaping error that cut the investigation short all reveal important lessons about the subtle ways that training and inference pipelines can diverge in large language model systems.

The Context: A Promising but Perplexing Fine-Tuning Attempt

To understand the significance of message 4978, we must trace the narrative arc that led to it. The team had been engaged in a multi-week effort to deploy the Kimi-K2.5 model (a 1-trillion-parameter Mixture-of-Experts architecture) with speculative decoding on an 8-GPU RTX PRO 6000 Blackwell server. Their earlier work had produced a from-scratch EAGLE-3 drafter that achieved 74.7% validation accuracy and ~60 tok/s in speculative decoding—a respectable result, but still below the 82 tok/s baseline without speculation.

Seeking improvement, the team turned to the AQ-MedAI K2 EAGLE-3 drafter, a publicly available checkpoint trained on the earlier Kimi-K2 model. The hypothesis was straightforward: since K2 and K2.5 share significant architectural DNA, the K2 drafter's weights might provide a better initialization than random, potentially accelerating convergence and yielding a higher-quality final drafter. Phase 0—a direct probe of the K2 drafter on K2.5 without any fine-tuning—confirmed this hypothesis partially: the drafter achieved an acceptance length of ~1.5, above the baseline of 1.0 (which would mean no speculation benefit). This was worse than their from-scratch drafter (accept_len ~2.0) but clearly demonstrated that the K2 representations were correlated with K2.5 hidden states.

Encouraged, the team launched Phase 1: fine-tuning the AQ-MedAI K2 drafter on their existing 37K K2.5 training samples. They used a conservative learning rate (5e-5) to gently adapt rather than destroy the pre-trained representations, and set the training to run for 5 epochs. The initial training metrics, however, were catastrophic. The loss started at approximately 18-20—essentially random for a 32,000-class classification problem—and the accuracy was 2-7%, barely above the 1/32000 ≈ 0.003% floor. The assistant initially dismissed this as expected startup noise, noting that "the learning rate is still very low (warming up)" and suggesting that the loss would drop once training progressed.

But the user, in message 4975, asked the critical 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 pierced through the assumption that the high loss was merely a warm-up artifact. The assistant immediately recognized the validity of the concern: if the drafter achieved accept_len ~1.5 in SGLang, it was clearly making non-random predictions, yet the training pipeline showed it as performing at chance level. Something was fundamentally wrong with how the training data was being fed to the model.

The Assistant's Systematic Reasoning

Message 4978 opens with the assistant acknowledging the contradiction and laying out a structured diagnostic plan. The reasoning is methodical and demonstrates a clear understanding of the EAGLE-3 architecture's data flow:

There are several possible causes. Let me check them systematically:

>

1. Vocab mapping mismatch: We're loading AQ-MedAI's lm_head (trained with their d2t) but the training data uses our t2d/d2t for target labels. AQ-MedAI's draft token 42 maps to a different target token than our draft token 42. 2. Hidden state format: Maybe something about how standardize_data_v1 feeds data vs how SGLang feeds it. 3. The verifier_lm_head: Training computes loss with verifier_lm_head (which we set from K2.5's masked lm_head with our t2d mask), but AQ-MedAI's lm_head was trained to match a different mask.

This triage is notable for its clarity and completeness. The assistant identifies three distinct mechanisms by which the training and inference pipelines could diverge, and ranks them by likelihood. The first hypothesis—vocab mapping mismatch—is identified as the most probable culprit, and the assistant proceeds to write a diagnostic script to test it.

Hypothesis 1: The Vocab Mapping Mismatch

The vocab mapping hypothesis requires understanding how EAGLE-3's draft model interfaces with the target model's vocabulary. In EAGLE-3, the draft model operates in a reduced "draft vocabulary" of 32,000 tokens, which is a subset of the target model's full vocabulary (163,840 tokens for Kimi-K2.5). Two mappings connect these vocabularies:

Hypothesis 2: Hidden State Format Mismatch

The second hypothesis concerns the hidden states fed to the drafter. In SGLang, the drafter receives hidden states from the target model's internal layers during actual inference. In training, the hidden states were pre-extracted and stored during a separate data generation phase (Step 2 of the pipeline). If the standardize_data_v1 function that loads these pre-extracted hidden states reshapes, normalizes, or selects them differently than SGLang does at inference time, the drafter would receive systematically different inputs, leading to poor predictions.

This hypothesis is less likely than the vocab mapping issue because the hidden state extraction pipeline was designed to mirror SGLang's behavior, but it remains a valid concern that the assistant correctly flags for investigation.

Hypothesis 3: The Verifier LM Head Mask

The third hypothesis involves the verifier_lm_head, which is used during training to compute the loss. In EAGLE-3 training, the loss is computed by passing the draft model's output logits through the verifier's lm_head (a masked version of the target model's lm_head that only considers draft-vocabulary tokens). The team constructed this verifier_lm_head from K2.5's masked lm_head using their own t2d mask. But AQ-MedAI's lm_head was trained to match a different mask (derived from K2's vocabulary statistics). If the masks differ, the verifier_lm_head would penalize the model for outputs that were actually correct according to AQ-MedAI's original training setup.

The Diagnostic Script: A Well-Designed but Flawed Execution

To test the vocab mapping hypothesis, the assistant wrote a Python script that compares AQ-MedAI's d2t/t2d mappings with the team's own mappings. The script is well-structured and demonstrates clear thinking about what metrics matter:

  1. Hot token ID comparison: It computes the actual target token IDs that each draft position maps to (by adding d2t[i] + i) and compares the first 20 entries between the two systems.
  2. Set overlap: It computes how many unique target tokens are covered by each mapping and how many overlap—a measure of whether the two vocabularies select the same tokens from the full 163,840-token vocabulary.
  3. Positional match count: It checks how many draft positions map to the same target token in both systems—a direct measure of mapping alignment.
  4. Remapping analysis: It builds a reverse mapping from target token IDs to draft positions in the team's system, then checks whether AQ-MedAI's draft positions can be remapped to the correct positions in the team's system. This is the most sophisticated analysis: it asks whether the information is preserved under a permutation of the draft vocabulary, even if the absolute positions differ. The script's design reveals the assistant's deep understanding of the EAGLE-3 architecture. The hot token ID computation (d2t[i] + i) is a non-obvious detail of how SGLang implements the draft-to-target mapping, and the remapping analysis anticipates the possibility that the two mappings might be isomorphic under permutation—in which case the training could be fixed by remapping the lm_head's output channels rather than retraining from scratch. However, the script execution fails with a shell error: zsh:32: parse error near ')'. This is a classic shell escaping problem. The Python code contains the expression set_aq & set_ours (using the bitwise AND operator), and the & character is being interpreted by the local zsh shell despite the command being wrapped in single quotes. The nested quoting structure—single quotes for the SSH command containing escaped double quotes for the Python string—creates a fragile escaping chain that breaks when special characters like & or ) appear in certain positions. This error is instructive. It highlights the practical challenges of running complex Python diagnostics through SSH on remote machines. The assistant's decision to inline the Python script within an SSH command rather than writing it to a file and executing it separately reflects a trade-off between speed (inline execution is faster for quick diagnostics) and robustness (file-based execution avoids shell escaping issues). In this case, the trade-off backfired: the diagnostic was cut short, and the team had to wait for the next message to see the results.

Assumptions and Their Validity

The assistant's reasoning in message 4978 rests on several assumptions, some explicit and some implicit:

Assumption 1: The vocab mapping mismatch is the most likely cause. This is a well-reasoned judgment. The assistant correctly identifies that the lm_head's output logits are interpreted relative to the d2t mapping, and if the training labels use a different mapping, the loss computation is invalid. The assumption is validated by the subsequent investigation (in later messages), which reveals that only 252 out of 32,000 draft positions map to the same target token between the two systems—a near-total mismatch.

Assumption 2: The hidden state format is less likely to be the issue. This assumption is reasonable but not definitively proven in this message. The assistant implicitly trusts that the standardize_data_v1 function correctly replicates SGLang's hidden state formatting. In practice, this assumption held—the hidden state format was not the root cause—but the assistant's willingness to deprioritize this check without verification could have led to wasted effort if the assumption had been wrong.

Assumption 3: The verifier_lm_head mask is a separate issue from the vocab mapping. The assistant lists this as a third hypothesis, but it is actually closely related to the first. The verifier_lm_head's mask is derived from the t2d mapping, so if the t2d mappings differ, the masks differ. The assistant's separation of these hypotheses suggests a nuanced understanding that even if the d2t/t2d mappings were identical, the verifier_lm_head mask could still differ due to how it was constructed from the base model's weights.

Assumption 4: The shell will correctly handle the inline Python script. This assumption proved incorrect. The assistant assumed that wrapping the command in single quotes would protect all special characters, but the nested escaping of double quotes within the Python code created a situation where zsh still attempted to parse certain character sequences. This is a common pitfall when embedding complex code in shell commands.

The Mistake: Shell Escaping and Its Consequences

The most concrete mistake in this message is the shell escaping error. The Python script contains the line overlap = set_aq & set_ours, which uses the & operator. In zsh, the & character outside of single quotes would start a background job. Inside single quotes, it should be literal—but the command structure is:

ssh host 'python3 -c "..."'

The outer quotes are single quotes, which should protect everything inside. However, the Python code contains escaped double quotes (\") which are being used to delimit strings within the Python code. When zsh processes the command line, it sees the pattern:

'...\"...\"...&...\"...\"...'

The single quotes should protect the &, but the error message "zsh:32: parse error near `)'" suggests that zsh is somehow interpreting the content. This might be because the command is being passed through multiple shells (e.g., the SSH command itself spawns a shell on the remote machine), or because the escaping of double quotes inside single quotes creates an ambiguity that zsh resolves incorrectly.

A more robust approach would have been to write the Python script to a temporary file on the remote machine and execute it directly, avoiding the nested quoting entirely. The assistant had already used scp in earlier messages to transfer files, so this pattern was available. The choice to use inline execution was likely driven by the desire for speed—writing a file, transferring it, and executing it takes more steps than a single SSH command—but it introduced fragility.

Input Knowledge Required to Understand This Message

To fully grasp the significance of message 4978, the reader needs:

  1. Understanding of EAGLE-3 architecture: Knowledge that the draft model operates in a reduced vocabulary, that d2t and t2d mappings connect draft and target vocabularies, and that the lm_head predicts logits in draft-vocabulary space.
  2. Knowledge of the AQ-MedAI K2 drafter: Awareness that this is a publicly available checkpoint trained on Kimi-K2 data, with its own vocabulary mappings derived from K2's token statistics.
  3. Awareness of the training pipeline: Understanding that training labels are constructed using the team's own d2t/t2d mappings, which may differ from AQ-MedAI's mappings.
  4. Context from earlier messages: Knowledge that Phase 0 (direct probe) achieved accept_len ~1.5, confirming the drafter works above chance in inference, and that Phase 1 (fine-tuning) showed random-level loss.
  5. Shell scripting knowledge: Understanding of how nested quoting works in shell commands, and how special characters like & and ) can cause parsing errors.

Output Knowledge Created by This Message

Message 4978 produces several important outputs:

  1. A structured diagnostic framework: The three-hypothesis triage provides a reusable template for debugging similar mismatches between training and inference pipelines in speculative decoding systems.
  2. A diagnostic script design: Even though the script failed to execute, its design—particularly the hot token ID comparison and the remapping analysis—represents a valuable template for future vocab mapping investigations.
  3. Identification of the vocab mapping mismatch as the primary suspect: The assistant's reasoning convincingly establishes that the vocab mapping is the most likely cause, which guides the subsequent investigation.
  4. A cautionary example of shell escaping issues: The failed execution serves as a practical lesson about the risks of inlining complex code in SSH commands.

The Thinking Process: A Window into Debugging Methodology

The assistant's thinking in this message reveals a sophisticated debugging methodology. Rather than randomly probing possible causes, the assistant:

  1. Articulates the contradiction explicitly: "The loss is ~18-20... 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 framing ensures the debugging target is clearly defined.
  2. Generates multiple hypotheses: Three distinct mechanisms are proposed, each grounded in a specific component of the EAGLE-3 architecture.
  3. Ranks hypotheses by likelihood: The vocab mapping mismatch is identified as the "most likely culprit" based on the reasoning that the lm_head and the training labels use different mappings.
  4. Designs a targeted test: The Python script directly measures the alignment between the two mappings, providing quantitative evidence to confirm or refute the hypothesis.
  5. Anticipates follow-up questions: The remapping analysis in the script asks not just "do the mappings differ?" but "can we fix this by remapping?"—anticipating the next step if the hypothesis is confirmed. This methodology—articulate the contradiction, generate hypotheses, rank by likelihood, design targeted tests, anticipate follow-ups—is a model of systematic debugging that could be applied to many problems in machine learning systems.

Conclusion

Message 4978 captures a pivotal moment in the EAGLE-3 fine-tuning effort: the transition from blind optimization to systematic diagnosis. The assistant's identification of the vocab mapping mismatch as the root cause of the training-inference contradiction was correct, as confirmed by subsequent analysis showing that only 252 out of 32,000 draft positions mapped to the same target token between the two systems. The shell escaping error that prevented the diagnostic script from running was a temporary setback, but the reasoning framework established in this message guided the team toward the correct fix.

More broadly, this message illustrates a fundamental challenge in transfer learning for large language models: vocabulary mappings are not portable across model versions. Even when the architecture is identical, the statistical properties of the training data (which determine which tokens are included in the reduced draft vocabulary) can differ enough to render pre-trained weights nearly useless without proper remapping. The lesson is that vocabulary alignment must be verified—not assumed—when transferring any component that operates in a reduced vocabulary space, from EAGLE-3 drafters to embedding layers to output heads.

The message also serves as a practical reminder that the simplest debugging tools—like a Python script that compares two arrays—are often the most effective, and that the biggest obstacles are sometimes not the ML complexity but the mundane challenges of shell escaping and command execution. In the high-stakes world of large-scale ML deployment, the difference between success and failure can hinge on whether a & character is properly quoted.