Avoid GPU syncs by reusing Pre-allocated Zero Tensor

[quanta42]

This commit improves the FullyShardedDataParallel (FSDP) class in PyTorch by reducing unnecessary GPU synchronizations by reusing a pre-allocated zero tensor.

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❌ 1 New Failure

As of commit 204485d with merge base 92151c8 ():

NEW FAILURE - The following job has failed:

• pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, amz2023.linux.2xlarge) (gh)
  'test/onnx/dynamo/test_exporter_api.py::TestDynamoExportAPI::test_save_succeeds_when_model_greater_than_2gb_and_destination_is_str'

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: quanta42 (24887e0, 204485d, 7726c7e)

[awgu]

Should we initialize this lazily (upon calling clip_grad_norm_ and actually needing to use this)?

Maybe we can put it in a private attribute like self._zero_scalar?

[awgu]

Thanks @quanta42 for the contribution! Could you sign the CLA?

[quanta42]

Thanks @quanta42 for the contribution! Could you sign the CLA?

I did sign it but I work for Adobe so I need the authorization from the administrators.
I also cc them in slack, hopefully, they will get back to me soon.

Thanks for the prompt review!

[awgu]

LGTM! Could you fix lint?

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

[quanta42]

@awgu I was wondering why this hasn't been merged after fixing the linting?
I don't have the authorization to merge this and it seems that now tests are failing probably because of code update?
What's the process to merge this type of changes?

[awgu]

sorry this fell through the cracks
let me rebase and try to remerge

[awgu]

@pytorchbot rebase -s

[pytorchmergebot]

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

[pytorchmergebot]

Successfully rebased patch-1 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout patch-1 && git pull --rebase)

[quanta42]

sorry this fell through the cracks let me rebase and try to remerge

Thank you

[awgu]

https://github.com/pytorch/pytorch/actions/runs/10476605434/job/29016485768
pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, amz2023.linux.2xlarge)
failure is not related

[awgu]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 1 checks: pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, amz2023.linux.2xlarge)

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