[DSD] Correctly handle _extra_state (#125336)

[mvpatel2000]

Pr for 2.3 cherry-pick

Summary:
distributed_state_dict should not try to use getattr to get _extra_state as this is not well-defined.

Pull Request resolved: #125336
Approved by: https://github.com/LucasLLC
ghstack dependencies: #125333, #125501, #125334, #125335

Fixes #ISSUE_NUMBER

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @LucasLLC

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/126567

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 67f0410 with merge base 86a2d67 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / linux-focal-py3.12-clang10 / test (default, 3, 3, linux.2xlarge) (gh) (disabled by #123655 but the issue was closed recently and a rebase is needed to make it pass)
  test_binary_ufuncs.py::TestBinaryUfuncsCPU::test_logaddexp_cpu_complex128

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[fegin]

Thanks for submitting the cherry pick.

[atalman]

@pytorchbot rebase -b release/2.3

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/release/2.3. Check the current status here

[pytorchmergebot]

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/release/2.3 pull/126567/head returned non-zero exit code 1

Rebasing (1/1)
Auto-merging test/distributed/checkpoint/test_state_dict.py
CONFLICT (content): Merge conflict in test/distributed/checkpoint/test_state_dict.py
Auto-merging torch/distributed/checkpoint/state_dict.py
error: could not apply 90619a2a174... [DSD] Correctly handle _extra_state (#125336)
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
hint: Disable this message with "git config advice.mergeConflict false"
Could not apply 90619a2a174... [DSD] Correctly handle _extra_state (#125336)

Raised by https://github.com/pytorch/pytorch/actions/runs/9225766192

[atalman]

@pytorchbot rebase -b release/2.3

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/release/2.3. Check the current status here

[pytorchmergebot]

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/release/2.3 pull/126567/head returned non-zero exit code 1

Rebasing (1/1)
Auto-merging test/distributed/checkpoint/test_state_dict.py
CONFLICT (content): Merge conflict in test/distributed/checkpoint/test_state_dict.py
Auto-merging torch/distributed/checkpoint/state_dict.py
error: could not apply 90619a2a174... [DSD] Correctly handle _extra_state (#125336)
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
hint: Disable this message with "git config advice.mergeConflict false"
Could not apply 90619a2a174... [DSD] Correctly handle _extra_state (#125336)

Raised by https://github.com/pytorch/pytorch/actions/runs/9225866547

[PaliC]

Validated this works for cpu on the 2.3 release branch