[DSD] Fix set_optimizer_state_dict() changes the parameters with some optimizers #125708
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125708
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 3d827e5 with merge base 196a0b1 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot merge -f "The failing tests are not related." |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
…5338) Summary: This is useful if users would like to avoid CPU memory OOM when loading from a full state_dict. Pull Request resolved: #125338 Approved by: https://github.com/weifengpy ghstack dependencies: #125708
…25339) Summary: This is useful if users would like to avoid CPU memory OOM when loading from a full state_dict. Pull Request resolved: #125339 Approved by: https://github.com/weifengpy ghstack dependencies: #125708, #125338
|
Is there going to be a PyTorch 2.3.2, and if so would it be possible to get this fix in it? I've spent all day running down slight parameter differences in my model when loading checkpoints as this is called in |
Stack from ghstack (oldest at bottom):
Summary:
Some optimizers, like AdamW, change the parameters even if gradients are zero. So
set_optimizer_state_dict()may affect the parameters values with these optimizers. This PR fixes the issue.This PR also fixes #121186.
cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @LucasLLC