The Orchestrator: Tying Together an EAGLE-3 Training Pipeline for Kimi-K2.5

Message Overview

The subject message (msg id=2528) is brief on the surface — the assistant writes a single file, run_pipeline.sh, described as "the orchestrator script that ties everything together." But this message represents the culmination of a significant engineering effort: the construction of a complete EAGLE-3 speculative decoding training pipeline for the 1-trillion-parameter Kimi-K2.5 model running on 8× RTX PRO 6000 Blackwell GPUs.

[assistant] Now let me write the orchestrator script that ties everything together and the run script:
[write] /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/run_pipeline.sh
Wrote file successfully.

Context: Why This Message Was Written

To understand this message, one must understand the journey that led to it. The session had spent dozens of hours deploying, profiling, and optimizing large language models on an 8-GPU Blackwell machine. The most recent profiling campaign ([msg 2471] through [msg 2504]) had revealed a stark bottleneck: AllReduce communication consumed 51.5% of decode time for the Kimi-K2.5 INT4 model. This was a hardware-imposed limit — the machine's PCIe topology simply could not move activation data between GPUs fast enough for dense allreduce operations during each decoding step.

The user and assistant then pivoted to investigate speculative decoding as a software-only optimization that could bypass this bottleneck. The reasoning was elegant: if a smaller "draft" model could generate multiple candidate tokens cheaply on a single GPU, and the large "target" model only needed to verify them in one forward pass, the number of expensive allreduce-bound verification steps would drop dramatically.

The assistant launched parallel research agents ([msg 2507]) covering speculative decoding fundamentals, framework support in vLLM and SGLang, candidate draft models, and training feasibility. The research was thorough: it examined the SpecForge training framework from the SGLang project, the Speculators library from the vLLM project, and the existing AQ-MedAI/Kimi-K2-Instruct-eagle3 draft model on HuggingFace. The key findings were sobering:

1. N-gram speculation (which requires no training) was tested empirically and found to be 9–26% slower than baseline — because the MoE (Mixture-of-Experts) architecture of Kimi-K2.5 means that even verification passes activate all experts, making the overhead of running the draft model negligible compared to the cost of the verification pass itself.
2. The only off-the-shelf EAGLE-3 draft model, AQ-MedAI/Kimi-K2-Instruct-eagle3, was trained for K2, not K2.5. Since K2.5 has a different tokenizer and vocabulary, acceptance rates would be degraded.
3. The most promising path was the Baseten approach: training a custom EAGLE-3 "head" — a small transformer that predicts the next N tokens' hidden states given the target model's penultimate-layer hidden states. The user then directed the assistant to begin implementing the training pipeline on the existing hardware ([msg 2505]), with the understanding that the "hero run" would be ported to rented B200/B300 machines. The assistant chose the Speculators library (from the vLLM project) over SpecForge, reasoning that since the deployment was already on vLLM, the Speculators pipeline would integrate more cleanly.

What Was Built Before This Message

The assistant constructed the pipeline incrementally across messages 2523–2528:

1. draft_config.json (msg 2523): A configuration file for the EAGLE-3 draft model, matching the K2 architecture with a 32K-token draft vocabulary, 4 layers, and specific feature dimensions.
2. 01_prepare_dataset.py (msg 2524): A dataset preparation script that loads a HuggingFace dataset, tokenizes it using the Kimi-K2.5 tokenizer, and saves tokenized sequences along with a token frequency distribution. This script runs without GPUs.
3. 02_extract_hidden_states.py (msg 2525): The core data generation script. It uses Speculators' VllmHiddenStatesGenerator to load the full target model via vLLM's MultiprocExecutor, run forward passes on the prepared data, and capture hidden states from specific intermediate layers (layers 2, 30, 58, 60 — matching the K2 EAGLE-3 configuration). This is the GPU-bound step and the most complex piece.
4. 03_build_vocab_mapping.py (msg 2526): A vocabulary mapping script that creates a mapping between the target model's full vocabulary (163,840 tokens) and the draft model's reduced vocabulary (32,000 tokens), selecting the most frequent tokens from the training data.
5. 04_train.py (msg 2527): The training script itself, which uses Speculators' Eagle3Trainer to train the draft model on the extracted hidden states. It supports DeepSpeed ZeRO-3, gradient checkpointing, and wandb logging.
6. run_pipeline.sh (msg 2528, the subject): The orchestrator that ties all five pieces together into a single end-to-end pipeline.

The Orchestrator's Role

The orchestrator script is the glue. It sequences the pipeline steps, passes the correct paths between them, and handles the critical server lifecycle: stopping the vLLM production server before hidden state extraction (which needs all GPU memory), then restarting it afterward. Without this script, running the pipeline would require manual intervention at each step — a recipe for errors when the hero run is eventually ported to rented hardware.

The assistant's comment "Now let me write the orchestrator script that ties everything together and the run script" reveals the thinking: the individual pieces were built, but they needed a conductor. The orchestrator encodes the workflow as a reproducible, automatable sequence — essential for the eventual port to B200/B300 machines where the user won't have the same interactive debugging environment.

Assumptions and Potential Issues

Several assumptions underpin this message and the pipeline it orchestrates:

Assumption 1: Speculators' VllmHiddenStatesGenerator works with vLLM 0.16. This turned out to be incorrect. The Speculators library (version 0.3.0) was designed for vLLM ≤0.15, and vLLM 0.16 introduced API changes — SchedulerConfig now requires is_encoder_decoder and max_model_len parameters, the supports_eagle3 protocol check fails on multimodal wrapper models, and KV cache utility APIs changed. The assistant would spend the next several messages (2529–2558) patching Speculators to work around these incompatibilities.

Assumption 2: The Kimi-K2.5 model's inner DeepseekV3ForCausalLM implements the SupportsEagle3 interface. While this was likely true for the inner model, the outer KimiK25ForConditionalGeneration multimodal wrapper did not expose the required interface. The custom_worker's supports_eagle3(model) check failed because it checked the wrapper, not the inner text model. This required patching the worker to drill into model.language_model.model.layers instead of model.model.layers.

Assumption 3: The existing hardware (8× RTX PRO 6000) is sufficient for a test run. With 921GB free on /shared/ and 765GB on /root/, storage was ample. But the 27-minute model loading time and the eventual training duration on modest data (10 samples) would reveal that the pipeline works but is slow — acceptable for a test, but the hero run would indeed need to be ported to faster hardware.

Assumption 4: The LSP errors shown in the message are irrelevant. The diagnostics block shows errors from server_args_sm120.py — a file unrelated to the pipeline being built. These are stale errors from a previous SGLang configuration file that the IDE's LSP server is still tracking. The assistant correctly ignores them.

Input Knowledge Required

To understand this message, one needs knowledge of:

• Speculative decoding and the EAGLE-3 architecture: how a small draft model predicts multiple candidate tokens from the target model's hidden states, and how the target model verifies them in parallel.
• The vLLM ecosystem: the relationship between vLLM (the serving framework), Speculators (the training library for draft models), and the internal worker extension API used for hidden state capture.
• The Kimi-K2.5 model architecture: a 1T-parameter MoE model with a multimodal wrapper (KimiK25ForConditionalGeneration) wrapping a text backbone (DeepseekV3ForCausalLM).
• The hardware topology: 8× RTX PRO 6000 Blackwell GPUs connected via PCIe, where AllReduce is the dominant bottleneck.
• The training data pipeline: how hidden states are extracted from intermediate layers of the target model during prefill, saved as .pt files, and used to train the draft model's prediction head.

Output Knowledge Created

This message and the pipeline it completes create:

1. A reproducible EAGLE-3 training pipeline for Kimi-K2.5, documented in the eagle3-train/ directory with all scripts and configuration.
2. A draft model configuration that can be used to initialize the EAGLE-3 head with the correct architecture (4 layers, 32K draft vocabulary, specific hidden dimensions matching the target model).
3. The orchestrator script that encodes the workflow as a single command, enabling the pipeline to be ported to rented hardware without manual step-by-step execution.
4. Documentation of the approach in next-steps-eagle.md, capturing the reasoning, the architecture decisions, and the plan for the hero run.

The Thinking Process

The assistant's thinking is visible in the progression of messages. After the profiling session revealed the AllReduce bottleneck, the assistant didn't immediately start coding — it researched. It launched parallel agents to understand both SpecForge and Speculators, examined the existing K2 EAGLE-3 model, and empirically tested n-gram speculation. Only after this research phase did it begin building.

The choice of Speculators over SpecForge was a deliberate architectural decision: "I'll use Speculators (vLLM project) since we're already on vLLM and it has a cleaner pipeline." This reveals a preference for integration simplicity over feature completeness.

The incremental construction — config first, then data preparation, then extraction, then mapping, then training, then orchestrator — shows a bottom-up approach. Each piece was testable independently (steps 1 and 3 don't need GPUs), and the orchestrator was written last to sequence them.

The LSP errors in the output are a reminder of the messy reality of development: the assistant is working in an environment with pre-existing code issues (the server_args_sm120.py file has syntax errors from earlier experimentation), but it correctly focuses on the task at hand rather than getting distracted by unrelated diagnostics.

Aftermath

The subsequent messages (2529–2558) would test this pipeline end-to-end. Steps 1 and 3 succeeded immediately. But step 2 (hidden state extraction) hit the vLLM 0.16 API incompatibilities. The assistant patched the Speculators code — first the trust_remote_code tokenizer issue, then the SchedulerConfig parameter change, then the supports_eagle3 protocol bypass for the multimodal wrapper. Each patch got the pipeline one step further, but new errors emerged deeper in the KV cache utility APIs. By the end of the chunk, the pipeline was partially working but still blocked by API mismatches between Speculators 0.3.0 and vLLM 0.16.

This message, then, stands at a pivot point: the completion of the pipeline's initial construction, just before the reality of software integration would test every assumption.