SPMD with scalar parameter

[giladxw]

🐛 Bug

In torch_xla.experimental.spmd_fully_sharded_data_parallel, the function "_prepare_spmd_partition_spec" handles
scalar wrongly, instead of returning a tuple it returns a list:
this should be:

  if len(partition_spec) == 0:
    return tuple(partition_spec)

otherwise we get the error:

  File "/usr/local/lib/python3.10/site-packages/torch_xla/experimental/spmd_fully_sharded_data_parallel.py", line 121, in __init__
    spmd.mark_sharding(
  File "/usr/local/lib/python3.10/site-packages/torch_xla/runtime.py", line 95, in wrapper
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch_xla/distributed/spmd/xla_sharding.py", line 582, in mark_sharding
    op_sharding = mesh.get_op_sharding(partition_spec)
TypeError: unhashable type: 'list'

To Reproduce

test SpmdFullyShardedDataParallel (i.e. FSDPv2) with any model that contains a scalar parameter

Environment

• Reproducible on XLA backend [CPU/TPU/CUDA]: ALL
• torch_xla version: 2.4.0, 2.5.1

[radna0]

@giladxw , I'm not on the xla team but from what I see, this already handles that, I think it's correct that it should return a tuple even if length is 0, But is it not unlikely that you will pass a param that has shape of 0?