When Averages Lie: The OOM That Exposed the Gap Between Theory and Practice in GPU Batch Sizing
Introduction
In the high-stakes world of large language model deployment, few moments are as instructive as a well-timed out-of-memory (OOM) error. Message [msg 7305] captures precisely such a moment: an assistant, acting on a user's confident directive to use batch sizes of 128–256 for hidden state extraction, runs a quick test and watches it crash immediately. The result is not a failure—it is a revelation. The message is a compact case study in how assumptions about data, hardware, and workload interact, and how even reasonable calculations can lead to wrong conclusions when they rest on averages that mask extreme variance.
The context is a DFlash drafter training pipeline. The team has already migrated to a new 4× RTX PRO 6000 Blackwell node after the previous instance was killed, rewritten the extraction script with batching support, and pushed all data and code to the new machine. The user's input—"Pretty sure we want 128-256 bigly batch" ([msg 7304])—reflects a natural instinct: to maximize GPU utilization, increase batch size. The assistant's response in [msg 7305] is a rapid, practical test of that hypothesis, and its outcome reshapes the entire approach to the extraction pipeline.
The Reasoning Behind the Message
The assistant's reasoning is laid out in the opening sentence: "with 96GB per GPU and 55GB model, we have ~40GB free for KV cache. With short sequences (~335 avg tokens), batch size 128-256 is totally viable and will saturate the GPU compute." This is a textbook example of capacity planning based on average-case analysis.
The calculation is straightforward. Each GPU on the new node has 96 GB of HBM2e memory. The Qwen3.6-27B model in BF16 consumes approximately 55 GB, leaving roughly 40 GB for the KV cache and activations. The average sequence length in the tokenized dataset is about 335 tokens. For a batch of 128 sequences, the total KV cache memory required would be roughly: batch_size × sequence_length × hidden_dim × num_layers × bytes_per_element × 2 (for key and value). With 5120 hidden dimension and 64 layers in BF16 (2 bytes per element), each token requires about 5120 × 64 × 2 × 2 = 1.3 MB of KV cache. For 128 sequences of 335 tokens, that's about 128 × 335 × 1.3 MB ≈ 56 GB—which would exceed the available 40 GB. But with flash attention and memory-efficient implementations, the actual peak memory might be lower, and the assistant clearly believed the margin was sufficient.
The reasoning is not reckless—it is grounded in real numbers from the actual hardware and model configuration. The mistake is not in the arithmetic but in the assumption that the average sequence length is a reliable predictor of peak memory demand. The assistant implicitly assumes a relatively uniform distribution of sequence lengths, where most samples cluster around the 335-token mean. This assumption is about to be demolished by empirical evidence.
The Test and Its Outcome
The assistant executes a bash command that runs the extraction script on GPU 0 with --max-samples 256 --batch-size 128. The script loads the model in 14.3 seconds—a testament to the Blackwell GPU's memory bandwidth—and then processes two batches of 128 samples each.
The output tells a stark story:
OOM on batch (max_len=3316, bs=128), falling back to bs=1
OOM on batch (max_len=4042, bs=128), falling back to bs=1
The maximum sequence lengths in the two batches are 3,316 and 4,042 tokens—approximately 10–12 times the average. A batch of 128 sequences where even one sequence is 4,042 tokens long requires KV cache memory proportional to that maximum length (because the KV cache is padded to the longest sequence in the batch). The memory required for a single 4,042-token sequence at batch size 128 is enormous: the KV cache for that one long sequence alone consumes roughly 4,042 × 1.3 MB ≈ 5.3 GB, and when multiplied across 128 sequences (each padded to 4,042 tokens), the total KV cache demand exceeds 600 GB—far beyond what any single GPU can provide.
The script's fallback mechanism kicks in, dropping to batch size 1, and the extraction proceeds at a crawl: 56.75 seconds per batch of 1, then 59.43 seconds for the next. The total time for 256 samples is nearly two minutes, which extrapolates to roughly 12 hours for the full 914K-sample dataset on a single GPU—far too slow for practical training.
Assumptions Made and Where They Went Wrong
The message reveals several layers of assumptions, some explicit and some implicit:
1. Average sequence length is representative. This is the critical error. The dataset of 913,786 samples, curated from diverse sources including OpenOrca, Evol-CodeAlpaca, Magicoder, Agentic-Coding-Trajectories, and Glaive Function Calling v2, contains highly variable sequence lengths. Tool-calling traces and agentic coding trajectories can be thousands of tokens long, while simple instruction-following samples may be only a few dozen tokens. The average of 335 tokens masks a long tail of very long sequences that dominate memory usage.
2. KV cache memory scales linearly with average length. In reality, KV cache memory scales with the maximum sequence length in the batch, not the average. A single outlier in a batch of 128 forces the entire batch to allocate memory for the maximum length. This is a fundamental property of transformer inference that is easy to overlook when thinking in terms of averages.
3. The model fits on one GPU with room to spare. While the 55 GB model does fit in 96 GB, the assistant underestimated the activation memory and KV cache overhead. The 40 GB "free" estimate did not account for the peak memory required by the largest sequences in the dataset.
4. The batch size heuristic from dense models transfers directly. The user's suggestion of 128–256 batch sizes likely comes from experience with dense transformer models on high-memory GPUs. But Qwen3.6-27B uses GDN (Gated Dense Network) hybrid attention, which may have different memory characteristics than standard transformers. The assistant's own earlier work showed that the model has 65 hidden states (64 layers plus embedding), and the GDN architecture includes sliding window attention layers that complicate KV cache management.
Input Knowledge Required to Understand This Message
To fully grasp what is happening in [msg 7305], several pieces of prior knowledge are necessary:
- Hardware configuration: The node has 4× NVIDIA RTX PRO 6000 Blackwell GPUs, each with 96 GB of HBM2e memory. This was established in [msg 7298].
- Model characteristics: Qwen3.6-27B is a 27-billion-parameter model in BF16, consuming approximately 55 GB of GPU memory. It uses GDN hybrid attention with 64 layers and a hidden dimension of 5120. These details were discovered during earlier debugging in [msg 7291] and [msg 7292].
- Dataset statistics: The tokenized dataset has 913,786 samples with an average sequence length of ~335 tokens. This was computed during the data preparation phase in chunk 1.
- Pipeline architecture: The extraction script loads the model via HuggingFace Transformers, iterates over dataset shards, and saves hidden states as safetensors files. The batching logic was added in [msg 7300].
- Prior low-utilization problem: In [msg 7296], the user noted "Really low gpu utilisation, 5-10%," which motivated the push for larger batch sizes. The assistant had been processing one sample at a time per GPU.
- The OOM fallback mechanism: The script includes a try-except that catches CUDA OOM errors and falls back to batch size 1. This was a defensive design choice that proves invaluable here.
Output Knowledge Created by This Message
The message produces several concrete outputs:
- Empirical refutation of the batch-size hypothesis: The test conclusively shows that batch size 128 is not viable for this dataset on this hardware, even with 96 GB GPUs. The bottleneck is not average memory but peak memory from long sequences.
- Quantitative data on sequence length extremes: The two batches had maximum lengths of 3,316 and 4,042 tokens, providing real data on the dataset's tail distribution. This is actionable information for designing the next iteration of the pipeline.
- Validation of the OOM fallback: The fallback-to-bs=1 mechanism works correctly, preventing a crash and allowing the extraction to complete (slowly). This confirms the defensive design is sound.
- A performance baseline: At ~59 seconds per sample (at bs=1), the single-GPU throughput is approximately 0.017 samples/second. This sets a lower bound that future optimizations must improve upon.
- A clear direction for the next optimization: The message implicitly defines the problem: how to batch efficiently when sequence lengths are highly variable. The solution space includes dynamic batching (grouping sequences of similar length), sequence length filtering (removing or truncating very long samples), gradient accumulation-style approaches, or using tensor parallelism to spread the memory load.
The Thinking Process Visible in the Message
The assistant's thinking process is remarkably transparent. The message opens with a reasoning statement that connects the user's suggestion to the hardware realities: "Right — with 96GB per GPU and 55GB model, we have ~40GB free for KV cache. With short sequences (~335 avg tokens), batch size 128-256 is totally viable and will saturate the GPU compute."
This is not blind agreement—it is an independent validation of the user's suggestion using concrete numbers. The assistant is saying, "Let me check if your intuition holds up against the math." The "~335 avg tokens" parenthetical reveals that the assistant has internalized the dataset statistics from earlier work and is applying them in real time.
The decision to test with --max-samples 256 --batch-size 128 is strategic: 256 samples across 2 batches of 128 gives a quick read on whether the approach works without committing to a full run. The assistant is thinking in terms of minimal experiments that maximize information per unit time.
The use of grep -v to filter out verbose transformer loading messages shows an awareness of signal-to-noise ratio in the output. The assistant wants to see only the critical information: OOM errors and timing.
The most telling aspect of the thinking process is what is not said. There is no post-hoc rationalization, no defensive justification of the failed assumption. The assistant simply presents the result and lets the numbers speak. The OOM messages are quoted verbatim, without commentary. This is the mark of an empirical mindset: the data is the authority, not the prior belief.
The Broader Significance
This message is a microcosm of a recurring pattern in ML engineering: the gap between theoretical capacity planning and empirical reality. The arithmetic said batch size 128 should work. The GPU disagreed. The resolution of this tension—through measurement, not argument—is the core methodology of the entire opencode session.
The message also illustrates a crucial property of production ML pipelines: the difference between average and worst-case performance. Many system design decisions are made based on average metrics (average sequence length, average request rate, average latency), but systems fail on worst-case inputs. The OOM error here is a worst-case failure: a few long sequences in the dataset consume disproportionate memory and break the batch.
Finally, the message demonstrates the value of fast experimentation. The test takes less than two minutes from start to finish. In that time, the assistant disproves a hypothesis, gathers quantitative data on sequence length distributions, validates a fallback mechanism, and establishes a performance baseline. This is the ideal cadence for ML engineering: measure, learn, iterate.
Conclusion
Message [msg 7305] is a small but perfect artifact of the engineering process. It contains a reasoned hypothesis, a quick experiment, an unambiguous result, and the seeds of the next iteration. The OOM error is not a setback—it is information. The assistant learns that the dataset's sequence length distribution has a long tail, that average-case memory planning is insufficient, and that the next optimization must address variable-length batching rather than simply increasing batch size. In the messages that follow, the assistant will redesign the extraction pipeline with dynamic batching, grouping sequences by length to avoid OOM while maintaining high GPU utilization. But that future work is only possible because of the clear signal produced in this message.
The user's "bigly batch" suggestion was reasonable. The assistant's test was correct. The OOM was inevitable. And the path forward is now visible.