Diagnosing the OOM: A Precision Strike on Memory Waste in Hidden State Extraction

The Message

OOM at batch=128 with long sequences (3316 tokens). The issue is output_hidden_states=True stores all 65 hidden state tensors — that's 65 × batch × seq_len × 5120 × 2 bytes. For batch=128, seq=3316: 65 × 128 × 3316 × 5120 × 2 = ~280GB. Way too much.

>

The fix: we need to use hooks to capture only the 5 layers we need, not all 65. This avoids the massive memory overhead: [edit] /data/dflash/scripts/extract_hidden_states.py Edit applied successfully.

Context: The Hidden State Extraction Pipeline

This message lands at a critical inflection point in a larger effort to train a better DFlash speculative decoding drafter for the Qwen3.6-27B model. The assistant had already spent considerable effort navigating the treacherous waters between research code and production deployment. The speculators' online vLLM pipeline proved incompatible with Qwen3.6-27B's GDN hybrid KV cache architecture, forcing a pivot to a custom offline extraction approach using HuggingFace Transformers. The user had just migrated to a new 4× RTX PRO 6000 Blackwell node after the previous instance was killed, and had explicitly requested large batch sizes: "Pretty sure we want 128-256 bigly batch" ([msg 7304]). The assistant complied, testing batch_size=128, only to be met with an out-of-memory (OOM) error and an automatic fallback to batch_size=1 ([msg 7305]). The subject message is the immediate post-mortem and corrective action.

Why This Message Was Written

The message exists because a straightforward approach — increasing batch size to saturate GPU compute — catastrophically failed. The assistant needed to understand why a 96GB GPU couldn't handle batch_size=128 with a 55GB model. The naive expectation was that with 40GB of free memory after loading the model, there would be ample room for batched inference. The OOM contradicted that expectation, forcing a diagnostic breakdown.

The reasoning is explicitly quantitative. The assistant doesn't guess or speculate about the memory pressure; it calculates exactly how much memory output_hidden_states=True consumes. The calculation — 65 × 128 × 3316 × 5120 × 2 = ~280GB — reveals that the hidden state storage alone exceeds the entire GPU memory by nearly 3×, before accounting for the model weights, KV cache, activations, and optimizer states. This is not a subtle memory leak or a configuration issue; it is a fundamental design flaw in the extraction approach.

The motivation is therefore twofold. First, the assistant must unblock the extraction pipeline to continue the DFlash training effort. Second, it must do so without sacrificing the batch size the user requested — the fix must preserve high throughput. The message is the pivot point between a broken approach and a working one.

The Thinking Process: A Quantitative Diagnosis

The assistant's reasoning is a model of systematic debugging. The thought process unfolds in three stages:

Stage 1: Observation. The extraction script OOM'd at batch_size=128 with a maximum sequence length of 3316 tokens. The fallback to batch_size=1 worked but at dramatically reduced throughput (~59 seconds for 256 samples).

Stage 2: Calculation. The assistant computes the exact memory footprint of the naive approach. The model has 65 hidden states (index 0 = embedding, indices 1-64 = the 64 transformer layers). Each hidden state is a tensor of shape [batch, seq_len, hidden_dim] where hidden_dim = 5120 for Qwen3.6-27B. With dtype=bfloat16 (2 bytes per element), the total is 65 × 128 × 3316 × 5120 × 2 = ~280 GB. This dwarfs the 96 GB available on each GPU.

Stage 3: Solution. The fix is conceptually simple but architecturally significant: instead of collecting all 65 hidden states and discarding 60 of them, capture only the 5 layers actually needed by the DFlash drafter. The DFlash training pipeline requires hidden states from specific layer indices — earlier in the conversation ([msg 7292]), the assistant had determined that the correct target layer IDs for Qwen3.6-27B are [1, 16, 31, 46, 61] (five layers out of sixty-four, plus the embedding layer at index 0 is implicitly handled). By using PyTorch's forward hooks to intercept the output of only these five layers, the memory requirement drops from 280 GB to roughly 5/65 × 280 GB ≈ 21.5 GB — well within the available 40 GB headroom.

This is a textbook example of matching the data collection strategy to the actual data requirements. The output_hidden_states=True flag is a blunt instrument that returns everything; hooks are a scalpel that extracts only what is needed.

Assumptions and Their Validity

Several assumptions underpin this message, some explicit and some implicit.

Assumption 1: The OOM is caused by hidden state storage, not by KV cache or activations. This is a reasonable first-order diagnosis. The calculation shows that hidden states alone exceed GPU memory by 3×, which is sufficient to explain the failure. However, in practice, the KV cache for long sequences (3316 tokens) with batch_size=128 also consumes significant memory — roughly 2 × 128 × 3316 × 5120 × 2 bytes ≈ 8.5 GB for a single layer's KV cache, multiplied by 64 layers. The assistant's diagnosis is correct but incomplete; the KV cache pressure compounds the problem. The hook-based fix addresses the dominant term, which is sufficient.

Assumption 2: Five layers are sufficient for DFlash training. This is inherited from the DFlash paper and the z-lab drafter configuration. The target_layer_ids [1, 16, 31, 46, 61] are spaced roughly evenly through the 64-layer network, providing a compressed representation of the model's hierarchical features. This assumption is validated by the DFlash architecture itself — the drafter is designed to predict the next token from a small set of hidden states, not from all of them.

Assumption 3: PyTorch hooks will not introduce significant overhead. Forward hooks add a function call per layer invocation. For 5 layers out of 64, the overhead is minimal — approximately 8% more Python function calls per forward pass. In practice, the overhead is dominated by the actual tensor operations, so this assumption holds.

Assumption 4: The user's batch size request (128-256) is achievable after the fix. The assistant implicitly assumes that reducing hidden state storage from 280 GB to ~21.5 GB will leave enough memory for batch_size=128. This is a reasonable assumption given the 96 GB budget: ~55 GB for model weights, ~21.5 GB for hidden states, leaving ~19.5 GB for KV cache, activations, and overhead. Whether this is sufficient depends on the actual KV cache size, which varies with sequence length. The assistant does not re-test in this message, but the subsequent conversation confirms the fix works.

Mistakes and Incorrect Assumptions

The most significant mistake is an earlier one that this message corrects: the assistant initially wrote the extraction script using output_hidden_states=True without considering its memory implications. This is a common pitfall — the flag is convenient and works well for small batches, but it scales poorly. The assistant's first extraction test with batch_size=1 and 5 samples worked fine ([msg 7293]), masking the memory issue until the user requested larger batches.

There is also a subtle error in the assistant's calculation. The formula 65 × batch × seq_len × hidden_dim × 2 bytes assumes that all 65 hidden states are simultaneously live in memory. In practice, PyTorch's output_hidden_states=True stores them as a tuple, and they are all retained until the tuple is garbage collected. However, the model's forward pass also requires memory for intermediate activations (attention scores, feed-forward activations, etc.), which are not accounted for in the calculation. The actual memory pressure is even higher than the 280 GB figure suggests. The assistant's diagnosis is directionally correct but conservatively underestimates the problem.

Another implicit assumption worth examining is that the hook-based approach is the only viable fix. An alternative would be to use gradient checkpointing or to move the hidden states to CPU incrementally during the forward pass. However, hooks are the cleanest solution — they require minimal code changes, introduce negligible overhead, and directly address the root cause (unnecessary data retention). The assistant's choice is pragmatic and correct.

Input Knowledge Required

To fully understand this message, one needs:

  1. The Qwen3.6-27B model architecture: 64 transformer layers, hidden dimension 5120, using GDN (Gated Differential Network) hybrid attention with both full and sliding window attention layers. The model uses bfloat16 precision.
  2. The DFlash speculative decoding architecture: DFlash uses a lightweight drafter model that predicts the next token from hidden states extracted at specific intermediate layers of the target model. The target_layer_ids for this model are [1, 16, 31, 46, 61].
  3. PyTorch's output_hidden_states behavior: When enabled, the model returns a tuple of all hidden states — one per layer plus the embedding layer. For a 64-layer model, this produces 65 tensors of shape [batch, seq, hidden_dim].
  4. PyTorch forward hooks: The register_forward_hook mechanism allows intercepting a layer's output during the forward pass without modifying the model code. Hooks can capture specific layer outputs and store them externally, avoiding the memory overhead of collecting all hidden states.
  5. GPU memory budgeting on Blackwell GPUs: The RTX PRO 6000 Blackwell has 96 GB of HBM3 memory. Loading a 55 GB model in bfloat16 leaves approximately 40 GB for activations, KV cache, and hidden state storage. The assistant's calculation shows that hidden states alone would consume 280 GB, making the OOM inevitable.
  6. The conversation history: The user's explicit request for batch sizes of 128-256 ([msg 7304]), the previous OOM test ([msg 7305]), and the broader context of building a training pipeline for a better DFlash drafter.

Output Knowledge Created

This message produces several valuable pieces of knowledge:

  1. A quantitative bound on hidden state memory: The explicit formula 65 × batch × seq_len × 5120 × 2 bytes provides a reusable heuristic for estimating memory requirements when extracting hidden states from large language models. This generalizes beyond Qwen3.6-27B to any transformer with known layer count and hidden dimension.
  2. A design pattern for memory-efficient extraction: The insight that hooks should be used instead of output_hidden_states=True when only a subset of layers is needed. This is a general principle for any training pipeline that requires intermediate representations from a subset of layers.
  3. A corrected extraction script: The edit to /data/dflash/scripts/extract_hidden_states.py transforms the script from a naive approach that collects everything to a targeted approach that collects only what the DFlash drafter needs. This is the immediate practical output.
  4. A debugging methodology: The message demonstrates a pattern of quantitative reasoning under memory pressure — calculate the exact memory footprint before optimizing. This contrasts with trial-and-error approaches that randomly adjust batch sizes or enable gradient checkpointing without understanding the root cause.

The Broader Significance

This message, though brief, represents a critical architectural decision in the hidden state extraction pipeline. Before this fix, the pipeline was fundamentally unscalable — it could handle small batches but would OOM on any production-relevant batch size. After this fix, the pipeline can scale to batch sizes of 128 or even 256, saturating the GPU compute and achieving the throughput the user requested.

The message also illustrates a recurring theme in this coding session: the gap between research convenience and production efficiency. The output_hidden_states=True flag is designed for research use cases where you want to inspect all layers of a model on a single sample. It is not designed for production data pipelines that process hundreds of thousands of samples at scale. The assistant's pivot to hooks is a recognition that production engineering requires different tools than research prototyping.

In the broader narrative of segment 43, this message is the turning point where the extraction pipeline transforms from a proof-of-concept (5 samples in 1.3 seconds per GPU) into a production system capable of processing 913,786 samples in under an hour. The subsequent messages in the chunk show the pipeline achieving 140–155 samples/s per GPU — a 20× improvement over the initial 7–11 samples/s. That improvement would have been impossible without the hook-based fix introduced in this message.