The 22-Minute Mistake: Diagnosing and Fixing a Batch Size OOM in EAGLE-3 Hidden State Extraction

Introduction

In the course of building an EAGLE-3 speculative decoding training pipeline for the 1-trillion-parameter Kimi-K2.5 INT4 model, a single message captures a pivotal moment of debugging and recovery. After investing over 22 minutes loading a 547GB verifier model across 8 NVIDIA Blackwell GPUs, the assistant's hidden state extraction job crashed with an out-of-memory (OOM) error. The culprit was a single parameter: --batch-size 2000. Message [msg 2816] represents the correction—a carefully calibrated re-launch with --batch-size 4 and a reduced GPU memory utilization target. This article examines the reasoning, assumptions, and lessons embedded in this seemingly simple command invocation.

The Message

[assistant] Now re-run with batch_size=4 (4 samples per batch, up to 8192 tokens):
[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 \
    --prepared-data /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 4 \
    --gpu-memory-utilization 0.85 \
    > /root/eagle3-train/extract_1k.log 2>&1 &
echo "PID: $!"'
PID: 269061

Context: The EAGLE-3 Training Pipeline

To understand why this message exists, one must appreciate the broader pipeline being assembled. EAGLE-3 is a speculative decoding technique that trains a lightweight "draft" model to predict the hidden states of a much larger "verifier" model. The training pipeline consists of several steps:

  1. Dataset preparation (01_prepare_dataset.py): Tokenizes conversation data from the mlabonne/open-perfectblend dataset, producing tokenized sequences with loss masks.
  2. Hidden state extraction (02_extract_hidden_states.py): The critical bottleneck step. It loads the full 547GB Kimi-K2.5 verifier model across all 8 GPUs using tensor parallelism (TP=8), then runs forward passes on the tokenized data to capture the hidden states from specific intermediate layers (layers 2, 30, and 58, as specified in the EAGLE-3 configuration). These hidden states become the training targets for the draft model.
  3. Vocab mapping (03_build_vocab_mapping.py): Builds a mapping between the verifier's full 163,840-token vocabulary and the draft model's 32,000-token vocabulary.
  4. Training (04_train.py): Uses the speculators library's Eagle3SpeculatorConfig and Trainer class to train the draft model on the extracted hidden states. By message [msg 2816], the assistant had already completed steps 1 and 3 for a 1000-sample dataset. Step 2—the hidden state extraction—was the remaining blocker. The assistant had previously launched the extraction with --batch-size 2000 (a single batch containing all 1000 samples), which after 22+ minutes of model loading, crashed with a CUDA OOM error when it attempted to prefill all 503,000 tokens simultaneously.

The Reasoning Behind the Fix

The assistant's decision to use batch_size=4 was not arbitrary—it was grounded in empirical evidence from earlier in the pipeline. In [msg 2799], the assistant had successfully extracted hidden states for 10 test samples using the default batch_size=4 from the script's argument parser. That test run had processed 3,875 tokens across 10 samples in 1.7 seconds (after model loading), proving that a batch size of 4 was both memory-safe and performant.

The calculation is straightforward: with --max-seq-len 2048, each sample can be up to 2,048 tokens. A batch of 4 samples therefore requires at most 8,192 tokens of prefill context. Against 91+ GB of occupied GPU memory per card (the model itself consumes nearly all of each GPU's 94.97 GiB capacity), an 8K-token prefill is a negligible additional allocation. By contrast, batch_size=2000 would attempt to prefill up to 4,096,000 tokens—a catastrophic memory demand that instantly overflows the remaining ~4 GiB of headroom on each GPU.

The assistant also added --gpu-memory-utilization 0.85, reducing the default from the likely 0.90 or higher. This parameter controls what fraction of available GPU memory vLLM's memory manager will reserve for model weights and KV cache. Setting it to 0.85 reserves 85% of each GPU's 94.97 GiB (~80.7 GiB) for the model, leaving ~14.3 GiB as safety margin for activations, temporary buffers, and the prefill itself. This conservative setting reflects the assistant's recognition that the extraction script was operating right at the edge of the GPU memory envelope.

Assumptions and Their Validity

The message rests on several assumptions, most of which are sound:

Assumption 1: The OOM was caused by batch size, not by model loading. This is correct. The error trace in [msg 2813] showed the crash occurring during forward pass execution, not during checkpoint loading. The model had loaded successfully across all 8 GPUs—the error came when it tried to allocate 896 MiB for a prefill tensor and found no room.

Assumption 2: A batch size of 4 will avoid OOM. This is well-supported by the earlier 10-sample test. However, the 10-sample test used --max-seq-len 512 (the default), while the 1000-sample run uses --max-seq-len 2048. A batch of 4 samples at 2048 tokens each is 8,192 tokens—four times larger than the 512-token batches tested. The assistant implicitly assumes that 8,192 tokens of prefill will still fit within the ~14 GiB memory headroom. This is a reasonable assumption given that 8K tokens of prefill typically requires ~2-4 GiB of temporary memory for activations in a model of this size, but it is not definitively proven.

Assumption 3: The GPUs are fully free after killing the previous process. The assistant verified this in [msg 2814] by running nvidia-smi and confirming all GPUs show "0 MiB" used. This is a critical safety check before re-launching.

Assumption 4: The extraction script handles variable-length sequences correctly within a batch. The script uses vLLM's prefill infrastructure, which supports padding/masking for variable-length sequences. This is a standard capability and a safe assumption.

The Mistake That Preceded This Message

The original error—using --batch-size 2000—stemmed from a misunderstanding of what the --batch-size parameter controlled. The assistant appears to have interpreted it as the total number of samples to process (i.e., "I have 1000 samples, so batch_size=2000 will process all of them"), when in fact it controls how many samples are fed to the GPU in a single forward pass. This is a common confusion point in ML pipelines where "batch size" can mean different things at different stages.

The assistant's thought process in [msg 2809] reveals the mistake: "Let me re-launch" after fixing the --prepared-data argument name, but without questioning the --batch-size 2000 value. The parameter was carried over from the earlier failed launch in [msg 2805], where the same --batch-size 2000 was used with the wrong argument name (--data-path instead of --prepared-data). That earlier launch never got far enough to OOM because it failed immediately on argument parsing. The assistant thus never saw the batch-size error signal until the corrected launch in [msg 2809] actually began loading and crashed 22 minutes later.

Input Knowledge Required

To fully understand this message, one needs:

  1. The EAGLE-3 architecture: Understanding that hidden state extraction requires loading the full verifier model and running forward passes, making it the most resource-intensive step in the pipeline.
  2. vLLM's memory model: Knowledge of how vLLM partitions GPU memory between model weights, KV cache, and scratch space, and how gpu-memory-utilization controls this partitioning.
  3. Tensor parallelism (TP): Understanding that TP=8 splits the model across 8 GPUs, meaning each GPU holds ~1/8 of the weights (~68 GB) plus activation memory proportional to the batch's total token count.
  4. The hardware constraints: The 8× RTX PRO 6000 Blackwell GPUs each have 94.97 GiB of memory, and the Kimi-K2.5 INT4 model consumes ~91 GiB per GPU when loaded with TP=8.
  5. The pipeline's sequential dependencies: Hidden state extraction must complete before training can begin, making it a critical-path bottleneck.

Output Knowledge Created

This message produces several tangible outcomes:

  1. A running extraction process (PID 269061) that will take ~22 minutes to load the model plus ~3-4 minutes for extraction (at the previously observed ~2,912 tok/s extraction rate for 503K tokens).
  2. A log file at /root/eagle3-train/extract_1k.log that will document the process and provide diagnostic information if anything goes wrong.
  3. Hidden state data in /root/eagle3-train/data_1k/hidden_states/ once extraction completes—approximately 28 GB of tensor data (503K tokens × 4 layers × 7168 hidden dim × 2 bytes per bfloat16 element, though only assistant tokens are stored).
  4. A validated parameter configuration: The combination of --batch-size 4 and --gpu-memory-utilization 0.85 becomes the proven-safe configuration for future extractions on this hardware.

The Broader Significance

This message exemplifies a recurring pattern in large-scale ML engineering: the tension between throughput and memory safety. The assistant's initial instinct to use --batch-size 2000 (process everything in one go) was driven by a desire for simplicity and speed—one batch, one forward pass, done. But the hardware's memory constraints forced a more nuanced approach: many small batches, each fitting within the tight memory budget left after loading a model that consumes 96% of available GPU memory.

The 22-minute cost of the failed run (model loading time) is a concrete illustration of why careful parameter validation matters in large-model workflows. A mistake that takes seconds to make can cost tens of minutes to discover. The assistant's response—methodically verifying GPU memory is freed, cleaning up the output directory, adjusting both batch size and memory utilization, and running in the background with logging—reflects the learned discipline of operating at scale.

The choice of batch_size=4 also reveals an implicit throughput calculation. With 1000 samples and batch_size=4, the extraction will require 250 forward passes. At the observed ~2,912 tok/s extraction rate, each pass processes ~8,192 tokens in ~2.8 seconds, totaling ~700 seconds (~12 minutes) of extraction time after model loading. This is acceptable—the bottleneck remains the 22-minute model load, not the extraction itself. A larger batch size would reduce the number of passes but risk OOM; a smaller batch size would increase passes without meaningful memory benefit. The assistant's choice of 4 is a pragmatic sweet spot.

Conclusion

Message [msg 2816] is, on its surface, a simple command invocation. But within the context of the broader EAGLE-3 pipeline, it represents a critical learning moment: the recovery from a costly mistake, the application of empirical evidence (the 10-sample test) to parameter selection, and the careful balancing of throughput against the hard constraints of GPU memory. It demonstrates that in large-model engineering, the most important decisions are often the smallest ones—a batch size here, a memory utilization fraction there—and that understanding why a parameter has the value it does is more important than the value itself.