From 0900878e74c7d523fb1ad252f7e25e7a8c60e1e3 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 1 Sep 2026 15:19:43 -0700 Subject: [PATCH 1/5] [JAX] Support compound EP with ETP1 in MoEBlock Signed-off-by: Jeremy Berchtold --- qa/L0_jax_distributed_unittest/test.sh | 6 + tests/jax/run_te_ep_moe.sh | 8 +- tests/jax/test_multi_process_ep.py | 80 +++++++- tests/jax/test_te_ep_moe.py | 194 ++++++++++++++++---- transformer_engine/jax/cpp_extensions/ep.py | 57 +++--- transformer_engine/jax/ep.py | 23 ++- transformer_engine/jax/flax/moe.py | 26 ++- transformer_engine/jax/moe.py | 109 +++++++---- transformer_engine/jax/sharding.py | 40 +++- 9 files changed, 423 insertions(+), 120 deletions(-) diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index a563e6908d..48c91455d3 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -47,6 +47,12 @@ wait # >=4 visible GPUs. TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ || test_fail "test_te_ep_moe.py" +# Focused four-rank mesh: dense DP2 x TP2, folded MoE EP4, explicit ETP1. +if [ "$(nvidia-smi -L | wc -l)" -ge 4 ]; then + TE_EP_MOE_COMPOUND_ETP1=1 NUM_GPUS=4 TE_PATH=$TE_PATH \ + bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ + || test_fail "test_te_ep_moe.py compound EP/ETP1" +fi # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs; # auto-skips otherwise). CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)" diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh index 9fcbb9dd9a..824ae752e5 100755 --- a/tests/jax/run_te_ep_moe.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -18,10 +18,15 @@ TEST_FILE="${TEST_FILE:-$TE_ROOT/tests/jax/test_te_ep_moe.py}" PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" +COMPOUND_EP_ETP1="${TE_EP_MOE_COMPOUND_ETP1:-0}" if [ "$NUM_GPUS" -lt 4 ]; then echo "[run_te_ep_moe.sh] need >=4 GPUs (got $NUM_GPUS); SKIPPING." exit 0 fi +if [ "$COMPOUND_EP_ETP1" = "1" ] && [ "$NUM_GPUS" -ne 4 ]; then + echo "[run_te_ep_moe.sh] compound EP/ETP1 mode requires exactly 4 processes (got $NUM_GPUS); SKIPPING." + exit 0 +fi export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}" export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.5}" @@ -31,6 +36,7 @@ echo "============================================================" echo "TE-EP MoE MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" echo " test file : $TEST_FILE" echo " coordinator : $TE_EP_MOE_COORDINATOR_ADDRESS" +echo " compound EP/ETP1 : $COMPOUND_EP_ETP1" echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" echo "============================================================" @@ -96,7 +102,7 @@ for i in "${!EXITS[@]}"; do done # Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which the -# file emits via pytest.skip(allow_module_level=True) on pre-Blackwell +# file emits via pytest.skip(allow_module_level=True) on pre-Hopper # GPUs) as success. FAILED=0 for e in "${EXITS[@]}"; do diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 47af0b0c39..3adbd2e20d 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -33,7 +33,13 @@ from jax.sharding import Mesh, NamedSharding, PartitionSpec from utils import is_devices_enough -from transformer_engine.jax.sharding import MeshResource, global_shard_guard +from transformer_engine.jax.cpp_extensions.ep import ( + ep_prepare, + ep_dispatch_fwd, + ep_combine_fwd, + get_ep_config, + _ep_outer_axis, +) from transformer_engine.jax.ep import ( EpLayerConfig, ep_bootstrap, @@ -42,11 +48,15 @@ ep_combine, _ep_domain_for_rank, ) -from transformer_engine.jax.cpp_extensions.ep import ( - ep_prepare, - ep_dispatch_fwd, - ep_combine_fwd, - get_ep_config, +from transformer_engine.jax.moe import _moe_etp_axis +from transformer_engine.jax.sharding import ( + BATCH_AXES, + W_TP_AXES, + MeshResource, + get_active_resource_axis, + get_mesh_axis_size, + get_sharding_map_logic_axis_to_mesh_axis, + global_shard_guard, ) from transformer_engine.jax.version_utils import is_collective_stream_supported @@ -960,6 +970,64 @@ def test_ep_tp_splits_domains(self): self.assertEqual(domains, {0: [0, 2, 4, 6], 1: [1, 3, 5, 7]}) + def test_compound_ep_with_explicit_etp1(self): + if not is_devices_enough(4): + self.skipTest("requires 4 devices") + mesh = Mesh( + np.asarray(jax.devices()[:4]).reshape(2, 2, 1), + ("expert", "tensor", "etp"), + ) + order = {int(d.id): i for i, d in enumerate(mesh.devices.reshape(-1))} + d2r = lambda d: order[int(d.id)] + + resource = MeshResource( + dp_resource="expert", + tp_resource="tensor", + ep_resource=("expert", "tensor"), + etp_resource="etp", + ) + with mesh, global_shard_guard(resource): + self.assertEqual(get_mesh_axis_size(resource.ep_resource), 4) + self.assertEqual(get_mesh_axis_size(resource.etp_resource), 1) + self.assertEqual(get_active_resource_axis("ep_resource"), ("expert", "tensor")) + self.assertIsNone(get_active_resource_axis("etp_resource")) + dense_rules = get_sharding_map_logic_axis_to_mesh_axis() + self.assertEqual(dense_rules[BATCH_AXES], "expert") + self.assertEqual(dense_rules[W_TP_AXES], "tensor") + # Dense DP is folded into compound EP here, not counted again as an + # outer group for the MoE communication buffers. + self.assertIsNone(_ep_outer_axis()) + + domains = {} + for rank in range(4): + root, col, ndom = _ep_domain_for_rank( + mesh, ("expert", "tensor"), rank, device_to_rank=d2r + ) + self.assertEqual(ndom, 1) + domains.setdefault(root, {})[col] = rank + domains = {root: [m[c] for c in sorted(m)] for root, m in domains.items()} + self.assertEqual(domains, {0: [0, 1, 2, 3]}) + + # A size-one member is legal in a compound resource even if JAX later + # elides it from a concrete PartitionSpec. + etp1_domains = {} + for rank in range(4): + root, col, ndom = _ep_domain_for_rank( + mesh, ("expert", "etp"), rank, device_to_rank=d2r + ) + self.assertEqual(ndom, 2) + etp1_domains.setdefault(root, {})[col] = rank + etp1_domains = { + root: [members[c] for c in sorted(members)] + for root, members in etp1_domains.items() + } + self.assertEqual(etp1_domains, {0: [0, 2], 1: [1, 3]}) + with mesh, global_shard_guard( + MeshResource(ep_resource=("expert", "etp"), etp_resource="etp") + ): + self.assertEqual(get_active_resource_axis("ep_resource"), ("expert", "etp")) + self.assertIsNone(_moe_etp_axis(("expert", "etp"), "etp")) + # ── Entry point ────────────────────────────────────────────────────────────── diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 015c73343a..d1555ca0e2 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -31,7 +31,7 @@ on the block are pytest parametrize values rather than separate test classes: -* ``test_forward`` covers BF16 and MXFP8 forward execution across a +* ``test_forward`` covers BF16 (Hopper+) and MXFP8 (Blackwell+) forward execution across a curated set of configurations (softmax/sigmoid scoring, optional non-zero expert_bias). Each config asserts shape, dtype, finiteness and numerical parity vs the same BF16 reference in one run. @@ -57,6 +57,7 @@ from jax.experimental import mesh_utils from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from flax import linen as flax_linen from flax.linen import partitioning as nn_partitioning @@ -109,12 +110,11 @@ def _read_mp_options(): from transformer_engine_jax import get_device_compute_capability -# Grouped GEMM in the MoE custom_vjp requires Blackwell (sm_100+). The -# TE EP NCCL primitives themselves need SM>=90, but the FFN body uses -# grouped_gemm, so the file as a whole gates on sm_100+. -if get_device_compute_capability(0) < 100: +# NCCL EP requires Hopper or newer. BF16 grouped GEMM falls back to the +# legacy implementation on Hopper; MXFP8 remains gated to Blackwell below. +if get_device_compute_capability(0) < 90: pytest.skip( - "MoE TE EP tests require Blackwell (sm_100+) for grouped GEMM", + "MoE TE EP tests require Hopper (sm_90+) or newer", allow_module_level=True, ) @@ -134,26 +134,55 @@ def _read_mp_options(): # Mesh / shape config # ----------------------------------------------------------------------------- -EP_AXIS = "ep" -FSDP_AXIS = "fsdp" -EP_SIZE = 2 -assert ( - jax.device_count() % EP_SIZE == 0 -), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" -FSDP_SIZE = jax.device_count() // EP_SIZE -NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE - -LOGICAL_AXIS_RULES = ( - ("exp", EP_AXIS), - ("embed", FSDP_AXIS), - ("mlp", None), - ("batch", (FSDP_AXIS, EP_AXIS)), -) +COMPOUND_EP_ETP1 = os.environ.get("TE_EP_MOE_COMPOUND_ETP1", "0") == "1" + +if COMPOUND_EP_ETP1: + EP_AXIS = ("expert", "tensor") + EP_SIZE = 4 + MESH_SHAPE = (2, 2, 1) + MESH_AXIS_NAMES = ("expert", "tensor", "etp") + # Preserve the dense-layer DP declaration. MoEBlock must recognize that + # this physical axis is already folded into compound EP and not count it twice. + DATA_PARALLELISM_AXES = ("expert",) + BATCH_MESH_AXIS = EP_AXIS + MESH_RESOURCE = MeshResource( + dp_resource="expert", + tp_resource="tensor", + ep_resource=EP_AXIS, + etp_resource="etp", + ) + LOGICAL_AXIS_RULES = ( + ("exp", EP_AXIS), + ("embed", None), + ("mlp", "etp"), + ("batch", EP_AXIS), + ) +else: + EP_AXIS = "ep" + FSDP_AXIS = "fsdp" + EP_SIZE = 2 + assert ( + jax.device_count() % EP_SIZE == 0 + ), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" + FSDP_SIZE = jax.device_count() // EP_SIZE + MESH_SHAPE = (FSDP_SIZE, EP_SIZE) + MESH_AXIS_NAMES = (FSDP_AXIS, EP_AXIS) + DATA_PARALLELISM_AXES = (FSDP_AXIS,) + BATCH_MESH_AXIS = (FSDP_AXIS, EP_AXIS) + MESH_RESOURCE = MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + LOGICAL_AXIS_RULES = ( + ("exp", EP_AXIS), + ("embed", FSDP_AXIS), + ("mlp", None), + ("batch", BATCH_MESH_AXIS), + ) + +NUM_DEVICES_REQUIRED = int(np.prod(MESH_SHAPE)) # Small shapes so the parity tests stay tight on bf16. The block still # has all four ranks participating in dispatch/combine. DTYPE = jnp.bfloat16 -BATCH = EP_SIZE * FSDP_SIZE * 2 # 8 on 4-GPU, 16 on 8-GPU +BATCH = NUM_DEVICES_REQUIRED * 2 SEQ = 32 HIDDEN = 128 INTER = 128 @@ -189,16 +218,20 @@ def _read_mp_options(): @pytest.fixture(scope="module") def mesh(): - if jax.device_count() < NUM_DEVICES_REQUIRED: + if jax.device_count() < NUM_DEVICES_REQUIRED or ( + COMPOUND_EP_ETP1 and jax.device_count() != NUM_DEVICES_REQUIRED + ): pytest.skip( - f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" - f" have {jax.device_count()}" + f"Need {'exactly' if COMPOUND_EP_ETP1 else '>='} {NUM_DEVICES_REQUIRED} devices for" + f" mesh={MESH_SHAPE}; have {jax.device_count()}" ) - # ``ep`` must be the inner axis: ``ep_bootstrap`` forms NCCL EP groups - # from consecutive global ranks via ``dp_color = rank // ep_size``, so - # only an (outer_fsdp, inner_ep) device layout groups ranks correctly. - devices = mesh_utils.create_device_mesh((FSDP_SIZE, EP_SIZE)) - mesh_obj = Mesh(devices, axis_names=(FSDP_AXIS, EP_AXIS)) + # Compound EP axes retain their declared order. With shape (2, 2, 1), + # (expert, tensor) flattens to contiguous ranks [0, 1, 2, 3]. + if COMPOUND_EP_ETP1: + devices = np.asarray(jax.devices()).reshape(MESH_SHAPE) + else: + devices = mesh_utils.create_device_mesh(MESH_SHAPE) + mesh_obj = Mesh(devices, axis_names=MESH_AXIS_NAMES) num_procs = jax.process_count() max_tokens_per_rank = (BATCH // num_procs) * SEQ @@ -215,7 +248,7 @@ def mesh(): # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather # and cannot run from inside jax.jit. Sized to the worst-case recv_pr # across _CONFIGS so every parametrized config is bootstrap-compatible. - with mesh_obj, global_shard_guard(MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS)): + with mesh_obj, global_shard_guard(MESH_RESOURCE): ep_bootstrap( world_size=num_procs, rank=jax.process_index(), @@ -364,7 +397,7 @@ def _make_block( num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, intermediate_size=INTER, - data_parallelism_axes=(FSDP_AXIS,), + data_parallelism_axes=DATA_PARALLELISM_AXES, apply_topk_weights_early=apply_topk_weights_early, aux_loss_coeff=aux_loss_coeff, use_expert_routing_bias=use_expert_routing_bias, @@ -394,9 +427,7 @@ def _strong_expert_bias_init(key, shape, dtype): def _shard_inputs(x, mesh): # Match the layout moe.py re-pins to: outer dp axes, then ep innermost. - return jax.lax.with_sharding_constraint( - x, NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None, None)) - ) + return jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P(BATCH_MESH_AXIS, None, None))) def _ctx(mesh): @@ -405,9 +436,7 @@ def _ctx(mesh): class _Combo: def __enter__(self_inner): self_inner._m = mesh.__enter__() - self_inner._gs = global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ) + self_inner._gs = global_shard_guard(MESH_RESOURCE) self_inner._gs.__enter__() self_inner._ar = nn_partitioning.axis_rules(LOGICAL_AXIS_RULES) self_inner._ar.__enter__() @@ -424,7 +453,18 @@ def __exit__(self_inner, *args): def _init_apply(block, mesh, x, key): with _ctx(mesh): x_sh = _shard_inputs(x, mesh) - variables = jax.jit(block.init)(key, x_sh) + if COMPOUND_EP_ETP1: + assert _axis_names(_spec_entry(x_sh.sharding.spec, 0)) == frozenset(EP_AXIS) + assert not any(_axis_names(entry) for entry in x_sh.sharding.spec[1:]) + if COMPOUND_EP_ETP1: + abstract_variables = jax.eval_shape(block.init, key, x_sh) + logical_specs = flax_linen.get_partition_spec(abstract_variables) + variable_shardings = flax_linen.logical_to_mesh_sharding( + logical_specs, mesh, LOGICAL_AXIS_RULES + ) + variables = jax.jit(block.init, out_shardings=variable_shardings)(key, x_sh) + else: + variables = jax.jit(block.init)(key, x_sh) jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) output, aux, _trt = jax.jit(block.apply)(variables, x_sh) jax.block_until_ready(output) @@ -500,6 +540,50 @@ def _params_global_numpy(variables, mesh): return {name: _to_global_numpy(_unwrap(p), mesh) for name, p in params.items()} +def _axis_names(entry): + if entry is None: + return frozenset() + return frozenset(entry if isinstance(entry, tuple) else (entry,)) + + +def _spec_entry(spec, index): + return spec[index] if index < len(spec) else None + + +def _slice_size(index, global_size): + start = 0 if index.start is None else index.start + stop = global_size if index.stop is None else index.stop + return stop - start + + +def _assert_compound_ep_etp1_sharding(variables, output): + """Assert TP is folded into EP while every local expert matrix is complete.""" + assert _axis_names(_spec_entry(output.sharding.spec, 0)) == frozenset(EP_AXIS) + assert not any(_axis_names(entry) for entry in output.sharding.spec[1:]) + + params = variables["params"] + gate = _unwrap(params["gate_kernel"]) + wi = _unwrap(params["wi"]) + wo = _unwrap(params["wo"]) + + assert not any(_axis_names(entry) for entry in gate.sharding.spec), ( + "gate weights should be replicated for global routing" + ) + assert _axis_names(_spec_entry(wi.sharding.spec, 0)) == frozenset(EP_AXIS) + assert _axis_names(_spec_entry(wo.sharding.spec, 0)) == frozenset(EP_AXIS) + # JAX may preserve or elide a size-one axis. Either representation must + # leave the expert's hidden/intermediate matrices physically complete. + assert _axis_names(_spec_entry(wi.sharding.spec, 2)) <= {"etp"} + assert _axis_names(_spec_entry(wo.sharding.spec, 1)) <= {"etp"} + + wi_index = wi.addressable_shards[0].index + wo_index = wo.addressable_shards[0].index + assert _slice_size(wi_index[0], NUM_EXPERTS) == NUM_EXPERTS // EP_SIZE + assert _slice_size(wo_index[0], NUM_EXPERTS) == NUM_EXPERTS // EP_SIZE + assert wi_index[1] == slice(None) and wi_index[2] == slice(None) + assert wo_index[1] == slice(None) and wo_index[2] == slice(None) + + def _make_inputs(key): """Generate a globally-identical input tensor on every process.""" return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) @@ -555,11 +639,16 @@ def _quantization_recipe(quantization): ), ] +if COMPOUND_EP_ETP1: + # This mode is a focused sharding qualification, not a repeat of the + # router feature matrix covered by the default topology. + _CONFIGS = [_CONFIGS[0]] + _QUANTIZATION_CASES = [ pytest.param("bf16", id="bf16"), ] -if get_device_compute_capability(0) >= 100: +if get_device_compute_capability(0) >= 100 and not COMPOUND_EP_ETP1: _QUANTIZATION_CASES.append(pytest.param("mxfp8", id="mxfp8")) @@ -586,6 +675,9 @@ def test_forward(self, mesh, config, quantization): x = _make_inputs(jax.random.PRNGKey(0)) variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) + if COMPOUND_EP_ETP1: + _assert_compound_ep_etp1_sharding(variables, output) + # Shape / dtype / finiteness (cheap; on the local shard). assert output.shape == x.shape assert output.dtype == x.dtype @@ -614,6 +706,27 @@ def test_forward(self, mesh, config, quantization): err_msg=f"forward parity breach for config={config}, quantization={quantization}", ) + def test_active_etp_is_rejected(self, mesh): + if not COMPOUND_EP_ETP1: + pytest.skip("only applies to the explicit ETP resource mode") + active_etp_mesh = Mesh(mesh.devices.reshape(1, 2, 2), MESH_AXIS_NAMES) + resource = MeshResource( + dp_resource="expert", + tp_resource="tensor", + ep_resource=("expert", "tensor"), + etp_resource="etp", + ) + rules = ( + ("exp", ("expert", "tensor")), + ("embed", None), + ("mlp", "etp"), + ("batch", ("expert", "tensor")), + ) + block = _make_block() + with active_etp_mesh, global_shard_guard(resource), nn_partitioning.axis_rules(rules): + with pytest.raises(NotImplementedError, match="ETP=1"): + block.init(jax.random.PRNGKey(40), _make_inputs(jax.random.PRNGKey(41))) + class TestTeEpMoeBackward: """Per-config backward correctness in a single run: per-tensor @@ -692,6 +805,7 @@ def loss_fn(params, x): ) +@pytest.mark.skipif(COMPOUND_EP_ETP1, reason="compound mode focuses on EP/ETP sharding") class TestTeEpMoeAuxLoss: """Aux-loss path. Consolidated into: * ``test_aux_loss``: one run that checks the returned scalar's diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index ca70ea145c..c184cf0977 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -24,7 +24,7 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive -from ..sharding import global_mesh_resource, get_mesh_axis_size +from ..sharding import global_mesh_resource, get_mesh_axis_size, normalize_mesh_axes from ..version_utils import is_collective_stream_supported @@ -139,18 +139,24 @@ def _leading_axis_ok(spec): first); all other dims must be replicated. """ gsr = global_mesh_resource() - ep_axis = gsr.ep_resource - outer_axes = tuple(a for a in (gsr.dp_resource, gsr.fsdp_resource) if a is not None) - if len(spec) < 2 or ep_axis is None: - return False, ep_axis, outer_axes + ep_axes = normalize_mesh_axes(gsr.ep_resource) + outer_axes = tuple( + a + for a in (gsr.dp_resource, gsr.fsdp_resource) + if a is not None and a not in ep_axes + ) + if len(spec) < 2 or not ep_axes: + return False, gsr.ep_resource, outer_axes if any(ax is not None for ax in spec[1:]): - return False, ep_axis, outer_axes + return False, gsr.ep_resource, outer_axes leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) - if ep_axis not in elts: - return False, ep_axis, outer_axes - allowed = set(outer_axes) | {ep_axis} - return all(a in allowed for a in elts), ep_axis, outer_axes + actual = set(a for a in elts if a is not None) + required_ep = {a for a in ep_axes if get_mesh_axis_size(a) > 1} + if not required_ep.issubset(actual): + return False, gsr.ep_resource, outer_axes + allowed = set(outer_axes) | set(ep_axes) + return actual.issubset(allowed), gsr.ep_resource, outer_axes def _ep_outer_axis(): @@ -163,11 +169,16 @@ def _ep_outer_axis(): we don't pin EP-output specs to a degenerate axis that JAX may collapse. """ gsr = global_mesh_resource() - if gsr.dp_resource is not None and get_mesh_axis_size(gsr.dp_resource) > 1: - return gsr.dp_resource - if gsr.fsdp_resource is not None and get_mesh_axis_size(gsr.fsdp_resource) > 1: - return gsr.fsdp_resource - return gsr.dp_resource or gsr.fsdp_resource + ep_axes = set(normalize_mesh_axes(gsr.ep_resource)) + candidates = tuple( + axis + for axis in (gsr.dp_resource, gsr.fsdp_resource) + if axis is not None and axis not in ep_axes + ) + for axis in candidates: + if get_mesh_axis_size(axis) > 1: + return axis + return candidates[0] if candidates else None def _ep_leading_dims(is_outer): @@ -184,9 +195,10 @@ def _ep_output_spec(*trailing): DP is set (compound leading axis on a single dim), else ``("ep",*trailing)``.""" gsr = global_mesh_resource() outer = _ep_outer_axis() - if outer is None: - return PartitionSpec(gsr.ep_resource, *trailing) - return PartitionSpec((outer, gsr.ep_resource), *trailing) + ep_axes = normalize_mesh_axes(gsr.ep_resource) + leading_axes = ep_axes if outer is None else (outer, *ep_axes) + leading = leading_axes[0] if len(leading_axes) == 1 else leading_axes + return PartitionSpec(leading, *trailing) def _ep_spec_ok(spec, trailing_count): @@ -195,7 +207,7 @@ def _ep_spec_ok(spec, trailing_count): so the leading entry is normalized to a set of named axes before comparing. """ gsr = global_mesh_resource() - ep_axis = gsr.ep_resource + ep_axes = normalize_mesh_axes(gsr.ep_resource) outer = _ep_outer_axis() if len(spec) != 1 + trailing_count: return False @@ -204,8 +216,11 @@ def _ep_spec_ok(spec, trailing_count): leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) actual = frozenset(a for a in elts if a is not None) - expected = {ep_axis} if outer is None else {ep_axis, outer} - return actual <= expected + expected = set(ep_axes) + if outer is not None: + expected.add(outer) + required = {axis for axis in expected if get_mesh_axis_size(axis) > 1} + return required.issubset(actual) and actual.issubset(expected) # ── ep_prepare ────────────────────────────────────────────────────────────── diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 2222a41e48..157cdd4d94 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -21,6 +21,7 @@ get_num_devices_in_mesh, global_mesh_resource, get_mesh_axis_size, + normalize_mesh_axes, with_sharding_constraint, ) @@ -85,11 +86,21 @@ def _ep_domain_for_rank(mesh, ep_resource, rank, device_to_rank=None): def device_to_rank(d): return d.process_index - ep_pos = mesh.axis_names.index(ep_resource) - ep_size = mesh.shape[ep_resource] + ep_axes = normalize_mesh_axes(ep_resource) + if not ep_axes: + raise ValueError("ep_bootstrap: ep_resource must contain at least one mesh axis.") + missing = tuple(axis for axis in ep_axes if axis not in mesh.axis_names) + if missing: + raise ValueError( + f"ep_bootstrap: EP axes {missing} are not present in mesh axes {mesh.axis_names}." + ) + ep_positions = tuple(mesh.axis_names.index(axis) for axis in ep_axes) + non_ep_positions = tuple(i for i in range(len(mesh.axis_names)) if i not in ep_positions) + ep_size = get_mesh_axis_size(ep_axes, mesh) ranks = np.vectorize(device_to_rank, otypes=[np.int64])(mesh.devices) - # Move ep last and flatten: each row is one domain (all non-ep coords fixed). - grid = np.moveaxis(ranks, ep_pos, -1).reshape(-1, ep_size) + # Move all EP axes last in the user-specified order and flatten them into + # one communicator dimension. Each row fixes every axis outside compound EP. + grid = np.transpose(ranks, non_ep_positions + ep_positions).reshape(-1, ep_size) loc = np.argwhere(grid == rank) if loc.shape[0] != 1: raise ValueError( @@ -248,7 +259,9 @@ def _default_out_partition_spec(): "ep_resource is not set on the active MeshResource; pass out_sharding=... explicitly." ) outer = _ep_outer_axis() - leading = (outer, gsr.ep_resource) if outer is not None else gsr.ep_resource + ep_axes = normalize_mesh_axes(gsr.ep_resource) + leading_axes = ep_axes if outer is None else (outer, *ep_axes) + leading = leading_axes[0] if len(leading_axes) == 1 else leading_axes return (leading,) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index cb10c7aa18..ecebe7ce60 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -13,7 +13,7 @@ as ``self.param`` slots (with the right :func:`flax.linen.with_logical_partitioning` annotations so JAX's sharding layer FSDPs the params correctly). -2. Resolve the EP axis name from the active +2. Resolve compound EP and ETP axis resources from the active :class:`transformer_engine.jax.sharding.MeshResource`. 3. Forward all knobs to :func:`moe`. @@ -34,10 +34,15 @@ from flax import linen as nn from transformer_engine.common.recipe import Recipe -from ..moe import moe +from ..moe import _moe_outer_axes, moe from ..quantize import QuantizerSet from ..router import ScoreFunction -from ..sharding import _get_mesh, get_active_resource_axis +from ..sharding import ( + _get_mesh, + get_active_resource_axis, + get_mesh_axis_size, + global_mesh_resource, +) from .module import TransformerEngineBase PRNGKey = Any @@ -97,7 +102,8 @@ class _MoEBlock(TransformerEngineBase): ADDITION to the EP axis. Empty (default) means activations are replicated across non-EP axes within an EP group; set e.g. ``("fsdp",)`` for true FSDP-of-batch where each device owns a - unique slice of the batch. + unique slice of the batch. Any dense DP axis already included in a + compound EP resource is ignored here rather than counted twice. apply_topk_weights_early : bool If ``True``, multiply expert outputs by their top-k weights *inside* each shard before ``ep_combine`` (saves one global @@ -255,8 +261,15 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: ep_axis = get_active_resource_axis("ep_resource") mesh = _get_mesh() + etp_axis = global_mesh_resource().etp_resource + if etp_axis is not None and get_mesh_axis_size(etp_axis, mesh) != 1: + raise NotImplementedError( + "_MoEBlock currently supports expert tensor parallelism only with ETP=1; " + f"axis {etp_axis!r} has size {get_mesh_axis_size(etp_axis, mesh)}." + ) data_parallel_size = 1 - for axis in self.data_parallelism_axes: + effective_data_parallelism_axes = _moe_outer_axes(ep_axis, self.data_parallelism_axes) + for axis in effective_data_parallelism_axes: data_parallel_size *= mesh.shape[axis] def make_grouped_quantizer_set(postfix): @@ -305,7 +318,8 @@ def make_grouped_quantizer_set(postfix): quantizer_sets=quantizer_sets, recv_capacity_per_rank=self.recv_capacity_per_rank, ep_axis=ep_axis, - data_parallelism_axes=self.data_parallelism_axes, + etp_axis=etp_axis, + data_parallelism_axes=effective_data_parallelism_axes, input_axes=self.input_axes, gate_kernel_axes=self.gate_kernel_axes, wi_kernel_axes=self.wi_kernel_axes, diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 55a85ebb2f..c520074eaa 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -16,8 +16,9 @@ Sharding model -------------- -* Inbound activations are 3D ``[B, S, H]`` sharded - ``((*data_parallelism_axes, ep_axis), None, None)``. The public +* Inbound activations are 3D ``[B, S, H]`` sharded on one leading + compound entry containing outer data-parallel axes followed by every axis + in ``ep_axis``. The public :func:`moe` soft-repins this on entry and warns when a reshard is inserted. * The EP, grouped-quantize, and grouped-GEMM primitives operate at global @@ -50,7 +51,7 @@ ) from .flax.module import _convert_to_activation_function from .router import ScoreFunction, _validate_score_function -from .sharding import _get_mesh +from .sharding import MeshAxis, _get_mesh, get_mesh_axis_size, normalize_mesh_axes __all__ = ["get_moe_recv_capacity_per_rank", "moe"] @@ -63,6 +64,34 @@ _ALIGN_SIZE = 128 +def _moe_outer_axes( + ep_axis: MeshAxis, data_parallelism_axes: Tuple[str, ...] +) -> Tuple[str, ...]: + """Drop dense DP axes that are folded into the compound EP resource.""" + ep_axes = set(normalize_mesh_axes(ep_axis)) + return tuple(axis for axis in data_parallelism_axes if axis not in ep_axes) + + +def _moe_leading_axis(ep_axis: MeshAxis, data_parallelism_axes: Tuple[str, ...]): + """Build one PartitionSpec entry with outer axes followed by compound EP.""" + axes = (*_moe_outer_axes(ep_axis, data_parallelism_axes), *normalize_mesh_axes(ep_axis)) + if not axes: + raise ValueError("moe(...) requires ep_axis to contain at least one mesh axis.") + return axes[0] if len(axes) == 1 else axes + + +def _moe_etp_axis(ep_axis: MeshAxis, etp_axis: Optional[str]) -> Optional[str]: + """Avoid assigning one physical axis to two tensor dimensions. + + A size-one ETP axis may legally appear inside a compound EP resource. JAX + can elide that axis, but it cannot name it on both the expert and matrix + dimensions of the same PartitionSpec. + """ + if etp_axis in normalize_mesh_axes(ep_axis): + return None + return etp_axis + + def get_moe_recv_capacity_per_rank( *, num_experts: int, @@ -560,6 +589,7 @@ def _moe_fwd_rule( scaling_factor, aux_loss_coeff, ep_axis, + etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -584,11 +614,12 @@ def _moe_fwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") if ep_axis is None: raise ValueError("moe(...) requires ep_axis to be set (TE EP backend).") - num_ep = mesh.shape[ep_axis] + num_ep = get_mesh_axis_size(ep_axis, mesh) if num_experts % num_ep != 0: raise ValueError(f"num_experts={num_experts} must be divisible by EP size={num_ep}") num_local_experts = num_experts // num_ep + data_parallelism_axes = _moe_outer_axes(ep_axis, data_parallelism_axes) dp_size = 1 for ax in data_parallelism_axes: dp_size *= mesh.shape[ax] @@ -630,13 +661,9 @@ def _moe_fwd_rule( ep_size=num_ep, ) - if not data_parallelism_axes: - batch_pspec_axis: Any = ep_axis - else: - # ep must be innermost: ep_bootstrap forms NCCL EP comms from - # consecutive global ranks (dp_color = rank // ep_size), so the - # comm only stays within one model replica under (outer_dp, ep). - batch_pspec_axis = (*data_parallelism_axes, ep_axis) + # EP axes must be innermost and retain their declared order: communicator + # ranks are formed by flattening the compound resource in that order. + batch_pspec_axis: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, ep3_spec)) @@ -750,12 +777,15 @@ def _moe_fwd_rule( # ---------------- FFN (per-shard via shard_map) ---------------- has_bias = wi_0_bias is not None - kernel_spec = P(ep_axis, None, None) - bias_spec = P(ep_axis, None) - ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec) + ffn_etp_axis = _moe_etp_axis(ep_axis, etp_axis) + wi_spec = P(ep_axis, None, ffn_etp_axis) + wo_spec = P(ep_axis, ffn_etp_axis, None) + wi_bias_spec = P(ep_axis, ffn_etp_axis) + wo_bias_spec = P(ep_axis, None) + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, wi_spec, wo_spec) ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi, wo] if has_bias: - ffn_in_specs += (bias_spec, bias_spec, bias_spec) + ffn_in_specs += (wi_bias_spec, wi_bias_spec, wo_bias_spec) ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) # Quantized grouped tensors store their data, scales, and group metadata @@ -885,6 +915,7 @@ def _moe_bwd_rule( scaling_factor, aux_loss_coeff, ep_axis, + etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -913,10 +944,8 @@ def _moe_bwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") B, S, _ = x_shape K = num_experts_per_tok - if not data_parallelism_axes: - batch_pspec_axis: Any = ep_axis - else: - batch_pspec_axis = (*data_parallelism_axes, ep_axis) + data_parallelism_axes = _moe_outer_axes(ep_axis, data_parallelism_axes) + batch_pspec_axis: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) out_partition_spec = (batch_pspec_axis, None, None) @@ -939,8 +968,11 @@ def _moe_bwd_rule( d_recv_w_from_combine = d_recv_w_from_combine.astype(ctx.recv_topk_weights.dtype) # ---------------- FFN bwd (per-shard via shard_map) ---------------- - kernel_spec = P(ep_axis, None, None) - bias_spec = P(ep_axis, None) + ffn_etp_axis = _moe_etp_axis(ep_axis, etp_axis) + wi_spec = P(ep_axis, None, ffn_etp_axis) + wo_spec = P(ep_axis, ffn_etp_axis, None) + wi_bias_spec = P(ep_axis, ffn_etp_axis) + wo_bias_spec = P(ep_axis, None) token_buffer_spec = P(batch_pspec_axis) token_matrix_spec = P(batch_pspec_axis, None) expert_buffer_spec = P(ep_axis) @@ -1005,14 +1037,14 @@ def _ffn_bwd_body(*args): bwd_out_specs = ( ep3_spec, ep2_spec, - kernel_spec, - kernel_spec, - bias_spec, - bias_spec, - bias_spec, + wi_spec, + wo_spec, + wi_bias_spec, + wi_bias_spec, + wo_bias_spec, ) else: - bwd_out_specs = (ep3_spec, ep2_spec, kernel_spec, kernel_spec, None, None, None) + bwd_out_specs = (ep3_spec, ep2_spec, wi_spec, wo_spec, None, None, None) ( d_sorted_x, @@ -1140,7 +1172,7 @@ def _ffn_bwd_body(*args): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 27))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 28))) def _moe( x, gate_kernel, @@ -1161,6 +1193,7 @@ def _moe( scaling_factor, aux_loss_coeff, ep_axis, + etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -1190,6 +1223,7 @@ def _moe( scaling_factor, aux_loss_coeff, ep_axis, + etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -1229,7 +1263,8 @@ def moe( noop_quantizer_set, noop_quantizer_set, ), - ep_axis: str, + ep_axis: MeshAxis, + etp_axis: Optional[str] = None, data_parallelism_axes: Tuple[str, ...] = (), input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), @@ -1278,9 +1313,11 @@ def moe( Axis-name parameters: - * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh - axis names* -- they index ``jax.sharding.Mesh.shape`` directly - (to compute ``num_ep`` / ``dp_size`` and to construct + * ``ep_axis``, ``etp_axis``, and ``data_parallelism_axes`` are *physical mesh + axis names*. ``ep_axis`` may be an ordered tuple whose sizes are + multiplied into one compound EP resource. ETP is currently accepted + only when its mesh size is one, leaving expert GEMM matrices complete. + These concrete axes are used to compute ``num_ep`` / ``dp_size`` and construct ``P((dp..., ep), None, None)`` for the physical ``jax.lax.with_sharding_constraint`` calls that JAX requires to refer to real mesh axes). @@ -1313,7 +1350,12 @@ def moe( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - expected_leading: Any = (*data_parallelism_axes, ep_axis) if data_parallelism_axes else ep_axis + if etp_axis is not None and get_mesh_axis_size(etp_axis, mesh) != 1: + raise NotImplementedError( + "moe(...) currently supports expert tensor parallelism only with ETP=1; " + f"axis {etp_axis!r} has size {get_mesh_axis_size(etp_axis, mesh)}." + ) + expected_leading: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) expected_spec = P(expected_leading, None, None) actual_spec = getattr(getattr(x, "sharding", None), "spec", None) if actual_spec is not None and tuple(actual_spec) != tuple(expected_spec): @@ -1354,6 +1396,7 @@ def moe( scaling_factor, float(aux_loss_coeff), ep_axis, + etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 2e8e611fa3..18b298d1f7 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -11,7 +11,7 @@ """ from contextlib import contextmanager from dataclasses import dataclass -from typing import Callable, Optional +from typing import Callable, Optional, Union import warnings import jax @@ -36,6 +36,20 @@ W_TP_AXES = "nvte_w_tp" W_JOINED_AXES = "nvte_w_joined" +MeshAxis = Union[str, tuple[str, ...]] + + +def normalize_mesh_axes(axis: Optional[MeshAxis]) -> tuple[str, ...]: + """Return a mesh resource as an ordered tuple of physical axis names.""" + if axis is None: + return () + axes = axis if isinstance(axis, tuple) else (axis,) + if not axes or any(not isinstance(name, str) or not name for name in axes): + raise ValueError(f"Mesh axes must be non-empty strings, got {axis!r}.") + if len(set(axes)) != len(axes): + raise ValueError(f"Mesh axes must not contain duplicates, got {axis!r}.") + return axes + def _get_mesh(): # Handle Mesh's set via `with mesh:` @@ -274,11 +288,14 @@ def get_mesh_axis_size(axis, mesh=None): if mesh is None: mesh = _get_mesh() - if axis is None: + axes = normalize_mesh_axes(axis) + if not axes: return 1 - - assert axis in mesh.shape, f"{axis} is not a axis of the given mesh {mesh.shape}" - return mesh.shape[axis] + size = 1 + for name in axes: + assert name in mesh.shape, f"{name} is not an axis of the given mesh {mesh.shape}" + size *= mesh.shape[name] + return size def get_mesh_axis_rank(axis: str, mesh=None): @@ -331,12 +348,18 @@ class MeshResource: fsdp_resource: Axis name for full-sharded data parallelism, default is None pp_resource: Axis name for pipeline parallelism (layer sharding), default is None cp_resource: Axis name for context parallelism (sequence sharding), default is None - ep_resource: Axis name for expert parallelism. Dispatch input tokens + ep_resource: Axis name or ordered tuple of axis names for expert + parallelism. A compound resource such as ``("expert", "tensor")`` + folds both physical axes into EP while preserving their order. + Dispatch input tokens must be sharded on their leading dim by ``ep_resource`` (alone or compound with ``dp_resource`` / ``fsdp_resource`` as outer, e.g. ``PartitionSpec(("dp", "ep"), None, None)``). Dispatch output ``[ep_size, recv_capacity, H]`` is always sharded by ``ep_resource`` on the leading ``ep_size`` dim. + etp_resource: Axis name for expert tensor parallelism. MoEBlock currently + supports this resource only when its mesh size is one, in which case + expert GEMM matrices remain complete. """ dp_resource: str = None @@ -345,7 +368,8 @@ class MeshResource: fsdp_resource: str = None pp_resource: str = None cp_resource: str = None - ep_resource: str = None + ep_resource: Optional[MeshAxis] = None + etp_resource: str = None _GLOBAL_MESH_RESOURCE = None @@ -385,7 +409,7 @@ def global_mesh_resource() -> MeshResource: return _GLOBAL_MESH_RESOURCE -def get_active_resource_axis(resource_name: str) -> Optional[str]: +def get_active_resource_axis(resource_name: str) -> Optional[MeshAxis]: """Resolve a :class:`MeshResource` attribute to its mesh axis name, or return ``None`` if that resource is not active. From 4dfb6b7c2dfc252d1d8ab9e2cb58944dff28cfa6 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 2 Sep 2026 09:19:57 -0700 Subject: [PATCH 2/5] [JAX] Keep TE EP MoE tests gated to Blackwell Signed-off-by: Jeremy Berchtold --- tests/jax/run_te_ep_moe.sh | 2 +- tests/jax/test_te_ep_moe.py | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh index 824ae752e5..8da242cc93 100755 --- a/tests/jax/run_te_ep_moe.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -102,7 +102,7 @@ for i in "${!EXITS[@]}"; do done # Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which the -# file emits via pytest.skip(allow_module_level=True) on pre-Hopper +# file emits via pytest.skip(allow_module_level=True) on pre-Blackwell # GPUs) as success. FAILED=0 for e in "${EXITS[@]}"; do diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index d1555ca0e2..7850b3fb29 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -31,7 +31,7 @@ on the block are pytest parametrize values rather than separate test classes: -* ``test_forward`` covers BF16 (Hopper+) and MXFP8 (Blackwell+) forward execution across a +* ``test_forward`` covers BF16 and MXFP8 forward execution across a curated set of configurations (softmax/sigmoid scoring, optional non-zero expert_bias). Each config asserts shape, dtype, finiteness and numerical parity vs the same BF16 reference in one run. @@ -110,11 +110,12 @@ def _read_mp_options(): from transformer_engine_jax import get_device_compute_capability -# NCCL EP requires Hopper or newer. BF16 grouped GEMM falls back to the -# legacy implementation on Hopper; MXFP8 remains gated to Blackwell below. -if get_device_compute_capability(0) < 90: +# Grouped GEMM in the MoE custom_vjp requires Blackwell (sm_100+). The +# TE EP NCCL primitives themselves need SM>=90, but the FFN body uses +# grouped_gemm, so the file as a whole gates on sm_100+. +if get_device_compute_capability(0) < 100: pytest.skip( - "MoE TE EP tests require Hopper (sm_90+) or newer", + "MoE TE EP tests require Blackwell (sm_100+) for grouped GEMM", allow_module_level=True, ) From be04a3824e1868c19e09152b0554b8a8167a99f7 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 2 Sep 2026 14:31:51 -0700 Subject: [PATCH 3/5] [JAX] Configure compound EP directly in MoEBlock Signed-off-by: Jeremy Berchtold --- qa/L0_jax_distributed_unittest/test.sh | 6 +- tests/jax/run_te_ep_moe.sh | 8 +- tests/jax/test_multi_process_ep.py | 36 +-------- tests/jax/test_te_ep_moe.py | 66 ++++++---------- transformer_engine/jax/cpp_extensions/ep.py | 68 ++++++++++------- transformer_engine/jax/ep.py | 31 ++++---- transformer_engine/jax/flax/moe.py | 30 ++++---- transformer_engine/jax/moe.py | 83 +++++++-------------- transformer_engine/jax/sharding.py | 16 ++-- 9 files changed, 128 insertions(+), 216 deletions(-) diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index 48c91455d3..0ab4e9a109 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -47,11 +47,11 @@ wait # >=4 visible GPUs. TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ || test_fail "test_te_ep_moe.py" -# Focused four-rank mesh: dense DP2 x TP2, folded MoE EP4, explicit ETP1. +# Focused four-rank mesh: dense DP2 x TP2, folded MoE EP4. if [ "$(nvidia-smi -L | wc -l)" -ge 4 ]; then - TE_EP_MOE_COMPOUND_ETP1=1 NUM_GPUS=4 TE_PATH=$TE_PATH \ + TE_EP_MOE_COMPOUND_EP=1 NUM_GPUS=4 TE_PATH=$TE_PATH \ bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ - || test_fail "test_te_ep_moe.py compound EP/ETP1" + || test_fail "test_te_ep_moe.py compound EP" fi # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs; # auto-skips otherwise). diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh index 8da242cc93..20e4973788 100755 --- a/tests/jax/run_te_ep_moe.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -18,13 +18,13 @@ TEST_FILE="${TEST_FILE:-$TE_ROOT/tests/jax/test_te_ep_moe.py}" PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" -COMPOUND_EP_ETP1="${TE_EP_MOE_COMPOUND_ETP1:-0}" +COMPOUND_EP="${TE_EP_MOE_COMPOUND_EP:-0}" if [ "$NUM_GPUS" -lt 4 ]; then echo "[run_te_ep_moe.sh] need >=4 GPUs (got $NUM_GPUS); SKIPPING." exit 0 fi -if [ "$COMPOUND_EP_ETP1" = "1" ] && [ "$NUM_GPUS" -ne 4 ]; then - echo "[run_te_ep_moe.sh] compound EP/ETP1 mode requires exactly 4 processes (got $NUM_GPUS); SKIPPING." +if [ "$COMPOUND_EP" = "1" ] && [ "$NUM_GPUS" -ne 4 ]; then + echo "[run_te_ep_moe.sh] compound EP mode requires exactly 4 processes (got $NUM_GPUS); SKIPPING." exit 0 fi @@ -36,7 +36,7 @@ echo "============================================================" echo "TE-EP MoE MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" echo " test file : $TEST_FILE" echo " coordinator : $TE_EP_MOE_COORDINATOR_ADDRESS" -echo " compound EP/ETP1 : $COMPOUND_EP_ETP1" +echo " compound EP : $COMPOUND_EP" echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" echo "============================================================" diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 3adbd2e20d..e7fb26dc63 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -48,19 +48,16 @@ ep_combine, _ep_domain_for_rank, ) -from transformer_engine.jax.moe import _moe_etp_axis from transformer_engine.jax.sharding import ( BATCH_AXES, W_TP_AXES, MeshResource, - get_active_resource_axis, get_mesh_axis_size, get_sharding_map_logic_axis_to_mesh_axis, global_shard_guard, ) from transformer_engine.jax.version_utils import is_collective_stream_supported - # ── Test config ───────────────────────────────────────────────────────────── # NCCL EP requires NUM_LOCAL_EXPERTS*ep % 4 == 0 (TMA alignment in # device/hybridep_adapter.cu:511). With NUM_LOCAL_EXPERTS=2, ep must be even. @@ -970,7 +967,7 @@ def test_ep_tp_splits_domains(self): self.assertEqual(domains, {0: [0, 2, 4, 6], 1: [1, 3, 5, 7]}) - def test_compound_ep_with_explicit_etp1(self): + def test_compound_ep_axis_group(self): if not is_devices_enough(4): self.skipTest("requires 4 devices") mesh = Mesh( @@ -983,20 +980,13 @@ def test_compound_ep_with_explicit_etp1(self): resource = MeshResource( dp_resource="expert", tp_resource="tensor", - ep_resource=("expert", "tensor"), - etp_resource="etp", ) with mesh, global_shard_guard(resource): - self.assertEqual(get_mesh_axis_size(resource.ep_resource), 4) - self.assertEqual(get_mesh_axis_size(resource.etp_resource), 1) - self.assertEqual(get_active_resource_axis("ep_resource"), ("expert", "tensor")) - self.assertIsNone(get_active_resource_axis("etp_resource")) + self.assertEqual(get_mesh_axis_size(("expert", "tensor")), 4) dense_rules = get_sharding_map_logic_axis_to_mesh_axis() self.assertEqual(dense_rules[BATCH_AXES], "expert") self.assertEqual(dense_rules[W_TP_AXES], "tensor") - # Dense DP is folded into compound EP here, not counted again as an - # outer group for the MoE communication buffers. - self.assertIsNone(_ep_outer_axis()) + self.assertIsNone(_ep_outer_axis(("expert", "tensor"))) domains = {} for rank in range(4): @@ -1008,26 +998,6 @@ def test_compound_ep_with_explicit_etp1(self): domains = {root: [m[c] for c in sorted(m)] for root, m in domains.items()} self.assertEqual(domains, {0: [0, 1, 2, 3]}) - # A size-one member is legal in a compound resource even if JAX later - # elides it from a concrete PartitionSpec. - etp1_domains = {} - for rank in range(4): - root, col, ndom = _ep_domain_for_rank( - mesh, ("expert", "etp"), rank, device_to_rank=d2r - ) - self.assertEqual(ndom, 2) - etp1_domains.setdefault(root, {})[col] = rank - etp1_domains = { - root: [members[c] for c in sorted(members)] - for root, members in etp1_domains.items() - } - self.assertEqual(etp1_domains, {0: [0, 2], 1: [1, 3]}) - with mesh, global_shard_guard( - MeshResource(ep_resource=("expert", "etp"), etp_resource="etp") - ): - self.assertEqual(get_active_resource_axis("ep_resource"), ("expert", "etp")) - self.assertIsNone(_moe_etp_axis(("expert", "etp"), "etp")) - # ── Entry point ────────────────────────────────────────────────────────────── diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 7850b3fb29..9525de03c8 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -130,27 +130,24 @@ def _read_mp_options(): from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.jax.sharding import MeshResource, global_shard_guard - # ----------------------------------------------------------------------------- # Mesh / shape config # ----------------------------------------------------------------------------- -COMPOUND_EP_ETP1 = os.environ.get("TE_EP_MOE_COMPOUND_ETP1", "0") == "1" +COMPOUND_EP = os.environ.get("TE_EP_MOE_COMPOUND_EP", "0") == "1" -if COMPOUND_EP_ETP1: +if COMPOUND_EP: EP_AXIS = ("expert", "tensor") EP_SIZE = 4 MESH_SHAPE = (2, 2, 1) MESH_AXIS_NAMES = ("expert", "tensor", "etp") - # Preserve the dense-layer DP declaration. MoEBlock must recognize that - # this physical axis is already folded into compound EP and not count it twice. - DATA_PARALLELISM_AXES = ("expert",) + # Both dense DP and TP axes are folded into EP for this MoE region, leaving + # no additional data-parallel axis outside the EP communicator. + DATA_PARALLELISM_AXES = () BATCH_MESH_AXIS = EP_AXIS MESH_RESOURCE = MeshResource( dp_resource="expert", tp_resource="tensor", - ep_resource=EP_AXIS, - etp_resource="etp", ) LOGICAL_AXIS_RULES = ( ("exp", EP_AXIS), @@ -220,15 +217,15 @@ def _read_mp_options(): @pytest.fixture(scope="module") def mesh(): if jax.device_count() < NUM_DEVICES_REQUIRED or ( - COMPOUND_EP_ETP1 and jax.device_count() != NUM_DEVICES_REQUIRED + COMPOUND_EP and jax.device_count() != NUM_DEVICES_REQUIRED ): pytest.skip( - f"Need {'exactly' if COMPOUND_EP_ETP1 else '>='} {NUM_DEVICES_REQUIRED} devices for" + f"Need {'exactly' if COMPOUND_EP else '>='} {NUM_DEVICES_REQUIRED} devices for" f" mesh={MESH_SHAPE}; have {jax.device_count()}" ) # Compound EP axes retain their declared order. With shape (2, 2, 1), # (expert, tensor) flattens to contiguous ranks [0, 1, 2, 3]. - if COMPOUND_EP_ETP1: + if COMPOUND_EP: devices = np.asarray(jax.devices()).reshape(MESH_SHAPE) else: devices = mesh_utils.create_device_mesh(MESH_SHAPE) @@ -258,6 +255,7 @@ def mesh(): recv_capacity_per_rank=recv_capacity_per_rank, hidden_dim=HIDDEN, max_token_dtype=DTYPE, + ep_axis=EP_AXIS if COMPOUND_EP else None, ) record_ep_bootstrap_signature_for_moe( num_experts=NUM_EXPERTS, @@ -398,6 +396,7 @@ def _make_block( num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, intermediate_size=INTER, + ep_axis=EP_AXIS if COMPOUND_EP else None, data_parallelism_axes=DATA_PARALLELISM_AXES, apply_topk_weights_early=apply_topk_weights_early, aux_loss_coeff=aux_loss_coeff, @@ -454,10 +453,10 @@ def __exit__(self_inner, *args): def _init_apply(block, mesh, x, key): with _ctx(mesh): x_sh = _shard_inputs(x, mesh) - if COMPOUND_EP_ETP1: + if COMPOUND_EP: assert _axis_names(_spec_entry(x_sh.sharding.spec, 0)) == frozenset(EP_AXIS) assert not any(_axis_names(entry) for entry in x_sh.sharding.spec[1:]) - if COMPOUND_EP_ETP1: + if COMPOUND_EP: abstract_variables = jax.eval_shape(block.init, key, x_sh) logical_specs = flax_linen.get_partition_spec(abstract_variables) variable_shardings = flax_linen.logical_to_mesh_sharding( @@ -557,8 +556,8 @@ def _slice_size(index, global_size): return stop - start -def _assert_compound_ep_etp1_sharding(variables, output): - """Assert TP is folded into EP while every local expert matrix is complete.""" +def _assert_compound_ep_sharding(variables, output): + """Assert TP is folded into EP while size-one ETP leaves matrices complete.""" assert _axis_names(_spec_entry(output.sharding.spec, 0)) == frozenset(EP_AXIS) assert not any(_axis_names(entry) for entry in output.sharding.spec[1:]) @@ -567,9 +566,9 @@ def _assert_compound_ep_etp1_sharding(variables, output): wi = _unwrap(params["wi"]) wo = _unwrap(params["wo"]) - assert not any(_axis_names(entry) for entry in gate.sharding.spec), ( - "gate weights should be replicated for global routing" - ) + assert not any( + _axis_names(entry) for entry in gate.sharding.spec + ), "gate weights should be replicated for global routing" assert _axis_names(_spec_entry(wi.sharding.spec, 0)) == frozenset(EP_AXIS) assert _axis_names(_spec_entry(wo.sharding.spec, 0)) == frozenset(EP_AXIS) # JAX may preserve or elide a size-one axis. Either representation must @@ -640,7 +639,7 @@ def _quantization_recipe(quantization): ), ] -if COMPOUND_EP_ETP1: +if COMPOUND_EP: # This mode is a focused sharding qualification, not a repeat of the # router feature matrix covered by the default topology. _CONFIGS = [_CONFIGS[0]] @@ -649,7 +648,7 @@ def _quantization_recipe(quantization): pytest.param("bf16", id="bf16"), ] -if get_device_compute_capability(0) >= 100 and not COMPOUND_EP_ETP1: +if get_device_compute_capability(0) >= 100 and not COMPOUND_EP: _QUANTIZATION_CASES.append(pytest.param("mxfp8", id="mxfp8")) @@ -676,8 +675,8 @@ def test_forward(self, mesh, config, quantization): x = _make_inputs(jax.random.PRNGKey(0)) variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) - if COMPOUND_EP_ETP1: - _assert_compound_ep_etp1_sharding(variables, output) + if COMPOUND_EP: + _assert_compound_ep_sharding(variables, output) # Shape / dtype / finiteness (cheap; on the local shard). assert output.shape == x.shape @@ -707,27 +706,6 @@ def test_forward(self, mesh, config, quantization): err_msg=f"forward parity breach for config={config}, quantization={quantization}", ) - def test_active_etp_is_rejected(self, mesh): - if not COMPOUND_EP_ETP1: - pytest.skip("only applies to the explicit ETP resource mode") - active_etp_mesh = Mesh(mesh.devices.reshape(1, 2, 2), MESH_AXIS_NAMES) - resource = MeshResource( - dp_resource="expert", - tp_resource="tensor", - ep_resource=("expert", "tensor"), - etp_resource="etp", - ) - rules = ( - ("exp", ("expert", "tensor")), - ("embed", None), - ("mlp", "etp"), - ("batch", ("expert", "tensor")), - ) - block = _make_block() - with active_etp_mesh, global_shard_guard(resource), nn_partitioning.axis_rules(rules): - with pytest.raises(NotImplementedError, match="ETP=1"): - block.init(jax.random.PRNGKey(40), _make_inputs(jax.random.PRNGKey(41))) - class TestTeEpMoeBackward: """Per-config backward correctness in a single run: per-tensor @@ -806,7 +784,7 @@ def loss_fn(params, x): ) -@pytest.mark.skipif(COMPOUND_EP_ETP1, reason="compound mode focuses on EP/ETP sharding") +@pytest.mark.skipif(COMPOUND_EP, reason="compound mode focuses on EP sharding") class TestTeEpMoeAuxLoss: """Aux-loss path. Consolidated into: * ``test_aux_loss``: one run that checks the returned scalar's diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index c184cf0977..7837940767 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -5,8 +5,8 @@ Sharding model: - EpPrepare / EpDispatch outputs carry a single leading ``num_procs`` dim. - Sharded compound ``(dp_resource, ep_resource)`` when DP is set, else - ``ep_resource`` alone. + Sharded by the bootstrapped ``ep_axis`` group, optionally preceded by an + outer DP/FSDP axis. - EpDispatch inputs are 2D ``[T, H]`` or 3D ``[B, S, H]``; only the first dim may be sharded, with axis in {ep, (dp, ep), dp, None}. Trailing dims must be replicated. ``dp`` alone gets ``ep`` folded in locally. @@ -24,7 +24,7 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive -from ..sharding import global_mesh_resource, get_mesh_axis_size, normalize_mesh_axes +from ..sharding import MeshAxis, global_mesh_resource, get_mesh_axis_size, normalize_mesh_axes from ..version_utils import is_collective_stream_supported @@ -73,12 +73,15 @@ def wrapper(*args, **kwargs): class EpConfig: """Snapshot of the EP bootstrap config (see ep_bootstrap). - num_ep_groups is the size of the outer dp/fsdp mesh axis (1 if neither - is set), captured at bootstrap so abstract-eval never reads the mesh. + ``ep_axis`` is the ordered physical mesh axis group used by the EP + communicator. ``num_ep_groups`` is the size of the outer dp/fsdp mesh + axis (1 if neither is set), captured at bootstrap so abstract-eval never + reads the mesh. """ world_size: int rank: int + ep_axis: MeshAxis ep_size: int num_ep_groups: int num_experts: int @@ -135,31 +138,31 @@ def ep_handle_mem_size(cfg: EpLayerConfig) -> int: def _leading_axis_ok(spec): """Validate an EP input spec; return ``(ok, ep_axis, outer_axes)``. - Leading dim is ``ep`` or a tuple ending in ``ep`` (outer dp/fsdp axes - first); all other dims must be replicated. + Leading dim contains the ordered EP axis group, optionally preceded by an + outer DP/FSDP axis; all other dims must be replicated. """ gsr = global_mesh_resource() - ep_axes = normalize_mesh_axes(gsr.ep_resource) + ep_axis = get_ep_config().ep_axis + ep_axes = normalize_mesh_axes(ep_axis) outer_axes = tuple( - a - for a in (gsr.dp_resource, gsr.fsdp_resource) - if a is not None and a not in ep_axes + a for a in (gsr.dp_resource, gsr.fsdp_resource) if a is not None and a not in ep_axes ) if len(spec) < 2 or not ep_axes: - return False, gsr.ep_resource, outer_axes + return False, ep_axis, outer_axes if any(ax is not None for ax in spec[1:]): - return False, gsr.ep_resource, outer_axes + return False, ep_axis, outer_axes leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) actual = set(a for a in elts if a is not None) - required_ep = {a for a in ep_axes if get_mesh_axis_size(a) > 1} - if not required_ep.issubset(actual): - return False, gsr.ep_resource, outer_axes + active_ep_axes = tuple(a for a in ep_axes if get_mesh_axis_size(a) > 1) + actual_ep_axes = tuple(a for a in elts if a in active_ep_axes) + if actual_ep_axes != active_ep_axes: + return False, ep_axis, outer_axes allowed = set(outer_axes) | set(ep_axes) - return actual.issubset(allowed), gsr.ep_resource, outer_axes + return actual.issubset(allowed), ep_axis, outer_axes -def _ep_outer_axis(): +def _ep_outer_axis(ep_axis=None): """The single dp/fsdp axis (if any) sitting outside ep on EP-output tensors. When set, EP-output globals carry an extra leading ``dp_size`` dim so SPMD @@ -169,7 +172,9 @@ def _ep_outer_axis(): we don't pin EP-output specs to a degenerate axis that JAX may collapse. """ gsr = global_mesh_resource() - ep_axes = set(normalize_mesh_axes(gsr.ep_resource)) + if ep_axis is None: + ep_axis = get_ep_config().ep_axis + ep_axes = set(normalize_mesh_axes(ep_axis)) candidates = tuple( axis for axis in (gsr.dp_resource, gsr.fsdp_resource) @@ -191,23 +196,22 @@ def _ep_leading_dims(is_outer): def _ep_output_spec(*trailing): - """PartitionSpec for an EP-output tensor: ``(("dp","ep"), *trailing)`` when - DP is set (compound leading axis on a single dim), else ``("ep",*trailing)``.""" - gsr = global_mesh_resource() + """Build an EP-output spec from the bootstrapped EP and outer axis groups.""" + ep_axis = get_ep_config().ep_axis outer = _ep_outer_axis() - ep_axes = normalize_mesh_axes(gsr.ep_resource) + ep_axes = normalize_mesh_axes(ep_axis) leading_axes = ep_axes if outer is None else (outer, *ep_axes) leading = leading_axes[0] if len(leading_axes) == 1 else leading_axes return PartitionSpec(leading, *trailing) def _ep_spec_ok(spec, trailing_count): - """Leading dim shards along ep (and outer dp/fsdp when set); trailing dims - are replicated. JAX may collapse size-1 mesh axes to ``None`` or drop them, - so the leading entry is normalized to a set of named axes before comparing. + """Validate an EP output's leading axis group and replicated trailing dims. + + JAX may collapse size-one mesh axes to ``None`` or drop them, so only + nontrivial EP axes are required. """ - gsr = global_mesh_resource() - ep_axes = normalize_mesh_axes(gsr.ep_resource) + ep_axes = normalize_mesh_axes(get_ep_config().ep_axis) outer = _ep_outer_axis() if len(spec) != 1 + trailing_count: return False @@ -216,10 +220,16 @@ def _ep_spec_ok(spec, trailing_count): leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) actual = frozenset(a for a in elts if a is not None) + active_ep_axes = tuple(a for a in ep_axes if get_mesh_axis_size(a) > 1) + actual_ep_axes = tuple(a for a in elts if a in active_ep_axes) + if actual_ep_axes != active_ep_axes: + return False expected = set(ep_axes) if outer is not None: expected.add(outer) - required = {axis for axis in expected if get_mesh_axis_size(axis) > 1} + required = set(active_ep_axes) + if outer is not None and get_mesh_axis_size(outer) > 1: + required.add(outer) return required.issubset(actual) and actual.issubset(expected) diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 157cdd4d94..9f5e14ee62 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -121,13 +121,14 @@ def ep_bootstrap( max_token_dtype=jnp.bfloat16, max_num_sms=0, drop_on_overflow=False, + ep_axis=None, ): """Initialize the EP communicator. Call once per process before any EP op. - Must run inside the active JAX Mesh and a global_shard_guard; ep_size and - num_ep_groups are read from the mesh axes named by MeshResource.ep_resource - and MeshResource.dp_resource/fsdp_resource. Axes orthogonal to EP (tp, pp, - cp, ...) are supported and replicated across EP tensors. + Must run inside the active JAX Mesh and a global_shard_guard. By default, + the EP axis is read from ``MeshResource.ep_resource``. ``ep_axis`` may + override it with an ordered tuple of physical mesh axes for a compound EP + communicator. DP/FSDP axes outside that group determine ``num_ep_groups``. Args: world_size: Total number of processes (product of all mesh axes). @@ -142,6 +143,8 @@ def ep_bootstrap( drop_on_overflow: Drop tokens exceeding recv_capacity_per_rank instead of trapping on overflow. Dropped tokens are still counted in total_recv_tokens, so callers can detect overflow from it. + ep_axis: Optional physical mesh axis name or ordered tuple of names. + Defaults to ``MeshResource.ep_resource``. """ if jnp.dtype(max_token_dtype) != jnp.bfloat16: raise NotImplementedError( @@ -161,12 +164,9 @@ def ep_bootstrap( ) gsr = global_mesh_resource() - ep_resource = gsr.ep_resource + ep_resource = gsr.ep_resource if ep_axis is None else ep_axis if ep_resource is None: - raise ValueError( - "ep_bootstrap requires MeshResource.ep_resource to be set; enter a" - " global_shard_guard(MeshResource(..., ep_resource=)) before bootstrap." - ) + raise ValueError("ep_bootstrap requires ep_axis or MeshResource.ep_resource to be set.") mesh = _get_mesh() if mesh.empty: raise ValueError( @@ -181,7 +181,7 @@ def ep_bootstrap( ep_size = get_mesh_axis_size(ep_resource) # num_ep_groups counts only the distinct-token (dp/fsdp) axes; replicated # axes (tp, pp, ...) do not create distinct EP-output slabs. - outer_axis = _ep_outer_axis() + outer_axis = _ep_outer_axis(ep_resource) num_ep_groups = 1 if outer_axis is None else get_mesh_axis_size(outer_axis) if num_experts % ep_size != 0: raise ValueError(f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size}).") @@ -226,6 +226,7 @@ def ep_bootstrap( tex.ep.EpConfig( world_size=world_size, rank=rank, + ep_axis=ep_resource, ep_size=ep_size, num_ep_groups=num_ep_groups, num_experts=num_experts, @@ -252,14 +253,10 @@ def ep_finalize(): def _default_out_partition_spec(): - """Leading-axis default: ``(("dp","ep"),)`` if DP/FSDP is set, else ``("ep",)``.""" - gsr = global_mesh_resource() - if gsr.ep_resource is None: - raise ValueError( - "ep_resource is not set on the active MeshResource; pass out_sharding=... explicitly." - ) + """Build the default leading-axis spec from the bootstrapped EP group.""" + ep_axis = tex.ep.get_ep_config().ep_axis outer = _ep_outer_axis() - ep_axes = normalize_mesh_axes(gsr.ep_resource) + ep_axes = normalize_mesh_axes(ep_axis) leading_axes = ep_axes if outer is None else (outer, *ep_axes) leading = leading_axes[0] if len(leading_axes) == 1 else leading_axes return (leading,) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index ecebe7ce60..29158e1c6b 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -13,7 +13,7 @@ as ``self.param`` slots (with the right :func:`flax.linen.with_logical_partitioning` annotations so JAX's sharding layer FSDPs the params correctly). -2. Resolve compound EP and ETP axis resources from the active +2. Resolve the EP axis from an explicit layer setting or the active :class:`transformer_engine.jax.sharding.MeshResource`. 3. Forward all knobs to :func:`moe`. @@ -34,14 +34,12 @@ from flax import linen as nn from transformer_engine.common.recipe import Recipe -from ..moe import _moe_outer_axes, moe +from ..moe import MeshAxis, moe from ..quantize import QuantizerSet from ..router import ScoreFunction from ..sharding import ( _get_mesh, get_active_resource_axis, - get_mesh_axis_size, - global_mesh_resource, ) from .module import TransformerEngineBase @@ -102,8 +100,11 @@ class _MoEBlock(TransformerEngineBase): ADDITION to the EP axis. Empty (default) means activations are replicated across non-EP axes within an EP group; set e.g. ``("fsdp",)`` for true FSDP-of-batch where each device owns a - unique slice of the batch. Any dense DP axis already included in a - compound EP resource is ignored here rather than counted twice. + unique slice of the batch. These axes must be outside ``ep_axis``. + ep_axis : Optional[str | tuple[str, ...]] + Physical mesh axis name, or ordered compound axis group, used for this + MoE block's EP communicator. Defaults to the active + ``MeshResource.ep_resource`` for backward compatibility. apply_topk_weights_early : bool If ``True``, multiply expert outputs by their top-k weights *inside* each shard before ``ep_combine`` (saves one global @@ -152,6 +153,7 @@ class _MoEBlock(TransformerEngineBase): input_axes: Tuple[Optional[str], ...] = () # Parallelism + ep_axis: Optional[MeshAxis] = None data_parallelism_axes: Tuple[str, ...] = () # MoE knobs forwarded to ``moe()`` @@ -259,17 +261,12 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: jnp.float32, ) - ep_axis = get_active_resource_axis("ep_resource") + ep_axis = self.ep_axis + if ep_axis is None: + ep_axis = get_active_resource_axis("ep_resource") mesh = _get_mesh() - etp_axis = global_mesh_resource().etp_resource - if etp_axis is not None and get_mesh_axis_size(etp_axis, mesh) != 1: - raise NotImplementedError( - "_MoEBlock currently supports expert tensor parallelism only with ETP=1; " - f"axis {etp_axis!r} has size {get_mesh_axis_size(etp_axis, mesh)}." - ) data_parallel_size = 1 - effective_data_parallelism_axes = _moe_outer_axes(ep_axis, self.data_parallelism_axes) - for axis in effective_data_parallelism_axes: + for axis in self.data_parallelism_axes: data_parallel_size *= mesh.shape[axis] def make_grouped_quantizer_set(postfix): @@ -318,8 +315,7 @@ def make_grouped_quantizer_set(postfix): quantizer_sets=quantizer_sets, recv_capacity_per_rank=self.recv_capacity_per_rank, ep_axis=ep_axis, - etp_axis=etp_axis, - data_parallelism_axes=effective_data_parallelism_axes, + data_parallelism_axes=self.data_parallelism_axes, input_axes=self.input_axes, gate_kernel_axes=self.gate_kernel_axes, wi_kernel_axes=self.wi_kernel_axes, diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index c520074eaa..204983776e 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -64,34 +64,14 @@ _ALIGN_SIZE = 128 -def _moe_outer_axes( - ep_axis: MeshAxis, data_parallelism_axes: Tuple[str, ...] -) -> Tuple[str, ...]: - """Drop dense DP axes that are folded into the compound EP resource.""" - ep_axes = set(normalize_mesh_axes(ep_axis)) - return tuple(axis for axis in data_parallelism_axes if axis not in ep_axes) - - def _moe_leading_axis(ep_axis: MeshAxis, data_parallelism_axes: Tuple[str, ...]): """Build one PartitionSpec entry with outer axes followed by compound EP.""" - axes = (*_moe_outer_axes(ep_axis, data_parallelism_axes), *normalize_mesh_axes(ep_axis)) + axes = (*data_parallelism_axes, *normalize_mesh_axes(ep_axis)) if not axes: raise ValueError("moe(...) requires ep_axis to contain at least one mesh axis.") return axes[0] if len(axes) == 1 else axes -def _moe_etp_axis(ep_axis: MeshAxis, etp_axis: Optional[str]) -> Optional[str]: - """Avoid assigning one physical axis to two tensor dimensions. - - A size-one ETP axis may legally appear inside a compound EP resource. JAX - can elide that axis, but it cannot name it on both the expert and matrix - dimensions of the same PartitionSpec. - """ - if etp_axis in normalize_mesh_axes(ep_axis): - return None - return etp_axis - - def get_moe_recv_capacity_per_rank( *, num_experts: int, @@ -589,7 +569,6 @@ def _moe_fwd_rule( scaling_factor, aux_loss_coeff, ep_axis, - etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -614,12 +593,18 @@ def _moe_fwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") if ep_axis is None: raise ValueError("moe(...) requires ep_axis to be set (TE EP backend).") + ep_axes = normalize_mesh_axes(ep_axis) + bootstrap_ep_axes = normalize_mesh_axes(tex.ep.get_ep_config().ep_axis) + if ep_axes != bootstrap_ep_axes: + raise ValueError( + f"moe(...) ep_axis={ep_axes} does not match the bootstrapped EP axes " + f"{bootstrap_ep_axes}." + ) num_ep = get_mesh_axis_size(ep_axis, mesh) if num_experts % num_ep != 0: raise ValueError(f"num_experts={num_experts} must be divisible by EP size={num_ep}") num_local_experts = num_experts // num_ep - data_parallelism_axes = _moe_outer_axes(ep_axis, data_parallelism_axes) dp_size = 1 for ax in data_parallelism_axes: dp_size *= mesh.shape[ax] @@ -777,15 +762,12 @@ def _moe_fwd_rule( # ---------------- FFN (per-shard via shard_map) ---------------- has_bias = wi_0_bias is not None - ffn_etp_axis = _moe_etp_axis(ep_axis, etp_axis) - wi_spec = P(ep_axis, None, ffn_etp_axis) - wo_spec = P(ep_axis, ffn_etp_axis, None) - wi_bias_spec = P(ep_axis, ffn_etp_axis) - wo_bias_spec = P(ep_axis, None) - ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, wi_spec, wo_spec) + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec) ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi, wo] if has_bias: - ffn_in_specs += (wi_bias_spec, wi_bias_spec, wo_bias_spec) + ffn_in_specs += (bias_spec, bias_spec, bias_spec) ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) # Quantized grouped tensors store their data, scales, and group metadata @@ -915,7 +897,6 @@ def _moe_bwd_rule( scaling_factor, aux_loss_coeff, ep_axis, - etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -944,7 +925,6 @@ def _moe_bwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") B, S, _ = x_shape K = num_experts_per_tok - data_parallelism_axes = _moe_outer_axes(ep_axis, data_parallelism_axes) batch_pspec_axis: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) @@ -968,11 +948,8 @@ def _moe_bwd_rule( d_recv_w_from_combine = d_recv_w_from_combine.astype(ctx.recv_topk_weights.dtype) # ---------------- FFN bwd (per-shard via shard_map) ---------------- - ffn_etp_axis = _moe_etp_axis(ep_axis, etp_axis) - wi_spec = P(ep_axis, None, ffn_etp_axis) - wo_spec = P(ep_axis, ffn_etp_axis, None) - wi_bias_spec = P(ep_axis, ffn_etp_axis) - wo_bias_spec = P(ep_axis, None) + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) token_buffer_spec = P(batch_pspec_axis) token_matrix_spec = P(batch_pspec_axis, None) expert_buffer_spec = P(ep_axis) @@ -1037,14 +1014,14 @@ def _ffn_bwd_body(*args): bwd_out_specs = ( ep3_spec, ep2_spec, - wi_spec, - wo_spec, - wi_bias_spec, - wi_bias_spec, - wo_bias_spec, + kernel_spec, + kernel_spec, + bias_spec, + bias_spec, + bias_spec, ) else: - bwd_out_specs = (ep3_spec, ep2_spec, wi_spec, wo_spec, None, None, None) + bwd_out_specs = (ep3_spec, ep2_spec, kernel_spec, kernel_spec, None, None, None) ( d_sorted_x, @@ -1172,7 +1149,7 @@ def _ffn_bwd_body(*args): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 28))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 27))) def _moe( x, gate_kernel, @@ -1193,7 +1170,6 @@ def _moe( scaling_factor, aux_loss_coeff, ep_axis, - etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -1223,7 +1199,6 @@ def _moe( scaling_factor, aux_loss_coeff, ep_axis, - etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, @@ -1264,7 +1239,6 @@ def moe( noop_quantizer_set, ), ep_axis: MeshAxis, - etp_axis: Optional[str] = None, data_parallelism_axes: Tuple[str, ...] = (), input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), @@ -1313,11 +1287,10 @@ def moe( Axis-name parameters: - * ``ep_axis``, ``etp_axis``, and ``data_parallelism_axes`` are *physical mesh - axis names*. ``ep_axis`` may be an ordered tuple whose sizes are - multiplied into one compound EP resource. ETP is currently accepted - only when its mesh size is one, leaving expert GEMM matrices complete. - These concrete axes are used to compute ``num_ep`` / ``dp_size`` and construct + * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh axis names*. + ``ep_axis`` may be an ordered tuple whose sizes are multiplied into one + compound EP group. These concrete axes are used to compute + ``num_ep`` / ``dp_size`` and construct ``P((dp..., ep), None, None)`` for the physical ``jax.lax.with_sharding_constraint`` calls that JAX requires to refer to real mesh axes). @@ -1350,11 +1323,6 @@ def moe( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - if etp_axis is not None and get_mesh_axis_size(etp_axis, mesh) != 1: - raise NotImplementedError( - "moe(...) currently supports expert tensor parallelism only with ETP=1; " - f"axis {etp_axis!r} has size {get_mesh_axis_size(etp_axis, mesh)}." - ) expected_leading: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) expected_spec = P(expected_leading, None, None) actual_spec = getattr(getattr(x, "sharding", None), "spec", None) @@ -1396,7 +1364,6 @@ def moe( scaling_factor, float(aux_loss_coeff), ep_axis, - etp_axis, data_parallelism_axes, input_axes, gate_kernel_axes, diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 18b298d1f7..10e439048e 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -9,6 +9,7 @@ parallelism (FSDP). It includes functions for sharding constraints, mesh management, and collective operations. """ + from contextlib import contextmanager from dataclasses import dataclass from typing import Callable, Optional, Union @@ -40,7 +41,7 @@ def normalize_mesh_axes(axis: Optional[MeshAxis]) -> tuple[str, ...]: - """Return a mesh resource as an ordered tuple of physical axis names.""" + """Return a mesh axis or axis group as an ordered tuple of physical names.""" if axis is None: return () axes = axis if isinstance(axis, tuple) else (axis,) @@ -348,18 +349,12 @@ class MeshResource: fsdp_resource: Axis name for full-sharded data parallelism, default is None pp_resource: Axis name for pipeline parallelism (layer sharding), default is None cp_resource: Axis name for context parallelism (sequence sharding), default is None - ep_resource: Axis name or ordered tuple of axis names for expert - parallelism. A compound resource such as ``("expert", "tensor")`` - folds both physical axes into EP while preserving their order. - Dispatch input tokens + ep_resource: Axis name for expert parallelism. Dispatch input tokens must be sharded on their leading dim by ``ep_resource`` (alone or compound with ``dp_resource`` / ``fsdp_resource`` as outer, e.g. ``PartitionSpec(("dp", "ep"), None, None)``). Dispatch output ``[ep_size, recv_capacity, H]`` is always sharded by ``ep_resource`` on the leading ``ep_size`` dim. - etp_resource: Axis name for expert tensor parallelism. MoEBlock currently - supports this resource only when its mesh size is one, in which case - expert GEMM matrices remain complete. """ dp_resource: str = None @@ -368,8 +363,7 @@ class MeshResource: fsdp_resource: str = None pp_resource: str = None cp_resource: str = None - ep_resource: Optional[MeshAxis] = None - etp_resource: str = None + ep_resource: str = None _GLOBAL_MESH_RESOURCE = None @@ -409,7 +403,7 @@ def global_mesh_resource() -> MeshResource: return _GLOBAL_MESH_RESOURCE -def get_active_resource_axis(resource_name: str) -> Optional[MeshAxis]: +def get_active_resource_axis(resource_name: str) -> Optional[str]: """Resolve a :class:`MeshResource` attribute to its mesh axis name, or return ``None`` if that resource is not active. From 990b223ab71fc75599f7113ad8d9ab751b89b311 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 2 Sep 2026 14:44:47 -0700 Subject: [PATCH 4/5] Revert compound EP MeshResource changes Signed-off-by: Jeremy Berchtold --- qa/L0_jax_distributed_unittest/test.sh | 6 - tests/jax/run_te_ep_moe.sh | 6 - tests/jax/test_multi_process_ep.py | 52 +------ tests/jax/test_te_ep_moe.py | 163 +++++--------------- transformer_engine/jax/cpp_extensions/ep.py | 89 ++++------- transformer_engine/jax/ep.py | 52 +++---- transformer_engine/jax/flax/moe.py | 20 +-- transformer_engine/jax/moe.py | 50 +++--- transformer_engine/jax/sharding.py | 28 +--- 9 files changed, 125 insertions(+), 341 deletions(-) diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index 0ab4e9a109..a563e6908d 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -47,12 +47,6 @@ wait # >=4 visible GPUs. TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ || test_fail "test_te_ep_moe.py" -# Focused four-rank mesh: dense DP2 x TP2, folded MoE EP4. -if [ "$(nvidia-smi -L | wc -l)" -ge 4 ]; then - TE_EP_MOE_COMPOUND_EP=1 NUM_GPUS=4 TE_PATH=$TE_PATH \ - bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ - || test_fail "test_te_ep_moe.py compound EP" -fi # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs; # auto-skips otherwise). CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)" diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh index 20e4973788..9fcbb9dd9a 100755 --- a/tests/jax/run_te_ep_moe.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -18,15 +18,10 @@ TEST_FILE="${TEST_FILE:-$TE_ROOT/tests/jax/test_te_ep_moe.py}" PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" -COMPOUND_EP="${TE_EP_MOE_COMPOUND_EP:-0}" if [ "$NUM_GPUS" -lt 4 ]; then echo "[run_te_ep_moe.sh] need >=4 GPUs (got $NUM_GPUS); SKIPPING." exit 0 fi -if [ "$COMPOUND_EP" = "1" ] && [ "$NUM_GPUS" -ne 4 ]; then - echo "[run_te_ep_moe.sh] compound EP mode requires exactly 4 processes (got $NUM_GPUS); SKIPPING." - exit 0 -fi export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}" export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.5}" @@ -36,7 +31,6 @@ echo "============================================================" echo "TE-EP MoE MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" echo " test file : $TEST_FILE" echo " coordinator : $TE_EP_MOE_COORDINATOR_ADDRESS" -echo " compound EP : $COMPOUND_EP" echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" echo "============================================================" diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index e7fb26dc63..47af0b0c39 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -33,13 +33,7 @@ from jax.sharding import Mesh, NamedSharding, PartitionSpec from utils import is_devices_enough -from transformer_engine.jax.cpp_extensions.ep import ( - ep_prepare, - ep_dispatch_fwd, - ep_combine_fwd, - get_ep_config, - _ep_outer_axis, -) +from transformer_engine.jax.sharding import MeshResource, global_shard_guard from transformer_engine.jax.ep import ( EpLayerConfig, ep_bootstrap, @@ -48,16 +42,15 @@ ep_combine, _ep_domain_for_rank, ) -from transformer_engine.jax.sharding import ( - BATCH_AXES, - W_TP_AXES, - MeshResource, - get_mesh_axis_size, - get_sharding_map_logic_axis_to_mesh_axis, - global_shard_guard, +from transformer_engine.jax.cpp_extensions.ep import ( + ep_prepare, + ep_dispatch_fwd, + ep_combine_fwd, + get_ep_config, ) from transformer_engine.jax.version_utils import is_collective_stream_supported + # ── Test config ───────────────────────────────────────────────────────────── # NCCL EP requires NUM_LOCAL_EXPERTS*ep % 4 == 0 (TMA alignment in # device/hybridep_adapter.cu:511). With NUM_LOCAL_EXPERTS=2, ep must be even. @@ -967,37 +960,6 @@ def test_ep_tp_splits_domains(self): self.assertEqual(domains, {0: [0, 2, 4, 6], 1: [1, 3, 5, 7]}) - def test_compound_ep_axis_group(self): - if not is_devices_enough(4): - self.skipTest("requires 4 devices") - mesh = Mesh( - np.asarray(jax.devices()[:4]).reshape(2, 2, 1), - ("expert", "tensor", "etp"), - ) - order = {int(d.id): i for i, d in enumerate(mesh.devices.reshape(-1))} - d2r = lambda d: order[int(d.id)] - - resource = MeshResource( - dp_resource="expert", - tp_resource="tensor", - ) - with mesh, global_shard_guard(resource): - self.assertEqual(get_mesh_axis_size(("expert", "tensor")), 4) - dense_rules = get_sharding_map_logic_axis_to_mesh_axis() - self.assertEqual(dense_rules[BATCH_AXES], "expert") - self.assertEqual(dense_rules[W_TP_AXES], "tensor") - self.assertIsNone(_ep_outer_axis(("expert", "tensor"))) - - domains = {} - for rank in range(4): - root, col, ndom = _ep_domain_for_rank( - mesh, ("expert", "tensor"), rank, device_to_rank=d2r - ) - self.assertEqual(ndom, 1) - domains.setdefault(root, {})[col] = rank - domains = {root: [m[c] for c in sorted(m)] for root, m in domains.items()} - self.assertEqual(domains, {0: [0, 1, 2, 3]}) - # ── Entry point ────────────────────────────────────────────────────────────── diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 9525de03c8..015c73343a 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -57,7 +57,6 @@ from jax.experimental import mesh_utils from jax.sharding import Mesh, NamedSharding, PartitionSpec as P -from flax import linen as flax_linen from flax.linen import partitioning as nn_partitioning @@ -130,57 +129,31 @@ def _read_mp_options(): from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.jax.sharding import MeshResource, global_shard_guard + # ----------------------------------------------------------------------------- # Mesh / shape config # ----------------------------------------------------------------------------- -COMPOUND_EP = os.environ.get("TE_EP_MOE_COMPOUND_EP", "0") == "1" - -if COMPOUND_EP: - EP_AXIS = ("expert", "tensor") - EP_SIZE = 4 - MESH_SHAPE = (2, 2, 1) - MESH_AXIS_NAMES = ("expert", "tensor", "etp") - # Both dense DP and TP axes are folded into EP for this MoE region, leaving - # no additional data-parallel axis outside the EP communicator. - DATA_PARALLELISM_AXES = () - BATCH_MESH_AXIS = EP_AXIS - MESH_RESOURCE = MeshResource( - dp_resource="expert", - tp_resource="tensor", - ) - LOGICAL_AXIS_RULES = ( - ("exp", EP_AXIS), - ("embed", None), - ("mlp", "etp"), - ("batch", EP_AXIS), - ) -else: - EP_AXIS = "ep" - FSDP_AXIS = "fsdp" - EP_SIZE = 2 - assert ( - jax.device_count() % EP_SIZE == 0 - ), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" - FSDP_SIZE = jax.device_count() // EP_SIZE - MESH_SHAPE = (FSDP_SIZE, EP_SIZE) - MESH_AXIS_NAMES = (FSDP_AXIS, EP_AXIS) - DATA_PARALLELISM_AXES = (FSDP_AXIS,) - BATCH_MESH_AXIS = (FSDP_AXIS, EP_AXIS) - MESH_RESOURCE = MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - LOGICAL_AXIS_RULES = ( - ("exp", EP_AXIS), - ("embed", FSDP_AXIS), - ("mlp", None), - ("batch", BATCH_MESH_AXIS), - ) - -NUM_DEVICES_REQUIRED = int(np.prod(MESH_SHAPE)) +EP_AXIS = "ep" +FSDP_AXIS = "fsdp" +EP_SIZE = 2 +assert ( + jax.device_count() % EP_SIZE == 0 +), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" +FSDP_SIZE = jax.device_count() // EP_SIZE +NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE + +LOGICAL_AXIS_RULES = ( + ("exp", EP_AXIS), + ("embed", FSDP_AXIS), + ("mlp", None), + ("batch", (FSDP_AXIS, EP_AXIS)), +) # Small shapes so the parity tests stay tight on bf16. The block still # has all four ranks participating in dispatch/combine. DTYPE = jnp.bfloat16 -BATCH = NUM_DEVICES_REQUIRED * 2 +BATCH = EP_SIZE * FSDP_SIZE * 2 # 8 on 4-GPU, 16 on 8-GPU SEQ = 32 HIDDEN = 128 INTER = 128 @@ -216,20 +189,16 @@ def _read_mp_options(): @pytest.fixture(scope="module") def mesh(): - if jax.device_count() < NUM_DEVICES_REQUIRED or ( - COMPOUND_EP and jax.device_count() != NUM_DEVICES_REQUIRED - ): + if jax.device_count() < NUM_DEVICES_REQUIRED: pytest.skip( - f"Need {'exactly' if COMPOUND_EP else '>='} {NUM_DEVICES_REQUIRED} devices for" - f" mesh={MESH_SHAPE}; have {jax.device_count()}" + f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" + f" have {jax.device_count()}" ) - # Compound EP axes retain their declared order. With shape (2, 2, 1), - # (expert, tensor) flattens to contiguous ranks [0, 1, 2, 3]. - if COMPOUND_EP: - devices = np.asarray(jax.devices()).reshape(MESH_SHAPE) - else: - devices = mesh_utils.create_device_mesh(MESH_SHAPE) - mesh_obj = Mesh(devices, axis_names=MESH_AXIS_NAMES) + # ``ep`` must be the inner axis: ``ep_bootstrap`` forms NCCL EP groups + # from consecutive global ranks via ``dp_color = rank // ep_size``, so + # only an (outer_fsdp, inner_ep) device layout groups ranks correctly. + devices = mesh_utils.create_device_mesh((FSDP_SIZE, EP_SIZE)) + mesh_obj = Mesh(devices, axis_names=(FSDP_AXIS, EP_AXIS)) num_procs = jax.process_count() max_tokens_per_rank = (BATCH // num_procs) * SEQ @@ -246,7 +215,7 @@ def mesh(): # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather # and cannot run from inside jax.jit. Sized to the worst-case recv_pr # across _CONFIGS so every parametrized config is bootstrap-compatible. - with mesh_obj, global_shard_guard(MESH_RESOURCE): + with mesh_obj, global_shard_guard(MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS)): ep_bootstrap( world_size=num_procs, rank=jax.process_index(), @@ -255,7 +224,6 @@ def mesh(): recv_capacity_per_rank=recv_capacity_per_rank, hidden_dim=HIDDEN, max_token_dtype=DTYPE, - ep_axis=EP_AXIS if COMPOUND_EP else None, ) record_ep_bootstrap_signature_for_moe( num_experts=NUM_EXPERTS, @@ -396,8 +364,7 @@ def _make_block( num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, intermediate_size=INTER, - ep_axis=EP_AXIS if COMPOUND_EP else None, - data_parallelism_axes=DATA_PARALLELISM_AXES, + data_parallelism_axes=(FSDP_AXIS,), apply_topk_weights_early=apply_topk_weights_early, aux_loss_coeff=aux_loss_coeff, use_expert_routing_bias=use_expert_routing_bias, @@ -427,7 +394,9 @@ def _strong_expert_bias_init(key, shape, dtype): def _shard_inputs(x, mesh): # Match the layout moe.py re-pins to: outer dp axes, then ep innermost. - return jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P(BATCH_MESH_AXIS, None, None))) + return jax.lax.with_sharding_constraint( + x, NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None, None)) + ) def _ctx(mesh): @@ -436,7 +405,9 @@ def _ctx(mesh): class _Combo: def __enter__(self_inner): self_inner._m = mesh.__enter__() - self_inner._gs = global_shard_guard(MESH_RESOURCE) + self_inner._gs = global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ) self_inner._gs.__enter__() self_inner._ar = nn_partitioning.axis_rules(LOGICAL_AXIS_RULES) self_inner._ar.__enter__() @@ -453,18 +424,7 @@ def __exit__(self_inner, *args): def _init_apply(block, mesh, x, key): with _ctx(mesh): x_sh = _shard_inputs(x, mesh) - if COMPOUND_EP: - assert _axis_names(_spec_entry(x_sh.sharding.spec, 0)) == frozenset(EP_AXIS) - assert not any(_axis_names(entry) for entry in x_sh.sharding.spec[1:]) - if COMPOUND_EP: - abstract_variables = jax.eval_shape(block.init, key, x_sh) - logical_specs = flax_linen.get_partition_spec(abstract_variables) - variable_shardings = flax_linen.logical_to_mesh_sharding( - logical_specs, mesh, LOGICAL_AXIS_RULES - ) - variables = jax.jit(block.init, out_shardings=variable_shardings)(key, x_sh) - else: - variables = jax.jit(block.init)(key, x_sh) + variables = jax.jit(block.init)(key, x_sh) jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) output, aux, _trt = jax.jit(block.apply)(variables, x_sh) jax.block_until_ready(output) @@ -540,50 +500,6 @@ def _params_global_numpy(variables, mesh): return {name: _to_global_numpy(_unwrap(p), mesh) for name, p in params.items()} -def _axis_names(entry): - if entry is None: - return frozenset() - return frozenset(entry if isinstance(entry, tuple) else (entry,)) - - -def _spec_entry(spec, index): - return spec[index] if index < len(spec) else None - - -def _slice_size(index, global_size): - start = 0 if index.start is None else index.start - stop = global_size if index.stop is None else index.stop - return stop - start - - -def _assert_compound_ep_sharding(variables, output): - """Assert TP is folded into EP while size-one ETP leaves matrices complete.""" - assert _axis_names(_spec_entry(output.sharding.spec, 0)) == frozenset(EP_AXIS) - assert not any(_axis_names(entry) for entry in output.sharding.spec[1:]) - - params = variables["params"] - gate = _unwrap(params["gate_kernel"]) - wi = _unwrap(params["wi"]) - wo = _unwrap(params["wo"]) - - assert not any( - _axis_names(entry) for entry in gate.sharding.spec - ), "gate weights should be replicated for global routing" - assert _axis_names(_spec_entry(wi.sharding.spec, 0)) == frozenset(EP_AXIS) - assert _axis_names(_spec_entry(wo.sharding.spec, 0)) == frozenset(EP_AXIS) - # JAX may preserve or elide a size-one axis. Either representation must - # leave the expert's hidden/intermediate matrices physically complete. - assert _axis_names(_spec_entry(wi.sharding.spec, 2)) <= {"etp"} - assert _axis_names(_spec_entry(wo.sharding.spec, 1)) <= {"etp"} - - wi_index = wi.addressable_shards[0].index - wo_index = wo.addressable_shards[0].index - assert _slice_size(wi_index[0], NUM_EXPERTS) == NUM_EXPERTS // EP_SIZE - assert _slice_size(wo_index[0], NUM_EXPERTS) == NUM_EXPERTS // EP_SIZE - assert wi_index[1] == slice(None) and wi_index[2] == slice(None) - assert wo_index[1] == slice(None) and wo_index[2] == slice(None) - - def _make_inputs(key): """Generate a globally-identical input tensor on every process.""" return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) @@ -639,16 +555,11 @@ def _quantization_recipe(quantization): ), ] -if COMPOUND_EP: - # This mode is a focused sharding qualification, not a repeat of the - # router feature matrix covered by the default topology. - _CONFIGS = [_CONFIGS[0]] - _QUANTIZATION_CASES = [ pytest.param("bf16", id="bf16"), ] -if get_device_compute_capability(0) >= 100 and not COMPOUND_EP: +if get_device_compute_capability(0) >= 100: _QUANTIZATION_CASES.append(pytest.param("mxfp8", id="mxfp8")) @@ -675,9 +586,6 @@ def test_forward(self, mesh, config, quantization): x = _make_inputs(jax.random.PRNGKey(0)) variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) - if COMPOUND_EP: - _assert_compound_ep_sharding(variables, output) - # Shape / dtype / finiteness (cheap; on the local shard). assert output.shape == x.shape assert output.dtype == x.dtype @@ -784,7 +692,6 @@ def loss_fn(params, x): ) -@pytest.mark.skipif(COMPOUND_EP, reason="compound mode focuses on EP sharding") class TestTeEpMoeAuxLoss: """Aux-loss path. Consolidated into: * ``test_aux_loss``: one run that checks the returned scalar's diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 7837940767..ca70ea145c 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -5,8 +5,8 @@ Sharding model: - EpPrepare / EpDispatch outputs carry a single leading ``num_procs`` dim. - Sharded by the bootstrapped ``ep_axis`` group, optionally preceded by an - outer DP/FSDP axis. + Sharded compound ``(dp_resource, ep_resource)`` when DP is set, else + ``ep_resource`` alone. - EpDispatch inputs are 2D ``[T, H]`` or 3D ``[B, S, H]``; only the first dim may be sharded, with axis in {ep, (dp, ep), dp, None}. Trailing dims must be replicated. ``dp`` alone gets ``ep`` folded in locally. @@ -24,7 +24,7 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive -from ..sharding import MeshAxis, global_mesh_resource, get_mesh_axis_size, normalize_mesh_axes +from ..sharding import global_mesh_resource, get_mesh_axis_size from ..version_utils import is_collective_stream_supported @@ -73,15 +73,12 @@ def wrapper(*args, **kwargs): class EpConfig: """Snapshot of the EP bootstrap config (see ep_bootstrap). - ``ep_axis`` is the ordered physical mesh axis group used by the EP - communicator. ``num_ep_groups`` is the size of the outer dp/fsdp mesh - axis (1 if neither is set), captured at bootstrap so abstract-eval never - reads the mesh. + num_ep_groups is the size of the outer dp/fsdp mesh axis (1 if neither + is set), captured at bootstrap so abstract-eval never reads the mesh. """ world_size: int rank: int - ep_axis: MeshAxis ep_size: int num_ep_groups: int num_experts: int @@ -138,31 +135,25 @@ def ep_handle_mem_size(cfg: EpLayerConfig) -> int: def _leading_axis_ok(spec): """Validate an EP input spec; return ``(ok, ep_axis, outer_axes)``. - Leading dim contains the ordered EP axis group, optionally preceded by an - outer DP/FSDP axis; all other dims must be replicated. + Leading dim is ``ep`` or a tuple ending in ``ep`` (outer dp/fsdp axes + first); all other dims must be replicated. """ gsr = global_mesh_resource() - ep_axis = get_ep_config().ep_axis - ep_axes = normalize_mesh_axes(ep_axis) - outer_axes = tuple( - a for a in (gsr.dp_resource, gsr.fsdp_resource) if a is not None and a not in ep_axes - ) - if len(spec) < 2 or not ep_axes: + ep_axis = gsr.ep_resource + outer_axes = tuple(a for a in (gsr.dp_resource, gsr.fsdp_resource) if a is not None) + if len(spec) < 2 or ep_axis is None: return False, ep_axis, outer_axes if any(ax is not None for ax in spec[1:]): return False, ep_axis, outer_axes leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) - actual = set(a for a in elts if a is not None) - active_ep_axes = tuple(a for a in ep_axes if get_mesh_axis_size(a) > 1) - actual_ep_axes = tuple(a for a in elts if a in active_ep_axes) - if actual_ep_axes != active_ep_axes: + if ep_axis not in elts: return False, ep_axis, outer_axes - allowed = set(outer_axes) | set(ep_axes) - return actual.issubset(allowed), ep_axis, outer_axes + allowed = set(outer_axes) | {ep_axis} + return all(a in allowed for a in elts), ep_axis, outer_axes -def _ep_outer_axis(ep_axis=None): +def _ep_outer_axis(): """The single dp/fsdp axis (if any) sitting outside ep on EP-output tensors. When set, EP-output globals carry an extra leading ``dp_size`` dim so SPMD @@ -172,18 +163,11 @@ def _ep_outer_axis(ep_axis=None): we don't pin EP-output specs to a degenerate axis that JAX may collapse. """ gsr = global_mesh_resource() - if ep_axis is None: - ep_axis = get_ep_config().ep_axis - ep_axes = set(normalize_mesh_axes(ep_axis)) - candidates = tuple( - axis - for axis in (gsr.dp_resource, gsr.fsdp_resource) - if axis is not None and axis not in ep_axes - ) - for axis in candidates: - if get_mesh_axis_size(axis) > 1: - return axis - return candidates[0] if candidates else None + if gsr.dp_resource is not None and get_mesh_axis_size(gsr.dp_resource) > 1: + return gsr.dp_resource + if gsr.fsdp_resource is not None and get_mesh_axis_size(gsr.fsdp_resource) > 1: + return gsr.fsdp_resource + return gsr.dp_resource or gsr.fsdp_resource def _ep_leading_dims(is_outer): @@ -196,22 +180,22 @@ def _ep_leading_dims(is_outer): def _ep_output_spec(*trailing): - """Build an EP-output spec from the bootstrapped EP and outer axis groups.""" - ep_axis = get_ep_config().ep_axis + """PartitionSpec for an EP-output tensor: ``(("dp","ep"), *trailing)`` when + DP is set (compound leading axis on a single dim), else ``("ep",*trailing)``.""" + gsr = global_mesh_resource() outer = _ep_outer_axis() - ep_axes = normalize_mesh_axes(ep_axis) - leading_axes = ep_axes if outer is None else (outer, *ep_axes) - leading = leading_axes[0] if len(leading_axes) == 1 else leading_axes - return PartitionSpec(leading, *trailing) + if outer is None: + return PartitionSpec(gsr.ep_resource, *trailing) + return PartitionSpec((outer, gsr.ep_resource), *trailing) def _ep_spec_ok(spec, trailing_count): - """Validate an EP output's leading axis group and replicated trailing dims. - - JAX may collapse size-one mesh axes to ``None`` or drop them, so only - nontrivial EP axes are required. + """Leading dim shards along ep (and outer dp/fsdp when set); trailing dims + are replicated. JAX may collapse size-1 mesh axes to ``None`` or drop them, + so the leading entry is normalized to a set of named axes before comparing. """ - ep_axes = normalize_mesh_axes(get_ep_config().ep_axis) + gsr = global_mesh_resource() + ep_axis = gsr.ep_resource outer = _ep_outer_axis() if len(spec) != 1 + trailing_count: return False @@ -220,17 +204,8 @@ def _ep_spec_ok(spec, trailing_count): leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) actual = frozenset(a for a in elts if a is not None) - active_ep_axes = tuple(a for a in ep_axes if get_mesh_axis_size(a) > 1) - actual_ep_axes = tuple(a for a in elts if a in active_ep_axes) - if actual_ep_axes != active_ep_axes: - return False - expected = set(ep_axes) - if outer is not None: - expected.add(outer) - required = set(active_ep_axes) - if outer is not None and get_mesh_axis_size(outer) > 1: - required.add(outer) - return required.issubset(actual) and actual.issubset(expected) + expected = {ep_axis} if outer is None else {ep_axis, outer} + return actual <= expected # ── ep_prepare ────────────────────────────────────────────────────────────── diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 9f5e14ee62..2222a41e48 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -21,7 +21,6 @@ get_num_devices_in_mesh, global_mesh_resource, get_mesh_axis_size, - normalize_mesh_axes, with_sharding_constraint, ) @@ -86,21 +85,11 @@ def _ep_domain_for_rank(mesh, ep_resource, rank, device_to_rank=None): def device_to_rank(d): return d.process_index - ep_axes = normalize_mesh_axes(ep_resource) - if not ep_axes: - raise ValueError("ep_bootstrap: ep_resource must contain at least one mesh axis.") - missing = tuple(axis for axis in ep_axes if axis not in mesh.axis_names) - if missing: - raise ValueError( - f"ep_bootstrap: EP axes {missing} are not present in mesh axes {mesh.axis_names}." - ) - ep_positions = tuple(mesh.axis_names.index(axis) for axis in ep_axes) - non_ep_positions = tuple(i for i in range(len(mesh.axis_names)) if i not in ep_positions) - ep_size = get_mesh_axis_size(ep_axes, mesh) + ep_pos = mesh.axis_names.index(ep_resource) + ep_size = mesh.shape[ep_resource] ranks = np.vectorize(device_to_rank, otypes=[np.int64])(mesh.devices) - # Move all EP axes last in the user-specified order and flatten them into - # one communicator dimension. Each row fixes every axis outside compound EP. - grid = np.transpose(ranks, non_ep_positions + ep_positions).reshape(-1, ep_size) + # Move ep last and flatten: each row is one domain (all non-ep coords fixed). + grid = np.moveaxis(ranks, ep_pos, -1).reshape(-1, ep_size) loc = np.argwhere(grid == rank) if loc.shape[0] != 1: raise ValueError( @@ -121,14 +110,13 @@ def ep_bootstrap( max_token_dtype=jnp.bfloat16, max_num_sms=0, drop_on_overflow=False, - ep_axis=None, ): """Initialize the EP communicator. Call once per process before any EP op. - Must run inside the active JAX Mesh and a global_shard_guard. By default, - the EP axis is read from ``MeshResource.ep_resource``. ``ep_axis`` may - override it with an ordered tuple of physical mesh axes for a compound EP - communicator. DP/FSDP axes outside that group determine ``num_ep_groups``. + Must run inside the active JAX Mesh and a global_shard_guard; ep_size and + num_ep_groups are read from the mesh axes named by MeshResource.ep_resource + and MeshResource.dp_resource/fsdp_resource. Axes orthogonal to EP (tp, pp, + cp, ...) are supported and replicated across EP tensors. Args: world_size: Total number of processes (product of all mesh axes). @@ -143,8 +131,6 @@ def ep_bootstrap( drop_on_overflow: Drop tokens exceeding recv_capacity_per_rank instead of trapping on overflow. Dropped tokens are still counted in total_recv_tokens, so callers can detect overflow from it. - ep_axis: Optional physical mesh axis name or ordered tuple of names. - Defaults to ``MeshResource.ep_resource``. """ if jnp.dtype(max_token_dtype) != jnp.bfloat16: raise NotImplementedError( @@ -164,9 +150,12 @@ def ep_bootstrap( ) gsr = global_mesh_resource() - ep_resource = gsr.ep_resource if ep_axis is None else ep_axis + ep_resource = gsr.ep_resource if ep_resource is None: - raise ValueError("ep_bootstrap requires ep_axis or MeshResource.ep_resource to be set.") + raise ValueError( + "ep_bootstrap requires MeshResource.ep_resource to be set; enter a" + " global_shard_guard(MeshResource(..., ep_resource=)) before bootstrap." + ) mesh = _get_mesh() if mesh.empty: raise ValueError( @@ -181,7 +170,7 @@ def ep_bootstrap( ep_size = get_mesh_axis_size(ep_resource) # num_ep_groups counts only the distinct-token (dp/fsdp) axes; replicated # axes (tp, pp, ...) do not create distinct EP-output slabs. - outer_axis = _ep_outer_axis(ep_resource) + outer_axis = _ep_outer_axis() num_ep_groups = 1 if outer_axis is None else get_mesh_axis_size(outer_axis) if num_experts % ep_size != 0: raise ValueError(f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size}).") @@ -226,7 +215,6 @@ def ep_bootstrap( tex.ep.EpConfig( world_size=world_size, rank=rank, - ep_axis=ep_resource, ep_size=ep_size, num_ep_groups=num_ep_groups, num_experts=num_experts, @@ -253,12 +241,14 @@ def ep_finalize(): def _default_out_partition_spec(): - """Build the default leading-axis spec from the bootstrapped EP group.""" - ep_axis = tex.ep.get_ep_config().ep_axis + """Leading-axis default: ``(("dp","ep"),)`` if DP/FSDP is set, else ``("ep",)``.""" + gsr = global_mesh_resource() + if gsr.ep_resource is None: + raise ValueError( + "ep_resource is not set on the active MeshResource; pass out_sharding=... explicitly." + ) outer = _ep_outer_axis() - ep_axes = normalize_mesh_axes(ep_axis) - leading_axes = ep_axes if outer is None else (outer, *ep_axes) - leading = leading_axes[0] if len(leading_axes) == 1 else leading_axes + leading = (outer, gsr.ep_resource) if outer is not None else gsr.ep_resource return (leading,) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 29158e1c6b..cb10c7aa18 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -13,7 +13,7 @@ as ``self.param`` slots (with the right :func:`flax.linen.with_logical_partitioning` annotations so JAX's sharding layer FSDPs the params correctly). -2. Resolve the EP axis from an explicit layer setting or the active +2. Resolve the EP axis name from the active :class:`transformer_engine.jax.sharding.MeshResource`. 3. Forward all knobs to :func:`moe`. @@ -34,13 +34,10 @@ from flax import linen as nn from transformer_engine.common.recipe import Recipe -from ..moe import MeshAxis, moe +from ..moe import moe from ..quantize import QuantizerSet from ..router import ScoreFunction -from ..sharding import ( - _get_mesh, - get_active_resource_axis, -) +from ..sharding import _get_mesh, get_active_resource_axis from .module import TransformerEngineBase PRNGKey = Any @@ -100,11 +97,7 @@ class _MoEBlock(TransformerEngineBase): ADDITION to the EP axis. Empty (default) means activations are replicated across non-EP axes within an EP group; set e.g. ``("fsdp",)`` for true FSDP-of-batch where each device owns a - unique slice of the batch. These axes must be outside ``ep_axis``. - ep_axis : Optional[str | tuple[str, ...]] - Physical mesh axis name, or ordered compound axis group, used for this - MoE block's EP communicator. Defaults to the active - ``MeshResource.ep_resource`` for backward compatibility. + unique slice of the batch. apply_topk_weights_early : bool If ``True``, multiply expert outputs by their top-k weights *inside* each shard before ``ep_combine`` (saves one global @@ -153,7 +146,6 @@ class _MoEBlock(TransformerEngineBase): input_axes: Tuple[Optional[str], ...] = () # Parallelism - ep_axis: Optional[MeshAxis] = None data_parallelism_axes: Tuple[str, ...] = () # MoE knobs forwarded to ``moe()`` @@ -261,9 +253,7 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: jnp.float32, ) - ep_axis = self.ep_axis - if ep_axis is None: - ep_axis = get_active_resource_axis("ep_resource") + ep_axis = get_active_resource_axis("ep_resource") mesh = _get_mesh() data_parallel_size = 1 for axis in self.data_parallelism_axes: diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 204983776e..55a85ebb2f 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -16,9 +16,8 @@ Sharding model -------------- -* Inbound activations are 3D ``[B, S, H]`` sharded on one leading - compound entry containing outer data-parallel axes followed by every axis - in ``ep_axis``. The public +* Inbound activations are 3D ``[B, S, H]`` sharded + ``((*data_parallelism_axes, ep_axis), None, None)``. The public :func:`moe` soft-repins this on entry and warns when a reshard is inserted. * The EP, grouped-quantize, and grouped-GEMM primitives operate at global @@ -51,7 +50,7 @@ ) from .flax.module import _convert_to_activation_function from .router import ScoreFunction, _validate_score_function -from .sharding import MeshAxis, _get_mesh, get_mesh_axis_size, normalize_mesh_axes +from .sharding import _get_mesh __all__ = ["get_moe_recv_capacity_per_rank", "moe"] @@ -64,14 +63,6 @@ _ALIGN_SIZE = 128 -def _moe_leading_axis(ep_axis: MeshAxis, data_parallelism_axes: Tuple[str, ...]): - """Build one PartitionSpec entry with outer axes followed by compound EP.""" - axes = (*data_parallelism_axes, *normalize_mesh_axes(ep_axis)) - if not axes: - raise ValueError("moe(...) requires ep_axis to contain at least one mesh axis.") - return axes[0] if len(axes) == 1 else axes - - def get_moe_recv_capacity_per_rank( *, num_experts: int, @@ -593,14 +584,7 @@ def _moe_fwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") if ep_axis is None: raise ValueError("moe(...) requires ep_axis to be set (TE EP backend).") - ep_axes = normalize_mesh_axes(ep_axis) - bootstrap_ep_axes = normalize_mesh_axes(tex.ep.get_ep_config().ep_axis) - if ep_axes != bootstrap_ep_axes: - raise ValueError( - f"moe(...) ep_axis={ep_axes} does not match the bootstrapped EP axes " - f"{bootstrap_ep_axes}." - ) - num_ep = get_mesh_axis_size(ep_axis, mesh) + num_ep = mesh.shape[ep_axis] if num_experts % num_ep != 0: raise ValueError(f"num_experts={num_experts} must be divisible by EP size={num_ep}") num_local_experts = num_experts // num_ep @@ -646,9 +630,13 @@ def _moe_fwd_rule( ep_size=num_ep, ) - # EP axes must be innermost and retain their declared order: communicator - # ranks are formed by flattening the compound resource in that order. - batch_pspec_axis: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) + if not data_parallelism_axes: + batch_pspec_axis: Any = ep_axis + else: + # ep must be innermost: ep_bootstrap forms NCCL EP comms from + # consecutive global ranks (dp_color = rank // ep_size), so the + # comm only stays within one model replica under (outer_dp, ep). + batch_pspec_axis = (*data_parallelism_axes, ep_axis) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, ep3_spec)) @@ -925,7 +913,10 @@ def _moe_bwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") B, S, _ = x_shape K = num_experts_per_tok - batch_pspec_axis: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) + if not data_parallelism_axes: + batch_pspec_axis: Any = ep_axis + else: + batch_pspec_axis = (*data_parallelism_axes, ep_axis) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) out_partition_spec = (batch_pspec_axis, None, None) @@ -1238,7 +1229,7 @@ def moe( noop_quantizer_set, noop_quantizer_set, ), - ep_axis: MeshAxis, + ep_axis: str, data_parallelism_axes: Tuple[str, ...] = (), input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), @@ -1287,10 +1278,9 @@ def moe( Axis-name parameters: - * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh axis names*. - ``ep_axis`` may be an ordered tuple whose sizes are multiplied into one - compound EP group. These concrete axes are used to compute - ``num_ep`` / ``dp_size`` and construct + * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh + axis names* -- they index ``jax.sharding.Mesh.shape`` directly + (to compute ``num_ep`` / ``dp_size`` and to construct ``P((dp..., ep), None, None)`` for the physical ``jax.lax.with_sharding_constraint`` calls that JAX requires to refer to real mesh axes). @@ -1323,7 +1313,7 @@ def moe( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - expected_leading: Any = _moe_leading_axis(ep_axis, data_parallelism_axes) + expected_leading: Any = (*data_parallelism_axes, ep_axis) if data_parallelism_axes else ep_axis expected_spec = P(expected_leading, None, None) actual_spec = getattr(getattr(x, "sharding", None), "spec", None) if actual_spec is not None and tuple(actual_spec) != tuple(expected_spec): diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 10e439048e..2e8e611fa3 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -9,10 +9,9 @@ parallelism (FSDP). It includes functions for sharding constraints, mesh management, and collective operations. """ - from contextlib import contextmanager from dataclasses import dataclass -from typing import Callable, Optional, Union +from typing import Callable, Optional import warnings import jax @@ -37,20 +36,6 @@ W_TP_AXES = "nvte_w_tp" W_JOINED_AXES = "nvte_w_joined" -MeshAxis = Union[str, tuple[str, ...]] - - -def normalize_mesh_axes(axis: Optional[MeshAxis]) -> tuple[str, ...]: - """Return a mesh axis or axis group as an ordered tuple of physical names.""" - if axis is None: - return () - axes = axis if isinstance(axis, tuple) else (axis,) - if not axes or any(not isinstance(name, str) or not name for name in axes): - raise ValueError(f"Mesh axes must be non-empty strings, got {axis!r}.") - if len(set(axes)) != len(axes): - raise ValueError(f"Mesh axes must not contain duplicates, got {axis!r}.") - return axes - def _get_mesh(): # Handle Mesh's set via `with mesh:` @@ -289,14 +274,11 @@ def get_mesh_axis_size(axis, mesh=None): if mesh is None: mesh = _get_mesh() - axes = normalize_mesh_axes(axis) - if not axes: + if axis is None: return 1 - size = 1 - for name in axes: - assert name in mesh.shape, f"{name} is not an axis of the given mesh {mesh.shape}" - size *= mesh.shape[name] - return size + + assert axis in mesh.shape, f"{axis} is not a axis of the given mesh {mesh.shape}" + return mesh.shape[axis] def get_mesh_axis_rank(axis: str, mesh=None): From 7514041d9fe5cf495d4ddcd5a54def6c6ee113fd Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 2 Sep 2026 14:50:23 -0700 Subject: [PATCH 5/5] [JAX] Support compound EP axes in MoEBlock Signed-off-by: Jeremy Berchtold --- tests/jax/test_multi_process_ep.py | 31 ++- transformer_engine/jax/cpp_extensions/ep.py | 251 ++++++++++++++------ transformer_engine/jax/ep.py | 95 +++++--- transformer_engine/jax/flax/moe.py | 10 +- transformer_engine/jax/moe.py | 32 ++- 5 files changed, 295 insertions(+), 124 deletions(-) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 47af0b0c39..122f03039a 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -47,10 +47,11 @@ ep_dispatch_fwd, ep_combine_fwd, get_ep_config, + _leading_axis_ok, + _ep_spec_ok, ) from transformer_engine.jax.version_utils import is_collective_stream_supported - # ── Test config ───────────────────────────────────────────────────────────── # NCCL EP requires NUM_LOCAL_EXPERTS*ep % 4 == 0 (TMA alignment in # device/hybridep_adapter.cu:511). With NUM_LOCAL_EXPERTS=2, ep must be even. @@ -960,6 +961,34 @@ def test_ep_tp_splits_domains(self): self.assertEqual(domains, {0: [0, 2, 4, 6], 1: [1, 3, 5, 7]}) + def test_compound_ep_axis(self): + if not is_devices_enough(4): + self.skipTest("requires 4 devices") + mesh = Mesh(np.asarray(jax.devices()[:4]).reshape(2, 2), ("expert", "tensor")) + order = {int(device.id): i for i, device in enumerate(mesh.devices.reshape(-1))} + + domains = {} + for rank in range(4): + root, col, num_domains = _ep_domain_for_rank( + mesh, + ("expert", "tensor"), + rank, + device_to_rank=lambda device: order[int(device.id)], + ) + self.assertEqual(num_domains, 1) + domains.setdefault(root, {})[col] = rank + self.assertEqual([domains[0][i] for i in range(4)], [0, 1, 2, 3]) + + with mesh: + ep_axes = ("expert", "tensor") + ok, ep_axis, outer_axes = _leading_axis_ok( + PartitionSpec(ep_axes, None, None), ep_axes=ep_axes + ) + self.assertTrue(ok) + self.assertEqual(ep_axis, ("expert", "tensor")) + self.assertEqual(outer_axes, ()) + self.assertTrue(_ep_spec_ok(PartitionSpec(ep_axes, None, None), 2, ep_axes)) + # ── Entry point ────────────────────────────────────────────────────────────── diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index ca70ea145c..ddb13bb331 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -16,6 +16,7 @@ import functools from dataclasses import dataclass +from typing import Tuple, Union import jax import jax.numpy as jnp @@ -91,6 +92,32 @@ class EpConfig: _ep_config: EpConfig = None +def _normalize_ep_axes(ep_axis: Union[str, Tuple[str, ...]]) -> Tuple[str, ...]: + """Normalize a physical EP mesh axis or axis tuple.""" + axes = (ep_axis,) if isinstance(ep_axis, str) else ep_axis + if not axes or any(not isinstance(axis, str) or not axis for axis in axes): + raise ValueError( + f"ep_axis must be a mesh axis name or non-empty tuple of names, got {ep_axis!r}" + ) + if len(set(axes)) != len(axes): + raise ValueError(f"ep_axis must not contain duplicate mesh axes, got {ep_axis!r}") + return axes + + +def _ep_axis_entry(ep_axes: Tuple[str, ...]): + """Return the PartitionSpec entry for normalized EP axes.""" + return ep_axes[0] if len(ep_axes) == 1 else ep_axes + + +def _resolve_ep_axes(ep_axes=None) -> Tuple[str, ...]: + """Resolve explicit EP axes, falling back to MeshResource for compatibility.""" + if ep_axes is None: + ep_axes = global_mesh_resource().ep_resource + if ep_axes is None: + raise ValueError("EP mesh axes must be passed explicitly or set as ep_resource.") + return _normalize_ep_axes(ep_axes) + + def set_ep_config(config: EpConfig) -> None: """Cache the EP config for abstract-eval / sharding helpers. Call once.""" global _ep_config @@ -132,28 +159,38 @@ def ep_handle_mem_size(cfg: EpLayerConfig) -> int: ) -def _leading_axis_ok(spec): +def _leading_axis_ok(spec, ep_axes=None): """Validate an EP input spec; return ``(ok, ep_axis, outer_axes)``. Leading dim is ``ep`` or a tuple ending in ``ep`` (outer dp/fsdp axes first); all other dims must be replicated. """ - gsr = global_mesh_resource() - ep_axis = gsr.ep_resource - outer_axes = tuple(a for a in (gsr.dp_resource, gsr.fsdp_resource) if a is not None) - if len(spec) < 2 or ep_axis is None: - return False, ep_axis, outer_axes + explicit_ep_axes = ep_axes is not None + ep_axes = _resolve_ep_axes(ep_axes) + if len(spec) < 2: + return False, _ep_axis_entry(ep_axes), () if any(ax is not None for ax in spec[1:]): - return False, ep_axis, outer_axes + return False, _ep_axis_entry(ep_axes), () leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) - if ep_axis not in elts: - return False, ep_axis, outer_axes - allowed = set(outer_axes) | {ep_axis} - return all(a in allowed for a in elts), ep_axis, outer_axes + if explicit_ep_axes: + outer_axes = elts[: -len(ep_axes)] + if tuple(elts[-len(ep_axes) :]) != ep_axes: + return False, _ep_axis_entry(ep_axes), outer_axes + else: + gsr = global_mesh_resource() + outer_axes = tuple( + axis + for axis in (gsr.dp_resource, gsr.fsdp_resource) + if axis is not None and axis not in ep_axes + ) + if not explicit_ep_axes and tuple(axis for axis in elts if axis in ep_axes) != ep_axes: + return False, _ep_axis_entry(ep_axes), outer_axes + allowed = set(outer_axes) | set(ep_axes) + return all(axis in allowed for axis in elts), _ep_axis_entry(ep_axes), outer_axes -def _ep_outer_axis(): +def _ep_outer_axis(ep_axes=None): """The single dp/fsdp axis (if any) sitting outside ep on EP-output tensors. When set, EP-output globals carry an extra leading ``dp_size`` dim so SPMD @@ -162,12 +199,19 @@ def _ep_outer_axis(): A dp/fsdp axis that is sized 1 in the active mesh is treated as absent so we don't pin EP-output specs to a degenerate axis that JAX may collapse. """ + if ep_axes is not None: + return None gsr = global_mesh_resource() - if gsr.dp_resource is not None and get_mesh_axis_size(gsr.dp_resource) > 1: - return gsr.dp_resource - if gsr.fsdp_resource is not None and get_mesh_axis_size(gsr.fsdp_resource) > 1: - return gsr.fsdp_resource - return gsr.dp_resource or gsr.fsdp_resource + ep_axes = _resolve_ep_axes(ep_axes) + candidates = tuple( + axis + for axis in (gsr.dp_resource, gsr.fsdp_resource) + if axis is not None and axis not in ep_axes + ) + for axis in candidates: + if get_mesh_axis_size(axis) > 1: + return axis + return candidates[0] if candidates else None def _ep_leading_dims(is_outer): @@ -179,24 +223,14 @@ def _ep_leading_dims(is_outer): return (cfg.num_ep_groups * cfg.ep_size,) -def _ep_output_spec(*trailing): - """PartitionSpec for an EP-output tensor: ``(("dp","ep"), *trailing)`` when - DP is set (compound leading axis on a single dim), else ``("ep",*trailing)``.""" - gsr = global_mesh_resource() - outer = _ep_outer_axis() - if outer is None: - return PartitionSpec(gsr.ep_resource, *trailing) - return PartitionSpec((outer, gsr.ep_resource), *trailing) - - -def _ep_spec_ok(spec, trailing_count): +def _ep_spec_ok(spec, trailing_count, ep_axes=None): """Leading dim shards along ep (and outer dp/fsdp when set); trailing dims are replicated. JAX may collapse size-1 mesh axes to ``None`` or drop them, so the leading entry is normalized to a set of named axes before comparing. """ - gsr = global_mesh_resource() - ep_axis = gsr.ep_resource - outer = _ep_outer_axis() + explicit_ep_axes = ep_axes is not None + outer = _ep_outer_axis(ep_axes) + ep_axes = _resolve_ep_axes(ep_axes) if len(spec) != 1 + trailing_count: return False if any(ax is not None for ax in spec[1:]): @@ -204,7 +238,9 @@ def _ep_spec_ok(spec, trailing_count): leading = spec[0] elts = leading if isinstance(leading, tuple) else (leading,) actual = frozenset(a for a in elts if a is not None) - expected = {ep_axis} if outer is None else {ep_axis, outer} + if explicit_ep_axes: + return tuple(elts[-len(ep_axes) :]) == ep_axes + expected = set(ep_axes) if outer is None else {*ep_axes, outer} return actual <= expected @@ -216,14 +252,15 @@ class EpPreparePrimitive(BasePrimitive): name = "te_ep_prepare_ffi" multiple_results = True - impl_static_args = (1, 2, 3) # top_k, dispatch_output_per_expert_alignment, is_outer + impl_static_args = (1, 2, 3, 4) # top_k, alignment, is_outer, ep_axes inner_primitive = None outer_primitive = None @staticmethod - def abstract(topk_idx_aval, *, top_k, dispatch_output_per_expert_alignment, is_outer): + def abstract(topk_idx_aval, *, top_k, dispatch_output_per_expert_alignment, is_outer, ep_axes): # is_outer=True: global leading dim = (dp*ep,) (or (ep,) with no DP); # False: per-shard = (1,). + del ep_axes cfg = get_ep_config() num_local_experts = cfg.num_local_experts assert ( @@ -247,8 +284,8 @@ def outer_abstract(*args, **kwargs): return EpPreparePrimitive.abstract(*args, **kwargs) # pylint: disable=missing-kwoa @staticmethod - def lowering(ctx, topk_idx, *, top_k, dispatch_output_per_expert_alignment, is_outer): - del is_outer + def lowering(ctx, topk_idx, *, top_k, dispatch_output_per_expert_alignment, is_outer, ep_axes): + del is_outer, ep_axes return ffi.ffi_lowering(EpPreparePrimitive.name)( ctx, topk_idx, @@ -257,27 +294,42 @@ def lowering(ctx, topk_idx, *, top_k, dispatch_output_per_expert_alignment, is_o ) @staticmethod - def impl(topk_idx, top_k, dispatch_output_per_expert_alignment, is_outer): + def impl(topk_idx, top_k, dispatch_output_per_expert_alignment, is_outer, ep_axes): assert EpPreparePrimitive.inner_primitive is not None token_counts, total_recv_tokens, handle_mem = EpPreparePrimitive.inner_primitive.bind( topk_idx, top_k=top_k, dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, is_outer=is_outer, + ep_axes=ep_axes, ) return token_counts, total_recv_tokens, handle_mem @staticmethod - def batcher(batched_args, batch_dims, *, top_k, dispatch_output_per_expert_alignment, is_outer): + def batcher( + batched_args, + batch_dims, + *, + top_k, + dispatch_output_per_expert_alignment, + is_outer, + ep_axes, + ): raise NotImplementedError("EpPreparePrimitive does not support vmap") @staticmethod def partition( - top_k, dispatch_output_per_expert_alignment, is_outer, mesh, arg_infos, result_infos + top_k, + dispatch_output_per_expert_alignment, + is_outer, + ep_axes, + mesh, + arg_infos, + result_infos, ): del is_outer, result_infos idx_spec = arg_infos[0].sharding.spec - ok, ep_axis, outer_axes = _leading_axis_ok(idx_spec) + ok, ep_axis, outer_axes = _leading_axis_ok(idx_spec, ep_axes) if not ok: raise NotImplementedError( "EpPrepare: topk_idx leading dim must include ep_resource" @@ -294,7 +346,7 @@ def partition( def sharded_impl(topk_idx): return EpPreparePrimitive.impl( - topk_idx, top_k, dispatch_output_per_expert_alignment, False + topk_idx, top_k, dispatch_output_per_expert_alignment, False, ep_axes ) return mesh, sharded_impl, (tc_sharding, trt_sharding, hm_sharding), arg_shardings @@ -302,7 +354,7 @@ def sharded_impl(topk_idx): @staticmethod def shardy_sharding_rule(*args): # Signature: (*static_args, mesh, value_types, result_types). Static args - # for this primitive are (top_k, dispatch_alignment, is_outer). + # for this primitive are (top_k, dispatch_alignment, is_outer, ep_axes). value_types = args[-2] topk_idx_rank = len(value_types[0].shape) in_axes = " ".join(f"L{i}" for i in range(topk_idx_rank - 1)) + " topk" @@ -320,8 +372,7 @@ class EpDispatchPrimitive(BasePrimitive): name = "te_ep_dispatch_ffi" multiple_results = True - impl_static_args = (4, 5, 6, 7) # top_k, dispatch_output_per_expert_alignment, - # recv_capacity_per_rank, is_outer + impl_static_args = (4, 5, 6, 7, 8) # top_k, alignment, recv capacity, is_outer, ep_axes inner_primitive = None outer_primitive = None @@ -336,11 +387,12 @@ def abstract( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): # is_outer=True: global leading dim = (dp*ep,) (or (ep,) with no DP); # False: per-shard = (1,). del topk_idx_aval, topk_weights_aval, top_k, dispatch_output_per_expert_alignment - del handle_mem_aval + del handle_mem_aval, ep_axes assert ( len(tokens_aval.shape) >= 2 ), f"tokens must be at least 2D [..., H], got shape {tokens_aval.shape}" @@ -369,8 +421,9 @@ def lowering( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): - del recv_capacity_per_rank, is_outer + del recv_capacity_per_rank, is_outer, ep_axes return ffi.ffi_lowering(EpDispatchPrimitive.name)( ctx, handle_mem, @@ -391,6 +444,7 @@ def impl( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): assert EpDispatchPrimitive.inner_primitive is not None recv_tokens, recv_topk_weights = EpDispatchPrimitive.inner_primitive.bind( @@ -402,6 +456,7 @@ def impl( dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, recv_capacity_per_rank=recv_capacity_per_rank, is_outer=is_outer, + ep_axes=ep_axes, ) return recv_tokens, recv_topk_weights @@ -414,6 +469,7 @@ def batcher( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): raise NotImplementedError("EpDispatchPrimitive does not support vmap") @@ -423,13 +479,14 @@ def partition( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, mesh, arg_infos, result_infos, ): del is_outer, result_infos tokens_spec = arg_infos[2].sharding.spec - ok, ep_axis, outer_axes = _leading_axis_ok(tokens_spec) + ok, ep_axis, outer_axes = _leading_axis_ok(tokens_spec, ep_axes) if not ok: raise NotImplementedError( "EpDispatch: tokens leading dim must include ep_resource" @@ -461,6 +518,7 @@ def sharded_impl(handle_mem, topk_idx, tokens, topk_weights): dispatch_output_per_expert_alignment, recv_capacity_per_rank, False, + ep_axes, ) return mesh, sharded_impl, out_shardings, arg_shardings @@ -468,7 +526,7 @@ def sharded_impl(handle_mem, topk_idx, tokens, topk_weights): @staticmethod def shardy_sharding_rule(*args): # Signature: (*static_args, mesh, value_types, result_types). Static args - # for this primitive are (top_k, dispatch_alignment, recv_capacity_per_rank, is_outer). + # Static args are (top_k, dispatch_alignment, recv capacity, is_outer, ep_axes). value_types = args[-2] # Inputs: handle_mem, topk_idx, tokens, topk_weights. idx_rank = len(value_types[1].shape) @@ -516,8 +574,7 @@ class EpCombinePrimitive(BasePrimitive): name = "te_ep_combine_ffi" multiple_results = False - impl_static_args = (2, 3, 4, 5) # top_k, dispatch_output_per_expert_alignment, - # out_leading_shape, out_partition_spec + impl_static_args = (2, 3, 4, 5, 6) # top_k, alignment, output shape/spec, ep_axes inner_primitive = None outer_primitive = None @@ -530,8 +587,15 @@ def abstract( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): - del top_k, dispatch_output_per_expert_alignment, out_partition_spec, handle_mem_aval + del ( + top_k, + dispatch_output_per_expert_alignment, + out_partition_spec, + handle_mem_aval, + ep_axes, + ) assert ( len(expert_out_aval.shape) == 3 ), f"expert_out must be 3D [num_procs, recv_pr, H], got shape {expert_out_aval.shape}" @@ -550,8 +614,9 @@ def lowering( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): - del out_leading_shape, out_partition_spec + del out_leading_shape, out_partition_spec, ep_axes return ffi.ffi_lowering(EpCombinePrimitive.name)( ctx, handle_mem, @@ -568,6 +633,7 @@ def impl( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): assert EpCombinePrimitive.inner_primitive is not None return EpCombinePrimitive.inner_primitive.bind( @@ -577,6 +643,7 @@ def impl( dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, out_leading_shape=out_leading_shape, out_partition_spec=out_partition_spec, + ep_axes=ep_axes, ) @staticmethod @@ -588,6 +655,7 @@ def batcher( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): raise NotImplementedError("EpCombinePrimitive does not support vmap") @@ -597,13 +665,14 @@ def partition( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, mesh, arg_infos, result_infos, ): del result_infos eo_spec = arg_infos[1].sharding.spec - if not _ep_spec_ok(eo_spec, trailing_count=2): + if not _ep_spec_ok(eo_spec, trailing_count=2, ep_axes=ep_axes): raise NotImplementedError( "EpCombine: expert_out must be sharded as PartitionSpec(ep_resource," " None, None) (or ((dp, ep), None, None) when dp/fsdp is set)" @@ -625,6 +694,7 @@ def sharded_impl(handle_mem, expert_out): dispatch_output_per_expert_alignment, per_shard_leading, out_partition_spec, + ep_axes, ) return mesh, sharded_impl, out_sharding, arg_shardings @@ -632,7 +702,7 @@ def sharded_impl(handle_mem, expert_out): @staticmethod def shardy_sharding_rule(*args): # Signature: (*static_args, mesh, value_types, result_types). Static args: - # (top_k, dispatch_alignment, out_leading_shape, out_partition_spec). + # (top_k, dispatch_alignment, out_leading_shape, out_partition_spec, ep_axes). result_types = args[-1] out_rank = len(result_types[0].shape) out_axes = " ".join(f"O{i}" for i in range(out_rank - 1)) + " H" @@ -650,8 +720,7 @@ class EpDispatchBwdPrimitive(BasePrimitive): name = "te_ep_dispatch_bwd_ffi" multiple_results = True - impl_static_args = (3, 4, 5, 6) # top_k, dispatch_output_per_expert_alignment, - # out_leading_shape, out_partition_spec + impl_static_args = (3, 4, 5, 6, 7) # top_k, alignment, output shape/spec, ep_axes inner_primitive = None outer_primitive = None @@ -665,9 +734,10 @@ def abstract( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): del dispatch_output_per_expert_alignment - del g_recv_topk_weights_aval, out_partition_spec, handle_mem_aval + del g_recv_topk_weights_aval, out_partition_spec, handle_mem_aval, ep_axes assert ( len(grad_aval.shape) == 3 ), f"grad must be 3D [num_procs, recv_pr, H], got shape {grad_aval.shape}" @@ -690,8 +760,9 @@ def lowering( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): - del out_leading_shape, out_partition_spec + del out_leading_shape, out_partition_spec, ep_axes return ffi.ffi_lowering(EpDispatchBwdPrimitive.name)( ctx, handle_mem, @@ -710,6 +781,7 @@ def impl( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): assert EpDispatchBwdPrimitive.inner_primitive is not None return EpDispatchBwdPrimitive.inner_primitive.bind( @@ -720,6 +792,7 @@ def impl( dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, out_leading_shape=out_leading_shape, out_partition_spec=out_partition_spec, + ep_axes=ep_axes, ) @staticmethod @@ -731,6 +804,7 @@ def batcher( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, ): raise NotImplementedError("EpDispatchBwdPrimitive does not support vmap") @@ -740,20 +814,21 @@ def partition( dispatch_output_per_expert_alignment, out_leading_shape, out_partition_spec, + ep_axes, mesh, arg_infos, result_infos, ): del result_infos g_spec = arg_infos[1].sharding.spec - if not _ep_spec_ok(g_spec, trailing_count=2): + if not _ep_spec_ok(g_spec, trailing_count=2, ep_axes=ep_axes): raise NotImplementedError( "EpDispatchBwd: grad must be sharded as PartitionSpec(ep_resource," " None, None) (or ((dp, ep), None, None) when dp/fsdp is set)" f" over [num_procs, recv_pr, H]; got spec={g_spec}." ) gw_spec = arg_infos[2].sharding.spec - if not _ep_spec_ok(gw_spec, trailing_count=1): + if not _ep_spec_ok(gw_spec, trailing_count=1, ep_axes=ep_axes): raise NotImplementedError( "EpDispatchBwd: g_recv_topk_weights must be sharded as" " PartitionSpec(ep_resource, None) (or ((dp, ep), None) when dp/fsdp is set)" @@ -782,6 +857,7 @@ def sharded_impl(handle_mem, grad, g_recv_topk_weights): dispatch_output_per_expert_alignment, per_shard_leading, out_partition_spec, + ep_axes, ) return mesh, sharded_impl, out_shardings, arg_shardings @@ -807,8 +883,7 @@ class EpCombineBwdPrimitive(BasePrimitive): name = "te_ep_combine_bwd_ffi" multiple_results = False - impl_static_args = (2, 3, 4, 5) # top_k, dispatch_output_per_expert_alignment, - # recv_capacity_per_rank, is_outer + impl_static_args = (2, 3, 4, 5, 6) # top_k, alignment, recv capacity, is_outer, ep_axes inner_primitive = None outer_primitive = None @@ -821,10 +896,11 @@ def abstract( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): # is_outer=True: global leading dim = (dp*ep,) (or (ep,) with no DP); # False: per-shard = (1,). - del top_k, dispatch_output_per_expert_alignment, handle_mem_aval + del top_k, dispatch_output_per_expert_alignment, handle_mem_aval, ep_axes assert ( len(grad_aval.shape) >= 2 ), f"grad must be at least 2D [..., H], got shape {grad_aval.shape}" @@ -848,8 +924,9 @@ def lowering( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): - del recv_capacity_per_rank, is_outer + del recv_capacity_per_rank, is_outer, ep_axes return ffi.ffi_lowering(EpCombineBwdPrimitive.name)( ctx, handle_mem, @@ -866,6 +943,7 @@ def impl( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): assert EpCombineBwdPrimitive.inner_primitive is not None return EpCombineBwdPrimitive.inner_primitive.bind( @@ -875,6 +953,7 @@ def impl( dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, recv_capacity_per_rank=recv_capacity_per_rank, is_outer=is_outer, + ep_axes=ep_axes, ) @staticmethod @@ -886,6 +965,7 @@ def batcher( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, ): raise NotImplementedError("EpCombineBwdPrimitive does not support vmap") @@ -895,14 +975,22 @@ def partition( dispatch_output_per_expert_alignment, recv_capacity_per_rank, is_outer, + ep_axes, mesh, arg_infos, result_infos, ): del is_outer, result_infos arg_shardings = tuple(a.sharding for a in arg_infos) - # EP-output leading (trailing dims auto-pad to None). - out_sharding = NamedSharding(mesh, _ep_output_spec()) + grad_spec = arg_infos[1].sharding.spec + ok, ep_axis, outer_axes = _leading_axis_ok(grad_spec, ep_axes) + if not ok: + raise NotImplementedError( + "EpCombineBwd: grad leading dim must include the EP axes" + f" ({ep_axis}), optionally preceded by {outer_axes}; got spec={grad_spec}." + ) + # The expert-output cotangent inherits the input's leading sharding. + out_sharding = NamedSharding(mesh, PartitionSpec(grad_spec[0])) def sharded_impl(handle_mem, grad): return EpCombineBwdPrimitive.impl( @@ -912,6 +1000,7 @@ def sharded_impl(handle_mem, grad): dispatch_output_per_expert_alignment, recv_capacity_per_rank, False, + ep_axes, ) return mesh, sharded_impl, out_sharding, arg_shardings @@ -932,23 +1021,32 @@ def shardy_sharding_rule(*args): @_on_collective_stream -def ep_prepare(cfg: EpLayerConfig, topk_idx): +def ep_prepare(cfg: EpLayerConfig, topk_idx, ep_axes=None): """Exchange routing metadata for ``cfg``; return ``(token_counts, total_recv_tokens, handle_mem)``. ``total_recv_tokens`` is the per-rank pre-drop recv-slot total (includes tokens dropped on overflow).""" + ep_axes = _resolve_ep_axes(ep_axes) return EpPreparePrimitive.outer_primitive.bind( topk_idx, top_k=int(cfg.top_k), dispatch_output_per_expert_alignment=int(cfg.dispatch_output_per_expert_alignment), is_outer=True, + ep_axes=ep_axes, ) @_on_collective_stream def ep_dispatch_fwd( - cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank + cfg: EpLayerConfig, + handle_mem, + topk_idx, + tokens, + topk_weights, + recv_capacity_per_rank, + ep_axes=None, ): """Scatter tokens and weights to expert ranks; returns (recv_tokens, recv_topk_weights).""" + ep_axes = _resolve_ep_axes(ep_axes) return EpDispatchPrimitive.outer_primitive.bind( handle_mem, topk_idx, @@ -958,15 +1056,22 @@ def ep_dispatch_fwd( dispatch_output_per_expert_alignment=int(cfg.dispatch_output_per_expert_alignment), recv_capacity_per_rank=recv_capacity_per_rank, is_outer=True, + ep_axes=ep_axes, ) @_on_collective_stream def ep_combine_fwd( - cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None + cfg: EpLayerConfig, + handle_mem, + expert_out, + num_local_tokens, + out_partition_spec=None, + ep_axes=None, ): """Gather expert outputs back to home ranks. expert_out is pre-weighted.""" out_leading = _normalize_leading_shape(num_local_tokens) + ep_axes = _resolve_ep_axes(ep_axes) return EpCombinePrimitive.outer_primitive.bind( handle_mem, expert_out, @@ -974,6 +1079,7 @@ def ep_combine_fwd( dispatch_output_per_expert_alignment=int(cfg.dispatch_output_per_expert_alignment), out_leading_shape=out_leading, out_partition_spec=out_partition_spec, + ep_axes=ep_axes, ) @@ -985,9 +1091,11 @@ def ep_dispatch_bwd( g_recv_topk_weights, num_local_tokens, out_partition_spec=None, + ep_axes=None, ): """Backward of dispatch; returns (grad_tokens, grad_topk_weights).""" out_leading = _normalize_leading_shape(num_local_tokens) + ep_axes = _resolve_ep_axes(ep_axes) return EpDispatchBwdPrimitive.outer_primitive.bind( handle_mem, grad, @@ -996,12 +1104,14 @@ def ep_dispatch_bwd( dispatch_output_per_expert_alignment=int(cfg.dispatch_output_per_expert_alignment), out_leading_shape=out_leading, out_partition_spec=out_partition_spec, + ep_axes=ep_axes, ) @_on_collective_stream -def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank): +def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank, ep_axes=None): """Backward of combine; returns grad_expert_out [num_procs, recv_capacity_per_rank, H].""" + ep_axes = _resolve_ep_axes(ep_axes) return EpCombineBwdPrimitive.outer_primitive.bind( handle_mem, grad, @@ -1009,4 +1119,5 @@ def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank) dispatch_output_per_expert_alignment=int(cfg.dispatch_output_per_expert_alignment), recv_capacity_per_rank=recv_capacity_per_rank, is_outer=True, + ep_axes=ep_axes, ) diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 2222a41e48..d5ff629b56 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -14,7 +14,7 @@ import transformer_engine_jax import transformer_engine.jax.cpp_extensions as tex -from transformer_engine.jax.cpp_extensions.ep import _ep_outer_axis +from transformer_engine.jax.cpp_extensions.ep import _ep_outer_axis, _normalize_ep_axes from transformer_engine.jax.cpp_extensions.misc import jax_dtype_to_te_dtype from transformer_engine.jax.sharding import ( _get_mesh, @@ -85,11 +85,13 @@ def _ep_domain_for_rank(mesh, ep_resource, rank, device_to_rank=None): def device_to_rank(d): return d.process_index - ep_pos = mesh.axis_names.index(ep_resource) - ep_size = mesh.shape[ep_resource] + ep_axes = _normalize_ep_axes(ep_resource) + ep_positions = tuple(mesh.axis_names.index(axis) for axis in ep_axes) + non_ep_positions = tuple(i for i in range(len(mesh.axis_names)) if i not in ep_positions) + ep_size = int(np.prod([mesh.shape[axis] for axis in ep_axes])) ranks = np.vectorize(device_to_rank, otypes=[np.int64])(mesh.devices) - # Move ep last and flatten: each row is one domain (all non-ep coords fixed). - grid = np.moveaxis(ranks, ep_pos, -1).reshape(-1, ep_size) + # Move EP axes last in their declared order and flatten them into one domain. + grid = np.transpose(ranks, non_ep_positions + ep_positions).reshape(-1, ep_size) loc = np.argwhere(grid == rank) if loc.shape[0] != 1: raise ValueError( @@ -110,13 +112,13 @@ def ep_bootstrap( max_token_dtype=jnp.bfloat16, max_num_sms=0, drop_on_overflow=False, + ep_axes=None, ): """Initialize the EP communicator. Call once per process before any EP op. - Must run inside the active JAX Mesh and a global_shard_guard; ep_size and - num_ep_groups are read from the mesh axes named by MeshResource.ep_resource - and MeshResource.dp_resource/fsdp_resource. Axes orthogonal to EP (tp, pp, - cp, ...) are supported and replicated across EP tensors. + Must run inside an active JAX Mesh. ``ep_axes`` overrides + ``MeshResource.ep_resource`` when provided; the global sharding context is + only required for the fallback. Args: world_size: Total number of processes (product of all mesh axes). @@ -131,6 +133,8 @@ def ep_bootstrap( drop_on_overflow: Drop tokens exceeding recv_capacity_per_rank instead of trapping on overflow. Dropped tokens are still counted in total_recv_tokens, so callers can detect overflow from it. + ep_axes: Physical mesh axis name or ordered tuple of names. Defaults to + ``MeshResource.ep_resource``. """ if jnp.dtype(max_token_dtype) != jnp.bfloat16: raise NotImplementedError( @@ -149,13 +153,12 @@ def ep_bootstrap( " support single-process multi-device setups." ) - gsr = global_mesh_resource() - ep_resource = gsr.ep_resource + explicit_ep_axes = ep_axes + ep_resource = ( + global_mesh_resource().ep_resource if explicit_ep_axes is None else explicit_ep_axes + ) if ep_resource is None: - raise ValueError( - "ep_bootstrap requires MeshResource.ep_resource to be set; enter a" - " global_shard_guard(MeshResource(..., ep_resource=)) before bootstrap." - ) + raise ValueError("ep_bootstrap requires ep_axes or MeshResource.ep_resource to be set.") mesh = _get_mesh() if mesh.empty: raise ValueError( @@ -167,10 +170,11 @@ def ep_bootstrap( f"ep_bootstrap: mesh device count ({get_num_devices_in_mesh(mesh)}) must equal" f" world_size ({world_size})." ) - ep_size = get_mesh_axis_size(ep_resource) + ep_axes = _normalize_ep_axes(ep_resource) + ep_size = int(np.prod([mesh.shape[axis] for axis in ep_axes])) # num_ep_groups counts only the distinct-token (dp/fsdp) axes; replicated # axes (tp, pp, ...) do not create distinct EP-output slabs. - outer_axis = _ep_outer_axis() + outer_axis = _ep_outer_axis(None if explicit_ep_axes is None else ep_axes) num_ep_groups = 1 if outer_axis is None else get_mesh_axis_size(outer_axis) if num_experts % ep_size != 0: raise ValueError(f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size}).") @@ -240,23 +244,20 @@ def ep_finalize(): tex.ep.reset_ep_config() -def _default_out_partition_spec(): +def _default_out_partition_spec(ep_axes=None): """Leading-axis default: ``(("dp","ep"),)`` if DP/FSDP is set, else ``("ep",)``.""" - gsr = global_mesh_resource() - if gsr.ep_resource is None: - raise ValueError( - "ep_resource is not set on the active MeshResource; pass out_sharding=... explicitly." - ) - outer = _ep_outer_axis() - leading = (outer, gsr.ep_resource) if outer is not None else gsr.ep_resource + outer = _ep_outer_axis(ep_axes) + resolved_ep_axes = tex.ep._resolve_ep_axes(ep_axes) + leading_axes = resolved_ep_axes if outer is None else (outer, *resolved_ep_axes) + leading = leading_axes[0] if len(leading_axes) == 1 else leading_axes return (leading,) # ── ep_dispatch (custom_vjp) ───────────────────────────────────────────────── -@partial(jax.custom_vjp, nondiff_argnums=(0, 4)) -def ep_dispatch(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank): +@partial(jax.custom_vjp, nondiff_argnums=(0, 4, 5)) +def ep_dispatch(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank, ep_axes=None): """Scatter tokens and weights to expert ranks. ``cfg`` is a per-layer ``EpLayerConfig``; distinct layers may share a @@ -271,29 +272,32 @@ def ep_dispatch(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank): ``drop_on_overflow`` is set. When ``recv_capacity_per_rank`` is not sized for the worst case, detect overflow by ``process_allgather``-ing it, then ``max(...) > recv_capacity_per_rank`` flags an overflowing step. + + ``ep_axes`` may be a mesh axis name or an ordered tuple of names. When it + is ``None``, the active ``MeshResource.ep_resource`` is used. """ - return _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank)[0] + return _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank, ep_axes)[0] -def _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank): +def _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank, ep_axes=None): if not jnp.issubdtype(topk_weights.dtype, jnp.floating): raise TypeError( f"ep_dispatch: topk_weights must be a floating dtype; got {topk_weights.dtype}." ) - token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx) + token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx, ep_axes=ep_axes) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( - cfg, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank + cfg, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank, ep_axes=ep_axes ) out_leading = tuple(tokens.shape[:-1]) primal = (recv_tokens, recv_topk_weights, handle_mem, token_counts, total_recv_tokens) return primal, (handle_mem, out_leading) -def _dispatch_bwd(cfg, recv_capacity_per_rank, res, g_outputs): +def _dispatch_bwd(cfg, recv_capacity_per_rank, ep_axes, res, g_outputs): del recv_capacity_per_rank handle_mem, out_leading = res # Re-pin cotangent: XLA transpose can drop the EP axis and feed the FFI a global tensor. - out_spec = _default_out_partition_spec() + out_spec = _default_out_partition_spec(ep_axes) spec = jax.sharding.PartitionSpec(*out_spec) g_recv_tokens = with_sharding_constraint(g_outputs[0], spec) g_recv_topk_weights = with_sharding_constraint(g_outputs[1], spec) @@ -304,6 +308,7 @@ def _dispatch_bwd(cfg, recv_capacity_per_rank, res, g_outputs): g_recv_topk_weights, out_leading, out_partition_spec=out_spec, + ep_axes=ep_axes, ) return (None, grad_tokens, grad_topk_weights) @@ -314,7 +319,7 @@ def _dispatch_bwd(cfg, recv_capacity_per_rank, res, g_outputs): # ── ep_combine (custom_vjp) ────────────────────────────────────────────────── -@partial(jax.custom_vjp, nondiff_argnums=(0, 4, 5)) +@partial(jax.custom_vjp, nondiff_argnums=(0, 4, 5, 6)) def ep_combine( cfg, handle_mem, @@ -322,6 +327,7 @@ def ep_combine( expert_out, num_local_tokens, out_sharding=None, + ep_axes=None, ): """Scatter-sum expert outputs back to source ranks. **Unweighted.** @@ -330,6 +336,8 @@ def ep_combine( not through this op. ``num_local_tokens`` is STATIC: int -> ``[T, H]``, tuple -> ``[*tuple, H]``. ``out_sharding`` defaults via ``_default_out_partition_spec``; only the leading dim may be sharded. + ``ep_axes`` accepts an explicit mesh axis name or tuple and otherwise falls + back to ``MeshResource.ep_resource``. """ return _combine_fwd( cfg, @@ -338,6 +346,7 @@ def ep_combine( expert_out, num_local_tokens, out_sharding, + ep_axes, )[0] @@ -348,24 +357,32 @@ def _combine_fwd( expert_out, num_local_tokens, out_sharding, + ep_axes=None, ): del token_counts if out_sharding is None: - out_sharding = _default_out_partition_spec() + out_sharding = _default_out_partition_spec(ep_axes) result = tex.ep_combine_fwd( - cfg, handle_mem, expert_out, num_local_tokens, out_partition_spec=out_sharding + cfg, + handle_mem, + expert_out, + num_local_tokens, + out_partition_spec=out_sharding, + ep_axes=ep_axes, ) return result, (handle_mem, expert_out.shape[-2]) -def _combine_bwd(cfg, _num_local_tokens, _out_sharding, res, g_result): +def _combine_bwd(cfg, _num_local_tokens, _out_sharding, ep_axes, res, g_result): handle_mem, recv_capacity_per_rank = res # Re-pin cotangent (same XLA-transpose workaround as _dispatch_bwd). if _out_sharding is None: - _out_sharding = _default_out_partition_spec() + _out_sharding = _default_out_partition_spec(ep_axes) spec = jax.sharding.PartitionSpec(*_out_sharding) g_result = with_sharding_constraint(g_result, spec) - grad_expert_out = tex.ep_combine_bwd(cfg, handle_mem, g_result, recv_capacity_per_rank) + grad_expert_out = tex.ep_combine_bwd( + cfg, handle_mem, g_result, recv_capacity_per_rank, ep_axes=ep_axes + ) return (None, None, grad_expert_out) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index cb10c7aa18..34b5049483 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -13,7 +13,7 @@ as ``self.param`` slots (with the right :func:`flax.linen.with_logical_partitioning` annotations so JAX's sharding layer FSDPs the params correctly). -2. Resolve the EP axis name from the active +2. Resolve the EP axes explicitly or from the active :class:`transformer_engine.jax.sharding.MeshResource`. 3. Forward all knobs to :func:`moe`. @@ -98,6 +98,9 @@ class _MoEBlock(TransformerEngineBase): replicated across non-EP axes within an EP group; set e.g. ``("fsdp",)`` for true FSDP-of-batch where each device owns a unique slice of the batch. + ep_axis : Optional[str | tuple[str, ...]] + Physical mesh axis or ordered tuple of axes used for expert parallelism. + Defaults to ``MeshResource.ep_resource``. apply_topk_weights_early : bool If ``True``, multiply expert outputs by their top-k weights *inside* each shard before ``ep_combine`` (saves one global @@ -146,6 +149,7 @@ class _MoEBlock(TransformerEngineBase): input_axes: Tuple[Optional[str], ...] = () # Parallelism + ep_axis: Optional[Union[str, Tuple[str, ...]]] = None data_parallelism_axes: Tuple[str, ...] = () # MoE knobs forwarded to ``moe()`` @@ -253,7 +257,9 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: jnp.float32, ) - ep_axis = get_active_resource_axis("ep_resource") + ep_axis = self.ep_axis + if ep_axis is None: + ep_axis = get_active_resource_axis("ep_resource") mesh = _get_mesh() data_parallel_size = 1 for axis in self.data_parallelism_axes: diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 55a85ebb2f..6b78730b89 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -16,8 +16,8 @@ Sharding model -------------- -* Inbound activations are 3D ``[B, S, H]`` sharded - ``((*data_parallelism_axes, ep_axis), None, None)``. The public +* Inbound activations are 3D ``[B, S, H]`` sharded over the combined + data-parallel and EP axes. The public :func:`moe` soft-repins this on entry and warns when a reshard is inserted. * The EP, grouped-quantize, and grouped-GEMM primitives operate at global @@ -584,7 +584,8 @@ def _moe_fwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") if ep_axis is None: raise ValueError("moe(...) requires ep_axis to be set (TE EP backend).") - num_ep = mesh.shape[ep_axis] + ep_axes = tex.ep._normalize_ep_axes(ep_axis) + num_ep = math.prod(mesh.shape[axis] for axis in ep_axes) if num_experts % num_ep != 0: raise ValueError(f"num_experts={num_experts} must be divisible by EP size={num_ep}") num_local_experts = num_experts // num_ep @@ -636,7 +637,7 @@ def _moe_fwd_rule( # ep must be innermost: ep_bootstrap forms NCCL EP comms from # consecutive global ranks (dp_color = rank // ep_size), so the # comm only stays within one model replica under (outer_dp, ep). - batch_pspec_axis = (*data_parallelism_axes, ep_axis) + batch_pspec_axis = (*data_parallelism_axes, *ep_axes) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, ep3_spec)) @@ -738,10 +739,10 @@ def _moe_fwd_rule( top_k=K, dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) - token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx_3d, ep_axes=ep_axes) token_counts = jax.lax.with_sharding_constraint(token_counts, NamedSharding(mesh, ep2_spec)) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( - cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr + cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr, ep_axes=ep_axes ) recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3_spec)) recv_topk_weights = jax.lax.with_sharding_constraint( @@ -816,6 +817,7 @@ def _ffn_fwd_body(*args): expert_outputs, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, + ep_axes=ep_axes, ) else: # HT combine is unweighted; apply routing weights before calling it. @@ -828,6 +830,7 @@ def _ffn_fwd_body(*args): weighted, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, + ep_axes=ep_axes, ) # output of MLP should be sharded the same way as the activation input output = with_sharding_constraint_by_logical_axes(output, input_axes) @@ -913,17 +916,20 @@ def _moe_bwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") B, S, _ = x_shape K = num_experts_per_tok + ep_axes = tex.ep._normalize_ep_axes(ep_axis) if not data_parallelism_axes: batch_pspec_axis: Any = ep_axis else: - batch_pspec_axis = (*data_parallelism_axes, ep_axis) + batch_pspec_axis = (*data_parallelism_axes, *ep_axes) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) out_partition_spec = (batch_pspec_axis, None, None) # ---------------- Combine bwd (global view) ---------------- d_output = jax.lax.with_sharding_constraint(d_output, NamedSharding(mesh, ep3_spec)) - grad_pre_combine = tex.ep_combine_bwd(ctx.cfg, ctx.handle_mem, d_output, recv_pr) + grad_pre_combine = tex.ep_combine_bwd( + ctx.cfg, ctx.handle_mem, d_output, recv_pr, ep_axes=ep_axes + ) grad_pre_combine = jax.lax.with_sharding_constraint( grad_pre_combine, NamedSharding(mesh, ep3_spec) ) @@ -1044,6 +1050,7 @@ def _ffn_bwd_body(*args): d_recv_w_total, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, + ep_axes=ep_axes, ) # ---------------- Routing bwd (global view) ---------------- @@ -1229,7 +1236,7 @@ def moe( noop_quantizer_set, noop_quantizer_set, ), - ep_axis: str, + ep_axis: Union[str, Tuple[str, ...]], data_parallelism_axes: Tuple[str, ...] = (), input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), @@ -1279,8 +1286,8 @@ def moe( Axis-name parameters: * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh - axis names* -- they index ``jax.sharding.Mesh.shape`` directly - (to compute ``num_ep`` / ``dp_size`` and to construct + axis names*. ``ep_axis`` may be an ordered tuple of names. They + are used to compute ``num_ep`` / ``dp_size`` and to construct ``P((dp..., ep), None, None)`` for the physical ``jax.lax.with_sharding_constraint`` calls that JAX requires to refer to real mesh axes). @@ -1313,7 +1320,8 @@ def moe( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - expected_leading: Any = (*data_parallelism_axes, ep_axis) if data_parallelism_axes else ep_axis + ep_axes = tex.ep._normalize_ep_axes(ep_axis) + expected_leading: Any = (*data_parallelism_axes, *ep_axes) if data_parallelism_axes else ep_axis expected_spec = P(expected_leading, None, None) actual_spec = getattr(getattr(x, "sharding", None), "spec", None) if actual_spec is not None and tuple(actual_spec) != tuple(expected_spec):