The Anatomy of an OOM: How Gradient Checkpointing Saved a Speculative Decoding Training Pipeline

Introduction

In the high-stakes world of large language model training, few errors strike as much dread into a practitioner's heart as the dreaded "CUDA Out of Memory" (OOM). When you're pushing the boundaries of what fits on a single GPU—juggling model weights, optimizer states, activations, and intermediate tensors across a complex training pipeline—the line between success and failure is measured in megabytes. This article examines a single message from an opencode coding session that captures, in excruciating detail, one such moment: the assistant's deep-dive debugging of an OOM error in the KL divergence computation of a DFlash (Draft-and-Flash) speculative decoding drafter training pipeline.

The message in question, <msg id=9305>, is remarkable not because it contains a clean solution, but because it documents the process of reaching that solution. It is a raw, unedited window into the assistant's reasoning as it grapples with a memory constraint that is just 430 MB short of fitting. The reasoning traces through multiple failed hypotheses, recalculates memory budgets from first principles, discovers a fundamental misunderstanding about how PyTorch's backward graph accumulates memory across chunked operations, and ultimately converges on gradient checkpointing as the correct fix. It is a masterclass in systematic debugging under tight resource constraints.

Let us examine this message in full, then unpack its layers of reasoning, assumptions, mistakes, and insights.

The Message in Full

The message begins with the assistant's agent reasoning section, which is an extended internal monologue triggered by an OOM error during training. The assistant has just deployed a fused chunked loss function (in the preceding messages <msg id=9296> through <msg id=9304>) designed to avoid materializing full [T, 248320] logit tensors. Despite this optimization, the training run crashes during the KL divergence computation. The reasoning opens with:

Still OOM in KL. The chunk_size=2048 is too big for KL. The kl_div with [2048, 248320] inputs needs 20482483202=1.0 GB for the output, plus the two input tensors are already ~1 GB each. So peak during KL is 3 GB.

>

But the issue is that at this point, there's already 93.4 GB allocated! The model + optimizer + fc + attention activations for 1024*32=32768 block tokens are already at the limit.

From this starting point, the assistant embarks on a systematic memory budget analysis that spans the entire reasoning section, ultimately arriving at the insight that gradient checkpointing is required to prevent the backward graph from accumulating 16+ GB of saved logit tensors across all chunks.

The message concludes with a bash command that kills the failed training session:

[bash] ssh -o ConnectTimeout=10 root@10.1.2.6 'pct exec 200 -- tmux kill-session -t dflash 2>/dev/null; echo stopped' 2>&1
stopped

Context: The DFlash Training Pipeline

To understand why this message was written, we must first understand what was being built. The DFlash (Draft-and-Flash) project is a speculative decoding system where a small "drafter" model learns to predict multiple tokens per forward pass of a large "target" model. The drafter is trained via distillation: it learns to match the target model's output distribution at multiple future positions simultaneously.

The training pipeline operates across 8 GPUs (5 for the target model, 3 for the drafter) and processes data in a complex arrangement of "blocks" and "anchors." Each training step processes sequences where the drafter must predict tokens at positions determined by a tree structure (DDTree). The key hyperparameters are block_size=32 (the spacing between anchor positions) and max_anchors=1024 (the number of anchor positions per sequence), yielding a total of 32,768 "block tokens" that need to be processed through the drafter's transformer layers.

The critical architectural detail is that the loss function blends two components: a hard cross-entropy (CE) loss that teaches the drafter to predict the exact target token, and a soft KL divergence loss that teaches the drafter to match the target model's full probability distribution. The KL divergence is computed over the entire vocabulary of 248,320 tokens, which means every position in the chunk produces a tensor of shape [chunk_size, 248320]—a 970 MB tensor at chunk_size 2048.

The fused chunked loss function, implemented in the immediately preceding messages, was designed to process these logit tensors in chunks to avoid holding the full [32768, 248320] tensor in memory (which would be ~16 GB). However, the chunking strategy, as the assistant is about to discover, has a subtle flaw: it only reduces peak instantaneous memory, not the total memory consumed by the backward computation graph.

The Surface Problem: KL Divergence at Scale

The immediate trigger for this reasoning is straightforward: the KL divergence computation runs out of memory. The assistant calculates that at chunk_size 2048, the KL divergence output alone requires 1.0 GB, and with the two input tensors (student logits and teacher probabilities) each at ~1 GB, the peak during KL computation reaches 3 GB. But the GPU already has 93.4 GB allocated before this point, leaving only about 543 MB free—far short of the 970 MB needed for the KL output.

This is the classic OOM scenario: you're not out of memory globally, but you're out of memory in the specific region where a particular operation needs to allocate. The assistant's first instinct is to trace through what's consuming the 93.4 GB.

The Memory Budget Analysis

The reasoning then launches into a detailed memory budget breakdown. This is the heart of the message and deserves careful examination. The assistant walks through:

  1. Model weights and optimizer states: ~20.76 GB combined. This includes the drafter's 5 transformer layers, the fully connected projection layers, the embedding and language model head, and their corresponding AdamW optimizer states (which typically double the parameter memory due to momentum and variance buffers).
  2. Flex attention computation: The attention mechanism processes 32,768 query tokens against 81,920 key/value tokens. The flex_attention kernel preallocates workspace memory, and the assistant realizes this is a significant but hard-to-quantify contributor.
  3. Backward pass activations per layer: Each transformer layer needs to save Q, K, V, and attention outputs for the backward pass (~0.82 GB per layer, ~4.1 GB across 5 layers). The MLP layers are even more demanding: the intermediate expansion to 17,408 dimensions creates a 1.09 GB footprint per layer, totaling ~8.5 GB across all layers.
  4. Context K/V projections: This is where the assistant initially makes a significant error in its mental model. It calculates ~37.7 GB for the obvious components, which is only ~40% of the 95 GB budget, and realizes it's missing something crucial: "the context K/V projections are being computed for all 49,152 context positions across every layer, which is a massive amount of memory I didn't account for initially." This moment of realization is instructive. The assistant had been focusing on the block tokens (the 32,768 positions being trained) but had forgotten about the context tokens (the 49,152 positions of the prefix sequence that the target model processes). These context tokens flow through the drafter's layers as key/value cache inputs, and their projections create substantial intermediate tensors that persist across all five layers.
  5. The backward graph accumulation: This is the critical insight that the assistant arrives at after multiple rounds of recalculation. The chunking strategy processes 16 chunks sequentially during the forward pass, computing logits, loss, and metrics for each chunk. But PyTorch's autograd system saves all intermediate tensors that are needed for backward computation. Each chunk's logits tensor (970 MB) is retained in the backward graph until the backward pass processes that chunk. With 16 chunks, that's 16 × 970 MB = ~15.5 GB of logit tensors saved simultaneously. This is the fundamental flaw in the naive chunking approach: it successfully reduces peak forward memory (you never hold all chunks' logits at once during forward), but it does nothing to reduce the backward graph memory, because PyTorch must retain all chunk outputs until the backward pass can process them.

The Backward Graph Revelation

The assistant's reasoning on this point is particularly sharp:

So the backward graph holds: logits_c, s_log, t_prob, kl_c output. That's 4 tensors of 970 MB each = 3.88 GB just for backward.

>

This is the real issue. Even with small chunk_size, the backward graph for KL grows.

>

Wait, but kl_c is aggregated into weighted_loss_sum = weighted_loss_sum + loss_c.sum(). The .sum() reduces to a scalar, so the backward graph only needs to reconstruct the per-position kl values, which requires s_log and t_prob.

The assistant is reasoning through PyTorch's autograd mechanics. When you compute loss_c.sum(), the resulting scalar is what's added to weighted_loss_sum. During backward, PyTorch needs to compute gradients through the chain: weighted_loss_sum → loss_c.sum() → kl_c → s_log, t_prob → logits_c, targets_c. The .sum() operation means the backward graph only needs to materialize the per-position KL values during the backward pass itself (when processing that chunk), not continuously throughout the forward pass.

But the key insight is that logits_c and s_log (the log-softmax of student logits) are both 970 MB tensors that must be retained because they're needed for the backward pass of the KL divergence. The t_prob tensor (teacher probabilities from softmax) can potentially be freed because it doesn't require gradients, but the student-side tensors are non-negotiable.

The assistant calculates: with 16 chunks, the backward graph contains 16 × 1.94 GB (logits_c + s_log per chunk) = ~31 GB. This is where the memory is actually going, and it explains why the naive chunking approach fails.

The Gradient Checkpointing Solution

The assistant's reasoning converges on gradient checkpointing as the correct solution:

The real issue is that chunking only reduces peak instantaneous memory, not the total backward graph memory. To actually save memory, I need gradient checkpointing to recompute logits during backward instead of saving them. Using torch.utils.checkpoint.checkpoint on the fused computation per chunk means logits_c gets freed after the forward pass and recomputed during backward from the saved normalized features.

Gradient checkpointing (also known as activation checkpointing) is a technique where intermediate tensors are freed during the forward pass and recomputed on-demand during the backward pass. The key trade-off is memory versus computation: you save memory by not storing intermediates, but you spend extra compute by recomputing them.

The assistant realizes that by checkpointing the entire lm_head → loss computation per chunk, only the normalized embeddings (the input to the checkpointed block) need to be saved. These are tiny: [chunk_size, 5120] in bf16 is only ~0.06 GB per chunk. During backward, PyTorch re-runs the checkpointed function to recompute logits, KL divergence, and their gradients, then frees them before moving to the next chunk.

This brings the backward graph memory from ~31 GB down to under 1 GB—a 30× reduction.

Assumptions and Mistakes

The reasoning reveals several assumptions and mistakes that are worth examining:

Mistake 1: Underestimating Context Token Memory

The assistant initially calculates ~37.7 GB for the obvious components and is puzzled that this is only 40% of the budget. It then realizes it forgot about the context K/V projections for all 49,152 context positions across all layers. This is a common blind spot: when designing a chunked loss function, it's easy to focus on the chunked path and forget about the unchunked parts of the pipeline (like context processing) that consume memory continuously.

Mistake 2: Assuming Chunking Solves the Backward Problem

The assistant initially assumes that chunking the forward pass automatically solves the memory problem. It takes several iterations of reasoning to realize that PyTorch's autograd retains all chunk outputs for backward, regardless of how they were produced. This is a subtle but critical point: chunking reduces peak allocation during forward, but the total allocated memory (including backward graph) can still exceed the budget.

Mistake 3: Overestimating KL Divergence Memory

The assistant initially calculates that KL divergence requires 4.85 GB peak (five copies of the chunk tensor simultaneously). After more careful reasoning, it realizes that t_prob (teacher probabilities) doesn't need gradients and can be freed after the forward pass, reducing the peak to 3.88 GB. This refinement shows the importance of understanding which tensors participate in the backward graph.

Mistake 4: The torch.no_grad() Interaction with Checkpointing

The assistant worries about a subtle interaction: the loss function uses torch.no_grad() for computing teacher probabilities (since the teacher model is frozen), but gradient checkpointing re-runs the function during backward with gradients enabled. The assistant reasons through this correctly: the no_grad() context still applies during the re-run, so gradients won't flow through the teacher path, which is the desired behavior. This is a nuanced understanding of PyTorch's autograd mechanics.

Assumption: The 95 GB Budget is Fixed

The assistant assumes a hard 95 GB memory budget (the GPU memory capacity) and treats any excess as a failure. This is correct for the immediate debugging context, but it's worth noting that techniques like memory pooling, fragmentation reduction, and expandable_segments (which was already enabled via PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True) can sometimes recover additional memory by reducing fragmentation.

Input Knowledge Required

To fully understand this message, one needs:

  1. PyTorch autograd mechanics: How the backward graph accumulates tensors, how .sum() reduces memory requirements, how gradient checkpointing works, and how torch.no_grad() interacts with checkpoint re-runs.
  2. Transformer architecture: The structure of attention (Q, K, V projections), MLP layers (gate, up, down projections), and the language model head (lm_head: a linear layer from hidden_dim to vocab_size).
  3. Memory accounting: How to estimate tensor sizes from shapes and data types (e.g., [2048, 248320] in bf16 = 2048 × 248320 × 2 bytes = ~970 MB), and how optimizer states double parameter memory.
  4. Speculative decoding and DFlash: The concept of drafters predicting multiple tokens, the DDTree structure for tree-based speculation, the block/anchor arrangement for training positions, and the distillation loss (CE + KL blend).
  5. The specific training configuration: block_size=32, max_anchors=1024, 5 drafter layers, vocab_size=248320, hidden_dim=5120, intermediate_dim=17408, chunk_size=2048, 16 chunks per step.
  6. CUDA memory management: How GPU memory is allocated and freed, the concept of memory fragmentation, and tools like PYTORCH_CUDA_ALLOC_CONF.

Output Knowledge Created

This message creates several important pieces of knowledge:

  1. The backward graph accumulation problem: A documented understanding that chunking the forward pass does not reduce backward graph memory, because PyTorch retains all chunk outputs until backward. This is a general principle applicable to any chunked training pipeline.
  2. The gradient checkpointing solution: A specific, implementable approach: wrap the lm_head → loss computation per chunk in torch.utils.checkpoint.checkpoint, saving only the normalized embeddings (~0.06 GB per chunk) and recomputing logits during backward.
  3. Memory budget breakdown: A detailed accounting of where memory goes in a DFlash training step: model weights (20.76 GB), attention activations (4.1 GB), MLP activations (8.5 GB), context K/V projections (substantial but unquantified), and backward graph logits (~31 GB without checkpointing).
  4. The 430 MB gap: A precise measurement that the pipeline is only 430 MB short of fitting, which informs the design of the checkpointing solution (any approach that saves more than 430 MB is sufficient).
  5. Multiple rejected alternatives: Documentation of approaches that were considered but rejected: reducing chunk_size to 256 (too slow), skipping KL entirely (loses distillation benefit), computing KL on top-K logits (complex, may lose signal), reducing block_size/max_anchors (reduces training signal), and freeing/recomputing targets tensor (insufficient savings).

The Thinking Process: A Window into Debugging Under Pressure

What makes this message exceptional is not the final answer (gradient checkpointing is a well-known technique), but the process by which the assistant arrives at it. Let us trace the arc of the reasoning:

Phase 1: Surface Analysis (lines 1-4) The assistant identifies the immediate problem: KL divergence at chunk_size 2048 needs 1.0 GB for its output, but only 543 MB is free. The initial framing is about the KL operation itself being too large.

Phase 2: Memory Budget Breakdown (lines 5-25) The assistant attempts to account for all 93.4 GB of allocated memory. This is a systematic audit that walks through model weights, optimizer states, attention activations, MLP activations, and flex_attention workspace. The initial calculation yields ~37.7 GB, which is only 40% of the budget—prompting the realization that context K/V projections were omitted.

Phase 3: The Chunk Size Trade-off (lines 26-40) The assistant explores reducing chunk_size from 2048 to 1024 or 512, calculating the resulting tensor sizes. It discovers that even at chunk_size 1024, the KL computation would still exceed the available 543 MB because five copies of the chunk tensor are needed simultaneously during KL computation. It calculates that chunk_size ~100 would be required to fit, which would require 300+ iterations per step—prohibitively slow.

Phase 4: Exploring Alternatives (lines 41-50) The assistant enumerates four options: reduce chunk_size for KL only, skip KL entirely, compute KL on top-K logits, or reduce block_size/max_anchors. Each is evaluated and found suboptimal.

Phase 5: The CE-only Fallacy (lines 51-60) The assistant considers using pure CE (which doesn't need full softmax), but realizes that even CE requires the logits tensor (970 MB at chunk_size 2048), which alone exceeds the 543 MB free. This eliminates the "just skip KL" option.

Phase 6: The Backward Graph Revelation (lines 61-85) This is the critical turning point. The assistant traces through what PyTorch actually saves for backward: logits_c, s_log (log-softmax of student logits), t_prob (teacher softmax), and kl_c output. It realizes that even though t_prob can be freed (no gradients), logits_c and s_log must be retained. With 16 chunks, this creates ~31 GB of backward graph memory.

Phase 7: The Checkpointing Solution (lines 86-120) The assistant designs the gradient checkpointing approach: wrap the lm_head → loss computation in torch.utils.checkpoint.checkpoint, saving only the normalized embeddings (0.06 GB per chunk) and recomputing logits during backward. It works through the interaction with torch.no_grad(), the handling of non-tensor outputs (pred_ids, tgt_ids for metrics), and the specific implementation details.

Phase 8: Implementation Decision (lines 121-125) The assistant concludes with a clear action plan: "The fix: gradient checkpoint each chunk so logits are recomputed during backward instead of saved. Let me implement this properly." It then kills the failed training session to prepare for the fix.

The Broader Significance

This message sits at a critical juncture in the DFlash project. The preceding messages had just implemented the fused chunked loss function, which was supposed to solve the memory problem. The discovery that it didn't—and the subsequent debugging journey—represents a significant learning moment.

The assistant's debugging approach exemplifies several best practices:

  1. First-principles memory accounting: Rather than guessing or using trial-and-error, the assistant calculates tensor sizes from shapes and data types, building a complete memory budget from the ground up.
  2. Understanding the tool: The assistant doesn't just use PyTorch—it understands how autograd works internally, how the backward graph accumulates tensors, and how gradient checkpointing interacts with torch.no_grad().
  3. Systematic elimination: Each hypothesis is tested against the memory budget. Options are evaluated quantitatively (e.g., "chunk_size ~100 would require 300+ iterations") rather than qualitatively.
  4. Knowing when to pivot: After exploring multiple alternatives, the assistant correctly identifies gradient checkpointing as the right solution and commits to implementing it properly.

Conclusion

Message <msg id=9305> is a remarkable document of real-time debugging under extreme memory pressure. It captures the moment when a promising approach (chunked loss computation) hits a fundamental limitation (backward graph accumulation), and the reasoning process that leads to the correct solution (gradient checkpointing).

The message teaches us several lessons: