Force use_reentrant=True for PEFT + ZeRO-3 + gradient checkpointing in all trainers#6356
Conversation
… in all trainers
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
![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: ab2a9c56f9
ℹ️ 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".
| # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see | ||
| # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703 | ||
| if ( | ||
| is_peft_model(model) | ||
| and args.deepspeed_plugin is not None | ||
| and args.deepspeed_plugin.zero_stage == 3 | ||
| and args.gradient_checkpointing | ||
| ): |
There was a problem hiding this comment.
This new guard is only added to the five main trainers, but the same PEFT + ZeRO-3 + gradient-checkpointing path still exists in experimental copies: repo-wide search shows trl/experimental/{cpo,orpo,bco,tpo}_trainer.py and trl/experimental/distillation/distillation_trainer.py still wrap with get_peft_model under args.deepspeed_plugin.zero_stage == 3, while several of those trainers later default use_reentrant to False. Those configurations will still hit the CheckpointError this change is meant to avoid, and AGENTS.md explicitly requires duplicated trainer logic to be updated consistently when one copy changes.
Useful? React with 👍 / 👎.
1 participant
#4951 added a guard to
SFTTrainerthat forcesuse_reentrant=Truewhen a PEFT model is combined with DeepSpeed ZeRO-3 and gradient checkpointing. This propagates the same guard to the other trainers that share the pattern: DPO, GRPO, RLOO, Reward, KTO.Why
Since #4811 flipped the gradient-checkpointing default to
use_reentrant=False, PEFT + ZeRO-3 + gradient checkpointing crashes in the backward pass withCheckpointError: Recomputed values ... have different metadata(frozen params get partitioned to shape[0]during the checkpoint recompute).This is a DeepSpeed-side bug (deepspeedai/DeepSpeed#4332), fixed by deepspeedai/DeepSpeed#8130; forcing reentrant checkpointing is the interim workaround until that lands. Closes #5217. Related huggingface/transformers#47254.
The guard is copied verbatim from
SFTTrainer(identical block, same placement) and warns if the user explicitly passeduse_reentrant=False.Note
Medium Risk
Changes training behavior only for a narrow stack (PEFT + ZeRO-3 + checkpointing) but directly affects backward-pass correctness and may surprise users who set
use_reentrant=False.Overview
Extends the
SFTTrainerworkaround to DPO, GRPO, RLOO, Reward, and KTO so PEFT models on DeepSpeed ZeRO-3 with gradient checkpointing always getgradient_checkpointing_kwargs["use_reentrant"] = True.That combination otherwise hits backward
CheckpointErrorwhen the default non-reentrant checkpointing is used (known DeepSpeed/PEFT interaction; upstream fix pending in DeepSpeed#8130). Users who explicitly passuse_reentrant=Falseget a warning that the setting is overridden.SFTTraineronly gains an updated comment noting the guard can be removed after DeepSpeed#8130 ships.Reviewed by Cursor Bugbot for commit 4c2d94a. Bugbot is set up for automated code reviews on this repo. Configure here.