The Hidden State Extraction: A Pivotal Milestone in EAGLE-3 Training

Introduction

In the course of building a speculative decoding pipeline for the 1-trillion-parameter Kimi-K2.5 language model on an 8× NVIDIA RTX PRO 6000 Blackwell GPU server, the assistant reached a critical juncture: the hidden state extraction step for the EAGLE-3 draft model training dataset. Message [msg 2819] captures the moment this extraction completed—a deceptively simple status report that represents the successful culmination of hours of debugging, resource management, and careful pipeline orchestration. The message reads:

Extraction completed successfully! 1000 samples, 502,983 tokens extracted in 2.9 minutes (after the 22.5 minute model load). 5.8 samples/s, 2912 tok/s.

This single message is a milestone marker. It signals that the most expensive and error-prone step of the EAGLE-3 training pipeline—loading a 547 GB verifier model across 8 GPUs and running forward passes to capture internal hidden states—has been executed at scale. Understanding why this message matters, what decisions it reflects, and what knowledge it creates requires tracing the path that led to it.

The Strategic Context: Why Extract Hidden States?

EAGLE-3 (Eagle Auto-regressive Generation with Linear-time Extraction) is a speculative decoding framework that accelerates inference by training a lightweight "draft" model to predict the next several tokens in parallel, using the hidden states from the main "verifier" model as conditioning signals. The draft model is small—only 2.5 billion parameters compared to the verifier's 1 trillion—but it needs to learn how the verifier's internal representations evolve across token positions.

The training data for this draft model is not text in the conventional sense. It consists of hidden states: the intermediate activations at specific transformer layers of the verifier model, captured while the verifier processes real text. For Kimi-K2.5, the assistant configured extraction at layers 2, 30, 58, and 60 (the final layer before the LM head), producing a rich 4-layer snapshot of the model's internal computation for each token position.

The extraction pipeline had four stages, as the assistant laid out in [msg 2801]:

  1. Prepare 1000 samples from the mlabonne/open-perfectblend dataset (CPU-only, minutes)
  2. Build a vocabulary mapping between the verifier's 163,840-token vocabulary and the draft model's 32,000-token vocabulary (CPU-only, fast)
  3. Extract hidden states by loading the 547 GB verifier model on all 8 GPUs and running forward passes (~22 min load + extraction time)
  4. Train the draft model on the extracted hidden states (single GPU, minutes) Step 3 was the bottleneck—the one that required all 8 GPUs, consumed 90+ GB of memory per GPU, and took over 22 minutes just to load the model weights. Message [msg 2819] reports the successful completion of this bottleneck step.

The Path to This Message: Debugging and Recovery

The extraction did not succeed on the first attempt. The assistant's earlier attempt in [msg 2813] failed with a CUDA out-of-memory (OOM) error:

RuntimeError: Worker failed with error 'CUDA out of memory. Tried to allocate 896.00 MiB. 
GPU 0 has a total capacity of 94.97 GiB of which...

The root cause was an overly aggressive batch size. The assistant had set --batch-size 2000, intending to process all 1000 samples in a single batch. But "batch" in this context means the number of samples to prefill simultaneously—and with 503,000 total tokens across 1000 samples, a single batch would attempt to prefill all of them at once, far exceeding the ~4 GB of free memory remaining after the model itself occupied ~91 GB per GPU.

This was a genuine mistake, but it reveals an important assumption: the assistant assumed that --batch-size controlled the number of samples to process sequentially (like a queue depth), when in fact it controlled the parallelism of the forward pass. The earlier 10-sample test had worked with the default batch_size=4, and the assistant had extrapolated incorrectly.

The recovery was methodical. The assistant killed the OOM'd process ([msg 2814]), verified that all GPUs were freed, cleaned up the partial output directory ([msg 2815]), and re-launched with --batch-size 4 and --gpu-memory-utilization 0.85 ([msg 2816]). This time, each batch would contain at most 4 samples (up to 8,192 tokens), which fit comfortably in the available memory.

The Thinking Process: Monitoring and Verification

The assistant's thinking process during the extraction is visible in the sequence of monitoring commands. After launching the extraction, the assistant waited 25 minutes ([msg 2817]: sleep 1500) and checked the log, finding the model 86% loaded. Then, after another 2 minutes ([msg 2818]: sleep 120), the extraction was already running at 2910 tok/s with 972 samples completed. The assistant was clearly watching the progress bar, computing throughput rates, and planning the next steps in parallel.

The throughput numbers are worth examining. The extraction ran at 2912 tokens per second across 8 GPUs with tensor parallelism. For a 1-trillion-parameter MoE model, this is remarkably efficient—it means the forward pass for 503,000 tokens completed in under 3 minutes. The 22.5-minute model load time, while long, is dominated by reading 64 safetensor shards from disk at approximately 27 seconds per shard, which is a disk I/O bottleneck rather than a compute bottleneck.

What the Message Reveals: Output Verification

The assistant did not simply declare success and move on. The message includes a verification step:

ssh root@10.1.230.174 'cat /root/eagle3-train/data_1k/hidden_states/data_config.json && echo "---" && 
ls /root/eagle3-train/data_1k/hidden_states/ && echo "---" && du -sh /root/eagle3-train/data_1k/hidden_states/'

The verification checks three things:

  1. The data_config.json confirms the extraction parameters: 4 layers captured (2, 30, 58, 60), 1000 total samples, 502,983 tokens, 172.7 seconds of extraction time, hidden size 7168, and data format version 1. This metadata is essential for the training script to correctly interpret the binary hidden state files.
  2. The directory listing shows rows_0-2000 (a directory containing individual sample files) and sample_lengths.json (a mapping of sample indices to their token counts). The naming rows_0-2000 is a vestige of the extraction script's batching logic—it allocates space for up to 2000 samples even though only 1000 were extracted.
  3. The disk usage reports 27 GB for the entire hidden states directory. This matches the assistant's earlier estimate in [msg 2800]: "Each sample is ~28 MB (512 tokens × 4 layers × 7168 × 2 bytes). For 1000 samples that's ~28 GB." The actual 27 GB is close to this estimate, confirming that the extraction produced the expected volume of data.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

Output Knowledge Created

This message creates several pieces of actionable knowledge:

  1. Extraction throughput benchmark: 2912 tok/s for hidden state extraction on 8× Blackwell GPUs with TP=8. This is a reference point for estimating future extraction runs: for a dataset of N tokens, expect approximately N/2912 seconds of extraction time after the 22.5-minute model load.
  2. Storage requirements: 27 GB for 1000 samples at 2048 max sequence length. This scales roughly linearly—10,000 samples would require ~270 GB, which the assistant later accounts for when the user redirects output to the 3 TB /data volume.
  3. Data format validation: The data_config.json confirms that the extraction script correctly captured the specified layers and produced the expected metadata. The sample_lengths.json file enables the training script to properly pack sequences of varying lengths into batches.
  4. Pipeline readiness: With the hidden states extracted, the training step (04_train.py) can proceed immediately. The assistant proceeds to launch training in [msg 2823], completing 10 epochs in 27.7 minutes.

The Broader Significance

Message [msg 2819] is a quiet triumph in a long debugging saga. It represents the point at which the EAGLE-3 training pipeline transitioned from theory to practice. Before this message, the pipeline had been tested on only 10 samples—enough to validate the code paths but not enough to produce a useful draft model. After this message, the pipeline had 1000 real samples, and the training could produce a checkpoint that might actually improve inference throughput.

The message also illustrates a key pattern in the assistant's methodology: verify everything. The assistant does not assume the extraction succeeded just because the process exited without errors. It reads the config file, checks the directory structure, measures the disk usage, and compares the results against its earlier estimates. This verification discipline is what separates a robust pipeline from a fragile one.

The 27 GB of hidden states sitting on disk at this moment represent the raw material for the next stage: training the draft model to predict the verifier's next-token hidden states, and ultimately achieving the 2×–3× inference speedup that speculative decoding promises. The message is a checkpoint in the truest sense—not just a log line, but a commitment that the most expensive step has been paid and the pipeline can move forward.