diff --git a/tests/jax/moe/test_moe_fsdp_layout.py b/tests/jax/moe/test_moe_fsdp_layout.py new file mode 100644 index 00000000000..2296dfe41ac --- /dev/null +++ b/tests/jax/moe/test_moe_fsdp_layout.py @@ -0,0 +1,368 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Isolate MoE weight layout errors using real packing and FSDP collectives. + +The fallback tests replace FP8 quantization and grouped GEMM with lossless +wrappers and JAX matmul to isolate layout/axis handling from kernel accuracy. +Weight preparation tests use real MXFP8 quantization, scales, and all-gathers. +""" + +import importlib +from types import SimpleNamespace + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from jax.experimental.shard_map import shard_map +from jax.sharding import Mesh, PartitionSpec as P +from transformer_engine.jax.quantize import QuantizerFactory, QuantizeLayout, ScalingMode +from transformer_engine.jax.quantize.dequantizer import Dequantizer + +moe = importlib.import_module("transformer_engine.jax.moe") + + +def fsdp_mesh(): + if len(jax.devices()) < 2: + pytest.skip("Requires two JAX devices for a real FSDP all-gather") + return Mesh(np.asarray(jax.devices()[:2]), ("fsdp",)) + + +def test_standard_gated_shards_match_unsharded_swiglu(): + mesh = fsdp_mesh() + # Each rank receives 128 contiguous columns: rank 0 gets only gates (1), + # rank 1 only ups (3). Each local half still meets pack's 32-column rule. + gate = jnp.ones((1, 128, 128), jnp.bfloat16) + up = 3 * gate + wi = jnp.concatenate((gate, up), axis=-1) + x = jnp.ones((1, 128), jnp.float32) / 128 + wo = jnp.ones((128, 1), jnp.float32) / 128 + quantizer = weight_quantizer(1) + + def output(tensor): + packed = jnp.concatenate(tensor.rowwise_tensor.dequantize()) + global_gate, global_up = moe.tex.unpack_gated_activation_pair(packed) + return (jax.nn.silu(x @ global_gate[0]) * (x @ global_up[0])) @ wo + + expected = output(prepare_weight(wi, quantizer, fused=True, native=False)) + + def local_forward(local_wi): + return prepare_weight(local_wi, quantizer, fused=True, native=False, axis=2) + + gathered = shard_map( + local_forward, + mesh=mesh, + in_specs=P(None, None, "fsdp"), + out_specs=P(), + check_rep=False, + )(wi) + actual = output(gathered) + # Both paths now yield 3*silu(1)=2.193176, rather than FSDP yielding 4.652113. + np.testing.assert_allclose(actual, expected, rtol=1e-6) + + +def weight_quantizer(experts): + return QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE_COLWISE, + n_groups=experts, + ) + + +def prepare_weight(wi, quantizer, *, fused, native, axis=None): + return moe._prepare_fc1_weight( + wi, + quantizer, + fused=fused, + native_layout=native, + quant_before_fsdp_ag=axis is not None, + fsdp_axis="fsdp", + fsdp_size=2, + sharded_axis=axis, + ) + + +def plain_weight_scales(tensor): + """Ignore unused allocation padding while checking every active scale byte.""" + kwargs = dict( + data_layout=tensor.data_layout, + is_colwise=tensor.is_colwise, + flatten_axis=tensor.flatten_axis - 1, + ) + padded = tensor.scaling_mode.get_scale_shape( + tensor.original_shape[1:], is_padded=True, **kwargs + ) + unpadded = tensor.scaling_mode.get_scale_shape( + tensor.original_shape[1:], is_padded=False, **kwargs + ) + scale_size = int(np.prod(padded)) + return jnp.stack( + [ + moe._unswizzle_mxfp8_grouped_scale( + tensor.scale_inv[e * scale_size : (e + 1) * scale_size], padded, tensor.is_colwise + )[: unpadded[0], : unpadded[1]] + for e in range(tensor.original_shape[0]) + ] + ).view(jnp.uint8) + + +def reference_quantize_weight(data, *args, **kwargs): + """Pure-JAX MXFP8 encoding for positive powers of two, including small shards. + + This avoids the legacy grouped quantizer's host-copy path for dimensions + below 128. The production gather, FP8 permutation, and swizzle still run. + """ + copies = [] + mode = ScalingMode.MXFP8_1D_SCALING + experts, rows, columns = data.shape + for colwise in (False, True): + blocks = ( + data.reshape(experts, rows // 32, 32, columns) + if colwise + else data.reshape(experts, rows, columns // 32, 32) + ) + maxima = jnp.max(blocks.astype(jnp.float32), axis=2 if colwise else 3) + exponent = jnp.log2(maxima).astype(jnp.int32) - 8 + inverse_scale = jnp.exp2(exponent.astype(jnp.float32)) + broadcast_scale = jnp.repeat(inverse_scale, 32, axis=1 if colwise else 2) + quantized = (data / broadcast_scale).astype(jnp.float8_e4m3fn) + scales = (exponent + 127).astype(jnp.uint8) + padded = mode.get_scale_shape((rows, columns), is_colwise=colwise, is_padded=True) + swizzled = [ + moe.swizzled_scale( + jnp.pad(scale, ((0, padded[0] - scale.shape[0]), (0, padded[1] - scale.shape[1]))), + 1, + colwise, + ).reshape(-1) + for scale in scales + ] + scale_inv = jnp.concatenate(swizzled) + scale_size = mode.get_grouped_scale_shape(data.shape, experts, colwise, flatten_axis=2)[0] + scale_inv = jnp.pad(scale_inv, (0, scale_size - scale_inv.size)) + copies.append( + moe.GroupedScaledTensor1x( + data=quantized.reshape(-1), + scale_inv=scale_inv, + amax=jnp.empty((0,), jnp.float32), + first_dims=None, + last_dims=None, + scaling_mode=mode, + dq_dtype=data.dtype, + _dq_func=Dequantizer.grouped_dequantize, + is_colwise=colwise, + data_layout="N", + flatten_axis=2, + original_shape=data.shape, + pre_swizzled=True, + ) + ) + return moe.ScaledTensor2x(*copies) + + +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("width", [64, 192]) +def test_small_gated_shards_preserve_quantized_pairs(monkeypatch, native, width): + mesh = fsdp_mesh() + source = jnp.exp2( + (jnp.arange(64) // 32)[None, :, None] + + (jnp.arange(width) // 32)[None, None, :] + + jnp.arange(2)[:, None, None] + ).astype(jnp.bfloat16) + gate, up = jnp.split(source, 2, axis=-1) + wi = moe.tex.pack_gated_activation_pair(gate, up).transpose(0, 2, 1) if native else source + monkeypatch.setattr(moe.tex, "grouped_quantize", reference_quantize_weight) + expected = prepare_weight(wi, None, fused=not native, native=native) + actual = shard_map( + lambda shard: prepare_weight( + shard, None, fused=not native, native=native, axis=1 if native else 2 + ), + mesh=mesh, + in_specs=P(None, "fsdp", None) if native else P(None, None, "fsdp"), + out_specs=P(), + check_rep=False, + )(wi) + for actual_copy, expected_copy in zip( + (actual.rowwise_tensor, actual.colwise_tensor), + (expected.rowwise_tensor, expected.colwise_tensor), + ): + np.testing.assert_array_equal( + actual_copy.data.view(jnp.uint8), expected_copy.data.view(jnp.uint8) + ) + np.testing.assert_array_equal( + plain_weight_scales(actual_copy), plain_weight_scales(expected_copy) + ) + decoded = jnp.concatenate(actual_copy.dequantize()) + execution_weight = moe.tex.pack_gated_activation_pair(gate, up) if not native else source + np.testing.assert_array_equal(decoded, execution_weight) + + +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("fused", [False, True]) +@pytest.mark.parametrize("shard_dimension", ["expert", "k", "gated"]) +@pytest.mark.parametrize("width", [256, 768]) +def test_quantized_fc1_gather_matches_unsharded(monkeypatch, native, fused, shard_dimension, width): + mesh = fsdp_mesh() + # Distinct column-block magnitudes, row magnitudes, and expert magnitudes + # expose scale permutations that equal-valued/constant-scale inputs hide. + source = jax.random.normal(jax.random.PRNGKey(9), (2, 256, width)) + column_magnitude = jnp.exp2(jnp.arange(width) // 32 % 7 - 3) + row_magnitude = jnp.exp2(jnp.arange(256) // 32 - 3) + source = ( + source + * column_magnitude + * row_magnitude[None, :, None] + * jnp.asarray([1.0, 4.0])[:, None, None] + ).astype(jnp.bfloat16) + gate, up = jnp.split(source, 2, axis=-1) + wi = moe.tex.pack_gated_activation_pair(gate, up).transpose(0, 2, 1) if native else source + axis = {"expert": 0, "k": 2 if native else 1, "gated": 1 if native else 2}[shard_dimension] + quantizer = weight_quantizer(2) + expected = prepare_weight(wi, quantizer, fused=fused, native=native) + spec = [None, None, None] + spec[axis] = "fsdp" + quantized_shapes = [] + original_quantize = moe.tex.grouped_quantize + + def record_quantize(data, *args, **kwargs): + quantized_shapes.append(data.shape) + return original_quantize(data, *args, **kwargs) + + monkeypatch.setattr(moe.tex, "grouped_quantize", record_quantize) + actual = shard_map( + lambda shard: prepare_weight(shard, quantizer, fused=fused, native=native, axis=axis), + mesh=mesh, + in_specs=P(*spec), + out_specs=P(), + check_rep=False, + )(wi) + # Only a local shard is quantized, once; no global BF16 requantization. + assert quantized_shapes + assert all(int(np.prod(shape)) == wi.size // 2 for shape in quantized_shapes) + for actual_copy, expected_copy in zip( + (actual.rowwise_tensor, actual.colwise_tensor), + (expected.rowwise_tensor, expected.colwise_tensor), + ): + np.testing.assert_array_equal( + actual_copy.data.view(jnp.uint8), expected_copy.data.view(jnp.uint8) + ) + np.testing.assert_array_equal( + plain_weight_scales(actual_copy), plain_weight_scales(expected_copy) + ) + np.testing.assert_array_equal( + jnp.concatenate(actual_copy.dequantize()), jnp.concatenate(expected_copy.dequantize()) + ) + + +@pytest.mark.parametrize("native_layout", [False, True]) +def test_unfused_k_sharded_weight_matches_unsharded(monkeypatch, native_layout): + mesh = fsdp_mesh() + + class LosslessTensor: + def __init__(self, data): + self.data = data + + def get_tensor(self, **kwargs): + return self + + def checkpoint(self, quantizer): + return self + + def quantize(data, *args, **kwargs): + return LosslessTensor(data) + + def gather(tensor, fsdp_axis, fsdp_size, sharded_axis): + assert fsdp_size == 2 + return LosslessTensor( + jax.lax.all_gather(tensor.data, fsdp_axis, axis=sharded_axis, tiled=True) + ) + + def gemm(lhs, rhs, *, contracting_dims, bias): + assert contracting_dims == ((1,), (1,)) + # One expert is sufficient to reproduce the matrix-axis error. + return lhs.data @ rhs.data[0] + + monkeypatch.setattr(moe.tex, "grouped_quantize", quantize) + monkeypatch.setattr(moe, "_gather_quantized_weight", gather) + monkeypatch.setattr(moe.tex, "grouped_gemm", gemm) + quantizers = SimpleNamespace(x=None, kernel=None) + gate = jnp.ones((1, 64, 64), jnp.float32) / 64 + up = 2 * gate + standard_wi = jnp.concatenate((gate, up), axis=-1) + wi = ( + moe.tex.pack_gated_activation_pair(gate, up).transpose(0, 2, 1) + if native_layout + else standard_wi + ) + x = jnp.ones((1, 1, 64), jnp.float32) + wo = jnp.ones((1, 64, 64), jnp.float32) / 64 + + def forward(weight, quant_before_fsdp_ag): + result, _ = moe._ffn_fwd_per_shard( + x, + jnp.ones((1, 1)), + jnp.asarray([1]), + weight, + wo, + None, + None, + None, + (quantizers, quantizers), + num_local_experts=1, + activation_type="silu", + apply_topk_weights_early=False, + use_cudnn_jax_fusion=False, + wi_0_checkpoint_name=None, + wi_1_checkpoint_name=None, + wo_checkpoint_name=None, + cudnn_native_weight_layout=native_layout, + quant_before_fsdp_ag=quant_before_fsdp_ag, + fsdp_axis="fsdp", + fsdp_size=2, + wi_fsdp_axis=(2 if native_layout else 1), + wo_fsdp_axis=None, + ) + return result + + expected = forward(wi, False) + independent_reference = (jax.nn.silu(x @ gate[0]) * (x @ up[0])) @ wo[0] + np.testing.assert_allclose(expected, independent_reference, rtol=1e-6) + actual = shard_map( + lambda local_wi: forward(local_wi, True), + mesh=mesh, + in_specs=P(None, None, "fsdp") if native_layout else P(None, "fsdp", None), + out_specs=P(), + check_rep=False, + )(wi) + np.testing.assert_allclose(actual, expected, rtol=1e-6) + + +@pytest.mark.parametrize( + "shape,flag,expected", + [ + ((2, 128, 256), None, False), + ((2, 256, 128), None, True), + ((2, 128, 128), False, False), + ((2, 128, 128), True, True), + ((2, 128, 128), None, ValueError), + ((2, 128, 256), True, ValueError), + ((2, 256, 128), False, ValueError), + ((2, 128, 256), "false", TypeError), + ], +) +def test_layout_resolution(shape, flag, expected): + wi = jnp.zeros(shape) + if isinstance(expected, type): + with pytest.raises(expected): + moe._resolve_cudnn_native_weight_layout(wi, 128, flag) + else: + assert moe._resolve_cudnn_native_weight_layout(wi, 128, flag) is expected + + +def test_layout_argument_is_optional_and_static(): + import inspect + + assert inspect.signature(moe.moe).parameters["cudnn_native_weight_layout"].default is None + assert inspect.signature(moe._moe).parameters["cudnn_native_weight_layout"].default is None + assert 33 in moe._moe.nondiff_argnums diff --git a/tests/jax/moe/test_moe_resource_api.py b/tests/jax/moe/test_moe_resource_api.py new file mode 100644 index 00000000000..f5b058e314c --- /dev/null +++ b/tests/jax/moe/test_moe_resource_api.py @@ -0,0 +1,176 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MoE resource resolution and MeshResource-only API without EP kernels.""" + +import importlib +import inspect + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from jax.sharding import Mesh + +from transformer_engine.jax.flax import _MoEBlock +from transformer_engine.jax.moe import _moe_mesh_axes, _resolve_moe_mesh_resource +from transformer_engine.jax.sharding import MeshResource, global_mesh_resource, global_shard_guard + + +@pytest.fixture(autouse=True) +def no_global_resource(): + with global_shard_guard(None): + yield + + +def test_resource_required_for_function_and_block(): + module = importlib.import_module("transformer_engine.jax.moe") + with pytest.raises(ValueError, match="active global_shard_guard"): + module.moe(None, None, None, None, num_experts=2, num_experts_per_tok=1) + with pytest.raises(ValueError, match="active global_shard_guard"): + _MoEBlock().init(jax.random.PRNGKey(0), jnp.ones((1, 1, 4))) + + +def test_explicit_resource_overrides_global_and_is_snapshotted(): + explicit = MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="ep") + enclosing = MeshResource(ep_resource="other") + with global_shard_guard(enclosing): + resolved, quantize = _resolve_moe_mesh_resource(explicit, True) + assert global_mesh_resource() is enclosing + explicit.ep_resource = "changed" + assert _moe_mesh_axes(resolved) == ("ep", ("dp", "fsdp")) + assert quantize is True + + +def test_global_default_and_duplicate_outer_axis(): + with global_shard_guard( + MeshResource(dp_resource="fsdp", fsdp_resource="fsdp", ep_resource="ep") + ): + resource, quantize = _resolve_moe_mesh_resource() + assert _moe_mesh_axes(resource) == ("ep", ("fsdp",)) + assert quantize is False + + +def test_ep_partitioning_retains_bind_time_resources(): + from transformer_engine.jax.cpp_extensions import ep + from jax.sharding import PartitionSpec as P + + with global_shard_guard(MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="ep")): + captured = ep._capture_ep_resource_axes() + spec = P(("dp", "fsdp", "ep"), None, None) + assert ep._leading_axis_ok(spec, captured) == (True, "ep", ("dp", "fsdp")) + assert ep._ep_output_spec(None, None, resource_axes=captured) == spec + assert ep._ep_spec_ok(spec, 2, resource_axes=captured) + assert not ep._ep_spec_ok(P(("wrong", "ep"), None, None), 2, resource_axes=captured) + + +def test_fused_ep_partitioning_retains_bind_time_resources(monkeypatch): + from types import SimpleNamespace + from jax.sharding import NamedSharding, PartitionSpec as P + from transformer_engine.jax.cpp_extensions import ep + + captured = {} + + def bind(*args, **kwargs): + captured.update(kwargs) + return () + + monkeypatch.setattr(ep.EpPrepareAndDispatchPrimitive.outer_primitive, "bind", bind) + cfg = SimpleNamespace(top_k=1, dispatch_output_per_expert_alignment=256, topk_idx_dtype=0) + with global_shard_guard(MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="ep")): + inspect.unwrap(ep._ep_fused_prepare_and_dispatch_fwd)(cfg, None, None, None, 256) + + mesh = Mesh(np.asarray(jax.devices()[:1]).reshape(1, 1, 1), ("dp", "fsdp", "ep")) + sharding = NamedSharding(mesh, P(("dp", "fsdp", "ep"), None, None)) + args = [SimpleNamespace(sharding=sharding)] * 3 + _, _, outputs, _ = ep.EpPrepareAndDispatchPrimitive.partition( + captured["top_k"], + captured["dispatch_output_per_expert_alignment"], + captured["topk_idx_dtype"], + captured["recv_capacity_per_rank"], + captured["is_outer"], + captured["resource_axes"], + mesh, + args, + None, + ) + assert all(output.spec[0] == ("dp", "fsdp", "ep") for output in outputs) + + +@pytest.mark.parametrize( + "resource,quantize,error", + [ + (MeshResource(), False, "ep_resource"), + (MeshResource(ep_resource="ep", dp_resource="ep"), False, "distinct"), + (MeshResource(ep_resource="ep"), True, "fsdp_resource"), + ], +) +def test_invalid_resource(resource, quantize, error): + with pytest.raises(ValueError, match=error): + _resolve_moe_mesh_resource(resource, quantize) + + +def test_bool_and_resource_types(): + with pytest.raises(TypeError, match="must be a bool"): + _resolve_moe_mesh_resource(MeshResource(ep_resource="ep"), "true") + with pytest.raises(TypeError, match="must be a MeshResource"): + _resolve_moe_mesh_resource("ep") + + +@pytest.mark.parametrize("quantize", [False, True]) +def test_public_api_uses_selected_resource(monkeypatch, quantize): + module = importlib.import_module("transformer_engine.jax.moe") + signature = inspect.signature(module._moe) + captured = {} + + def fake_vjp(*args): + captured.update(signature.bind(*args).arguments) + assert global_mesh_resource().ep_resource == "ep" + return args[0], None, jnp.zeros((1,), jnp.int32) + + monkeypatch.setattr(module, "_moe", fake_vjp) + mesh = Mesh(np.asarray(jax.devices()[:1]).reshape(1, 1, 1), ("dp", "fsdp", "ep")) + kwargs = dict( + num_experts=2, + num_experts_per_tok=1, + mesh_resource=MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="ep"), + quant_before_fsdp_ag=quantize, + ) + with jax.set_mesh(mesh): + module.moe( + jnp.ones((1, 1, 4)), + jnp.ones((4, 2)), + jnp.ones((2, 4, 8)), + jnp.ones((2, 4, 4)), + **kwargs, + ) + assert _moe_mesh_axes(captured["mesh_resource"]) == ("ep", ("dp", "fsdp")) + assert captured["quant_before_fsdp_ag"] is quantize + assert "ep_axis" not in captured + with pytest.raises(AssertionError, match="Global mesh resource is not set"): + global_mesh_resource() + + +@pytest.mark.parametrize("quantize", [False, True]) +def test_flax_block_forwards_boolean(monkeypatch, quantize): + module = importlib.import_module("transformer_engine.jax.flax.moe") + resource = MeshResource(fsdp_resource="fsdp", ep_resource="ep") + captured = {} + + def fake_moe(inputs, *args, **kwargs): + captured.update(kwargs) + return inputs, None, jnp.zeros((1,), jnp.int32) + + monkeypatch.setattr(module, "moe", fake_moe) + mesh = Mesh(np.asarray(jax.devices()[:1]).reshape(1, 1), ("fsdp", "ep")) + block = _MoEBlock( + num_experts=2, + intermediate_size=4, + mesh_resource=resource, + quant_before_fsdp_ag=quantize, + ) + with jax.set_mesh(mesh): + block.init(jax.random.PRNGKey(0), jnp.ones((1, 1, 4))) + assert captured["quant_before_fsdp_ag"] is quantize + assert _moe_mesh_axes(captured["mesh_resource"]) == ("ep", ("fsdp",)) diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh index 9fcbb9dd9ad..ee1c55fc33b 100755 --- a/tests/jax/run_te_ep_moe.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -23,6 +23,13 @@ if [ "$NUM_GPUS" -lt 4 ]; then exit 0 fi +# Bitwise fallback/layout comparisons need stable expert-row ordering in the +# split dispatch path. Bootstrap previously selected this scan mode globally. +# Keep it local to these tests; the fused dispatch path uses its own ordering. +if [ "${NVTE_EP_FUSED_PREPARE_DISPATCH:-0}" != "1" ]; then + export NCCL_EP_HT_EM_AG_SCAN_MODE="${NCCL_EP_HT_EM_AG_SCAN_MODE:-1}" +fi + export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}" export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.5}" export TE_EP_MOE_COORDINATOR_ADDRESS="${TE_EP_MOE_COORDINATOR_ADDRESS:-127.0.0.1:13457}" @@ -33,6 +40,9 @@ echo " test file : $TEST_FILE" echo " coordinator : $TE_EP_MOE_COORDINATOR_ADDRESS" echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" +if [ "$#" -gt 0 ]; then + echo " extra pytest args : $*" +fi echo "============================================================" if [ -n "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then @@ -44,6 +54,8 @@ fi echo "Per-process logs: $LOG_DIR" PIDS=() +EXITS=() +FAILED=0 cleanup() { for pid in "${PIDS[@]:-}"; do @@ -61,25 +73,25 @@ cleanup() { trap cleanup EXIT INT TERM for i in $(seq 0 $((NUM_GPUS - 1))); do - LOG_FILE="$LOG_DIR/proc_${i}.log" - PYTEST_CMD=( + log_file="$LOG_DIR/proc_${i}.log" + pytest_cmd=( python3 -m pytest -c "$PYTEST_INI" "$TEST_FILE" -p no:typeguard -v -s --num-process="$NUM_GPUS" --process-id="$i" + "$@" ) if [ "$i" -eq 0 ]; then echo "=== Live output from process 0 ===" - "${PYTEST_CMD[@]}" 2>&1 | tee "$LOG_FILE" & + "${pytest_cmd[@]}" 2>&1 | tee "$log_file" & else - "${PYTEST_CMD[@]}" > "$LOG_FILE" 2>&1 & + "${pytest_cmd[@]}" > "$log_file" 2>&1 & fi PIDS+=("$!") done -EXITS=() for pid in "${PIDS[@]}"; do if wait "$pid"; then EXITS+=("0") @@ -89,16 +101,11 @@ for pid in "${PIDS[@]}"; do done echo -echo "============================================================" echo "Per-process exit codes:" for i in "${!EXITS[@]}"; do echo " proc $i -> ${EXITS[$i]}" 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 -# GPUs) as success. -FAILED=0 for e in "${EXITS[@]}"; do if [ "$e" != "0" ] && [ "$e" != "5" ]; then FAILED=1 @@ -117,6 +124,4 @@ fi echo "[run_te_ep_moe.sh] at least one process FAILED" echo " retaining logs at $LOG_DIR for diagnosis" -echo " process 0 tail:" -tail -20 "$LOG_DIR/proc_0.log" 2>/dev/null || true exit 1 diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 9c3b489c1fc..7ee4cb8d8fb 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -2,7 +2,13 @@ # # See LICENSE for license information. +import importlib +import inspect +import warnings +from types import SimpleNamespace + import jax +import numpy as np import jax.numpy as jnp import pytest from jax import jit, value_and_grad @@ -49,6 +55,10 @@ from transformer_engine.jax.dense import dense, grouped_dense from transformer_engine.jax.layernorm_dense import layernorm_dense from transformer_engine.jax.cpp_extensions.topk import topk +import transformer_engine_jax + +_glu_adapter = importlib.import_module("transformer_engine.jax.cpp_extensions.grouped_gemm_glu") +_moe_module = importlib.import_module("transformer_engine.jax.moe") GEMM_CASES = [ (256, 256, 512), @@ -2128,3 +2138,404 @@ def test_topk_2d(self, dtype, problem_size): prim_gathered = jnp.take_along_axis(x, prim_idx, axis=1) ref_gathered = jnp.take_along_axis(x, ref_idx, axis=1) assert_allclose(prim_gathered, ref_gathered, dtype=dtype) + + +class TestGroupedGemmGluFallback: + """MoE API compatibility and dispatch with mocked kernels on one JAX device.""" + + @staticmethod + def _api(arguments, change=None): + def function(**kwargs): + raise AssertionError("Probes must not execute kernels") + + parameters = [ + inspect.Parameter( + name, + ( + inspect.Parameter.POSITIONAL_ONLY + if change == "positional" and name == "a_tensor" + else inspect.Parameter.POSITIONAL_OR_KEYWORD + ), + ) + for name in arguments + if not (change == "renamed" and name == "discrete_col_sfd") + ] + if change in ("optional", "mandatory"): + parameters.append( + inspect.Parameter( + "new_arg", + inspect.Parameter.KEYWORD_ONLY, + default=None if change == "optional" else inspect.Parameter.empty, + ) + ) + function.__signature__ = inspect.Signature(parameters) + return function + + @pytest.fixture + def frontend(self, monkeypatch): + module = SimpleNamespace( + __name__="cudnn", + __version__="1.32.0", + grouped_gemm_glu_wrapper_sm100=self._api(_glu_adapter._FORWARD_ARGS), + grouped_gemm_dglu_wrapper_sm100=self._api(_glu_adapter._BACKWARD_ARGS), + grouped_gemm_glu_jax_sm100=self._api(_glu_adapter._JAX_FORWARD_ARGS), + grouped_gemm_dglu_jax_sm100=self._api(_glu_adapter._JAX_BACKWARD_ARGS), + ) + original = importlib.import_module + + def import_module(name, package=None): + if name == "cudnn": + return module + if name == "cutlass.jax": + return SimpleNamespace(is_available=lambda: True) + return original(name, package) + + monkeypatch.setattr(_glu_adapter.importlib, "import_module", import_module) + return module + + @pytest.mark.parametrize( + "version,expected", + [ + ("1.31.0", False), + ("1.32.0.dev1", False), + ("1.32.0", True), + ("1.33.0", True), + ("invalid", False), + (None, False), + ], + ) + def test_frontend_version(self, frontend, monkeypatch, version, expected): + # Matching signatures alone are insufficient for older/debug frontends. + if version is None: + del frontend.__version__ + else: + frontend.__version__ = version + available, reason = _glu_adapter.grouped_gemm_glu_dependencies_available() + assert available is expected + assert bool(reason) is not expected + if not expected: + monkeypatch.setattr( + transformer_engine_jax, "get_device_compute_capability", lambda _: 100 + ) + with pytest.warns(UserWarning, match="falling back"): + assert _moe_module._select_cudnn_jax_fusion([]) is False + + @pytest.mark.parametrize("fc1_dgrad", [jnp.float8_e4m3fn, jnp.float8_e5m2]) + @pytest.mark.parametrize("fc2_dgrad", [jnp.float8_e4m3fn, jnp.float8_e5m2]) + def test_fused_dgrad_dtypes(self, monkeypatch, fc1_dgrad, fc2_dgrad): + monkeypatch.setattr(_moe_module, "_get_mesh", lambda: None) + quantizers = tuple( + QuantizerFactory.create_set( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + fwd_dtype=jnp.float8_e4m3fn, + bwd_dtype=dtype, + is_2x2x=True, + n_groups=2, + ) + for dtype in (fc1_dgrad, fc2_dgrad) + ) + errors = _moe_module._cudnn_jax_fusion_rejection_reasons( + jnp.ones((1, 128), jnp.bfloat16), + jnp.ones((2, 128, 256), jnp.bfloat16), + None, + None, + quantizers, + num_experts=2, + activation_type="silu", + ep_axis="ep", + ) + assert any("FC1 dgrad" in error for error in errors) == (fc1_dgrad == jnp.float8_e5m2) + assert any("FC2 dgrad" in error for error in errors) == (fc2_dgrad == jnp.float8_e5m2) + assert bool(errors) == (jnp.float8_e5m2 in (fc1_dgrad, fc2_dgrad)) + + @pytest.mark.parametrize( + "name,args", + [ + ("grouped_gemm_glu_wrapper_sm100", _glu_adapter._FORWARD_ARGS), + ("grouped_gemm_dglu_wrapper_sm100", _glu_adapter._BACKWARD_ARGS), + ("grouped_gemm_glu_jax_sm100", _glu_adapter._JAX_FORWARD_ARGS), + ("grouped_gemm_dglu_jax_sm100", _glu_adapter._JAX_BACKWARD_ARGS), + ], + ) + @pytest.mark.parametrize( + "change", + [ + "missing", + "not_callable", + "renamed", + "mandatory", + "optional", + "positional", + "opaque", + ], + ) + def test_api_signature(self, frontend, name, args, change): + if change == "missing": + delattr(frontend, name) + elif change == "not_callable": + setattr(frontend, name, 42) + elif change == "opaque": + setattr(frontend, name, lambda **kwargs: None) + else: + setattr(frontend, name, self._api(args, change)) + available, reason = _glu_adapter.grouped_gemm_glu_dependencies_available() + assert available == (change == "optional") + assert bool(reason) == (change != "optional") + + @pytest.mark.parametrize("name", ["grouped_gemm_glu_jax_sm100", "grouped_gemm_dglu_jax_sm100"]) + def test_bf16_only_jax_api_falls_back(self, frontend, monkeypatch, name): + # Frontend 1.29 exposes JAX entry points for BF16 discrete weights, but + # its shared MXFP8 wrapper cannot dispatch traced arrays to them. + arguments = ("a_tensor", "padded_offsets", "alpha_tensor", "b_ptrs", "n", "prob_tensor") + if "dglu" in name: + arguments += ("c_tensor", "beta_tensor") + setattr(frontend, name, self._api(arguments)) + monkeypatch.setattr(transformer_engine_jax, "get_device_compute_capability", lambda _: 100) + available, reason = _glu_adapter.grouped_gemm_glu_dependencies_available() + assert not available + assert name in reason + for argument in ("b_tensor", "sfa_tensor", "sfb_tensor"): + assert argument in reason + with pytest.warns(UserWarning, match="falling back"): + assert _moe_module._select_cudnn_jax_fusion([]) is False + + @pytest.mark.parametrize( + "capability,missing,expected", + [ + (90, None, False), + (100, None, True), + (103, None, True), + (107, None, True), + (120, None, True), + (100, "grouped_gemm_glu_wrapper_sm100", False), + (100, "grouped_gemm_dglu_wrapper_sm100", False), + (100, "grouped_gemm_glu_jax_sm100", False), + (100, "grouped_gemm_dglu_jax_sm100", False), + ], + ) + def test_fused_or_unfused_selection(self, frontend, monkeypatch, capability, missing, expected): + if missing: + delattr(frontend, missing) + monkeypatch.setattr( + transformer_engine_jax, "get_device_compute_capability", lambda _: capability + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert _moe_module._select_cudnn_jax_fusion([]) == expected + if expected: + assert not caught + else: + assert len(caught) == 1 + message = str(caught[0].message) + assert ( + "falling back" in message + and "shared grouped_gemm_glu/grouped_gemm_dglu JAX APIs" in message + ) + assert "unfused TE" in message + + @pytest.mark.parametrize("name", ["cudnn", "cutlass.jax"]) + def test_missing_dependency(self, monkeypatch, name): + original = importlib.import_module + + def import_module(module, package=None): + if module == name: + raise ModuleNotFoundError(f"No module named {module}") + return original(module, package) + + monkeypatch.setattr(_glu_adapter.importlib, "import_module", import_module) + available, reason = _glu_adapter.grouped_gemm_glu_dependencies_available() + assert not available and name in reason + + def test_ineligible_call_and_device_query_failure(self, frontend, monkeypatch): + monkeypatch.setattr(transformer_engine_jax, "get_device_compute_capability", lambda _: 107) + with pytest.warns(UserWarning, match="bias unsupported"): + assert _moe_module._select_cudnn_jax_fusion(["bias unsupported"]) is False + + def fail(_): + raise RuntimeError("device query failed") + + monkeypatch.setattr(transformer_engine_jax, "get_device_compute_capability", fail) + with pytest.warns(UserWarning, match="device query failed"): + assert _moe_module._select_cudnn_jax_fusion([]) is False + + @pytest.mark.parametrize( + "request_fusion,missing,native_layout,square,expected", + [ + (None, None, False, False, True), + (True, None, False, False, True), + (False, None, False, False, False), + (True, "grouped_gemm_glu_wrapper_sm100", False, False, False), + (True, "grouped_gemm_dglu_wrapper_sm100", False, False, False), + (False, "grouped_gemm_glu_wrapper_sm100", False, False, False), + (False, "grouped_gemm_dglu_wrapper_sm100", False, False, False), + (True, None, True, False, True), + (True, None, False, True, True), + (True, None, True, True, True), + ], + ) + def test_public_moe_passes_fusion_bool( + self, frontend, monkeypatch, native_layout, square, missing, expected, request_fusion + ): + from jax.sharding import Mesh + from transformer_engine.jax.sharding import MeshResource + + if missing: + delattr(frontend, missing) + monkeypatch.setattr(transformer_engine_jax, "get_device_compute_capability", lambda _: 107) + monkeypatch.setattr( + _moe_module, "_cudnn_jax_fusion_rejection_reasons", lambda *args, **kwargs: [] + ) + mesh = Mesh(np.asarray(jax.devices()[:1]), ("ep",)) + monkeypatch.setattr(_moe_module, "_get_mesh", lambda: mesh) + monkeypatch.setattr(_moe_module, "_with_sharding_constraint_cast_bwd", lambda x, _: x) + received = [] + signature = inspect.signature(_moe_module._moe) + + def execute(*args): + bound = signature.bind(*args).arguments + assert bound["use_cudnn_fusion"] is (request_fusion is not False) + assert bound["cudnn_native_weight_layout"] is native_layout + received.append(bound["use_cudnn_jax_fusion"]) + return args[0], None, jnp.asarray(0) + + monkeypatch.setattr(_moe_module, "_moe", execute) + x = jnp.ones((1, 1, 128), jnp.bfloat16) + wi = jnp.ones((2, 256, 128) if native_layout else (2, 128, 256), jnp.bfloat16) + kwargs = {} if request_fusion is None else {"use_cudnn_fusion": request_fusion} + if square: + wi = jnp.ones((2, 128, 128), jnp.bfloat16) + kwargs["cudnn_native_weight_layout"] = native_layout + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + output, _, _ = _moe_module.moe( + x, + jnp.ones((128, 2)), + wi, + jnp.ones((2, 128, 128)), + num_experts=2, + num_experts_per_tok=1, + mesh_resource=MeshResource(ep_resource="ep"), + **kwargs, + ) + assert received == [expected] + assert output is x + assert len(caught) == (0 if request_fusion is False or expected else 1) + + @pytest.mark.parametrize("native_layout", [False, True]) + @pytest.mark.parametrize("gather_gated_dimension", [False, True]) + def test_fused_forward_uses_shared_glu( + self, monkeypatch, native_layout, gather_gated_dimension + ): + + class SelectedKernel(Exception): + pass + + def selected(*args, **kwargs): + assert args[1].shape == (1, 256, 128) + raise SelectedKernel + + monkeypatch.setattr(_moe_module.tex, "grouped_gemm_glu", selected) + + def quantize(data, *args, **kwargs): + return SimpleNamespace( + get_tensor=lambda **kwargs: SimpleNamespace( + data=data, scale_inv=jnp.ones((1,), dtype=jnp.uint8) + ) + ) + + monkeypatch.setattr(_moe_module.tex, "grouped_quantize", quantize) + + def gather(tensor, fsdp_axis, fsdp_size, sharded_axis): + data = tensor.get_tensor().data + assert sharded_axis == (1 if native_layout else 2) + return quantize(jnp.concatenate([data, data], axis=sharded_axis)) + + monkeypatch.setattr(_moe_module, "_gather_quantized_weight", gather) + + def reorder(tensor, *, interleave): + assert interleave and gather_gated_dimension and not native_layout + gate, up = jnp.split(tensor.get_tensor().data, 2, axis=-1) + return quantize(_moe_module.tex.pack_gated_activation_pair(gate, up)) + + monkeypatch.setattr(_moe_module, "_reorder_quantized_swiglu_weight", reorder) + quantizers = SimpleNamespace(x=SimpleNamespace(q_dtype=jnp.float8_e4m3fn), kernel=None) + kwargs = dict.fromkeys(inspect.signature(_moe_module._ffn_fwd_per_shard).parameters) + combined = 128 if gather_gated_dimension else 256 + kwargs.update( + recv_tokens_local=jnp.ones((1, 256, 128), jnp.bfloat16), + recv_topk_weights_local=jnp.ones((1, 256)), + token_counts_local=jnp.asarray([256]), + wi=jnp.ones((1, combined, 128) if native_layout else (1, 128, combined), jnp.bfloat16), + wo=jnp.ones((1, 128, 128), jnp.bfloat16), + quantizer_sets=(quantizers, quantizers), + num_local_experts=1, + use_cudnn_jax_fusion=True, + cudnn_native_weight_layout=native_layout, + quant_before_fsdp_ag=gather_gated_dimension, + wi_fsdp_axis=(1 if native_layout else 2), + fsdp_axis="fsdp", + fsdp_size=2, + ) + with pytest.raises(SelectedKernel): + _moe_module._ffn_fwd_per_shard(**kwargs) + + @pytest.mark.parametrize("request_fusion", [None, True, False]) + def test_flax_forwards_fusion_bool(self, monkeypatch, request_fusion): + from transformer_engine.jax.flax import _MoEBlock + from transformer_engine.jax.sharding import MeshResource + + flax_moe = importlib.import_module("transformer_engine.jax.flax.moe") + received = [] + + def execute(inputs, *args, **kwargs): + received.append(kwargs["use_cudnn_fusion"]) + return inputs, None, jnp.asarray(0) + + monkeypatch.setattr(flax_moe, "moe", execute) + kwargs = {} if request_fusion is None else {"use_cudnn_fusion": request_fusion} + block = _MoEBlock( + num_experts=2, + intermediate_size=32, + mesh_resource=MeshResource(ep_resource="ep"), + **kwargs, + ) + block.init(jax.random.PRNGKey(0), jnp.ones((1, 1, 32))) + assert received == [request_fusion is not False] + + @pytest.mark.parametrize("request_fusion", [True, False]) + def test_capacity_follows_explicit_fusion_bool(self, request_fusion): + kwargs = dict(num_experts=8, num_experts_per_tok=2, max_tokens_per_rank=64, ep_size=2) + alignment = _moe_module._CUDNN_JAX_ALIGN_SIZE if request_fusion else _moe_module._ALIGN_SIZE + expected = _moe_module.get_moe_recv_capacity_per_rank(**kwargs, alignment=alignment) + assert ( + _moe_module.get_moe_recv_capacity_per_rank(**kwargs, use_cudnn_fusion=request_fusion) + == expected + ) + if request_fusion: + assert _moe_module.get_moe_recv_capacity_per_rank(**kwargs) == expected + + @pytest.mark.parametrize("invalid", [0, 1, None, "true"]) + def test_fusion_argument_requires_bool(self, invalid): + with pytest.raises(TypeError, match="use_cudnn_fusion must be a bool"): + _moe_module.moe( + None, + None, + None, + None, + num_experts=2, + num_experts_per_tok=1, + use_cudnn_fusion=invalid, + ) + with pytest.raises(TypeError, match="use_cudnn_fusion must be a bool"): + _moe_module.get_moe_recv_capacity_per_rank( + num_experts=2, + num_experts_per_tok=1, + max_tokens_per_rank=16, + ep_size=1, + use_cudnn_fusion=invalid, + ) + + def test_vjp_bool_is_static_and_defaults_to_true(self): + assert inspect.signature(_moe_module._moe).parameters["use_cudnn_fusion"].default is True + assert inspect.signature(_moe_module.moe).parameters["use_cudnn_fusion"].default is True + assert 32 in _moe_module._moe.nondiff_argnums diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 015c73343aa..752fa680885 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -29,7 +29,8 @@ single run supports — shape, dtype, finiteness AND numerical parity vs a pure-JAX reference. Variations on the block are pytest parametrize values rather than separate test -classes: +classes. The module-scoped fusion parameter runs all unfused cases first, +then tears down EP and runs all cuDNN cases in the same processes: * ``test_forward`` covers BF16 and MXFP8 forward execution across a curated set of configurations (softmax/sigmoid scoring, optional @@ -42,6 +43,7 @@ main+aux grads stay finite) in two consolidated tests. """ +import importlib import os os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") @@ -55,7 +57,7 @@ import numpy as np import pytest -from jax.experimental import mesh_utils +from jax.experimental import mesh_utils, multihost_utils from jax.sharding import Mesh, NamedSharding, PartitionSpec as P from flax.linen import partitioning as nn_partitioning @@ -120,16 +122,14 @@ def _read_mp_options(): from transformer_engine.jax.flax import _MoEBlock as MoEBlock from transformer_engine.jax.moe import ( - _ALIGN_SIZE, get_moe_recv_capacity_per_rank, moe, record_ep_bootstrap_signature_for_moe, ) -from transformer_engine.jax.ep import ep_bootstrap -from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.jax.ep import ep_bootstrap, ep_finalize +from transformer_engine.common.recipe import Format, MXFP8BlockScaling from transformer_engine.jax.sharding import MeshResource, global_shard_guard - # ----------------------------------------------------------------------------- # Mesh / shape config # ----------------------------------------------------------------------------- @@ -145,6 +145,8 @@ def _read_mp_options(): LOGICAL_AXIS_RULES = ( ("exp", EP_AXIS), + ("expert_weight_fsdp", (EP_AXIS, FSDP_AXIS)), + ("gated_weight_fsdp", FSDP_AXIS), ("embed", FSDP_AXIS), ("mlp", None), ("batch", (FSDP_AXIS, EP_AXIS)), @@ -187,8 +189,14 @@ def _read_mp_options(): # ----------------------------------------------------------------------------- +@pytest.fixture(scope="module", params=[False, True], ids=["unfused", "cudnn"], autouse=True) +def use_cudnn_fusion(request): + """Group every test by fusion mode before changing EP's cached alignment.""" + return request.param + + @pytest.fixture(scope="module") -def mesh(): +def mesh(use_cudnn_fusion): if jax.device_count() < NUM_DEVICES_REQUIRED: pytest.skip( f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" @@ -210,6 +218,7 @@ def mesh(): num_experts_per_tok=TOPK, max_tokens_per_rank=max_tokens_per_rank, ep_size=EP_SIZE, + use_cudnn_fusion=use_cudnn_fusion, ) # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather @@ -232,7 +241,15 @@ def mesh(): hidden_dim=HIDDEN, ep_size=EP_SIZE, ) - return mesh_obj + mode = "cudnn" if use_cudnn_fusion else "unfused" + try: + yield mesh_obj + finally: + # Every rank finishes the mode before releasing cached EP executables. + jax.effects_barrier() + multihost_utils.sync_global_devices(f"te_moe_{mode}_finished") + ep_finalize() + multihost_utils.sync_global_devices(f"te_moe_{mode}_finalized") # ----------------------------------------------------------------------------- @@ -359,12 +376,16 @@ def _make_block( expert_bias_init=None, input_axes=("batch", None, None), quantization_recipe=None, + dispatch_checkpoint_name=None, + quant_before_fsdp_ag=False, + mesh_resource=None, + use_cudnn_fusion, ): kwargs = dict( num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, intermediate_size=INTER, - data_parallelism_axes=(FSDP_AXIS,), + mesh_resource=mesh_resource, apply_topk_weights_early=apply_topk_weights_early, aux_loss_coeff=aux_loss_coeff, use_expert_routing_bias=use_expert_routing_bias, @@ -372,6 +393,9 @@ def _make_block( dtype=DTYPE, input_axes=input_axes, quantization_recipe=quantization_recipe, + dispatch_checkpoint_name=dispatch_checkpoint_name, + quant_before_fsdp_ag=quant_before_fsdp_ag, + use_cudnn_fusion=use_cudnn_fusion, ) # Custom expert_bias_init lets tests inject a non-zero expert_bias without # poking variables['params'] post-init. @@ -380,6 +404,11 @@ def _make_block( return MoEBlock(**kwargs) +@pytest.fixture +def make_block(use_cudnn_fusion): + return partial(_make_block, use_cudnn_fusion=use_cudnn_fusion) + + def _strong_expert_bias_init(key, shape, dtype): """Half +5, half -5 — large enough to force topk onto the +ve half.""" del key @@ -563,6 +592,121 @@ def _quantization_recipe(quantization): _QUANTIZATION_CASES.append(pytest.param("mxfp8", id="mxfp8")) +@pytest.mark.parametrize("native_weight_layout", [False, True]) +@pytest.mark.parametrize("fsdp_dimension", ["k", "gated", "expert"]) +def test_quantized_weight_gather_matches_full_precision_gather( + mesh, monkeypatch, native_weight_layout, fsdp_dimension, make_block +): + """The FP8 weight gather retains forward and backward MoE semantics.""" + if native_weight_layout or fsdp_dimension != "k": + from transformer_engine.jax import cpp_extensions as tex + + flax_moe_module = importlib.import_module("transformer_engine.jax.flax.moe") + original_moe = flax_moe_module.moe + + def layout_moe(*args, **kwargs): + args = list(args) + if native_weight_layout: + gate, up = jnp.split(args[2], 2, axis=-1) + args[2] = tex.pack_gated_activation_pair(gate, up).transpose(0, 2, 1) + kwargs["wi_kernel_axes"] = ("exp", "mlp", "embed") + kwargs["cudnn_native_weight_layout"] = True + if fsdp_dimension == "gated": + kwargs["wi_kernel_axes"] = ( + ("exp", "gated_weight_fsdp", None) + if native_weight_layout + else ("exp", None, "gated_weight_fsdp") + ) + spec = ( + P(EP_AXIS, FSDP_AXIS, None) + if native_weight_layout + else P(EP_AXIS, None, FSDP_AXIS) + ) + args[2] = jax.lax.with_sharding_constraint(args[2], NamedSharding(mesh, spec)) + if fsdp_dimension == "expert": + kwargs["wi_kernel_axes"] = ("expert_weight_fsdp", None, None) + kwargs["wo_kernel_axes"] = ("expert_weight_fsdp", None, None) + spec = P((EP_AXIS, FSDP_AXIS), None, None) + args[2] = jax.lax.with_sharding_constraint(args[2], NamedSharding(mesh, spec)) + args[3] = jax.lax.with_sharding_constraint(args[3], NamedSharding(mesh, spec)) + return original_moe(*args, **kwargs) + + monkeypatch.setattr(flax_moe_module, "moe", layout_moe) + x = _make_inputs(jax.random.PRNGKey(41)) + baseline = make_block(quantization_recipe=MXFP8BlockScaling()) + quantized_ag = make_block(quantization_recipe=MXFP8BlockScaling(), quant_before_fsdp_ag=True) + variables, baseline_out, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(42)) + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + quantized_out, _, _ = jax.jit(quantized_ag.apply)(variables, x_sh) + quantized_out.block_until_ready() + np.testing.assert_allclose( + _to_global_numpy(quantized_out, mesh).astype(np.float32), + _to_global_numpy(baseline_out, mesh).astype(np.float32), + **FWD_TOLERANCE["mxfp8"], + ) + baseline_grads, baseline_dx = _grad_step(baseline, variables, mesh, x) + quantized_grads, quantized_dx = _grad_step(quantized_ag, variables, mesh, x) + for name in ("wi", "wo"): + np.testing.assert_allclose( + _to_global_numpy(_unwrap(quantized_grads["params"][name]), mesh).astype(np.float32), + _to_global_numpy(_unwrap(baseline_grads["params"][name]), mesh).astype(np.float32), + **GRAD_FFN_TOLERANCE["mxfp8"], + ) + np.testing.assert_allclose( + _to_global_numpy(quantized_dx, mesh).astype(np.float32), + _to_global_numpy(baseline_dx, mesh).astype(np.float32), + **GRAD_FFN_TOLERANCE["mxfp8"], + ) + + +@pytest.mark.parametrize("quant_before_fsdp_ag", [False, True]) +def test_mesh_resource_api_forward_and_backward(mesh, quant_before_fsdp_ag, make_block): + """Explicit and global resources give identical forward and backward results.""" + baseline = make_block( + quantization_recipe=MXFP8BlockScaling(), quant_before_fsdp_ag=quant_before_fsdp_ag + ) + candidate = baseline.clone( + mesh_resource=MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ) + x = _make_inputs(jax.random.PRNGKey(51)) + variables, baseline_out, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(52)) + baseline_grads, baseline_dx = _grad_step(baseline, variables, mesh, x) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): + with global_shard_guard(None): + x_sh = _shard_inputs(x, mesh) + candidate_out, _, _ = jax.jit(candidate.apply)(variables, x_sh) + + def loss_fn(variables, inputs): + output, _, _ = candidate.apply(variables, inputs) + return jnp.mean(output.astype(jnp.float32) ** 2) + + candidate_grads, candidate_dx = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))( + variables, x_sh + ) + jax.block_until_ready((candidate_out, candidate_grads, candidate_dx)) + np.testing.assert_allclose( + _to_global_numpy(candidate_out, mesh).astype(np.float32), + _to_global_numpy(baseline_out, mesh).astype(np.float32), + **FWD_TOLERANCE["mxfp8"], + ) + for name in ("gate_kernel", "wi", "wo"): + np.testing.assert_allclose( + _to_global_numpy(_unwrap(candidate_grads["params"][name]), mesh).astype(np.float32), + _to_global_numpy(_unwrap(baseline_grads["params"][name]), mesh).astype(np.float32), + **( + GRAD_GATE_TOLERANCE["mxfp8"] + if name == "gate_kernel" + else GRAD_FFN_TOLERANCE["mxfp8"] + ), + ) + np.testing.assert_allclose( + _to_global_numpy(candidate_dx, mesh).astype(np.float32), + _to_global_numpy(baseline_dx, mesh).astype(np.float32), + **GRAD_FFN_TOLERANCE["mxfp8"], + ) + + def _reference_kwargs_from_config(config, params_np): """Pick out the reference-relevant pieces of a parametrize config.""" return dict( @@ -581,8 +725,8 @@ class TestTeEpMoeForward: @pytest.mark.parametrize("config", _CONFIGS) @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) - def test_forward(self, mesh, config, quantization): - block = _make_block(**config, quantization_recipe=_quantization_recipe(quantization)) + def test_forward(self, mesh, config, quantization, make_block): + block = make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(0)) variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) @@ -621,8 +765,8 @@ class TestTeEpMoeBackward: @pytest.mark.parametrize("config", _CONFIGS) @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) - def test_backward(self, mesh, config, quantization): - block = _make_block(**config, quantization_recipe=_quantization_recipe(quantization)) + def test_backward(self, mesh, config, quantization, make_block): + block = make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(2)) variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) grads_te, grad_x_te = _grad_step(block, variables, mesh, x) @@ -692,6 +836,325 @@ def loss_fn(params, x): ) +def test_ep_checkpoint_names(mesh, monkeypatch, make_block): + moe_module = importlib.import_module("transformer_engine.jax.moe") + named_values = {} + original_checkpoint_name = moe_module.checkpoint_name + + def recorded_checkpoint_name(value, name): + named_values.setdefault(name, []).append(value) + return original_checkpoint_name(value, name) + + monkeypatch.setattr(moe_module, "checkpoint_name", recorded_checkpoint_name) + block = make_block( + dispatch_checkpoint_name="saved_dispatch", + ) + inputs = _make_inputs(jax.random.PRNGKey(34)) + variables, output, _ = _init_apply(block, mesh, inputs, jax.random.PRNGKey(35)) + grads, grad_inputs = _grad_step(block, variables, mesh, inputs) + + assert len(named_values["saved_dispatch"]) >= 2 + assert all(hasattr(value, "shape") for values in named_values.values() for value in values) + assert np.all(np.isfinite(_to_global_numpy(output, mesh))) + assert np.all(np.isfinite(_to_global_numpy(grad_inputs, mesh))) + for name in ("gate_kernel", "wi", "wo"): + assert np.all(np.isfinite(_to_global_numpy(_unwrap(grads["params"][name]), mesh))) + + +def test_explicitly_disabled_fusion_with_native_weights( + mesh, monkeypatch, make_block, use_cudnn_fusion +): + """Disabling fusion bypasses dependency probing and preserves native gradients.""" + if use_cudnn_fusion: + pytest.skip("Requires the ordinary 128-token EP bootstrap") + from transformer_engine.jax import cpp_extensions as tex + + moe_module = importlib.import_module("transformer_engine.jax.moe") + + def unexpected_selection(*args, **kwargs): + raise AssertionError("Explicitly disabled fusion must not probe cuDNN") + + monkeypatch.setattr(moe_module, "_select_cudnn_jax_fusion", unexpected_selection) + block = make_block(quantization_recipe=MXFP8BlockScaling(), use_cudnn_fusion=False) + x = _make_inputs(jax.random.PRNGKey(53)) + variables, baseline_output, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(54)) + baseline_grads, baseline_dx = _grad_step(block, variables, mesh, x) + flax_moe = importlib.import_module("transformer_engine.jax.flax.moe") + original_moe = flax_moe.moe + + def native_moe(*args, **kwargs): + args = list(args) + gate, up = jnp.split(args[2], 2, axis=-1) + args[2] = tex.pack_gated_activation_pair(gate, up).transpose(0, 2, 1) + kwargs["wi_kernel_axes"] = ("exp", "mlp", "embed") + kwargs["cudnn_native_weight_layout"] = True + return original_moe(*args, **kwargs) + + monkeypatch.setattr(flax_moe, "moe", native_moe) + with _ctx(mesh): + output, _, _ = jax.jit(block.apply)(variables, _shard_inputs(x, mesh)) + output.block_until_ready() + grads, dx = _grad_step(block, variables, mesh, x) + np.testing.assert_array_equal( + _to_global_numpy(output, mesh), _to_global_numpy(baseline_output, mesh) + ) + np.testing.assert_array_equal(_to_global_numpy(dx, mesh), _to_global_numpy(baseline_dx, mesh)) + for name in ("gate_kernel", "wi", "wo"): + np.testing.assert_array_equal( + _to_global_numpy(_unwrap(grads["params"][name]), mesh), + _to_global_numpy(_unwrap(baseline_grads["params"][name]), mesh), + ) + + +class TestTeEpMoeCudnnCutedslFusion: + """End-to-end MXFP8 coverage for cuDNN's grouped GLU JAX APIs.""" + + @pytest.mark.parametrize("fallback", ["missing_api", "old_version", "hybrid"]) + @pytest.mark.parametrize("native_layout", [False, True]) + def test_unsupported_fusion_unfused_forward_and_backward( + self, mesh, monkeypatch, native_layout, make_block, use_cudnn_fusion, fallback + ): + """Unfused fallback retains bootstrap capacity and native-layout gradients.""" + if not use_cudnn_fusion: + pytest.skip("Requires fusion requested at bootstrap") + import cudnn + from transformer_engine.jax import cpp_extensions as tex + + if fallback == "missing_api": + # Avoid triggering a lazy optional import just to disable the API. + monkeypatch.setitem(cudnn.__dict__, "grouped_gemm_glu_wrapper_sm100", None) + elif fallback == "old_version": + monkeypatch.setattr(cudnn, "__version__", "1.31.0") + recipe = MXFP8BlockScaling( + fp8_format=Format.HYBRID if fallback == "hybrid" else Format.E4M3 + ) + + def unexpected_fused_kernel(*args, **kwargs): + raise AssertionError("Unsupported fusion must use unfused kernels") + + monkeypatch.setattr(tex, "grouped_gemm_glu", unexpected_fused_kernel) + monkeypatch.setattr(tex, "grouped_gemm_dglu", unexpected_fused_kernel) + baseline = make_block(quantization_recipe=recipe) + block = make_block(quantization_recipe=recipe) + x = _make_inputs(jax.random.PRNGKey(51)) + # Keep the fusion-requested EP alignment and bootstrap capacity while + # independently forcing the baseline to execute the unfused FFN. + moe_module = importlib.import_module("transformer_engine.jax.moe") + with monkeypatch.context() as baseline_patch: + baseline_patch.setattr(moe_module, "_select_cudnn_jax_fusion", lambda _: False) + variables, baseline_output, _ = _init_apply(baseline, mesh, x, jax.random.PRNGKey(52)) + baseline_grads, baseline_dx = _grad_step(baseline, variables, mesh, x) + + if native_layout: + flax_moe = importlib.import_module("transformer_engine.jax.flax.moe") + original_moe = flax_moe.moe + + def native_moe(*args, **kwargs): + args = list(args) + gate, up = jnp.split(args[2], 2, axis=-1) + args[2] = tex.pack_gated_activation_pair(gate, up).transpose(0, 2, 1) + kwargs["wi_kernel_axes"] = ("exp", "mlp", "embed") + kwargs["cudnn_native_weight_layout"] = True + return original_moe(*args, **kwargs) + + monkeypatch.setattr(flax_moe, "moe", native_moe) + with _ctx(mesh): + output, _, _ = jax.jit(block.apply)(variables, _shard_inputs(x, mesh)) + output.block_until_ready() + grads, dx = _grad_step(block, variables, mesh, x) + np.testing.assert_array_equal( + _to_global_numpy(output, mesh), _to_global_numpy(baseline_output, mesh) + ) + np.testing.assert_array_equal( + _to_global_numpy(dx, mesh), _to_global_numpy(baseline_dx, mesh) + ) + for name in ("gate_kernel", "wi", "wo"): + np.testing.assert_array_equal( + _to_global_numpy(_unwrap(grads["params"][name]), mesh), + _to_global_numpy(_unwrap(baseline_grads["params"][name]), mesh), + ) + + @pytest.mark.parametrize("native_layout", [False, True]) + @pytest.mark.parametrize("apply_topk_weights_early", [False, True]) + def test_mxfp8_forward_and_backward( + self, + mesh, + apply_topk_weights_early, + monkeypatch, + make_block, + use_cudnn_fusion, + native_layout, + ): + if not use_cudnn_fusion: + pytest.skip("requires the cuDNN fusion fixture parameter") + from transformer_engine.jax import cpp_extensions as tex + + available, reason = tex.grouped_gemm_glu_dependencies_available() + if not available: + pytest.skip(reason) + + if native_layout: + flax_moe = importlib.import_module("transformer_engine.jax.flax.moe") + original_moe = flax_moe.moe + + def native_moe(*args, **kwargs): + args = list(args) + gate, up = jnp.split(args[2], 2, axis=-1) + args[2] = tex.pack_gated_activation_pair(gate, up).transpose(0, 2, 1) + kwargs["wi_kernel_axes"] = ("exp", "mlp", "embed") + kwargs["cudnn_native_weight_layout"] = True + return original_moe(*args, **kwargs) + + monkeypatch.setattr(flax_moe, "moe", native_moe) + + shared_calls = [] + original_glu = tex.grouped_gemm_glu + original_dglu = tex.grouped_gemm_dglu + + def checked_glu(*args, **kwargs): + shared_calls.append("glu") + return original_glu(*args, **kwargs) + + def checked_dglu(*args, **kwargs): + shared_calls.append("dglu") + return original_dglu(*args, **kwargs) + + monkeypatch.setattr(tex, "grouped_gemm_glu", checked_glu) + monkeypatch.setattr(tex, "grouped_gemm_dglu", checked_dglu) + block = make_block( + apply_topk_weights_early=apply_topk_weights_early, + quantization_recipe=MXFP8BlockScaling(), + ) + x = _make_inputs(jax.random.PRNGKey(30)) + variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(31)) + assert "glu" in shared_calls, "MoE did not select the shared cuDNN GLU JAX API" + grads, grad_x = _grad_step(block, variables, mesh, x) + assert "dglu" in shared_calls, "MoE VJP did not select the shared cuDNN dGLU JAX API" + + assert output.shape == x.shape + assert output.dtype == x.dtype + assert aux is None + output_np = _to_global_numpy(output, mesh).astype(np.float32) + assert np.all(np.isfinite(output_np)) + assert np.any(output_np != 0) + + for name in ("gate_kernel", "wi", "wo"): + grad = _to_global_numpy(_unwrap(grads["params"][name]), mesh).astype(np.float32) + assert np.all(np.isfinite(grad)), f"{name} fused MXFP8 grad has NaN/Inf" + assert np.any(grad != 0), f"{name} fused MXFP8 grad is identically zero" + grad_x_np = _to_global_numpy(grad_x, mesh).astype(np.float32) + assert np.all(np.isfinite(grad_x_np)) + assert np.any(grad_x_np != 0) + + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + + def loss_fn(params, inputs): + out, _ = _pure_jax_moe_reference( + inputs, + params["gate_kernel"], + params["wi"], + params["wo"], + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + score_function="softmax", + expert_bias=None, + ) + return jnp.mean(out.astype(jnp.float32) ** 2) + + grads_ref, grad_x_ref = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))( + {k: jnp.asarray(v) for k, v in params_np.items() if k != "expert_bias"}, + jnp.asarray(x_np), + ) + for name in ("gate_kernel", "wi", "wo"): + np.testing.assert_allclose( + _to_global_numpy(_unwrap(grads["params"][name]), mesh).astype(np.float32), + np.asarray(jax.device_get(grads_ref[name])).astype(np.float32), + **( + GRAD_GATE_TOLERANCE["mxfp8"] + if name == "gate_kernel" + else GRAD_FFN_TOLERANCE["mxfp8"] + ), + err_msg=f"{name} fused MXFP8 gradient parity breach", + ) + grad_x_ref_np = np.asarray(jax.device_get(grad_x_ref)).astype(np.float32) + np.testing.assert_allclose( + grad_x_np, + grad_x_ref_np, + **GRAD_FFN_TOLERANCE["mxfp8"], + err_msg="d_x fused MXFP8 gradient parity breach", + ) + # The mean-square loss makes dX smaller than the absolute tolerance. + # Also bound relative error to catch a wrong MXFP8 quantization axis. + relative_dx_error = np.linalg.norm(grad_x_np - grad_x_ref_np) / np.linalg.norm( + grad_x_ref_np + ) + assert ( + relative_dx_error < 0.1 + ), f"d_x MXFP8 relative L2 error {relative_dx_error:.3%} >= 10%" + + def test_cudnn_fused_with_checkpoint_names( + self, mesh, monkeypatch, make_block, use_cudnn_fusion + ): + if not use_cudnn_fusion: + pytest.skip("cuDNN grouped GEMM fusion is disabled") + from transformer_engine.jax import cpp_extensions as tex + + available, reason = tex.grouped_gemm_glu_dependencies_available() + if not available: + pytest.skip(reason) + + moe_module = importlib.import_module("transformer_engine.jax.moe") + flax_moe_module = importlib.import_module("transformer_engine.jax.flax.moe") + selected_calls = [] + fused_op_name = "grouped_gemm_glu" + original_fused_op = getattr(tex, fused_op_name) + + def checked_fused_op(*args, **kwargs): + selected_calls.append(True) + return original_fused_op(*args, **kwargs) + + monkeypatch.setattr(tex, fused_op_name, checked_fused_op) + named_values = {} + original_checkpoint_name = moe_module.checkpoint_name + + def recorded_checkpoint_name(value, name): + named_values.setdefault(name, []).append(value) + return original_checkpoint_name(value, name) + + monkeypatch.setattr(moe_module, "checkpoint_name", recorded_checkpoint_name) + original_moe = flax_moe_module.moe + + def checkpointed_moe(*args, **kwargs): + kwargs.update( + wi_0_checkpoint_name="moe_mlpwi_0", + wi_1_checkpoint_name="moe_mlpwi_1", + wo_checkpoint_name="moe_mlpwo", + ) + return original_moe(*args, **kwargs) + + monkeypatch.setattr(flax_moe_module, "moe", checkpointed_moe) + block = make_block( + quantization_recipe=MXFP8BlockScaling(), + dispatch_checkpoint_name="saved_dispatch", + ) + inputs = _make_inputs(jax.random.PRNGKey(32)) + variables, output, _ = _init_apply(block, mesh, inputs, jax.random.PRNGKey(33)) + grads, grad_inputs = _grad_step(block, variables, mesh, inputs) + + assert selected_calls, f"{fused_op_name} was not selected" + for label in ("moe_mlpwi_0", "moe_mlpwi_1"): + assert len(named_values[label]) == 5 * len(selected_calls) + assert all(hasattr(value, "shape") for value in named_values[label]) + assert len(named_values["moe_mlpwo"]) == len(selected_calls) + assert all(hasattr(value, "shape") for value in named_values["moe_mlpwo"]) + assert len(named_values["saved_dispatch"]) >= 2 * len(selected_calls) + assert np.all(np.isfinite(_to_global_numpy(output, mesh))) + assert np.all(np.isfinite(_to_global_numpy(grad_inputs, mesh))) + for name in ("gate_kernel", "wi", "wo"): + assert np.all(np.isfinite(_to_global_numpy(_unwrap(grads["params"][name]), mesh))) + + class TestTeEpMoeAuxLoss: """Aux-loss path. Consolidated into: * ``test_aux_loss``: one run that checks the returned scalar's @@ -701,9 +1164,9 @@ class TestTeEpMoeAuxLoss: finite + non-zero per tensor. """ - def test_aux_loss(self, mesh): + def test_aux_loss(self, mesh, make_block): coeff = 1e-2 - block = _make_block(aux_loss_coeff=coeff) + block = make_block(aux_loss_coeff=coeff) x = _make_inputs(jax.random.PRNGKey(20)) variables, _, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(21)) @@ -743,10 +1206,10 @@ def test_aux_loss(self, mesh): assert np.all(np.isfinite(g_gate)), "gate grad NaN/Inf under aux-only loss" assert np.any(g_gate != 0.0), "aux bwd should propagate to gate_kernel" - def test_combined_loss_grads(self, mesh): + def test_combined_loss_grads(self, mesh, make_block): """Joint main + aux loss bwd: per-tensor finite + non-zero in one pass.""" - block = _make_block(aux_loss_coeff=1e-2) + block = make_block(aux_loss_coeff=1e-2) x = _make_inputs(jax.random.PRNGKey(22)) variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(23)) grads, _ = _grad_step(block, variables, mesh, x, include_aux=True) diff --git a/transformer_engine/jax/cpp_extensions/__init__.py b/transformer_engine/jax/cpp_extensions/__init__.py index c9647afb826..c9637838577 100644 --- a/transformer_engine/jax/cpp_extensions/__init__.py +++ b/transformer_engine/jax/cpp_extensions/__init__.py @@ -10,6 +10,7 @@ from .quantization import * from .softmax import * from .gemm import * +from .grouped_gemm_glu import * from .router import * from .ep import * from .topk import * diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 456b4b5124f..6f939ab8f7f 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -201,15 +201,22 @@ def ep_handle_mem_size(cfg: EpLayerConfig) -> int: ) -def _leading_axis_ok(spec): - """Validate an EP input spec; return ``(ok, ep_axis, outer_axes)``. +def _capture_ep_resource_axes(): + """Capture physical axes at bind time for asynchronous SPMD partitioning.""" + resource = global_mesh_resource() + outer_axes = tuple( + dict.fromkeys( + axis for axis in (resource.dp_resource, resource.fsdp_resource) if axis is not None + ) + ) + return resource.ep_resource, 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) + +def _leading_axis_ok(spec, resource_axes=None): + """Validate replicated trailing dimensions and EP plus optional outer axes.""" + ep_axis, outer_axes = ( + resource_axes if resource_axes is not None else _capture_ep_resource_axes() + ) 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:]): @@ -223,20 +230,13 @@ def _leading_axis_ok(spec): 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 - sees each DP color's slab as distinct (rather than replicated across DP). - - 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. - """ - 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 + """Legacy single outer axis used by the standalone EP custom VJP wrapper.""" + resource = global_mesh_resource() + if resource.dp_resource is not None and get_mesh_axis_size(resource.dp_resource) > 1: + return resource.dp_resource + if resource.fsdp_resource is not None and get_mesh_axis_size(resource.fsdp_resource) > 1: + return resource.fsdp_resource + return resource.dp_resource or resource.fsdp_resource def _ep_leading_dims(is_outer): @@ -248,24 +248,20 @@ 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_output_spec(*trailing, resource_axes=None): + """Output sharding uses every DP/FSDP outer axis, with EP innermost.""" + ep_axis, outer_axes = ( + resource_axes if resource_axes is not None else _capture_ep_resource_axes() + ) + leading = (*outer_axes, ep_axis) if outer_axes else ep_axis + 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. - """ - gsr = global_mesh_resource() - ep_axis = gsr.ep_resource - outer = _ep_outer_axis() +def _ep_spec_ok(spec, trailing_count, resource_axes=None): + """Validate EP-output axes, allowing JAX to drop size-one mesh axes.""" + ep_axis, outer_axes = ( + resource_axes if resource_axes is not None else _capture_ep_resource_axes() + ) if len(spec) != 1 + trailing_count: return False if any(ax is not None for ax in spec[1:]): @@ -273,8 +269,7 @@ 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 + return actual <= set((*outer_axes, ep_axis)) # ── ep_prepare ────────────────────────────────────────────────────────────── @@ -285,8 +280,8 @@ class EpPreparePrimitive(BasePrimitive): name = "te_ep_prepare_ffi" multiple_results = True - impl_static_args = (1, 2, 3, 4) # top_k, dispatch_output_per_expert_alignment, - # topk_idx_dtype, is_outer + impl_static_args = (1, 2, 3, 4, 5) # top_k, dispatch_output_per_expert_alignment, + # topk_idx_dtype, is_outer, resource_axes inner_primitive = None outer_primitive = None @@ -298,10 +293,11 @@ def abstract( dispatch_output_per_expert_alignment, topk_idx_dtype, is_outer, + resource_axes, ): # is_outer=True: global leading dim = (dp*ep,) (or (ep,) with no DP); # False: per-shard = (1,). - del topk_idx_dtype + del topk_idx_dtype, resource_axes cfg = get_ep_config() num_local_experts = cfg.num_local_experts assert ( @@ -326,8 +322,16 @@ def outer_abstract(*args, **kwargs): @staticmethod def lowering( - ctx, topk_idx, *, top_k, dispatch_output_per_expert_alignment, topk_idx_dtype, is_outer + ctx, + topk_idx, + *, + top_k, + dispatch_output_per_expert_alignment, + topk_idx_dtype, + is_outer, + resource_axes, ): + del resource_axes del is_outer return ffi.ffi_lowering(EpPreparePrimitive.name)( ctx, @@ -338,7 +342,14 @@ def lowering( ) @staticmethod - def impl(topk_idx, top_k, dispatch_output_per_expert_alignment, topk_idx_dtype, is_outer): + def impl( + topk_idx, + top_k, + dispatch_output_per_expert_alignment, + topk_idx_dtype, + is_outer, + resource_axes, + ): assert EpPreparePrimitive.inner_primitive is not None token_counts, total_recv_tokens, handle_mem = EpPreparePrimitive.inner_primitive.bind( topk_idx, @@ -346,6 +357,7 @@ def impl(topk_idx, top_k, dispatch_output_per_expert_alignment, topk_idx_dtype, dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, topk_idx_dtype=topk_idx_dtype, is_outer=is_outer, + resource_axes=resource_axes, ) return token_counts, total_recv_tokens, handle_mem @@ -358,6 +370,7 @@ def batcher( dispatch_output_per_expert_alignment, topk_idx_dtype, is_outer, + resource_axes, ): raise NotImplementedError("EpPreparePrimitive does not support vmap") @@ -367,13 +380,14 @@ def partition( dispatch_output_per_expert_alignment, topk_idx_dtype, is_outer, + resource_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, resource_axes) if not ok: raise NotImplementedError( "EpPrepare: topk_idx leading dim must include ep_resource" @@ -390,7 +404,12 @@ def partition( def sharded_impl(topk_idx): return EpPreparePrimitive.impl( - topk_idx, top_k, dispatch_output_per_expert_alignment, topk_idx_dtype, False + topk_idx, + top_k, + dispatch_output_per_expert_alignment, + topk_idx_dtype, + False, + resource_axes=resource_axes, ) return mesh, sharded_impl, (tc_sharding, trt_sharding, hm_sharding), arg_shardings @@ -416,8 +435,8 @@ class EpDispatchPrimitive(BasePrimitive): name = "te_ep_dispatch_ffi" multiple_results = True - impl_static_args = (4, 5, 6, 7, 8) # top_k, dispatch_output_per_expert_alignment, - # topk_idx_dtype, recv_capacity_per_rank, is_outer + impl_static_args = (4, 5, 6, 7, 8, 9) # top_k, dispatch_output_per_expert_alignment, + # topk_idx_dtype, recv_capacity_per_rank, is_outer, resource_axes inner_primitive = None outer_primitive = None @@ -433,9 +452,11 @@ def abstract( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): # is_outer=True: global leading dim = (dp*ep,) (or (ep,) with no DP); # False: per-shard = (1,). + del resource_axes del topk_idx_aval, topk_weights_aval, top_k, dispatch_output_per_expert_alignment del topk_idx_dtype assert ( @@ -467,7 +488,9 @@ def lowering( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): + del resource_axes del recv_capacity_per_rank, is_outer return ffi.ffi_lowering( EpDispatchPrimitive.name, @@ -494,6 +517,7 @@ def impl( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): assert EpDispatchPrimitive.inner_primitive is not None recv_tokens, recv_topk_weights, handle_mem_out = EpDispatchPrimitive.inner_primitive.bind( @@ -506,6 +530,7 @@ def impl( topk_idx_dtype=topk_idx_dtype, recv_capacity_per_rank=recv_capacity_per_rank, is_outer=is_outer, + resource_axes=resource_axes, ) return recv_tokens, recv_topk_weights, handle_mem_out @@ -519,6 +544,7 @@ def batcher( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): raise NotImplementedError("EpDispatchPrimitive does not support vmap") @@ -529,13 +555,14 @@ def partition( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_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, resource_axes) if not ok: raise NotImplementedError( "EpDispatch: tokens leading dim must include ep_resource" @@ -569,6 +596,7 @@ def sharded_impl(handle_mem, topk_idx, tokens, topk_weights): topk_idx_dtype, recv_capacity_per_rank, False, + resource_axes=resource_axes, ) return mesh, sharded_impl, out_shardings, arg_shardings @@ -600,8 +628,9 @@ class EpPrepareAndDispatchPrimitive(BasePrimitive): name = "te_ep_prepare_and_dispatch_ffi" multiple_results = True - impl_static_args = (3, 4, 5, 6, 7) # top_k, dispatch_output_per_expert_alignment, - # topk_idx_dtype, recv_capacity_per_rank, is_outer + impl_static_args = (3, 4, 5, 6, 7, 8) # top_k, dispatch_output_per_expert_alignment, + # topk_idx_dtype, recv_capacity_per_rank, is_outer, + # resource_axes inner_primitive = None outer_primitive = None @@ -616,7 +645,9 @@ def abstract( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): + del resource_axes del topk_idx_aval, topk_weights_aval, topk_idx_dtype assert ( len(tokens_aval.shape) >= 2 @@ -665,7 +696,9 @@ def lowering( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): + del resource_axes del recv_capacity_per_rank, is_outer return ffi.ffi_lowering(EpPrepareAndDispatchPrimitive.name)( ctx, @@ -687,6 +720,7 @@ def impl( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): assert EpPrepareAndDispatchPrimitive.inner_primitive is not None return tuple( @@ -699,6 +733,7 @@ def impl( topk_idx_dtype=topk_idx_dtype, recv_capacity_per_rank=recv_capacity_per_rank, is_outer=is_outer, + resource_axes=resource_axes, ) ) @@ -712,6 +747,7 @@ def batcher( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): raise NotImplementedError("EpPrepareAndDispatchPrimitive does not support vmap") @@ -722,13 +758,14 @@ def partition( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, mesh, arg_infos, result_infos, ): del is_outer, result_infos tokens_spec = arg_infos[1].sharding.spec - ok, ep_axis, outer_axes = _leading_axis_ok(tokens_spec) + ok, ep_axis, outer_axes = _leading_axis_ok(tokens_spec, resource_axes) if not ok: raise NotImplementedError( "EpPrepareAndDispatch: tokens leading dim must include ep_resource" @@ -756,6 +793,7 @@ def sharded_impl(topk_idx, tokens, topk_weights): topk_idx_dtype, recv_capacity_per_rank, False, + resource_axes=resource_axes, ) return mesh, sharded_impl, out_shardings, arg_shardings @@ -763,7 +801,7 @@ def sharded_impl(topk_idx, tokens, topk_weights): @staticmethod def shardy_sharding_rule(*args): # Static args: (top_k, dispatch_alignment, topk_idx_dtype, recv_capacity_per_rank, - # is_outer). + # is_outer, resource_axes). value_types = args[-2] idx_rank = len(value_types[0].shape) tok_rank = len(value_types[1].shape) @@ -813,8 +851,8 @@ class EpCombinePrimitive(BasePrimitive): name = "te_ep_combine_ffi" multiple_results = False - impl_static_args = (2, 3, 4, 5, 6) # top_k, dispatch_output_per_expert_alignment, - # topk_idx_dtype, out_leading_shape, out_partition_spec + impl_static_args = (2, 3, 4, 5, 6, 7) # top_k, dispatch_output_per_expert_alignment, + # topk_idx_dtype, out_leading_shape, out_partition_spec, resource_axes inner_primitive = None outer_primitive = None @@ -828,7 +866,9 @@ def abstract( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): + del resource_axes del top_k, dispatch_output_per_expert_alignment, topk_idx_dtype del out_partition_spec, handle_mem_aval assert ( @@ -850,7 +890,9 @@ def lowering( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): + del resource_axes del out_leading_shape, out_partition_spec return ffi.ffi_lowering(EpCombinePrimitive.name)( ctx, @@ -870,6 +912,7 @@ def impl( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): assert EpCombinePrimitive.inner_primitive is not None return EpCombinePrimitive.inner_primitive.bind( @@ -880,6 +923,7 @@ def impl( topk_idx_dtype=topk_idx_dtype, out_leading_shape=out_leading_shape, out_partition_spec=out_partition_spec, + resource_axes=resource_axes, ) @staticmethod @@ -892,6 +936,7 @@ def batcher( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): raise NotImplementedError("EpCombinePrimitive does not support vmap") @@ -902,13 +947,14 @@ def partition( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_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, resource_axes=resource_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)" @@ -931,6 +977,7 @@ def sharded_impl(handle_mem, expert_out): topk_idx_dtype, per_shard_leading, out_partition_spec, + resource_axes=resource_axes, ) return mesh, sharded_impl, out_sharding, arg_shardings @@ -956,8 +1003,8 @@ class EpDispatchBwdPrimitive(BasePrimitive): name = "te_ep_dispatch_bwd_ffi" multiple_results = True - impl_static_args = (3, 4, 5, 6, 7) # top_k, dispatch_output_per_expert_alignment, - # topk_idx_dtype, out_leading_shape, out_partition_spec + impl_static_args = (3, 4, 5, 6, 7, 8) # top_k, dispatch_output_per_expert_alignment, + # topk_idx_dtype, out_leading_shape, out_partition_spec, resource_axes inner_primitive = None outer_primitive = None @@ -972,7 +1019,9 @@ def abstract( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): + del resource_axes del dispatch_output_per_expert_alignment, topk_idx_dtype del g_recv_topk_weights_aval, out_partition_spec, handle_mem_aval assert ( @@ -998,7 +1047,9 @@ def lowering( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): + del resource_axes del out_leading_shape, out_partition_spec return ffi.ffi_lowering(EpDispatchBwdPrimitive.name)( ctx, @@ -1020,6 +1071,7 @@ def impl( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): assert EpDispatchBwdPrimitive.inner_primitive is not None return EpDispatchBwdPrimitive.inner_primitive.bind( @@ -1031,6 +1083,7 @@ def impl( topk_idx_dtype=topk_idx_dtype, out_leading_shape=out_leading_shape, out_partition_spec=out_partition_spec, + resource_axes=resource_axes, ) @staticmethod @@ -1043,6 +1096,7 @@ def batcher( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_axes, ): raise NotImplementedError("EpDispatchBwdPrimitive does not support vmap") @@ -1053,20 +1107,21 @@ def partition( topk_idx_dtype, out_leading_shape, out_partition_spec, + resource_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, resource_axes=resource_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, resource_axes=resource_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)" @@ -1096,6 +1151,7 @@ def sharded_impl(handle_mem, grad, g_recv_topk_weights): topk_idx_dtype, per_shard_leading, out_partition_spec, + resource_axes=resource_axes, ) return mesh, sharded_impl, out_shardings, arg_shardings @@ -1121,8 +1177,8 @@ class EpCombineBwdPrimitive(BasePrimitive): name = "te_ep_combine_bwd_ffi" multiple_results = False - impl_static_args = (2, 3, 4, 5, 6) # top_k, dispatch_output_per_expert_alignment, - # topk_idx_dtype, recv_capacity_per_rank, is_outer + impl_static_args = (2, 3, 4, 5, 6, 7) # top_k, dispatch_output_per_expert_alignment, + # topk_idx_dtype, recv_capacity_per_rank, is_outer, resource_axes inner_primitive = None outer_primitive = None @@ -1136,9 +1192,11 @@ def abstract( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): # is_outer=True: global leading dim = (dp*ep,) (or (ep,) with no DP); # False: per-shard = (1,). + del resource_axes del top_k, dispatch_output_per_expert_alignment, topk_idx_dtype, handle_mem_aval assert ( len(grad_aval.shape) >= 2 @@ -1164,7 +1222,9 @@ def lowering( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): + del resource_axes del recv_capacity_per_rank, is_outer return ffi.ffi_lowering(EpCombineBwdPrimitive.name)( ctx, @@ -1184,6 +1244,7 @@ def impl( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): assert EpCombineBwdPrimitive.inner_primitive is not None return EpCombineBwdPrimitive.inner_primitive.bind( @@ -1194,6 +1255,7 @@ def impl( topk_idx_dtype=topk_idx_dtype, recv_capacity_per_rank=recv_capacity_per_rank, is_outer=is_outer, + resource_axes=resource_axes, ) @staticmethod @@ -1206,6 +1268,7 @@ def batcher( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, ): raise NotImplementedError("EpCombineBwdPrimitive does not support vmap") @@ -1216,6 +1279,7 @@ def partition( topk_idx_dtype, recv_capacity_per_rank, is_outer, + resource_axes, mesh, arg_infos, result_infos, @@ -1223,7 +1287,7 @@ def partition( 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()) + out_sharding = NamedSharding(mesh, _ep_output_spec(resource_axes=resource_axes)) def sharded_impl(handle_mem, grad): return EpCombineBwdPrimitive.impl( @@ -1234,6 +1298,7 @@ def sharded_impl(handle_mem, grad): topk_idx_dtype, recv_capacity_per_rank, False, + resource_axes=resource_axes, ) return mesh, sharded_impl, out_sharding, arg_shardings @@ -1268,6 +1333,7 @@ def _ep_prepare(cfg: EpLayerConfig, topk_idx): dispatch_output_per_expert_alignment=int(cfg.dispatch_output_per_expert_alignment), topk_idx_dtype=int(cfg.topk_idx_dtype), is_outer=True, + resource_axes=_capture_ep_resource_axes(), ) @@ -1287,6 +1353,7 @@ def _ep_dispatch_fwd( topk_idx_dtype=int(cfg.topk_idx_dtype), recv_capacity_per_rank=recv_capacity_per_rank, is_outer=True, + resource_axes=_capture_ep_resource_axes(), ) @@ -1344,6 +1411,7 @@ def _ep_fused_prepare_and_dispatch_fwd( topk_idx_dtype=int(cfg.topk_idx_dtype), recv_capacity_per_rank=recv_capacity_per_rank, is_outer=True, + resource_axes=_capture_ep_resource_axes(), ) @@ -1380,6 +1448,7 @@ def ep_combine_fwd( topk_idx_dtype=int(cfg.topk_idx_dtype), out_leading_shape=out_leading, out_partition_spec=out_partition_spec, + resource_axes=_capture_ep_resource_axes(), ) @@ -1403,6 +1472,7 @@ def ep_dispatch_bwd( topk_idx_dtype=int(cfg.topk_idx_dtype), out_leading_shape=out_leading, out_partition_spec=out_partition_spec, + resource_axes=_capture_ep_resource_axes(), ) @@ -1417,4 +1487,5 @@ def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank) topk_idx_dtype=int(cfg.topk_idx_dtype), recv_capacity_per_rank=recv_capacity_per_rank, is_outer=True, + resource_axes=_capture_ep_resource_axes(), ) diff --git a/transformer_engine/jax/cpp_extensions/grouped_gemm_glu.py b/transformer_engine/jax/cpp_extensions/grouped_gemm_glu.py new file mode 100644 index 00000000000..86c497218ce --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/grouped_gemm_glu.py @@ -0,0 +1,272 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""cuDNN Frontend JAX API adapters for shared grouped GEMM + GLU/dGLU.""" + +from __future__ import annotations + +import importlib +import inspect + +import jax +import jax.numpy as jnp +from packaging.version import Version + +__all__ = [ + "grouped_gemm_dglu", + "grouped_gemm_glu", + "grouped_gemm_glu_dependencies_available", + "pack_gated_activation_pair", + "unpack_gated_activation_pair", +] + + +def pack_gated_activation_pair(gate: jax.Array, up: jax.Array) -> jax.Array: + """Interleave 32-column gate/up blocks as required by the cuDNN kernel.""" + if gate.shape != up.shape: + raise ValueError(f"gate shape {gate.shape} must match up shape {up.shape}") + if gate.shape[-1] % 32: + raise ValueError( + f"Gated activation intermediate dimension {gate.shape[-1]} must be divisible by 32" + ) + blocks = gate.shape[-1] // 32 + return jnp.stack( + ( + gate.reshape(*gate.shape[:-1], blocks, 32), + up.reshape(*up.shape[:-1], blocks, 32), + ), + axis=-2, + ).reshape(*gate.shape[:-1], 2 * gate.shape[-1]) + + +def unpack_gated_activation_pair(interleaved: jax.Array) -> tuple[jax.Array, jax.Array]: + """Undo :func:`pack_gated_activation_pair`.""" + if interleaved.shape[-1] % 64: + raise ValueError( + f"Interleaved gated activation dimension {interleaved.shape[-1]} must be divisible" + " by 64" + ) + intermediate = interleaved.shape[-1] // 2 + blocks = intermediate // 32 + paired = interleaved.reshape(*interleaved.shape[:-1], blocks, 2, 32) + return ( + paired[..., 0, :].reshape(*interleaved.shape[:-1], intermediate), + paired[..., 1, :].reshape(*interleaved.shape[:-1], intermediate), + ) + + +def _ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def _sf_atom_shape(groups: int, rows: int, cols: int) -> tuple[int, ...]: + """Return the physical row-major scale-factor shape used by cudnn.jax.call.""" + return (groups, _ceil_div(rows, 128), _ceil_div(_ceil_div(cols, 32), 4), 32, 4, 4) + + +def _compact_sf(scale: jax.Array, shape: tuple[int, ...], name: str) -> jax.Array: + """Trim TE's conservative grouped-scale allocation to cudnn's dense atom view.""" + size = 1 + for extent in shape: + size *= extent + if scale.size < size: + raise ValueError(f"{name} has {scale.size} elements, but cuDNN requires {size}") + return scale.reshape(-1)[:size].reshape(shape) + + +# These contracts use cuDNN Frontend's shared GLU API with JAX dispatch. +# Require both the release containing MXFP8 JAX support and compatible APIs: +# newer optional parameters are compatible, but TE must supply every required one. +_FORWARD_ARGS = ( + "a_tensor", + "b_tensor", + "sfa_tensor", + "sfb_tensor", + "padded_offsets", + "alpha_tensor", + "prob_tensor", + "norm_const_tensor", + "c_dtype", + "d_dtype", + "discrete_col_sfd", +) +_BACKWARD_ARGS = tuple(arg for arg in _FORWARD_ARGS if arg != "c_dtype") + ( + "c_tensor", + "beta_tensor", + "dprob_tensor", + "sf_vec_size", + "act_func", +) +_FORWARD_ARGS += ("sf_vec_size", "act_func", "generate_c") +# The shared wrappers predate MXFP8 JAX dispatch and have compatible signatures +# even in frontends that only support eager execution of block-scaled inputs. +# The JAX entry points also existed for BF16 discrete weights, so require their +# dense-weight and scale arguments rather than checking only their names. +_JAX_FORWARD_ARGS = tuple(arg for arg in _FORWARD_ARGS if arg != "sf_vec_size") +_JAX_BACKWARD_ARGS = tuple( + arg for arg in _BACKWARD_ARGS if arg not in ("sf_vec_size", "dprob_tensor") +) + + +def grouped_gemm_glu_dependencies_available() -> tuple[bool, str]: + """Require Frontend >= 1.32.0 and shared JAX APIs supporting dense MXFP8.""" + try: + cutlass_jax = importlib.import_module("cutlass.jax") + if not cutlass_jax.is_available(): + return False, "CuTeDSL JAX support is unavailable" + cudnn = importlib.import_module("cudnn") + if Version(cudnn.__version__) < Version("1.32.0"): + return False, f"cuDNN Frontend >= 1.32.0 is required, got {cudnn.__version__}" + for module, name, arguments in ( + (cudnn, "grouped_gemm_glu_wrapper_sm100", _FORWARD_ARGS), + (cudnn, "grouped_gemm_dglu_wrapper_sm100", _BACKWARD_ARGS), + (cudnn, "grouped_gemm_glu_jax_sm100", _JAX_FORWARD_ARGS), + (cudnn, "grouped_gemm_dglu_jax_sm100", _JAX_BACKWARD_ARGS), + ): + api = getattr(module, name, None) + if not callable(api): + return False, f"{module.__name__}.{name} is unavailable" + signature = inspect.signature(api) + missing = set(arguments) - signature.parameters.keys() + if missing: + return ( + False, + f"{module.__name__}.{name} is missing parameters: {', '.join(sorted(missing))}", + ) + # Binding detects missing/renamed keywords, positional-only arguments, + # and new mandatory arguments, while allowing new optional arguments. + signature.bind(**dict.fromkeys(arguments)) + except (ImportError, RuntimeError, AttributeError, TypeError, ValueError) as exc: + return False, str(exc) + return True, "" + + +def grouped_gemm_glu( + a: jax.Array, + b: jax.Array, + sfa: jax.Array, + sfb: jax.Array, + padded_offsets: jax.Array, + prob: jax.Array, + *, + compute_dtype, + output_dtype, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Run cuDNN's shared grouped MXFP8 GLU API in SwiGLU mode.""" + cudnn_grouped_gemm_glu = getattr( + importlib.import_module("cudnn"), "grouped_gemm_glu_wrapper_sm100" + ) + + return _grouped_gemm_forward( + cudnn_grouped_gemm_glu, + a, + b, + sfa, + sfb, + padded_offsets, + prob, + compute_dtype=compute_dtype, + output_dtype=output_dtype, + ) + + +def _grouped_gemm_forward( + cudnn_forward, + a, + b, + sfa, + sfb, + padded_offsets, + prob, + *, + compute_dtype, + output_dtype, +): + rows, hidden, _ = a.shape + experts, combined, b_hidden = b.shape + if hidden != b_hidden: + raise ValueError(f"A K={hidden} does not match B K={b_hidden}") + alpha = jnp.ones((experts,), dtype=jnp.float32) + norm_const = jnp.ones((1,), dtype=jnp.float32) + result = cudnn_forward( + a_tensor=a.reshape(rows, hidden), + b_tensor=b, + sfa_tensor=_compact_sf(sfa, _sf_atom_shape(1, rows, hidden), "sfa"), + sfb_tensor=_compact_sf(sfb, _sf_atom_shape(experts, combined, hidden), "sfb"), + padded_offsets=padded_offsets.astype(jnp.int32), + alpha_tensor=alpha, + prob_tensor=prob.reshape(rows).astype(jnp.float32), + norm_const_tensor=norm_const, + c_dtype=jnp.dtype(compute_dtype), + d_dtype=jnp.dtype(output_dtype), + discrete_col_sfd=True, + sf_vec_size=32, + act_func="swiglu", + generate_c=True, + ) + return ( + result["c_tensor"].reshape(rows, combined), + result["d_tensor"].reshape(rows, combined // 2), + result["d_col_tensor"].reshape(rows, combined // 2), + result["sfd_row_tensor"].reshape(-1), + result["sfd_col_tensor"].reshape(-1), + ) + + +def grouped_gemm_dglu( + a: jax.Array, + b: jax.Array, + c: jax.Array, + sfa: jax.Array, + sfb: jax.Array, + padded_offsets: jax.Array, + prob: jax.Array, + *, + output_dtype, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Run cuDNN's shared grouped MXFP8 dGLU API in dSwiGLU mode.""" + if a.ndim != 2: + raise ValueError(f"Expected A[M,K], got {a.shape}") + if b.ndim != 3: + raise ValueError(f"Expected physical B[E,N,K], got {b.shape}") + if c.ndim != 2: + raise ValueError(f"Expected C[M,2N], got {c.shape}") + + cudnn_grouped_gemm_dglu = getattr( + importlib.import_module("cudnn"), "grouped_gemm_dglu_wrapper_sm100" + ) + + rows, hidden = a.shape + experts, intermediate, b_hidden = b.shape + if hidden != b_hidden: + raise ValueError(f"A K={hidden} does not match B K={b_hidden}") + if c.shape != (rows, 2 * intermediate): + raise ValueError(f"Expected C shape {(rows, 2 * intermediate)}, got {c.shape}") + + alpha = jnp.ones((experts,), dtype=jnp.float32) + beta = jnp.ones((experts,), dtype=jnp.float32) + norm_const = jnp.ones((1,), dtype=jnp.float32) + result = cudnn_grouped_gemm_dglu( + a_tensor=a, + b_tensor=b, + c_tensor=c, + sfa_tensor=_compact_sf(sfa, _sf_atom_shape(1, rows, hidden), "sfa"), + sfb_tensor=_compact_sf(sfb, _sf_atom_shape(experts, intermediate, hidden), "sfb"), + padded_offsets=padded_offsets.astype(jnp.int32), + alpha_tensor=alpha, + beta_tensor=beta, + prob_tensor=prob.reshape(rows).astype(jnp.float32), + norm_const_tensor=norm_const, + d_dtype=jnp.dtype(output_dtype), + discrete_col_sfd=True, + dprob_tensor=None, + sf_vec_size=32, + act_func="dswiglu", + ) + return ( + result["d_row_tensor"], + result["d_col_tensor"], + result["dprob_tensor"], + result["sfd_row_tensor"].reshape(-1), + result["sfd_col_tensor"].reshape(-1), + ) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 58573966edf..6c4d5a2c8cf 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -34,10 +34,10 @@ from flax import linen as nn from transformer_engine.common.recipe import Recipe -from ..moe import moe +from ..moe import _moe_mesh_axes, _resolve_moe_mesh_resource, moe from ..quantize import QuantizerSet from ..router import ScoreFunction -from ..sharding import _get_mesh, get_active_resource_axis +from ..sharding import MeshResource, _get_mesh, global_shard_guard from .module import TransformerEngineBase PRNGKey = Any @@ -92,12 +92,19 @@ class _MoEBlock(TransformerEngineBase): Logical sharding axis tuples (consumed by Flax's :func:`with_logical_partitioning` and our internal :func:`with_sharding_constraint_by_logical_axes`). - data_parallelism_axes : tuple[str, ...] - FSDP axes over which the input *batch* dim is sharded IN - 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. + mesh_resource : Optional[MeshResource] + Physical DP, FSDP and EP mesh axes. ``None`` resolves the active global + MeshResource context. An explicit resource takes precedence; one of + these is required. Batch sharding uses DP and FSDP as outer axes and + EP innermost. + quant_before_fsdp_ag : bool + Quantize MXFP8 weight shards before their FSDP all-gather. Defaults to + ``False``; ``True`` requires ``mesh_resource.fsdp_resource``. + use_cudnn_fusion : bool + Defaults to ``True``: try shared cuDNN GLU/dGLU fusion on SM100+ GPUs. + Ineligible calls warn and use unfused TE grouped GEMM. + ``False`` selects unfused execution. Use the same value when calculating + EP bootstrap receive capacity with ``get_moe_recv_capacity_per_rank``. apply_topk_weights_early : bool If ``True``, multiply expert outputs by their top-k weights *inside* each shard before ``ep_combine`` (saves one global @@ -105,11 +112,12 @@ class _MoEBlock(TransformerEngineBase): recv_capacity_per_rank : Optional[int] Exact aligned receive capacity per EP rank. ``None`` reserves the dropless worst case. + dispatch_checkpoint_name : Optional[str] + JAX rematerialization checkpoint name for the EP dispatch outputs. + ``None`` leaves them unnamed. - The per-expert dispatch-slot alignment is fixed internally at 128 - tokens (see ``moe._ALIGN_SIZE``) -- the value required by NCCL EP - HT and satisfied by every current TE grouped-GEMM recipe -- and is - therefore not exposed as a per-instance knob. + Per-expert dispatch-slot alignment is 256 tokens when fusion is requested, + including fallbacks, and 128 tokens with ``use_cudnn_fusion=False``. dtype : jnp.dtype Compute / parameter dtype. @@ -146,11 +154,14 @@ class _MoEBlock(TransformerEngineBase): input_axes: Tuple[Optional[str], ...] = () # Parallelism - data_parallelism_axes: Tuple[str, ...] = () + mesh_resource: Optional[MeshResource] = None + quant_before_fsdp_ag: bool = False # MoE knobs forwarded to ``moe()`` + use_cudnn_fusion: bool = True apply_topk_weights_early: bool = False recv_capacity_per_rank: Optional[int] = None + dispatch_checkpoint_name: Optional[str] = None # Dtypes / init / misc dtype: DType = jnp.float32 @@ -191,6 +202,11 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: Non-differentiable per-rank pre-drop recv-slot total; flags overflow when ``drop_on_overflow`` is set at ep_bootstrap. """ + mesh_resource, quant_before_fsdp_ag = _resolve_moe_mesh_resource( + self.mesh_resource, + self.quant_before_fsdp_ag, + ) + _, data_parallelism_axes = _moe_mesh_axes(mesh_resource) assert ( inputs.ndim == 3 ), f"_MoEBlock expects [batch, sequence, hidden] input, got shape {inputs.shape}" @@ -253,62 +269,66 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: jnp.float32, ) - ep_axis = get_active_resource_axis("ep_resource") mesh = _get_mesh() data_parallel_size = 1 - for axis in self.data_parallelism_axes: + for axis in data_parallelism_axes: data_parallel_size *= mesh.shape[axis] - def make_grouped_quantizer_set(postfix): - # Dispatched token groups span every data-parallel replica, - # whereas expert kernels have one group per global expert. - token_set = self.generate_quantizer_set( - f"{postfix}_token", - fp8_recipe=self.quantization_recipe, - n_groups=data_parallel_size * self.num_experts, - ) - expert_set = self.generate_quantizer_set( - f"{postfix}_expert", - fp8_recipe=self.quantization_recipe, - n_groups=self.num_experts, - ) - return QuantizerSet( - x=token_set.x, - kernel=expert_set.kernel, - dgrad=token_set.dgrad, + with global_shard_guard(mesh_resource): + + def make_grouped_quantizer_set(postfix): + # Dispatched token groups span every data-parallel replica, + # whereas expert kernels have one group per global expert. + token_set = self.generate_quantizer_set( + f"{postfix}_token", + fp8_recipe=self.quantization_recipe, + n_groups=data_parallel_size * self.num_experts, + ) + expert_set = self.generate_quantizer_set( + f"{postfix}_expert", + fp8_recipe=self.quantization_recipe, + n_groups=self.num_experts, + ) + return QuantizerSet( + x=token_set.x, + kernel=expert_set.kernel, + dgrad=token_set.dgrad, + ) + + quantizer_sets = ( + make_grouped_quantizer_set("_fc1"), + make_grouped_quantizer_set("_fc2"), ) - quantizer_sets = ( - make_grouped_quantizer_set("_fc1"), - make_grouped_quantizer_set("_fc2"), - ) - - return moe( - inputs, - gate_kernel, - wi, - wo, - wi_0_bias, - wi_1_bias, - wo_bias, - expert_bias, - num_experts=self.num_experts, - num_experts_per_tok=self.num_experts_per_tok, - activation_type=self.activation_type, - score_function=self.score_function, - use_pre_softmax=self.use_pre_softmax, - num_groups=self.num_groups, - group_topk=self.group_topk, - scaling_factor=self.scaling_factor, - aux_loss_coeff=self.aux_loss_coeff, - apply_topk_weights_early=self.apply_topk_weights_early, - quantizer_sets=quantizer_sets, - recv_capacity_per_rank=self.recv_capacity_per_rank, - ep_axis=ep_axis, - 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, - wo_kernel_axes=self.wo_kernel_axes, - dtype=self.dtype, - ) + return moe( + inputs, + gate_kernel, + wi, + wo, + wi_0_bias, + wi_1_bias, + wo_bias, + expert_bias, + num_experts=self.num_experts, + num_experts_per_tok=self.num_experts_per_tok, + activation_type=self.activation_type, + score_function=self.score_function, + use_pre_softmax=self.use_pre_softmax, + num_groups=self.num_groups, + group_topk=self.group_topk, + scaling_factor=self.scaling_factor, + aux_loss_coeff=self.aux_loss_coeff, + use_cudnn_fusion=self.use_cudnn_fusion, + cudnn_native_weight_layout=False, + apply_topk_weights_early=self.apply_topk_weights_early, + quantizer_sets=quantizer_sets, + recv_capacity_per_rank=self.recv_capacity_per_rank, + quant_before_fsdp_ag=quant_before_fsdp_ag, + mesh_resource=mesh_resource, + input_axes=self.input_axes, + gate_kernel_axes=self.gate_kernel_axes, + wi_kernel_axes=self.wi_kernel_axes, + wo_kernel_axes=self.wo_kernel_axes, + dtype=self.dtype, + dispatch_checkpoint_name=self.dispatch_checkpoint_name, + ) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 825d2788fde..d762719f9ae 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -32,10 +32,12 @@ import math import warnings +from dataclasses import replace from functools import partial from typing import Any, Optional, Tuple, Union import flax.struct +import flax.linen as nn import jax import jax.numpy as jnp from jax.ad_checkpoint import checkpoint_name @@ -44,25 +46,170 @@ from . import cpp_extensions as tex from .cpp_extensions.misc import jax_dtype_to_te_dtype from .quantize import ( + GroupedScaledTensor1x, GroupedQuantizer, QuantizerSet, + ScaledTensorFactory, + ScalingMode, + ScaledTensor2x, TensorUsage, noop_quantizer_set, with_sharding_constraint_by_logical_axes, ) +from .quantize.dequantizer import _unswizzle_mxfp8_grouped_scale +from .cpp_extensions.gemm import swizzled_scale from .flax.module import _convert_to_activation_function from .router import ScoreFunction, _validate_score_function -from .sharding import _get_mesh +from .sharding import MeshResource, _get_mesh, global_mesh_resource, global_shard_guard __all__ = ["get_moe_recv_capacity_per_rank", "moe"] +def _moe_mesh_axes(resource: MeshResource): + """Resolve physical MoE axes, keeping EP innermost in the batch shard.""" + if not isinstance(resource.ep_resource, str) or not resource.ep_resource: + raise ValueError("TE MoE requires MeshResource.ep_resource to name a physical mesh axis.") + outer_axes = tuple( + dict.fromkeys( + axis for axis in (resource.dp_resource, resource.fsdp_resource) if axis is not None + ) + ) + if any(not isinstance(axis, str) or not axis for axis in outer_axes): + raise ValueError("TE MoE DP and FSDP resources must name physical mesh axes.") + if resource.ep_resource in outer_axes: + raise ValueError("TE MoE EP and outer data-parallel axes must be distinct.") + return resource.ep_resource, outer_axes + + +def _resolve_moe_mesh_resource( + mesh_resource=None, + quant_before_fsdp_ag=False, +): + """Resolve and snapshot the explicit or active mesh resource.""" + if not isinstance(quant_before_fsdp_ag, bool): + raise TypeError("quant_before_fsdp_ag must be a bool.") + if mesh_resource is not None and not isinstance(mesh_resource, MeshResource): + raise TypeError("mesh_resource must be a MeshResource or None.") + if mesh_resource is None: + try: + mesh_resource = global_mesh_resource() + except AssertionError: + mesh_resource = None + + if mesh_resource is None: + raise ValueError( + "TE MoE requires mesh_resource=MeshResource(...) or an active global_shard_guard" + " context." + ) + # Snapshot the mutable resource so the custom VJP retains its forward-time axes. + mesh_resource = replace(mesh_resource) + _moe_mesh_axes(mesh_resource) + if quant_before_fsdp_ag and not mesh_resource.fsdp_resource: + raise ValueError("quant_before_fsdp_ag=True requires MeshResource.fsdp_resource.") + return mesh_resource, quant_before_fsdp_ag + + # Per-expert dispatch-slot alignment fed to ``tex.ep_prepare`` as # ``dispatch_output_per_expert_alignment``. NCCL EP HT requires the # per-expert recv block to be at least 128-token aligned, and all current # TE grouped-GEMM recipes (bf16/fp16/fp8/mxfp8) are satisfied by the # same 128-token tile, so a single constant covers every supported path. _ALIGN_SIZE = 128 +_CUDNN_JAX_ALIGN_SIZE = 256 + + +def _select_cudnn_jax_fusion(rejection_reasons: list[str]) -> bool: + """Select shared cuDNN GLU/dGLU fusion or the unfused TE path.""" + from transformer_engine_jax import get_device_compute_capability + + reasons = list(rejection_reasons) + try: + capability = get_device_compute_capability(0) + except RuntimeError as exc: + reasons.append(f"could not query GPU compute capability: {exc}") + else: + if capability < 100: + reasons.append(f"fused kernels require SM100+, got SM{capability}") + if not reasons: + available, error = tex.grouped_gemm_glu_dependencies_available() + if available: + return True + reasons.append(f"shared fused API is incompatible: {error}") + warnings.warn( + "use_cudnn_fusion=True: falling back to the unfused TE grouped-GEMM path: " + + "; ".join(reasons) + + ". Install cuDNN Frontend with shared grouped_gemm_glu/grouped_gemm_dglu " + "JAX APIs for dense MXFP8 and CuTeDSL JAX support; use supported GPU hardware.", + UserWarning, + stacklevel=2, + ) + return False + + +def _cudnn_jax_fusion_rejection_reasons( + x, + wi, + wi_0_bias, + wi_1_bias, + quantizer_sets, + *, + num_experts, + activation_type, + ep_axis, +) -> list[str]: + """Return reasons this call cannot use cuDNN's shared grouped GLU/dGLU JAX APIs.""" + errors = [] + if str(activation_type).lower() != "silu": + errors.append("requires activation_type='silu'") + if wi_0_bias is not None or wi_1_bias is not None: + errors.append("does not support FC1 gate/up bias") + hidden = x.shape[-1] + standard_layout = wi.ndim == 3 and wi.shape[-2] == hidden and wi.shape[-1] % 64 == 0 + native_layout = wi.ndim == 3 and wi.shape[-1] == hidden and wi.shape[-2] % 64 == 0 + if not standard_layout and not native_layout: + errors.append( + "requires rank-3 wi in standard [E,K,2N] or cuDNN-native [E,2N,K] " + f"layout with a 64-aligned gated dimension and K={hidden}, got {wi.shape}" + ) + if x.dtype not in (jnp.bfloat16, jnp.float16): + errors.append(f"requires BF16 or FP16 activations, got {x.dtype}") + + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets + required_quantizers = { + "fc1.x": fc1_quantizer_set.x, + "fc1.kernel": fc1_quantizer_set.kernel, + "fc1.dgrad": fc1_quantizer_set.dgrad, + "fc2.x": fc2_quantizer_set.x, + "fc2.kernel": fc2_quantizer_set.kernel, + "fc2.dgrad": fc2_quantizer_set.dgrad, + } + for name, quantizer in required_quantizers.items(): + if not isinstance(quantizer, GroupedQuantizer): + errors.append(f"requires a grouped MXFP8 quantizer for {name}") + elif quantizer.scaling_mode != ScalingMode.MXFP8_1D_SCALING: + errors.append(f"requires MXFP8_1D_SCALING for {name}") + elif not quantizer.q_layout.is_rowwise_colwise: + errors.append(f"requires rowwise+colwise quantization for {name}") + + if all(isinstance(q, GroupedQuantizer) for q in required_quantizers.values()): + if fc1_quantizer_set.x.q_dtype != fc1_quantizer_set.kernel.q_dtype: + errors.append("requires identical FC1 activation and weight MXFP8 payload dtypes") + if fc1_quantizer_set.dgrad.q_dtype != jnp.float8_e4m3fn: + errors.append("requires FC1 dgrad MXFP8 payload dtype float8_e4m3fn") + if fc2_quantizer_set.dgrad.q_dtype != fc2_quantizer_set.kernel.q_dtype: + errors.append("requires identical FC2 dgrad and weight MXFP8 payload dtypes") + supported = (jnp.float8_e4m3fn, jnp.float8_e5m2) + for name, quantizer in required_quantizers.items(): + if quantizer.q_dtype not in supported: + errors.append(f"unsupported MXFP8 payload dtype {quantizer.q_dtype} for {name}") + + mesh = _get_mesh() + if mesh is not None and not mesh.empty and ep_axis in mesh.shape: + num_local_experts = num_experts // mesh.shape[ep_axis] + if num_local_experts > 1024: + errors.append(f"requires at most 1024 local experts, got {num_local_experts}") + + return errors def get_moe_recv_capacity_per_rank( @@ -72,15 +219,23 @@ def get_moe_recv_capacity_per_rank( max_tokens_per_rank: int, ep_size: int, recv_capacity_factor: Optional[float] = None, - alignment: int = _ALIGN_SIZE, + alignment: Optional[int] = None, + use_cudnn_fusion: bool = True, ) -> int: """Return the aligned receive capacity for one EP rank. ``recv_capacity_factor=None`` reserves the dropless worst case. A finite factor >= 1 scales the capacity needed by perfectly balanced routing and is capped at the worst case. The balanced baseline includes the independent - per-local-expert alignment required by NCCL EP. + per-local-expert alignment required by NCCL EP. When ``alignment`` is not + supplied, ``use_cudnn_fusion=True`` (default) reserves 256-token alignment, + including fallbacks; ``False`` reserves 128 for regular TE grouped GEMM. This keeps + eager bootstrap callers in sync with the later compiled ``moe()`` call. """ + if not isinstance(use_cudnn_fusion, bool): + raise TypeError("use_cudnn_fusion must be a bool") + if alignment is None: + alignment = _CUDNN_JAX_ALIGN_SIZE if use_cudnn_fusion else _ALIGN_SIZE if num_experts <= 0 or num_experts_per_tok <= 0 or max_tokens_per_rank <= 0: raise ValueError( "num_experts, num_experts_per_tok, and max_tokens_per_rank must be positive" @@ -337,6 +492,241 @@ def _validate_moe_quantizer_sets( ) +def _with_grouped_weight_buffers(tensor, data, scale_inv, original_shape): + """Rebuild a dense grouped weight while preserving its quantization metadata.""" + return GroupedScaledTensor1x( + data=data, + scale_inv=scale_inv, + amax=tensor.amax, + first_dims=None, + last_dims=None, + scaling_mode=tensor.scaling_mode, + dq_dtype=tensor.dq_dtype, + _dq_func=tensor._dq_func, + is_colwise=tensor.is_colwise, + data_layout=tensor.data_layout, + flatten_axis=tensor.flatten_axis, + original_shape=original_shape, + pre_swizzled=True, + ) + + +def _gather_quantized_weight(tensor, fsdp_axis: str, fsdp_size: int, sharded_axis: int): + """Gather an MXFP8 grouped weight without gathering its BF16 source. + + Grouped tensor data and scales are flat, with scales independently padded + and swizzled for each expert. Reassemble each expert in logical order, + then pad and swizzle its gathered scales for grouped GEMM. Whole-expert + shards can instead concatenate their already-swizzled scale blocks. + """ + if isinstance(tensor, ScaledTensor2x): + return ScaledTensor2x( + _gather_quantized_weight(tensor.rowwise_tensor, fsdp_axis, fsdp_size, sharded_axis), + _gather_quantized_weight(tensor.colwise_tensor, fsdp_axis, fsdp_size, sharded_axis), + ) + if not isinstance(tensor, GroupedScaledTensor1x): + raise TypeError("Quantized FSDP gather requires grouped MXFP8 weight tensors.") + + local_shape = tensor.original_shape + num_experts = local_shape[0] + # The T layout swaps the two matrix dimensions within each expert. + data_axis = ( + 3 - sharded_axis if tensor.data_layout == "T" and sharded_axis != 0 else sharded_axis + ) + global_shape = list(local_shape) + global_shape[data_axis] *= fsdp_size + global_shape = tuple(global_shape) + data = jax.lax.all_gather( + tensor.data.reshape(local_shape), fsdp_axis, axis=data_axis, tiled=True + ).reshape(-1) + + local_matrix = local_shape[1:] + global_matrix = global_shape[1:] + # Scales use a block-wise 2D view. The local shape can include padding + # that differs from the padding required after the gather. + local_scale_shape = tensor.scaling_mode.get_scale_shape( + local_matrix, + data_layout=tensor.data_layout, + is_colwise=tensor.is_colwise, + is_padded=True, + flatten_axis=tensor.flatten_axis - 1, + ) + if sharded_axis == 0: + # The grouped allocation includes a worst-case padding tail. Gather + # only actual expert scale blocks, then allocate the gathered tail. + local_scale_size = math.prod(local_scale_shape) + scale_inv = jax.lax.all_gather( + tensor.scale_inv[: num_experts * local_scale_size], + fsdp_axis, + axis=0, + tiled=True, + ) + else: + scale_inv = _gather_quantized_matrix_scales( + tensor, fsdp_axis, data_axis, local_scale_shape, global_matrix + ) + expected_scale_size = tensor.scaling_mode.get_grouped_scale_shape( + global_shape, + global_shape[0], + tensor.is_colwise, + is_padded=True, + flatten_axis=tensor.flatten_axis, + )[0] + scale_inv = jnp.pad(scale_inv, (0, expected_scale_size - scale_inv.size)) + return _with_grouped_weight_buffers(tensor, data, scale_inv, global_shape) + + +def _weight_scale_shapes(tensor, matrix_shape): + """Return padded and active per-expert MXFP8 scale shapes.""" + kwargs = { + "data_layout": tensor.data_layout, + "is_colwise": tensor.is_colwise, + "flatten_axis": tensor.flatten_axis - 1, + } + return ( + tensor.scaling_mode.get_scale_shape(matrix_shape, is_padded=True, **kwargs), + tensor.scaling_mode.get_scale_shape(matrix_shape, is_padded=False, **kwargs), + ) + + +def _read_weight_scale(tensor, expert, padded_shape, plain_shape): + """Read one expert's active scales in matrix order.""" + scale_size = math.prod(padded_shape) + swizzled = jax.lax.dynamic_slice_in_dim(tensor.scale_inv, expert * scale_size, scale_size) + plain = _unswizzle_mxfp8_grouped_scale(swizzled, padded_shape, tensor.is_colwise) + return plain[: plain_shape[0], : plain_shape[1]] + + +def _swizzle_weight_scale(plain, padded_shape, is_colwise): + """Pad and swizzle one expert's scale matrix for grouped GEMM.""" + padded = jnp.pad( + plain, ((0, padded_shape[0] - plain.shape[0]), (0, padded_shape[1] - plain.shape[1])) + ) + return swizzled_scale(padded, 1, is_colwise).reshape(-1) + + +def _gather_quantized_matrix_scales(tensor, fsdp_axis, data_axis, local_scale_shape, global_matrix): + """Reassemble scales when FSDP splits each expert's matrix dimension.""" + _, local_unpadded_shape = _weight_scale_shapes(tensor, tensor.original_shape[1:]) + global_padded_shape, global_unpadded_shape = _weight_scale_shapes(tensor, global_matrix) + scale_axis = data_axis - 1 + gathered_scales = [] + for expert in range(tensor.original_shape[0]): + local_plain = _read_weight_scale(tensor, expert, local_scale_shape, local_unpadded_shape) + full_plain = jax.lax.all_gather(local_plain, fsdp_axis, axis=scale_axis, tiled=True) + assert full_plain.shape == global_unpadded_shape + gathered_scales.append( + _swizzle_weight_scale(full_plain, global_padded_shape, tensor.is_colwise) + ) + return jnp.concatenate(gathered_scales) + + +def _resolve_cudnn_native_weight_layout(wi, hidden, layout=None): + """Resolve FC1 parameter storage using its global shape or an explicit flag.""" + if layout is not None and not isinstance(layout, bool): + raise TypeError("cudnn_native_weight_layout must be a bool or None") + standard = wi.ndim == 3 and wi.shape[1] == hidden and wi.shape[2] % 2 == 0 + native = wi.ndim == 3 and wi.shape[2] == hidden and wi.shape[1] % 64 == 0 + if layout is None: + if standard and native: + raise ValueError( + "Ambiguous square FC1 weight layout; pass cudnn_native_weight_layout=" + "False for standard [E,K,2N] or True for native [E,2N,K] storage." + ) + if standard or native: + return native + elif native if layout else standard: + return layout + raise ValueError(f"Invalid FC1 weight shape {wi.shape} for K={hidden} and layout={layout}") + + +def _weight_fsdp_axis(spec, fsdp_axis): + """Find the tensor dimension partitioned by the physical FSDP resource.""" + return next( + ( + i + for i, axes in enumerate(spec) + if fsdp_axis in (axes if isinstance(axes, tuple) else (axes,)) + ), + None, + ) + + +def _swiglu_column_indices(width, interleave): + """Map contiguous gate/up halves to 32-column pairs, or undo that mapping.""" + if width % 64: + raise ValueError("The global gated dimension must be divisible by 64") + blocks = jnp.arange(width // 32) + blocks = ( + blocks.reshape(2, -1).T.reshape(-1) if interleave else blocks.reshape(-1, 2).T.reshape(-1) + ) + columns = (blocks[:, None] * 32 + jnp.arange(32)[None, :]).reshape(-1) + return blocks, columns + + +def _reorder_quantized_swiglu_weight(tensor, *, interleave): + """Permute gathered MXFP8 columns and their scales without requantization. + + Input logical storage is [E,K,2N]. Both rowwise and colwise copies may use + physical N or T layout. Gate/up boundaries and permutations are 32-aligned, + so each quantization block keeps its original data and inverse scale. + """ + if isinstance(tensor, ScaledTensor2x): + return ScaledTensor2x( + _reorder_quantized_swiglu_weight(tensor.rowwise_tensor, interleave=interleave), + _reorder_quantized_swiglu_weight(tensor.colwise_tensor, interleave=interleave), + ) + if not isinstance(tensor, GroupedScaledTensor1x): + raise TypeError("SwiGLU weight reordering requires grouped MXFP8 tensors") + if tensor.scaling_mode != ScalingMode.MXFP8_1D_SCALING or not tensor.pre_swizzled: + raise ValueError("SwiGLU weight reordering requires pre-swizzled MXFP8 scales") + + shape = tensor.original_shape + data_axis = 1 if tensor.data_layout == "T" else 2 + block_indices, column_indices = _swiglu_column_indices(shape[data_axis], interleave) + data = jnp.take(tensor.data.reshape(shape), column_indices, axis=data_axis).reshape(-1) + padded_shape, plain_shape = _weight_scale_shapes(tensor, shape[1:]) + scale_axis = data_axis - 1 + # Rowwise scales cover 32 columns; colwise scales have one entry per column. + scale_indices = ( + block_indices if plain_shape[scale_axis] == shape[data_axis] // 32 else column_indices + ) + reordered_scales = [] + for expert in range(shape[0]): + plain = _read_weight_scale(tensor, expert, padded_shape, plain_shape) + reordered = jnp.take(plain, scale_indices, axis=scale_axis) + reordered_scales.append(_swizzle_weight_scale(reordered, padded_shape, tensor.is_colwise)) + scale_inv = jnp.concatenate(reordered_scales) + scale_inv = jnp.pad(scale_inv, (0, tensor.scale_inv.size - scale_inv.size)) + return _with_grouped_weight_buffers(tensor, data, scale_inv, shape) + + +def _prepare_fc1_weight( + wi, quantizer, *, fused, native_layout, quant_before_fsdp_ag, fsdp_axis, fsdp_size, sharded_axis +): + """Convert FC1 storage and gather quantized shards in the correct logical order.""" + gated_axis = 1 if native_layout else 2 + convert_layout = bool(fused) != native_layout + convert_after_gather = convert_layout and quant_before_fsdp_ag and sharded_axis == gated_axis + if native_layout and not fused: + wi = wi.transpose(0, 2, 1) + if sharded_axis in (1, 2): + sharded_axis = 3 - sharded_axis + if not convert_after_gather: + gate, up = tex.unpack_gated_activation_pair(wi) + wi = jnp.concatenate((gate, up), axis=-1) + elif fused and not native_layout and not convert_after_gather: + gate, up = jnp.split(wi, 2, axis=-1) + wi = tex.pack_gated_activation_pair(gate, up) + + casted = tex.grouped_quantize(wi, quantizer, flatten_axis=-1) + if quant_before_fsdp_ag and sharded_axis is not None: + casted = _gather_quantized_weight(casted, fsdp_axis, fsdp_size, sharded_axis) + if convert_after_gather: + casted = _reorder_quantized_swiglu_weight(casted, interleave=bool(fused)) + return casted + + def _ffn_fwd_per_shard( recv_tokens_local: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, @@ -351,9 +741,16 @@ def _ffn_fwd_per_shard( num_local_experts: int, activation_type: str, apply_topk_weights_early: bool, + use_cudnn_jax_fusion: bool, wi_0_checkpoint_name: Optional[str], wi_1_checkpoint_name: Optional[str], wo_checkpoint_name: Optional[str], + cudnn_native_weight_layout: bool, + quant_before_fsdp_ag: bool, + fsdp_axis: Optional[str], + fsdp_size: int, + wi_fsdp_axis: Optional[int], + wo_fsdp_axis: Optional[int], ): """Run the grouped FFN on one shard's EP receive buffer.""" hidden = recv_tokens_local.shape[-1] @@ -364,8 +761,6 @@ def _ffn_fwd_per_shard( wi = wi.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) - # ``wi`` is stored in its gated-SwiGLU layout [expert, hidden, 2*mlp]. - # Keeping it contiguous lets grouped quantize/GEMM consume it directly. wi_combined_bias = ( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) @@ -377,18 +772,118 @@ def _ffn_fwd_per_shard( group_sizes, flatten_axis=-1, ) - casted_wi = tex.grouped_quantize(wi, fc1_quantizer_set.kernel, flatten_axis=-1) - combined_out = tex.grouped_gemm( - casted_sorted_x.get_tensor(usage=TensorUsage.LHS), - casted_wi.get_tensor(usage=TensorUsage.RHS), - contracting_dims=((1,), (1,)), - bias=wi_combined_bias, + casted_wi = _prepare_fc1_weight( + wi, + fc1_quantizer_set.kernel, + fused=use_cudnn_jax_fusion, + native_layout=cudnn_native_weight_layout, + quant_before_fsdp_ag=quant_before_fsdp_ag, + fsdp_axis=fsdp_axis, + fsdp_size=fsdp_size, + sharded_axis=wi_fsdp_axis, ) - gate_proj_out, up_proj_out = jnp.split(combined_out, 2, axis=-1) - if wi_0_checkpoint_name is not None: - gate_proj_out = checkpoint_name(gate_proj_out, wi_0_checkpoint_name) - if wi_1_checkpoint_name is not None: - up_proj_out = checkpoint_name(up_proj_out, wi_1_checkpoint_name) + casted_intermediate = None + if use_cudnn_jax_fusion: + casted_sorted_x_lhs = casted_sorted_x.get_tensor(usage=TensorUsage.LHS) + casted_wi_rhs = casted_wi.get_tensor( + usage=TensorUsage.LHS if cudnn_native_weight_layout else TensorUsage.RHS + ) + # Quantized gathering may have expanded the gated dimension since packing. + combined = casted_wi_rhs.data.size // (num_local_experts * hidden) + padded_offsets = jnp.cumsum(group_sizes, dtype=jnp.int32) + prob = ( + recv_w_flat[:, None, None] + if apply_topk_weights_early + else jnp.ones((sorted_x.shape[0], 1, 1), dtype=jnp.float32) + ) + ( + combined_out_3d, + intermediate_row, + intermediate_col, + intermediate_scale_row, + intermediate_scale_col, + ) = tex.grouped_gemm_glu( + casted_sorted_x_lhs.data.reshape(sorted_x.shape[0], hidden, 1), + ( + casted_wi_rhs.data.reshape(num_local_experts, combined, hidden) + if cudnn_native_weight_layout + else casted_wi_rhs.data.reshape(num_local_experts, hidden, combined).transpose( + 0, 2, 1 + ) + ), + casted_sorted_x_lhs.scale_inv, + casted_wi_rhs.scale_inv, + padded_offsets, + prob, + compute_dtype=sorted_x.dtype, + output_dtype=fc2_quantizer_set.x.q_dtype, + ) + combined_out = combined_out_3d.reshape(sorted_x.shape[0], combined) + if wi_0_checkpoint_name is not None: + combined_out = checkpoint_name(combined_out, wi_0_checkpoint_name) + if wi_1_checkpoint_name is not None: + combined_out = checkpoint_name(combined_out, wi_1_checkpoint_name) + # The fused backward consumes interleaved C directly. Keep both residual + # slots as aliases instead of materializing an otherwise-unused unpack. + gate_proj_out = up_proj_out = combined_out + + intermediate_shape = (sorted_x.shape[0], combined // 2) + scaling_mode = fc2_quantizer_set.x.scaling_mode + row_scale_size = scaling_mode.get_grouped_scale_shape( + intermediate_shape, + num_local_experts, + False, + is_padded=True, + flatten_axis=1, + )[0] + col_scale_size = scaling_mode.get_grouped_scale_shape( + intermediate_shape, + num_local_experts, + True, + is_padded=True, + flatten_axis=1, + )[0] + intermediate_scale_row = jnp.pad( + intermediate_scale_row, + (0, row_scale_size - intermediate_scale_row.size), + ) + intermediate_scale_col = jnp.pad( + intermediate_scale_col, + (0, col_scale_size - intermediate_scale_col.size), + ) + for checkpoint_label in (wi_0_checkpoint_name, wi_1_checkpoint_name): + if checkpoint_label is not None: + intermediate_row = checkpoint_name(intermediate_row, checkpoint_label) + intermediate_col = checkpoint_name(intermediate_col, checkpoint_label) + intermediate_scale_row = checkpoint_name(intermediate_scale_row, checkpoint_label) + intermediate_scale_col = checkpoint_name(intermediate_scale_col, checkpoint_label) + casted_intermediate = ScaledTensorFactory.create( + data=intermediate_row.reshape(-1), + scale_inv=intermediate_scale_row, + colwise_data=intermediate_col.reshape(-1), + colwise_scale_inv=intermediate_scale_col, + scaling_mode=scaling_mode, + dq_dtype=sorted_x.dtype, + data_layout=fc2_quantizer_set.x.data_layout, + q_layout=fc2_quantizer_set.x.q_layout, + flatten_axis=1, + first_dims=group_sizes, + original_shape=intermediate_shape, + pre_swizzled=True, + ) + else: + combined_out = tex.grouped_gemm( + casted_sorted_x.get_tensor(usage=TensorUsage.LHS), + casted_wi.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((1,), (1,)), + bias=wi_combined_bias, + ) + gate_proj_out, up_proj_out = jnp.split(combined_out, 2, axis=-1) + if not use_cudnn_jax_fusion: + if wi_0_checkpoint_name is not None: + gate_proj_out = checkpoint_name(gate_proj_out, wi_0_checkpoint_name) + if wi_1_checkpoint_name is not None: + up_proj_out = checkpoint_name(up_proj_out, wi_1_checkpoint_name) # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM # output dtype; the activation output (`intermediate`) stays in the @@ -396,21 +891,25 @@ def _ffn_fwd_per_shard( # that's all bf16; for FP8/FP4 the downstream grouped_quantize is what # transitions to the target precision. act_fn = _convert_to_activation_function(activation_type) - intermediate = act_fn(gate_proj_out) * up_proj_out - if apply_topk_weights_early: + if not use_cudnn_jax_fusion: + intermediate = act_fn(gate_proj_out) * up_proj_out + if apply_topk_weights_early and not use_cudnn_jax_fusion: # Fold the per-token combine weights into the FFN intermediate; # the downstream wo GEMM is linear so this is equivalent to the # late-weighting path. Grouped GEMM skips overallocation tail padding automatically. # Padding between groups is padded with zeros by NCCL EP. intermediate = intermediate * recv_w_flat[:, None].astype(intermediate.dtype) - casted_intermediate = tex.grouped_quantize( - intermediate, - fc2_quantizer_set.x, - group_sizes, - flatten_axis=-1, - ) + if not use_cudnn_jax_fusion: + casted_intermediate = tex.grouped_quantize( + intermediate, + fc2_quantizer_set.x, + group_sizes, + flatten_axis=-1, + ) casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) + if quant_before_fsdp_ag and wo_fsdp_axis is not None: + casted_wo = _gather_quantized_weight(casted_wo, fsdp_axis, fsdp_size, wo_fsdp_axis) expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), casted_wo.get_tensor(usage=TensorUsage.RHS), @@ -421,10 +920,17 @@ def _ffn_fwd_per_shard( expert_outputs = checkpoint_name(expert_outputs, wo_checkpoint_name) expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) group_sizes_2d = group_sizes.reshape(1, num_local_experts) + ffn_activation_residual = combined_out if use_cudnn_jax_fusion else gate_proj_out residuals = ( casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(fc1_quantizer_set.x), - casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(fc1_quantizer_set.kernel), - gate_proj_out, + casted_wi.get_tensor( + usage=( + TensorUsage.RHS + if use_cudnn_jax_fusion and cudnn_native_weight_layout + else TensorUsage.RHS_TRANS + ) + ).checkpoint(fc1_quantizer_set.kernel), + ffn_activation_residual, up_proj_out, casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(fc2_quantizer_set.x), casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(fc2_quantizer_set.kernel), @@ -448,6 +954,8 @@ def _ffn_bwd_per_shard( activation_type: str, apply_topk_weights_early: bool, has_bias: bool, + use_cudnn_jax_fusion: bool, + cudnn_native_weight_layout: bool, ): """Backward mirror of :func:`_ffn_fwd_per_shard`.""" group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) @@ -464,11 +972,6 @@ def _ffn_bwd_per_shard( ) _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) - d_intermediate = tex.grouped_gemm( - _casted_d_eo_lhs, - casted_wo_rhs_trans, - contracting_dims=((1,), (2,)), - ) d_wo = tex.grouped_gemm( casted_intermediate_lhs_trans, _casted_d_eo_rhs, @@ -476,58 +979,137 @@ def _ffn_bwd_per_shard( ) d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None - act_fn = _convert_to_activation_function(activation_type) - if apply_topk_weights_early: - # intermediate' = intermediate * w. - # Masking is not required as: - # 1. Padding between groups is zero padded due to NCCL EP. - # 2. Overallocated padding past all groups is uninitialized, but subsequent GEMMs and EP are all group-size aware and will not read past the final group. - w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) - gate_proj_for_bwd = gate_proj_out - up_proj_for_bwd = up_proj_out - intermediate_unweighted = act_fn(gate_proj_out) * up_proj_out - d_recv_w_from_intermediate = jnp.sum( - d_intermediate * intermediate_unweighted, - axis=-1, - ).astype(recv_w_flat.dtype) - d_intermediate = d_intermediate * w_b + if use_cudnn_jax_fusion: + # The forward residual uses this slot for cuDNN's interleaved pre-activation C. + combined_out = gate_proj_out + rows, combined = combined_out.shape + intermediate = combined // 2 + hidden = d_eo_2d.shape[-1] + num_local_experts = group_sizes.size + padded_offsets = jnp.cumsum(group_sizes, dtype=jnp.int32) + prob = recv_w_flat if apply_topk_weights_early else jnp.ones((rows,), dtype=jnp.float32) + ( + d_combined_row, + d_combined_col, + dprob, + d_combined_scale_row, + d_combined_scale_col, + ) = tex.grouped_gemm_dglu( + _casted_d_eo_lhs.data.reshape(rows, hidden), + casted_wo_rhs_trans.data.reshape(num_local_experts, intermediate, hidden), + combined_out, + _casted_d_eo_lhs.scale_inv, + casted_wo_rhs_trans.scale_inv, + padded_offsets, + prob, + output_dtype=fc1_quantizer_set.dgrad.q_dtype, + ) + d_recv_w_from_intermediate = ( + dprob.astype(recv_w_flat.dtype) + if apply_topk_weights_early + else jnp.zeros_like(recv_w_flat) + ) + + combined_shape = (rows, combined) + scaling_mode = fc1_quantizer_set.dgrad.scaling_mode + row_scale_size = scaling_mode.get_grouped_scale_shape( + combined_shape, + num_local_experts, + False, + is_padded=True, + flatten_axis=1, + )[0] + col_scale_size = scaling_mode.get_grouped_scale_shape( + combined_shape, + num_local_experts, + True, + is_padded=True, + flatten_axis=1, + )[0] + d_combined_scale_row = jnp.pad( + d_combined_scale_row, + (0, row_scale_size - d_combined_scale_row.size), + ) + d_combined_scale_col = jnp.pad( + d_combined_scale_col, + (0, col_scale_size - d_combined_scale_col.size), + ) + casted_d_combined = ScaledTensorFactory.create( + data=d_combined_row.reshape(-1), + scale_inv=d_combined_scale_row, + colwise_data=d_combined_col.reshape(-1), + colwise_scale_inv=d_combined_scale_col, + scaling_mode=scaling_mode, + dq_dtype=d_eo_2d.dtype, + data_layout=fc1_quantizer_set.dgrad.data_layout, + q_layout=fc1_quantizer_set.dgrad.q_layout, + flatten_axis=1, + first_dims=group_sizes, + original_shape=combined_shape, + pre_swizzled=True, + ) + d_combined_for_bias = None else: - gate_proj_for_bwd = gate_proj_out - up_proj_for_bwd = up_proj_out - d_recv_w_from_intermediate = jnp.zeros_like(recv_w_flat) - - # Activation bwd, symmetric with the fwd: silu' and the two - # elementwise products run in the GEMM dtype (no fp32 island), so - # the chain rule composes through at the same precision the wi/wo - # GEMMs consume. - act_gp, dact_pullback = jax.vjp(act_fn, gate_proj_for_bwd) - d_up_proj_out = d_intermediate * act_gp - (d_gate_proj_out,) = dact_pullback(d_intermediate * up_proj_for_bwd) - - # wi bwd (fused gate/up via concat). Mirror the fused fwd: pack the - # gate/up cotangents along the trailing axis, run a single - # grouped_quantize + two grouped_gemm pair (one dgrad, one wgrad) - # against the fused casted_wi_rhs_trans residual, then split the - # wgrad result remains in the contiguous gated-SwiGLU ``wi`` layout. - d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) - casted_d_combined = tex.grouped_quantize( - d_combined, - fc1_quantizer_set.dgrad, - group_sizes, - flatten_axis=-1, - ) + d_intermediate = tex.grouped_gemm( + _casted_d_eo_lhs, + casted_wo_rhs_trans, + contracting_dims=((1,), (2,)), + ) + act_fn = _convert_to_activation_function(activation_type) + if apply_topk_weights_early: + # intermediate' = intermediate * w. + # Masking is not required as grouped GEMMs consume only group-size rows. + w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) + intermediate_unweighted = act_fn(gate_proj_out) * up_proj_out + d_recv_w_from_intermediate = jnp.sum( + d_intermediate * intermediate_unweighted, + axis=-1, + ).astype(recv_w_flat.dtype) + d_intermediate = d_intermediate * w_b + else: + d_recv_w_from_intermediate = jnp.zeros_like(recv_w_flat) + + # Activation bwd stays in the GEMM dtype, matching the forward path. + act_gp, dact_pullback = jax.vjp(act_fn, gate_proj_out) + d_up_proj_out = d_intermediate * act_gp + (d_gate_proj_out,) = dact_pullback(d_intermediate * up_proj_out) + d_combined_for_bias = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) + casted_d_combined = tex.grouped_quantize( + d_combined_for_bias, + fc1_quantizer_set.dgrad, + group_sizes, + flatten_axis=-1, + ) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), casted_wi_rhs_trans, - contracting_dims=((1,), (2,)), - ) - d_wi_combined = tex.grouped_gemm( - casted_sorted_x_lhs_trans, - casted_d_combined.get_tensor(usage=TensorUsage.RHS), - contracting_dims=((0,), (0,)), + contracting_dims=( + (1,), + (1 if use_cudnn_jax_fusion and cudnn_native_weight_layout else 2,), + ), ) + if use_cudnn_jax_fusion and cudnn_native_weight_layout: + # dY^T @ X directly produces [E,2N,K], matching the persistent native + # parameter, without a post-GEMM transpose or de-interleave/repack. + d_wi_combined = tex.grouped_gemm( + casted_d_combined.get_tensor(usage=TensorUsage.LHS_TRANS), + casted_sorted_x_lhs_trans, + contracting_dims=((0,), (0,)), + ) + else: + d_wi_combined = tex.grouped_gemm( + casted_sorted_x_lhs_trans, + casted_d_combined.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((0,), (0,)), + ) + if use_cudnn_jax_fusion and not cudnn_native_weight_layout: + d_wi_gate, d_wi_up = tex.unpack_gated_activation_pair(d_wi_combined) + d_wi_combined = jnp.concatenate([d_wi_gate, d_wi_up], axis=-1) + elif not use_cudnn_jax_fusion and cudnn_native_weight_layout: + d_wi_gate, d_wi_up = jnp.split(d_wi_combined, 2, axis=-1) + d_wi_combined = tex.pack_gated_activation_pair(d_wi_gate, d_wi_up).transpose(0, 2, 1) if has_bias: - d_wi_combined_bias = tex.grouped_dbias(d_combined, group_sizes) + d_wi_combined_bias = tex.grouped_dbias(d_combined_for_bias, group_sizes) d_wi_0_bias, d_wi_1_bias = jnp.split(d_wi_combined_bias, 2, axis=-1) else: d_wi_0_bias = None @@ -570,8 +1152,7 @@ def _moe_fwd_rule( group_topk, scaling_factor, aux_loss_coeff, - ep_axis, - data_parallelism_axes, + mesh_resource, input_axes, gate_kernel_axes, wi_kernel_axes, @@ -579,123 +1160,148 @@ def _moe_fwd_rule( dtype, apply_topk_weights_early, recv_capacity_per_rank, + use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, wo_checkpoint_name, + dispatch_checkpoint_name, + quant_before_fsdp_ag, + use_cudnn_fusion: bool = True, + cudnn_native_weight_layout: Optional[bool] = None, ): """Forward: gate -> topk -> ep_dispatch -> FFN -> ep_combine. Returns ``(output, aux_loss)``. ``aux_loss`` is a zero scalar when ``aux_loss_coeff == 0``. """ - del gate_kernel_axes, wi_kernel_axes, wo_kernel_axes # used in bwd only - from jax.experimental.shard_map import shard_map - - x = with_sharding_constraint_by_logical_axes(x, input_axes) - - mesh = _get_mesh() - if mesh is None or mesh.empty: - 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] - 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 - - dp_size = 1 - for ax in data_parallelism_axes: - dp_size *= mesh.shape[ax] - num_procs = num_ep * dp_size - _validate_moe_quantizer_sets( - quantizer_sets, - num_token_groups=dp_size * num_experts, - num_expert_groups=num_experts, - ) - - B, S, H = x.shape - K = num_experts_per_tok - if B % num_procs != 0: - raise ValueError(f"batch={B} not divisible by ep*dp={num_procs}") - - # Per-rank send capacity: B/num_procs rows x S tokens per rank. - max_tokens_per_rank = (B // num_procs) * S - worst_case_recv_pr = get_moe_recv_capacity_per_rank( - num_experts=num_experts, - num_experts_per_tok=K, - max_tokens_per_rank=max_tokens_per_rank, - ep_size=num_ep, - ) - if recv_capacity_per_rank is None: - recv_pr = worst_case_recv_pr - else: - recv_pr = int(recv_capacity_per_rank) - if recv_pr <= 0 or recv_pr % _ALIGN_SIZE != 0: - raise ValueError( - f"recv_capacity_per_rank must be a positive multiple of {_ALIGN_SIZE}, got" - f" {recv_pr}" - ) + with global_shard_guard(mesh_resource): + ep_axis, data_parallelism_axes = _moe_mesh_axes(mesh_resource) + del gate_kernel_axes # used in bwd only + from jax.experimental.shard_map import shard_map + + x = with_sharding_constraint_by_logical_axes(x, input_axes) + + mesh = _get_mesh() + if mesh is None or mesh.empty: + 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] + 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 + + dp_size = 1 + for ax in data_parallelism_axes: + dp_size *= mesh.shape[ax] + num_procs = num_ep * dp_size + _validate_moe_quantizer_sets( + quantizer_sets, + num_token_groups=dp_size * num_experts, + num_expert_groups=num_experts, + ) + B, S, H = x.shape + K = num_experts_per_tok + cudnn_native_weight_layout = _resolve_cudnn_native_weight_layout( + wi, H, cudnn_native_weight_layout + ) + kernel_spec = P(ep_axis, None, None) + wi_input_spec = wo_input_spec = kernel_spec + wi_fsdp_axis = wo_fsdp_axis = None + if quant_before_fsdp_ag: + # Logical axes describe parameter storage even under Auto mode, + # where tracer types do not expose the physical input sharding. + if nn.get_logical_axis_rules(): + wi_input_spec = nn.logical_to_mesh_axes(wi_kernel_axes) + wo_input_spec = nn.logical_to_mesh_axes(wo_kernel_axes) + else: + # Preserve the existing no-Flax-rules convention. + wi_input_spec = ( + P(ep_axis, None, mesh_resource.fsdp_resource) + if cudnn_native_weight_layout + else P(ep_axis, mesh_resource.fsdp_resource, None) + ) + wo_input_spec = P(ep_axis, None, mesh_resource.fsdp_resource) + wi_fsdp_axis = _weight_fsdp_axis(wi_input_spec, mesh_resource.fsdp_resource) + wo_fsdp_axis = _weight_fsdp_axis(wo_input_spec, mesh_resource.fsdp_resource) + if mesh_resource.fsdp_resource not in data_parallelism_axes: + raise ValueError( + "Quantized weight all-gather requires its FSDP axis among the outer batch axes." + ) + if any(quantizer_set.kernel is None for quantizer_set in quantizer_sets): + raise ValueError("Quantized weight all-gather requires MXFP8 kernel quantizers.") + if any( + axis is not None + and axis != 0 + and weight.shape[axis] % (mesh.shape[mesh_resource.fsdp_resource] * 32) + for weight, axis in ((wi, wi_fsdp_axis), (wo, wo_fsdp_axis)) + ): + raise ValueError("FSDP weight shards must be divisible by the MXFP8 block size 32.") + + if B % num_procs != 0: + raise ValueError(f"batch={B} not divisible by ep*dp={num_procs}") + + # Per-rank send capacity: B/num_procs rows x S tokens per rank. + max_tokens_per_rank = (B // num_procs) * S + # Keep capacity and alignment consistent with an EP bootstrap sized + # for requested fusion, even when API/hardware checks choose unfused. + dispatch_alignment = _CUDNN_JAX_ALIGN_SIZE if use_cudnn_fusion else _ALIGN_SIZE + worst_case_recv_pr = get_moe_recv_capacity_per_rank( + num_experts=num_experts, + num_experts_per_tok=K, + max_tokens_per_rank=max_tokens_per_rank, + ep_size=num_ep, + alignment=dispatch_alignment, + ) + if recv_capacity_per_rank is None: + recv_pr = worst_case_recv_pr + else: + recv_pr = int(recv_capacity_per_rank) + if recv_pr <= 0 or recv_pr % dispatch_alignment != 0: + raise ValueError( + "recv_capacity_per_rank must be a positive multiple of " + f"{dispatch_alignment}, got" + f" {recv_pr}" + ) - _te_ep_assert_compatible_bootstrap( - num_experts=num_experts, - max_tokens_per_rank=max_tokens_per_rank, - recv_capacity_per_rank=recv_pr, - hidden_dim=H, - ep_size=num_ep, - ) + _te_ep_assert_compatible_bootstrap( + num_experts=num_experts, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_pr, + hidden_dim=H, + 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) - 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)) - - # ---------------- Gate (global view) ---------------- - # tex.fused_topk_with_score_function is only validated against its - # pytorch reference at fp32 (see tests/pytorch/test_fused_router.py: - # parametrize gates dtype on torch.float32 only; the tolerance helper - # raises NotImplementedError for any other dtype). Keeping logits in - # the activation dtype (e.g. bf16) lets sigmoid / softmax / topk - # accumulate at low precision and silently produce NaNs on tokens - # whose normalised weights underflow. Cast to fp32 here to stay in - # the validated regime. - gate_kernel_cast = gate_kernel.astype(x.dtype) - gate_logits = jnp.einsum("bsh,he->bse", x, gate_kernel_cast) - logits_2d = gate_logits.reshape(-1, num_experts).astype(jnp.float32) - - # ---------------- Routing (global view) ---------------- - # expert_bias is an empty (shape-(0,)) sentinel when the caller did - # not enable it; the primitive treats that as "no bias". - eb_arg = expert_bias if expert_bias.shape != (0,) else jnp.zeros((0,), dtype=jnp.float32) - sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( - logits_2d, - topk=K, - use_pre_softmax=use_pre_softmax, - num_groups=-1 if num_groups is None else num_groups, - group_topk=-1 if group_topk is None else group_topk, - scaling_factor=scaling_factor, - score_function=score_function, - expert_bias=eb_arg, - compute_aux_scores=False, - ) - sparse_probs = sparse_probs.astype(dtype) - - # ---------------- Aux loss (global view, replicated) ---------------- - # ``fused_moe_aux_loss_fwd`` sums probs and tokens_per_expert across - # all tokens, which is wrong when T is sharded. Force-replicate the - # gate logits and recompute the routing map at global view so the - # kernel sees a complete [T_global, E] tensor. The replication is a - # single all-gather over (*dp, ep) and lives off the dispatch - # critical path. - if aux_loss_coeff > 0.0: - global_logits_2d = jax.lax.with_sharding_constraint(logits_2d, NamedSharding(mesh, P())) - _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( - global_logits_2d, + 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)) + + # ---------------- Gate (global view) ---------------- + # tex.fused_topk_with_score_function is only validated against its + # pytorch reference at fp32 (see tests/pytorch/test_fused_router.py: + # parametrize gates dtype on torch.float32 only; the tolerance helper + # raises NotImplementedError for any other dtype). Keeping logits in + # the activation dtype (e.g. bf16) lets sigmoid / softmax / topk + # accumulate at low precision and silently produce NaNs on tokens + # whose normalised weights underflow. Cast to fp32 here to stay in + # the validated regime. + gate_kernel_cast = gate_kernel.astype(x.dtype) + gate_logits = jnp.einsum("bsh,he->bse", x, gate_kernel_cast) + logits_2d = gate_logits.reshape(-1, num_experts).astype(jnp.float32) + + # ---------------- Routing (global view) ---------------- + # expert_bias is an empty (shape-(0,)) sentinel when the caller did + # not enable it; the primitive treats that as "no bias". + eb_arg = expert_bias if expert_bias.shape != (0,) else jnp.zeros((0,), dtype=jnp.float32) + sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( + logits_2d, topk=K, use_pre_softmax=use_pre_softmax, num_groups=-1 if num_groups is None else num_groups, @@ -705,190 +1311,228 @@ def _moe_fwd_rule( expert_bias=eb_arg, compute_aux_scores=False, ) - aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) - # compute_aux_scores=True takes a separate kernel path: clean - # per-expert softmax, no grouping / bias / scaling. - aux_probs, _aux_rm, aux_saved_scores = tex.fused_topk_with_score_function_fwd( - global_logits_2d.astype(jnp.float32), - topk=K, - use_pre_softmax=False, - num_groups=-1, - group_topk=-1, - scaling_factor=1.0, - score_function=score_function, - expert_bias=jnp.zeros((0,), dtype=jnp.float32), - compute_aux_scores=True, + sparse_probs = sparse_probs.astype(dtype) + + # ---------------- Aux loss (global view, replicated) ---------------- + # ``fused_moe_aux_loss_fwd`` sums probs and tokens_per_expert across + # all tokens, which is wrong when T is sharded. Force-replicate the + # gate logits and recompute the routing map at global view so the + # kernel sees a complete [T_global, E] tensor. The replication is a + # single all-gather over (*dp, ep) and lives off the dispatch + # critical path. + if aux_loss_coeff > 0.0: + global_logits_2d = jax.lax.with_sharding_constraint(logits_2d, NamedSharding(mesh, P())) + _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( + global_logits_2d, + topk=K, + use_pre_softmax=use_pre_softmax, + num_groups=-1 if num_groups is None else num_groups, + group_topk=-1 if group_topk is None else group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=eb_arg, + compute_aux_scores=False, + ) + aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) + # compute_aux_scores=True takes a separate kernel path: clean + # per-expert softmax, no grouping / bias / scaling. + aux_probs, _aux_rm, aux_saved_scores = tex.fused_topk_with_score_function_fwd( + global_logits_2d.astype(jnp.float32), + topk=K, + use_pre_softmax=False, + num_groups=-1, + group_topk=-1, + scaling_factor=1.0, + score_function=score_function, + expert_bias=jnp.zeros((0,), dtype=jnp.float32), + compute_aux_scores=True, + ) + aux_loss, aux_const_buf = tex.fused_moe_aux_loss_fwd( + aux_probs.astype(jnp.float32), + aux_tokens_per_expert.astype(jnp.int32), + topk=K, + coeff=aux_loss_coeff, + ) + aux_loss = aux_loss.astype(dtype) + else: + aux_loss = jnp.zeros((), dtype=dtype) + aux_const_buf = None + aux_tokens_per_expert = None + aux_saved_scores = None + + # ---------------- Routing -> (topk_idx, topk_w) at 3D ---------------- + # argsort on a bool tensor places True last (False=0 < True=1), so the + # last K indices are the selected expert IDs. + selected_experts = jnp.argsort(routing_map, axis=-1)[..., -K:] + routing_weights = jnp.take_along_axis(sparse_probs, selected_experts, axis=-1) + topk_idx_3d = selected_experts.reshape(B, S, K).astype(jnp.int32) + topk_w_3d = routing_weights.reshape(B, S, K).astype(jnp.float32) + # tex.ep_prepare/dispatch's partition only folds ep_axis into a replicated + # leading dim, not the outer dp/fsdp axes, so a replicated topk_idx makes + # each rank see B/ep rows (not B/num_procs) and overrun the bootstrap-sized + # send buffer. Pin both routing tensors to the (outer, ep) leading sharding + # so per-rank token counts match max_tokens_per_rank. + topk_idx_3d = jax.lax.with_sharding_constraint(topk_idx_3d, NamedSharding(mesh, ep3_spec)) + topk_w_3d = jax.lax.with_sharding_constraint(topk_w_3d, NamedSharding(mesh, ep3_spec)) + + # ---------------- TE EP dispatch (global view) ---------------- + cfg = tex.EpLayerConfig( + top_k=K, + dispatch_output_per_expert_alignment=dispatch_alignment, + topk_idx_dtype=int(jax_dtype_to_te_dtype(topk_idx_3d.dtype)), ) - aux_loss, aux_const_buf = tex.fused_moe_aux_loss_fwd( - aux_probs.astype(jnp.float32), - aux_tokens_per_expert.astype(jnp.int32), - topk=K, - coeff=aux_loss_coeff, + recv_tokens, recv_topk_weights, token_counts, total_recv_tokens, handle_mem = ( + tex.ep_prepare_and_dispatch_fwd(cfg, topk_idx_3d, x, topk_w_3d, recv_pr) + ) + token_counts = jax.lax.with_sharding_constraint(token_counts, NamedSharding(mesh, ep2_spec)) + if dispatch_checkpoint_name is not None: + recv_tokens = checkpoint_name(recv_tokens, dispatch_checkpoint_name) + recv_topk_weights = checkpoint_name(recv_topk_weights, dispatch_checkpoint_name) + recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3_spec)) + recv_topk_weights = jax.lax.with_sharding_constraint( + recv_topk_weights, NamedSharding(mesh, ep2_spec) ) - aux_loss = aux_loss.astype(dtype) - else: - aux_loss = jnp.zeros((), dtype=dtype) - aux_const_buf = None - aux_tokens_per_expert = None - aux_saved_scores = None - - # ---------------- Routing -> (topk_idx, topk_w) at 3D ---------------- - # argsort on a bool tensor places True last (False=0 < True=1), so the - # last K indices are the selected expert IDs. - selected_experts = jnp.argsort(routing_map, axis=-1)[..., -K:] - routing_weights = jnp.take_along_axis(sparse_probs, selected_experts, axis=-1) - topk_idx_3d = selected_experts.reshape(B, S, K).astype(jnp.int32) - topk_w_3d = routing_weights.reshape(B, S, K).astype(jnp.float32) - # tex.ep_prepare/dispatch's partition only folds ep_axis into a replicated - # leading dim, not the outer dp/fsdp axes, so a replicated topk_idx makes - # each rank see B/ep rows (not B/num_procs) and overrun the bootstrap-sized - # send buffer. Pin both routing tensors to the (outer, ep) leading sharding - # so per-rank token counts match max_tokens_per_rank. - topk_idx_3d = jax.lax.with_sharding_constraint(topk_idx_3d, NamedSharding(mesh, ep3_spec)) - topk_w_3d = jax.lax.with_sharding_constraint(topk_w_3d, NamedSharding(mesh, ep3_spec)) - - # ---------------- TE EP dispatch (global view) ---------------- - cfg = tex.EpLayerConfig( - top_k=K, - dispatch_output_per_expert_alignment=_ALIGN_SIZE, - topk_idx_dtype=int(jax_dtype_to_te_dtype(topk_idx_3d.dtype)), - ) - recv_tokens, recv_topk_weights, token_counts, total_recv_tokens, handle_mem = ( - tex.ep_prepare_and_dispatch_fwd(cfg, topk_idx_3d, x, topk_w_3d, recv_pr) - ) - token_counts = jax.lax.with_sharding_constraint(token_counts, NamedSharding(mesh, ep2_spec)) - recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3_spec)) - recv_topk_weights = jax.lax.with_sharding_constraint( - recv_topk_weights, NamedSharding(mesh, ep2_spec) - ) - - # ---------------- 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_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_args.extend([wi_0_bias, wi_1_bias, wo_bias]) - - # Quantized grouped tensors store their data, scales, and group metadata - # as physical buffers rather than in the source tensor's logical shape. - # A PartitionSpec used as a pytree prefix applies the same ownership to - # every array leaf of the grouped tensor: dispatched-token buffers belong - # to the compound batch shard, while expert-weight buffers belong to EP. - token_buffer_spec = P(batch_pspec_axis) - token_matrix_spec = P(batch_pspec_axis, None) - expert_buffer_spec = P(ep_axis) - residuals_spec = ( - token_buffer_spec, - expert_buffer_spec, - token_matrix_spec, - token_matrix_spec, - token_buffer_spec, - expert_buffer_spec, - ep2_spec, - ) - def _ffn_fwd_body(*args): + # ---------------- FFN (per-shard via shard_map) ---------------- + has_bias = wi_0_bias is not None + bias_spec = P(ep_axis, None) + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, wi_input_spec, wo_input_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi, wo] if has_bias: - r_tok, r_w, tc, local_wi, local_wo, w0b, w1b, wob = args - else: - r_tok, r_w, tc, local_wi, local_wo = args - w0b = w1b = wob = None - return _ffn_fwd_per_shard( - r_tok, - r_w, - tc, - local_wi, - local_wo, - w0b, - w1b, - wob, - quantizer_sets, - num_local_experts=num_local_experts, - activation_type=activation_type, - apply_topk_weights_early=apply_topk_weights_early, - wi_0_checkpoint_name=wi_0_checkpoint_name, - wi_1_checkpoint_name=wi_1_checkpoint_name, - wo_checkpoint_name=wo_checkpoint_name, + 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 + # as physical buffers rather than in the source tensor's logical shape. + # A PartitionSpec used as a pytree prefix applies the same ownership to + # every array leaf of the grouped tensor: dispatched-token buffers belong + # to the compound batch shard, while expert-weight buffers belong to EP. + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) + residuals_spec = ( + token_buffer_spec, + expert_buffer_spec, + token_matrix_spec, + token_matrix_spec, + token_buffer_spec, + expert_buffer_spec, + ep2_spec, ) - expert_outputs, ffn_residuals = shard_map( - _ffn_fwd_body, - mesh=mesh, - in_specs=ffn_in_specs, - out_specs=(ep3_spec, residuals_spec), - check_rep=False, - )(*ffn_in_args) - expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) - - # ---------------- TE EP combine (global view) ---------------- - out_partition_spec = (batch_pspec_axis, None, None) - if apply_topk_weights_early: - # expert_outputs is already weighted upstream. - output = tex.ep_combine_fwd( - cfg, - handle_mem, - expert_outputs, - num_local_tokens=(B, S), - out_partition_spec=out_partition_spec, - ) - else: - # HT combine is unweighted; apply routing weights before calling it. - # Padded recv slots are ignored by combine via handle_mem metadata. - w = recv_topk_weights[..., None].astype(expert_outputs.dtype) - weighted = expert_outputs * w - output = tex.ep_combine_fwd( - cfg, - handle_mem, - weighted, - num_local_tokens=(B, S), - out_partition_spec=out_partition_spec, + def _ffn_fwd_body(*args): + if has_bias: + r_tok, r_w, tc, local_wi, local_wo, w0b, w1b, wob = args + else: + r_tok, r_w, tc, local_wi, local_wo = args + w0b = w1b = wob = None + return _ffn_fwd_per_shard( + r_tok, + r_w, + tc, + local_wi, + local_wo, + w0b, + w1b, + wob, + quantizer_sets, + num_local_experts=num_local_experts, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + use_cudnn_jax_fusion=use_cudnn_jax_fusion, + wi_0_checkpoint_name=wi_0_checkpoint_name, + wi_1_checkpoint_name=wi_1_checkpoint_name, + wo_checkpoint_name=wo_checkpoint_name, + cudnn_native_weight_layout=cudnn_native_weight_layout, + quant_before_fsdp_ag=quant_before_fsdp_ag, + fsdp_axis=mesh_resource.fsdp_resource, + fsdp_size=( + mesh.shape[mesh_resource.fsdp_resource] + if mesh_resource.fsdp_resource is not None + else 1 + ), + wi_fsdp_axis=wi_fsdp_axis, + wo_fsdp_axis=wo_fsdp_axis, + ) + + expert_outputs, ffn_residuals = shard_map( + _ffn_fwd_body, + mesh=mesh, + in_specs=ffn_in_specs, + out_specs=(ep3_spec, residuals_spec), + check_rep=False, + )(*ffn_in_args) + expert_outputs = jax.lax.with_sharding_constraint( + expert_outputs, NamedSharding(mesh, ep3_spec) ) - # output of MLP should be sharded the same way as the activation input - output = with_sharding_constraint_by_logical_axes(output, input_axes) - ( - casted_sorted_x_lhs_trans, - casted_wi_rhs_trans, - gate_proj_out, - up_proj_out, - casted_intermediate_lhs_trans, - casted_wo_rhs_trans, - local_group_sizes, - ) = ffn_residuals - - ctx = _Ctx( - x=x, - gate_kernel=gate_kernel, - expert_bias=expert_bias, - logits_2d=logits_2d, - saved_scores=saved_scores, - routing_map=routing_map, - cfg=cfg, - handle_mem=handle_mem, - recv_topk_weights=recv_topk_weights, - casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, - casted_wi_rhs_trans=casted_wi_rhs_trans, - gate_proj_out=gate_proj_out, - up_proj_out=up_proj_out, - casted_intermediate_lhs_trans=casted_intermediate_lhs_trans, - casted_wo_rhs_trans=casted_wo_rhs_trans, - expert_outputs=expert_outputs, - local_group_sizes=local_group_sizes, - quantizer_sets=quantizer_sets, - aux_const_buf=aux_const_buf, - aux_tokens_per_expert=aux_tokens_per_expert, - aux_saved_scores=aux_saved_scores, - ) - static = { - "has_bias": has_bias, - "x_shape": x.shape, - "recv_pr": recv_pr, - } - # total_recv_tokens is a non-differentiable overflow signal (see moe()). - return (output, aux_loss, total_recv_tokens), (ctx, static) + # ---------------- TE EP combine (global view) ---------------- + out_partition_spec = (batch_pspec_axis, None, None) + if apply_topk_weights_early: + # expert_outputs is already weighted upstream. + output = tex.ep_combine_fwd( + cfg, + handle_mem, + expert_outputs, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + else: + # HT combine is unweighted; apply routing weights before calling it. + # Padded recv slots are ignored by combine via handle_mem metadata. + w = recv_topk_weights[..., None].astype(expert_outputs.dtype) + weighted = expert_outputs * w + output = tex.ep_combine_fwd( + cfg, + handle_mem, + weighted, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + # output of MLP should be sharded the same way as the activation input + output = with_sharding_constraint_by_logical_axes(output, input_axes) + + ( + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out, + up_proj_out, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes, + ) = ffn_residuals + + ctx = _Ctx( + x=x, + gate_kernel=gate_kernel, + expert_bias=expert_bias, + logits_2d=logits_2d, + saved_scores=saved_scores, + routing_map=routing_map, + cfg=cfg, + handle_mem=handle_mem, + recv_topk_weights=recv_topk_weights, + casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, + casted_wi_rhs_trans=casted_wi_rhs_trans, + gate_proj_out=gate_proj_out, + up_proj_out=up_proj_out, + casted_intermediate_lhs_trans=casted_intermediate_lhs_trans, + casted_wo_rhs_trans=casted_wo_rhs_trans, + expert_outputs=expert_outputs, + local_group_sizes=local_group_sizes, + quantizer_sets=quantizer_sets, + aux_const_buf=aux_const_buf, + aux_tokens_per_expert=aux_tokens_per_expert, + aux_saved_scores=aux_saved_scores, + ) + static = { + "has_bias": has_bias, + "x_shape": x.shape, + "recv_pr": recv_pr, + "cudnn_native_weight_layout": cudnn_native_weight_layout, + } + # total_recv_tokens is a non-differentiable overflow signal (see moe()). + return (output, aux_loss, total_recv_tokens), (ctx, static) def _moe_bwd_rule( @@ -901,8 +1545,7 @@ def _moe_bwd_rule( group_topk, scaling_factor, aux_loss_coeff, - ep_axis, - data_parallelism_axes, + mesh_resource, input_axes, gate_kernel_axes, wi_kernel_axes, @@ -910,257 +1553,282 @@ def _moe_bwd_rule( dtype, apply_topk_weights_early, recv_capacity_per_rank, + use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, wo_checkpoint_name, + dispatch_checkpoint_name, + quant_before_fsdp_ag, + use_cudnn_fusion, + cudnn_native_weight_layout, residuals, cotangents, ): """Backward mirror of :func:`_moe_fwd_rule`.""" - del ( - num_groups, - group_topk, - dtype, - recv_capacity_per_rank, - wi_0_checkpoint_name, - wi_1_checkpoint_name, - wo_checkpoint_name, - ) # captured / unused in bwd - from jax.experimental.shard_map import shard_map - - # total_recv_tokens is a non-differentiable output; its cotangent is unused. - d_output, d_aux_loss, _d_total_recv_tokens = cotangents - - ctx, static = residuals - has_bias = static["has_bias"] - x_shape = static["x_shape"] - recv_pr = static["recv_pr"] - - mesh = _get_mesh() - if mesh is None or mesh.empty: - 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) - 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 = jax.lax.with_sharding_constraint( - grad_pre_combine, NamedSharding(mesh, ep3_spec) - ) - if apply_topk_weights_early: - # combine_fwd consumed already-weighted expert_outputs; the recv_w - # cotangent flows through the early-weighting step inside the FFN bwd. - d_expert_outputs = grad_pre_combine - d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) - else: - w = ctx.recv_topk_weights[..., None].astype(grad_pre_combine.dtype) - d_expert_outputs = grad_pre_combine * w - d_recv_w_from_combine = (grad_pre_combine * ctx.expert_outputs).sum(axis=-1) - 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) - token_buffer_spec = P(batch_pspec_axis) - token_matrix_spec = P(batch_pspec_axis, None) - expert_buffer_spec = P(ep_axis) - residuals_specs = ( - token_buffer_spec, - expert_buffer_spec, - token_matrix_spec, - token_matrix_spec, - token_buffer_spec, - expert_buffer_spec, - ep2_spec, - ) - bwd_in_specs = (ep3_spec, *residuals_specs, ep2_spec) - bwd_in_args = [ - d_expert_outputs, - ctx.casted_sorted_x_lhs_trans, - ctx.casted_wi_rhs_trans, - ctx.gate_proj_out, - ctx.up_proj_out, - ctx.casted_intermediate_lhs_trans, - ctx.casted_wo_rhs_trans, - ctx.local_group_sizes, - ctx.recv_topk_weights, - ] - - def _ffn_bwd_body(*args): - grads = _ffn_bwd_per_shard( - *args, - ctx.quantizer_sets, - activation_type=activation_type, - apply_topk_weights_early=apply_topk_weights_early, - has_bias=has_bias, - ) - ( - d_sorted_x_local, - d_recv_w_local, - d_wi_local, - d_wo_local, - d_wi_0_bias_local, - d_wi_1_bias_local, - d_wo_bias_local, - ) = grads - if data_parallelism_axes: - dp_axes = tuple(data_parallelism_axes) - d_wi_local = jax.lax.psum(d_wi_local, axis_name=dp_axes) - d_wo_local = jax.lax.psum(d_wo_local, axis_name=dp_axes) - if has_bias: - d_wi_0_bias_local = jax.lax.psum(d_wi_0_bias_local, axis_name=dp_axes) - d_wi_1_bias_local = jax.lax.psum(d_wi_1_bias_local, axis_name=dp_axes) - d_wo_bias_local = jax.lax.psum(d_wo_bias_local, axis_name=dp_axes) - return ( - d_sorted_x_local, - d_recv_w_local, - d_wi_local, - d_wo_local, - d_wi_0_bias_local, - d_wi_1_bias_local, - d_wo_bias_local, + with global_shard_guard(mesh_resource): + ep_axis, data_parallelism_axes = _moe_mesh_axes(mesh_resource) + del ( + num_groups, + group_topk, + dtype, + recv_capacity_per_rank, + wi_0_checkpoint_name, + wi_1_checkpoint_name, + wo_checkpoint_name, + dispatch_checkpoint_name, + quant_before_fsdp_ag, + use_cudnn_fusion, + cudnn_native_weight_layout, + ) # captured / unused in bwd + from jax.experimental.shard_map import shard_map + + # total_recv_tokens is a non-differentiable output; its cotangent is unused. + d_output, d_aux_loss, _d_total_recv_tokens = cotangents + + ctx, static = residuals + has_bias = static["has_bias"] + x_shape = static["x_shape"] + recv_pr = static["recv_pr"] + + mesh = _get_mesh() + if mesh is None or mesh.empty: + 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) + 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 = jax.lax.with_sharding_constraint( + grad_pre_combine, NamedSharding(mesh, ep3_spec) ) - - if has_bias: - bwd_out_specs = ( - ep3_spec, + if apply_topk_weights_early: + # combine_fwd consumed already-weighted expert_outputs; the recv_w + # cotangent flows through the early-weighting step inside the FFN bwd. + d_expert_outputs = grad_pre_combine + d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) + else: + w = ctx.recv_topk_weights[..., None].astype(grad_pre_combine.dtype) + d_expert_outputs = grad_pre_combine * w + d_recv_w_from_combine = (grad_pre_combine * ctx.expert_outputs).sum(axis=-1) + 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) + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) + residuals_specs = ( + token_buffer_spec, + expert_buffer_spec, + token_matrix_spec, + token_matrix_spec, + token_buffer_spec, + expert_buffer_spec, ep2_spec, - kernel_spec, - kernel_spec, - bias_spec, - bias_spec, - bias_spec, ) - else: - bwd_out_specs = (ep3_spec, ep2_spec, kernel_spec, kernel_spec, None, None, None) + bwd_in_specs = (ep3_spec, *residuals_specs, ep2_spec) + bwd_in_args = [ + d_expert_outputs, + ctx.casted_sorted_x_lhs_trans, + ctx.casted_wi_rhs_trans, + ctx.gate_proj_out, + ctx.up_proj_out, + ctx.casted_intermediate_lhs_trans, + ctx.casted_wo_rhs_trans, + ctx.local_group_sizes, + ctx.recv_topk_weights, + ] + + def _ffn_bwd_body(*args): + grads = _ffn_bwd_per_shard( + *args, + ctx.quantizer_sets, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + has_bias=has_bias, + use_cudnn_jax_fusion=use_cudnn_jax_fusion, + cudnn_native_weight_layout=static["cudnn_native_weight_layout"], + ) + ( + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, + ) = grads + if data_parallelism_axes: + dp_axes = tuple(data_parallelism_axes) + d_wi_local = jax.lax.psum(d_wi_local, axis_name=dp_axes) + d_wo_local = jax.lax.psum(d_wo_local, axis_name=dp_axes) + if has_bias: + d_wi_0_bias_local = jax.lax.psum(d_wi_0_bias_local, axis_name=dp_axes) + d_wi_1_bias_local = jax.lax.psum(d_wi_1_bias_local, axis_name=dp_axes) + d_wo_bias_local = jax.lax.psum(d_wo_bias_local, axis_name=dp_axes) + return ( + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, + ) - ( - d_sorted_x, - d_recv_w_from_intermediate, - d_wi, - d_wo, - d_wi_0_bias, - d_wi_1_bias, - d_wo_bias, - ) = shard_map( - _ffn_bwd_body, - mesh=mesh, - in_specs=bwd_in_specs, - out_specs=bwd_out_specs, - check_rep=False, - )( - *bwd_in_args - ) + if has_bias: + bwd_out_specs = ( + ep3_spec, + ep2_spec, + kernel_spec, + kernel_spec, + bias_spec, + bias_spec, + bias_spec, + ) + else: + bwd_out_specs = ( + ep3_spec, + ep2_spec, + kernel_spec, + kernel_spec, + None, + None, + None, + ) - d_recv_w_total = d_recv_w_from_combine + d_recv_w_from_intermediate - - # ---------------- Dispatch bwd (global view) ---------------- - d_sorted_x = jax.lax.with_sharding_constraint(d_sorted_x, NamedSharding(mesh, ep3_spec)) - d_recv_w_total = jax.lax.with_sharding_constraint(d_recv_w_total, NamedSharding(mesh, ep2_spec)) - d_x_from_dispatch, d_topk_w = tex.ep_dispatch_bwd( - ctx.cfg, - ctx.handle_mem, - d_sorted_x, - d_recv_w_total, - num_local_tokens=(B, S), - out_partition_spec=out_partition_spec, - ) + ( + d_sorted_x, + d_recv_w_from_intermediate, + d_wi, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = shard_map( + _ffn_bwd_body, + mesh=mesh, + in_specs=bwd_in_specs, + out_specs=bwd_out_specs, + check_rep=False, + )( + *bwd_in_args + ) - # ---------------- Routing bwd (global view) ---------------- - # The cotangent on routing_weights is a sparse scatter into sparse_probs - # at the selected_experts indices. - selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] - d_topk_w_flat = d_topk_w.reshape(-1, K) - d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_topk_w_flat.dtype) - d_sparse_probs = d_sparse_probs.at[ - jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts - ].set(d_topk_w_flat) - - d_logits_2d = tex.fused_topk_with_score_function_bwd( - ctx.routing_map, - ctx.saved_scores, - d_sparse_probs.astype(ctx.saved_scores.dtype), - topk=K, - use_pre_softmax=use_pre_softmax, - scaling_factor=scaling_factor, - score_function=score_function, - compute_aux_scores=False, - ) + d_recv_w_total = d_recv_w_from_combine + d_recv_w_from_intermediate - # ---------------- Aux loss bwd (global view, replicated) ---------------- - # Reverse the fwd's all-gather/aux pipeline: aux_loss_bwd produces - # d_aux_probs, then topk_bwd(compute_aux_scores=True) produces the - # extra d_logits contribution. The replicated tensor adds into the - # T-sharded routing-side d_logits via JAX's normal broadcast. - if aux_loss_coeff > 0.0: - T_global = ctx.logits_2d.shape[0] - d_aux_loss_scalar = d_aux_loss.reshape(()).astype(jnp.float32) - d_aux_probs = tex.fused_moe_aux_loss_bwd( - ctx.aux_const_buf, - ctx.aux_tokens_per_expert.astype(jnp.int32), - d_aux_loss_scalar, - num_tokens=int(T_global), + # ---------------- Dispatch bwd (global view) ---------------- + d_sorted_x = jax.lax.with_sharding_constraint(d_sorted_x, NamedSharding(mesh, ep3_spec)) + d_recv_w_total = jax.lax.with_sharding_constraint( + d_recv_w_total, NamedSharding(mesh, ep2_spec) + ) + d_x_from_dispatch, d_topk_w = tex.ep_dispatch_bwd( + ctx.cfg, + ctx.handle_mem, + d_sorted_x, + d_recv_w_total, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, ) - # routing_map is ignored by the kernel when compute_aux_scores=True, - # so pass a zero placeholder of the right shape/dtype. - zero_routing_map = jnp.zeros(ctx.aux_saved_scores.shape, dtype=ctx.routing_map.dtype) - d_logits_aux = tex.fused_topk_with_score_function_bwd( - zero_routing_map, - ctx.aux_saved_scores, - d_aux_probs.astype(ctx.aux_saved_scores.dtype), + + # ---------------- Routing bwd (global view) ---------------- + # The cotangent on routing_weights is a sparse scatter into sparse_probs + # at the selected_experts indices. + selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] + d_topk_w_flat = d_topk_w.reshape(-1, K) + d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_topk_w_flat.dtype) + d_sparse_probs = d_sparse_probs.at[ + jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts + ].set(d_topk_w_flat) + + d_logits_2d = tex.fused_topk_with_score_function_bwd( + ctx.routing_map, + ctx.saved_scores, + d_sparse_probs.astype(ctx.saved_scores.dtype), topk=K, - use_pre_softmax=False, - scaling_factor=1.0, + use_pre_softmax=use_pre_softmax, + scaling_factor=scaling_factor, score_function=score_function, - compute_aux_scores=True, + compute_aux_scores=False, ) - d_logits_2d = d_logits_2d + d_logits_aux.astype(d_logits_2d.dtype) - - # ---------------- Gate bwd (global view) ---------------- - d_gate_logits = d_logits_2d.reshape(B, S, num_experts) - gate_kernel_cast = ctx.gate_kernel.astype(ctx.x.dtype) - d_x_from_gate = jnp.einsum("bse,he->bsh", d_gate_logits, gate_kernel_cast) - d_gate_kernel = jnp.einsum("bsh,bse->he", ctx.x, d_gate_logits).astype(ctx.gate_kernel.dtype) - d_x = d_x_from_gate + d_x_from_dispatch - - # Pin output grads to the declared logical axes so downstream - # optimizers see consistent shardings. - d_x = with_sharding_constraint_by_logical_axes(d_x, input_axes) - d_gate_kernel = with_sharding_constraint_by_logical_axes(d_gate_kernel, gate_kernel_axes) - d_wi = with_sharding_constraint_by_logical_axes(d_wi, wi_kernel_axes) - d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_axes) - if has_bias: - wi_bias_axes = (wi_kernel_axes[0], *wi_kernel_axes[2:]) - wo_bias_axes = (wo_kernel_axes[0], *wo_kernel_axes[2:]) - d_wi_0_bias = with_sharding_constraint_by_logical_axes(d_wi_0_bias, wi_bias_axes) - d_wi_1_bias = with_sharding_constraint_by_logical_axes(d_wi_1_bias, wi_bias_axes) - d_wo_bias = with_sharding_constraint_by_logical_axes(d_wo_bias, wo_bias_axes) - - # expert_bias has no learnable bwd path through fused_topk: the - # primitive's bwd returns None for the bias slot. Match that with a - # zero cotangent of the right shape so custom_vjp's arity check - # passes. - d_expert_bias = jnp.zeros_like(ctx.expert_bias) - return ( - d_x, - d_gate_kernel, - d_wi, - d_wo, - d_wi_0_bias if has_bias else None, - d_wi_1_bias if has_bias else None, - d_wo_bias if has_bias else None, - d_expert_bias, - ctx.quantizer_sets, - ) + # ---------------- Aux loss bwd (global view, replicated) ---------------- + # Reverse the fwd's all-gather/aux pipeline: aux_loss_bwd produces + # d_aux_probs, then topk_bwd(compute_aux_scores=True) produces the + # extra d_logits contribution. The replicated tensor adds into the + # T-sharded routing-side d_logits via JAX's normal broadcast. + if aux_loss_coeff > 0.0: + T_global = ctx.logits_2d.shape[0] + d_aux_loss_scalar = d_aux_loss.reshape(()).astype(jnp.float32) + d_aux_probs = tex.fused_moe_aux_loss_bwd( + ctx.aux_const_buf, + ctx.aux_tokens_per_expert.astype(jnp.int32), + d_aux_loss_scalar, + num_tokens=int(T_global), + ) + # routing_map is ignored by the kernel when compute_aux_scores=True, + # so pass a zero placeholder of the right shape/dtype. + zero_routing_map = jnp.zeros(ctx.aux_saved_scores.shape, dtype=ctx.routing_map.dtype) + d_logits_aux = tex.fused_topk_with_score_function_bwd( + zero_routing_map, + ctx.aux_saved_scores, + d_aux_probs.astype(ctx.aux_saved_scores.dtype), + topk=K, + use_pre_softmax=False, + scaling_factor=1.0, + score_function=score_function, + compute_aux_scores=True, + ) + d_logits_2d = d_logits_2d + d_logits_aux.astype(d_logits_2d.dtype) + + # ---------------- Gate bwd (global view) ---------------- + d_gate_logits = d_logits_2d.reshape(B, S, num_experts) + gate_kernel_cast = ctx.gate_kernel.astype(ctx.x.dtype) + d_x_from_gate = jnp.einsum("bse,he->bsh", d_gate_logits, gate_kernel_cast) + d_gate_kernel = jnp.einsum("bsh,bse->he", ctx.x, d_gate_logits).astype( + ctx.gate_kernel.dtype + ) + d_x = d_x_from_gate + d_x_from_dispatch + + # Pin output grads to the declared logical axes so downstream + # optimizers see consistent shardings. + d_x = with_sharding_constraint_by_logical_axes(d_x, input_axes) + d_gate_kernel = with_sharding_constraint_by_logical_axes(d_gate_kernel, gate_kernel_axes) + d_wi = with_sharding_constraint_by_logical_axes(d_wi, wi_kernel_axes) + d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_axes) + if has_bias: + wi_bias_axes = (wi_kernel_axes[0], *wi_kernel_axes[2:]) + wo_bias_axes = (wo_kernel_axes[0], *wo_kernel_axes[2:]) + d_wi_0_bias = with_sharding_constraint_by_logical_axes(d_wi_0_bias, wi_bias_axes) + d_wi_1_bias = with_sharding_constraint_by_logical_axes(d_wi_1_bias, wi_bias_axes) + d_wo_bias = with_sharding_constraint_by_logical_axes(d_wo_bias, wo_bias_axes) + + # expert_bias has no learnable bwd path through fused_topk: the + # primitive's bwd returns None for the bias slot. Match that with a + # zero cotangent of the right shape so custom_vjp's arity check + # passes. + d_expert_bias = jnp.zeros_like(ctx.expert_bias) + + return ( + d_x, + d_gate_kernel, + d_wi, + d_wo, + d_wi_0_bias if has_bias else None, + d_wi_1_bias if has_bias else None, + d_wo_bias if has_bias else None, + d_expert_bias, + ctx.quantizer_sets, + ) # ============================================================================= @@ -1168,7 +1836,7 @@ def _ffn_bwd_body(*args): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 30))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 34))) def _moe( x, gate_kernel, @@ -1188,8 +1856,7 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - ep_axis, - data_parallelism_axes, + mesh_resource, input_axes, gate_kernel_axes, wi_kernel_axes, @@ -1197,9 +1864,14 @@ def _moe( dtype, apply_topk_weights_early, recv_capacity_per_rank, + use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, wo_checkpoint_name, + dispatch_checkpoint_name, + quant_before_fsdp_ag, + use_cudnn_fusion: bool = True, + cudnn_native_weight_layout: Optional[bool] = None, ): primal, _ = _moe_fwd_rule( x, @@ -1220,8 +1892,7 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - ep_axis, - data_parallelism_axes, + mesh_resource, input_axes, gate_kernel_axes, wi_kernel_axes, @@ -1229,9 +1900,14 @@ def _moe( dtype, apply_topk_weights_early, recv_capacity_per_rank, + use_cudnn_jax_fusion, wi_0_checkpoint_name, wi_1_checkpoint_name, wo_checkpoint_name, + dispatch_checkpoint_name, + quant_before_fsdp_ag, + use_cudnn_fusion, + cudnn_native_weight_layout, ) return primal @@ -1263,8 +1939,8 @@ def moe( noop_quantizer_set, noop_quantizer_set, ), - ep_axis: str, - data_parallelism_axes: Tuple[str, ...] = (), + mesh_resource: Optional[MeshResource] = None, + quant_before_fsdp_ag: bool = False, input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), @@ -1274,6 +1950,9 @@ def moe( wi_0_checkpoint_name: Optional[str] = None, wi_1_checkpoint_name: Optional[str] = None, wo_checkpoint_name: Optional[str] = None, + dispatch_checkpoint_name: Optional[str] = None, + use_cudnn_fusion: bool = True, + cudnn_native_weight_layout: Optional[bool] = None, ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: """Run a full MoE block under a single fused custom_vjp on the TE EP path. @@ -1310,109 +1989,162 @@ def moe( ``total_recv_tokens`` when bootstrap used ``drop_on_overflow=True``. wi_0_checkpoint_name : Optional[str] JAX rematerialization checkpoint name for the gate projection output. + With cuDNN fusion, also names the shared fused outputs, including its quantized output. ``None`` leaves the value unnamed. wi_1_checkpoint_name : Optional[str] JAX rematerialization checkpoint name for the up projection output. + With cuDNN fusion, also names the same fused outputs, including its quantized output. ``None`` leaves the value unnamed. wo_checkpoint_name : Optional[str] JAX rematerialization checkpoint name for the per-expert down projection output. ``None`` leaves the value unnamed. - - Note that the per-expert dispatch-slot alignment is fixed internally - at 128 tokens (``_ALIGN_SIZE``); see that constant's docstring for - rationale and how to extend if a future recipe needs >128. - - 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 - ``P((dp..., ep), None, None)`` for the physical - ``jax.lax.with_sharding_constraint`` calls that JAX requires - to refer to real mesh axes). - * ``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, - ``wo_kernel_axes`` are *logical axis names* (e.g. - ``"batch"``, ``"embed"``, ``"mlp"``, ``"exp"``) -- they get - resolved via the active Flax logical-axis rules and consumed - by ``with_sharding_constraint_by_logical_axes``. They are - ``Optional[str]`` tuples so a rule of ``None`` means - "replicated on this axis". - - Logical-axis support for ``ep_axis`` / ``data_parallelism_axes`` - is intentionally out of scope: the EP comm-group construction - (``dp_color = rank // ep_size``) and the bootstrap signature - check both require concrete integer sizes, so a logical name - would have to be resolved to a physical one anyway before any - EP primitive is called. If a downstream pipeline needs to plumb - logical names all the way to ``moe()``, do the rule lookup at - the call site. + dispatch_checkpoint_name : Optional[str] + JAX rematerialization checkpoint name for the EP dispatch outputs. + ``None`` leaves these values unnamed; the opaque prepare handle is never named. + mesh_resource : Optional[MeshResource] + Physical parallelism resources. ``None`` uses the active + ``global_shard_guard`` context; an explicit resource takes precedence. + One of these is required. Batch sharding combines ``dp_resource`` then + ``fsdp_resource`` (omitting unset/duplicate axes), with ``ep_resource`` + innermost. EP must be set and distinct from the outer axes. + quant_before_fsdp_ag : bool + Quantize expert-weight shards before gathering their MXFP8 data and + scales on ``mesh_resource.fsdp_resource``. Default ``False`` gathers + full-precision weights first. ``True`` requires an FSDP resource and + MXFP8 kernel quantizers. With active Flax logical-axis rules, the weight + input specs follow ``wi_kernel_axes`` and ``wo_kernel_axes``; FSDP may + shard whole experts or a matrix dimension within each expert. When + FSDP shards the gated dimension, layout conversion permutes gathered + FP8 data and inverse scales, preserving global gate/up pairing without + gathering or requantizing full-precision weights. + Per-expert dispatch-slot alignment defaults to 128 tokens (``_ALIGN_SIZE``). + Requesting cuDNN fusion reserves 256 tokens, also when falling back, to + preserve compatibility with EP bootstrap buffer sizing. + + use_cudnn_fusion : bool + Defaults to ``True``: try cuDNN's shared JAX grouped MXFP8 GLU/dGLU APIs + on SM100+ GPUs; cuDNN selects the architecture-specific kernels. ``False`` uses unfused TE grouped GEMM. + Requires cuDNN Frontend >= 1.32.0 with compatible JAX APIs. + Frontends whose JAX GLU/dGLU entry points only support BF16 discrete + weights fall back to unfused TE grouped GEMM. + Ineligible calls warn and fall back to TE's regular grouped-GEMM + implementation. API signatures and GPU capability determine support. + + cudnn_native_weight_layout : Optional[bool] + ``False`` selects standard contiguous gate/up [E,K,2N] storage; + ``True`` selects native [E,2N,K] storage with alternating 32-column + gate/up blocks. Defaults to ``None``, inferring storage from the global + shape. Ambiguous square shapes require an explicit flag. This option + describes parameter storage independently of the selected execution path. + + MeshResource fields name physical mesh axes, not Flax logical axes. + ``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes`` and + ``wo_kernel_axes`` remain logical-axis tuples resolved through the active + Flax rules. The selected MeshResource is also used by internal TE sharding + in both forward and backward, independently of an enclosing global context. See module docstring for the rest of the parameter semantics and the surrounding design rationale. """ + if not isinstance(use_cudnn_fusion, bool): + raise TypeError("use_cudnn_fusion must be a bool") + mesh_resource, quant_before_fsdp_ag = _resolve_moe_mesh_resource( + mesh_resource, quant_before_fsdp_ag + ) + ep_axis, data_parallelism_axes = _moe_mesh_axes(mesh_resource) score_function = _validate_score_function(score_function) - # Enforce ((outer_dp..., ep), None, None) on inbound activations. The - # EP comm groups consecutive global ranks (dp_color = rank // ep_size), - # so ep MUST be innermost in the partition spec. Soft re-pin: free if - # upstream already matches, single reshard otherwise. - 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 - 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): - warnings.warn( - f"moe(...): inbound x sharding {actual_spec} does not match expected " - f"{expected_spec}; inserting a reshard. Apply " - "jax.lax.with_sharding_constraint upstream to avoid this overhead.", - UserWarning, - stacklevel=2, + with global_shard_guard(mesh_resource): + # Enforce ((outer_dp..., ep), None, None) on inbound activations. The + # EP comm groups consecutive global ranks (dp_color = rank // ep_size), + # so ep MUST be innermost in the partition spec. Soft re-pin: free if + # upstream already matches, single reshard otherwise. + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + for axis in (ep_axis, *data_parallelism_axes): + if axis not in mesh.shape: + raise ValueError(f"TE MoE resource axis {axis!r} is not in the active mesh.") + expected_leading: Any = ( + (*data_parallelism_axes, ep_axis) if data_parallelism_axes else ep_axis ) - x = _with_sharding_constraint_cast_bwd(x, NamedSharding(mesh, expected_spec)) + 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): + warnings.warn( + f"moe(...): inbound x sharding {actual_spec} does not match expected " + f"{expected_spec}; inserting a reshard. Apply " + "jax.lax.with_sharding_constraint upstream to avoid this overhead.", + UserWarning, + stacklevel=2, + ) + x = _with_sharding_constraint_cast_bwd(x, NamedSharding(mesh, expected_spec)) - # custom_vjp can't trace through None args; lower expert_bias to an - # empty shape-(0,) tensor that fused_topk_with_score_function treats - # as "no bias". - if expert_bias is None: - expert_bias_arg = jnp.zeros((0,), dtype=jnp.float32) - else: - expert_bias_arg = expert_bias.astype(jnp.float32) + cudnn_native_weight_layout = _resolve_cudnn_native_weight_layout( + wi, x.shape[-1], cudnn_native_weight_layout + ) - output, aux_loss, total_recv_tokens = _moe( - x, - gate_kernel, - wi, - wo, - wi_0_bias, - wi_1_bias, - wo_bias, - expert_bias_arg, - quantizer_sets, - num_experts, - num_experts_per_tok, - activation_type, - score_function, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - float(aux_loss_coeff), - ep_axis, - data_parallelism_axes, - input_axes, - gate_kernel_axes, - wi_kernel_axes, - wo_kernel_axes, - dtype, - apply_topk_weights_early, - recv_capacity_per_rank, - wi_0_checkpoint_name, - wi_1_checkpoint_name, - wo_checkpoint_name, - ) - if aux_loss_coeff <= 0.0: - aux_loss = None - assert output.dtype == x.dtype, f"moe() output dtype {output.dtype} != input dtype {x.dtype}" - return output, aux_loss, total_recv_tokens + # custom_vjp can't trace through None args; lower expert_bias to an + # empty shape-(0,) tensor that fused_topk_with_score_function treats + # as "no bias". + if expert_bias is None: + expert_bias_arg = jnp.zeros((0,), dtype=jnp.float32) + else: + expert_bias_arg = expert_bias.astype(jnp.float32) + + use_cudnn_jax_fusion = False + if use_cudnn_fusion: + rejection_reasons = _cudnn_jax_fusion_rejection_reasons( + x, + wi, + wi_0_bias, + wi_1_bias, + quantizer_sets, + num_experts=num_experts, + activation_type=activation_type, + ep_axis=ep_axis, + ) + use_cudnn_jax_fusion = _select_cudnn_jax_fusion(rejection_reasons) + + output, aux_loss, total_recv_tokens = _moe( + x, + gate_kernel, + wi, + wo, + wi_0_bias, + wi_1_bias, + wo_bias, + expert_bias_arg, + quantizer_sets, + num_experts, + num_experts_per_tok, + activation_type, + score_function, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + float(aux_loss_coeff), + mesh_resource, + input_axes, + gate_kernel_axes, + wi_kernel_axes, + wo_kernel_axes, + dtype, + apply_topk_weights_early, + recv_capacity_per_rank, + use_cudnn_jax_fusion, + wi_0_checkpoint_name, + wi_1_checkpoint_name, + wo_checkpoint_name, + dispatch_checkpoint_name, + quant_before_fsdp_ag, + use_cudnn_fusion, + cudnn_native_weight_layout, + ) + if aux_loss_coeff <= 0.0: + aux_loss = None + assert ( + output.dtype == x.dtype + ), f"moe() output dtype {output.dtype} != input dtype {x.dtype}" + return output, aux_loss, total_recv_tokens