[JAX] Fix GEMM partitioning to reduce over nested contracting mesh axes

tests/jax/test_gemm_partitioning.py

@@ -0,0 +1,315 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Unit tests for the GEMM custom-partitioning spec inference.
+
+These tests exercise ``GemmPrimitive._parse_operand_output_specs`` directly on
+sharding specs, covering the FWD/DGRAD/WGRAD GEMMs of an MoE FFN block. They run
+on CPU and do not require GPUs.
+
+Cases come in two self-consistent parallelism configs where every backward output
+sharding matches its forward input (dX~X, dW~W, dH~H):
+
+  Megatron TP (tp shards feature dims, token dim is (fsdp, expert)):
+    X=((fsdp,expert),None), W1=(None,tp), H=Y1=((fsdp,expert),tp), W2=(tp,None).
+    - FWD/DGRAD contracting the tp-sharded ffn dim reduce over tp.
+    - WGRAD contracting the token dim reduces over (fsdp, expert).
+
+  TPSP (tpsp shards the token dim, weights replicated):
+    X=H=Y=((fsdp,tpsp,expert),None), W1=W2=(None,None).
+    - FWD/DGRAD contract replicated dims, no reduction.
+    - WGRAD contracting the nested token dim reduces over (fsdp, tpsp, expert).
+"""
+
+import os
+
+if "xla_force_host_platform_device_count" not in os.environ.get("XLA_FLAGS", ""):
+    os.environ["XLA_FLAGS"] = (
+        os.environ.get("XLA_FLAGS", "") + " --xla_force_host_platform_device_count=8"
+    ).strip()
+
+from collections import namedtuple
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import jax
+from jax.sharding import Mesh, PartitionSpec
+
+from transformer_engine.jax.cpp_extensions.gemm import CollectiveOp, GemmPrimitive
+from transformer_engine.jax.quantize import ScalingMode
+from transformer_engine.jax.sharding import MeshResource, global_shard_guard
+
+pytestmark = pytest.mark.skipif(len(jax.devices("cpu")) < 8, reason="requires 8 CPU devices")
+
+
+def _mesh(axes):
+    """Build a CPU mesh with a size-2 device grid over the named axes."""
+    devices = np.asarray(jax.devices("cpu")[: 2 ** len(axes)]).reshape((2,) * len(axes))
+    return Mesh(devices, axes)
+
+
+def _arg_info(shape, spec):
+    """Minimal arg_info stub exposing ndim, size and sharding.spec."""
+    sharding = None if spec is None else SimpleNamespace(spec=PartitionSpec(*spec))
+    return SimpleNamespace(ndim=len(shape), size=int(np.prod(shape)), sharding=sharding)
+
+
+def _parse(lhs_shape, lhs_spec, rhs_shape, rhs_spec, contracting_dims):
+    """Run the partition spec inference for a plain (non-collective) GEMM."""
+    scalar = _arg_info((0,), None)  # scales / bias / alpha / beta (unused for NO_SCALING)
+    arg_infos = (
+        _arg_info(lhs_shape, lhs_spec),
+        scalar,
+        _arg_info(rhs_shape, rhs_spec),
+        scalar,
+        scalar,
+        scalar,
+        scalar,
+    )
+    operand_specs, out_specs, reduce_spec, _ = GemmPrimitive._parse_operand_output_specs(
+        arg_infos,
+        contracting_dims,
+        transpose_batch_sequence=False,
+        collective_op=CollectiveOp.NONE,
+        scaling_mode=ScalingMode.NO_SCALING,
+    )
+    lhs_specs, _, rhs_specs, *_ = operand_specs
+    return lhs_specs, rhs_specs, tuple(out_specs), reduce_spec
+
+
+# Representative MoE FFN GEMMs (FFN1: hidden->ffn, FFN2: ffn->hidden) for the two
+# parallelism configs described in the module docstring. Expected: reduce over the mesh
+# axes that shard the contracting dim of both operands, gather any axis that shards it on
+# only one operand.
+Case = namedtuple(
+    "Case",
+    "axes, resource, lhs_shape, lhs_spec, rhs_shape, rhs_spec, cdims,"
+    " exp_lhs, exp_rhs, exp_out, exp_reduce",
+)
+
+_MR_TP = dict(fsdp_resource="fsdp", tp_resource="tp", ep_resource="expert")
+_MR_TPSP = dict(fsdp_resource="fsdp", tpsp_resource="tpsp", ep_resource="expert")
+
+_TP_AXES = ("fsdp", "tp", "expert")
+_TPSP_AXES = ("fsdp", "tpsp", "expert")
+
+CASES = {
+    # --- Megatron TP: tp shards feature dims, token dim is (fsdp, expert). ---
+    # FFN1 FWD: Y1 = X @ W1, contract the replicated hidden dim -> no reduction.
+    "megatron_ffn1_fwd": Case(
+        _TP_AXES,
+        _MR_TP,
+        (524288, 7168),
+        (("fsdp", "expert"), None),
+        (7168, 2048),
+        (None, "tp"),
+        ((1,), (0,)),
+        (("fsdp", "expert"), None),
+        (None, "tp"),
+        (("fsdp", "expert"), "tp"),
+        None,
+    ),
+    # FFN1 DGRAD: dX = dY1 @ W1^T, contract the tp-sharded ffn dim -> reduce over tp.
+    "megatron_ffn1_dgrad": Case(
+        _TP_AXES,
+        _MR_TP,
+        (524288, 2048),
+        (("fsdp", "expert"), "tp"),
+        (7168, 2048),
+        (None, "tp"),
+        ((1,), (1,)),
+        (("fsdp", "expert"), "tp"),
+        (None, "tp"),
+        (("fsdp", "expert"), None),
+        "tp",
+    ),
+    # FFN1 WGRAD: dW1 = X^T @ dY1, contract the token dim -> reduce over (fsdp, expert).
+    "megatron_ffn1_wgrad": Case(
+        _TP_AXES,
+        _MR_TP,
+        (7168, 524288),
+        (None, ("fsdp", "expert")),
+        (2048, 524288),
+        ("tp", ("fsdp", "expert")),
+        ((1,), (1,)),
+        (None, ("fsdp", "expert")),
+        ("tp", ("fsdp", "expert")),
+        (None, "tp"),
+        ("fsdp", "expert"),
+    ),
+    # FFN2 FWD: Y2 = H @ W2, contract the tp-sharded ffn dim -> reduce over tp.
+    "megatron_ffn2_fwd": Case(
+        _TP_AXES,
+        _MR_TP,
+        (524288, 2048),
+        (("fsdp", "expert"), "tp"),
+        (2048, 7168),
+        ("tp", None),
+        ((1,), (0,)),
+        (("fsdp", "expert"), "tp"),
+        ("tp", None),
+        (("fsdp", "expert"), None),
+        "tp",
+    ),
+    # FFN2 DGRAD: dH = dY2 @ W2^T, contract the replicated hidden dim -> no reduction.
+    "megatron_ffn2_dgrad": Case(
+        _TP_AXES,
+        _MR_TP,
+        (524288, 7168),
+        (("fsdp", "expert"), None),
+        (2048, 7168),
+        ("tp", None),
+        ((1,), (1,)),
+        (("fsdp", "expert"), None),
+        ("tp", None),
+        (("fsdp", "expert"), "tp"),
+        None,
+    ),
+    # FFN2 WGRAD: dW2 = H^T @ dY2, contract the token dim -> reduce over (fsdp, expert).
+    "megatron_ffn2_wgrad": Case(
+        _TP_AXES,
+        _MR_TP,
+        (2048, 524288),
+        ("tp", ("fsdp", "expert")),
+        (7168, 524288),
+        (None, ("fsdp", "expert")),
+        ((1,), (1,)),
+        ("tp", ("fsdp", "expert")),
+        (None, ("fsdp", "expert")),
+        ("tp", None),
+        ("fsdp", "expert"),
+    ),
+    # --- Sequence-parallel TP: tpsp shards the token dim, weights replicated. ---
+    # FFN1 FWD: Y1 = X @ W1, contract the replicated hidden dim -> no reduction.
+    "tpsp_ffn1_fwd": Case(
+        _TPSP_AXES,
+        _MR_TPSP,
+        (524288, 7168),
+        (("fsdp", "tpsp", "expert"), None),
+        (7168, 2048),
+        (None, None),
+        ((1,), (0,)),
+        (("fsdp", "tpsp", "expert"), None),
+        (None, None),
+        (("fsdp", "tpsp", "expert"), None),
+        None,
+    ),
+    # FFN1 DGRAD: dX = dY1 @ W1^T, contract the replicated ffn dim -> no reduction.
+    "tpsp_ffn1_dgrad": Case(
+        _TPSP_AXES,
+        _MR_TPSP,
+        (524288, 2048),
+        (("fsdp", "tpsp", "expert"), None),
+        (7168, 2048),
+        (None, None),
+        ((1,), (1,)),
+        (("fsdp", "tpsp", "expert"), None),
+        (None, None),
+        (("fsdp", "tpsp", "expert"), None),
+        None,
+    ),
+    # FFN1 WGRAD: dW1 = X^T @ dY1, contract the nested token dim -> reduce (fsdp, tpsp, expert).
+    "tpsp_ffn1_wgrad": Case(
+        _TPSP_AXES,
+        _MR_TPSP,
+        (7168, 524288),
+        (None, ("fsdp", "tpsp", "expert")),
+        (2048, 524288),
+        (None, ("fsdp", "tpsp", "expert")),
+        ((1,), (1,)),
+        (None, ("fsdp", "tpsp", "expert")),
+        (None, ("fsdp", "tpsp", "expert")),
+        (None, None),
+        ("fsdp", "tpsp", "expert"),
+    ),
+    # FFN2 FWD: Y2 = H @ W2, contract the replicated ffn dim -> no reduction.
+    "tpsp_ffn2_fwd": Case(
+        _TPSP_AXES,
+        _MR_TPSP,
+        (524288, 2048),
+        (("fsdp", "tpsp", "expert"), None),
+        (2048, 7168),
+        (None, None),
+        ((1,), (0,)),
+        (("fsdp", "tpsp", "expert"), None),
+        (None, None),
+        (("fsdp", "tpsp", "expert"), None),
+        None,
+    ),
+    # FFN2 DGRAD: dH = dY2 @ W2^T, contract the replicated hidden dim -> no reduction.
+    "tpsp_ffn2_dgrad": Case(
+        _TPSP_AXES,
+        _MR_TPSP,
+        (524288, 7168),
+        (("fsdp", "tpsp", "expert"), None),
+        (2048, 7168),
+        (None, None),
+        ((1,), (1,)),
+        (("fsdp", "tpsp", "expert"), None),
+        (None, None),
+        (("fsdp", "tpsp", "expert"), None),
+        None,
+    ),
+    # FFN2 WGRAD: dW2 = H^T @ dY2, contract the nested token dim -> reduce (fsdp, tpsp, expert).
+    "tpsp_ffn2_wgrad": Case(
+        _TPSP_AXES,
+        _MR_TPSP,
+        (2048, 524288),
+        (None, ("fsdp", "tpsp", "expert")),
+        (7168, 524288),
+        (None, ("fsdp", "tpsp", "expert")),
+        ((1,), (1,)),
+        (None, ("fsdp", "tpsp", "expert")),
+        (None, ("fsdp", "tpsp", "expert")),
+        (None, None),
+        ("fsdp", "tpsp", "expert"),
+    ),
+    # Single-axis contracting shared by both operands (backward compat).
+    "single_axis": Case(
+        ("fsdp",),
+        dict(fsdp_resource="fsdp"),
+        (128, 512),
+        (None, "fsdp"),
+        (256, 512),
+        (None, "fsdp"),
+        ((1,), (1,)),
+        (None, "fsdp"),
+        (None, "fsdp"),
+        (None, None),
+        "fsdp",
+    ),
+    # No shared contracting axis -> gather both, no reduction.
+    "no_shared_axis": Case(
+        ("fsdp", "tp"),
+        dict(fsdp_resource="fsdp", tp_resource="tp"),
+        (128, 512),
+        ("fsdp", None),
+        (256, 512),
+        ("tp", None),
+        ((1,), (1,)),
+        ("fsdp", None),
+        ("tp", None),
+        ("fsdp", "tp"),
+        None,
+    ),
+}
+
+
+class TestGemmPartitioning:
+    """Spec inference for the plain GEMM partition rule."""
+
+    @pytest.mark.parametrize("case", CASES.values(), ids=CASES.keys())
+    def test_partition_specs(self, case):
+        with _mesh(case.axes), global_shard_guard(MeshResource(**case.resource)):
+            lhs_specs, rhs_specs, out_specs, reduce_spec = _parse(
+                case.lhs_shape, case.lhs_spec, case.rhs_shape, case.rhs_spec, case.cdims
+            )
+        assert lhs_specs == case.exp_lhs
+        assert rhs_specs == case.exp_rhs
+        assert out_specs == case.exp_out
+        assert reduce_spec == case.exp_reduce
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

transformer_engine/jax/cpp_extensions/gemm.py

@@ -913,13 +913,31 @@ def _parse_operand_output_specs(
             (lhs_non_cdims, lhs_cdims, rhs_non_cdims, rhs_cdims),
         )
 
+        # A contracting-dim spec element can be a single mesh axis or a nested tuple of axes
+        # (e.g. ("fsdp", "tp", "expert")). Convert to a list so we can reason per mesh axis.
+        def _convert_axis_spec_to_list(spec):
+            if spec is None:
+                return []
+            return list(spec) if isinstance(spec, tuple) else [spec]
+
+        def _retain_axes(spec, keep):
+            axes = tuple(a for a in _convert_axis_spec_to_list(spec) if a in keep)
+            if len(axes) == 0:
+                return None
+            return axes[0] if len(axes) == 1 else axes
+
+        # The GEMM reduces over the mesh axes that shard the contracting dims of both operands.
+        # Axes sharding the contracting dim of only one operand must be gathered before the GEMM.
+        lhs_c_axes = [a for s in lhs_cspecs for a in _convert_axis_spec_to_list(s)]
+        rhs_c_axes = [a for s in rhs_cspecs for a in _convert_axis_spec_to_list(s)]
+        reduce_axes = tuple(a for a in lhs_c_axes if a in rhs_c_axes)
+        if len(set(reduce_axes)) != len(reduce_axes):
+            raise RuntimeError("Multiple reduce dimension is detected!")
         reduce_spec = None
-        for l in lhs_cspecs:
-            for r in rhs_cspecs:
-                if l is not None and l == r:
-                    if reduce_spec is not None:
-                        raise RuntimeError("Multiple reduce dimension is detected!")
-                    reduce_spec = l
+        if len(reduce_axes) == 1:
+            reduce_spec = reduce_axes[0]
+        elif len(reduce_axes) > 1:
+            reduce_spec = reduce_axes
 
         sequence_dim = None
 
@@ -951,15 +969,14 @@ def _parse_operand_output_specs(
             sequence_dim = int(not transpose_batch_sequence)
 
         if reduce_spec is not None:
-            # Other non-reduce cdims (if exists) need to be unsharded
-            lhs_cspecs = tuple(s if s == reduce_spec else None for s in lhs_cspecs)
+            # Non-reduce contracting axes (if any) need to be gathered, i.e. set to unsharded.
+            lhs_cspecs = tuple(_retain_axes(s, reduce_axes) for s in lhs_cspecs)
             # Only do AG Sequence dim if not Overlap
             if collective_op.is_all_gather:
-                rhs_cspecs = tuple(
-                    s if s in (reduce_spec, gsr.tpsp_resource) else None for s in rhs_cspecs
-                )
+                keep = set(reduce_axes) | {gsr.tpsp_resource}
+                rhs_cspecs = tuple(_retain_axes(s, keep) for s in rhs_cspecs)
             else:
-                rhs_cspecs = tuple(s if s == reduce_spec else None for s in rhs_cspecs)
+                rhs_cspecs = tuple(_retain_axes(s, reduce_axes) for s in rhs_cspecs)
 
             # Non-contracting dims of RHS always needs to be gathered, i.e. for TP + activation_hidden
             # No batch-dim check needed as `rhs_non_cspecs` never contains batch-dim.