ZeRO-3: don't partition frozen params during an activation-checkpoint recompute by qgallouedec · Pull Request #8130 · deepspeedai/DeepSpeed

qgallouedec · 2026-07-10T21:27:35Z

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".

… recompute Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>

…dParameterCoordinator Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>

delock · 2026-07-11T13:48:28Z

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>

qgallouedec · 2026-07-11T23:25:41Z

Thanks, done

Before the fix:

$ python -m pytest v1/zero/test_zero_user_backward.py -k FrozenParamCheckpointing -v
================================================================ test session starts ================================================================
platform linux -- Python 3.13.13, pytest-9.1.0, pluggy-1.6.0 -- /fsx/qgallouedec/miniconda3/envs/trl/bin/python
cachedir: .pytest_cache
rootdir: /fsx/qgallouedec/DeepSpeed/tests
configfile: pytest.ini
plugins: anyio-4.13.0, datadir-1.8.0, xdist-3.8.0, rerunfailures-15.1, cov-7.1.0
collected 61 items / 55 deselected / 6 selected                                                                                                     

v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1] PASSED               [ 16%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-2] PASSED               [ 33%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-3] PASSED               [ 50%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-1] PASSED              [ 66%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-2] PASSED              [ 83%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-3] FAILED              [100%]

===================================================================== FAILURES ======================================================================
_______________________________ TestZeroUserBackwardFrozenParamCheckpointing.test_checkpointed_frozen_param[False-3] ________________________________
multiprocessing.pool.RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
                    ~~~~^^^^^^^^^^^^^^^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/multiprocessing/pool.py", line 51, in starmapstar
    return list(itertools.starmap(args[0], args[1]))
  File "/fsx/qgallouedec/DeepSpeed/tests/unit/common.py", line 338, in _dist_run
    raise e
  File "/fsx/qgallouedec/DeepSpeed/tests/unit/common.py", line 330, in _dist_run
    self.run(**self._fixture_kwargs)
    ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/qgallouedec/DeepSpeed/tests/unit/common.py", line 481, in run
    self._current_test(**fixture_kwargs)
    ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^
  File "/fsx/qgallouedec/DeepSpeed/tests/unit/v1/zero/test_zero_user_backward.py", line 1516, in test_checkpointed_frozen_param
    output_ds.backward(torch.ones_like(output_ds))
    ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/_tensor.py", line 631, in backward
    torch.autograd.backward(
    ~~~~~~~~~~~~~~~~~~~~~~~^
        self, gradient, retain_graph, create_graph, inputs=inputs
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/autograd/__init__.py", line 381, in backward
    _engine_run_backward(
    ~~~~~~~~~~~~~~~~~~~~^
        tensors,
        ^^^^^^^^
    ...<5 lines>...
        accumulate_grad=True,
        ^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/autograd/graph.py", line 869, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        t_outputs, *args, **kwargs
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
    )  # Calls into the C++ engine to run the backward pass
    ^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/utils/checkpoint.py", line 1177, in unpack_hook
    frame.check_recomputed_tensors_match(gid)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/utils/checkpoint.py", line 921, in check_recomputed_tensors_match
    raise CheckpointError(
    ...<4 lines>...
    )
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
tensor at position 0:
saved metadata: {'shape': torch.Size([8]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
tensor at position 2:
saved metadata: {'shape': torch.Size([8]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
.

Tip: To see a more detailed error message, either pass `debug=True` to
`torch.utils.checkpoint.checkpoint(...)` or wrap the code block
with `with torch.utils.checkpoint.set_checkpoint_debug_enabled(True):` to
enable checkpoint‑debug mode globally.

"""

The above exception was the direct cause of the following exception:

item = <Function test_checkpointed_frozen_param[False-3]>

    @pytest.hookimpl(tryfirst=True)
    def pytest_runtest_call(item):
        # We want to use our own launching function for distributed tests
        if getattr(item.cls, "is_dist_test", False):
            dist_test_class = item.cls()
>           dist_test_class(item._request)

../conftest.py:69: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
common.py:498: in __call__
    self._launch_with_file_store(request, world_size)
common.py:350: in _launch_with_file_store
    self._launch_procs(procs, init_method)
common.py:293: in _launch_procs
    self._launch_daemonic_procs(num_procs, init_method)
common.py:203: in _launch_daemonic_procs
    skip_msgs = skip_msgs_async.get(self.exec_timeout)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <multiprocessing.pool.MapResult object at 0x7fea84d15a90>, timeout = 600

    def get(self, timeout=None):
        self.wait(timeout)
        if not self.ready():
            raise TimeoutError
        if self._success:
            return self._value
        else:
>           raise self._value
E           torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
E           tensor at position 0:
E           saved metadata: {'shape': torch.Size([8]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
E           recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
E           tensor at position 2:
E           saved metadata: {'shape': torch.Size([8]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
E           recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
E           .
E           
E           Tip: To see a more detailed error message, either pass `debug=True` to
E           `torch.utils.checkpoint.checkpoint(...)` or wrap the code block
E           with `with torch.utils.checkpoint.set_checkpoint_debug_enabled(True):` to
E           enable checkpoint‑debug mode globally.

../../../miniconda3/envs/trl/lib/python3.13/multiprocessing/pool.py:774: CheckpointError
--------------------------------------------------------------- Captured stderr call ----------------------------------------------------------------
/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/autograd/graph.py:869: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:335.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
================================================================= warnings summary ==================================================================
../../../miniconda3/envs/trl/lib/python3.13/site-packages/torch/jit/_script.py:365: 14 warnings
  /fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/jit/_script.py:365: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

<string>:8
  <string>:8: PytestDeprecationWarning: A private pytest class or function was used.

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
  /fsx/qgallouedec/DeepSpeed/tests/conftest.py:47: UserWarning: Running test without verifying torch version, please provide an expected torch version with --torch_ver
    warnings.warn(

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
  /fsx/qgallouedec/DeepSpeed/tests/conftest.py:54: UserWarning: Running test without verifying cuda version, please provide an expected cuda version with --cuda_ver
    warnings.warn(

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
  /fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/_pytest/fixtures.py:1312: PytestRemovedIn10Warning: Class-scoped fixture defined as instance method is deprecated.
  Instance attributes set in this fixture will NOT be visible to test methods,
  as each test gets a new instance while the fixture runs only once per class.
  Use @classmethod decorator and set attributes on cls instead.
  See https://docs.pytest.org/en/stable/deprecations.html#class-scoped-fixture-as-instance-method
    fixturefunc = resolve_fixture_function(fixturedef, request)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================================================================= slowest durations =================================================================
29.27s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-3]
29.06s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-2]
28.96s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
28.58s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-1]
28.39s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-2]
28.28s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-3]

(12 durations < 1s hidden.)
============================================================== short test summary info ==============================================================
FAILED v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-3] - torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during ...
======================================== 1 failed, 5 passed, 55 deselected, 18 warnings in 193.17s (0:03:13) ========================================

After the fix:

$ python -m pytest v1/zero/test_zero_user_backward.py -k FrozenParamCheckpointing -v
=============================================== test session starts ===============================================
platform linux -- Python 3.13.13, pytest-9.1.0, pluggy-1.6.0 -- /fsx/qgallouedec/miniconda3/envs/trl/bin/python
cachedir: .pytest_cache
rootdir: /fsx/qgallouedec/DeepSpeed/tests
configfile: pytest.ini
plugins: anyio-4.13.0, datadir-1.8.0, xdist-3.8.0, rerunfailures-15.1, cov-7.1.0
collected 61 items / 55 deselected / 6 selected                                                                   

v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1] PASSED  [ 16%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-2] PASSED  [ 33%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-3] PASSED  [ 50%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-1] PASSED [ 66%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-2] PASSED [ 83%]
v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-3] PASSED [100%]

========================================================== warnings summary ===========================================================
../../../miniconda3/envs/trl/lib/python3.13/site-packages/torch/jit/_script.py:365: 14 warnings
  /fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/torch/jit/_script.py:365: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

<string>:8
  <string>:8: PytestDeprecationWarning: A private pytest class or function was used.

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
  /fsx/qgallouedec/DeepSpeed/tests/conftest.py:47: UserWarning: Running test without verifying torch version, please provide an expected torch version with --torch_ver
    warnings.warn(

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
  /fsx/qgallouedec/DeepSpeed/tests/conftest.py:54: UserWarning: Running test without verifying cuda version, please provide an expected cuda version with --cuda_ver
    warnings.warn(

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
  /fsx/qgallouedec/miniconda3/envs/trl/lib/python3.13/site-packages/_pytest/fixtures.py:1312: PytestRemovedIn10Warning: Class-scoped fixture defined as instance method is deprecated.
  Instance attributes set in this fixture will NOT be visible to test methods,
  as each test gets a new instance while the fixture runs only once per class.
  Use @classmethod decorator and set attributes on cls instead.
  See https://docs.pytest.org/en/stable/deprecations.html#class-scoped-fixture-as-instance-method
    fixturefunc = resolve_fixture_function(fixturedef, request)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================================= slowest durations =============================================================
30.55s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-1]
29.84s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-2]
29.73s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[True-3]
29.65s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-1]
29.25s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-2]
29.12s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardFrozenParamCheckpointing::test_checkpointed_frozen_param[False-3]

(12 durations < 1s hidden.)
========================================= 6 passed, 55 deselected, 18 warnings in 199.92s (0:03:19) =========================================

qgallouedec requested review from tjruwase and tohtana as code owners July 10, 2026 21:27

qgallouedec mentioned this pull request Jul 10, 2026

CheckpointError with PEFT + DeepSpeed ZeRO-3 + gradient checkpointing huggingface/transformers#47254

Open

4 tasks

chatgpt-codex-connector Bot reviewed Jul 10, 2026

View reviewed changes

Comment thread deepspeed/runtime/zero/partitioned_param_coordinator.py Outdated

qgallouedec mentioned this pull request Jul 10, 2026

Force use_reentrant=True for PEFT + ZeRO-3 + gradient checkpointing in all trainers huggingface/trl#6356

Open

qgallouedec added 2 commits July 10, 2026 22:01

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

6346875

… recompute Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>

Fix checkpoint recompute handling for frozen parameters in Partitione…

d652fd9

…dParameterCoordinator Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>

qgallouedec force-pushed the fix/zero3-frozen-param-checkpoint-recompute branch from 9d2cb7f to d652fd9 Compare July 11, 2026 01:36

qgallouedec and others added 2 commits July 11, 2026 18:31

Merge branch 'master' into fix/zero3-frozen-param-checkpoint-recompute

b1095c7

Add regression test for ZeRO-3 gradient checkpointing with frozen params

4ff4998

Signed-off-by: Quentin Gallouédec <gallouedec.quentin@gmail.com>

qgallouedec requested a review from loadams as a code owner July 11, 2026 23:11

delock approved these changes Jul 12, 2026

View reviewed changes

delock enabled auto-merge July 12, 2026 02:18

delock added this pull request to the merge queue Jul 12, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 12, 2026

Merge branch 'master' into fix/zero3-frozen-param-checkpoint-recompute

8c4a47b

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

ZeRO-3: don't partition frozen params during an activation-checkpoint recompute#8130
qgallouedec wants to merge 5 commits into
deepspeedai:masterfrom
qgallouedec:fix/zero3-frozen-param-checkpoint-recompute

qgallouedec commented Jul 10, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

delock commented Jul 11, 2026

Uh oh!

qgallouedec commented Jul 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

qgallouedec commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

MRE

References

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

delock commented Jul 11, 2026

Uh oh!

qgallouedec commented Jul 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

qgallouedec commented Jul 10, 2026 •

edited

Loading