The Milestone: EAGLE-3 Training Succeeds for Kimi-K2.5

A Pivotal Moment in Speculative Decoding Development

In the course of a complex engineering session spanning multiple days, one message stands out as a genuine milestone: the moment when the assistant announces, with justified emphasis, "TRAINING WORKS!" This message ([msg 2776]) marks the first successful end-to-end run of the EAGLE-3 training pipeline for the Kimi-K2.5 model on an 8× Blackwell GPU machine. But this message is far more than a simple success notification—it is a carefully calibrated checkpoint that simultaneously celebrates a breakthrough, validates weeks of debugging, and sets the stage for the next phase of work.

The Context: Why This Message Was Written

To understand the significance of this message, one must appreciate the journey that led to it. The assistant had been working for days to deploy and optimize 1-trillion-parameter models on 8× NVIDIA RTX PRO 6000 Blackwell GPUs. After successfully deploying Kimi-K2.5 in INT4 format and profiling its inference bottlenecks, the team identified that AllReduce operations consumed 51.5% of decode time—a hardware limitation of PCIe-bound multi-GPU communication. Speculative decoding emerged as the most promising software-only optimization path, and EAGLE-3 was chosen as the draft model architecture.

The immediate preceding messages ([msg 2752] through [msg 2775]) reveal a grueling debugging session. The assistant had to:

  1. Reverse-engineer the speculators library API: The speculators library (v0.3.0) had its own Trainer class, Eagle3SpeculatorConfig, and Eagle3DraftModel—but the documentation was sparse. The assistant had to inspect source code, trace constructor calls, and understand the library's internal conventions.
  2. Monkey-patch for Kimi-K2.5's nested config: The Kimi-K2.5 verifier model uses a nested language_model prefix in its weight keys (language_model.model.embed_tokens.weight instead of model.embed_tokens.weight). The speculators library's _setup_embeddings_and_lm_heads method hardcoded the flat key names, requiring a surgical monkey-patch to extract weights correctly.
  3. Fix dtype mismatches: The hidden states extracted from the verifier were in bfloat16, but the EAGLE-3 draft model was being initialized in float32 by default. This caused a silent type mismatch that crashed training.
  4. Set the correct attention implementation: The EAGLE-3 draft model uses flex_attention with BlockMask for its causal masking, but the LlamaConfig created by the speculators library had _attn_implementation = None. The assistant had to explicitly set it to "flex_attention" to avoid a transformers routing error. Each of these fixes was applied iteratively—edit the script, copy it to the remote machine, run it, observe the error, debug, and repeat. The message at index 2776 is the first time the assistant sees the pipeline complete without an error.

What the Message Actually Contains

The message itself is deceptively simple. It opens with the bold declaration of success, then immediately transitions to a bash command that inspects the training output:

ssh root@10.1.230.174 'cat /root/eagle3-train/output_test/train_config.json'

The output shows a complete training configuration: 10 samples (8 train, 2 validation), 3 epochs, a learning rate of 3e-5, cosine scheduler, 24 total steps, and most importantly—2.5 billion parameters in the draft model. The checkpoint weighs 4.6 GB, consistent with a 2.5B-parameter model in bfloat16 precision.

The structure of this message reveals the assistant's disciplined engineering approach. Rather than simply declaring victory, it immediately pivots to verification. The success is provisional—it needs to be confirmed by inspecting the output artifacts. This pattern of "celebrate briefly, then verify rigorously" characterizes the entire session.

Input Knowledge Required

To fully understand this message, several pieces of background knowledge are essential:

EAGLE-3 Architecture: EAGLE-3 is a speculative decoding framework where a lightweight "draft" model predicts multiple future tokens in parallel, which are then verified by the full "verifier" model. The draft model is a single-layer transformer with a fusion layer (fc.weight) that combines hidden states from multiple verifier layers. The key hyperparameters include eagle_aux_hidden_state_layer_ids (which verifier layers to extract hidden states from), draft_vocab_size (the compressed vocabulary for draft predictions), and ttt_steps (the number of tokens to speculate per step).

The speculators Library: This is a third-party library that provides the EAGLE-3 implementation, including model definitions, configuration classes, and a training framework. The assistant had to learn its API through source code inspection, discovering that the Trainer class expects a pre-constructed model, a TransformTensors noise transform, and a data directory with pre-extracted hidden states.

Kimi-K2.5 Model Architecture: This is a 1-trillion-parameter Mixture-of-Experts model with a DeepSeekV2-style architecture. Its config is nested (with a language_model wrapper), which caused the weight extraction issues. The model uses a vocabulary of 163,840 tokens and special tokens like <| thinking |> (token 163606) and <| response |> (token 163607) for reasoning output.

The Hardware Context: The training ran on a single RTX PRO 6000 Blackwell GPU (96 GB memory), using CUDA 13.1 and PyTorch 2.9.1. The full 8-GPU machine was available, but the test deliberately used CUDA_VISIBLE_DEVICES=0 to constrain to one GPU.

Assumptions Made

The message embodies several assumptions, some explicit and some implicit:

The training pipeline is correct: The assistant assumes that a successful 10-sample, 3-epoch run validates the pipeline architecture. This is a reasonable assumption for a smoke test, but it doesn't guarantee convergence or generalization to the full dataset.

The checkpoint format is usable: The 4.6 GB model.safetensors file is assumed to be in the correct format for vLLM integration. This assumption is tested in the very next messages ([msg 2780]–[msg 2786]), where the assistant discovers that the head_dim parameter is 112 instead of the expected 128, and that the weight key naming (layers.0.* vs midlayer.*) may differ from the AQ-MedAI reference.

Single-GPU training is sufficient: The test uses only one GPU, but production training would likely benefit from data parallelism across multiple GPUs. The assistant's assumption is that the pipeline works correctly on one GPU and can be scaled later.

The speculators library API is stable: The assistant assumes that the Trainer class, Eagle3SpeculatorConfig, and Eagle3DraftModel will continue to work as observed. This is a reasonable assumption for a local development environment, but API changes in future library versions could break the pipeline.

Mistakes and Incorrect Assumptions

The most significant issue lurking beneath this success message is the head_dim discrepancy. The assistant later discovers ([msg 2780]) that the AQ-MedAI reference model uses head_dim=128 with num_attention_heads=64, which implies an attention dimension of 64 × 128 = 8192—larger than the hidden_size of 7168. This means the Q/K/V projections project up from 7168 to 8192 dimensions. The assistant's test run used the default head_dim = hidden_size / num_attention_heads = 7168 / 64 = 112, which gives attention dimension 64 × 112 = 7168—matching the hidden size exactly. This architectural difference could affect the draft model's capacity and compatibility with the AQ-MedAI reference.

Another subtle issue is the weight key naming convention. The speculators library saves weights under layers.0.* keys, while the AQ-MedAI reference uses midlayer.*. The assistant later confirms ([msg 2785]) that vLLM accepts both naming conventions by remapping midlayer to layers.0, but this is a discovered fact rather than an assumption validated at the time of message 2776.

The assistant also assumes that 10 samples are sufficient for a meaningful test. While this is adequate for catching pipeline errors (dtype mismatches, shape errors, API issues), it doesn't test for convergence, overfitting, or generalization. The training loss on 10 samples with 3 epochs is likely near-zero for the training set, which could mask issues with the noise transform or data preprocessing.

Output Knowledge Created

This message creates several important pieces of knowledge:

  1. Pipeline validation: The EAGLE-3 training pipeline for Kimi-K2.5 is confirmed to work end-to-end. The 04_train.py script correctly loads the verifier model, initializes the draft model, loads pre-extracted hidden states, applies noise transforms, and runs the speculators Trainer.
  2. Training metrics: The run completes 3 epochs on 10 samples in approximately 1 minute on a single Blackwell GPU. This gives a baseline for scaling estimates: 1,000 samples at 10 epochs would take roughly 33 minutes (linear scaling), and the full dataset would take proportionally longer.
  3. Checkpoint characteristics: The output is a 4.6 GB model.safetensors file containing 16 weight tensors including embed_tokens.weight, fc.weight, attention projections, MLP layers, and the d2t (draft-to-target) vocabulary mapping buffer. The config is saved in HuggingFace-compatible format with architectures: ["Eagle3DraftModel"] and the critical eagle_aux_hidden_state_layer_ids: [2, 30, 58] field.
  4. Resource requirements: The training uses approximately 96 GB of GPU memory (the full capacity of one RTX PRO 6000) and cleans up properly after completion, as confirmed in the subsequent message ([msg 2778]).

The Thinking Process Revealed

The message's structure reveals the assistant's mental model. The opening "TRAINING WORKS!" is an emotional release—a moment of satisfaction after hours of debugging. But the very next action is to run a bash command to inspect the output. This reveals a hypothesis-testing mindset: the assistant doesn't trust the success until the artifacts are verified.

The choice of which artifact to inspect first—train_config.json—is telling. The assistant wants to confirm that the training hyperparameters were correctly recorded, that the data split (8 train, 2 validation) is correct, and that the total parameter count (2.5B) matches expectations. This is a sanity check before diving into the model weights themselves.

The message also serves as a communication checkpoint for the user. By announcing success and immediately showing the output config, the assistant invites the user to verify the results and make decisions about next steps. This conversational rhythm—announce, show evidence, ask implicitly for direction—characterizes the entire session.

The Aftermath

The messages immediately following ([msg 2777]–[msg 2786]) show the assistant executing a systematic verification protocol: checking the checkpoint file listing, verifying GPU memory cleanup, updating the todo list, and then—critically—comparing the output against the AQ-MedAI reference model. This comparison reveals the head_dim issue and leads to a fix in the training script.

The user then redirects the approach ([msg 2787] onward) toward generating higher-quality synthetic training data by capturing Kimi-K2.5's actual reasoning outputs through the vLLM inference server. This pivot acknowledges that the 10-sample test, while validating the pipeline, doesn't produce a useful draft model—the real value comes from training on thousands of high-quality samples that capture the model's actual reasoning patterns.

Conclusion

Message 2776 is a milestone that encapsulates the engineering discipline of the entire session: celebrate the breakthrough, but immediately verify, document, and prepare for the next challenge. The EAGLE-3 training pipeline works, but the work is far from over. The head_dim must be fixed, the checkpoint must be validated against vLLM, and most importantly, the training data must be dramatically expanded and improved. This message marks the transition from "can we train?" to "can we train well?"—a shift from infrastructure validation to quality optimization.