torch.distributed.pipelining tutorial #2962
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/2962
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit de14717 with merge base cad4839 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Also, since it's .rst, we can target to merge directly to main. |
|
Thanks so much for the comments @svekars, that was very helpful! Also updated the base branch to merge directly into main |
| device = torch.device(f"cuda:{rank}") if torch.cuda.is_available() else torch.device("cpu") | ||
| dist.init_process_group() | ||
|
|
||
| pp_group = dist.new_group() |
There was a problem hiding this comment.
its a little funny that we show creating a new group for pp usage, but we don't explain why, and it's the same size as the default group. I think its good to leave pp_group here but maybe add a comment explaining that in this example its trivial but it could be a sub-group in N-D parallel cases
PyTorch 2.4 release of
torch.distributed.pipelining, pytorch docs: https://pytorch.org/docs/main/distributed.pipelining.html#pipeline-parallelismhttps://fburl.com/workplace/9w8z3das
tutorial preview: https://docs-preview.pytorch.org/pytorch/tutorials/2962/intermediate/pipelining_tutorial.html