Skip to content

DeviceMesh helper functions#2228

Merged
cowanmeg merged 5 commits intoNVIDIA:mainfrom
cowanmeg:mesh_cleanup
May 14, 2024
Merged

DeviceMesh helper functions#2228
cowanmeg merged 5 commits intoNVIDIA:mainfrom
cowanmeg:mesh_cleanup

Conversation

@cowanmeg
Copy link
Collaborator

@cowanmeg cowanmeg commented May 9, 2024

This PR introduces some common convenience helper functions for DeviceMesh include size, at, etc. so that we rely less on accessing the underlying vector of DeviceIdxType. There are still a few places in csrc/multidevice/lower_communications.cpp that access the vector directly.

Eventually, these helper functions will be parameterized by the axis of the DeviceMesh.

@cowanmeg
Copy link
Collaborator Author

cowanmeg commented May 9, 2024

Note: I will update this after #2185 is merged.

@cowanmeg cowanmeg requested review from samnordmann and wujingyue May 9, 2024 20:55
Copy link
Collaborator

@wujingyue wujingyue left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you!

@cowanmeg
Copy link
Collaborator Author

!build --dist

@cowanmeg
Copy link
Collaborator Author

Failures are unrelated from the missing imports.

@cowanmeg cowanmeg merged commit 03716fc into NVIDIA:main May 14, 2024
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.

3 participants