Scaling the EAGLE-3 Training Pipeline: From 10 Samples to 1000

In the long arc of deploying speculative decoding for a 1-trillion-parameter language model on 8x Blackwell GPUs, few messages capture the transition from proof-of-concept to production quite like this one. At message index 2801, the assistant makes a deliberate decision to scale the EAGLE-3 draft model training pipeline from a 10-sample validation run to a 1000-sample real training run. This message is the hinge point where the pipeline stops being a prototype and starts becoming a serious data-generation and training operation.

The Message in Full

Here is the message exactly as written:

Good. Let me now run the full pipeline for 1000 samples. This will take a while since the hidden state extraction needs to load the 547GB model. Let me plan the steps:

>

1. Step 1: Prepare 1000 samples (~minutes, CPU only) 2. Step 3: Build vocab mapping (fast, CPU only) 3. Step 2: Extract hidden states (~22 min load + extraction time, needs all 8 GPUs) 4. Step 4: Train (single GPU, minutes)

>

Steps 1 and 3 can run without GPUs. Step 2 is the bottleneck. Let me start step 1: [bash command to run 01_prepare_dataset.py with 1000 samples]

The message is deceptively simple. It contains a plan, a prioritization judgment, and the launch of the first step. But beneath this surface lies a dense web of reasoning about resource constraints, pipeline dependencies, and the economics of GPU time.

Why This Message Was Written

The immediate context is that the assistant had just completed a full end-to-end validation of the EAGLE-3 training pipeline on 10 samples. In the preceding messages ([msg 2785] through [msg 2798]), the assistant had:

  1. Verified that the trained checkpoint's weight shapes exactly match the AQ-MedAI reference model
  2. Confirmed that vLLM can parse the flat config format with LlamaForCausalLMEagle3 architecture
  3. Simulated vLLM's weight key mapping and confirmed all 16 weight tensors map correctly
  4. Validated that the config loads as a proper LlamaConfig with all EAGLE-3 fields With the pipeline validated, the natural next step was to scale up. But scaling up from 10 samples to 1000 samples is not a trivial change — it involves a 22-minute model load on 8 GPUs consuming ~91 GB each, 27 GB of disk storage for hidden states, and careful resource management to avoid OOM errors. The deeper motivation is that EAGLE-3 draft models need meaningful amounts of training data to learn the verifier model's output distribution. Ten samples might validate the pipeline mechanics, but they cannot produce a useful draft model. The assistant needed to demonstrate that the pipeline could handle real data volumes before the user would invest time in generating high-quality synthetic training data — which is precisely what happens in the next chunk of the conversation.

The Thinking Process Visible in the Message

The assistant's reasoning is laid out explicitly in the message's structure. The four steps are numbered 1, 3, 2, 4 — notably out of order. This is not a typo; it's a deliberate reordering based on resource dependencies.

Steps 1 and 3 (dataset preparation and vocab mapping) are CPU-only and can run immediately without any GPU contention. Step 2 (hidden state extraction) is the bottleneck — it requires loading the 547GB verifier model across all 8 GPUs, a process that takes ~22 minutes just for model loading. Step 4 (training) uses a single GPU and is fast once the data is ready.

The assistant recognizes that by launching step 1 first, it can overlap the CPU-bound work with the planning for the GPU-bound bottleneck. The message explicitly flags step 2 as "the bottleneck" — a clear signal that the assistant is thinking about critical path scheduling.

The parenthetical estimates are also revealing: "~minutes, CPU only" for step 1, "fast, CPU only" for step 3, "~22 min load + extraction time, needs all 8 GPUs" for step 2, and "single GPU, minutes" for step 4. These estimates come from the assistant's prior experience: the 10-sample extraction took 1.7 seconds after model load, and the model load was observed to take ~22 minutes in earlier runs.

Assumptions Made

The message makes several assumptions, some of which prove incorrect:

Assumption 1: The batch size used for 10 samples will work for 1000 samples. The previous successful extraction used --batch-size 2000 (which was actually the total token count, not the sample count — a confusion that becomes critical). The assistant assumes that because 10 samples worked, 1000 samples will work with the same parameters. This assumption is wrong, as we see in the following messages ([msg 2813]) where the extraction OOMs with a "CUDA out of memory" error because batch_size=2000 tried to prefill all 503K tokens at once.

Assumption 2: The argument names in the extraction script match what was used before. The assistant initially launches the extraction with --data-path instead of the correct --prepared-data argument ([msg 2805]), which causes the script to print its help message instead of running. This is caught and corrected in [msg 2809].

Assumption 3: Disk space is sufficient. The assistant checks disk space in the preceding message ([msg 2799]) and finds 736 GB free on /root/. The estimate of ~112 MB per sample for 2048-token sequences leads to ~112 GB for 1000 samples, which fits comfortably. This assumption holds.

Assumption 4: The extraction rate scales linearly. The assistant estimates extraction at ~2280 tok/s based on the 10-sample run (3875 tokens in 1.7 seconds). The actual extraction achieves 2912 tok/s ([msg 2819]), which is faster than estimated — a pleasant surprise.

Mistakes and Incorrect Assumptions

The most significant mistake is the batch size assumption. The 10-sample test used --batch-size 2000 as a single batch, but the 10 samples had only 3875 total tokens. When scaled to 1000 samples with 503K tokens, the same batch size tries to prefill half a million tokens in one shot, which exceeds the ~4 GB of free memory per GPU after the model uses ~91 GB.

This mistake reveals an important subtlety about the extraction script's --batch-size parameter. The parameter name is ambiguous — it could mean "samples per batch" or "tokens per batch." The script's default value of 4 suggests samples per batch, but the assistant initially uses 2000, which works for 10 samples only because 2000 > 10, so all samples fit in one batch. The OOM only manifests at scale.

The argument name error (--data-path vs --prepared-data) is a more mundane mistake but costs time — the first extraction launch fails silently (it prints help text to the log), requiring a kill and re-launch.

Input Knowledge Required

To understand this message fully, one needs:

  1. The EAGLE-3 training pipeline architecture: That hidden state extraction requires loading the full verifier model (547 GB, 8 GPUs, TP=8), while training is a single-GPU operation on extracted features.
  2. The resource profile of the verifier model: Kimi-K2.5 INT4 at 547 GB distributed across 8 Blackwell GPUs with ~95 GB each, leaving ~4 GB free per GPU for computation.
  3. The data pipeline: That dataset preparation tokenizes conversations, vocab mapping builds a compressed vocabulary from token frequencies, hidden state extraction runs the model forward to capture intermediate activations, and training learns a lightweight draft model from those activations.
  4. The previous validation results: That 10 samples trained successfully in ~1 minute, weight shapes matched AQ-MedAI's reference, and vLLM could parse the checkpoint config.
  5. The disk and memory constraints: 736 GB free on the root volume, 27 GB expected for 1000 samples of hidden states at 2048 tokens each.

Output Knowledge Created

This message produces several concrete outputs:

  1. A launched dataset preparation job that processes 1000 samples from the mlabonne/open-perfectblend dataset, tokenizing them with the Kimi-K2.5 tokenizer and computing loss masks.
  2. A documented plan for the remaining pipeline steps that serves as a coordination artifact for the user and the assistant's own future actions.
  3. Resource scheduling decisions that prioritize CPU-bound work before GPU-bound work, minimizing idle GPU time.
  4. Estimates for each step's duration that set expectations for the overall timeline (~22 minutes for model load, ~3-4 minutes for extraction, ~minutes for training). The downstream results (visible in subsequent messages) include: 1000 samples processed with 503K total tokens and 360K assistant tokens ([msg 2802]), 100% vocab coverage ([msg 2803]), successful hidden state extraction at 2912 tok/s producing 27 GB of data ([msg 2819]), and training completing 10 epochs in 27.7 minutes at 6 steps/s.

The Broader Significance

This message represents the moment when the EAGLE-3 project transitions from "can we build this?" to "let's build this at scale." The 10-sample validation proved the pipeline mechanics; the 1000-sample run proves the pipeline economics. The 22-minute model load is a fixed cost that becomes worthwhile only when amortized over enough samples — and 1000 samples at 2912 tok/s extraction makes that amortization work.

The message also illustrates a pattern that recurs throughout the session: the assistant thinks in terms of critical path scheduling, resource bottlenecks, and overlapping CPU and GPU work. This is not just a technical skill but a mindset shaped by working with extremely large models where GPU time is the scarcest resource. Every minute of GPU idle time is a minute wasted, and the assistant's planning reflects that awareness.

Finally, the mistakes in this message — the batch size OOM, the wrong argument name — are instructive. They show that even with careful planning, scaling up a pipeline by two orders of magnitude reveals hidden assumptions that only manifest at the new scale. The assistant's response to these failures (killing processes, freeing GPU memory, correcting parameters, and re-launching) is as important as the initial plan. In production ML engineering, the ability to recover quickly from failures is often more valuable than getting everything right on the first try.