_globalize_single_replica_arrays fails with "incompatible devices" under active mesh context (JAX 0.9/Orbax 0.11.33)

[evelyn22chen]

Description

SingleReplicaArrayHandler checkpoint restore fails when a multi-device mesh is active via jax.set_mesh(). The jnp.expand_dims and jnp.zeros calls in _globalize_single_replica_arrays are single-device operations, but JAX 0.9's implicit mesh context causes them to target the full mesh.

Environment

• JAX 0.9, multi-host GPU (4 workers × 8 H200 GPUs = 32 devices)
• Orbax 0.11.33
• S3 checkpointing
• Mesh shape: (pp=1, dp=16, tp=2)
• Active mesh set via jax.set_mesh() in training loop

Error

File "orbax/checkpoint/_src/multihost/multislice.py", line 272, in _globalize_single_replica_arrays
    source_device_map[s.device] = jnp.expand_dims(s.data, axis=0)

ValueError: Received incompatible devices for jitted computation. Got argument args[0] of
broadcast_in_dim with shape float32[1] and device ids [0] on platform GPU and jit's context
mesh with device ids [0, 1, 2, ..., 31] on platform GPU

Root Cause

https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/_src/multihost/multislice.py#L250-L280

jnp.expand_dims(s.data, axis=0) and jnp.zeros(..., device=d) are per-device ops that pick up the active multi-device mesh context in JAX 0.9, causing a device mismatch.

Suggested Fix

Scope single-device ops to a temporary per-device mesh:

python
if is_source:
    for s in inp.addressable_shards:
        sd_mesh = jax.sharding.Mesh(np.array([s.device]), ('_single',))
        with jax.set_mesh(sd_mesh):
            source_device_map[s.device] = jnp.expand_dims(s.data, axis=0)

...

else:
    slice_shape = _get_slice_shape(index, global_shape)
    sd_mesh = jax.sharding.Mesh(np.array([d]), ('_single',))
    with jax.set_mesh(sd_mesh):
        zero_data = jnp.zeros(slice_shape, dtype=inp.dtype, device=d)
    device_buffers.append(zero_data)

Notes

• The save path had a similar "incompatible devices" issue under active mesh context, which was fixed in saving with default/global mesh is broken #2545. This is the same class of problem on the restore/broadcast side.

Questions

• Is this a known issue? Is scoping to a temporary per-device mesh via jax.set_mesh() the right fix here, or is there a preferred approach for isolating single-device ops from the active mesh context?

[lyglst]

Hi @evelyn22chen ,

I wonder how SingleReplicaArrayHandler is used in your current restoration setup? I'm trying to localize where the problem is.

Usually, it is more recommended to configure the sharding for each PyTree leave specifically, as suggested here. jax.set_mesh() might introduce some uncertainties during restoration, if the sharding is not configured properly for the PyTree leaf.

[evelyn22chen]

SingleReplicaArrayHandler is registered globally as the type handler for jax.Array:

def _register_single_replica_handler(replica_axis_index) -> None:
    array_handler = ocp.type_handlers.SingleReplicaArrayHandler(
        replica_axis_index=replica_axis_index,
        broadcast_memory_limit_bytes=1024 * 1024 * 1000,  # 1000 MB
    )
    ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True)

Sharding is configured per pytree leaf via SingleReplicaArrayRestoreArgs:

def _make_single_replica_restore_args(abstract_shape, sharding):
    mesh = sharding.mesh
    pspec = sharding.spec
    return ocp.type_handlers.SingleReplicaArrayRestoreArgs(
        sharding=jax.sharding.NamedSharding(mesh, pspec),
        global_shape=abstract_shape.shape,
        dtype=abstract_shape.dtype,
    )

def _apply_single_replica_restore_args(state, shardings, is_leaf=None):
    def _maybe_make_sr_args(abstract_shape, sharding):
        if hasattr(abstract_shape, "shape") and hasattr(abstract_shape, "dtype") and sharding is not None:
            return _make_single_replica_restore_args(abstract_shape, sharding)
        return abstract_shape

    return jax.tree.map(_maybe_make_sr_args, state, shardings, is_leaf=is_leaf)

This is called from both model and optimizer restore paths to build per-leaf restore args:

model restore path

if enable_single_replica_ckpt_restoring:
    state = _apply_single_replica_restore_args(state, shardings)
    return _stringify_int_keys(state.to_pure_dict())

The restore is then invoked via ocp.args.PyTreeRestore(restore_args). So each leaf has explicit sharding, global_shape, and dtype configured. The error occurs inside Orbax's _globalize_single_replica_arrays where
jnp.expand_dims(s.data, axis=0) operates on raw device buffers, which doesn't consult the per-leaf restore args and instead picks up the ambient mesh context.

[lyglst]

Hi @evelyn22chen

Feel free to patch #3177 and verify if it fixes your problem.

Meanwhile, thank you for reporting this issue to us!