Extend MeshTensor read/write APIs to support uneven sharding #18735
+275
−52
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
omilyutin-tt
requested review from
cfjchu,
ayerofieiev-tt,
dmakoviichuk-tt,
rfurko-tt,
TT-BrianLiu,
razorback3,
dongjin-na,
bbradelTT,
abhullar-tt,
pgkeller,
aliuTT,
tt-aho,
tt-dma,
tt-asaigal and
ubcheema
as code owners
cfjchu
reviewed
Mar 7, 2025
Comment on lines
+878
to
+880
|
|
||
| // Set storage with the populated metadata. | ||
| mesh_tensor.set_storage(mesh_storage); |
There was a problem hiding this comment.
errm... was this just not doing anything before? 😱
There was a problem hiding this comment.
It was using the same mesh buffer behind the scenes, but there was no metadata to populate, so just worked :)
cfjchu
reviewed
Mar 7, 2025
| std::shared_ptr<Buffer> buffer; | ||
| std::shared_ptr<distributed::MeshBuffer> mesh_buffer; | ||
| std::map<distributed::MeshCoordinate, TensorSpec> specs; |
There was a problem hiding this comment.
why not add it as a constructor argument?
There was a problem hiding this comment.
There is a path that creates DeviceStorage for a single device... I think there are more changes upcoming that touch this, I think we need to tackle this separately.
cfjchu
approved these changes
Mar 7, 2025
Comment on lines
+91
to
+97
| bool operator==(const MeshCoordinate& lhs, const MeshCoordinate& rhs); | ||
| bool operator!=(const MeshCoordinate& lhs, const MeshCoordinate& rhs); | ||
| bool operator<(const MeshCoordinate& lhs, const MeshCoordinate& rhs); | ||
| bool operator>(const MeshCoordinate& lhs, const MeshCoordinate& rhs); | ||
| bool operator<=(const MeshCoordinate& lhs, const MeshCoordinate& rhs); | ||
| bool operator>=(const MeshCoordinate& lhs, const MeshCoordinate& rhs); | ||
| std::ostream& operator<<(std::ostream& os, const MeshCoordinate& shape); |
Comment on lines
+807
to
+810
| DeviceStorage device_storage(mesh_buffer); | ||
| device_storage.specs = specs; | ||
| device_storage.strategy = storage.strategy; | ||
| return device_storage; |
There was a problem hiding this comment.
we should probably move away from direct access like this
There was a problem hiding this comment.
+1, I was just going to say. Should be class.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Ticket
N/A
Problem description
MeshBuffer-backed Tensors should support uneven shards.
What's changed
Changes to read/write paths + tests.
Checklist