[Experimental] Add initial implementation of GSPMD->Shardy pass within PyTorch/XLA #1
Conversation
This file contains hidden or 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
…sues. Will fix later
… into hshah/add-gspmd-to-shardy-pass
… into hshah/add-gspmd-to-shardy-pass
…xla into hshah/add-gspmd-to-shardy-pass
…sues. Will fix later
…xla into hshah/add-gspmd-to-shardy-pass
…d-gspmd-to-shardy-pass
| return t | ||
|
|
||
| op_sharding = mesh.get_op_sharding(partition_spec) | ||
| if os.environ.get('CONVERT_SHLO_TO_SHARDY', False): |
There was a problem hiding this comment.
should this be hidden under a different env var or does shardy inherently only understand V2?
There was a problem hiding this comment.
Shardy only understands V2. I wasn't able to get the pass working with a V1 graph, and also Kevin mentioned that V2 is a required work item for getting the Shardy pass working: pytorch#9348 (comment)
There was a problem hiding this comment.
ok, awesome, then it can remain under the CONVERT_SHLO_TO_SHARDY flag.
AleksKnezevic
approved these changes
Jul 18, 2025
hshahTT
added a commit
that referenced
this pull request
Aug 30, 2025
…n PyTorch/XLA (#1) Adds an environment variable CONVERT_SHLO_TO_SHARDY that does 2 things: - Uses V2 sharding annotations when generating the GSPMD SHLO module (i.e., in V1 a mesh annotation string like: devices=[2,1,4]0,1,2,3,4,5,6,7 becomes this in V2: devices=[2,1,4]<=[8]). - Converts the new GSPMD module with the V2 annotations into a Shardy module.
hshahTT
added a commit
that referenced
this pull request
Sep 7, 2025
…n PyTorch/XLA (#1) Adds an environment variable CONVERT_SHLO_TO_SHARDY that does 2 things: - Uses V2 sharding annotations when generating the GSPMD SHLO module (i.e., in V1 a mesh annotation string like: devices=[2,1,4]0,1,2,3,4,5,6,7 becomes this in V2: devices=[2,1,4]<=[8]). - Converts the new GSPMD module with the V2 annotations into a Shardy module.
hshahTT
added a commit
that referenced
this pull request
Oct 3, 2025
…n PyTorch/XLA (#1) Adds an environment variable CONVERT_SHLO_TO_SHARDY that does 2 things: - Uses V2 sharding annotations when generating the GSPMD SHLO module (i.e., in V1 a mesh annotation string like: devices=[2,1,4]0,1,2,3,4,5,6,7 becomes this in V2: devices=[2,1,4]<=[8]). - Converts the new GSPMD module with the V2 annotations into a Shardy module.
hshahTT
added a commit
that referenced
this pull request
Oct 3, 2025
…n PyTorch/XLA (#1) Adds an environment variable CONVERT_SHLO_TO_SHARDY that does 2 things: - Uses V2 sharding annotations when generating the GSPMD SHLO module (i.e., in V1 a mesh annotation string like: devices=[2,1,4]0,1,2,3,4,5,6,7 becomes this in V2: devices=[2,1,4]<=[8]). - Converts the new GSPMD module with the V2 annotations into a Shardy module.
sshonTT
pushed a commit
that referenced
this pull request
Oct 3, 2025
…n PyTorch/XLA (#1) Adds an environment variable CONVERT_SHLO_TO_SHARDY that does 2 things: - Uses V2 sharding annotations when generating the GSPMD SHLO module (i.e., in V1 a mesh annotation string like: devices=[2,1,4]0,1,2,3,4,5,6,7 becomes this in V2: devices=[2,1,4]<=[8]). - Converts the new GSPMD module with the V2 annotations into a Shardy module.
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.
Adds an environment variable
CONVERT_SHLO_TO_SHARDYthat does 2 things:devices=[2,1,4]0,1,2,3,4,5,6,7becomes this in V2:devices=[2,1,4]<=[8]).