Supported sharded GPU simulation with mjwarp

[jjyyxx]

The feature, motivation and pitch

It would be very useful to support multi-GPU mjwarp simulations.

Problem

Running mjwarp simulations across multiple GPUs is currently difficult due to limitations in sharding and device handling.

Challenges

1. Device-side model/data creation: The current API is not compatible with replicated or data-parallel sharding.
2. Sharding awareness in JAX FFI integration: The mjx/third_party/warp/jax_experimental/ffi.py file (and other code locations) assumes the first visible device instead of handling different devices properly.
3. Collision data layout: mjwarp stores collision data in a flattened layout, making partitioning and sharding more complex.

Alternatives

No response

Additional context

No response

[btaba]

This is being worked on here google-deepmind/mujoco_warp#475

I'm debugging issues with 2, not 1/3 (pmap works fine with the current data layout).

[jjyyxx]

Thanks! Great to hear the progress. Though I believe the unified pjit API is more modern, user-friendly, and easier to work with.

[btaba]

Can you expand on point 1/3 in the OP? What do you mean by flattened layout? The collision data should simply be replicated for all devices.

[jjyyxx]

Here’s an expansion on points 1 and 3 from my original post.

---

Point 1: API Support for jax.sharding.Sharding

The core issue is that the mjx.put_model, mjx.make_data, and mjx.put_data function signatures don't explicitly support modern JAX sharding APIs, which is crucial for pjit-based workflows.

• Problem: The current API accepts device: Optional[jax.Device] = None. While this might silently pass a jax.sharding.Sharding object to the underlying jax.device_put, it's not documented, typed, or guaranteed to work.
• Proposal:
  1. Update API Signatures: Modify the type hints and docstrings to explicitly show that a jax.sharding.Sharding object is a valid input, similar to the standard jax.device_put API.
  2. Verify Internal Compatibility: Ensure that internal logic, such as the "Warm-up the warp kernel cache" code, is compatible with sharded model/data inputs.
• Scope of Sharding: We likely only need to support two specific sharding strategies:
  • Replicated Sharding for MjModel: The model structure is identical across all devices.
  • Data-Parallel Sharding for MjData: The data is partitioned along the batch (nworld) axis.

A typical setup would look like this:

devices = jax.devices()
mesh = jax.sharding.Mesh(devices, ('data',))

# MjModel should be fully replicated across all devices
replicated_sharding = jax.sharding.NamedSharding(mesh, jax.P())

# MjData should be sharded along the batch axis
data_sharding = jax.sharding.NamedSharding(mesh, jax.P('data'))

A design question remains whether MJX should enforce these specific shardings or leave validation to the user.

---

Point 3: Data Layout for Sharded Collisions

The "flattened layout" I mentioned refers to the inconsistent leading dimensions within the MjData struct, which complicates data-parallel sharding.

• Problem: My core assumption is that multi-GPU simulations should not require inter-device communication during a step. While simply replicating works for pmap, it's not the goal for pjit. With pjit, we want to partition data. However, MjData fields are not uniform.

  • Most fields, like geom_xpos, have a leading batch dimension of nworld.
  • Collision-related fields (e.g., contact) have a leading dimension of nconmax.
  • Several fields lack a batch dimension entirely, e.g., ncollision.

  This mismatch makes it difficult to apply a single, consistent data-parallel sharding strategy across the entire MjData object.

• Proposed Solutions:

  Option A: Partition nconmax Dimension
  Shard all MjData arrays along the batch axis. For collision data, this means partitioning the nconmax dimension.

  • Pros: The user-facing shape of MjData remains consistent with the single-device case.
  • Cons: This would require careful internal handling to adjust nconmax on a per-device basis to prevent out-of-bounds memory access.

  Option B: Add a device Axis (pmap-style)
  Add a new leading axis with size num_devices to all fields in MjData and shard along this new axis.

  • Pros: Potentially simpler to implement inside mjx, as the core mj_step logic would operate on a standard-looking MjData slice.
  • Cons: This approach is less idiomatic for pjit and pushes complexity to the user. They would need to manually add/remove this device axis, and handle the extra axis in their code.

Hope this clarifies the two points.