JAX multi-host data gather not working

[AshishKumar4]

I have been trying to build a multi-host training pipeline using jax and flax for training on TPU-v4-16/32+
I firstly tried with pmap, and then with shard_map, but both times, on the arrays that are returned from the jitted pmapped/shmapped functions I am unable to do any operation. The mere act of doing something like this:

def something(x, state):
   loss = loss_fn(x, state)
   return loss

mapped = shard_map(something, in_spec=(P('data'), P()), out_spec=P('data'))
loss = mapped(some_data, some_state)
loss = jax.experimental.multihost_utils.process_allgather(loss)  ----> THROWS ERROR

Now, I actually did make some example code and on this its working-->

import jax
import jax.numpy as jnp

from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
from jax.experimental.shard_map import shard_map
from jax.experimental.multihost_utils import process_allgather
import flax
import jax.tree_util as jtu
import numpy as np
from functools import partial
from flax import struct          
from flax.training import train_state
import optax

jax.distributed.initialize()

class RandomMarkovState(struct.PyTreeNode):
    rng: jax.random.PRNGKey
    def get_random_key(self):
        rng, subkey = jax.random.split(self.rng)
        return RandomMarkovState(rng), subkey

# Define the TrainState
class TrainState(train_state.TrainState):
    pass

data = jnp.arange(16).reshape((16, 1))

def func(x, rngs, state:TrainState):
    # rngs = RandomMarkovState(rngs.rng.reshape((2,)))
    print(rngs, state)
    rngs, subkey = rngs.get_random_key()
    subkey = jax.random.fold_in(subkey, 1)
    return jax.lax.psum(x + jax.random.normal(subkey, x.shape), 'i'), rngs, state# RandomMarkovState(rngs.rng.reshape((1, 2))), state

rngs = RandomMarkovState(jax.random.PRNGKey(0))
rngs_mapped = flax.jax_utils.replicate(rngs, jax.local_devices())

sample_params = {'params': {'dense': {'kernel': jnp.ones((1, 1)), 'bias': jnp.zeros((1,))}}}
print("Sample params:", sample_params)

state = TrainState.create(
    apply_fn=lambda : None,
    params=sample_params,
    tx=optax.adam(1e-3),
)
state_mapped = flax.jax_utils.replicate(state, jax.local_devices())

############################################################################
# Pmap
############################################################################

out1 = jax.pmap(func, axis_name='i')(data.reshape((jax.local_device_count(), -1)), rngs_mapped, state_mapped)
print("Output pmap:", out1, jax.debug.visualize_array_sharding(out1[0]))

############################################################################
# Shard map
############################################################################

def _build_global_shape_and_sharding(
    local_shape: tuple[int, ...], global_mesh: Mesh
) -> tuple[tuple[int, ...], NamedSharding]:
  sharding = NamedSharding(global_mesh, P(global_mesh.axis_names))
  global_shape = (jax.process_count() * local_shape[0],) + local_shape[1:]
  return global_shape, sharding

def form_global_array(path, array: np.ndarray, global_mesh: Mesh) -> jax.Array:
  """Put local sharded array into local devices"""
  global_shape, sharding = _build_global_shape_and_sharding(np.shape(array), global_mesh)
  try:
    local_device_arrays = np.split(array, len(global_mesh.local_devices), axis=0)
  except ValueError as array_split_error:
    raise ValueError(
        f"Unable to put to devices shape {array.shape} with "
        f"local device count {len(global_mesh.local_devices)} "
        f"at {jtu.keystr(path)}"
    ) from array_split_error
  local_device_buffers = jax.device_put(local_device_arrays, global_mesh.local_devices)
  return jax.make_array_from_single_device_arrays(global_shape, sharding, local_device_buffers)

def convert_to_global_tree(global_mesh, pytree):
    return jax.tree_util.tree_map_with_path(partial(form_global_array, global_mesh=global_mesh), pytree)

mesh = jax.sharding.Mesh(jax.devices(), 'i')
sharding = NamedSharding(mesh, P('i'))
rngs_mapped = convert_to_global_tree(mesh, rngs_mapped)
state_mapped = convert_to_global_tree(mesh, state_mapped)
out2 = jax.jit(shard_map(func, mesh=mesh, in_specs=(P('i'), P(), P()), out_specs=(P('i'), P(), P())))(data, rngs, state)
print("Output Shard map:", jax.experimental.multihost_utils.process_allgather(out2[0]), jax.debug.visualize_array_sharding(out2[0]))

The output is printed for both pmap case and shard_map cases without issues.
However when I try to actually do the same in training code, it fails :(
I even commented out all the actual training and returned hard coded values and that failed too. Below is the screenshot for the training code that I have

The fit function:

Mesh is simply defined as:
self.mesh = jax.sharding.Mesh(jax.devices(), 'data')

The training code works for single host TPU-v4s, but doesn't for multi. And it happens if I switch to pmap as well.

The error:

processing batch
Got the batch, going to train
Got till here
		Epoch 1:   0%|                                                          | 0/127 [00:13<?, ?step/s]
Traceback (most recent call last):
  File "/home/mrwhite0racle/training.py", line 1093, in <module>
    main(args)
  File "/home/mrwhite0racle/training.py", line 1089, in main
    final_state = trainer.fit(data, batches, epochs=CONFIG['epochs'])
  File "/home/mrwhite0racle/training.py", line 874, in fit
    super().fit(data, steps_per_epoch, epochs, {
  File "/home/mrwhite0racle/training.py", line 654, in fit
    loss = jax.experimental.multihost_utils.process_allgather(loss)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/experimental/multihost_utils.py", line 154, in process_allgather
    return jax.tree.map(_pjit, in_tree)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/tree.py", line 155, in map
    return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/tree_util.py", line 343, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/tree_util.py", line 343, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/experimental/multihost_utils.py", line 153, in _pjit
    return _handle_array_process_allgather(inp, tiled)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/experimental/multihost_utils.py", line 130, in _handle_array_process_allgather
    return np.asarray(out.addressable_data(0))
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/array.py", line 414, in __array__
    return np.asarray(self._value, dtype=dtype, **kwds)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 336, in wrapper
    return func(*args, **kwargs)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/array.py", line 613, in _value
    self._npy_value = self._single_device_array_to_np_array()
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:1:0x177 (from TensorCoreSequencer:1:0x221): scheckne: 
***************
An unexpected peer shows up in the launch group with a different launch id than the current group leader. This signals a bug in the runtime scheduling system or above. no HLO mapping
=== Source Location Trace: === 
learning/45eac/tpu/runtime/hal/internal/tpu_program_termination_validation.cc:113

Traceback (most recent call last):
  File "/home/mrwhite0racle/training.py", line 1093, in <module>
    main(args)
  File "/home/mrwhite0racle/training.py", line 1089, in main
    final_state = trainer.fit(data, batches, epochs=CONFIG['epochs'])
  File "/home/mrwhite0racle/training.py", line 874, in fit
    super().fit(data, steps_per_epoch, epochs, {
  File "/home/mrwhite0racle/training.py", line 654, in fit
    loss = jax.experimental.multihost_utils.process_allgather(loss)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/experimental/multihost_utils.py", line 154, in process_allgather
    return jax.tree.map(_pjit, in_tree)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/tree.py", line 155, in map
    return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/tree_util.py", line 343, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/tree_util.py", line 343, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/experimental/multihost_utils.py", line 153, in _pjit
    return _handle_array_process_allgather(inp, tiled)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/experimental/multihost_utils.py", line 130, in _handle_array_process_allgather
    return np.asarray(out.addressable_data(0))
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/array.py", line 414, in __array__
    return np.asarray(self._value, dtype=dtype, **kwds)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 336, in wrapper
    return func(*args, **kwargs)
  File "/home/mrwhite0racle/.local/lib/python3.10/site-packages/jax/_src/array.py", line 613, in _value
    self._npy_value = self._single_device_array_to_np_array()
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:1:0x177 (from TensorCoreSequencer:1:0x221): scheckne: 
***************
An unexpected peer shows up in the launch group with a different launch id than the current group leader. This signals a bug in the runtime scheduling system or above. no HLO mapping
=== Source Location Trace: === 
learning/45eac/tpu/runtime/hal/internal/tpu_program_termination_validation.cc:113

What am I doing wrong? I lost a lot of hairs on this :(

[yashk2810]

Can you please paste what the error is?

[AshishKumar4]

I am extremely sorry, totally forgot putting the error here. I have added it

[AshishKumar4]

I think I might have figured out the issue. I was doing

subkey = jax.random.fold_in(subkey, jax.process_index())

as a way to make sure every process gets a unique random key. Turns out, doing this process_index() operation inside of the shard_mapped jitted function caused this issue. Removing this and instead supplying global indexes solved the issue, as follows:

global_device_indexes = jnp.arange(jax.device_count())
.
.
.
<pass the global device indexes array as argument to the shmapped train function>
.
.
[inside the jax shmapped jitted function]
subkey = jax.random.fold_in(subkey, local_device_index.reshape())

Finally figured this out!