Trouble saving checkpoint on a multi-gpu node

[jrabary]

Hi all,

I have a code based on the latest version of Flax/NNX and orbax that fails to save checkpoint on a multi-gpu node with the following error

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

Here is the code:

from pathlib import Path

import jax
import jax.numpy as jnp
import optax
import orbax.checkpoint as ocp
from flax import nnx
from jax.sharding import NamedSharding
from jax.sharding import PartitionSpec as P


class SimpleModel(nnx.Module):
    def __init__(self, rngs: nnx.Rngs):
        self.linear1 = nnx.Linear(128, 128, rngs=rngs)
        self.linear2 = nnx.Linear(128, 1, rngs=rngs)

    def __call__(self, x):
        x = self.linear1(x)
        x = jax.nn.relu(x)
        return self.linear2(x)

@nnx.jit
def train_step(model, optimizer, x, y):
    def loss_fn(model):
        pred = model(x)
        return jnp.mean((pred - y) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model)

    # This update triggers an all-reduce on grads because parameters are replicated
    optimizer.update(model, grads)
    return loss

def main():
    num_devices = jax.device_count()

    # Simple data parallelism: shards batch dim across devices
    mesh = jax.make_mesh((num_devices,), ("data",))
    print(f"Mesh created: {mesh}")

    ckpt_manager_options = ocp.CheckpointManagerOptions(max_to_keep=3, create=True)
    checkpoint_manager = ocp.CheckpointManager(
        Path("tmp_checkpoints").resolve(),
        options=ckpt_manager_options,
    )

    with jax.set_mesh(mesh):

        model = SimpleModel(rngs=nnx.Rngs(0))
       
        optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)
      
        total_batch_size = 32 * num_devices
        x_data = jax.random.normal(jax.random.PRNGKey(1), (total_batch_size, 128))
        y_data = jax.random.normal(jax.random.PRNGKey(2), (total_batch_size, 1))

        data_sharding = NamedSharding(mesh, P("data",))
        x_sharded = jax.device_put(x_data, data_sharding)
        y_sharded = jax.device_put(y_data, data_sharding)

        print("Starting training loop...")
        # Run a few steps
        for i in range(5):
            loss = train_step(model, optimizer, x_sharded, y_sharded)
            # Block until ready to catch async errors
            print(f"Step {i}, Loss: {loss}")

        checkpoint_manager.save(optimizer.step, args=ocp.args.Composite(
            state=ocp.args.PyTreeSave(nnx.state(model)),
        ))

if __name__ == "__main__":
    main()

Any ideas will be welcome.

Best

[ChromeHearts]

The error indicates some data only exists on device0.

Could you give us more details?

1. the full stack trace of the above mentioned
2. please help print the sharding being initiated before and after training.
   For example

def get_sharding(x):
  if hasattr(x, "sharding"):
    return x.sharding
  else:
    return x

...
model = SimpleModel(rngs=nnx.Rngs(0))

sharding_tree = jax.tree.map(print_sharding, nnx.state(model))
print("before training")
print(sharding_tree)
  
  ...

  for i in range(5):
    loss = train_step(model, optimizer, x_sharded, y_sharded)
    # Block until ready to catch async errors
    print(f"Step {i}, Loss: {loss}")
  
  sharding_tree = jax.tree.map(print_sharding, nnx.state(model))

  print("after training")
  print(sharding_tree)

  checkpoint_manager.save(optimizer.step, args=ocp.args.Composite(
            state=ocp.args.PyTreeSave(nnx.state(model)),
        ))