From 8ae2dfe91e286af966c51ac7f9878ae9822ed687 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 26 Aug 2026 15:53:41 -0700 Subject: [PATCH 1/3] [JAX] Fix GEMM partitioning to reduce over nested contracting mesh axes Signed-off-by: Phuong Nguyen --- tests/jax/test_gemm_partitioning.py | 169 ++++++++++++++++++ transformer_engine/jax/cpp_extensions/gemm.py | 41 +++-- 2 files changed, 198 insertions(+), 12 deletions(-) create mode 100644 tests/jax/test_gemm_partitioning.py diff --git a/tests/jax/test_gemm_partitioning.py b/tests/jax/test_gemm_partitioning.py new file mode 100644 index 0000000000..a8cec6a389 --- /dev/null +++ b/tests/jax/test_gemm_partitioning.py @@ -0,0 +1,169 @@ +# 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. +""" +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 + + +def _mesh(axes): + """Build a CPU mesh with a size-2 device grid over the named axes.""" + devices = np.asarray(jax.devices()[: 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 on a fsdp x tp x expert mesh. Hidden dims are sharded +# over tp; the token dim is sharded over (fsdp, expert) and may additionally carry tp +# on only one operand (the two WGrad orientations seen in a real 256-GPU HLO). Expected: +# gather the odd tp axis on whichever operand carries it, then reduce over the axes that +# shard the contracting dim of both operands. +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="tp", ep_resource="expert") + +CASES = { + # WGrad, tp leaked onto X's token dim only -> gather tp on X, reduce (fsdp, expert). + "wgrad_nested_tp_on_x": Case( + ("fsdp", "tp", "expert"), + _MR_TP, + (7168, 524288), + (None, ("fsdp", "tp", "expert")), + (256, 524288), + ("tp", ("fsdp", "expert")), + ((1,), (1,)), + (None, ("fsdp", "expert")), + ("tp", ("fsdp", "expert")), + (None, "tp"), + ("fsdp", "expert"), + ), + # Mirror orientation: tp leaked onto dY's token dim only -> gather tp on dY. + "wgrad_mirror_tp_on_dy": Case( + ("fsdp", "tp", "expert"), + _MR_TP, + (7168, 524288), + (None, ("fsdp", "expert")), + (256, 524288), + ("tp", ("fsdp", "tp", "expert")), + ((1,), (1,)), + (None, ("fsdp", "expert")), + ("tp", ("fsdp", "expert")), + (None, "tp"), + ("fsdp", "expert"), + ), + # Forward: contract the tp-sharded hidden dim -> reduce over tp only. + "fwd_reduce_over_tp": Case( + ("fsdp", "tp", "expert"), + _MR_TPSP, + (524288, 7168), + (("fsdp", "expert"), "tp"), + (7168, 2048), + ("tp", None), + ((1,), (0,)), + (("fsdp", "expert"), "tp"), + ("tp", None), + (("fsdp", "expert"), None), + "tp", + ), + # 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"]) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index fe3515ba03..699664ca03 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/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")). Flatten so we can reason per mesh axis. + def _flatten_spec(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 _flatten_spec(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 _flatten_spec(s)] + rhs_c_axes = [a for s in rhs_cspecs for a in _flatten_spec(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. From bad3f055a9ec529e514372c60e86c1b516f90d67 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 26 Aug 2026 16:59:54 -0700 Subject: [PATCH 2/3] [JAX] Use CPU devices explicitly in gemm partitioning test mesh Signed-off-by: Phuong Nguyen --- tests/jax/test_gemm_partitioning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/jax/test_gemm_partitioning.py b/tests/jax/test_gemm_partitioning.py index a8cec6a389..e020f66e58 100644 --- a/tests/jax/test_gemm_partitioning.py +++ b/tests/jax/test_gemm_partitioning.py @@ -29,7 +29,7 @@ def _mesh(axes): """Build a CPU mesh with a size-2 device grid over the named axes.""" - devices = np.asarray(jax.devices()[: 2 ** len(axes)]).reshape((2,) * len(axes)) + devices = np.asarray(jax.devices("cpu")[: 2 ** len(axes)]).reshape((2,) * len(axes)) return Mesh(devices, axes) From f36703b487dc23f47002be6f219bb8fe70951071 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 27 Aug 2026 16:58:22 -0700 Subject: [PATCH 3/3] [JAX] Expand GEMM partitioning tests with consistent FFN fwd/bwd sharding cases Signed-off-by: Phuong Nguyen --- tests/jax/test_gemm_partitioning.py | 202 +++++++++++++++--- transformer_engine/jax/cpp_extensions/gemm.py | 10 +- 2 files changed, 179 insertions(+), 33 deletions(-) diff --git a/tests/jax/test_gemm_partitioning.py b/tests/jax/test_gemm_partitioning.py index e020f66e58..e0a3ba5dd3 100644 --- a/tests/jax/test_gemm_partitioning.py +++ b/tests/jax/test_gemm_partitioning.py @@ -6,7 +6,21 @@ 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", ""): @@ -26,6 +40,8 @@ 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.""" @@ -51,7 +67,7 @@ def _parse(lhs_shape, lhs_spec, rhs_shape, rhs_spec, contracting_dims): scalar, scalar, ) - (operand_specs, out_specs, reduce_spec, _) = GemmPrimitive._parse_operand_output_specs( + operand_specs, out_specs, reduce_spec, _ = GemmPrimitive._parse_operand_output_specs( arg_infos, contracting_dims, transpose_batch_sequence=False, @@ -62,11 +78,10 @@ def _parse(lhs_shape, lhs_spec, rhs_shape, rhs_spec, contracting_dims): return lhs_specs, rhs_specs, tuple(out_specs), reduce_spec -# Representative MoE FFN GEMMs on a fsdp x tp x expert mesh. Hidden dims are sharded -# over tp; the token dim is sharded over (fsdp, expert) and may additionally carry tp -# on only one operand (the two WGrad orientations seen in a real 256-GPU HLO). Expected: -# gather the odd tp axis on whichever operand carries it, then reduce over the axes that -# shard the contracting dim of both operands. +# 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," @@ -74,44 +89,62 @@ def _parse(lhs_shape, lhs_spec, rhs_shape, rhs_spec, contracting_dims): ) _MR_TP = dict(fsdp_resource="fsdp", tp_resource="tp", ep_resource="expert") -_MR_TPSP = dict(fsdp_resource="fsdp", tpsp_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 = { - # WGrad, tp leaked onto X's token dim only -> gather tp on X, reduce (fsdp, expert). - "wgrad_nested_tp_on_x": Case( - ("fsdp", "tp", "expert"), + # --- 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, - (7168, 524288), - (None, ("fsdp", "tp", "expert")), - (256, 524288), - ("tp", ("fsdp", "expert")), + (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,)), - (None, ("fsdp", "expert")), - ("tp", ("fsdp", "expert")), + (("fsdp", "expert"), "tp"), (None, "tp"), - ("fsdp", "expert"), + (("fsdp", "expert"), None), + "tp", ), - # Mirror orientation: tp leaked onto dY's token dim only -> gather tp on dY. - "wgrad_mirror_tp_on_dy": Case( - ("fsdp", "tp", "expert"), + # 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")), - (256, 524288), - ("tp", ("fsdp", "tp", "expert")), + (2048, 524288), + ("tp", ("fsdp", "expert")), ((1,), (1,)), (None, ("fsdp", "expert")), ("tp", ("fsdp", "expert")), (None, "tp"), ("fsdp", "expert"), ), - # Forward: contract the tp-sharded hidden dim -> reduce over tp only. - "fwd_reduce_over_tp": Case( - ("fsdp", "tp", "expert"), - _MR_TPSP, - (524288, 7168), + # 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"), - (7168, 2048), + (2048, 7168), ("tp", None), ((1,), (0,)), (("fsdp", "expert"), "tp"), @@ -119,6 +152,119 @@ def _parse(lhs_shape, lhs_spec, rhs_shape, rhs_spec, contracting_dims): (("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",), diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 699664ca03..68c72d5059 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -914,22 +914,22 @@ def _parse_operand_output_specs( ) # A contracting-dim spec element can be a single mesh axis or a nested tuple of axes - # (e.g. ("fsdp", "tp", "expert")). Flatten so we can reason per mesh axis. - def _flatten_spec(spec): + # (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 _flatten_spec(spec) if a in 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 _flatten_spec(s)] - rhs_c_axes = [a for s in rhs_cspecs for a in _flatten_spec(s)] + 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!")