pytorch/torch/distributed/pipelining
Howard Huang 9b6d680131 Remove stage_index_to_group_rank from schedule (#146217)
This PR allows schedules loaded via CSV to automatically set their `stage_index_to_group_rank ` and removes the `stage_index_to_group_rank ` argument from the `PipelineScheduleMulti` constructor

Pull Request resolved: https://github.com/pytorch/pytorch/pull/146217
Approved by: https://github.com/wconstab
ghstack dependencies: #146193
2025-02-05 21:26:45 +00:00
..
__init__.py
_backward.py [NFC] Fix some minor typos. (#145599) 2025-01-24 18:58:59 +00:00
_debug.py remove allow-untyped-defs from torch/distributed/pipelining/_debug.py (#143871) 2024-12-27 01:20:26 +00:00
_IR.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
_unflatten.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
_utils.py Remove stage_index_to_group_rank from schedule (#146217) 2025-02-05 21:26:45 +00:00
microbatch.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
README.md
schedules.py Remove stage_index_to_group_rank from schedule (#146217) 2025-02-05 21:26:45 +00:00
stage.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00

Pipeline Parallelism for PyTorch

torch.distributed.pipelining is a package for implementing pipeline parallelism on your model.

Our documentation is available here.

pipeline_diagram_web