Error with device_put on TPUv3-32 pod with NamedSharding

[magicknight]

I'm trying to shard my JAX model on a TPUv3-32 pod using device_put. The code works fine on a single TPUv3-8 device, but I'm getting the following error on the TPUv3-32 pod:

device_put's second argument must be a Device or a Sharding which represents addressable devices, but got NamedSharding(mesh={'model': 32, 'data': 1}, spec=PartitionSpec('model', 'data'))

It seems like the issue is with the NamedSharding object that I'm passing as the second argument to device_put. Here's a simplified version of my code:

import jax
from jax.experimental import mesh_utils
from jax.sharding import Mesh, PartitionSpec, NamedSharding
x = jax.random.normal(jax.random.PRNGKey(0), (8192, 8192))
devices = mesh_utils.create_device_mesh(mesh_shape)
mesh = Mesh(devices, axis_names=('model', 'data'))
y = jax.device_put(x, NamedSharding(mesh, PartitionSpec('model', 'data')))

I would appreciate any help in resolving this issue. Thank you!

Steps to reproduce:

Run the code snippet above on a TPUv3-32 pod.
Expected behavior:

The code should successfully shard the input array and place it on the TPUv3-32 devices using device_put.

Actual behavior:

The code throws an error: device_put's second argument must be a Device or a Sharding which represents addressable devices, but got NamedSharding(mesh={'model': 32, 'data': 1}, spec=PartitionSpec('model', 'data')).

[yashk2810]

device_put does not work across multiple processes. You can use jax.make_array_from_callback or jax.make_array_from_single_device_arrays or an identity pjit to do this!

[yashk2810]

In OP's code isn't the array the same on all processes?

We basically run multihost_utils.assert_equal(inp) on all processes to check if the array is the same. You can run that and see what it returns. If it return True and you still see that error, a repro would be nice so that I can debug :)

[j-towns]

Thanks for the info @yashk2810. It seems in my case the error was caused by applying device_put to a Python scalar. On a single process, device_put works fine with scalars. My intention was to replicate the scalar to all devices (and cast it to a JAX array).

Minimal repro on CPU...

save device_put_repro.py:

import sys
import numpy as np
import jax
from jax.sharding import PartitionSpec, NamedSharding

proc_id = int(sys.argv[1])
jax.distributed.initialize('localhost:10000', 2, proc_id)
mesh = jax.make_mesh((2,), ("a",))

# WORKS
jax.device_put(np.array(0), NamedSharding(mesh, PartitionSpec()))

#FAILS
jax.device_put(0, NamedSharding(mesh, PartitionSpec()))

save device_put_repro.sh:

num_processes=2

for i in $(seq 0 $(($num_processes - 1))); do
  python device_put_repro.py $i &
done

wait

then run

$ bash device_put_repro.sh

[yashk2810]

Ahh, that seems like an easy fix. #28736 fixes it. Thanks for reporting!

[yashk2810]

The fix has been submitted.

[j-towns]

Brilliant, thanks!