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!

[qgallouedec]

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) =========================================

[tohtana]

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?

[qgallouedec]

Thanks @tohtana, confirmed. With a no-grad block input the frozen params stay AVAILABLE across grad-accumulation microbatches, as you describe.

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-reset_step, approach and couldn't get it to engage). So I'll leave the call to you: merge this as-is (it fixes the CheckpointError crash on the common HF PEFT path) and address the release lifetime in a follow-up, fold it into a more complete fix, or close this in favor of a global solution, your call happy to close it if best for you