ZeRO-3: don't partition frozen params during an activation-checkpoint recompute

[qgallouedec]

Problem

ZeRO-3 + gradient checkpointing (torch non-reentrant, the PyTorch/HF default) + any frozen params (e.g. LoRA adapters or a quantized base) crashes in backward:

torch.utils.checkpoint.CheckpointError: Recomputed values ... have different metadata ...
tensor at position N: saved {'shape': [1024]} vs recomputed {'shape': [0]}

The recompute re-fires ZeRO-3's forward hooks; the post-hook partitions params in place (param.data = torch.empty(0)). torch's checkpoint keeps a live, undetached reference to non-grad tensors (x = x.detach() if x.requires_grad else x), so freeing a frozen param shrinks the very tensor the recompute is about to validate → shape [0] → error. Trainable params are saved detached and are unaffected, which is why the bug only appears with frozen params.

Fix

In PartitionedParameterCoordinator.release_sub_module, skip partitioning a frozen (non-grad) param when a forward release fires inside a backward (forward and torch._C._current_graph_task_id() != -1), i.e. a checkpoint recompute; it is released normally by the ensuing backward. Trainable params still partition module-by-module, so full finetuning (and its memory profile) is unchanged. Correctness-preserving — the recompute runs on full-shape weights — and a no-op outside recompute.

MRE

# deepspeed --num_gpus 2 mre.py
import torch, torch.nn as nn, deepspeed
from torch.utils.checkpoint import checkpoint

D = 1024

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.LayerNorm(D)   # frozen below -> recomputes to shape [0]
        self.lin = nn.Linear(D, D)    # trainable -> keeps the block on the autograd graph
    def forward(self, x):
        return self.lin(self.norm(x))

class Net(nn.Module):
    def __init__(self, n=2):
        super().__init__()
        self.blocks = nn.ModuleList(Block() for _ in range(n))
    def forward(self, x):
        for b in self.blocks:
            x = checkpoint(b, x, use_reentrant=False)   # non-reentrant (PyTorch default)
        return x

deepspeed.init_distributed()
model = Net()
for name, p in model.named_parameters():
    if "norm" in name:                # freeze the norms, as LoRA/quantization freezes the base
        p.requires_grad_(False)

config = {
    "train_micro_batch_size_per_gpu": 1,
    "optimizer": {"type": "AdamW", "params": {"lr": 1e-4}},
    "zero_optimization": {"stage": 3, "stage3_param_persistence_threshold": 0},
}
engine, *_ = deepspeed.initialize(
    model=model, model_parameters=[p for p in model.parameters() if p.requires_grad], config=config)

x = torch.randn(1, 8, D, device=engine.device, requires_grad=True)
engine.backward(engine(x).float().sum())
engine.step()

• before: CheckpointError (recomputed shape [0])
• after: ok

Verified on 2×H100 (torch 2.11, deepspeed 0.19.2; identical code path on master). Full finetune (no frozen params) is unaffected.

References

This is not an edge case! This has been reported many times in the wild, and is a known limitation of ZeRO-3 + non-reentrant checkpointing.

Root cause / trackers:

• [BUG][stage3] using torch.utils.checkpoint in unet_3d_blocks causes weights size [0] error #4332 root-cause tracker (sept 2023!)
• CheckpointError with PEFT + DeepSpeed ZeRO-3 + gradient checkpointing huggingface/transformers#47254 minimal PEFT + ZeRO-3 repro.
• Recomputed tensor size does not match when using quantization + LoRA with the DPO Trainer with gradient checkpointing using deepspeed and accelerate huggingface/trl#5217 original DPO + QLoRA report.
• Recomputed tensor size does not match when using activation checkpointing when using FSDP and accelerate huggingface/transformers#34928 umbrella issue for this error string — the ZeRO-3 frozen-param reports live here (mixed in with unrelated FSDP / use_cache root causes).

It's now the default failure in huggingface/transformers:

• Switch gradient checkpointing default to use_reentrant=False (PyTorch recommended) huggingface/transformers#43203 switched the gradient-checkpointing default to use_reentrant=False (Switch gradient checkpointing default to use_reentrant=False (PyTorch recommended) huggingface/trl#4811 did the same in TRL).

Prior mitigation / adjacent fixes:

• Fix ZeRO-3 + PEFT + gradient checkpointing huggingface/trl#4951 forces use_reentrant=True for PEFT + ZeRO-3 + gradient checkpointing (downstream guard working around this bug).
• feat(activation_checkpointing): add non_reentrant_checkpoint to support inputs require no grad #4118 + use non_reentrant_checkpoint fix requires_grad of input must be true for activation checkpoint layer in pipeline train. #4224 added non_reentrant_checkpoint (pipeline-only; not used by the HF path); [REQUEST] How to use deepspeed.checkpointing.non_reentrant_checkpoint() properly with Stage3? #4595 shows it still fails under stage 3.
• Fix gradient checkpointing with use_reentrant=True / PyTorch-style backward / ZeRO-3 #7780 fixed the separate use_reentrant=True stage-3 path.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: acd54c5a83

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[delock]

Hi @qgallouedec thanks for your fix! Can you also add a UT in this PR to expose the issue? Thanks!