The 1000-Sample Extraction Launch: A Pivotal Moment in EAGLE-3 Training

In the course of building a speculative decoding pipeline for the 1-trillion-parameter Kimi-K2.5 model on 8× NVIDIA Blackwell GPUs, one message stands out as a critical inflection point: the launch of hidden state extraction for 1000 training samples. This single bash command, issued via SSH with nohup, represents the transition from a validated proof-of-concept to a production-scale data pipeline. But it also encapsulates the subtle assumptions and mistakes that make real-world ML engineering so challenging.

The Message

Here is the exact message under analysis, from message index 2805 in the conversation:

[bash] ssh root@10.1.230.174 'nohup /root/ml-env/bin/python3 /root/eagle3-train/02_extract_hidden_states.py \
    --model-path /shared/kimi-k2.5-int4 \
    --data-path /root/eagle3-train/data_1k/prepared/tokenized_data.jsonl \
    --output-dir /root/eagle3-train/data_1k/hidden_states \
    --tp-size 8 \
    --max-seq-len 2048 \
    --batch-size 2000 \
    > /root/eagle3-train/extract_1k.log 2>&1 &
echo "PID: $!"'
PID: 268082

At first glance, this looks like a routine command: launch a Python script in the background on a remote machine, redirect output to a log file, and print the process ID. But this message sits at the convergence of weeks of infrastructure work, model deployment, and algorithmic research. Understanding why it was written, what assumptions it encodes, and what went wrong reveals the texture of cutting-edge ML engineering.

Context and Motivation

To understand why this message exists, we need to trace the arc of the broader session. The assistant and user had been working together for over 2800 messages across 22 segments. They had deployed multiple 1T-parameter models (GLM-5, Kimi-K2.5 NVFP4, MiniMax-M2.5, Kimi-K2.5 INT4), benchmarked them extensively, identified PCIe AllReduce as the dominant bottleneck at 51.5% of decode time, and pivoted to speculative decoding as a software-only optimization path.

The EAGLE-3 training pipeline was the culmination of this pivot. EAGLE-3 is a speculative decoding technique that trains a small "draft" model to predict the next several tokens of the large "verifier" model, using hidden states from intermediate layers of the verifier as auxiliary features. The assistant had already:

1. Built the complete EAGLE-3 training pipeline from scratch
2. Patched the speculators library for API compatibility with vLLM 0.16
3. Validated the pipeline end-to-end on 10 test samples (3 epochs in ~1 minute)
4. Verified that the output checkpoint was vLLM-compatible with identical weight shapes to the AQ-MedAI reference model But 10 samples is not enough to train a useful draft model. The assistant needed to scale up to a meaningful dataset. The user had redirected the approach toward generating higher-quality training data by capturing Kimi-K2.5's actual reasoning outputs via the vLLM inference server. However, before that synthetic data generation pipeline could be built, the assistant needed to complete a baseline training run on the existing dataset — the mlabonne/open-perfectblend conversation dataset. This message launches the critical data preparation step: extracting hidden states from the verifier model for 1000 training samples. Without these hidden states, the EAGLE-3 draft model cannot be trained, because the auxiliary features (hidden states from layers 2, 30, 58, and 60 of the verifier) are what the draft model learns to condition on.

The Two Mistakes

The command contains two significant errors, both of which are caught in subsequent messages ([msg 2806] through [msg 2809]).

Mistake 1: Wrong argument name. The script 02_extract_hidden_states.py defines its input argument as --prepared-data, not --data-path. The assistant used --data-path, which Python's argparse does not recognize. This means the script would fail immediately with an unrecognized argument error — but because the command was launched with nohup and output redirected to a log file, the error would not appear in the terminal. The assistant only discovered this mistake in the next message when checking the log.

This is a classic "works on my machine" problem compounded by remote execution. The assistant had previously run the extraction script for the 10-sample test, but that test used default arguments. The first time they tried to pass --data-path, they guessed the wrong parameter name. The script's argument names were defined in a file that the assistant had read earlier but not internalized.

Mistake 2: Catastrophically wrong batch size. The --batch-size 2000 parameter is far too large. The script's default is 4 (samples per batch), and the assistant assumed --batch-size referred to token count — a reasonable assumption for someone coming from training code where batch size typically means tokens or samples per gradient step. But in this extraction script, batch-size means the number of samples to process in a single forward pass. With 1000 samples at up to 2048 tokens each, a batch of 2000 would attempt to prefill up to 4 million tokens at once across 8 GPUs. The 547GB verifier model already consumes ~91 GB per GPU with tensor parallelism. Adding a 4M-token prefill would overflow GPU memory by orders of magnitude.

The subsequent OOM crash ([msg 2813]) confirms this: "CUDA out of memory. Tried to allocate 896.00 MiB. GPU 0 has a total capacity of 94.97 GiB." The assistant had to kill the process, free the GPUs, and re-run with --batch-size 4 and --gpu-memory-utilization 0.85.

Assumptions Embedded in the Command

Beyond the specific mistakes, this message reveals several implicit assumptions:

Assumption 1: The extraction is the bottleneck. The assistant structured the pipeline so that hidden state extraction is a blocking step before training. The command is launched with nohup specifically because model loading takes ~22 minutes — the assistant knew this from the 10-sample test. The assumption is that extraction time dominates the pipeline, so it should be started as early as possible.

Assumption 2: All 8 GPUs are needed. The --tp-size 8 flag loads the model across all 8 GPUs using tensor parallelism. This is necessary because the 547GB model cannot fit on a single GPU (95 GB each). The assistant assumes the GPUs are idle and available — a reasonable assumption given that the previous vLLM server had been stopped.

Assumption 3: The data is ready. The assistant had just finished preparing the 1000 samples in the previous message ([msg 2802]), confirming 503K total tokens with 360K assistant tokens. The vocab mapping was also built ([msg 2803]). The assumption is that the data pipeline is complete and the extraction can proceed immediately.

Assumption 4: The remote environment is stable. The command uses nohup and background execution, assuming the SSH connection will not be the limiting factor. The 22-minute model load time makes this necessary — the assistant cannot keep an SSH session open for that duration.

Input Knowledge Required

To understand this message, one needs to know:

• The EAGLE-3 architecture: EAGLE-3 trains a draft model using hidden states from specific intermediate layers of the verifier model. Layer IDs 2, 30, 58, and 60 are chosen because they correspond to early, middle, late, and final layers of the 60-layer Kimi-K2.5 model.
• Tensor parallelism (TP): The model is too large for a single GPU, so it must be sharded across 8 GPUs. TP=8 means each GPU holds 1/8 of each layer's parameters.
• The speculators library: This is the third-party library providing the EAGLE-3 training infrastructure. The assistant had to patch it extensively for compatibility.
• The verifier model: Kimi-K2.5 INT4, a 1T-parameter Mixture-of-Experts model with 60 layers, 163840 vocabulary size, and 7168 hidden dimension. The INT4 quantization reduces memory from ~2TB to ~547GB.
• The open-perfectblend dataset: A conversation dataset used for training, containing 1000 samples with 503K tokens total.

Output Knowledge Created

This message does not produce immediate output — it launches a background process. The output (hidden states) will be created over the next ~25 minutes. But the message itself creates several important artifacts:

1. A background process (PID 268082) running on the remote machine, consuming all 8 GPUs for model loading.
2. A log file at /root/eagle3-train/extract_1k.log that will capture the extraction progress.
3. A scheduling decision: The assistant has committed to this extraction run, blocking other GPU workloads for the next ~25 minutes. The eventual output (after the mistakes are corrected in subsequent messages) is 27 GB of hidden state data for 1000 samples, extracted at 2912 tok/s after a 22.5-minute model load.

The Thinking Process

The reasoning visible in this message is about efficiency and pipeline orchestration. The assistant is thinking several steps ahead:

1. "The extraction takes 22 minutes to load the model, then ~220 seconds for the actual extraction at 2280 tok/s."
2. "I should launch this in the background so I can monitor it without blocking the conversation."
3. "I need to save the PID so I can check on it later."
4. "The batch size of 2000 should be fine because the extraction is fast — 503K tokens at 2280 tok/s." The last point reveals the flawed mental model. The assistant was thinking in terms of total throughput (2280 tok/s from the 10-sample test) and scaled the batch size proportionally. But the batch size controls peak memory, not throughput. The 10-sample test used batch_size=4 (the default) and achieved 2280 tok/s. Increasing batch_size to 2000 doesn't increase throughput linearly — it just causes an OOM crash because the prefill memory scales with batch size. This is a subtle but important distinction. In inference, batch size affects both throughput and memory. In training, batch size affects gradient quality and memory. The assistant was thinking like a trainer ("larger batch = faster processing") but the extraction script was doing inference-style forward passes where batch size directly controls peak memory allocation.

Significance

This message is a microcosm of the entire session. It represents the moment when theory meets practice — when a validated pipeline meets real-world scale. The mistakes are not failures; they are learning opportunities that reveal the true semantics of the tools. The assistant corrected both errors within minutes ([msg 2809]), re-launched with --batch-size 4, and successfully extracted all 1000 samples' hidden states in 2.9 minutes.

The 1000-sample extraction enabled the full training run (10 epochs, 27.7 minutes, 6 steps/s) that produced a vLLM-compatible EAGLE-3 draft model checkpoint. This checkpoint was the foundation for the subsequent pivot to synthetic data generation, where the assistant built a pipeline to capture Kimi-K2.5's actual reasoning outputs for higher-quality training data.

In the end, this single bash command — with its wrong argument name and absurd batch size — is a testament to the iterative nature of ML engineering. Every mistake is a signal. Every crash is a correction. And the 1000-sample extraction, once fixed, became the backbone of the entire EAGLE-3 training effort.