Calibrating Against the Reference: How One Message Bridged Training and Inference for EAGLE-3 on Kimi-K2.5

Introduction

In the high-stakes world of large language model deployment, training a speculative decoding draft model is only half the battle. The other half—arguably the more treacherous half—is ensuring that the trained checkpoint actually works when plugged into the inference engine. This article examines a single pivotal message ([msg 2784]) in an opencode coding session where the assistant, having just completed an end-to-end EAGLE-3 training run on the Kimi-K2.5 model, paused to compare its freshly trained checkpoint against the reference implementation from AQ-MedAI. What emerged was a meticulous forensic analysis that revealed three critical differences, culminating in a decision to investigate how vLLM loads EAGLE-3 speculators—a decision that would determine whether the entire training effort had produced something usable or something doomed to fail at inference time.

The Moment of Calibration

The message arrives at a triumphant but precarious moment. In the preceding messages ([msg 2776] through [msg 2783]), the assistant had successfully completed the first end-to-end EAGLE-3 training run on 10 samples, producing a 4.6 GB checkpoint in approximately one minute on a single GPU. The training pipeline—built from scratch by patching the speculators library to work with vLLM 0.16—was now functional. But rather than celebrating and scaling up, the assistant did something more valuable: it checked its work against an external standard.

The subject message opens with a direct comparison between the AQ-MedAI reference checkpoint (downloaded from Hugging Face in [msg 2783]) and the locally trained output. The assistant lists three differences:

1. Key naming: AQ-MedAI uses midlayer.* while ours uses layers.0.* — this is a naming convention difference 2. No embed_tokens: AQ-MedAI doesn't save embed_tokens.weight (it's frozen and loaded from verifier at runtime) 3. Attention shapes with head_dim=128: - q_proj: [8192, 14336] = [64*128, 2*7168] - k_proj: [8192, 14336] = [64*128, 2*7168] - v_proj: [8192, 14336] = [64*128, 2*7168] - o_proj: [7168, 8192] = [hidden_size, 64*128]

This comparison is the intellectual core of the message. Each difference represents a potential compatibility problem, and the assistant must determine which are cosmetic and which are deal-breakers.

The head_dim=128 Discovery: A Subtle but Critical Architectural Choice

The most significant finding is the head_dim=128 discrepancy. This discovery actually began in the message immediately prior ([msg 2780]), where the assistant first noticed that the locally trained checkpoint had head_dim=112—the default computed from hidden_size / num_attention_heads = 7168 / 64. Curious about the reference, the assistant queried the AQ-MedAI config and found it used head_dim=128 explicitly. This means the attention projections in the reference model project up from 7168 to 8192 dimensions (64 heads × 128), rather than the default 7168 to 7168 (64 heads × 112).

This is not a trivial detail. The head_dim parameter controls the dimensionality of the query, key, and value projections in the attention mechanism. Using head_dim=128 instead of head_dim=112 means the draft model's attention operates in a higher-dimensional space, giving it more capacity to model the verifier's behavior. The AQ-MedAI team made a deliberate architectural choice here, and the assistant correctly recognized that matching it was essential for compatibility.

In [msg 2781], the assistant immediately patched 04_train.py to set head_dim=128 in the LlamaConfig. But the subject message reveals that even after this fix, the shape verification was still pending—the locally trained checkpoint still showed the old head_dim=112 shapes because it hadn't been re-run yet. The assistant was working with stale data from the first run while planning the next.

The Naming Convention Puzzle: midlayer vs layers.0

The second difference—key naming—is more subtle. The AQ-MedAI checkpoint uses midlayer.* as the prefix for the single transformer layer's weights, while the speculators library's save_pretrained() uses layers.0.*. This is a pure naming convention difference, but naming conventions in deep learning frameworks are rarely arbitrary. The assistant's immediate thought was: "The naming difference (midlayer vs layers.0) might cause compatibility issues with vLLM."

This concern is well-founded. When vLLM loads a speculator checkpoint, it performs key-by-key matching between the state dictionary and the model's expected parameter names. If the names don't match, the load fails silently or produces a randomly initialized model. The assistant correctly identified this as a risk worth investigating before scaling up training.

The embed_tokens Observation

The third difference—the absence of embed_tokens.weight in the AQ-MedAI checkpoint—is actually reassuring rather than alarming. The assistant notes that the embedding weights are "frozen and loaded from verifier at runtime," meaning the draft model shares the verifier's embedding layer rather than learning its own. This is standard practice in EAGLE-3: the draft model's embedding is tied to the verifier's, so saving it in the checkpoint is redundant. The assistant's checkpoint includes it (as embed_tokens.weight: [163840, 7168]), which is harmless but unnecessary.

The Decision to Investigate vLLM Loading

The most important decision in this message is the assistant's response to these findings. Rather than guessing which differences matter, the assistant launches a sub-task to investigate exactly how vLLM loads EAGLE-3 speculator models:

Let me check how vLLM loads an EAGLE-3 speculator: [task] {"description":"Find vLLM EAGLE-3 loading code", ...}

This is a textbook example of the right way to handle uncertainty in ML engineering. The assistant could have:

  1. Ignored the differences and hoped they didn't matter (risking wasted compute on incompatible checkpoints)
  2. Changed the training code to match AQ-MedAI's naming blindly (risking breakage if vLLM actually expects layers.0.*)
  3. Asked the user for guidance (slowing down the autonomous workflow) Instead, the assistant chose option 4: investigate the inference engine's expectations directly. By reading vLLM's source code, the assistant would get definitive answers about what naming conventions, weight shapes, and config fields are required.

Assumptions and Potential Pitfalls

Several assumptions underpin this message, and it's worth examining them critically:

Assumption 1: AQ-MedAI's checkpoint is the canonical reference. The assistant implicitly assumes that the AQ-MedAI checkpoint represents the "correct" format that vLLM expects. This is a reasonable assumption—AQ-MedAI is the original publisher of the Kimi-K2 EAGLE-3 model—but it's not guaranteed that vLLM's loading code was written specifically for their format. There could be version mismatches.

Assumption 2: The head_dim=128 choice is mandatory. The assistant assumes that matching AQ-MedAI's head_dim=128 is necessary for compatibility. But it's possible that vLLM's EAGLE-3 loader is flexible about head_dim, computing attention shapes dynamically from the config. The task investigation would confirm this.

Assumption 3: The naming difference is a potential problem. The assistant assumes that layers.0.* might not be recognized by vLLM. This is a conservative assumption—better to verify than to discover a silent failure after 1000-sample training.

Assumption 4: The local environment has access to vLLM source code. The task assumes that vLLM is installed on the remote machine at /root/ml-env/lib/python3.12/site-packages/vllm/. This is confirmed by the successful task execution.

Input Knowledge Required

To fully understand this message, the reader needs:

  1. EAGLE-3 architecture knowledge: Understanding that EAGLE-3 is a speculative decoding method where a lightweight "draft" model predicts the verifier's hidden states, and that the draft model has a single transformer layer with custom attention.
  2. The speculators library: Knowledge that the speculators Python package provides EAGLE-3 model definitions, config classes, and a Trainer, and that save_pretrained() produces HuggingFace-compatible checkpoints.
  3. vLLM speculative decoding: Understanding that vLLM supports EAGLE-3 via SpeculativeConfig and loads draft models from HuggingFace-style directories.
  4. Transformer attention mechanics: Understanding head_dim, num_attention_heads, and how Q/K/V projection shapes are computed from these parameters.
  5. The Kimi-K2.5 model: Knowing that this is a 1T-parameter MoE model with hidden_size=7168, and that the EAGLE-3 draft model is trained to predict its hidden states.
  6. The conversation history: The message sits at the end of a long debugging chain where the assistant resolved API incompatibilities between speculators v0.3.0 and vLLM 0.16, rewrote the training script, and got the pipeline working on 10 samples.

Output Knowledge Created

This message produces several valuable outputs:

  1. A structured comparison table: The three differences are clearly enumerated with specific weight shapes, providing a reference for future checkpoint validation.
  2. A decision to investigate: The task launched to examine vLLM's EAGLE-3 loading code will produce actionable knowledge about required naming conventions, optional weights, and config format.
  3. A corrected head_dim value: The discovery that head_dim=128 is needed (already patched in the previous message) ensures the next training run produces attention projections with the correct dimensions.
  4. A methodology for checkpoint validation: The message demonstrates a reusable pattern: compare against a reference, enumerate differences, investigate the inference engine's expectations, then adjust.

The Thinking Process

The reasoning in this message follows a clear arc:

  1. Recognition: The assistant recognizes that a checkpoint that trains successfully is not necessarily a checkpoint that loads successfully into vLLM.
  2. Comparison: The assistant fetches the reference checkpoint and performs a side-by-side comparison of weight keys and shapes.
  3. Categorization: Each difference is categorized by its nature (naming convention, optional weight, architectural parameter).
  4. Risk assessment: The naming difference is flagged as a "might cause compatibility issues" risk, while the embed_tokens difference is implicitly lower-risk.
  5. Investigation: Rather than acting on assumptions, the assistant launches a targeted code search in vLLM's source. The thinking is methodical and conservative. The assistant doesn't assume that matching the reference is necessary—it recognizes that the reference might have its own quirks. Instead, it goes to the source of truth: the inference engine's loading code.

Conclusion

Message [msg 2784] represents a critical inflection point in the EAGLE-3 training pipeline. It's the moment when the assistant shifts from "can I train a model?" to "will the trained model actually work?" The careful comparison against the AQ-MedAI reference, the identification of three specific differences, and the decision to investigate vLLM's loading code all demonstrate a disciplined approach to ML engineering. The message's true value lies not in any single answer it provides, but in the questions it asks—and the methodology it establishes for answering them. In the subsequent messages ([msg 2785]), the task results would confirm that vLLM accepts both naming conventions, that embed_tokens is optional, and that head_dim=128 is indeed the correct setting, validating the assistant's investigative approach and clearing the path for scaled-up training.