Add FSDP2 activation_checkpointing_offload: offload checkpointed layer inputs to pinned CPU memory

[qgallouedec]

With activation checkpointing, the surviving GPU activation cost is one tensor per layer - the checkpointed layer's input, held in the recompute closure for the whole forward+backward: num_layers × seq_len × hidden bytes. At long sequence lengths this dominates: ~39 GB for an 8B model at 131K tokens/rank.
This PR adds fsdp_activation_checkpointing_offload (FSDP2 only): each checkpointed layer's input is copied to pinned host memory after the layer's forward and its GPU storage freed (untyped_storage().resize_(0)); a shim refills the storage just-in-time when the checkpoint recompute calls back into the layer during backward.

Why not saved_tensors_hooks or reentrant checkpointing

• Non-reentrant checkpointing stores the boundary inputs in a closure: saved_tensors_hooks never sees them, so hook-based offloaders (e.g. torchtune-style) cannot reach the dominant cost.
• The only way to expose them to hooks is reentrant checkpointing, which re-enters the autograd engine per layer and which torch has put on a deprecation path.

This implementation keeps torch's non-reentrant checkpoint untouched (grad-enabled forward), so gradients are exactly those of plain activation checkpointing.

Correctness

One GPU, one released MoE (OLMoE-1B-7B: 7B total, 1B active, 64 experts), one fixed batch, run once per configuration. A MoE is the sensitive case: each layer records its router logits for the load-balancing loss, so those are captured outside the checkpointed region and are the first thing that would break if the wrapper interfered with the graph.

L="accelerate launch --num_processes 1 --use_fsdp --fsdp_version 2 --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP"
$L repro_offload_correctness.py
$L --fsdp_activation_checkpointing true repro_offload_correctness.py
$L --fsdp_activation_checkpointing true --fsdp_activation_checkpointing_offload true repro_offload_correctness.py
$L repro_offload_correctness.py --reentrant

     none  loss 12.206604003906  grad_norm 22.679054379899  router_grad 0.439933836460
       ac  loss 12.206604003906  grad_norm 22.679054202441  router_grad 0.439933687449
   ac+off  loss 12.206604003906  grad_norm 22.679053958828  router_grad 0.439933747053
reentrant  loss 12.206604003906  grad_norm 22.660604998358  router_grad 0.438178092241

Identical loss, and none / ac / ac+off agree on both gradient norms to ~1e-8 relative, which is fp32 recompute rounding. The offload changes nothing about what the backward computes.

The reentrant row is the alternative design, and it is not in that band: its router gradient is 0.40% low, five orders of magnitude beyond the noise between the other three, and it deviates five times more on the router than on the total gradient norm. Small, but not rounding.

repro_offload_correctness.py

import argparse

import torch
from accelerate import Accelerator
from transformers import AutoModelForCausalLM

MODEL, SEQ = "allenai/OLMoE-1B-7B-0924", 1024


def norm(tensor):
    """Gradient norm, unwrapping the DTensor that FSDP2 hands back."""
    grad = tensor.grad
    return (grad.to_local() if hasattr(grad, "to_local") else grad).float().norm().item()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--reentrant", action="store_true", help="use torch's reentrant checkpointing instead")
    args = parser.parse_args()

    accelerator = Accelerator()
    plugin = accelerator.state.fsdp_plugin
    label = "reentrant" if args.reentrant else (
        "ac+off" if plugin.activation_checkpointing_offload else "ac" if plugin.activation_checkpointing else "none"
    )

    torch.manual_seed(0)
    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32)
    model.train()  # transformers only applies its own gradient checkpointing in training mode
    if args.reentrant:
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True})
    # FSDP2 requires the optimizer to be prepared alongside the model; it is never stepped here.
    optimizer = torch.optim.SGD(model.parameters(), lr=0.0)
    model, optimizer = accelerator.prepare(model, optimizer)

    ids = torch.randint(0, model.config.vocab_size, (1, SEQ), generator=torch.Generator().manual_seed(1))
    # output_router_logits puts the load-balancing loss, and with it the router, in the graph
    output = model(input_ids=ids.to(accelerator.device), labels=ids.to(accelerator.device),
                   output_router_logits=True)
    accelerator.backward(output.loss)

    grads = {name: norm(p) for name, p in model.named_parameters() if p.grad is not None}
    total = sum(value**2 for value in grads.values()) ** 0.5
    router = next(value for name, value in grads.items() if "router" in name or name.endswith("gate.weight"))
    print(f"{label:>9}  loss {output.loss.item():.12f}  grad_norm {total:.12f}  router_grad {router:.12f}",
          flush=True)


if __name__ == "__main__":
    main()

Measured

One GPU, the transformer backbone of a released 1.3B MoE (Granite 3.1 1b-a400m), a 32768-token sequence, run with and without the flag. The script prints memory at the end of the forward as well as the peak, plus the size of the checkpoint boundary inputs the wrapper is supposed to move (num_layers × seq × hidden × 2 bytes), so the saving can be checked against the mechanism rather than taken on faith.

      none  peak  79.43 GB  (after forward  79.03)  step 1.08 s  wrapped_layers  0  (boundary inputs: 1.61 GB)
        ac  peak  24.73 GB  (after forward  13.71)  step 1.38 s  wrapped_layers 24  (boundary inputs: 1.61 GB)
ac+offload  peak  24.73 GB  (after forward  12.17)  step 1.86 s  wrapped_layers 24  (boundary inputs: 1.61 GB)

The forward keeps 1.54 GB less against a predicted 1.61 GB, so the mechanism does exactly what it says. The peak is unchanged, and that is worth being explicit about: at this length the peak sits in the backward's recompute, not in what survives the forward, so removing the boundary inputs has nothing to bite on. The flag pays off once the boundary inputs are what sets the peak, which is the long-context case: at 1,048,576 tokens with cp_size=8, Qwen3-0.6B goes from 27.9 GB to 20.8 GB (−25%) for +0.6% step time, and the 7.1 GB saved matches the 7.5 GB of boundary inputs at that length.

The short-sequence row also shows the cost: +0.5 s on a 1.4 s step. The copy volume grows with sequence length but so does the compute it hides behind, so the trade only becomes free at long context. Async double-buffering is the natural follow-up for the short-sequence case.

repro_offload_memory.py

import time

import torch
from accelerate import Accelerator
from transformers import AutoModel

MODEL, SEQ, STEPS = "ibm-granite/granite-3.1-1b-a400m-instruct", 32768, 2


def main():
    accelerator = Accelerator()
    plugin = accelerator.state.fsdp_plugin
    label = ("ac+offload" if plugin.activation_checkpointing_offload
             else "ac" if plugin.activation_checkpointing else "none")

    model = AutoModel.from_pretrained(MODEL, dtype=torch.bfloat16, attn_implementation="sdpa")
    config = model.config.get_text_config()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
    model, optimizer = accelerator.prepare(model, optimizer)

    ids = torch.randint(0, config.vocab_size, (1, SEQ), generator=torch.Generator().manual_seed(0))
    ids = ids.to(accelerator.device)
    for step in range(STEPS):
        torch.cuda.reset_peak_memory_stats()
        start = time.perf_counter()
        loss = model(input_ids=ids).last_hidden_state.float().square().mean()
        after_forward = torch.cuda.max_memory_allocated() / 1e9
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start

    boundary = config.num_hidden_layers * SEQ * config.hidden_size * 2 / 1e9
    wrapped = sum("Checkpoint" in type(module).__name__ for module in model.modules())
    print(f"{label:>10}  peak {torch.cuda.max_memory_allocated() / 1e9:6.2f} GB"
          f"  (after forward {after_forward:6.2f})  step {elapsed:5.2f} s"
          f"  wrapped_layers {wrapped}  (boundary inputs: {boundary:.2f} GB)", flush=True)


if __name__ == "__main__":
    main()

At scale, with everything else identical:

setup	peak alloc	note
Qwen3-8B @ 1M tokens, cp_size=8	56.2 GB, 364 s/step	OOM without offload
Qwen3-30B-A3B (MoE) @ 1M tokens	46.0 GB, 483 s/step
Qwen3-8B @ 1M, driven purely from the YAML flag	40.2 GB, 377 s/step	with fsdp_offload_params as well; no custom code
Qwen3-8B @ 2M / 4M, 2 / 4 nodes	48.2 / 47.2 GB	696 s / 1346 s per step

Together with parameter offload this is what makes >=8B models fit at 1M-token sequences at all: without it the same configuration OOMs in the ring-attention backward.

Robustness checks

I also ran a few sanity checks on the wrapper's behavior, to make sure it does not pin host memory unnecessarily:

case	result
gradient accumulation (4 micro-steps per optimizer step)	trains, loss decreases, no host-memory growth
checkpoint save + resume at 1M tokens	saves and resumes at the same speed/memory (see below)
forward with no backward (torch.no_grad, evaluation)	no stash entries created
a forward whose backward never runs	host copy released with the tensor (weakref stash)

About checkpointing

The wrapper subclasses torch's ActivationWrapper, so it inherits the state-dict hooks that hide the wrapper from parameter names: a checkpoint written from a wrapped model has exactly the same keys as an unwrapped one and loads into it.
Without that (eg, subclass nn.Module directly), saving a checkpoint fails with RuntimeError: An unexpected key, model.layers.0._checkpoint_wrapped_module.self_attn.q_proj.weight, exists. Covered by a unit test that compares state-dict keys against an unwrapped model and round-trips a load_state_dict, and verified at scale: a Qwen3-8B checkpoint saved from a 1M-token run contains 399 tensors and zero keys mentioning the wrapper.

Notes

• Requires fsdp_activation_checkpointing: true and fsdp_version: 2 (validated in the plugin __post_init__).
• Only the first positional tensor argument (hidden_states) is offloaded, uniquely owned by its layer's closure; rope embeddings/masks are shared across layers and stay resident.
• The stash holds a weakref to the GPU tensor, so a forward that is never followed by a backward (an aborted step) releases its host copy instead of pinning it.
• The auto-wrap policy look-through is generalized to any wrapper exposing _checkpoint_wrapped_module so wrapped layers keep their own FSDP group.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.