ZeRO-3: don't partition frozen params during an activation-checkpoint recompute#8130
Open
qgallouedec wants to merge 7 commits into
Open
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
… recompute Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>
…dParameterCoordinator Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>
qgallouedec
force-pushed
the
fix/zero3-frozen-param-checkpoint-recompute
branch
from
July 11, 2026 01:36
9d2cb7f to
d652fd9
Compare
Collaborator
|
Hi @qgallouedec thanks for your fix! Can you also add a UT in this PR to expose the issue? Thanks! |
Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>
Contributor
Author
|
Thanks, done Before the fix: After the fix: |
delock
approved these changes
Jul 12, 2026
delock
enabled auto-merge
July 12, 2026 02:18
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Jul 12, 2026
tohtana
requested changes
Jul 12, 2026
Collaborator
There was a problem hiding this comment.
Hi @qgallouedec,
Thank you for submitting a PR! The behavior of Z3+AC has always been a bit confusing and often thrown errors that are hard to understand for many users. This PR will significanly improves the issue.
The direction of the fix looks good to me. However, I see there is still a release-lifetime issue when the checkpointed block receives an input that does not require gradients. ZeRO only registers the module post-backward callback when that module input requires gradients, but the new recompute path skips the frozen parameter's forward release unconditionally within the graph task.
In a two-rank ZeRO-3 test with
- stage3_param_persistence_threshold=0
- no-grad input
- frozen layer followed by a trainable layer
both release-eligible frozen parameters remained AVAILABLE after the first backward. Running a second gradient-accumulation microbatch before step() confirmed that they remained gathered across microbatches and caused the coordinator accounting to diverge. They were only partitioned by the optimizer-step safety net.
Can you double check this behavior and add this no-grad-input/delayed-optimizer-step case?
Contributor
Author
|
Thanks @tohtana, confirmed. With a no-grad block input the frozen params stay Honestly, I don't understand ZeRO-3's release lifecycle deeply enough to land a clean fix for the no-grad case here (I tried, helped by Claude, a deferred-release-at- |
Collaborator
|
@qgallouedec thanks so much for tackling this high-impact issue. I think the root cause is that metadata mapping between params and modules (such as ds_active_sub_modules) does not properly handle nested forward/backward execution of activation-checkpoint. I am exploring a fix along this line and will update this thread asap. |
Collaborator
|
Hi @qgallouedec, @sfc-gh-truwase, I've also been exploring the same direction, using a unique The hard part is determining when each invocation has actually finished using its gathered parameters: releasing them in the recompute forward post-hook is too early, while keeping all of them until GraphTask completion is unnecessarily late. My current attempt is here. (supplemental PR to this PR) I ran four order-counterbalanced benchmark runs in a master → candidate → candidate → master sequence (codex suggested) using my benchmark script. The workload used 8×H100, Llama 3 8B, BF16, ZeRO-3, sequence length 1024, micro-batch size 1, activation checkpointing, 10 warmup steps, and 30 measured steps.
The performance impact seems okay, but I'm not still sure this is the best approach. I needed to use some PyTorch internal APIs though we had already done it. If @sfc-gh-truwase can suggest a better approach, I'm happy to follow it. |
Collaborator
|
@tohtana thanks for sharing WIP and results. A slight adjustment to the module/param mapping seems to work on the original case shared by @qgallouedec. I further extended that case as below to assert param sharding on module exit. # deepspeed --num_gpus 2 mre.py
import torch, torch.nn as nn, deepspeed
from torch.utils.checkpoint import checkpoint
from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
D = 1024
def dump_params(model):
for name, p in model.named_parameters():
print(name, p.ds_id, p.shape, p.requires_grad, p.ds_status)
def assert_all_not_available(model):
for name, p in model.named_parameters():
assert p.ds_status == ZeroParamStatus.NOT_AVAILABLE, f"{name} is not not available"
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)
for _ in range(4):
x = torch.randn(1, 8, D, device=engine.device, requires_grad=True)
engine.backward(engine(x).float().sum())
dump_params(model)
assert_all_not_available(model)
engine.step()
|
Collaborator
|
Specifically, the key change is to map the frozen params to the backward module. that uses them, so the backward module is responsible for releasing them. @tohtana, @qgallouedec I will appreciate your feedback. @tohtana, I lack your extensive testing framework, so could you please test my solution. @qgallouedec, it would be helpful if you could test as well. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
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
CheckpointError(recomputed shape[0])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:
CheckpointErrorwith PEFT + DeepSpeed ZeRO-3 + gradient checkpointing huggingface/transformers#47254 minimal PEFT + ZeRO-3 repro.It's now the default failure in huggingface/transformers:
use_reentrant=False(PyTorch recommended) huggingface/transformers#43203 switched the gradient-checkpointing default touse_reentrant=False(Switch gradient checkpointing default to use_reentrant=False (PyTorch recommended) huggingface/trl#4811 did the same in TRL).Prior mitigation / adjacent fixes:
use_reentrant=Truefor PEFT + ZeRO-3 + gradient checkpointing (downstream guard working around this bug).non_reentrant_checkpointto support inputs require no grad #4118 + usenon_reentrant_checkpointfix requires_grad of input must be true for activation checkpoint layer in pipeline train. #4224 addednon_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.use_reentrant=Truestage-3 path.