Skip to content

Comments

Delay device mesh import#561

Merged
2015aroras merged 4 commits intomainfrom
shanea/delay-device-mesh-import
Apr 26, 2024
Merged

Delay device mesh import#561
2015aroras merged 4 commits intomainfrom
shanea/delay-device-mesh-import

Conversation

@2015aroras
Copy link
Collaborator

DeviceMesh is introduced in torch 2.2, and so trying to import it in torch 2.1 is caused the training code to break. This PR delays its import until it is needed (for hybrid sharding).

Fixes #559

@2015aroras 2015aroras requested a review from dumitrac April 26, 2024 16:03
@2015aroras 2015aroras marked this pull request as ready for review April 26, 2024 16:15
@2015aroras 2015aroras merged commit 295d309 into main Apr 26, 2024
@2015aroras 2015aroras deleted the shanea/delay-device-mesh-import branch April 26, 2024 19:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

No module named 'torch.distributed.device_mesh'

1 participant