Add V2 sharding support and improve partition spec handling for multichip training

[sshonTT]

These changes are required to support multi-chip training for real models on the torch-xla side.

• Added V2 OpSharding support in XlaShardingSpec, which is used internally by MpLoader for parallel input loading. The original implementation only supported V1 shardings.
• Fixed environment variable parsing for CONVERT_SHLO_TO_SHARDY - previous logic treated values like "0" or "false" as truthy.
• Added logic to compute dims, reshape_dims, and transpose_perm for V2 sharding based on mesh_shape and partition_spec.

The new logic now correctly handles cases that were previously unsupported:

case 1: mesh_shape=(2,1,1,1), partition_spec=(0,None,None,None)
-> dims=[2,1,1,1], reshape_dims=[2,1,1,1], transpose_perm=[0,1,2,3]

case 2: mesh_shape=(2,1,1,1), partition_spec=(0,)
-> dims=[2], reshape_dims=[2,1,1,1], transpose_perm=[0,1,2,3]

case 3: mesh_shape=(2,4), partition_spec=(0,None)
-> dims=[2,1,4], reshape_dims=[2,4], transpose_perm=[0,1]

[hshahTT]

Hey Sungjoon, could you add more details to the PR about why this is needed and an overview of the high-level changes? Also, are you sure this is needed "for GSPMD support"? The purpose of the V2 shardings is to enable Shardy support within torch-xla. Torch-xla should already natively support GSPMD without needing any changes.

[sshonTT]

Hey Sungjoon, could you add more details to the PR about why this is needed and an overview of the high-level changes? Also, are you sure this is needed "for GSPMD support"? The purpose of the V2 shardings is to enable Shardy support within torch-xla. Torch-xla should already natively support GSPMD without needing any changes.

Hey Het, thanks for the feedback — and sorry for the lack of detail earlier. I’ll update the commit message to include more context.

These changes are required to support multi-chip training for real models from the torch-xla side. Specifically:

We use MpLoader for parallel input loading, which internally calls xla_spec and _XLAC.XlaShardingSpec. That code currently creates an OpSharding using the v1 format, so I added v2 support alongside it.

I also updated how we handle os.environ.get('CONVERT_SHLO_TO_SHARDY', False), since the previous logic treated values like "0" or "false" as truthy, which was causing unexpected behavior.

And you're right — GSPMD itself is already supported in torch-xla. What I meant was GSPMD-based training, and I’ll clarify that in the updated message.

[hshahTT]

I think you can just rename initial_dims to `dims and remove L179, since you're not using initial dims after this line.

[sshonTT]

Agreed! FWIW, it was working on the older branch (link), but not on the new one (link). I’m looking into it now. Unless it turns out to require changes in torch-xla, I probably won’t need anything, but if it does, I’ll ping you for a review.