The Capacity Planning Pivot: Scaling EAGLE-3 Training from 10 to 1000 Samples

In the long arc of building a speculative decoding system for the 1-trillion-parameter Kimi-K2.5 model, there comes a quiet moment where the proof of concept ends and the real work begins. Message 2800 in this opencode session is that moment. It is a short, deceptively simple message — barely a few lines of reasoning followed by a file read — but it represents a critical pivot point in the engineering workflow. The assistant has just validated that the EAGLE-3 training pipeline runs end-to-end on 10 test samples and that the resulting checkpoint is compatible with vLLM. Now it must decide: do we stay at 10 samples for a quick experiment, or do we scale up to something that might actually produce a useful draft model? The answer, delivered in this message, is a deliberate, resource-conscious decision to scale to 1000 samples.

The Message in Full

The subject message reads:

Plenty of space. Each sample is ~28 MB (512 tokens × 4 layers × 7168 × 2 bytes). For 1000 samples that's ~28 GB. For longer samples (2048 tokens), it's ~112 MB each, so 1000 samples = ~112 GB. We have 736 GB free on /root/.

>

Let me check the extraction rate and plan for 1000 samples. From the data_config.json, we extracted 10 samples (3875 tokens) in 1.7 seconds, but that's after the model was loaded. The model load takes ~22 minutes. For 1000 samples we need a larger dataset. Let me prepare one: [read] /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/01_prepare_dataset.py

This is not a message that contains tool calls in the traditional sense — it is a message that is a tool call. The assistant issues a read operation on the dataset preparation script, which will be executed and its results returned in the next round. But the reasoning embedded in the text before that read operation is where the real substance lies.

The Reasoning: Capacity Planning as Engineering Discipline

The assistant begins with a seemingly mundane observation: "Plenty of space." But this observation is the result of a rapid mental calculation that reveals deep understanding of the data pipeline. Let's unpack each term in the storage formula:

"512 tokens × 4 layers × 7168 × 2 bytes"

Each component reflects a design decision or architectural constraint:

• 512 tokens: This is the sequence length used in the 10-sample validation run. The assistant uses it as a baseline for the per-sample storage estimate. However, the training pipeline supports sequences up to 2048 tokens, so the assistant immediately provides a worst-case estimate as well.
• 4 layers: The EAGLE-3 architecture for Kimi-K2.5 captures hidden states from four specific layers of the verifier model. Earlier in the conversation ([msg 2786]), the assistant discovered that the AQ-MedAI reference model uses auxiliary layers at indices 2, 30, and 58. The fourth layer is likely the final hidden state or an additional layer required by the EAGLE-3 training objective. This is not a random choice — the EAGLE-3 algorithm needs hidden states from intermediate layers to train the draft model's feature extraction network (the fc layer that maps verifier hidden states to the draft model's input).
• 7168: This is the hidden dimension of Kimi-K2.5, a 1-trillion-parameter Mixture-of-Experts model. Every token's representation at every captured layer is a 7168-dimensional vector.
• 2 bytes: The storage format is bfloat16, a 16-bit floating-point format that uses 2 bytes per scalar value. This is the standard precision for modern ML training and inference. Multiplying these together: 512 × 4 × 7168 × 2 = 29,360,128 bytes ≈ 28 MB per sample. For 1000 samples at this length, that's 28 GB. But the pipeline supports sequences up to 2048 tokens, which quadruples the per-sample storage to ~112 MB, bringing the 1000-sample total to ~112 GB. The assistant then checks this against the available disk space: 736 GB free on /root/. The headroom is substantial — even the worst case uses only ~15% of available space. This calculation is not trivial; it reflects an engineering mindset that checks resource constraints before launching a long-running job. A less careful approach might have started the extraction and discovered hours later that the disk was full.

The Extraction Rate Calculation

The assistant also performs a throughput estimate: "we extracted 10 samples (3875 tokens) in 1.7 seconds." This gives an extraction rate of approximately 2279 tokens per second. For 1000 samples at an average of, say, 1000 tokens each (1 million tokens total), the extraction itself would take about 7-8 minutes. But the dominant cost is the model load: 22 minutes to load the 547 GB verifier model across 8 GPUs. This means the total time for hidden state extraction is dominated by the model load, not the extraction itself — a classic fixed-cost vs. marginal-cost tradeoff. Scaling from 10 to 1000 samples adds only marginal extraction time but doesn't increase the 22-minute load time. This makes the 1000-sample run economically sensible: the marginal cost of extracting 990 more samples is small compared to the fixed cost of loading the model.

Assumptions Embedded in the Message

Several assumptions underlie this message, some explicit and some implicit:

1. The extraction rate scales linearly. The 2279 tok/s rate was measured on a 10-sample batch with short sequences (3875 tokens total). For larger batches with longer sequences, throughput might differ due to memory bandwidth, PCIe transfer costs, or batching efficiency. The assistant implicitly assumes the rate holds.
2. The 4-layer capture is correct. The assistant assumes that 4 layers of hidden states are needed per token. This is based on the EAGLE-3 architecture design, but the exact number could vary depending on how the training script is configured.
3. The model load time is stable. The 22-minute estimate comes from the 10-sample run. In practice, load time can vary based on GPU memory fragmentation, NCCL initialization, and filesystem cache state.
4. The dataset preparation script is the right starting point. The assistant reads 01_prepare_dataset.py assuming it's the correct entry point for scaling up. This is a reasonable assumption given the numbered pipeline (01, 02, 03, 04), but the assistant hasn't verified that the script handles 1000 samples correctly.
5. The draft model training will benefit from more data. This is the fundamental assumption behind the entire scale-up: that training on 1000 samples will produce a meaningfully better draft model than training on 10. For EAGLE-3, which learns to predict verifier hidden states, more diverse training data should improve the draft model's ability to generalize to unseen prompts.

What Knowledge Was Required

To understand and produce this message, the assistant needed:

• The EAGLE-3 architecture: Knowledge that the draft model is trained on hidden states extracted from specific verifier layers, and that 4 layers are captured per token.
• The Kimi-K2.5 model dimensions: Hidden size of 7168, vocabulary size of 163840, and the bfloat16 storage format.
• The pipeline structure: Understanding that the workflow has numbered steps (prepare dataset → extract hidden states → build vocab mapping → train), and that step 1 produces tokenized data that step 2 consumes.
• The 10-sample test results: Knowledge of the extraction rate (3875 tokens in 1.7s), model load time (~22 min), and output size from the previous validation run.
• Disk space information: The df -h output from [msg 2799] showing 736 GB free on /root/.
• The dataset source: The open-perfectblend dataset from mlabonne, which was used in the 10-sample test and would be reused for the 1000-sample run.

What Knowledge Was Created

This message produces several important outputs:

1. A storage budget: 1000 samples at max length = ~112 GB, well within the 736 GB available. This de-risks the scale-up.
2. A time budget: ~22 minutes for model load + ~7 minutes for extraction = ~30 minutes total for step 2.
3. A decision: Scale to 1000 samples. This sets the direction for the next several hours of work.
4. A concrete next action: Read and execute 01_prepare_dataset.py to create the 1000-sample tokenized dataset.

The Thinking Process Visible in the Message

The message reveals a clear chain of reasoning:

1. Check resources first: "Plenty of space" — the assistant has just received disk space information and immediately evaluates it against the anticipated storage needs.
2. Calculate worst case: The assistant doesn't just use the 512-token baseline; it also computes the 2048-token worst case (112 GB vs 28 GB). This is defensive engineering — planning for the maximum possible resource consumption.
3. Consider the bottleneck: The assistant identifies that model load (22 min) dominates extraction time (1.7s for 10 samples). This understanding justifies scaling up: the fixed cost is already paid, so extracting more samples adds little time.
4. Plan the next step: Having decided to scale, the assistant immediately reads the script that starts the pipeline. There's no hesitation or second-guessing — the decision is made and execution begins.

How This Message Shaped Subsequent Events

The decision made in this message set off a chain of events that consumed the next several hours of the session:

• The assistant ran 01_prepare_dataset.py with --max-samples 1000, which processed 1000 samples (skipping 58 that were too long) and produced 503K total tokens ([msg 2802]).
• The vocab mapping step (03_build_vocab_mapping.py) showed 100% coverage of the observed tokens by the 32K draft vocabulary ([msg 2803]).
• The hidden state extraction (02_extract_hidden_states.py) hit a critical OOM failure when the assistant naively used --batch-size 2000 (trying to process all 1000 samples as a single batch), requiring a re-launch with --batch-size 4 (<msg id=2813-2816>).
• After the OOM fix, extraction completed successfully: 22.5 minutes for model load, then 2.9 minutes of extraction at 2912 tok/s ([msg 2818]).
• Training on 1000 samples took 27.7 minutes for 10 epochs at 6 steps/s ([msg 2820]). The OOM failure is particularly instructive. The assistant's assumption that --batch-size referred to tokens (as in many ML pipelines) rather than samples led to a costly mistake — 22 minutes of model loading wasted. This is a classic example of an API misunderstanding that a more careful reading of the argument parser would have prevented. The assistant later discovered the correct argument name (--prepared-data vs --data-path) and the correct batch size semantics.

Conclusion

Message 2800 is a masterclass in engineering decision-making under uncertainty. It demonstrates how to evaluate resource constraints, estimate costs, and make a data-driven decision about scaling up a complex ML pipeline. The assistant's reasoning is transparent, quantitative, and grounded in empirical measurements from the 10-sample validation run. While some assumptions proved incorrect (the batch size semantics), the core decision — to scale from 10 to 1000 samples — was sound and ultimately successful. The message captures the moment when a proof of concept transforms into a production-scale experiment, and it reveals the engineering discipline required to make that transition smoothly.