pytorch
38b687ed - [PTD][Checkpoint] Add checkpointing support for DTensor submesh (#96802)

Commit
1 year ago
[PTD][Checkpoint] Add checkpointing support for DTensor submesh (#96802) DTensor submesh support is added in https://github.com/pytorch/pytorch/pull/95458. This PR adds support for DTensor submesh by adding an extra check when create local save/load plan. If the rank is not participating in the mesh, we simply skip creating WriteItem/ReadItem for the local SavePlan/LoadPlan. Updated the associated test as well. cc. @wanchaol, @kumpera Pull Request resolved: https://github.com/pytorch/pytorch/pull/96802 Approved by: https://github.com/wanchaol
Author
Committer
Parents
Loading