4.10 Memory-Efficient Training: Checkpointing, Offloading & LoRA Math¶

Training large language models is fundamentally a memory management problem. A 7-billion-parameter model in full bf16 precision occupies about 14 GB just for its weights — but the actual GPU memory consumed during a training step can easily be 10× that figure once you account for gradients, optimizer states, and the intermediate tensors produced by the forward pass. Understanding where that memory goes, and which techniques claw it back, is essential knowledge for anyone who trains or fine-tunes LLMs.

This chapter builds a precise memory budget for a training step from first principles, then covers the three main families of solutions: activation/gradient checkpointing, CPU and NVMe offloading, and parameter-efficient fine-tuning (PEFT). We quantify the tradeoffs with real numbers, show working code, and answer the interview questions that consistently trip up strong candidates.

For background on GPU memory hierarchies and how data moves between HBM, L2, and SRAM, see GPU Architecture & The Memory Hierarchy. For the distributed-training memory picture (ZeRO, FSDP), see Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP. The mixed-precision regime in which all of this happens is covered in Mixed Precision, bf16 & FP8 Training.

The Full Training Memory Budget¶

Before optimizing anything, we need a precise accounting. At any training step, GPU HBM holds six categories of tensors:

Let \(P\) be the number of parameters. In standard Adam mixed-precision training:

So for a 7B model: \(16 \times 7 \times 10^9 = 112\,\text{GB}\) — already beyond a single 80 GB H100. And we haven’t counted activations yet.

The Activation Memory Equation¶

For a transformer trained with a batch of \(B\) sequences of length \(T\), with \(L\) layers, hidden dimension \(d\), and \(h\) attention heads:

During the forward pass, each transformer block needs to store activations for the backward pass. The dominant contributors per layer, per token are:

• Attention QKV projections: \(3 \times d\) (one fp16 tensor per projection)
• Attention scores and softmax output: \(B \times h \times T \times T\) (the full \(T \times T\) attention matrix per head)
• Post-attention projections: \(B \times T \times d\)
• MLP intermediate: \(B \times T \times 4d\) (for a standard 4× MLP expansion)
• Layer norms: \(B \times T \times d\)

A common rule of thumb aggregates this to roughly:

so in fp16/bf16 (2 bytes per element) this is \(M_{\text{act}} \approx 24\,B\,T\,d\,L\) bytes. The coefficient 12 counts stored elements per token per layer; it is a rounded-down version of Megatron-LM’s exact per-layer accounting \(34\,s\,b\,h + 5\,a\,s^2\,b\) bytes (with \(s=T\), \(b=B\), \(h=d\), \(a\) = number of heads), whose first term is \(\approx 17\) elements/token/layer. The second term \(5\,a\,s^2\,b\) is the \(T\times T\) attention-score matrix; it is absent when FlashAttention is used (it never materializes the scores), leaving only the term linear in \(T\). The exact coefficient depends on architecture details (e.g., GQA reduces the attention term). The key observation is that activation memory scales as \(B \times T \times L\) — it can dwarf the static weight cost for long sequences or large batches.

The 7B numbers make optimizer state look like the whole story. Run the same accounting at the scale you can actually afford to pretrain — a ~100M-parameter model — and the ranking inverts: \(16P = 1.6\) GB of static state disappears into a corner of any GPU, while a micro-batch of 32 sequences at \(T=2048\) puts tens of GB into activations. Two consequences follow, and both are why the capstone run is engineered the way it is. First, at 100M the memory levers that matter are micro-batch size, activation checkpointing, and the loss head — the \(B \times T \times V\) logits tensor, which lives outside the transformer blocks and is therefore untouched by block-level checkpointing, and which for \(B\,T = 65{,}536\) and \(V = 32{,}768\) is over a gigabyte in fp32 before you have counted a single block. Second, PEFT is the wrong tool here: you are training from scratch, so there is no pretrained base to freeze. The Pretraining Run: A Complete Single-GPU Training Loop does this budget line by line for Stack-100M, including the chunked cross-entropy head that shrinks the logits term ~14–30×.

Activation Checkpointing: Recompute vs. Store¶

Activation checkpointing (also called gradient checkpointing) is the oldest and most universally applicable memory-reduction technique. The idea, introduced in the systems literature as “rematerialization” and popularized for deep learning by Chen et al. (2016) in Training Deep Nets with Sublinear Memory Cost, is simple:

Do not store every intermediate tensor during the forward pass. Instead, recompute them on demand during the backward pass.

The Recompute–Storage Tradeoff¶

Without checkpointing, storing all activations for an \(L\)-layer network costs \(O(L)\) memory but zero extra compute. With full recomputation, you store only the input to each layer, paying one extra forward pass: memory drops to \(O(1)\) (or \(O(\sqrt{L})\) with optimal placement), compute goes up by roughly 33%.

The memory–compute tradeoff is:

where \(k\) is the number of “checkpoints” (layer boundaries where you store a tensor). Choosing \(k = \sqrt{L}\) minimizes the product, giving \(O(\sqrt{L})\) memory and \(O(\sqrt{L})\) extra cost — the classic sublinear memory result.

In practice, modern frameworks let you checkpoint at the granularity of an entire transformer block, so the +33% compute overhead estimate is approximately correct for full-checkpointing of all blocks.

PyTorch Activation Checkpointing in Practice¶

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint, checkpoint_sequential

# -----------------------------------------------------------------------
# Simple demonstration: a single transformer block with checkpointing.
# We wrap the forward in torch.utils.checkpoint.checkpoint so PyTorch
# will NOT save intermediate activations, and will recompute them during
# backward instead.
# -----------------------------------------------------------------------

class TransformerBlock(nn.Module):
    """A minimal causal transformer block (MHA + FFN + layer norms)."""

    def __init__(self, d_model: int, n_heads: int, ffn_mult: int = 4):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_mult * d_model),
            nn.GELU(),
            nn.Linear(ffn_mult * d_model, d_model),
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Self-attention with residual
        normed = self.ln1(x)
        attn_out, _ = self.attn(normed, normed, normed, need_weights=False)
        x = x + attn_out
        # FFN with residual
        x = x + self.ffn(self.ln2(x))
        return x


class CheckpointedModel(nn.Module):
    """Wraps a stack of transformer blocks and applies gradient checkpointing."""

    def __init__(self, n_layers: int, d_model: int, n_heads: int,
                 use_checkpointing: bool = True):
        super().__init__()
        self.blocks = nn.ModuleList(
            [TransformerBlock(d_model, n_heads) for _ in range(n_layers)]
        )
        self.use_checkpointing = use_checkpointing

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            if self.use_checkpointing and x.requires_grad:
                # checkpoint() replaces the saved activations with a
                # recomputation graph.  use_reentrant=False is recommended
                # in PyTorch >= 2.0 for compatibility with compiled graphs.
                x = checkpoint(block, x, use_reentrant=False)
            else:
                x = block(x)
        return x


# -----------------------------------------------------------------------
# Memory comparison: illustrate the difference at inference scale.
# -----------------------------------------------------------------------
def measure_peak_memory(n_layers: int, use_checkpointing: bool,
                         batch: int = 2, seq_len: int = 512,
                         d_model: int = 1024, n_heads: int = 8) -> float:
    """Returns peak HBM usage in MB after one forward+backward pass."""
    torch.cuda.reset_peak_memory_stats()
    model = CheckpointedModel(n_layers, d_model, n_heads, use_checkpointing).cuda()
    model = model.to(torch.bfloat16)
    x = torch.randn(batch, seq_len, d_model, device="cuda",
                     dtype=torch.bfloat16, requires_grad=True)
    loss = model(x).mean()
    loss.backward()
    return torch.cuda.max_memory_allocated() / 1e6  # MB


if __name__ == "__main__":
    for ckpt in [False, True]:
        mb = measure_peak_memory(n_layers=24, use_checkpointing=ckpt)
        print(f"Checkpointing={ckpt}: peak memory = {mb:.1f} MB")
    # Typical output:
    # Checkpointing=False: peak memory = 3421.3 MB
    # Checkpointing=True:  peak memory = 1108.7 MB  (~3x reduction)

Selective Recomputation¶

Full checkpointing of all blocks is conservative. Modern implementations (e.g., Megatron-LM’s --recompute-granularity selective) let you choose which operations to recompute. The cost-benefit analysis:

• Attention QK softmax (the \(T \times T\) matrix): large memory footprint, cheap to recompute since it is memory-bound not compute-bound.
• MLP GELU activations: moderate size, fast recompute.
• Layer norm outputs: small, fast recompute — often not worth the overhead.

The rule of thumb: recompute anything whose memory cost exceeds its FLOPs cost. FlashAttention (see FlashAttention I: IO-Awareness & The Online Softmax) is itself an extreme form of selective recomputation — it does not materialize the \(N \times N\) attention matrix at all, recomputing softmax statistics on-the-fly, saving \(O(N^2)\) memory per layer.

# Selective checkpointing: only checkpoint the attention sub-block,
# not the cheaper FFN.  Saves ~60% of attention-related activation memory
# at a small recompute cost.

from torch.utils.checkpoint import checkpoint

class SelectiveCheckpointBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ffn  = nn.Sequential(
            nn.Linear(d_model, 4 * d_model), nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)

    def _attn_only(self, x: torch.Tensor) -> torch.Tensor:
        """The sub-computation we want to recompute in backward."""
        normed = self.ln1(x)
        out, _ = self.attn(normed, normed, normed, need_weights=False)
        return x + out

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Recompute attention activations; store FFN activations normally.
        if x.requires_grad:
            x = checkpoint(self._attn_only, x, use_reentrant=False)
        else:
            x = self._attn_only(x)
        return x + self.ffn(self.ln2(x))

The Real Library APIs: SAC, apply_activation_checkpointing, and One-Line Trainer Flags¶

Hand-rolling the split above is instructive but coarse — it forces you to carve the block into functions. Since PyTorch 2.4, torch.utils.checkpoint exposes selective activation checkpointing (SAC) at operator granularity through a policy function, which is what torchtitan uses as its default recompute mode. The policy sees each ATen op as it is traced and decides MUST_SAVE (keep the output — use for expensive matmuls and SDPA) or PREFER_RECOMPUTE (throw it away — use for cheap elementwise/normalization ops that produce big tensors):

import torch
from torch.utils.checkpoint import (
    checkpoint, create_selective_checkpoint_contexts, CheckpointPolicy,
)

# Ops whose outputs are expensive to recompute -> save them.
# Everything else (GELU/SiLU, mul, add, LayerNorm internals) is recomputed:
# those are memory-bound elementwise ops, so recompute is nearly free.
_SAVE = {
    torch.ops.aten.mm.default,
    torch.ops.aten._scaled_dot_product_flash_attention.default,
}

def _policy(ctx, op, *args, **kwargs):
    return (CheckpointPolicy.MUST_SAVE if op in _SAVE
            else CheckpointPolicy.PREFER_RECOMPUTE)

def sac_forward(block, x):
    """Run `block` with per-op selective recompute instead of all-or-nothing."""
    return checkpoint(
        block, x,
        use_reentrant=False,
        context_fn=lambda: create_selective_checkpoint_contexts(_policy),
    )

This typically recovers most of full checkpointing’s memory saving for a fraction of its compute overhead, because the SwiGLU/GELU intermediates — the single largest activation term in a modern block — are exactly the cheap-to-recompute ones, while the matmul outputs you would otherwise pay 33% to redo are kept.

Two more library entry points you should reach for before writing your own wrapper:

# 1) FSDP-composable block wrapping (torch.distributed). Preferred over a
#    hand-rolled nn.Module wrapper because it composes with sharding.
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing, checkpoint_wrapper, CheckpointImpl,
)
apply_activation_checkpointing(
    model,
    checkpoint_wrapper_fn=lambda m: checkpoint_wrapper(
        m, checkpoint_impl=CheckpointImpl.NO_REENTRANT),
    check_fn=lambda m: isinstance(m, TransformerBlock),   # which modules to wrap
)

# 2) HuggingFace Transformers / TRL: one flag, on any supported architecture.
model.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs={"use_reentrant": False})
# ...or, from TrainingArguments / SFTConfig:
#   gradient_checkpointing=True,
#   gradient_checkpointing_kwargs={"use_reentrant": False},
#   gradient_accumulation_steps=16,

Megatron-LM exposes the same idea as launcher flags — --recompute-granularity selective (the Korthikanti et al. attention-only policy) or --recompute-granularity full --recompute-method uniform|block --recompute-num-layers N for the coarse version; DeepSpeed as "activation_checkpointing" in its JSON config. See Megatron-LM, DeepSpeed & Parallelism in Practice.

CPU and NVMe Offloading¶

When GPU memory is exhausted even after checkpointing, the next option is to spill tensors to cheaper, larger memory tiers.

The Memory Hierarchy Revisited¶

PCIe bandwidth is ~50× slower than HBM bandwidth. This means CPU offloading is only viable if the tensor being offloaded is not needed every step, or the compute on GPU is long enough to hide the transfer latency.

DeepSpeed ZeRO-Infinity and Offload¶

DeepSpeed’s ZeRO-Offload and ZeRO-Infinity implement optimizer-state and parameter offloading. The key insight is that Adam optimizer states (\(m_t\), \(v_t\), and the fp32 master weights) are updated only once per step, after the gradient has been reduced. They are read once (to compute the weight update) and written once (with the new values). This access pattern tolerates the latency of a PCIe transfer because the Adam update itself is memory-bound and can be executed on CPU cheaply.

# deepspeed_offload_config.json — enable ZeRO-3 with CPU offload
# Drop this into your DeepSpeed config to offload optimizer states + params.

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto"
  },
  "bf16": { "enabled": true },
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 16
}

PyTorch-native FSDP offers the same lever without DeepSpeed: CPUOffload(offload_params=True) in FSDP1, and fully_shard(module, offload_policy=CPUOffloadPolicy()) in FSDP2, which keeps the sharded parameters, gradients and optimizer state in host DRAM and streams each shard to the GPU as its all-gather comes due. ZeRO-Infinity extends the same idea one tier further, to NVMe, where the relevant budget is SSD read bandwidth (a few GB/s per drive) rather than PCIe.

With ZeRO-3 + CPU offload, the GPU holds only a working subset of parameters and optimizer states at any time, allowing effective training of 30B+ models on a single GPU — at the cost of significantly reduced throughput (roughly 5–10× slowdown versus in-memory training due to PCIe bandwidth saturation). Newer offload engines attack this stall directly: DeepSpeed’s ZenFlow (2025) updates only the highest-importance gradients synchronously on GPU and defers the rest to asynchronous, double-buffered CPU accumulation, overlapping PCIe transfer with GPU compute; DeepSpeed reports up to 5× end-to-end speedup over ZeRO-Offload with over 85% fewer GPU stalls on A100/H100 nodes.

Gradient Accumulation as an Offloading Strategy¶

Gradient accumulation is not traditionally called “offloading,” but it achieves the same goal of separating the memory cost of a large effective batch from the peak memory of a single forward–backward pass. With accumulation steps \(A\):

The gradient buffer costs \(2P\) bytes (fp16), held across all accumulation steps. But the activation peak at any step is determined by \(B/A\), not \(B\). This is cheap (no data movement), at the cost of \(A\) forward–backward passes per parameter update.

# Gradient accumulation — manual implementation in pure PyTorch.
# Using a micro-batch of 1 to simulate an effective batch of 8.

model.train()
optimizer.zero_grad()

ACCUMULATION_STEPS = 8
for step, (x, y) in enumerate(dataloader):
    x, y = x.cuda(), y.cuda()

    # Optionally use autocast for mixed precision.
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(x)
        # Scale loss by accumulation factor so gradients are averaged, not summed.
        loss = criterion(logits, y) / ACCUMULATION_STEPS

    loss.backward()  # Accumulates .grad on parameters — no optimizer step yet.

    if (step + 1) % ACCUMULATION_STEPS == 0:
        # Gradient clipping before the optimizer step.
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()
        scheduler.step()

Two corrections this minimal loop deliberately leaves for you to add.

Under DDP, wrap all but the last micro-batch in model.no_sync(). Otherwise DDP’s autograd hooks all-reduce the full gradient after every micro-batch, so an 8-step accumulation does 8× the necessary communication — the classic “why is accumulation not free?” bug. The pattern is in Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP; HuggingFace accelerate wraps it portably as with accelerator.accumulate(model):.

Dividing the loss by ACCUMULATION_STEPS is only correct when every micro-batch contributes the same number of loss terms. With padded or variable-length sequences it is not: a mean-reduced cross-entropy already divides by that micro-batch’s own token count, so short micro-batches get up-weighted and the accumulated gradient is a mean-of-means rather than the true token mean. (This is a real bug that shipped in several trainers and was fixed in HuggingFace Transformers/TRL in late 2024.) The fix is to normalize by the total unmasked tokens in the window:

import itertools
import torch.nn.functional as F

IGNORE_INDEX = -100  # label id excluded from the loss (padding / prompt tokens)

# Correct normalization under variable token counts: sum the per-token losses
# in each micro-batch and divide once by the window's total token count.
window = list(itertools.islice(loader, ACCUMULATION_STEPS))
total_tokens = sum(int((y != IGNORE_INDEX).sum()) for _, y in window)

for x, y in window:
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(x.cuda())
        loss = F.cross_entropy(                       # reduction="sum", not "mean"
            logits.flatten(0, 1).float(), y.cuda().flatten(),
            ignore_index=IGNORE_INDEX, reduction="sum",
        ) / total_tokens
    loss.backward()

The Math Behind LoRA and Why PEFT Slashes Memory¶

Parameter-efficient fine-tuning (PEFT) methods attack the memory problem from a different angle: instead of reducing the cost of training all parameters, they freeze most parameters and only train a small adapter. PEFT I: LoRA, QLoRA, DoRA & The Adapter Family covers the practical usage in depth; here we focus on the memory mathematics. Note that PEFT is a fine-tuning tool with a base-model prerequisite: it pays off when the frozen base is large relative to your GPU, or when you want to serve many task adapters over one shared base. Below a few hundred million parameters, full-parameter fine-tuning is usually both cheaper to reason about and better — which is exactly the call Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M makes for Stack-100M.

The LoRA Factorization¶

LoRA (Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, 2021) approximates a weight update \(\Delta W \in \mathbb{R}^{d \times k}\) as a product of two low-rank matrices:

Only \(A\) and \(B\) are trained; the original \(W\) is frozen. The number of trainable parameters per adapted weight matrix is:

The compression ratio is:

For a 7B model with \(d = k = 4096\) and rank \(r = 16\): \(\rho \approx 16/4096 = 0.4\%\). Only 0.4% of each weight matrix’s parameters are trained.

Memory Impact: Exact Accounting¶

With LoRA, the memory budget changes dramatically:

For a frozen weight tensor, PyTorch does not allocate a gradient buffer, so frozen parameters contribute 0 bytes of gradient or optimizer state. The savings are enormous: if LoRA covers all linear layers in a 7B model with rank 16, the optimizer state shrinks from \(\sim\)56 GB (fp32 Adam) to roughly \(56 \times 0.004 = 0.22\) GB.

Note the last row carefully, because it is the single most common misconception about LoRA. Because adapters sit at every depth, the backward pass still traverses the whole network and every layer still saves the input it needs to form \(\partial\mathcal{L}/\partial A\). LoRA cuts the gradient and optimizer lines by ~99% and the activation line by ~0%. That is why the standard single-GPU recipe is QLoRA plus gradient checkpointing, not QLoRA alone: the two techniques attack disjoint line items. See PEFT I: LoRA, QLoRA, DoRA & The Adapter Family for the gradient-flow derivation.

The frozen base model’s weights still occupy \(2P\) bytes, but they require no gradient storage. With 4-bit quantization of the base model (QLoRA), these compress further to \(\frac{P}{2}\) bytes:

For LLaMA-7B with rank 16 covering all four attention projections (q, k, v, o), \(|\theta_{\text{LoRA}}| \approx 16.8\)M: approximately \(3.5 + 0.03 + 0.13 \approx 3.7\) GB — fitting in a 6 GB GPU.

LoRA From Scratch: A Full Implementation¶

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class LoRALinear(nn.Module):
    """
    Drop-in replacement for nn.Linear with a LoRA side path.

    During forward: output = x @ W^T + (x @ A^T) @ B^T * scale
    where W is frozen and only A, B are updated.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 16,
        alpha: float = 32.0,   # LoRA scaling hyper-param; scale = alpha/rank
        dropout: float = 0.05,
        bias: bool = True,
    ):
        super().__init__()
        self.in_features  = in_features
        self.out_features = out_features
        self.rank         = rank
        self.scale        = alpha / rank  # Hu et al. use this to keep LR independent of r

        # Frozen base weight (will be loaded from pretrained model)
        self.weight = nn.Parameter(
            torch.empty(out_features, in_features), requires_grad=False
        )
        self.bias_param = nn.Parameter(
            torch.zeros(out_features), requires_grad=False
        ) if bias else None

        # Trainable LoRA matrices
        # A is initialized from N(0, 1/sqrt(r)) to give unit-variance init.
        # B is initialized to zero so ΔW = 0 at the start of training.
        self.lora_A = nn.Parameter(
            torch.empty(rank, in_features)
        )
        self.lora_B = nn.Parameter(
            torch.zeros(out_features, rank)
        )
        self.lora_dropout = nn.Dropout(dropout)

        # Kaiming init for A (matches standard linear init scale)
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

    @classmethod
    def from_linear(cls, linear: nn.Linear, rank: int = 16,
                    alpha: float = 32.0) -> "LoRALinear":
        """Convert an existing nn.Linear to LoRALinear, preserving its weights."""
        bias = linear.bias is not None
        lora = cls(linear.in_features, linear.out_features,
                   rank=rank, alpha=alpha, bias=bias)
        with torch.no_grad():
            lora.weight.copy_(linear.weight)
            if bias:
                lora.bias_param.copy_(linear.bias)
        return lora

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Standard linear (no grad flows here since weight.requires_grad=False)
        base_out = F.linear(x, self.weight, self.bias_param)

        # LoRA side path: (x @ A^T) @ B^T, scaled
        # shapes: x[..., in_features] -> A[rank, in] -> B[out, rank]
        lora_out = F.linear(
            F.linear(self.lora_dropout(x), self.lora_A),  # [..., rank]
            self.lora_B                                     # [..., out]
        )
        return base_out + self.scale * lora_out

    def merge_weights(self) -> nn.Linear:
        """
        Merge the LoRA update into W for efficient inference.
        Returns a standard nn.Linear with merged weights.
        """
        merged_weight = self.weight + self.scale * (self.lora_B @ self.lora_A)
        linear = nn.Linear(self.in_features, self.out_features,
                           bias=self.bias_param is not None)
        with torch.no_grad():
            linear.weight.copy_(merged_weight)
            if self.bias_param is not None:
                linear.bias.copy_(self.bias_param)
        return linear

    def extra_repr(self) -> str:
        return (f"in={self.in_features}, out={self.out_features}, "
                f"rank={self.rank}, scale={self.scale:.3f}")


# -----------------------------------------------------------------------
# Utility: inject LoRA into all attention Q, V projections of a GPT-style
# model.  This is the most common recipe (K and O are often frozen).
# -----------------------------------------------------------------------

def inject_lora(model: nn.Module, rank: int = 16, alpha: float = 32.0,
                target_modules: tuple = ("q_proj", "v_proj")) -> nn.Module:
    """
    Walk the module tree and replace named Linear sub-modules
    whose name ends with any string in target_modules with LoRALinear.
    Freezes all non-LoRA parameters.
    """
    # First, freeze everything
    for param in model.parameters():
        param.requires_grad_(False)

    # Replace target projections with LoRA versions
    for name, module in list(model.named_modules()):
        for target in target_modules:
            if name.endswith(target) and isinstance(module, nn.Linear):
                # Navigate to parent and set child
                parts = name.split(".")
                parent = model
                for part in parts[:-1]:
                    parent = getattr(parent, part)
                lora_module = LoRALinear.from_linear(module, rank=rank, alpha=alpha)
                setattr(parent, parts[-1], lora_module)
                break  # Found target for this module; move on

    # Report trainable parameter count
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total     = sum(p.numel() for p in model.parameters())
    print(f"LoRA injection complete: {trainable:,} / {total:,} params trainable "
          f"({100 * trainable / total:.3f}%)")
    return model

QLoRA: Quantized Base + LoRA Adapters¶

QLoRA (Dettmers et al., QLoRA: Efficient Finetuning of Quantized LLMs, 2023) combines two ideas:

1. NF4 (NormalFloat4): A 4-bit data type optimized for normally-distributed weights. Instead of linear quantization, NF4 assigns quantization levels at equal-probability points of a standard normal distribution, minimizing quantization error for the typical weight distribution.

2. Double quantization: The NF4 quantization constants themselves are quantized to 8 bits, saving roughly 0.37 bits per parameter on top of the base NF4 savings.

3. Paged optimizer: Uses NVIDIA’s unified memory to page optimizer states to CPU DRAM on demand, preventing OOM crashes from memory spikes. In bitsandbytes this is a drop-in optimizer class — bnb.optim.PagedAdamW8bit(...) (or PagedAdamW32bit) — which is what TRL’s optim="paged_adamw_8bit" selects.

Be precise about what “frozen 4-bit base” means for autograd. Each base weight is dequantized to the bf16 compute dtype on the fly for its matmul, and the gradient of the loss with respect to the layer input does propagate back through that dequantized weight — otherwise no adapter below the top layer could learn. What is never formed is \(\partial\mathcal{L}/\partial W_0\): the base weight carries requires_grad=False, so no weight gradient and no optimizer state is ever allocated for it. There is no straight-through estimator involved, because nothing is being quantized in the forward path of a trainable parameter — the adapters \(A\) and \(B\) stay in bf16 end to end.

# QLoRA with bitsandbytes and HuggingFace Transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType
import torch

# 4-bit quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",         # NormalFloat4 — best for LLM weights
    bnb_4bit_use_double_quant=True,    # Quantize the quantization constants too
    bnb_4bit_compute_dtype=torch.bfloat16,  # Activations/LoRA in bf16
)

# Load the model in 4-bit; the base weights are frozen automatically
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto",
)

# Prepare for k-bit training (adds gradient checkpointing + layer-norm fixes)
from peft import prepare_model_for_kbit_training
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)

# Attach LoRA adapters to Q and V projections
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],   # expand to k_proj, o_proj for more capacity
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Typical output: trainable params: 4,194,304 || all params: 6,738,415,616
# || trainable%: 0.0623

Practical Memory Accounting: A Step-by-Step Recipe¶

When you encounter an OOM error, the correct debugging strategy is systematic, not random. Here is the procedure we use in practice.

Step 1: Measure the Static Footprint¶

import torch

def memory_snapshot(label: str = "") -> None:
    """Print current and peak allocated GPU memory."""
    alloc = torch.cuda.memory_allocated() / 1e9
    peak  = torch.cuda.max_memory_allocated() / 1e9
    res   = torch.cuda.memory_reserved() / 1e9
    print(f"[{label}] allocated={alloc:.2f}GB  peak={peak:.2f}GB  reserved={res:.2f}GB")

# Profile a forward pass step-by-step:
torch.cuda.reset_peak_memory_stats()
memory_snapshot("start")

model = model.cuda()
memory_snapshot("after model load")  # Should be ~2P bytes for bf16 weights

x = x.cuda()
memory_snapshot("after input")

with torch.no_grad():
    out = model(x)
memory_snapshot("after forward (no_grad)")  # No activation retention

del out
torch.cuda.empty_cache()

out = model(x)                   # With grad enabled
memory_snapshot("after forward (with_grad)")  # Activations retained here!

out.sum().backward()
memory_snapshot("after backward")  # Grads allocated, activations freed

Step 2: Use torch.cuda.memory_stats for Detailed Breakdown¶

# Dump the full memory stats dictionary to identify the largest allocation.
stats = torch.cuda.memory_stats()
for key, val in sorted(stats.items(), key=lambda kv: -kv[1]):
    if val > 0:
        print(f"  {key}: {val / 1e6:.1f} MB")

Two numbers in that dict deserve special attention. allocated_bytes.all.peak is what your tensors actually need; reserved_bytes.all.peak is what the caching allocator holds from the driver. A large gap between them is fragmentation: the allocator has the bytes but not in a contiguous block of the right size, and you OOM with “GB free” in the error message. The single highest-value fix is one environment variable, which lets the allocator grow a segment in place instead of pooling fixed-size blocks:

# Cuts fragmentation-driven OOMs, especially with variable sequence lengths
# or gradient checkpointing (whose alloc/free pattern is very bursty).
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Step 2b: Record a Memory Snapshot and Look at It¶

Aggregate counters tell you how much; the snapshot recorder tells you which tensor, allocated by which line of Python. This is the tool to reach for when the budget on paper disagrees with reality:

# Record every allocation with its Python stack, then dump a snapshot.
torch.cuda.memory._record_memory_history(max_entries=100_000)

for step, batch in enumerate(loader):   # a handful of steps is plenty
    train_step(batch)
    if step == 3:
        break

torch.cuda.memory._dump_snapshot("mem_snapshot.pickle")
torch.cuda.memory._record_memory_history(enabled=None)   # stop recording

Drag mem_snapshot.pickle onto https://pytorch.org/memory_viz (a purely client-side page) to get a flame-graph timeline of every allocation. The signature shapes are easy to read once you know them: a staircase that never comes back down is a leak (usually a tensor with grad_fn retained in a Python list — accumulate loss.detach() or loss.item(), never loss); a tall plateau between forward and backward is activation memory (checkpoint it); and a spike that appears only at optimizer.step() is optimizer state (8-bit Adam, offload, or PEFT).

Step 3: The Decision Tree¶

Combining Techniques: The Memory Stack¶

These techniques are not mutually exclusive. In practice, large-scale fine-tuning stacks several simultaneously. The table below shows approximate peak memory for a 70B-parameter Llama-style model with a batch of 1 at sequence length 2048:

Read the LoRA row twice: at 70B, freezing everything is not enough, because the frozen bf16 base alone is 140 GB. Once optimizer state is gone, the resident base weights become the floor, and the only lever that moves that floor is quantizing the base — which is precisely the gap QLoRA was invented to close (Dettmers et al. report 65B fine-tuning on a single 48 GB GPU). Everything below the base-weight floor is then an activation problem, i.e. a checkpointing and micro-batch problem.

Optimizer State Memory Reduction¶

Optimizer state is often the largest single memory consumer in full fine-tuning. Beyond LoRA, several optimizer designs intrinsically reduce state:

Adafactor¶

Adafactor (Shazeer & Stern, 2018) replaces the full second-moment matrix with a factored rank-1 approximation:

This reduces optimizer state from \(O(dk)\) (Adam’s \(v_t\) for a \(d \times k\) weight) to \(O(d + k)\). For a 4096×4096 linear layer, that is 16.7 million → 8,192 values: a 2,048× compression. Adafactor also omits the first moment \(m_t\) (relying on relative step size), further halving the state. The tradeoff is that Adafactor can be less stable for fine-tuning on small datasets; many practitioners use it for pretraining but fall back to Adam for RLHF.

8-bit Adam¶

Dettmers et al. introduced 8-bit Adam in 8-bit Optimizers via Block-wise Quantization (2022). The fp32 \(m_t\) and \(v_t\) are stored in 8-bit integers with per-block scaling factors. Memory for optimizer states drops from \(8P\) to roughly \(2P\) bytes — a 4× reduction — with negligible accuracy loss.

# pip install bitsandbytes
import bitsandbytes as bnb

optimizer = bnb.optim.Adam8bit(
    [p for p in model.parameters() if p.requires_grad],
    lr=2e-4,
    betas=(0.9, 0.999),
    eps=1e-8,
)
# Drop-in replacement for torch.optim.Adam; uses ~4x less optimizer state memory.

GaLore: Low-Rank Optimizer State Without Low-Rank Weights¶

LoRA constrains the weights to a low-rank update; GaLore (Zhao et al., 2024) instead constrains only the optimizer state. For a gradient \(G \in \mathbb{R}^{m \times n}\) it computes a projector \(P \in \mathbb{R}^{m \times r}\) from an SVD of \(G\) (refreshed only every few hundred steps, so the SVD cost amortizes to near zero), runs Adam entirely on the projected gradient \(P^\top G \in \mathbb{R}^{r \times n}\), then projects the resulting update back with \(P\) before applying it. Optimizer state drops from \(O(mn)\) to \(O(rn)\) — the same asymptotic win as LoRA — but because \(P\) is recomputed periodically, the accumulated weight update spans the full space over training rather than being pinned to one rank-\(r\) subspace. That makes it usable for pretraining from scratch, where LoRA is not; the paper reports pretraining a 7B model on a single 24 GB consumer GPU. The cost is that the projection sits on the critical path of every step, and the weights themselves are still full-rank and full-size in memory.

Muon: A Single-Buffer Alternative to Adam¶

A different lever, orthogonal to compressing Adam’s states, is to replace Adam itself. Muon (an orthogonalized-momentum optimizer scaled to production LLM training by Liu et al., Muon is Scalable for LLM Training, 2025) applies Newton–Schulz iteration to orthogonalize the momentum matrix of each 2D weight, and needs only a single momentum buffer rather than Adam’s first and second moments — roughly halving optimizer-state memory for the hidden-weight matrices that dominate transformer parameter counts. It saw rapid adoption in 2025–2026 pretraining and full-fine-tuning runs as a memory- and compute-efficient AdamW alternative; it is complementary to, not a substitute for, the PEFT techniques in this chapter, since LoRA already eliminates almost all optimizer state on frozen layers regardless of which optimizer updates the adapters.

Gradient Accumulation and FP16 Gradients¶

When gradient accumulation is used, the gradient buffer persists across every micro-batch of the window, so its dtype matters. Keeping gradients in bf16/fp16 costs \(2P\) bytes versus \(4P\) for fp32 — a 2× saving on that line item, but one you should take deliberately: accumulating many micro-batch gradients into a 10-bit-mantissa bf16 buffer loses small contributions to rounding, which is why DDP/FSDP expose an explicit reduce_dtype and why long accumulation windows usually keep fp32 gradient accumulation on (MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) in FSDP). With fp16 — not bf16 — you additionally need loss scaling to keep small gradients from flushing to zero; PyTorch’s autocast + GradScaler handles that:

scaler = torch.amp.GradScaler("cuda")  # torch.cuda.amp.GradScaler is deprecated in 2.4+

for micro_batch_x, micro_batch_y in batches:
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = model(micro_batch_x, labels=micro_batch_y).loss / N_ACCUM
    # Backward in fp16; GradScaler prevents underflow
    scaler.scale(loss).backward()

scaler.step(optimizer)   # Unscales grads, then steps
scaler.update()
optimizer.zero_grad()

Key Implementation Patterns and Pitfalls¶

Memory Pinning for CPU Offload¶

When offloading tensors to CPU DRAM, pinned (page-locked) memory dramatically accelerates PCIe transfers by allowing DMA without CPU involvement:

# Allocate CPU tensor in pinned memory for fast GPU<->CPU transfer
cpu_tensor = torch.empty_like(gpu_tensor, device="cpu", pin_memory=True)

# Non-blocking copy from GPU to CPU
cpu_tensor.copy_(gpu_tensor, non_blocking=True)

# ... do GPU compute on other tensors while transfer completes ...

# When needed again, copy back
gpu_tensor.copy_(cpu_tensor, non_blocking=True)
torch.cuda.synchronize()  # Ensure transfer is complete before using gpu_tensor

DeepSpeed’s ZeRO-Offload uses this pattern internally, pre-fetching the next layer’s parameters while the current layer’s backward pass is running.

Gradient Checkpointing with Compiled Models¶

torch.compile and gradient checkpointing interact non-trivially. The recomputed sub-graph is traced separately from the outer graph. As of PyTorch 2.2+, the recommended pattern is:

# Compile the model AFTER enabling checkpointing, not before.
model = CheckpointedModel(n_layers=32, d_model=4096, n_heads=32,
                          use_checkpointing=True)
# torch.compile will trace through checkpoint boundaries correctly
# with use_reentrant=False
model = torch.compile(model, mode="reduce-overhead")

If you compile first and then toggle checkpointing, the compiled graph may not include the recompute branches, and you will silently fall back to full activation storage.

The no_grad vs. detach Distinction¶

A common confusion: torch.no_grad() prevents creation of the autograd graph but does not free existing activation tensors. tensor.detach() severs the graph at a specific point. When implementing custom checkpointing, always use:

# Correct: save only the input, discard intermediate activations
saved_input = input.detach()  # Severs autograd graph; no grad fn stored
# ... run forward normally (intermediates are freed) ...
# In backward: rerun from saved_input (now re-attaches to the graph)

Further Reading¶

• Chen, T., Xu, B., Zhang, C., Guestrin, C. — Training Deep Nets with Sublinear Memory Cost (2016). The original activation checkpointing paper for deep networks.
• Hu, E., et al. — LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022. The foundational PEFT paper.
• Dettmers, T., et al. — QLoRA: Efficient Finetuning of Quantized LLMs, NeurIPS 2023. Combines NF4 quantization, double quantization, and paged optimizers.
• Dettmers, T., et al. — 8-bit Optimizers via Block-wise Quantization, ICLR 2022.
• Rajbhandari, S., et al. — ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (DeepSpeed), SC 2020. Covers ZeRO-½/3 and the theoretical memory analysis.
• Rajbhandari, S., et al. — ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning, SC 2021. Extends ZeRO to NVMe offloading.
• Shazeer, N., Stern, M. — Adafactor: Adaptive Learning Rates with Sublinear Memory Cost, ICML 2018.
• Dao, T., et al. — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022. Eliminates \(O(T^2)\) attention materialization.
• HuggingFace PEFT library (github.com/huggingface/peft) — canonical implementation of LoRA, QLoRA, and adapter variants.

Exercises¶

1. (Static budget.) You want to fully fine-tune a 3B-parameter model with standard Adam mixed-precision training. Using the chapter’s \(16P\) accounting, write out the size (in GB) of each of the four static categories, give the total, and state what fraction of the static footprint is pure optimizer state (\(m_t\) and \(v_t\)). Then say what the total becomes if you swap fp32 Adam for 8-bit Adam.

2. (Activation memory and checkpointing.) Take a LLaMA-style model with \(L = 32\) layers and \(d = 4096\), trained in bf16. Use the chapter’s rule of thumb \(M_{\text{act}} \approx 24\,B\,T\,d\,L\) bytes. (a) Compute the stored-activation memory for a batch of \(B = 2\) at sequence length \(T = 4096\). (b) If you switch to full gradient checkpointing that stores only each block’s input (\(L\,B\,T\,d \times 2\) bytes), how much activation memory remains? © Which of these two numbers does FlashAttention affect, and why?

3. (LoRA parameter budget.) Consider a 13B-style model with hidden size \(d = k = 5120\) and \(L = 40\) layers. You attach rank-\(r = 8\) LoRA adapters to all four attention projections (q, k, v, o), each a \(5120 \times 5120\) matrix. (a) How many trainable parameters does this add in total? (b) What is the per-matrix compression ratio \(\rho = r(d+k)/(dk)\)? © Compare the Adam optimizer-state memory for these adapters (fp32, 8 bytes/param) against the \(8P\) optimizer state of full fine-tuning.

4. (Why the tricks work.) Answer each briefly, grounding your reasoning in the chapter. (a) Optimal activation checkpointing places \(k = \sqrt{L}\) checkpoints. Why does this particular \(k\) minimize the memory-compute product, rather than \(k = 1\) or \(k = L\)? (b) CPU offloading of optimizer states is described as viable, but offloading activations is generally not. What property of the access pattern makes optimizer-state offload tolerate PCIe latency? © Gradient accumulation is called an “offloading strategy” even though no tensor moves off the GPU. In what sense does it offload?

5. (Implementation.) Write a function training_memory_gb(P, r, d, n_adapted, mode) that returns the approximate static training memory in GB for mode in {"full", "lora", "qlora"}, using only the chapter’s formulas. Full training is \(16P\) bytes. LoRA keeps the base in bf16 (\(2P\)) with no base gradients/optimizer state, plus bf16 adapters (\(2|\theta|\)) and fp32 Adam on adapters (\(8|\theta|\)). QLoRA replaces the bf16 base with a 4-bit base (\(P/2\) bytes). Assume each adapted matrix is \(d \times d\) so \(|\theta| = r \cdot 2d \cdot n_{\text{adapted}}\). Verify it reproduces the chapter’s LLaMA-7B QLoRA figure (\(\approx 3.7\) GB) for \(P = 7\times10^9\), \(r = 16\), \(d = 4096\), \(n_{\text{adapted}} = 128\).

def training_memory_gb(P: float, r: int, d: int,
                       n_adapted: int, mode: str) -> float:
    """Approximate static training memory (GB) using the chapter's formulas.

    P          : total base-model parameter count
    r          : LoRA rank
    d          : hidden size (each adapted matrix is d x d)
    n_adapted  : number of adapted matrices (e.g. 4 proj * L layers)
    mode       : "full" | "lora" | "qlora"
    """
    GB = 1e9  # 1 GB = 1e9 bytes (chapter's convention)

    if mode == "full":
        return 16 * P / GB                      # 2P+4P+2P+8P

    # LoRA adapter parameter count: r*(d_in+d_out) per matrix, d_in=d_out=d
    theta = r * 2 * d * n_adapted
    adapters = 2 * theta                        # bf16 adapter weights
    adam     = 8 * theta                        # fp32 Adam m,v on adapters

    if mode == "lora":
        base = 2 * P                            # frozen bf16 base, no grad/opt
    elif mode == "qlora":
        base = P / 2                            # 4-bit NF4 base
    else:
        raise ValueError(mode)

    return (base + adapters + adam) / GB


if __name__ == "__main__":
    P, r, d, n = 7e9, 16, 4096, 128   # LLaMA-7B: 4 attn proj * 32 layers
    for m in ("full", "lora", "qlora"):
        print(f"{m:6s}: {training_memory_gb(P, r, d, n, m):6.2f} GB")