The Head_Dim Correction: A Pivotal Training Re-Run in the EAGLE-3 Pipeline

Introduction

In the course of building an EAGLE-3 speculative decoding pipeline for the massive Kimi-K2.5 language model (a 1-trillion-parameter Mixture-of-Experts architecture deployed across 8 NVIDIA Blackwell GPUs), a seemingly minor architectural detail — the head_dim parameter in the attention mechanism — threatened to create an invisible incompatibility between the trained draft model and the inference engine that would ultimately use it. Message 2789 in this coding session captures the moment when that detail was corrected and the training pipeline was re-executed, transforming a subtly broken checkpoint into one that matched the reference architecture.

This article examines that single message in depth: a bash command that launches a training run on a remote server, and the truncated output that confirms it is executing. Behind this simple action lies a chain of reasoning that spans architecture comparisons, HuggingFace config parsing, weight shape verification, and careful consideration of vLLM's model loading internals. Understanding this message reveals how the assistant navigated the gap between a library's default behavior and a reference implementation's explicit choices.

The Message

The subject message, message index 2789 in the conversation, reads as follows:

[assistant] [bash] ssh root@10.1.230.174 'CUDA_VISIBLE_DEVICES=0 /root/ml-env/bin/python3 /root/eagle3-train/04_train.py \
    --verifier-path /shared/kimi-k2.5-int4 \
    --data-dir /root/eagle3-train/data_test/hidden_states/rows_0-2000 \
    --vocab-mapping-dir /root/eagle3-train/data_test/vocab_mapping \
    --output-dir /root/eagle3-train/output_test \
    --epochs 3 \
    --lr 3e-5 \
    --max-seq-len 2048 \
    --ttt-steps 3 \
    --val-ratio 0.2 \
    --num-workers 0 2>&1'
/root/ml-env/lib/python3.12/site-packages/speculators/models/mlp.py:9: UserWarning: Field name "torch_dtype" in "MLPSpeculatorConfig" shadows an attribute in parent "SpeculatorModelConfig"
  class MLPSpeculatorConfig(SpeculatorModelConfig):
/root/ml-env/lib/python3.12/site-packages/torch/_inductor/compile_fx.py:321: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performa...

The output is truncated, showing only the initial warnings before the training progress would have appeared. But the key fact is that this command ran to completion successfully — the very next message (index 2790) confirms "Training works with head_dim=128."

Context: Why This Message Was Written

To understand why this message exists, we must trace backward through the preceding messages. The assistant had already completed a successful training run at message 2775, using the same 04_train.py script with 10 samples and 3 epochs. That run produced a 4.6 GB checkpoint and completed in approximately one minute on a single GPU. On the surface, everything worked.

But the assistant was not satisfied with mere execution success. It immediately began a thorough compatibility verification process, asking: Will this checkpoint actually work when loaded by vLLM for speculative decoding? This question drove a multi-step investigation spanning messages 2780 through 2788.

At message 2780, the assistant noticed a discrepancy: the saved config had head_dim = 112, which is the default computation from LlamaConfig (hidden_size / num_attention_heads = 7168 / 64 = 112). But the reference model from AQ-MedAI (the established EAGLE-3 implementation for Kimi-K2) used head_dim = 128. This is a non-standard choice — it means the attention projections map from 7168-dimensional hidden states to 8192-dimensional attention spaces (64 heads × 128 dims), an expansion rather than the typical identity projection.

At message 2781, the assistant confirmed this by loading the AQ-MedAI config directly from HuggingFace Hub and verifying the explicit head_dim: 128 setting. It then edited 04_train.py to override the default head_dim calculation.

At message 2782, the assistant inspected the weight shapes from the first training run and confirmed they used the smaller head_dim=112 projection shapes (e.g., q_proj.weight: [7168, 14336] instead of the expected [8192, 14336]).

At messages 2783–2786, the assistant conducted a deep investigation into vLLM's EAGLE-3 loading code via a subagent task, discovering that vLLM accepts both midlayer.* and layers.0.* weight naming conventions, and that the config format (nested speculators format vs. flat AQ-MedAI format) is handled by vLLM's SpeculatorsConfig parser.

At message 2787, the assistant weighed the trade-offs between the two config formats, decided that both would work, and pivoted to the more critical issue: re-running training with the corrected head_dim=128. It cleaned up the old test output and prepared for the re-run.

Message 2788 copied the updated script to the remote server. Then came message 2789: the re-run itself.

The Head_Dim Discovery: An Architectural Insight

The head_dim discrepancy is more than a simple parameter mismatch — it reflects a genuine architectural choice in the EAGLE-3 draft model design. In standard transformer attention, head_dim is implicitly hidden_size / num_attention_heads. For Kimi-K2.5's hidden size of 7168 and 64 attention heads, this gives 112. But the AQ-MedAI reference explicitly sets head_dim = 128, which means the QKV projections have output dimension 64 × 128 = 8192, which is larger than the hidden size of 7168.

This design choice has concrete implications:

• Wider attention projections: The query, key, and value projections each have 8192 output dimensions rather than 7168, requiring approximately 14% more parameters in those layers.
• Different representational capacity: The expanded attention dimension may provide better gradient flow or representational capacity for the draft model's task of predicting the verifier's next token distribution.
• Compatibility requirement: If the trained draft model uses head_dim=112 but vLLM's EAGLE-3 loader expects head_dim=128 (matching the reference), the weight shapes would be incompatible, causing a silent loading failure or incorrect behavior. The assistant's decision to match the reference architecture, rather than keeping the default, was a judgment call. The default (112) would have produced a smaller, faster draft model that might still work — but it would deviate from the only known working EAGLE-3 implementation for this model family. By matching AQ-MedAI's choices, the assistant maximized the probability of seamless integration with vLLM's speculative decoding engine.

Assumptions and Potential Issues

Several assumptions underpin this re-run:

1. The reference architecture is correct: The assistant assumes that AQ-MedAI's head_dim=128 is the intended design, not an artifact or mistake. Given that AQ-MedAI's model is a known working EAGLE-3 implementation for Kimi-K2, this is a reasonable assumption, but it is still an assumption — there is no official documentation explaining why 128 was chosen.
2. vLLM compatibility requires matching head_dim: The assistant assumes that vLLM's EAGLE-3 loader will check or depend on head_dim matching between the draft model config and the verifier config. If vLLM simply loads the weights by name without validating head_dim consistency, the mismatch might go undetected but cause subtle degradation.
3. The training will converge with the new architecture: Changing head_dim changes the parameter count and the initialization of the attention projections. The assistant implicitly assumes that the training hyperparameters (learning rate 3e-5, cosine scheduler, 3 epochs) that worked for head_dim=112 will also work for head_dim=128. The next message confirms this was correct, but it was not guaranteed.
4. The flat config format is not required: At message 2787, the assistant considered converting to AQ-MedAI's flat config format but decided the speculators nested format would work because vLLM has a handler for it. This is an assumption that would only be fully validated during actual inference testing.

Input Knowledge Required

To understand and execute this message, the assistant needed knowledge spanning multiple domains:

• Transformer attention mechanics: Understanding how head_dim, num_attention_heads, and hidden_size relate, and how they determine projection matrix shapes.
• HuggingFace Transformers internals: Knowing that LlamaConfig computes head_dim automatically from hidden_size / num_attention_heads, and how to override it.
• The speculators library API: Understanding how Eagle3SpeculatorConfig and Eagle3DraftModel use the underlying LlamaConfig, and where to inject the head_dim override.
• vLLM's model loading architecture: Knowing that vLLM has separate handling for EAGLE-3 speculators, that it accepts both naming conventions, and that it has a SpeculatorsConfig parser.
• Safetensors format: Being able to inspect weight files directly to verify shapes and compare against reference checkpoints.
• Remote execution patterns: Using SSH with CUDA_VISIBLE_DEVICES to run on a specific GPU, piping stderr to stdout for capture.

Output Knowledge Created

This message produced several concrete outputs:

1. A corrected checkpoint: The training run produced a new model.safetensors with head_dim=128 attention projections, verified in message 2790 to have the correct shapes (e.g., q_proj.weight: [8192, 14336] instead of [7168, 14336]).
2. Validation of the architecture change: The successful training run confirmed that the head_dim override did not break the training pipeline — the model converged and produced a valid checkpoint.
3. Confidence in vLLM compatibility: By matching AQ-MedAI's weight shapes exactly, the assistant could be confident that vLLM's EAGLE-3 loader would accept the checkpoint without shape mismatches.
4. A reusable training configuration: The 04_train.py script, now updated with the head_dim override, could be used for future training runs at scale without the same bug recurring.

The Thinking Process

The assistant's reasoning in the messages leading up to 2789 reveals a methodical, verification-first approach. After the initial training success, the assistant did not declare victory and move on. Instead, it asked a deeper question: "Will this actually work when deployed?" This led to a cascade of investigations:

First, it noticed the head_dim discrepancy by inspecting the saved config. Then it cross-referenced against the AQ-MedAI reference. Then it inspected the actual weight shapes to confirm the mismatch was real. Then it investigated vLLM's loading code to understand what format was expected. Only after all these checks did it decide to re-run.

This pattern — train, verify, discover mismatch, investigate root cause, fix, re-train — is characteristic of robust ML engineering. The assistant treated the first successful training not as an endpoint but as a hypothesis that needed validation against the deployment target.

Broader Significance

Message 2789 represents a microcosm of the challenges in building speculative decoding pipelines for large language models. The draft model is not an independent component — it must interface precisely with both the verifier model (whose hidden states it consumes) and the inference engine (which orchestrates the speculative decoding process). Every architectural parameter must be consistent across all three systems.

The head_dim correction also illustrates a common pitfall in ML engineering: relying on library defaults without checking whether they match the reference implementation. The speculators library's Eagle3SpeculatorConfig inherited LlamaConfig's default head_dim computation, which produced 112. But the reference implementation (AQ-MedAI) explicitly overrode this to 128. Without the assistant's diligence in cross-checking, this subtle mismatch would have propagated into the production checkpoint, potentially causing failures only at inference time — the worst possible moment for a bug to surface.

In the end, this message is about the invisible work of compatibility engineering: the checks, cross-references, and re-runs that ensure a trained model will actually work when deployed. The training itself is the easy part. Getting the architecture right is where the real engineering happens.