Issue 1:
Hello Jax team,

I have noticed that jax.distributed currently only supports a subset of GPU cores, and I was wondering if it would be possible to extend this support to a subset of TPU cores as well.

I have tried running a code snippet with jax.pmap on TPUs using the suggested snippet from @skye in #8809 (comment), but it only recognizes the local TPU device and not other devices, resulting in a failed psum operation.

For comparison, I have also included the output of the same code snippet when run on GPUs for reference.

import jax
import jax.numpy as jnp

jax.distributed.initialize()

print("process_count", jax.process_count())
print("process_index", jax.process_index())
print("local_devices", jax.local_devices())
print("local_device_count", jax.local_device_count())
print("devices", jax.devices())

def my_function(x):
return jax.lax.psum(x, 'i')

x = jnp.ones((1, 2))
print(x)
parallel_fn = jax.pmap(my_function, devices=jax.devices(), axis_name='i')
result = parallel_fn(x)
print(result)

On GPUs, when running with

CUDA_VISIBLE_DEVICES="0" SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=0 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py
CUDA_VISIBLE_DEVICES="1" SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=1 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py

produces

process_count 2
process_index 1
local_devices [StreamExecutorGpuDevice(id=1, process_index=1, slice_index=0)]
local_device_count 1
devices [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=1, slice_index=0)]
[[1. 1.]]
[[2. 2.]]

on TPUs, when running with the snippet suggested by @skye at #8809 (comment)

TPU_CHIPS_PER_HOST_BOUNDS=1,2,1 TPU_HOST_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=0 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=0 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py
TPU_CHIPS_PER_HOST_BOUNDS=1,2,1 TPU_HOST_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=1 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=1 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py

we get the following.

process_count 1
process_index 0
local_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]
local_device_count 1
devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]
[[1. 1.]]
[[1. 1.]]

Note that the TPU example does not recognize other devices, thus making the psum fails.

Use cases

We are using jax.distributed this way in our new distributed RL platform Cleanba, which supports different ways to leverage accelerators. Specifically, the users can only use 1 accelerator core for the actor, multiple accelerator cores for the learner, and replicate the training process through jax.distributed. Below is a figure showing our results

The format a0_l1 2 3 4_d1 means the actor runs on GPU 0 (a0), the learner runs on GPU 1 2 3 4 (l1 2 3 4), and the computation is replicated/distributed one time (d1). The blue baseline curve corresponds to a0_l1 2 3_d2.

TPUv4 is fast, and the results on a0_l1 2_d1 is pretty great. However, because of the problem listed in this issue, I couldn't test out the a0_l1_d2 to fully utilize 4 chips on TPUv4. If a0_l1_d2 were to work, I can expect it to be 10-20% faster than our current fastest 8 A100 setting, which would be very exciting 😬

Check for duplicate requests.

Describe your goal, and if possible provide a code snippet with a motivating example.

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

import jax.numpy as jnp
print(jnp.arange(10))
# [0 1 2 3 4 5 6 7 8 9]

If applicable, include full error messages/tracebacks.