pytorch
a02a58d8 - [FSDP][1/N]Add device_mesh to FSDPstate (#102317) (#102551)

Commit
1 year ago
[FSDP][1/N]Add device_mesh to FSDPstate (#102317) (#102551) This PR creates a device_mesh and share it across all FSDP state. The device_mesh will later be used to test out dtensor state_dict (1d device_mesh). Approved by: https://github.com/awgu Add device mesh to fsdp state skip dist.get_world_size(pg) != dist.get_world_size() address test_fake_pg.py test failure fix test_fake_py.py failure Pull Request resolved: https://github.com/pytorch/pytorch/pull/102551 Approved by: https://github.com/fegin
Author
Committer
Parents
Loading