Shardy support

[nsmithtt]

🚀 Feature

Support shardy instead of GSPMD.

Motivation

GSPMD will soon become a deprecated IR in openxla in favor of shardy.

Pitch

There are probably two paths forward:

• Support the GSPMD -> Shardy conversion pass from OpenXLA, this can probably enable shardy based IR faster.
• Support shardy as a first class backend IR in torch-xla.

[nsmithtt]

Would love to get some input from torch-xla maintainers who'd have a much better feel for the scope of work here. Also, if it's already been thought about / considered.

[pgmoka]

Hello, me and @GleasonK both took a look at this a while ago. I will add here what I know, and I will let Kevin correct me or add anything:

The biggest hurdle on the way to achieve Shardy is that it consumes MLIR. Currently we output both HLO and StableHLO. Migrating fully to StableHLO should be relatively trivial. Migrating the entire system to use MLIR is desirable, but initial estimates for achieving this is really long as it requires a lot of changes accross PyTorchXLA.

Some of the potential work being done related to #9019 will likely see PyTorchXLA eventually use IFRT, but the timing on that is not settled in stone at the moment.

The "GSPMD -> Shardy conversion pass from OpenXLA" initial implementation might be what needs to be done at first depending on the timing of deprecation for GSPMD.

[nsmithtt]

Awesome, thank you for the quick response! This is a good insight, we're going to investigate what it takes to invoke the GSPMD -> Shardy conversion pass.

[GleasonK]

That's probably the best path. I'm curious what your PTXLA ingestion pipeline looks like? AFAIK today torch_xla pass HLO to PJRT plugins, are you using the XLA_STABLEHLO_COMPILE flag? If so adding the SDY support you want shouldn't be too difficult!

cc @tomnatan30 @bartchr808

[nsmithtt]

Yes we are using that flag:

os.environ["PJRT_DEVICE"] = "TT"
os.environ["XLA_STABLEHLO_COMPILE"] = "1"

I think we also tried, without luck:

config.update("jax_use_shardy_partitioner", True)

[GleasonK]

Chatted with the Shardy folks - there's going to be two work items for getting these passes to work:

1. The conversion should happen in the PT/XLA layer right around the XLA_STABLEHLO_COMPILE callsite before passing the IR to the PJRT plugin - this will save you from needing to take a dependency on XLA since the passes require an XLA dep for the sharding proto.
2. We'll need PTXLA to build OpShardingV2 annotations (V1: devices=[4,2]0,1,2,3,4,5,6,7,8 vs V2: devices=[4,2]<=[8]) -- this isn't a hard change in theory since PT/XLA is already marking sharding based on a mesh, and that's the only requirement of V2.

So it wont be as trivial as I initially thought, but we should do this in the near future IMO (both changes independent of eachother are improvements, (2) is a compile time improvement independent of SDY support since it reduces program size on large distributed programs by a lot.

[hshahTT]

@GleasonK Thanks so much for looking into this! Regarding (2), could you give me a code pointer to where PTXLA is currently building the V1 annotations?

Also, regarding (1), we found this pipeline within the OpenXLA repo that converts modules with mhlo.shardings into the Shardy dialect: https://github.com/openxla/xla/blob/d6beb849cf6ac93f419d6be27a9d1eac1b72420a/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc#L678

We tested it with the SHLO output of some TorchXLA programs and see that (2) needs to happen before it converts the GSPMD module to Shardy (otherwise it just errors out). Is this also what you had in mind for the GSPMD->Shardy conversion path?

[GleasonK]

Code pointer on (2) is probably around here (I'd assume contributions are welcome here, its functionally the same and JAX has already migrated to use OpShardingV2):

xla/torch_xla/distributed/spmd/xla_sharding.py

Line 652 in 3a1ed62

op_sharding = mesh.get_op_sharding(partition_spec)

It generates an OpSharding of the old format. I think @sdasgup3 may have ideas on how this would look at a V2.

And RE (1) yeah that's exactly the pass I think. But tt-xla shouldn't need a dep on XLA just for that pass - PTXLA already has a dep so we can instead call that pass here once (2) is addressed.

[qihqi]

config.update("jax_use_shardy_partitioner", True)

fwiw, after setting the option above, it only changes what IR jax emits and not what torchxla emits. However, given that we have call_jax feature launched recently, we can probably get shardy support with composition of call_jax and jax.lax.with_sharding_constraint.

This is akin to how shard_as is implemented in torch_xla:

xla/torch_xla/distributed/spmd/xla_sharding.py

Line 726 in bea86ee

x, y = xb.call_jax(shard_alike, (x, y))

cc. @tengyifei

[tengyifei]

Furthermore, if you wrap a PyTorch function with https://docs.pytorch.org/xla/master/perf/assume_pure.html, and enable the jax Shardy config, then any xs.mark_sharding call within that PyTorch function will be staged out with Shardy annotations.

That's because @assume_pure lowers your PyTorch function using torchax to JAX then to HLO.

[hshahTT]

Hey everyone, sorry for the long message.

I'm working on adding OpShardingV2 annotations to torch/xla (mentioned by @GleasonK as work item 2 here).

I found the following documentation that gives some examples and explanations of how to create the V2 sharding annotations:

• https://github.com/jax-ml/jax/blob/84af8a8e74c05ce4196079e145d50f0c9504ff16/jax/_src/named_sharding.py#L414
• https://github.com/openxla/xla/blob/bdfcd696c1c4311b40ccd081f5286011199d931b/xla/hlo/ir/tile_assignment.h#L50

To make sure that I understand the format correctly, could someone with more knowledge please confirm whether the dims, reshape_dims, and transpose_perm lists for the following examples are correct (so I know if I'm on the right track)?

In all cases we define a 2D device mesh with x and y axes in a [2, 4] shape (the x axis has 2 devices and y has 4):

mesh = Mesh([0,...,7], (2, 4), ("x", "y"))

Also, in all cases a 2D tensor is being sharded. The reshape_dims will also be [2, 4] everywhere.

1. PartitionSpec = ('x', 'y')
   a. dims = [2, 4]
   b. transpose_perm = [0, 1]

2. PartitionSpec = ('y', 'x')
   a. dims = [4, 2]
   b. transpose_perm = [1, 0]

3. PartitionSpec = ('x', None)
   a. dims = [2, 1, 4]
   b. transpose_perm = [0, 1]

4. PartitionSpec = ('y', None)
   a. dims = [4, 1, 2]
   b. transpose_perm = [1, 0]

5. PartitionSpec = (None, 'x')
   a. dims = [1, 2, 4]
   b. transpose_perm = [0, 1]

6. PartitionSpec = (None, 'y')
   a. dims = [1, 4, 2]
   b. transpose_perm = [1, 0]

7. PartitionSpec = (('x', 'y'), None)
   a. dims = [8, 1]
   b. transpose_perm = [0, 1]

8. PartitionSpec = (('y', 'x'), None)
   a. dims = [8, 1]
   b. transpose_perm = [1, 0]

9. PartitionSpec = (None, ('x', 'y'))
   a. dims = [1, 8]
   b. transpose_perm = [0, 1]

10. PartitionSpec = (None, ('y', 'x'))
    a. dims = [1, 8]
    b. transpose_perm = [1, 0]

[GleasonK]

I think dims should always be [2,4], and the rest is reshape and transpose ('y', 'x') -> [2,4]T(1,0). But not fully sure how the translation from named axes into numbered device mesh works or what None does to axes dims, cc @tomnatan30 for that.

[tomnatan30]

Hey! The transpose perm is used to reorder the axes in the mesh so they can be used in any order. So in the above example 4 you actually need transpose perm [1, 0] because y is used before x: PartitionSpec = ('y', None) a. dims = [4, 1, 2] b. transpose_perm = [0, 1] There is also the reshape dims (which must be specified), which can merge consecutive axes in the mesh if they are kept consecutive, for example: Mesh - [x:2, y:2, z:2] Pspec = (z, (x, y)) a. dims = [2, 4] b. reshape = [4, 2] c. Transpose = [1, 0] Finally, this allows you to use sub axes, e.g. for axis x of size 4 you can use a sub axis x:(1)2 (offset 1, size 2) or x:(2)2 (offset 2, size 2). Shardy also supports this with the notation “x”:(offset)size. Hope this helps!

…

On Thu, 26 Jun 2025 at 22:53, Kevin Gleason ***@***.***> wrote: *GleasonK* left a comment (pytorch/xla#9348) <#9348 (comment)> *I think* dims should always be [2,4], and the rest is reshape and transpose ('y', 'x') -> [2,4]T(1,0). But not fully sure how the translation from named axes into numbered device mesh works or what None does to axes dims, cc @tomnatan30 <https://github.com/tomnatan30> for that. — Reply to this email directly, view it on GitHub <#9348 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/A7DIFHZN3DRO4DPH6V36PTT3FRTWVAVCNFSM6AAAAAB7FP35K2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTAMJQGI2TQNZSGM> . You are receiving this because you were mentioned.Message ID: ***@***.***>