The Activation Memory Problem: Debugging GPU OOM in Hidden State Extraction

Introduction

In the course of building a high-throughput hidden state extraction pipeline for training a DFlash speculative decoding drafter, the assistant encountered a persistent GPU out-of-memory (OOM) error that resisted two rounds of attempted fixes. Message [msg 7309] captures the moment of insight when the assistant correctly identified the root cause and pivoted to a solution that would ultimately unlock the pipeline's full performance. This message is a microcosm of the broader engineering challenge: bridging the gap between research code and production-grade ML infrastructure, where theoretical understanding of memory usage must be matched against the concrete constraints of GPU hardware.

Context: The DFlash Drafter Training Pipeline

The broader session (Segment 43) had migrated from deploying existing speculative decoding methods to building the infrastructure required to train better draft models. After discovering that DFlash and DDTree speculative decoding suffered from low acceptance rates—primarily because the underlying drafter model was labeled "still under training"—the assistant pivoted to training a new drafter. This required extracting hidden states from the Qwen3.6-27B target model across a curated 913K-sample dataset.

The extraction pipeline initially ran at only 7–11 samples/s per GPU with high CPU overhead. The key breakthrough had been batching hidden state capture entirely on GPU—concatenating all samples in a batch before a single .cpu() transfer—which eliminated thousands of individual copies and boosted throughput to 140–155 samples/s per GPU. However, this batching approach introduced a new problem: GPU memory pressure.

The OOM Problem: Two Failed Fixes

The assistant had attempted to use batch size 128–256 as suggested by the user ("Pretty sure we want 128-256 bigly batch," [msg 7304]), reasoning that with 96GB per GPU and a 55GB model, approximately 40GB would be available for activations and KV cache. With short sequences averaging ~335 tokens, this seemed viable.

First attempt ([msg 7305]): Testing batch_size=128 resulted in immediate OOM for long sequences (3316–4042 tokens), with a fallback to batch_size=1 that reduced throughput to ~2 samples/s. The assistant correctly identified that output_hidden_states=True was storing all 65 hidden state tensors (embedding + 64 layers), computing the memory cost as 65 × batch × seq_len × 5120 × 2 bytes ≈ 280GB for batch=128, seq=3316.

Second attempt (<msg id=7306-7308>): The fix was to use PyTorch forward hooks to capture only the 5 specific layers needed for DFlash training (layers 1, 16, 31, 46, 61), rather than all 65. This eliminated the ~280GB hidden state storage overhead. However, testing batch=256 with hooks still produced an OOM error: "CUDA out of memory. Tried to allocate 8.10 GiB. GPU 0 has a total capacity of 94.97 GiB of which 4.95 GiB is free."

The Insight in Message 7309

The target message reads:

Still OOM even with hooks — the problem is the long sequences. Max len 4042-4096 tokens × batch 128-256 is too much. The forward pass activations (not just hidden states) are what's eating memory. With 64 layers of compute, batch × seq_len × hidden_size activations per layer add up fast.

>

The solution: sort samples by length and use dynamic batch sizing — small batches for long sequences, large batches for short ones. Most of our samples are short (avg 335 tokens), so we'll get huge batches there.

This message represents a significant refinement in the assistant's mental model of GPU memory usage. The key realization is that the OOM was not caused by storing hidden states (which had been fixed by hooks), but by the forward pass activations themselves. Every one of the model's 64 layers produces intermediate activations of shape batch × seq_len × hidden_size that must be kept in memory for the backward pass (or, in this inference-only case, for the subsequent layers to consume). Even without output_hidden_states=True, the forward pass through 64 transformer layers with batch=256 and seq_len=4042 generates enormous activation memory.

The Reasoning Process

The assistant's reasoning follows a clear diagnostic pattern:

  1. Observation: Hooks eliminated hidden state storage, but OOM persists.
  2. Hypothesis refinement: The remaining memory consumer must be forward pass activations, not stored outputs.
  3. Quantitative reasoning: 64 layers × batch × seq_len × hidden_size activations per layer. For batch=256, seq=4042, hidden=5120, each layer's activations are ~5GB in BF16. Across 64 layers, even with activation checkpointing or selective retention, the peak memory is substantial.
  4. Solution design: Rather than reducing batch size uniformly (which would hurt throughput on the majority of short sequences), sort samples by length and use dynamic batch sizing. Short sequences (avg 335 tokens) can use large batches; long sequences (up to 4096 tokens) use small batches. This is a classic engineering tradeoff: throughput vs. memory. The insight is that the dataset has a highly skewed length distribution—most samples are short, but a few are very long. A fixed batch size either wastes GPU capacity (if set conservatively) or OOMs (if set aggressively). Dynamic batching based on sorted lengths optimizes for both.

Assumptions Made

The assistant makes several assumptions in this message:

  1. That the dataset's length distribution is heavily skewed toward short sequences. This is based on the observed average of 335 tokens and the presence of samples up to 4096 tokens. This assumption is correct for the curated dataset, which mixes instruction-following, code generation, and tool-calling data—all typically short.
  2. That forward pass activations are the dominant memory consumer after hidden state storage is eliminated. This is technically correct for transformer inference: each layer's attention and MLP computations produce activations proportional to batch × seq_len × hidden_size. However, the assistant does not quantify exactly how much memory each layer consumes, nor does it account for PyTorch's memory fragmentation or CUDA caching allocator behavior.
  3. That sorting by length and dynamic batching will not introduce significant overhead. Sorting the dataset is O(N log N) and is done once. Dynamic batching requires grouping samples of similar lengths, which is straightforward with sorted data. The assumption is that the overhead of sorting is negligible compared to the throughput gain.
  4. That the model fits in 55GB with hooks active. The model itself consumes ~55GB in BF16. With hooks registered, the assistant assumes the model can still load within the 96GB GPU memory, leaving ~40GB for activations. This proved approximately correct—the OOM occurred only with the largest sequences.

Mistakes and Incorrect Assumptions

The primary mistake in the preceding messages was the initial assumption that hidden state storage was the sole cause of OOM. The assistant computed that output_hidden_states=True required ~280GB for a single batch, which was clearly impossible. Switching to hooks eliminated this storage, but the OOM persisted. The assistant had implicitly assumed that without storing all hidden states, the remaining memory would be sufficient—an assumption that overlooked the activation memory of the forward pass itself.

A secondary error was underestimating the memory cost of long sequences. The assistant's mental model focused on the average sequence length (335 tokens) and assumed batch size 128-256 would be viable. The presence of sequences 10-12× longer than average meant that even a single batch of long sequences could exhaust memory. This is a common pitfall in sequence modeling: averages can be misleading when the distribution has a long tail.

The assistant also did not initially account for PyTorch's memory fragmentation. The OOM message showed "4.95 GiB is free" out of 94.97 GiB total, with "90.02 GiB memory in use." This suggests that even after the OOM, 4.95 GiB was free but unavailable as a contiguous block—a classic fragmentation issue that dynamic batching would partially mitigate by using smaller, more predictable allocations.

Input Knowledge Required

To understand this message, the reader needs:

  1. Transformer architecture knowledge: Understanding that each layer produces intermediate activations of shape [batch, seq_len, hidden_size] and that 64 layers of these accumulate significant memory even without explicit storage.
  2. GPU memory model: Knowledge that GPU memory is shared between model weights, activations, optimizer states (for training), and CUDA caching allocator overhead. For inference, the primary consumers are weights and activations.
  3. PyTorch hooks: Understanding that register_forward_hook can capture intermediate layer outputs without storing all hidden states via output_hidden_states=True.
  4. BF16 memory sizing: 5120 hidden dimension × 2 bytes per parameter = ~10KB per token per layer for activations.
  5. The DFlash training context: DFlash requires hidden states from specific layers (1, 16, 31, 46, 61) of the target model, not all layers.

Output Knowledge Created

This message creates:

  1. A concrete solution: Sort samples by length and use dynamic batch sizing. This is implemented immediately in the edit to extract_hidden_states.py.
  2. A refined mental model of GPU memory: The distinction between explicit storage (hidden states saved to tensors) and implicit storage (forward pass activations that exist transiently but consume memory). The assistant now correctly attributes OOM to the latter.
  3. A design pattern for robust batch processing: Dynamic batching based on sorted lengths is a general technique applicable to any sequence processing pipeline with variable-length inputs.

The Broader Significance

This message exemplifies a pattern that recurs throughout the opencode session: the gap between theoretical understanding and practical deployment. The assistant knew that forward pass activations consume memory, but had not fully accounted for them in the mental model until empirical testing forced the realization. The solution—sorting by length and dynamic batching—is elegant precisely because it works with the data's natural distribution rather than against it.

The dynamic batching approach also has implications beyond this single pipeline. It represents a shift from a uniform, one-size-fits-all processing strategy to an adaptive one that respects the heterogeneity of real-world data. In production ML systems, this kind of adaptive resource management is essential for achieving high throughput without sacrificing reliability.

Conclusion

Message [msg 7309] captures a critical diagnostic breakthrough in the hidden state extraction pipeline. The assistant correctly identified that forward pass activations—not hidden state storage—were causing GPU OOM, and proposed sorting samples by length with dynamic batch sizing as the solution. This insight, while technically straightforward, required synthesizing knowledge of transformer architecture, GPU memory management, and dataset statistics into a coherent diagnosis. The message demonstrates the iterative nature of debugging ML infrastructure: each failed hypothesis refines the mental model until the true cause is identified.