pytorch
b96bb2f1 - [spmd] Introduce ParallelMode and add DTensorExpandMode (#98452)

Commit
1 year ago
[spmd] Introduce ParallelMode and add DTensorExpandMode (#98452) This PR introduces a ParallelMode interface to define how to do SPMD expansion and optimize the captured graph. This would be beneifical for different parallelisms to expand differently and apply different optimization passes Put DTensorExpandMode as the first parallel mode that does the existing dtensor_expand functionality. Differential Revision: [D45174399](https://our.internmc.facebook.com/intern/diff/D45174399) Pull Request resolved: https://github.com/pytorch/pytorch/pull/98452 Approved by: https://github.com/mrshenli
Author
Committer
Parents
Loading