The 15-Gigabyte Ghost: Debugging a Persistent OOM in DFlash Training on Blackwell GPUs

Introduction

In the high-stakes world of large language model training, few experiences are as frustrating as a crash that reproduces identically regardless of what parameter you change. When the assistant in this opencode session encountered an out-of-memory (OOM) error that consistently allocated exactly 15.09 GiB on GPU 2—whether using 512, 256, or 128 anchors—it triggered a deep investigation that would span multiple hypotheses, memory accounting exercises, and ultimately reveal a fundamental limitation of unfused attention implementations on NVIDIA's Blackwell architecture.

This article examines message 7885 from a lengthy coding session focused on training a DFlash (Drafting with Flash Attention) speculative decoding drafter for the Qwen3.6-27B language model. The message captures a pivotal moment in the debugging process: the assistant realizes that the OOM is not coming from the target model (as previously suspected) but from the drafter's flex_attention backward pass, and that the unfused implementation materializes a full score matrix regardless of the anchor count. What follows is a masterclass in systems-level debugging—tracing through GPU memory allocation, understanding PyTorch's attention backend selection, and grappling with the tension between algorithmic efficiency and hardware constraints.

The Context: Training a Speculative Decoding Drafter

To understand the significance of this message, we need to situate it within the broader project. The team is training a DFlash drafter—a lightweight model that predicts multiple future tokens in parallel to accelerate inference of the main Qwen3.6-27B model. The training setup spans four NVIDIA RTX PRO 6000 Blackwell GPUs (each with 96 GB of VRAM), organized into two data-parallel pairs: GPUs 0/1 run the frozen target model, and GPUs 2/3 run the trainable drafter.

The drafter uses a specialized attention mechanism called flex_attention, which supports block-sparse attention masks. This is critical for the DFlash architecture, which processes anchor positions—a subset of token positions where the drafter predicts future tokens. The attention pattern combines a prefix lookahead with a "same block" document isolation mask, a complex pattern that cannot be expressed with standard causal masking.

The training pipeline had already survived a gauntlet of issues: six bugs in the training script (incorrect config copying, missing sequence packing, absent noise augmentation, per-document anchor boundary violations, incorrect position IDs, and lack of torch.compile), followed by FLA Triton autotuner crashes on the Blackwell architecture (sm_120). After upgrading Triton to 3.7.0 and implementing sequential warmup to avoid concurrent autotuner calls, the team believed they had stabilized the training run. But a new adversary emerged: a persistent OOM that struck on the very first training step.

The Message: A Debugging Odyssey

Message 7885 opens with the assistant in the midst of confusion. The previous message (msg 7884) had shown a crash trace pointing to flex_attention.py—the flex_attention backward pass on GPU 2. But earlier in the log, the assistant had seen a crash on line 253 of the training script, which is the target model's forward pass. This ambiguity sets the stage for the reasoning chain that follows.

Step 1: Disambiguating the Crash Source

The assistant's first move is to question which process generated the log file:

"Wait, now the error is BACK to being on the flex_attention backward on GPU 2 (drafter), not the target model. But earlier I saw the crash was on line 253 (target_model forward). Let me look more carefully—the log file /workspace/train.log may be from the background setsid process, not the interactive one I just ran."

This is a crucial moment of meta-cognition. The assistant realizes it may have been looking at stale log output from a previous background process, not the freshly launched interactive run. The 2>&1 | head -50 pipe in the previous command meant stdout and stderr were piped through head, which could have terminated early or timed out. The assistant is essentially reconstructing the timeline of events from fragmented evidence.

This realization is important because it explains why the assistant had been oscillating between two different failure modes. The target model OOM (from large batches of short sequences on GPUs 0/1) and the drafter flex_attention OOM (on GPUs 2/3) are separate problems, but they appeared interleaved in the logs because different processes were writing to the same file.

Step 2: Recognizing the Two Distinct OOM Problems

With this clarified, the assistant identifies the dual nature of the problem:

"Looking at the background runs, they're hitting different failure points depending on batch composition—large batches crash the target model on GPU 0/1, while the drafter hits unfused flex_attention backward OOM on GPU 2/3."

The target model OOM occurs because batches can contain up to 67 samples (when sequences are very short, around 122 tokens each). With token_budget=8192, the padded batch tensor is [67, 122], and the target model's full attention layers have quadratic memory scaling with batch size. The GDN (Global Dense Network) state is also proportional to batch size, pushing memory past the 96 GB limit.

The drafter OOM, however, is more mysterious. It consistently allocates 15.09 GiB regardless of the anchor count—whether 512, 256, or 128 anchors. This invariance is the key clue that drives the investigation.

Step 3: Exploring Fix Options for the Drafter

The assistant considers three approaches to fix the drafter's OOM:

(a) Getting torch.compile working for flex_attention. The fused Triton kernel generated by torch.compile would avoid materializing the full score matrix, reducing memory from O(Q·KV·heads) to O(Q·block_size·heads). This is the ideal fix but requires understanding why compilation isn't working.

(b) Replacing flex_attention with standard SDPA and manual masking. This seems simpler but has a fatal flaw: the complex block mask pattern (prefix + same block with document isolation) cannot be expressed with SDPA's limited causal mask alone. Constructing a full boolean mask would require materializing the Q×KV attention matrix anyway, defeating the purpose.

(c) Manually chunking the attention computation. This would involve splitting the query and key-value sequences into smaller blocks and computing attention incrementally, avoiding the full score matrix. This is feasible but requires significant code changes.

The assistant correctly identifies that option (b) is a dead end:

"SDPA with custom boolean masks falls back to the math backend when the mask doesn't fit standard patterns, which means the attention scores get materialized anyway—leading to the same OOM problem."

This demonstrates a deep understanding of PyTorch's attention dispatch mechanism. The math backend is the fallback that computes attention via explicit matrix multiplication, which is exactly what causes the OOM.

Step 4: The torch.compile Hypothesis

The assistant zeroes in on option (a) as the most promising path:

"I think flex_attention with torch.compile is still the right path, but I need to figure out why the compilation isn't working properly for the drafter. The problem might be that create_block_mask is dynamic and depends on runtime anchor positions, causing torch.compile to graph-break and fall back to eager execution."

This is a sophisticated insight. torch.compile works by tracing the computation graph and generating fused kernels. If the graph contains dynamic operations that change shape or structure at runtime (like create_block_mask, which depends on the specific anchor positions selected for each batch), torch.compile may "graph-break"—falling back to eager-mode execution for those portions. The fix would be to isolate the flex_attention call in its own compiled function, taking pre-computed Q, K, V, and the BlockMask as inputs, so that the dynamic mask creation happens outside the compiled region.

Step 5: The Memory Accounting Exercise

Perhaps the most impressive section of the reasoning is the detailed memory accounting. The assistant attempts to reconstruct exactly what is consuming GPU memory on the drafter (GPU 2) by tracing through every tensor allocation:

Model weights (BF16):

Step 6: The 15.09 GiB Mystery

The assistant repeatedly returns to the 15.09 GiB allocation, trying to reverse-engineer what tensor it represents:

"15.09 * 1024^3 = 16,199,917,568 bytes. That's either 4 billion elements in float32 or 8 billion in BF16."

Working backward from the assumption that this is an attention score matrix with 32 heads:

"If I assume this is attention-related with 32 heads, I can back-calculate what the sequence length and KV cache dimensions would need to be, which gives me either Q=8192 with KV=15,449 or Q=2048 with KV=61,796."

Neither of these makes sense for the configured parameters (max_anchors=128, block_size=16, giving Q=2048). The assistant also considers whether it could be the gradient of lm_head.weight (248320 × 5120 × 12 bytes ≈ 15.26 GB) but can't explain the factor of 12.

This mystery is the central puzzle of the message. The invariance of the allocation size across different anchor counts strongly suggests that the tensor in question is not the attention score matrix at all, but something else that doesn't depend on the anchor count. The assistant's inability to identify the tensor from first principles leads to the decision to add memory profiling directly to the training step.

Step 7: The Bash Command

The message concludes with a bash command that attempts to probe flex_attention's backward behavior:

ssh -o StrictHostKeyChecking=no -p 10638 root@154.59.156.41 'source /root/venv/bin/activate && python3 -c "
import torch
print(f\"flex_attention backward modes:\")
from torch.nn.attention import SDPBackend
print(f\"Available: {[b.name for b in SDPBackend]}\")
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
from torch._higher_order_ops.flex_attention import flex_attention_backward
compiled_fa = torch.compile(flex_attention)
print(f\"Compiled flex_attention: {type(compiled_fa)}\")
import triton
print(f\"Triton version: {triton.__version__}\")
"'

The command fails with TypeError: 'pybind11_type' object is not iterable on line 6, where it tries to iterate over SDPBackend. This is a minor API issue—SDPBackend is a pybind11 enum, not a Python iterable. The failure is itself informative: it shows that the assistant is working at the edge of PyTorch's API, where documentation is sparse and behavior is platform-dependent.

Assumptions Made by the Assistant

Throughout the reasoning, the assistant makes several assumptions that shape the investigation:

  1. The 15.09 GiB allocation is invariant across anchor counts. This is the central observation that drives the investigation away from anchor-related parameters. The assistant assumes this invariance means the tensor is not the attention score matrix, since the score matrix dimensions depend on Q_len (which is proportional to max_anchors).
  2. The unfused backward is the culprit. The assistant assumes that the OOM is caused by flex_attention's unfused backward pass materializing the full score matrix, rather than some other tensor allocation. This is supported by the stack trace pointing to flex_attention.py.
  3. torch.compile would fix the problem if it worked. The assistant assumes that the fused Triton kernel generated by torch.compile would avoid the score matrix materialization. This is correct in principle, but the assumption that compilation can be made to work on Blackwell (sm_120) is less certain.
  4. The two OOM problems are independent. The assistant treats the target model OOM and the drafter OOM as separate issues requiring separate fixes. This is a reasonable decomposition but risks missing interactions (e.g., memory fragmentation from one affecting the other via CUDA's memory allocator).
  5. The batch composition after shuffling could explain the variance. The assistant considers that shuffled batches might contain long sequences that exceed the token budget, but then notes that the 15.09 GiB allocation is consistent across runs, which argues against a data-dependent explanation.

Mistakes and Incorrect Assumptions

Several aspects of the reasoning contain errors or incomplete understanding:

  1. The SDPBackend iteration error. The assistant attempts to iterate over SDPBackend as if it were a list-like enum, but it's a pybind11 type that doesn't support iteration. This is a minor API mistake but reflects the difficulty of working with cutting-edge PyTorch features where the API is still evolving.
  2. The log file confusion. The assistant initially misattributes the crash to the target model based on stale log output. While it eventually corrects this, the confusion costs several rounds of debugging effort. This is a classic systems debugging pitfall: trusting log output without verifying which process generated it.
  3. The memory accounting gap. The assistant calculates ~46 GB of expected memory usage but observes 84 GB in practice. The 38 GB gap is never fully explained within this message. Possible explanations include: CUDA memory fragmentation, cached allocator retention, duplicated tensors from the ThreadPoolExecutor, or intermediate activations from the verifier logits computation that the assistant underestimated.
  4. The assumption that the 15.09 GiB allocation is the same tensor each time. While the allocation size is identical, it could be different tensors that happen to have the same size. The assistant doesn't consider this possibility explicitly.
  5. Overlooking the torch.no_grad() interaction. The assistant notes that the verifier logits are computed under torch.no_grad(), but doesn't fully explore whether the intermediate activations from this computation are being retained. The no_grad() context prevents gradient computation but doesn't automatically free intermediate tensors—they're still allocated until they go out of scope.

Input Knowledge Required

To fully understand this message, the reader needs:

  1. Understanding of speculative decoding and DFlash architecture. Knowledge of how drafters predict multiple future tokens and how anchor-based attention works is essential.
  2. Familiarity with flex_attention and block-sparse attention. Understanding that flex_attention supports custom score_mod and mask_mod functions, and that it can use either a fused Triton kernel (with torch.compile) or an unfused implementation that materializes the full score matrix.
  3. Knowledge of PyTorch's attention dispatch mechanism. Understanding the hierarchy of backends (FlashAttention, EfficientAttention, Math) and when each is selected.
  4. GPU memory management concepts. Understanding of model weights, optimizer states, intermediate activations, and gradient buffers in the context of training large models.
  5. The Blackwell GPU architecture (sm_120). Knowledge that Blackwell is a new architecture with specific limitations, including potential issues with Triton kernel compilation.
  6. CUDA memory allocation and the expandable_segments feature. Understanding of how PyTorch's CUDA allocator works and how PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True affects memory management.

Output Knowledge Created

This message creates several important pieces of knowledge:

  1. The identification of two distinct OOM failure modes in the DFlash training setup: target model OOM from large batches, and drafter OOM from unfused flex_attention backward.
  2. The invariance of the 15.09 GiB allocation across anchor counts, which rules out anchor-dependent explanations and points to a fundamental issue with the unfused attention implementation.
  3. The infeasibility of replacing flex_attention with SDPA for complex block mask patterns, documented with specific reasoning about mask construction and backend fallback behavior.
  4. A detailed memory budget analysis for the drafter GPU, tracing through weights, optimizer states, forward intermediates, and backward buffers.
  5. The torch.compile graph-breaking hypothesis for why flex_attention isn't using its fused kernel, with a proposed fix of isolating the attention call in its own compiled function.
  6. A practical debugging methodology that combines log analysis, memory accounting, hypothesis formation, and targeted probing commands.

The Thinking Process: A Window into Systems Debugging

What makes this message exceptional is not any single insight but the process of reasoning. The assistant demonstrates several hallmarks of expert systems debugging:

Iterative hypothesis refinement. The assistant doesn't commit to a single explanation but continuously updates its understanding based on new evidence. The shift from "target model OOM" to "drafter OOM" to "two independent OOM problems" is a textbook example of Bayesian updating in debugging.

Memory as a first-principles constraint. Rather than relying on profiling tools alone, the assistant computes memory budgets from scratch, tracing through every tensor allocation. This reveals the 38 GB gap that profiling alone might not explain.

Understanding the toolchain stack. The assistant reasons about multiple layers of the software stack simultaneously: the Triton compiler, PyTorch's autograd engine, the CUDA memory allocator, and the Blackwell hardware architecture. This systems-level thinking is essential for debugging issues at the intersection of these components.

Pragmatic trade-off analysis. When evaluating fix options, the assistant considers not just correctness but feasibility, development cost, and risk. The analysis of SDPA vs. flex_attention vs. manual chunking weighs the benefits of each against the practical constraints of implementation.

Meta-cognition about data sources. The assistant questions whether the log file it's reading is from the correct process, recognizing that stale or interleaved output can mislead. This attention to data provenance is a hallmark of experienced debuggers.

Conclusion

Message 7885 captures a pivotal moment in a challenging debugging session. The assistant confronts a persistent OOM error that resists simple parameter tuning, systematically works through the possible explanations, and arrives at a refined understanding of the problem. The 15.09 GiB ghost—a tensor allocation that appears identically regardless of configuration—remains unidentified within this message, but the investigation has narrowed the possibilities and charted a path forward.

The broader significance of this message lies in what it reveals about the state of large model training on cutting-edge hardware. The Blackwell GPU architecture (sm_120) is so new that many software components—Triton, PyTorch's flex_attention, FLA kernels—are still being adapted to support it. Debugging on this frontier requires not just knowledge of deep learning but deep systems expertise: understanding GPU memory architectures, compiler internals, and the intricate dance between eager-mode execution and graph compilation.

For practitioners training large models on new hardware, this message offers both a cautionary tale and a methodological guide. The caution: even after fixing six training bugs and resolving autotuner crashes, new failure modes emerge at the boundary between software and hardware. The guide: systematic reasoning, memory accounting, and hypothesis-driven debugging can navigate this complexity, one 15-gigabyte mystery at a time.