The Permutation That Saved a Fine-Tuning Run: Diagnosing and Fixing a Vocab Mapping Catastrophe in EAGLE-3 Transfer Learning

In the high-stakes world of large language model inference optimization, few things are more disorienting than watching a model that performed reasonably well in production suddenly produce random noise during training. This was precisely the situation facing the assistant in message [msg 4983] of an extended opencode session devoted to improving speculative decoding throughput for the Kimi-K2.5 model. The message, deceptively brief at first glance, represents the culmination of a focused diagnostic effort and the application of a critical fix that would determine the viability of an entire fine-tuning strategy.

The Context: A Promising Shortcut Turns Sour

The session had been exploring whether the AQ-MedAI K2 EAGLE-3 drafter—a pre-trained draft model for speculative decoding—could be adapted to work with the Kimi-K2.5 base model. This was an attractive shortcut: rather than training a drafter from scratch (which had already been done and achieved 74.7% validation accuracy), fine-tuning an existing K2 drafter might yield better performance with less data and compute. The initial probe was encouraging: when deployed directly in SGLang with K2.5, the AQ-MedAI drafter achieved an acceptance length of ~1.5 tokens and 52 tok/s throughput. While worse than the from-scratch drafter (60 tok/s) and the baseline without speculation (82 tok/s), the acceptance length above 1.0 confirmed that the K2 representations shared meaningful structure with K2.5's hidden state distribution.

However, when the assistant launched fine-tuning in message [msg 4971], the results were catastrophic. The loss hovered around 18–20 (cross-entropy over a 32,000-token vocabulary), and accuracy was essentially zero at 2–7%. This was random-guess territory. Yet the same drafter had just demonstrated above-chance prediction in SGLang. Something was fundamentally wrong with how the training pipeline was feeding data to the model.

The Investigation: A User's Insight Triggers a Deep Dive

The user's question in [msg 4975] cut to the heart of the matter: "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 prompted the assistant to halt training immediately and begin a systematic investigation.

The assistant identified three possible causes in [msg 4978]:

  1. Vocab mapping mismatch: The AQ-MedAI drafter's lm_head was trained with their own draft-to-target (d2t) and target-to-draft (t2d) mappings, but the training labels were computed using the local K2.5 mappings. If draft token 42 maps to target token A in AQ-MedAI's system but target token B in the local system, the labels are scrambled.
  2. Hidden state format differences: Perhaps the training data preprocessing (standardize_data_v1) fed hidden states differently than SGLang did at inference time.
  3. Verifier lm_head mismatch: The loss computation used the verifier's lm_head (masked with the local t2d mapping), which might not align with what the AQ-MedAI drafter was trained to predict. The assistant correctly prioritized the vocab mapping as the most likely culprit and wrote a diagnostic script ([msg 4979]) to compare the two mappings. The results, returned in [msg 4980], were stark: only 252 out of 32,000 draft-to-target token positions matched between the AQ-MedAI and local mappings. The first ~50 positions (low token IDs for special tokens and byte representations) happened to align, but after that the orderings diverged completely. While 25,750 of AQ-MedAI's target tokens existed somewhere in the local vocab, they were at different positions—meaning the lm_head weights were systematically misaligned with the training labels.

The Target Message: A Surgical Code Intervention

Message [msg 4983] is the moment of action. The assistant states its intent plainly: "I need to replace the entire finetune-from block. The key change is: when loading lm_head.weight, remap it from AQ-MedAI's draft vocab ordering to ours by building a permutation from their d2t to our d2t."

This single sentence encapsulates the entire fix. The lm_head is a linear layer that maps from the drafter's hidden dimension (7,168) to the draft vocabulary (32,000 positions). Each row of the weight matrix corresponds to a draft vocabulary position, and during training, the loss is computed by comparing the logit at each position against the label at that same position. If the label says "position 500 should predict target token X" but the weight at row 500 was trained to predict a completely different target token Y, the gradient is meaningless—hence the random loss.

The fix is conceptually elegant: build a permutation matrix that reorders the rows of the AQ-MedAI lm_head weight matrix so that row i in the remapped matrix corresponds to the same target token as row i in the local mapping. For draft positions where AQ-MedAI's target token doesn't exist in the local vocab (6,250 tokens), those rows would need random initialization or copying from the verifier's lm_head.

The assistant applies the edit using the edit tool, and the operation succeeds. The LSP errors that follow are a red herring—they report that torch, transformers, and speculators cannot be resolved, but this is expected because the development environment doesn't have these ML packages installed locally. They only exist in the remote container where training runs. The assistant correctly ignores these errors.

Decisions and Assumptions

The assistant made several key decisions in this message:

Decision to remap rather than retrain from scratch: Rather than abandoning the fine-tuning approach or trying to retrain the entire lm_head from random initialization (which would require many epochs to converge), the assistant chose to preserve the pre-trained weights by permuting them. This was the right call: the lm_head contains learned associations between hidden states and token probabilities that took significant compute to acquire. A permutation preserves all that information while fixing the label alignment.

Decision to only remap lm_head: The assistant implicitly assumed that only the lm_head layer was affected by the vocab mapping mismatch. This was validated in the next message ([msg 4984]) where the assistant checked whether the fc (fully connected) layer also needed remapping and correctly concluded it did not, since the fc layer operates on concatenated hidden states (21,504 dimensions → 7,168 dimensions) and is independent of token ordering.

Assumption about the permutation approach: The assistant assumed that a simple row permutation would suffice—that for each AQ-MedAI draft position, there exists a corresponding local draft position with the same target token. The diagnostic data showed this was true for 25,750 out of 32,000 positions. The remaining 6,250 positions (where AQ-MedAI's target token doesn't appear in the local vocab) would need special handling, but the message doesn't detail this—it's left for the implementation.

What the Message Reveals About the Thinking Process

The reasoning visible in this message and its immediate predecessors reveals a methodical, hypothesis-driven approach to debugging. The assistant didn't just guess at the cause of the random loss—it formulated three specific hypotheses, prioritized them by likelihood, wrote a targeted diagnostic script, interpreted the results, and only then applied the fix. This is textbook scientific debugging.

The message also reveals the assistant's deep understanding of the EAGLE-3 architecture. The lm_head in an EAGLE-3 drafter is not the same as the base model's lm_head—it's a separate linear layer that maps from the drafter's hidden states to a draft vocabulary that is a subset of the full model vocabulary, determined by the t2d (target-to-draft) mapping. The draft vocabulary ordering is arbitrary and can differ between independently trained drafters. The assistant understood that the AQ-MedAI and local drafters had independently chosen different orderings for their draft vocabularies, and that this ordering difference was the root cause of the label scrambling.

Input and Output Knowledge

To understand this message, the reader needs to know:

The Aftermath

The fix was applied, and training was relaunched. The results, visible in subsequent messages, showed immediate improvement: loss dropped from ~18 to ~9, and accuracy rose to ~24%. However, the fine-tuning ultimately plateaued at ~38% accuracy, converging much slower than the from-scratch model (which reached 75% by epoch 5). The assistant concluded that the K2 weights were a poor initialization for K2.5 and abandoned the approach. But the fix itself was correct—it transformed a completely broken training run into one that made meaningful progress, even if that progress wasn't sufficient to justify continuing down this path.

Message [msg 4983] stands as a testament to the importance of understanding data representations in deep learning. A seemingly minor detail—the ordering of rows in a vocabulary mapping—can completely destroy training if misaligned. The assistant's systematic diagnosis and surgical fix saved what could have been hours of wasted compute and provided a reusable technique for future cross-drafter transfer learning attempts.