Repository navigation
Reshard FSDP2 modules before creating the optimizer - #50535
Open
qgallouedec wants to merge 3 commits into
Open
qgallouedec wants to merge 3 commits into
qgallouedec wants to merge 3 commits into
Conversation
|
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. |
SunMarc
approved these changes
Nov 12, 2026
SunMarc
left a comment
SunMarc
left a comment
Member
There was a problem hiding this comment.
Thanks ! just a couple of nits
Comment on lines
35
to
+37
| PEFT_SAVE_SCRIPT = os.path.join(SCRIPTS_DIR, "peft_save_distributed_config.py") | ||
| GRAD_ACCUM_SCRIPT = os.path.join(SCRIPTS_DIR, "grad_accum_distributed_config.py") | ||
| EVAL_BEFORE_TRAIN_SCRIPT = os.path.join(SCRIPTS_DIR, "eval_before_train_distributed_config.py") |
Member
There was a problem hiding this comment.
for all these, maybe just move them in the test directly, i prefer not having global var + we will probably have much more test file in the future
| parser.add_argument("--output_dir", required=True) | ||
| script_args = parser.parse_args() | ||
|
|
||
| dist.init_process_group(backend=get_distributed_backend(torch_device), timeout=timedelta(seconds=60)) |
Member
There was a problem hiding this comment.
we shouldn't need this no ?
Suggested change
| dist.init_process_group(backend=get_distributed_backend(torch_device), timeout=timedelta(seconds=60)) |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Found while fixing #48841.
After
evaluate(), the FSDP2 modules sharded withreshard_after_forward=False(the final norm and LM head underDistributedConfig, the root under Accelerate) keep their unsharded parameters registered.train()then builds the optimizer from those, the next forward registers the sharded ones again, and these weights never train. This reshards every FSDP2 module before creating the optimizer.evaluate()thentrain()underDistributedConfig: on mainmodel.norm.weightandlm_head.weightnever change, with this PR every parameter updates. Same under Accelerate FSDP2, on top of Allow preparing a model without an optimizer under FSDP2 accelerate#4519, which fixes the FSDP2: evaluate/predict on a fresh Trainer raises in Accelerator.prepare #48841 crash.