5 PyTorch Memory Tactics for Bigger, Faster Models

Practical moves — KV cache reuse, gradient checkpointing, BF16, selective offload, and FlashAttention — that squeeze more sequence length and batch size out of the same GPU.

Five PyTorch memory tactics — KV cache, gradient checkpointing, BF16 mixed precision, offloading, and FlashAttention — to fit longer contexts and train faster.

Let’s be real: most of us don’t have unlimited H100s lying around. We have one or two GPUs and an ambition that keeps outgrowing VRAM. The good news? With a few disciplined choices, you can pull real headroom out of the same hardware. Below are five the-trenches tactics that consistently cut memory while keeping throughput high and code sane.

1) Pre-allocate and Reuse the KV Cache

When: Transformer inference or long-sequence training with cached attention.

Why it helps: Keys/values for each layer are the big-ticket item at long context. Reusing a pre-allocated KV cache avoids churn in the allocator and lets you expand sequence length without fragmentation.

How:

• Pre-allocate [layers, batch, heads, tokens, head_dim] for k and v.
• Pass a kv_cache handle through your attention module.
• For generation, append tokens in-place instead of re-materializing.

import torch

def allocate_kv_cache(n_layers, bsz, n_heads, max_t, head_dim, dtype=torch.bfloat16, device="cuda"):
    shape = (n_layers, bsz, n_heads, max_t, head_dim)
    k_cache = torch.empty(shape, dtype=dtype, device=device)
    v_cache = torch.empty(shape, dtype=dtype, device=device)
    return k_cache, v_cache

@torch.no_grad()
def append_kv(layer_id, t_idx, K, V, k_cache, v_cache):
    # K,V: [bsz, heads, 1, head_dim] for the new token
    k_cache[layer_id, :, :, t_idx:t_idx+1, :] = K
    v_cache[layer_id, :, :, t_idx:t_idx+1, :] = V

Pro tip: Keep the cache in BF16 (or FP16 if safe) even if model weights use a higher precision. Add a small ring-buffer if you do sliding-window attention.

2) Gradient Checkpointing Where It Pays

When: Training with long sequences or deep stacks where activations dominate memory.

Why it helps: Checkpointing trades compute for memory by dropping intermediate activations and recomputing them during backward. The win is often 30–50% activation savings for a ~10–20% compute penalty, which is a great trade if you’re memory-bound.

from torch.utils.checkpoint import checkpoint

class Block(torch.nn.Module):
    def __init__(self, attn, mlp): super().__init__(); self.attn, self.mlp = attn, mlp
    def forward(self, x, attn_mask=None):
        def fn(inp): 
            y = self.attn(inp, attn_mask) + inp
            return self.mlp(y) + y
        return checkpoint(fn, x, use_reentrant=False)

Placement: Don’t wrap everything. Checkpoint the memory-heavy blocks (attention + MLP) and leave light layers (layernorms, small projections) as-is to avoid pointless recompute.

Combine with: torch.compile(mode="reduce-overhead") can claw back some of the recompute cost with fused kernels.

3) BF16 Mixed Precision, Not Blind FP16

When: Training or inference on Ampere or newer (BF16-native) GPUs.

Why it helps: BF16 gives FP32-range with 16-bit storage — fewer NaNs, fewer losses exploding, and almost the same memory savings as FP16. You keep stability and gain headroom.

import torch
from torch.cuda.amp import autocast, GradScaler

model = model.to(torch.bfloat16).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler(enabled=False)  # BF16 typically doesn’t need scaling

for batch in loader:
    inputs, targets = (t.cuda() for t in batch)
    with autocast(dtype=torch.bfloat16):
        logits = model(inputs)
        loss = torch.nn.functional.cross_entropy(logits, targets)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

Where to keep FP32: optimizer state (Adam moments) and critical reductions (softmax max-subtraction, layernorm variance) if your stack doesn’t already handle them stably.

Sanity check: If you’re stuck on Turing/V100 without BF16, try FP16 with dynamic loss scaling and keep layernorms in FP32.

4) Selective Offload: Optimizer & Activations to CPU (or NVMe)

When: You need “just a bit more” than VRAM allows, but can tolerate modest latency.

Why it helps: Moving rarely-used tensors off-device frees VRAM for activations and KV, which gates batch and context size. Done carefully, the throughput hit is small relative to the capacity unlocked.

Two practical flavors:

4.1 ZeRO-style optimizer offload (Adam moments)

Offload optimizer states (and optionally gradients) to CPU RAM.

# Conceptual: move optimizer states to CPU after each step
for group in optimizer.param_groups:
    for p in group['params']:
        state = optimizer.state[p]
        for k, v in state.items():
            if torch.is_tensor(v):
                state[k] = v.cpu()  # keep params on GPU, moments on CPU

Use a library (DeepSpeed ZeRO Offload, bitsandbytes) in production; they batch transfers to avoid PCIe thrash.

4.2 Activation/attention offload for outlier windows

For occasional long-context batches, spill a subset of activations or KV pages to host and fetch only for backward. A simple pattern is a “long batch lane” with fewer sequences that gets special treatment.

Rule of thumb: Prefer offloading states over constantly used weights. And batch your DMA — many small copies kill you.

5) Use FlashAttention (or Equivalent Fused Attention)

When: Any modern transformer with long context windows.

Why it helps: FlashAttention computes attention in tiles, keeping Q/K/V blocks in SRAM and avoiding materializing the full n x n attention matrix. That slashes both memory and time, especially beyond a few thousand tokens.

# Example: xFormers' memory-efficient attention (drop-in)
import torch
import xformers.ops as xops

def attn(q, k, v, attn_bias=None):
    # q, k, v: [bsz, seq, heads, head_dim], bfloat16
    return xops.memory_efficient_attention(q, k, v, attn_bias=attn_bias, p=0.0)

# Or FlashAttention-2 attention if using its kernels

Bonus synergy: FlashAttention + BF16 + checkpointing often unlocks another 2–4× context length on mid-range GPUs, without unstable gradients.

Minimal “Memory-First” Training Skeleton

import torch, math
from torch.utils.checkpoint import checkpoint
from torch.cuda.amp import autocast

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

class TinyTransformer(torch.nn.Module):
    def __init__(self, L, H, D, FF, heads):
        super().__init__()
        self.emb = torch.nn.Embedding(H, D)
        self.blocks = torch.nn.ModuleList([Block(D, FF, heads) for _ in range(L)])
        self.ln = torch.nn.LayerNorm(D)
        self.head = torch.nn.Linear(D, H, bias=False)
    def forward(self, x):
        x = self.emb(x)
        for b in self.blocks:
            x = checkpoint(b, x, use_reentrant=False)  # tactic #2
        return self.head(self.ln(x))

model = TinyTransformer(L=24, H=32000, D=1024, FF=4096, heads=16).cuda().to(torch.bfloat16)  # tactic #3
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

for step, batch in enumerate(loader):
    x, y = (t.cuda(non_blocking=True) for t in batch)
    with autocast(dtype=torch.bfloat16):
        logits = model(x)
        loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    opt.step()

Add FlashAttention inside Block and consider offloading optimizer states if VRAM is tight.

An ASCII Mental Model (what we just optimized)

VRAM Budget
├─ Weights (BF16)           ↓  memory
├─ Optimizer states (CPU)   →  offload
├─ Activations (chkpt)      ↓↓ recompute not store
├─ Attention (FlashAttn)    ↓ avoid n×n
└─ KV Cache (prealloc/reuse) ↔ stable, fragmentation-free

• BF16 shrinks almost everything without catastrophes.
• Checkpointing shrinks activations; your GPU does a small redo during backward.
• FlashAttention shrinks attention intermediates and speeds them up.
• Offload moves cold state away from hot VRAM.
• KV cache reuse keeps long-context runs from fragmenting memory.

Quick, Practical Benchmarks to Run

• Max batch @ fixed seq length: switch BF16 on/off; measure peak batch, step time, and loss stability.
• Max seq length @ fixed batch: toggle checkpointing + FlashAttention; record OOM boundary.
• Throughput with/without offload: ensure PCIe isn’t your new bottleneck; check p95 step time.
• Fragmentation test: run 100 inference requests of variable lengths with pre-allocated KV vs. ad-hoc allocations; watch “Active/Reserved” in torch.cuda.memory_summary().

Numbers vary by model and GPU, but the shape of wins is consistent: more context and batch on the same silicon, plus calmer step-time variance.

Wrap-Up

You don’t need new GPUs to train longer or serve faster — you need less surprise work in memory. Pre-allocate the KV cache. Checkpoint the heavy blocks. Default to BF16. Offload what’s cold. And let FlashAttention avoid the n×n cliff. Do those five, and VRAM stops being a hard wall and becomes a set of dials you can tune.