The Training Script That Almost Worked: Building an EAGLE-3 Pipeline for Kimi-K2.5

A Single Message in an OpenCode Coding Session

"Now the training script — this uses speculators' trainer directly"

This unassuming line, buried in a long session of ML infrastructure work, marks the moment when an ambitious speculative decoding pipeline reached its architectural apex. The message — <msg id=2527> — is the assistant writing 04_train.py, the fourth and most critical component of a five-script EAGLE-3 training pipeline for the 1-trillion-parameter Kimi-K2.5 model running on 8× RTX PRO 6000 Blackwell GPUs. It is simultaneously a moment of triumph and a harbinger of the incompatibility issues that would soon derail the effort.

The Message Itself

The complete message reads:

[assistant] Now the training script — this uses speculators' trainer directly: [write] /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/04_train.py Wrote file successfully.

>

LSP errors detected in other files: <diagnostics file="/home/theuser/glm-kimi-sm120-rtx6000bw/source/server_args_sm120.py"> ERROR [1:1] Unexpected indentation ...

The message is deceptively short. It reports a single action — writing a file — and then displays LSP diagnostics from an unrelated file. But this brevity conceals the weight of the decision behind it. The assistant is committing to a specific architectural strategy for the most performance-critical component of the pipeline: the training loop itself.

Context: Why This Message Exists

To understand why this message was written, we must trace back through the session's narrative arc. The user and assistant had been profiling the Kimi-K2.5 INT4 model on 8× Blackwell GPUs and discovered that AllReduce communication was consuming 51.5% of decode time — a bottleneck that no amount of GEMM optimization could fix. The user then pivoted to speculative decoding as a software-only path to higher throughput.

Three parallel research agents were dispatched ([msg 2507]) to investigate:

1. SpecForge — SGLang's training framework for EAGLE-3
2. Speculators — vLLM's official library for speculative decoding training and inference
3. AQ-MedAI/Kimi-K2-Instruct-eagle3 — an existing EAGLE-3 draft model trained for the predecessor model K2 The research revealed a clear landscape. The existing K2 EAGLE-3 draft model would have lower acceptance rates on K2.5 due to architectural differences. N-gram speculation was empirically proven to be 9–26% slower than baseline on MoE models due to expert activation overhead during verification. The most promising path was the "Baseten approach": training a custom EAGLE-3 head on the target model's own hidden states. The user then issued the directive that set this message in motion ([msg 2505]):

"Start implementing the training scripts, on our existing machine with lowered numbers, then should be easy to port to much more expensive b300 machine for a simple train run. Write down notes as you go"

This instruction carried several implicit assumptions: that the pipeline could be developed and tested on the existing 8× RTX PRO 6000 hardware (with reduced data sizes), that it would be portable to B200/B300 machines with minimal changes, and that the training scripts themselves would be the primary deliverable rather than the trained model.

The Architecture Decision: Why Speculators?

The assistant had to choose between SpecForge (from the SGLang project) and Speculators (from the vLLM project). The decision was made in [msg 2509]:

"I'll use Speculators (vLLM project) since we're already on vLLM and it has a cleaner pipeline."

This was a pragmatic choice driven by the existing environment. The machine already had vLLM 0.16.0rc2 running the Kimi-K2.5 service. Using Speculators meant staying within the same ecosystem, reusing the same model loading code, and avoiding the complexity of integrating a second framework. The assistant had already verified that speculators could be installed (pip install speculators succeeded in [msg 2512]) and that its core components imported correctly (Eagle3DraftModel and VllmHiddenStatesGenerator both loaded without errors in <msg id=2513-2515>).

However, this decision carried a hidden risk. Speculators was designed for vLLM ≤0.15, and the installed version was 0.16.0rc2. The assistant acknowledged this tension in [msg 2511]:

"speculators requires vllm<=0.15.0 for data generation which is a problem — our vLLM is 0.16.0. Let me check if we can work around this"

The assistant's strategy was to isolate the risk: use speculators' training infrastructure (which doesn't depend on vLLM internals) directly, while writing custom data generation code that could work around the vLLM version mismatch. This is precisely what 04_train.py represents — the part of the pipeline that should work because it only depends on speculators' pure-PyTorch training code, not on its vLLM integration layer.

The Pipeline Architecture

By the time message 2527 is written, the assistant has already created three preceding scripts:

1. draft_config.json ([msg 2523]) — The draft model configuration, matching the K2 EAGLE-3 architecture with a 32K draft vocabulary, 4 layers, 8 KV heads, and a hidden dimension of 7168. This configuration defines the "student" model that will learn to predict the target model's hidden states.
2. 01_prepare_dataset.py ([msg 2524]) — Dataset preparation using HuggingFace's datasets library. This script tokenizes raw text data, splits it into training sequences, and saves the tokenized output for later processing.
3. 02_extract_hidden_states.py ([msg 2525]) — The core data generation script. This is where the assistant anticipated the most difficulty. It uses speculators' VllmHiddenStatesGenerator to run the target model (Kimi-K2.5) on the prepared data and capture the hidden states from intermediate layers. These hidden states become the training targets for the draft model.
4. 03_build_vocab_mapping.py ([msg 2526]) — Vocabulary mapping between the target model's vocabulary and the draft model's 32K sub-vocabulary. EAGLE-3 uses a smaller "draft vocabulary" that the head predicts, and this mapping bridges the gap. And now, 04_train.py — the script that ties everything together by loading the prepared hidden states and training the draft model using speculators' Eagle3Trainer.

What the Training Script Does

The training script (04_train.py) is the payoff for all the preceding work. It:

1. Loads the draft model configuration from draft_config.json, instantiating an Eagle3DraftModel with the specified architecture.
2. Loads the prepared training data — the .pt files containing hidden states, input IDs, and vocabulary mappings saved by the previous pipeline steps.
3. Configures the trainer using speculators' Eagle3Trainer, which wraps PyTorch's training loop with EAGLE-3-specific loss functions (typically a combination of cross-entropy loss for token prediction and a regression loss for hidden state prediction).
4. Runs the training loop with configurable batch size, learning rate, number of epochs, and distributed training settings (DeepSpeed for multi-GPU). The key architectural decision captured in this message is the choice to use speculators' trainer directly rather than writing a custom training loop. This was a bet on API stability — the assumption that speculators' training code, being pure PyTorch without vLLM dependencies, would be immune to the version compatibility issues plaguing the data generation step.

The LSP Diagnostics: A Window into the Development Environment

The message includes an interesting artifact: LSP errors from an unrelated file (server_args_sm120.py). These diagnostics are displayed because the editor's language server is running in the project root and reports errors across all open files. The file in question is a SGLang server configuration that has pre-existing syntax errors (unexpected indentation, unclosed brackets, unresolved imports).

This detail reveals several things about the development environment:

• The assistant is working in a code editor with LSP integration (likely VS Code or a similar IDE)
• The project directory contains artifacts from multiple experiments (SGLang, vLLM, custom scripts)
• The LSP errors are noise — they come from an abandoned configuration file, not from the training script being written
• The assistant is comfortable ignoring these diagnostics, focusing instead on the successful file write This is a realistic picture of ML engineering work: messy project directories, stale configuration files, and the constant hum of LSP warnings that must be filtered out to focus on the task at hand.

Assumptions and Their Consequences

The message and its surrounding context reveal several assumptions, some of which would prove incorrect:

Assumption 1: Speculators' trainer would be API-compatible. The assistant assumed that because the trainer doesn't depend on vLLM internals, it would work with vLLM 0.16. This turned out to be partially true — the trainer itself loaded, but the data format produced by the custom extraction scripts didn't perfectly match what the trainer expected.

Assumption 2: The pipeline could be tested end-to-end with 10 samples. The assistant planned to validate the full pipeline with a tiny dataset (mlabonne/open-perfectblend, 10 samples, sequence length 512) before scaling up. This was sound engineering practice, but it meant that certain integration issues (like the hidden state extraction failing at runtime) would only surface during testing.

Assumption 3: The B300 port would be straightforward. The user's directive assumed that scripts developed on 8× RTX PRO 6000 would "be easy to port to much more expensive b300 machine." This ignored potential differences in CUDA capability (SM120 vs SM100), memory bandwidth, and NCCL topology that could affect both data generation and training performance.

Assumption 4: The draft model architecture from K2 would work for K2.5. The assistant copied the EAGLE-3 architecture from the K2 draft model (4 layers, 7168 hidden dimension, 32K draft vocab). While this was a reasonable starting point, K2.5 has a different hidden state distribution, and the optimal draft model depth and width might differ.

The Thinking Process

The assistant's reasoning, visible across the sequence of messages leading to 2527, follows a clear pattern:

1. Research first: Before writing any code, the assistant dispatched three parallel research agents to understand the available frameworks, their architectures, and the existing draft model.
2. Choose the path of least resistance: Speculators was chosen over SpecForge because it was already compatible with the vLLM ecosystem in use. The assistant explicitly noted the "cleaner pipeline" as a deciding factor.
3. Isolate risk: The assistant recognized the vLLM version incompatibility and designed the pipeline to isolate it. The data generation step (which depends on vLLM internals) was separated from the training step (which doesn't). This way, even if data generation needed custom workarounds, the training step could remain clean.
4. Build modularly: Each script in the pipeline has a single responsibility — dataset prep, hidden state extraction, vocab mapping, training. This modularity makes it possible to test each step independently and swap out components if needed.
5. Test small, scale later: The pipeline was designed to work with a tiny test dataset (10 samples, 512 tokens each) before scaling to the full training corpus. This is textbook ML engineering practice.

Input Knowledge Required

To understand this message, one needs:

• EAGLE-3 architecture: Understanding that EAGLE-3 is a speculative decoding framework where a small "draft" model predicts hidden states of a large "target" model, enabling speculative token generation with verification.
• Speculators library: Knowledge that speculators is vLLM's official library providing both data generation utilities and training infrastructure for EAGLE-3 and other speculative decoding methods.
• The session's history: Awareness that the team had been profiling Kimi-K2.5, discovered AllReduce as the primary bottleneck, and pivoted to speculative decoding as the optimization path.
• The hardware context: Understanding that 8× RTX PRO 6000 Blackwell GPUs with PCIe interconnect create a communication-bound environment where speculative decoding's compute-verification tradeoff is particularly attractive.

Output Knowledge Created

This message creates:

• 04_train.py: A complete training script that can train an EAGLE-3 draft model on Kimi-K2.5 hidden states. This script is the core deliverable of the pipeline — the component that will eventually produce the draft model weights.
• An architectural precedent: The decision to use speculators' trainer directly establishes a pattern for the pipeline. Future developers extending this work will know that the training component is a thin wrapper around the speculators library.
• A testable artifact: The script, combined with the preceding pipeline steps, enables end-to-end testing of the training pipeline with small datasets, validating the data format and training configuration before committing to a full-scale run.

The Broader Significance

Message 2527 represents the moment when abstract planning crystallizes into executable code. The assistant had spent several messages researching, installing dependencies, and designing the pipeline architecture. Now, with the writing of 04_train.py, the pipeline becomes concrete. The training script is the component that justifies all the preceding work — without it, the draft model config, the dataset prep, the hidden state extraction, and the vocab mapping are all preparation for nothing.

Yet this message also carries the seeds of the pipeline's eventual difficulties. The same version incompatibility that the assistant had identified as a risk for data generation would also manifest in the training step, albeit in subtler ways. The speculators trainer, while nominally independent of vLLM, expects data in a specific format that the custom extraction scripts may not perfectly reproduce. The API surface area between pipeline components is where integration bugs hide.

In the broader narrative of the session, this message is the calm before the storm. The next messages would show the pipeline being tested end-to-end, with dataset preparation and vocabulary mapping succeeding, but hidden state extraction failing at runtime due to the vLLM 0.16 API mismatch. The assistant would patch the speculators code to handle Kimi-K2.5's multimodal wrapper architecture (model.language_model.model.layers instead of model.model.layers), but further KV cache utility API mismatches would remain unresolved by the chunk's end.

The training script itself — 04_train.py — would never be executed in this session. It sits in the eagle3-train/ directory as a promise unfulfilled, waiting for the data generation step to be fixed so that it can receive the hidden states it needs. This is the nature of pipeline development: each component depends on the one before it, and a single broken link halts the entire chain.

Conclusion

Message 2527 is a study in the tension between architectural vision and integration reality. The assistant made a sound decision — use speculators' trainer directly, isolate the vLLM version risk, build modularly — but the decision was made on incomplete information. The speculators library, while well-designed for its intended use case, had not been tested with vLLM 0.16 or with the Kimi-K2.5 model architecture. The assistant's confidence that "this uses speculators' trainer directly" would be tested against the messy reality of API drift, architectural quirks, and the inevitable gap between what a library promises and what it delivers in an unanticipated context.

The message stands as a monument to the most difficult part of ML engineering: not the algorithms, not the math, but the integration — the thousand small decisions about which library to use, which version to trust, and which assumptions to make, all of which can be rendered invalid by a single API change in a dependency.