Fix DeepSpeed ZeRO-3 + PEFT + Trainer scheduler mismatch on PyTorch 2.10+
Author: qgallouedec
Pull request: yes

## Summary

On PyTorch 2.10+, `Trainer.train()` crashes at the first `lr_scheduler.step()` when using DeepSpeed ZeRO-3 with a PEFT model.

This PR provides fix, alothough I'm sure it's not the ideal one.

The failure only appears with torch 2.10+ probably because PyTorch made the scheduler update strict (`zip(..., strict=True)`), surfacing a latent mismatch between the scheduler’s `base_lrs` and the optimizer’s param groups after DeepSpeed wraps the optimizer. This PR rebuilds the scheduler after DeepSpeed/Accelerate preparation if a mismatch is detected.

## Reproduction

From the `transformers` repo:

```sh
accelerate launch --config_file zero3.yaml mre.py
```

`mre.py`:

```python
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
from peft import LoraConfig, get_peft_model

model_id = "Qwen/Qwen3-0.6B"
data = [
    [46518, 374, 2664, 1091, 27261],
    [98923, 374, 2664, 1091, 17995],
    [16374, 374, 2664, 1091, 6351],
    [31137, 374, 2664, 1091, 16965],
]
Trainer(
    model=get_peft_model(AutoModelForCausalLM.from_pretrained(model_id), LoraConfig()),
    args=TrainingArguments(bf16=True),  # not related: bf16=False fails
    train_dataset=Dataset.from_dict({"input_ids": data, "labels": data}),
    processing_class=AutoTokenizer.from_pretrained(model_id),
).train()
```

`zero3.yaml`:

```yaml
distributed_type: DEEPSPEED
deepspeed_config:
  zero_stage: 3
num_processes: 2
```

**Observed error (torch 2.10+)**

```
W0202 21:38:37.706000 2150915 site-packages/torch/distributed/run.py:852] 
W0202 21:38:37.706000 2150915 site-packages/torch/distributed/run.py:852] *****************************************
W0202 21:38:37.706000 2150915 site-packages/torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
W0202 21:38:37.706000 2150915 site-packages/torch/distributed/run.py:852] *****************************************
Loading weights: 100%|████████████████████████████████████████████████| 311/311 [00:00<00:00, 4280.70it/s, Materializing param=model.norm.weight]
Loading weights: 100%|████████████████████████████████████████████████| 311/311 [00:00<00:00, 4247.09it/s, Materializing param=model.norm.weight]
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.
The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.
Stage 3 initialize beginning
MA 1.4 GB         Max_MA 1.4 GB         CA 1.43 GB         Max_CA 1 GB 
CPU Virtual Memory:  used = 63.8 GB, percent = 3.2%
DeepSpeedZeRoOffload initialize [begin]
MA 1.4 GB         Max_MA 1.4 GB         CA 1.43 GB         Max_CA 1 GB 
CPU Virtual Memory:  used = 63.94 GB, percent = 3.2%
Parameter Offload - Persistent parameters statistics: param_count = 225, numel = 1212416
DeepSpeedZeRoOffload initialize [end]
MA 0.73 GB         Max_MA 1.55 GB         CA 1.62 GB         Max_CA 2 GB 
CPU Virtual Memory:  used = 63.98 GB, percent = 3.2%
Before creating fp16 partitions
MA 0.73 GB         Max_MA 0.73 GB         CA 1.62 GB         Max_CA 2 GB 
CPU Virtual Memory:  used = 63.99 GB, percent = 3.2%
After creating fp16 partitions: 1
MA 0.73 GB         Max_MA 0.73 GB         CA 1.33 GB         Max_CA 2 GB 
CPU Virtual Memory:  used = 63.99 GB, percent = 3.2%
Before creating fp32 partitions
MA 0.73 GB         Max_MA 0.73 GB         CA 1.33 GB         Max_CA 1 GB 
CPU Virtual Memory:  used = 63.99 GB, percent = 3.2%
After creating fp32 partitions
MA 0.73 GB         Max_MA 0.73 GB         CA 1.33 GB         Max_CA 1 GB 
CPU Virtual Memory:  used = 63.99 GB, percent = 3.2%
Before initializing optimizer states
MA 0.73 GB         Max_MA 0.73 GB         CA 1.33 GB         Max_CA 1 GB 
CPU Virtual Memory:  used = 63.99 GB, percent = 3.2%
After initializing optimizer states
MA 0.73 GB         Max_MA 0.73 GB         CA 1.33 GB         Max_CA 1 GB 
CPU Virtual Memory:  used = 63.99 GB, percent = 3.2%
After initializing ZeRO optimizer
MA 1.66 GB         Max_MA 1.66 GB         CA 2.27 GB         Max_CA 2 GB 
CPU Virtual Memory:  used = 64.03 GB, percent = 3.2%
  0%|                                                                                                                      | 0/3 [00:00<?, ?it/s][rank1]: Traceback (most recent call last):
[rank1]:   File "/fsx/qgallouedec/transformers/mre.py", line 17, in <module>
[rank1]:     ).train()
[rank1]:       ^^^^^^^
[rank1]:   File "/fsx/qgallouedec/transformers/src/transformers/trainer.py", line 2174, in train
[rank1]:     return inner_training_loop(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/fsx/qgallouedec/transformers/src/transformers/trainer.py", line 2614, in _inner_training_loop
[rank1]:     self.lr_scheduler.step()
[rank1]:   File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/optim/lr_scheduler.py", line 282, in step
[rank1]:     self._update_lr(epoch)
[rank1]:   File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/optim/lr_scheduler.py", line 296, in _update_lr
[rank1]:     for param_group, lr in zip(self.optimizer.param_groups, values, strict=True):
[rank1]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: ValueError: zip() argument 2 is longer than argument 1
[rank0]: Traceback (most recent call last):
[rank0]:   File "/fsx/qgallouedec/transformers/mre.py", line 17, in <module>
[rank0]:     ).train()
[rank0]:       ^^^^^^^
[rank0]:   File "/fsx/qgallouedec/transformers/src/transformers/trainer.py", line 2174, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/fsx/qgallouedec/transformers/src/transformers/trainer.py", line 2614, in _inner_training_loop
[rank0]:     self.lr_scheduler.step()
[rank0]:   File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/optim/lr_scheduler.py", line 282, in step
[rank0]:     self._update_lr(epoch)
[rank0]:   File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/optim/lr_scheduler.py", line 296, in _update_lr
[rank0]:     for param_group, lr in zip(self.optimizer.param_groups, values, strict=True):
[rank0]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: ValueError: zip() argument 2 is longer than argument 1
  0%|                                                                                                                      | 0/3 [00:01<?, ?it/s]
W0202 21:39:11.776000 2150915 site-packages/torch/distributed/elastic/multiprocessing/api.py:1010] Sending process 2151022 closing signal SIGTERM
E0202 21:39:11.788000 2150915 site-packages/torch/distributed/elastic/multiprocessing/api.py:984] failed (exitcode: 1) local_rank: 1 (pid: 2151023) of binary: /fsx/qgallouedec/miniconda3/envs/trl/bin/python3
Traceback (most recent call last):
  File "/fsx/qgallouedec/miniconda3/envs/trl/bin/accelerate", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/fsx/qgallouedec/accelerate/src/accelerate/commands/accelerate_cli.py", line 50, in main
    args.func(args)
  File "/fsx/qgallouedec/accelerate/src/accelerate/commands/launch.py", line 1384, in launch_command
    deepspeed_launcher(args)
  File "/fsx/qgallouedec/accelerate/src/accelerate/commands/launch.py", line 1076, in deepspeed_launcher
    distrib_run.run(args)
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/distributed/run.py", line 982, in run
    elastic_launch(
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 170, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 317, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
mre.py FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2026-02-02_21:39:11
  host      : ip-26-0-164-45.ec2.internal
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 2151022)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2026-02-02_21:39:11
  host      : ip-26-0-164-45.ec2.internal
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 2151023)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
```

**My understanding**

- Trainer builds optimizer + scheduler early (based on original param groups: decay/no‑decay).
- DeepSpeed ZeRO-3 wraps/rewrites optimizer param groups during `accelerator.prepare`.
- Scheduler still holds the old `base_lrs`, so `len(base_lrs) != len(param_groups)`.
- PyTorch 2.10+ uses `zip(..., strict=True)` and raises immediately.

**Fix**

After `accelerator.prepare`, if DeepSpeed is enabled and a Trainer‑created scheduler exists, compare `len(optimizer.param_groups)` with `len(lr_scheduler.base_lrs)` and rebuild the scheduler if they differ. This keeps the scheduler consistent with the final optimizer state.

I'm quite sure this is not the ideal fix (and I'm fine with not merging fhis PR), but this one works and will probably help figuring out a better long‑term solution.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@SunMarc 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_43692). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

closing in favour of #43711