[16/N] Add _allgather_base custom op with CPU/CUDA implementation #88889
Conversation
[ghstack-poisoned]
H-Huang
requested review from
mrshenli,
pritamdamania87,
zhaojuanmao,
rohan-varma,
awgu and
kwen2501
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/88889
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 16fab8b: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
release notes: distributed (c10d)
| &::c10d::ProcessGroup::_allgather_base, | ||
| [](const c10::intrusive_ptr<::c10d::ProcessGroup>& self, | ||
| at::Tensor& output_tensor, | ||
| at::Tensor& input_tensor, |
There was a problem hiding this comment.
i think input_tensor can be const &
There was a problem hiding this comment.
I had a look at init.cpp, it seems the convention has not been using const. Not sure why. Can go with convention I guess.
| &::c10d::ProcessGroup::_allgather_base, | ||
| [](const c10::intrusive_ptr<::c10d::ProcessGroup>& self, | ||
| at::Tensor& output_tensor, | ||
| at::Tensor& input_tensor, |
There was a problem hiding this comment.
I had a look at init.cpp, it seems the convention has not been using const. Not sure why. Can go with convention I guess.
test/distributed/test_c10d_nccl.py
Outdated
| device = "cuda" | ||
| tensor = torch.ones(10, 10, device=torch.device(device)) | ||
| output_tensor = torch.zeros(10, 10, device=torch.device(device)) | ||
| dist._all_gather_base(output_tensor, tensor) |
There was a problem hiding this comment.
Please use all_gather_into_tensor at Python level.
We will also rename _all_gather_base at the binding and C++ levels to all_gather_into_tensor when we introduce extension-breaking changes in 2.0.
|
@H-Huang has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
…ntation" Differential Revision: [D41227739](https://our.internmc.facebook.com/intern/diff/D41227739) [ghstack-poisoned]
|
@pytorchbot merge |
Merge startedYour 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 |
…torch#88889) Differential Revision: [D41227739](https://our.internmc.facebook.com/intern/diff/D41227739) Pull Request resolved: pytorch#88889 Approved by: https://github.com/kwen2501
Stack from ghstack:
Context
#86225
Differential Revision: D41227739