fix: reset prefetch flag upon reshard

[0x804d8000]

The prefetched flag should be reset upon reshard. Otherwise, for zero2, next access to the corresponding parameter will skip "unshard" operation, and results in wrong parameter shape.

The need of unsharding is also metioned in the comment of FlatParameterHandle.unshard.

As FlatParameterHandle already guarded it against unnecessary all gather, this shouldn't incur extra communication overhead.

Personally I also find _prefetched a bit of mis-named, it should really be _unsharded.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 adef22c with merge base e0e15a4 ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• trunk / linux-focal-rocm5.6-py3.8 / test (default, 1, 3, linux.rocm.gpu, unstable) (gh)

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

[0x804d8000]

I updated the UT to add a case when no_grad is used in forward pass. This is not needed to reproduce the failure, though without it we'd see a different error saying loss mismatch between FSDP and DDP. I suspect the reason is somehow related to https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_flat_param.py#L1681 .

The thing that does make a difference is forward_prefetch. Formerly test_hooks_multi_traversal did not specify this argument and left it to False. This prevented _prefetch_handle from ever running, and kept handle._prefetched always False. Therefore, _unshard was always called by _pre_forward_unshard, and thus no error was raised. Now with forward_prefetch=True added to the test, the error shows up.

[awgu]

Thanks for the fix!

I think the usage of the _prefetched flag has changed slowly over time. If we were to rename it to _unsharded, we may want to check its usage more precisely since I think it is not currently used if we do not do any prefetching.

[awgu]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here