Update 26.01.18

  - 9e20ec9dc2b556db2faba1990bf6579da80bf6ea by Anxhelo Xhebraj <axhebraj@nvidia.com>
  - de11bd8aa863acbab840e33368ffdb199893e2e4 by Anxhelo Xhebraj <axhebraj@nvidia.com>
  - 2d398e9802ab2ad6016e385c9a4f18c86f6381d0 by Sean Lee <selee@nvidia.com>
  - 960973450466b0a29b8c40e56799bc8f34e76025 by Anxhelo Xhebraj <axhebraj@nvidia.com>
  - 1d19d8f7a05de81c2efd3214d4f050421dbaec2c by Anxhelo Xhebraj <axhebraj@nvidia.com>
  - aaec562a093047ad2847b576ed1113587ae77f17 by Anxhelo Xhebraj <axhebraj@nvidia.com>
  - 7164423d1ea990ea482612631042f88569ca857e by Anxhelo Xhebraj <axhebraj@nvidia.com>
  - 00f246f52cbfb8d3086f995b8d3dfb7b4c1d861e by Anxhelo Xhebraj <axhebraj@nvidia.com>
  - c360203d29269207ced92451746b5f3fe08d5250 by Anxhelo Xhebraj <axhebraj@nvidia.com>

Co-authored-by: Anxhelo Xhebraj <axhebraj@nvidia.com>
Signed-off-by: Sean Lee <selee@nvidia.com>
GitOrigin-RevId: 9e20ec9dc2b556db2faba1990bf6579da80bf6ea

Author: seanprime7

pyproject.toml

@@ -99,7 +99,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "jax[cuda12]>=0.5.1,<=0.8.1",
+    "jax[cuda12]>=0.6.2,<=0.8.2",
     "cupy-cuda12x",
     "portpicker==1.6.0",
 ]

scripts/local_mc.sh

@@ -15,14 +15,13 @@
 
 set -u
 set -o pipefail
+set -m
 
 if [ -z "${N_PROCS:-}" ] || [ -z "${N_GPUS:-}" ] || [ -z "${COMMAND:-}" ]; then
   echo "N_PROCS, N_GPUS, and COMMAND must be set"
   exit 1
 fi
 
-
-
 # Default coordinator setup
 export JAX_COORDINATOR_IP="${JAX_COORDINATOR_IP:-localhost}"
 export JAX_COORDINATOR_PORT="${JAX_COORDINATOR_PORT:-1234}"
@@ -35,9 +34,7 @@ PIDS=()
 cleanup() {
   echo "Cleaning up..."
   for pid in "${PIDS[@]}"; do
-    if kill -0 "$pid" 2>/dev/null; then
-      kill "$pid" 2>/dev/null
-    fi
+    kill -9 -- -"$pid" 2>/dev/null || true
   done
 }
 

scripts/test_jax_versions.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e
 
-JAX_VERSIONS=("0.6.1" "0.6.2" "0.7.0" "0.7.1" "0.7.2" "0.8.0" "0.8.1")
+JAX_VERSIONS=("0.6.2" "0.7.0" "0.7.1" "0.7.2" "0.8.0" "0.8.1" "0.8.2")
 
 for version in "${JAX_VERSIONS[@]}"
 do
@@ -13,6 +13,6 @@ done
 
 # Test nightly
 pip install -U --pre jax jaxlib "jax-cuda12-plugin[with-cuda]" jax-cuda12-pjrt -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
-# JAXPP_DEBUG_FORCE_MPMDIFY=True JAXPP_ENABLE_LICM=True python examples/basic.py --train_steps=10
+JAXPP_DEBUG_FORCE_MPMDIFY=True JAXPP_ENABLE_LICM=True python examples/basic.py --train_steps=10
 
 exit 0

src/jaxpp/api.py

@@ -23,8 +23,8 @@
     mpmd_to_spmd_reshard,
     spmd_to_mpmd_reshard,
 )
-from jaxpp.core import mpmd_jit_by_yield, mpmd_jit_rev, mpmd_jit_with_loop
-from jaxpp.jax_primitives import add_multi_p
+from jaxpp.core import mpmd_jit_by_yield, mpmd_jit_with_loop
+from jaxpp.jax_primitives import add_multi_p, gather_multi_p
 from jaxpp.mesh import MpmdMesh
 from jaxpp.pipelining import pipeline_enter_stage
 from jaxpp.schedules import (
@@ -45,3 +45,10 @@ def cross_mpmd_all_reduce(*args):
             f"All arguments must have the same dtype, got {[a.dtype for a in args]}"
         )
     return add_multi_p.bind(*args)
+
+
+def cross_mpmd_stack(arrays, axis: int = 0):
+    import jax
+
+    expanded = [jax.numpy.expand_dims(a, axis=axis) for a in arrays]
+    return gather_multi_p.bind(*expanded, axis=axis)

src/jaxpp/array.py

@@ -14,15 +14,15 @@
 # limitations under the License.
 
 from collections import OrderedDict, defaultdict
-from typing import Any, cast, overload
+from typing import Any, cast
 
 import jax
-import jax._src.core as jcore
 import numpy as np
 
+from jaxpp.jax_compat import core as jcore
 from jaxpp.mesh import MpmdMesh
 from jaxpp.types import MpmdSharding
-from jaxpp.utils import get_named_sharding, update_named_sharding
+from jaxpp.utils import filter_axes, get_named_sharding, update_named_sharding
 
 
 class MpmdArray:
@@ -83,8 +83,8 @@ def __init__(
             mesh = get_named_sharding(arr).mesh
             if (mpmd_idx := mpmd_mesh.mpmd_idx_for_mesh.get(mesh)) is None:
                 raise ValueError(
-                    f"Argument array {idx} {arr.shape} is not on a mesh that is part "
-                    f"mpmd_mesh={mpmd_mesh.jax_mesh}"
+                    f"Argument array {idx} {arr.shape} is not on a mesh that is part"
+                    f" of mpmd_mesh={mpmd_mesh.jax_mesh}"
                 )
 
             if mpmd_idx not in mpmd_idxs:
@@ -118,8 +118,9 @@ def __init__(
             assert all(_ == shape for _ in shapes), (shape, shapes)
             dtypes = [a.dtype for a in self._partially_addressable_arrays.values()]
             assert all(_ == dtype for _ in dtypes), (dtype, dtypes)
+            mpmd_axis = mpmd_sharding.mpmd_mesh.mpmd_axis_name
             specs = [
-                get_named_sharding(a).spec
+                filter_axes(get_named_sharding(a).spec, {mpmd_axis})
                 for a in self._partially_addressable_arrays.values()
             ]
             assert all(_ == mpmd_sharding.spec for _ in specs), (
@@ -281,56 +282,6 @@ def _to_global_jax_array(mpmd_array: MpmdArray) -> jax.Array | None:
     )
 
 
-@overload
-def filter_axes(
-    sharding_or_pspec: jax.sharding.NamedSharding, axes: set[str]
-) -> jax.sharding.NamedSharding: ...
-
-
-@overload
-def filter_axes(
-    sharding_or_pspec: jax.sharding.PartitionSpec, axes: set[str]
-) -> jax.sharding.PartitionSpec: ...
-
-
-def filter_axes(
-    sharding_or_pspec: jax.sharding.NamedSharding | jax.sharding.PartitionSpec,
-    axes: set[str],
-) -> jax.sharding.NamedSharding | jax.sharding.PartitionSpec:
-    """Filter out specified axes from a sharding or partition spec.
-
-    Args:
-        sharding_or_pspec: Either a NamedSharding or PartitionSpec to filter.
-        axes: Set of axis names to remove from the spec.
-
-    Returns:
-        Same type as input with specified axes filtered out.
-    """
-    if isinstance(sharding_or_pspec, jax.sharding.NamedSharding):
-        return jax.sharding.NamedSharding(
-            sharding_or_pspec.mesh, filter_axes(sharding_or_pspec.spec, axes)
-        )
-
-    assert isinstance(sharding_or_pspec, jax.sharding.PartitionSpec)
-
-    new_spec = []
-    for axis in sharding_or_pspec:
-        if axis is None:
-            new_spec.append(None)
-        elif isinstance(axis, str):
-            if axis not in axes:
-                new_spec.append(axis)
-            else:
-                new_spec.append(None)
-        elif isinstance(axis, (list, tuple)):
-            new_axis = [a for a in axis if a not in axes]
-            new_spec.append(type(axis)(new_axis))
-        else:
-            raise ValueError(f"Unsupported_axis_type: {type(axis)}")
-
-    return jax.sharding.PartitionSpec(*new_spec)
-
-
 def _id(*xs):
     return xs
 
@@ -346,8 +297,9 @@ def _spmd_to_mpmd_reshard(
 
     for spmd_value, dist_sharding in zip(spmd_values, dist_shardings):
         assert isinstance(
-            spmd_value.sharding, jax.sharding.NamedSharding
-        ), spmd_value.sharding
+            spmd_value.sharding,
+            (jax.sharding.NamedSharding, jax.sharding.SingleDeviceSharding),
+        ), f"Unsupported sharding type: {spmd_value.sharding}"
 
     # NOTE: We filter out the mpmd axis from the sharding so that
     # the output is replicated across all mpmd ranks.
@@ -362,11 +314,7 @@ def _spmd_to_mpmd_reshard(
         for dist_sharding in dist_shardings
     )
 
-    res: list[jax.Array] = jax.jit(
-        _id,
-        in_shardings=tuple(_.sharding for _ in spmd_values),
-        out_shardings=_actual_shardings,
-    )(*spmd_values)
+    res: list[jax.Array] = jax.jit(_id, out_shardings=_actual_shardings)(*spmd_values)
 
     for spmd_value, donated in zip(spmd_values, donate, strict=True):
         if donated:
@@ -398,17 +346,19 @@ def _spmd_to_mpmd_reshard(
     _res = []
     for arr, dsh in zip(res, dist_shardings, strict=True):
         mesh_ids = dsh.mesh_ids
-        filtered_sharding = MpmdSharding(
+        # MpmdSharding.__post_init__ canonicalizes the spec by filtering out
+        # the mpmd axis, so we can just use dsh.spec directly
+        mpmd_sharding = MpmdSharding(
             mpmd_mesh=dsh.mpmd_mesh,
             mesh_ids=dsh.mesh_ids,
-            spec=filter_axes(dsh.sharding.spec, {mpmd_mesh.mpmd_axis_name}),
+            spec=dsh.spec,
         )
 
         if mpmd_mesh.my_mpmd_axis_index not in mesh_ids:
             _res.append(
                 MpmdArray(
                     partially_addressable_arrays=[],
-                    mpmd_sharding=filtered_sharding,
+                    mpmd_sharding=mpmd_sharding,
                     shape=arr.shape,
                     dtype=arr.dtype,
                 )
@@ -425,7 +375,7 @@ def _spmd_to_mpmd_reshard(
             _res.append(
                 MpmdArray(
                     partially_addressable_arrays=[new_arr],
-                    mpmd_sharding=filtered_sharding,
+                    mpmd_sharding=mpmd_sharding,
                 )
             )
     return _res
@@ -481,6 +431,9 @@ def spmd_to_mpmd_reshard(
     The specs of the returned arrays will _not_ have `mpmd_mesh.mpmd_axis_name` in
     them.
 
+    Limitations: same constraints as jax.jit apply (e.g. _device_assignment must be the
+    same for all arrays)
+
     Args:
         mpmd_mesh: The MPMD mesh definition.
         spmd_arrays: A pytree of source SPMD arrays.
@@ -506,28 +459,6 @@ def spmd_to_mpmd_reshard(
 
     assert spmd_tree_def == mpmd_tree_def
 
-    # Verify all arrays are on the same mesh
-    first_path, first_leaf = spmd_arrays_with_path[0]
-    first_sharding = first_leaf.sharding
-    assert isinstance(first_sharding, jax.sharding.NamedSharding), first_sharding
-    mesh = first_sharding.mesh
-
-    # This check is the same as the one performed by jax.jit
-    assert mesh._flat_devices_tuple == mpmd_mesh.jax_mesh._flat_devices_tuple, (
-        mesh,
-        mpmd_mesh.jax_mesh,
-    )
-    for path, leaf in spmd_arrays_with_path:
-        assert isinstance(leaf.sharding, jax.sharding.NamedSharding), (
-            path,
-            leaf.sharding,
-        )
-        assert leaf.sharding.mesh._flat_devices_tuple == mesh._flat_devices_tuple, (
-            path,
-            mesh,
-            leaf.sharding.mesh,
-        )
-
     _, spmd_arrays_flat = jax._src.util.unzip2(spmd_arrays_with_path)
     spmd_arrays_flat_list = list(spmd_arrays_flat)
 
@@ -605,7 +536,11 @@ def _axis_name_in_spec(axis_name: str, spec) -> bool:
 
 
 def logically_stacked(
-    array: jax.Array, comm_mesh: jax.sharding.Mesh, axis_name: str, strict: bool = False
+    array: jax.Array,
+    comm_mesh: jax.sharding.Mesh,
+    mesh_axis_name: str,
+    array_axis: int = 0,
+    strict: bool = False,
 ):
     """
     Logically stacks an array along a new axis corresponding to the MPMD dimension.
@@ -619,18 +554,25 @@ def logically_stacked(
     if strict:
         spec = array.sharding.spec
         assert not _axis_name_in_spec(
-            axis_name, spec
-        ), f"axis_name {axis_name!r} already exists in spec {spec}"
+            mesh_axis_name, spec
+        ), f"axis_name {mesh_axis_name!r} already exists in spec {spec}"
     else:
-        spec = filter_axes(array.sharding.spec, {axis_name})
+        spec = filter_axes(array.sharding.spec, {mesh_axis_name})
 
-    expanded_array = jax.numpy.expand_dims(array, 0)
+    expanded_array = jax.numpy.expand_dims(array, array_axis)
     in_sharding = jax.sharding.NamedSharding(
-        comm_mesh, jax.sharding.PartitionSpec(axis_name, *spec)
+        comm_mesh,
+        jax.sharding.PartitionSpec(
+            *spec[:array_axis], mesh_axis_name, *spec[array_axis:]
+        ),
     )
 
     global_array = jax.make_array_from_single_device_arrays(
-        (comm_mesh.shape[axis_name], *array.shape),
+        (
+            *array.shape[:array_axis],
+            comm_mesh.shape[mesh_axis_name],
+            *array.shape[array_axis:],
+        ),
         in_sharding,
         [s.data for s in expanded_array.addressable_shards],
     )
@@ -664,6 +606,12 @@ def mpmd_to_spmd_reshard(
     Returns:
         A pytree of JAX arrays with the same structure as mpmd_arrays.
     """
+
+    if not mpmd_mesh.jax_mesh.is_multi_process:
+        return jax.device_put(
+            jax.tree.map(lambda _: _.first_mpmd_replica, mpmd_arrays), spmd_shardings
+        )
+
     mpmd_arrays_with_path, mpmd_tree_def = jax.tree.flatten_with_path(mpmd_arrays)
     mpmd_arrays_with_path: list[tuple[Any, MpmdArray]]
     spmd_shardings_flat, spmd_tree_def = jax.tree.flatten(spmd_shardings)