`CheckpointError` with PEFT + DeepSpeed ZeRO-3 + gradient checkpointing
Author: qgallouedec
Pull request: no
Labels: name: bug

### System Info


- Platform: Linux-5.15.0-1048-aws-x86_64-with-glibc2.31
- Python version: 3.13.13
- PyTorch version: 2.11.0+cu128
- Transformers version: 5.13.0.dev0
- Accelerate version: 1.15.0.dev0
- DeepSpeed version: 0.19.2
- PEFT version: 0.19.1
- GPUs: 2x NVIDIA H100 80GB


### Who can help?

@SunMarc 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Training a PEFT (LoRA) model with `gradient_checkpointing=True` under DeepSpeed ZeRO-3 (≥2 processes) crashes in the backward pass with:

```
torch.utils.checkpoint.CheckpointError: Recomputed values for the following tensors
have different metadata than during the forward pass.
tensor at position 4: saved {'shape': [1024]} vs recomputed {'shape': [0]}
```

The `[0]`-shaped tensors are **frozen** parameters (RMSNorm weights: `hidden_size` and the per-head q/k norm `head_dim`). ZeRO-3 all-gathers them during the forward, but they are back in the partitioned state (empty) during the checkpoint recompute, so non-reentrant (now default) checkpointing's metadata check rejects them.

This is **not** specific to TRL / bitsandbytes / DoRA / VLMs: The original report (huggingface/trl#5217) used all of those, but the MRE below reproduces with none of them: plain `transformers.Trainer` + PEFT + ZeRO-3 + gradient checkpointing.

## Reproduction

`mre.py`

```python
from datasets import Dataset
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")
model = get_peft_model(model, LoraConfig())

dataset = Dataset.from_dict({"input_ids": [[1, 2, 3, 4]], "labels": [[1, 2, 3, 4]]})

trainer = Trainer(model, TrainingArguments(gradient_checkpointing=True), train_dataset=dataset)
trainer.train()
```

`ds_z3.yaml`

```yaml
distributed_type: DEEPSPEED
num_processes: 2
deepspeed_config:
  zero_stage: 3
```

```bash
accelerate launch --config_file ds_z3.yaml mre.py
```


## Workaround

Set `use_reentrant=True` for gradient checkpointing:

```python
TrainingArguments(gradient_checkpointing=True,
                  gradient_checkpointing_kwargs={"use_reentrant": True})
```

but it's old, and not recommended. in other words, not a good solution, because `use_reentrant=False` is the modern config for checkpointing.

## Full traceback

```
File ".../torch/utils/checkpoint.py", line 1177, in unpack_hook
    frame.check_recomputed_tensors_match(gid)
File ".../torch/utils/checkpoint.py", line 921, in check_recomputed_tensors_match
    raise CheckpointError(
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for
the following tensors have different metadata than during the forward pass.
tensor at position 4:
saved metadata:      {'shape': torch.Size([1024]), 'dtype': torch.bfloat16, 'device': cuda:0}
recomputed metadata: {'shape': torch.Size([0]),    'dtype': torch.bfloat16, 'device': cuda:0}
tensor at position 14:
saved metadata:      {'shape': torch.Size([128]),  'dtype': torch.bfloat16, 'device': cuda:0}
recomputed metadata: {'shape': torch.Size([0]),    'dtype': torch.bfloat16, 'device': cuda:0}
...
```

Related: huggingface/transformers#34928, huggingface/trl#5217, huggingface/trl#2514


### Expected behavior

To work

Comments:
I traced the exact mechanism, and it's a bit different from "frozen params aren't re-gathered". It's a PyTorch ↔ DeepSpeed interaction, and the `requires_grad` discriminator is actually on the PyTorch side.

## Root cause

During the non-reentrant checkpoint recompute of a block:

1. DeepSpeed does re-gather the frozen param. The ZeRO-3 forward pre-hook fires on the recompute forward and all-gathers every param with `ds_status == NOT_AVAILABLE`. So the frozen RMSNorm weights are briefly full-shape (`[1024]`/`[128]`), same as trainable ones.
2. At the end of that submodule's recompute forward, the ZeRO-3 post-hook re-partitions the param in place: `param.data = torch.empty(0, ...)`  → [`partition_parameters.py:315`](https://github.com/deepspeedai/DeepSpeed/blob/v0.19.2/deepspeed/runtime/zero/partition_parameters.py#L315). This mutates the existing `Parameter` object; anything still holding a reference to it now sees shape `[0]`.
3. PyTorch's checkpoint saves the recomputed tensors, but only detaches the ones that require grad: `x = x.detach() if x.requires_grad else x`  → [`torch/utils/checkpoint.py:1084`](https://github.com/pytorch/pytorch/blob/v2.11.0/torch/utils/checkpoint.py#L1084).
   - trainable param → a detached copy, shape frozen at save time → immune to step 2.
   - frozen param → the live `Parameter` object itself → shrinks to `[0]` when step 2 runs.
4. `check_recomputed_tensors_match` runs *after* recompute (i.e. after step 2) and reads the stored reference live → frozen params report `[0]` → `CheckpointError`. Trainable params match.

So both frozen and trainable params are re-gathered; the asymmetry is entirely that torch keeps a live, undetached reference for non-grad tensors while DeepSpeed mutates that same object in place.

Empirical confirmation: unfreezing the norm weights (making them `requires_grad=True`) makes the error disappear.

## Where a real fix belongs

The interaction is the bug; either side could break it:
- PyTorch: snapshot metadata (or detach unconditionally) when packing recomputed tensors, so a later in-place `.data` swap can't corrupt the saved reference. The strict `check_recomputed_tensors_match` is intentional and has no relax-for-sharded-params option.
- DeepSpeed: don't `free_param` a param that was gathered *for a recompute forward* (or keep it gathered for the checkpoint region). The ZeRO-3 shape-`[0]` root cause is tracked in deepspeedai/DeepSpeed#4332 (open since 2023-09, unfixed). DeepSpeed's `non_reentrant_checkpoint` (deepspeedai/DeepSpeed#4118) is pipeline-only and still breaks under stage 3 (deepspeedai/DeepSpeed#4595).

## Workaround

`gradient_checkpointing_kwargs={"use_reentrant": True}`: works, validated end-to-end (trains), and this what we've in trl, see https://github.com/huggingface/trl/pull/4951

## Possible transformers-side mitigation

Since #43203 made this the default, the Trainer could, when ZeRO-3 + gradient checkpointing + a frozen/PEFT model + no explicit `gradient_checkpointing_kwargs` are all present, warn (or auto-set `use_reentrant=True`). Mirroring the existing FSDP-vs-GC guard. Would be the same as https://github.com/huggingface/trl/pull/4951



The fix is here: https://github.com/deepspeedai/DeepSpeed/pull/8130