The Moment of Failure: When a CUDA OOM Revealed a Critical Assumption About Batch Size

Introduction

In the high-stakes world of large language model deployment, few moments are as jarring as watching a 22-minute model loading process culminate in a CUDA out of memory error. Message [msg 2813] captures precisely such a moment. It is the instant when an assistant, having patiently monitored the progress of a critical hidden state extraction pipeline for the EAGLE-3 training workflow, discovers that instead of nearing completion, the process has silently crashed. This message is a study in how assumptions about parameter semantics can cascade into costly failures, and how the recovery from those failures produces the most valuable debugging knowledge.

The Scene: Scaling Up the EAGLE-3 Pipeline

To understand the gravity of this message, one must appreciate the context. The assistant had spent the better part of a day building an end-to-end EAGLE-3 training pipeline for the Kimi-K2.5 model, a 1-trillion-parameter Mixture-of-Experts architecture running on 8 NVIDIA RTX PRO 6000 Blackwell GPUs. The pipeline had been validated on a toy dataset of 10 samples, and the assistant was now scaling it to 1000 samples—a necessary step toward producing a useful speculative decoding draft model.

The bottleneck in this pipeline is Step 2: hidden state extraction. This step requires loading the full 547 GB verifier model across all 8 GPUs using tensor parallelism (TP=8), then running every training sample through the model to capture the internal hidden states from specific layers. The model loading alone takes approximately 22 minutes. The assistant had launched this extraction job in the background (see [msg 2809]) with a critical parameter choice: --batch-size 2000.

What the Message Says

The message begins with the assistant's internal monologue:

78% done with loading, about 5 more minutes. Let me wait a bit more:

This is the assistant speaking to itself, processing the status it observed from a previous check. It then issues a bash command to check the log after a 600-second (10-minute) sleep:

[bash] sleep 600 && ssh root@[REDACTED_IP] 'tail -30 /root/eagle3-train/extract_1k.log 2>/dev/null'

But instead of the expected progress update showing completed extraction, the command returns a catastrophic error trace:

  File "/root/ml-env/lib/python3.12/site-packages/vllm/v1/executor/multiproc_executor.py", line 375, in collective_rpc
    return aggregate(get_response())
                     ^^^^^^^^^^^^^^
  File "/root/ml-env/lib/python3.12/site-packages/vllm/v1/executor/multiproc_executor.py", line 358, in get_response
    raise RuntimeError(
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 wh...

The error originates from vLLM's multiproc_executor.py, specifically the collective_rpc function—the mechanism vLLM uses for inter-worker communication in the tensor-parallel setup. One of the GPU worker processes has crashed with an out-of-memory error, and the collective_rpc call has detected the failure and propagated it upward.

The Thinking Process: What the Assistant Was Doing

The assistant's reasoning in this message reveals several layers of cognitive activity. First, there is the monitoring loop: the assistant had previously checked the extraction log at [msg 2812] and seen loading progress at 50–55%. Now, at [msg 2813], it references a more recent observation of "78% done." This tells us the assistant was actively tracking the process, checking in periodically to gauge progress.

Second, there is the estimation: "about 5 more minutes." This is a reasonable extrapolation from the observed loading rate of approximately 1 shard per 18–24 seconds. With 64 shards total and 78% (50 shards) complete, 14 shards remaining at ~22 seconds each would take about 5 minutes. The assistant then decides to wait 10 minutes (sleep 600) to give the process time to finish loading and begin extraction.

Third, there is the surprise: the assistant expected to see extraction progress but instead found a crash. The error trace is presented without commentary in this message—the assistant is still processing what it sees. The diagnosis will come in the next message ([msg 2814]), where the assistant immediately identifies the root cause: "OOM! The batch size of 2000 (all 1000 samples at once as a single batch) is too large."

The Critical Assumption: What Does --batch-size Mean?

The central mistake in this message is an assumption about the semantics of the --batch-size parameter in the extraction script. The assistant had set --batch-size 2000, reasoning (presumably) that this would process tokens efficiently. In many ML frameworks, batch size refers to the number of tokens or samples processed in a single forward pass, and larger batches are generally more efficient.

However, the extraction script's --batch-size parameter controlled the number of samples per batch, not tokens. With 1000 samples in the dataset and a --batch-size of 2000, the script attempted to process all 1000 samples—potentially up to 503,000 tokens—in a single batch. The model, already consuming approximately 91 GB of each GPU's 95 GB capacity for the INT4 quantized weights and KV cache, had only about 4 GB of headroom. Trying to allocate an additional 896 MB for the batch activations pushed it over the edge.

This assumption is understandable. The previous successful test with 10 samples had used the default batch_size=4, which was small enough to fit comfortably. When scaling to 1000 samples, the assistant increased the batch size dramatically, likely intending to process all data in one go for simplicity. But the parameter's semantics—samples, not tokens—meant that batch_size=2000 was not just large; it was catastrophically large.

Input Knowledge Required

To fully understand this message, a reader needs several pieces of context:

  1. The EAGLE-3 pipeline architecture: Hidden state extraction requires loading the full verifier model (547 GB, INT4 quantized) across 8 GPUs with tensor parallelism. This consumes nearly all available GPU memory.
  2. vLLM's multiprocess executor: The error trace references multiproc_executor.py and collective_rpc. This is vLLM's mechanism for coordinating multiple GPU worker processes. When one worker crashes (e.g., from OOM), the collective_rpc call detects the failure and raises a RuntimeError.
  3. GPU memory budgeting: Each RTX PRO 6000 Blackwell GPU has 96 GB of memory. The Kimi-K2.5 INT4 model with TP=8 uses ~91 GB, leaving only ~4 GB of headroom for activations, temporary buffers, and batch processing.
  4. The batch-size parameter: The extraction script's --batch-size controls how many samples are processed per forward pass, not how many tokens. This distinction is crucial.
  5. The previous successful run: The 10-sample test used batch_size=4 and worked fine. The assistant scaled this to 2000 without considering the memory implications.

Output Knowledge Created

Despite being a failure, this message produces valuable knowledge:

  1. The memory ceiling is real and tight: With the INT4 model loaded at TP=8, each GPU has approximately 4 GB of usable headroom. Any additional allocation above ~896 MB will trigger an OOM.
  2. Batch size means samples, not tokens: The extraction script processes batches of samples, each of which may contain up to 2048 tokens. A batch of 2000 samples means up to 4 million tokens of prefill—far exceeding available memory.
  3. The crash signature is diagnostic: The specific error location—collective_rpc in multiproc_executor.py—is characteristic of a worker crash during distributed execution. This pattern will be recognizable in future debugging.
  4. Small batches are mandatory: The correct batch size for this setup is small (e.g., 4 samples per batch), which keeps the activation memory within the available headroom.

The Recovery

In the immediate aftermath ([msg 2814]), the assistant diagnoses the problem correctly within seconds: "OOM! The batch size of 2000 (all 1000 samples at once as a single batch) is too large." It then kills the stuck processes, frees the GPU memory, and re-launches with --batch-size 4 and --gpu-memory-utilization 0.85 ([msg 2816]). The second attempt succeeds: 1000 samples extracted in 2.9 minutes after the 22.5-minute model load ([msg 2819]).

Broader Significance

This message exemplifies a common pattern in ML engineering: the gap between small-scale validation and production-scale execution. The pipeline worked perfectly with 10 samples and batch_size=4. Scaling to 1000 samples required a seemingly innocuous parameter change that violated an invisible constraint—the GPU memory ceiling. The assistant's assumption that "bigger batch = faster" was correct in principle but fatal in practice because it didn't account for the razor-thin memory margin.

The message also illustrates the value of the opencode session format. The assistant's reasoning is fully visible: we see it monitoring progress, estimating completion time, deciding to wait, and then discovering the failure. We see the raw error trace. In the next message, we see the immediate diagnosis and correction. This transparency turns a simple error into a rich learning opportunity about distributed ML system behavior.