The Complete EAGLE-3 Pipeline: From Hidden State Extraction to Speculative Decoding Deployment

Introduction

In the span of a single opencode session segment, an AI assistant and a user accomplished what many ML engineering teams would consider a multi-week endeavor: they completed the full EAGLE-3 speculative decoding training pipeline for the Kimi-K2.5 language model, from hidden state extraction through distributed training optimization to SGLang deployment. The journey spanned approximately 130 messages (index 4220 through 4347) and involved surviving a VM crash and disk migration, extracting 37,312 hidden state samples (87.8 million tokens, ~4.6 TB), training a draft model across 5 epochs on 4 GPUs for ~10.8 hours, debugging a Triton shared-memory OOM on Blackwell GPUs, fixing a DataLoader batch-packing bug that unlocked 8× training efficiency, correcting SGLang argument names through iterative crash-and-inspect cycles, applying a weight key remapping fix for framework compatibility, and finally deploying the server with EAGLE3 speculation at 16 draft tokens — ready for the benchmarking that would validate the entire effort.

This article synthesizes the work across this chunk, tracing the narrative arc from data completion through deployment, examining the key technical decisions, debugging breakthroughs, and architectural insights that emerged along the way.

Part I: The Moment of Completion — Hidden State Extraction Crosses the Finish Line

The chunk opens with a moment of quiet triumph. After a catastrophic VM crash caused by the underlying Ceph cluster running out of disk space ([msg 4192]), the assistant had recovered the pipeline — killing an auto-started vLLM server, verifying data integrity of 18,421 previously extracted samples, re-applying the hidden state dump patch to SGLang, and resuming extraction. In message 4221, the assistant reports: "It's done! The extraction finished — it's now in the final step building sample_lengths.json." The log output shows samples 37,260 through 37,300, with the ETA reading "0 min" and an error count of zero since the resume point.

This moment represents the payoff of careful infrastructure design. The extraction script's resume capability — checking for existing .pt files before processing — meant that the crash did not waste the 18,421 samples already extracted. As analyzed in [2], the assistant's brief exclamation — "It's done!" — is the only indication that something significant has occurred, but it marks the pivot point between the uncertain recovery phase and the productive training phase.

However, the transition was not seamless. The assistant assumed that building sample_lengths.json would be a quick metadata step, but discovered in message 4223 that it involved scanning all 37,312 .pt files — totaling 3.5 TB — and loading each one to measure its sequence length. This led to a brief detour where the assistant investigated whether the file was even needed for training, ultimately determining that the training pipeline reads .pt files directly and doesn't require sample_lengths.json. The scan was killed, and the pipeline proceeded.

Part II: The Architecture of Frozen Knowledge — Understanding Parameter Freezing

Before training could begin, a critical architectural question needed resolution. In message 4231, the user asked: "Is anything frozen?" This seemingly simple question triggered a deep investigation into the EAGLE-3 training script's parameter freezing strategy.

The assistant examined the monkey-patched training script and discovered a deliberate design: embed_tokens (the embedding layer) and verifier_lm_head (the verifier's language modeling head) were frozen, while the draft model's lm_head, projection layer (fc), and decoder layer were kept trainable [14]. This matched the upstream speculators library convention exactly. The reasoning was sound: the drafter should learn to predict tokens using the verifier's hidden states as conditioning, but the verifier's own embedding and output layers should remain fixed to prevent the drafter from drifting away from the verifier's representation space.

The user's response in message 4235 — "Makes sense, proceed" — was the green light that launched the 100K-scale training run. As [16] notes, those three words carried the weight of the entire pipeline: the data was ready, the architecture was validated, and the user had signed off on the approach.

Part III: The First Breath of Training — Monitoring Early Convergence

With the launch command issued in message 4236 using torchrun --nproc_per_node=4 with a learning rate of 3e-5, cosine scheduler, 5% validation split, and 3 TTT steps, the assistant began monitoring. Message 4239 captures the first meaningful health check, taken roughly three minutes after launch.

The loss trajectory told an encouraging story. A total loss of ~28 across three TTT steps was "starting high" but "already dropping" — loss_0 from ~9.2 to ~8.9, loss_1 from ~9.1 to ~8.2, loss_2 from ~9.0 to ~8.0 — all at a warmup learning rate of just 1.9e-6 (barely 6% of the target). As [20] analyzes, this downward trajectory at a fraction of the target learning rate was a strong signal that the optimization landscape was well-behaved.

The GPU utilization probe revealed a healthy system: four GPUs running at 100% compute utilization with 28.6 GB of memory allocated and power draw of 227–297W out of 600W TDP. The remaining four GPUs sat idle in the P8 power state — exactly what one would expect from a --nproc_per_node=4 invocation on an 8-GPU node. The 100% utilization was particularly significant: it meant the GPUs were not bottlenecked by data loading, CPU preprocessing, or inter-GPU communication, validating the training configuration.

Part IV: The 35-Hour Miscalculation — Discovering the FSDP Bottleneck

The initial training run revealed a shocking estimate: 177,230 total steps across 5 epochs, implying a training time of approximately 35 hours. This was far longer than expected for a 37,312-sample dataset. The assistant diagnosed the issue in message 4242: FSDP (Fully Sharded Data Parallelism) was sharding the data across 4 GPUs, meaning each GPU only processed a quarter of the samples per epoch. Combined with the DataLoader's batch_size=1 and max_seq_len=8192, each step processed a single packed sequence — resulting in 35,446 batches per epoch.

As [24] and [25] document, this was a subtle interaction between FSDP's automatic data sharding and the training script's batch configuration. The assistant initially considered increasing max_seq_len to pack more tokens per batch, but this led down a rabbit hole of OOM errors and Triton compilation failures that would consume the next several hours of debugging.

Part V: The Triton Shared-Memory Barrier — A Compiler-Level OOM on Blackwell GPUs

The quest to increase GPU utilization through larger sequence lengths hit a wall at max_seq_len=16384. The training didn't crash with a familiar CUDA out-of-memory error — it crashed with something far more subtle: a Triton shared-memory overflow.

The error message read: "RuntimeError: No valid triton configs. OutOfMemoryError: out of resource: triton_per_fused__to_copy_add_div_expand_mul_pow_sum_view_1 Required: 163912 Hardware limit:101376." As [60] explains, this was not a GPU VRAM issue but a compiler-level failure: torch.compile with the Triton backend had generated a kernel configuration for RMSNorm that required 163,912 bytes of shared memory per block, but the SM120 architecture (Blackwell) only provides 101,376 bytes.

This distinction between HBM (VRAM) and shared memory (on-chip SRAM) is critical and often misunderstood. The assistant's diagnosis in message 4279 was precise: "This is a Triton autotuning issue on SM120, not a memory issue." The proposed workaround — disabling torch.compile — was pragmatic but had implications: it would break flex_attention, which is required for the block-diagonal masking in packed training. The assistant correctly chose not to pursue this path, instead stepping back to re-examine the actual data pipeline.

Part VI: The Packing Breakthrough — How Fixing Batch Size Unlocked 8× Efficiency

The breakthrough came when the assistant re-examined the collate function in the speculators library. The critical insight, documented across messages 4282-4288, was that the DataLoader's batch_size parameter controls how many samples are passed to the collate function. With batch_size=1, the collate function received exactly one sample per call — no packing occurred at all. The 35,446 batches/epoch were each processing a single sample, which explained the low GPU utilization: the GPUs were spending most of their time on overhead rather than computation.

The fix was elegant: increase batch_size to 8 while keeping max_seq_len at 8192. The collate function would receive 8 samples, concatenate them (~18K tokens total), then truncate/pad to 8192 tokens. This produced tightly packed batches where multiple shorter sequences filled the attention window, with block-diagonal masking keeping them independent.

Message 4289 captures the triumphant result: "Now we're talking." The numbers told the story:

Part VII: The Patient Wait — Training to Convergence

With the packing fix in place, the training ran for approximately 10.8 hours across 5 epochs. The assistant monitored periodically, checking GPU utilization, loss trajectories, and convergence patterns. Message 4313 captured the first validation milestone at epoch 1, showing full_acc_0 at 73.2% and an estimated acceptance length of 2.73 — already better than the previous 10K-sample drafter's 2.1.

The final validation metrics, reported in message 4321, showed a model that had converged well:

| TTT Step | Loss | Full Acc | Cond Acc | |----------|------|----------|----------| | 0 | 0.920 | 74.7% | 74.7% | | 1 | 1.919 | 50.6% | 67.0% | | 2 | 2.705 | 33.4% | 64.3% | | 3 | 3.298 | 22.0% | 63.2% | | 4 | 3.772 | 14.3% | 61.9% |

The estimated acceptance length was ~2.95 tokens, up from 2.73 at epoch 1 and 2.1 with the earlier 10K-sample model. As [102] notes, the assistant's qualitative observations — "no overfitting," "plateaued around epoch 3-4," "well-utilized" — reflected diagnostic thinking that went beyond mere metric reporting. The model had converged to its capacity limit given the data.

Part VIII: The Weight Key Fix — Bridging Training and Deployment

With training complete, the deployment phase began. The first obstacle was a weight key mismatch between the training framework and the inference engine. The speculators library saved weights under the prefix layers.0.*, but SGLang expected midlayer.*. This mismatch had caused a zero-acceptance-rate bug earlier in the session and was a known integration issue.

In message 4325, the assistant checked for the existing fix script on the remote machine. The output was simply "MISSING" — the script, stored in /tmp, had been wiped by the VM crash and reboot. The assistant wrote a new version locally, transferred it via scp in message 4327, and executed it in message 4328. The script loaded 16 tensors from the safetensors file, identified 10 keys that needed renaming, and performed the transformation. As [109] documents, this was a quintessential "last mile" problem: without this simple key rename, the entire training investment would have been worthless.

Part IX: The SGLang Argument Odyssey — Three Crashes to Correctness

Deploying the SGLang server with EAGLE3 speculation proved surprisingly treacherous. The first attempt crashed because --num-speculative-tokens was not a recognized flag. The assistant consulted the help output and discovered the correct name: --speculative-num-draft-tokens. The second attempt crashed with an AssertionError in _handle_speculative_decoding() because the validation logic required that --speculative-num-steps be explicitly set when using EAGLE3. The assistant traced through SGLang's source code to understand the three-parameter constraint [124].

Message 4337 captures the correction: "It's --speculative-num-draft-tokens, not --num-speculative-tokens. Also --speculative-num-steps controls how many draft rounds per verification." As [118] notes, this seemingly trivial typo — a flag name difference of a few characters — brought the entire deployment to a halt and required three crash-and-inspect cycles to resolve.

Part X: The Moment of Truth — Server Up and Benchmark Ready

After 890 seconds of server startup (including CUDA graph capture for the speculative decoding pathway), the SGLang server was finally healthy. Message 4346 reports: "Server is up (~15 min startup with CUDA graph capture for speculation). Now let me write a benchmark script and run it."

This message marks the transition from building to measuring. The entire pipeline — hidden state extraction, distributed training, weight key fixing, server configuration — had been a means to this end. The benchmark script, benchmark_eagle3.py, was designed to test multiple configurations (16, 10, and 5 draft tokens) against a no-speculation baseline, providing the data needed to determine whether the drafter actually accelerated inference.

The chunk concludes with an empty user message (4347) — a silence that, as [128] argues, speaks volumes. It signals trust, delegation, and forward momentum. The assistant responds by producing a comprehensive status document, re-establishing shared context before proceeding to the benchmarking that would validate the entire effort.

Conclusion: Themes and Lessons

Several themes emerge from this chunk's narrative arc. First, the data pipeline is the bottleneck — the single most impactful optimization was not a model architecture change or a hyperparameter tweak, but fixing the DataLoader's batch_size to enable proper sequence packing. Second, compiler-level failures require hardware-level understanding — the Triton shared-memory OOM on Blackwell GPUs could only be diagnosed by understanding the distinction between HBM and shared memory, and the specific constraints of the SM120 architecture. Third, the last mile is the hardest — weight key mismatches, argument name typos, and assertion failures consumed more debugging time than the training itself. Fourth, infrastructure resilience must be engineered — the VM crash could have been catastrophic, but the resume capability in the extraction script and the careful verification of data integrity saved days of work.

The EAGLE-3 drafter that emerged from this pipeline — trained on 37,312 samples across 5 epochs, achieving 74.7% validation accuracy and an estimated acceptance length of 2.95 tokens — represents not just a trained model, but a validated end-to-end system for speculative decoding on the Kimi-K2.5 architecture. The server was deployed and ready for benchmarking. The rest was measurement.