[PT-D][Tensor parallelism] Add documentations for TP #94421
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/94421
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 8d99683: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
wanchaol
left a comment
wanchaol
left a comment
There was a problem hiding this comment.
first pass, let's add the experimental line as this is prototype release.
| :members: | ||
|
|
||
| We also enabled 2D parallelism to integrate with ``FullyShardedDataParallel``. | ||
| Users just need to call the following API explicitly: |
There was a problem hiding this comment.
I remembered we have a FSDP extension, Is TP automatically register the extension now?
Also, I wonder if we should give a small code snippet showing how the 2-D parallel look like
There was a problem hiding this comment.
The registrations is in the is_available. Let me send a follow-up PR for this one.
|
|
||
|
|
||
| .. currentmodule:: torch.distributed.tensor.parallel.fsdp | ||
| .. autofunction:: is_available |
There was a problem hiding this comment.
do we really need to add this API to the doc? I remembered is_available is introduced when we are in tau, but since now it's pytorch I think fsdp should always be available?
There was a problem hiding this comment.
Yes, because of 2D hook registration.
There was a problem hiding this comment.
Will send a follow-up PR to address the naming of this one.
This is far from completed and we will definitely polish it down the road. [ghstack-poisoned]
fduwjj
requested review from
H-Huang,
awgu,
kwen2501,
mrshenli,
rohan-varma and
zhaojuanmao
as code owners
|
@pytorchbot rebase |
|
@pytorchbot successfully started a rebase job. Check the current status here |
This is far from completed and we will definitely polish it down the road. [ghstack-poisoned]
|
Successfully rebased |
|
@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 |
Stack from ghstack (oldest at bottom):
This is far from completed and we will definitely polish it down the road.