From cba9717fedc1f3351e6f77fa0b280a794318454c Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 10 Jun 2026 14:58:09 -0700 Subject: [PATCH 01/23] [JAX] Resync onto upstream PR #3036, restore TE-EP-only MoE block Reset 33 local commits onto phuong/ep-3-jax @ c34771d4 (her latest with EpConfig + EpLayerConfig API, NCCL bumped to 808d2433) and re-applied the three deltas uniquely ours: * transformer_engine/jax/moe.py: replaces upstream's multi-backend MoE block with our TE-EP-only single-custom-vjp rewrite. Adapted to her new API surface: tex.EpLayerConfig replaces tex.ep_make_handle (no more EpHandle pool/cache); 5 EP callsites rewired (cfg passed in place of handle, ep_prepare arg order swapped, top_k= dropped from ep_dispatch_bwd since it's now in cfg. * tests/jax/test_te_ep_moe.py: TE-EP MoE test (kept), with ep_bootstrap kwargs ep_size= and allow_handle_mem_reloc= dropped (no longer supported; ep_size is derived from mesh axes and the handle_mem reloc gating is gone). * tests/jax/run_te_ep_moe.sh: multi-process launcher (kept). Pre-sync state preserved at branch teddy/te_ep_integration.backup-pre-phuong-sync. EOF ) Signed-off-by: Teddy Do --- tests/jax/run_te_ep_moe.sh | 122 ++ tests/jax/test_te_ep_moe.py | 813 ++++++++++ transformer_engine/jax/moe.py | 2879 ++++++++++++--------------------- 3 files changed, 1973 insertions(+), 1841 deletions(-) create mode 100755 tests/jax/run_te_ep_moe.sh create mode 100644 tests/jax/test_te_ep_moe.py diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh new file mode 100755 index 0000000000..32d5f21956 --- /dev/null +++ b/tests/jax/run_te_ep_moe.sh @@ -0,0 +1,122 @@ +#!/usr/bin/env bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Multiprocess (one-GPU-per-process) launcher for the TE-EP MoE custom_vjp +# test suite. Forks one pytest invocation per visible GPU, passing each +# its own --num-process=N --process-id=i, and waits for all of them. Each +# child calls jax.distributed.initialize(..., local_device_ids=process_id) +# so each Python process only sees its one GPU as a local device and the +# participating processes form a global (ep, fsdp) mesh. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TE_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +TEST_FILE="$TE_ROOT/tests/jax/test_te_ep_moe.py" +PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" + +NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" +if [ "$NUM_GPUS" -lt 4 ]; then + echo "[run_te_ep_moe.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2 + exit 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}" + +echo "============================================================" +echo "TE-EP MoE MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" +echo " test file : $TEST_FILE" +echo " coordinator : $TE_EP_MOE_COORDINATOR_ADDRESS" +echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" +echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" +echo "============================================================" + +if [ -n "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then + LOG_DIR="$TE_EP_MOE_MP_LOG_DIR" + mkdir -p "$LOG_DIR" +else + LOG_DIR=$(mktemp -d -t te_ep_moe_mp_XXXXXX) +fi +echo "Per-process logs: $LOG_DIR" + +PIDS=() + +cleanup() { + for pid in "${PIDS[@]:-}"; do + if kill -0 "$pid" 2>/dev/null; then + kill -TERM "$pid" 2>/dev/null || true + fi + done + sleep 1 + for pid in "${PIDS[@]:-}"; do + if kill -0 "$pid" 2>/dev/null; then + kill -KILL "$pid" 2>/dev/null || true + fi + done +} +trap cleanup EXIT INT TERM + +for i in $(seq 0 $((NUM_GPUS - 1))); do + 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" & + else + "${PYTEST_CMD[@]}" > "$LOG_FILE" 2>&1 & + fi + PIDS+=("$!") +done + +EXITS=() +for pid in "${PIDS[@]}"; do + if wait "$pid"; then + EXITS+=("0") + else + EXITS+=("$?") + fi +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 + break + fi +done + +echo +if [ "$FAILED" -eq 0 ]; then + echo "[run_te_ep_moe.sh] all processes PASSED" + if [ -z "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then + rm -rf "$LOG_DIR" + fi + exit 0 +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_te_ep_moe.py b/tests/jax/test_te_ep_moe.py new file mode 100644 index 0000000000..a5ab1c266b --- /dev/null +++ b/tests/jax/test_te_ep_moe.py @@ -0,0 +1,813 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-process (one-GPU-per-process) tests for the TE-EP MoE custom_vjp. + +The launcher ``tests/jax/run_te_ep_moe.sh`` forks one pytest process per +visible GPU (mirroring ``run_multiprocess_moe_vjp.sh``). Each process binds +to exactly one device via +``jax.distributed.initialize(..., local_device_ids=process_id)``; the +participating processes form a global ``(ep, fsdp)`` mesh through JAX's +distributed runtime. + +How to run +---------- + +You typically do NOT invoke pytest on this file directly -- use the +launcher, which passes ``--num-process=N --process-id=i`` to each +forked process. Driving it directly with only one process will skip +every test because :func:`jax.distributed.initialize` requires +multiple participants, and the TE EP NCCL primitives require at +least four ranks. + + bash tests/jax/run_te_ep_moe.sh + +What this suite covers +---------------------- + +This file is the TE-EP-only successor to ``test_moe_vjp.py`` and +``test_multiprocess_moe_vjp.py``. Each test exercises one MoE-block +run and bundles every check that 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: + +* ``test_forward`` covers the forward across a curated set of + configurations (apply_topk_weights_early on/off, softmax/sigmoid + scoring, optional expert_bias). Each config asserts shape, dtype, + finiteness and numerical parity vs the reference in one run. +* ``test_backward`` mirrors that for gradients. +* ``TestTeEpMoeAuxLoss`` covers the second return value end-to-end + (returned + parity + aux-only grad propagates to gate + combined + main+aux grads stay finite) in two consolidated tests. +* ``TestTeEpMoEBlockFlax`` exercises the Flax wrapper with the same + parity reference. +* ``TestZZZTeEpMoeBootstrap`` verifies the per-process NCCL bootstrap + rejects a mismatched signature. + +FP8 / MXFP8 recipes are deferred — the ``quantizer_sets`` plumbing +has not yet been re-wired across the TE-EP ``shard_map`` boundary +(see ``.pr3036-review/INTEGRATION_DESIGN.md``). +""" + +import os + +os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") +os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") + +import sys +from functools import partial + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from jax.experimental import mesh_utils +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from flax.linen import partitioning as nn_partitioning + + +def _init_distributed(num_process: int, process_id: int) -> bool: + """Initialize jax.distributed for this pytest process. + + Returns True on a real multi-process launch, False otherwise so + the module can fast-skip when pytest collects it without the + launcher. + """ + if num_process <= 1: + return False + coord = os.environ.get("TE_EP_MOE_COORDINATOR_ADDRESS", "127.0.0.1:13457") + jax.distributed.initialize( + coordinator_address=coord, + num_processes=num_process, + process_id=process_id, + local_device_ids=process_id, + ) + assert jax.local_device_count() == 1, "one GPU per process is required for TE EP" + assert ( + jax.device_count() == num_process + ), f"global device_count {jax.device_count()} != num_process {num_process}" + return True + + +def _read_mp_options(): + num = int(os.environ.get("MP_NUM_PROCESS", "0") or "0") + pid = int(os.environ.get("MP_PROCESS_ID", "0") or "0") + for i, a in enumerate(sys.argv): + if a.startswith("--num-process="): + num = int(a.split("=", 1)[1]) + elif a == "--num-process" and i + 1 < len(sys.argv): + num = int(sys.argv[i + 1]) + elif a.startswith("--process-id="): + pid = int(a.split("=", 1)[1]) + elif a == "--process-id" and i + 1 < len(sys.argv): + pid = int(sys.argv[i + 1]) + return num, pid + + +_MP_NUM_PROCESS, _MP_PROCESS_ID = _read_mp_options() +_MP_ACTIVE = _init_distributed(_MP_NUM_PROCESS, _MP_PROCESS_ID) + +if not _MP_ACTIVE: + pytest.skip( + "test_te_ep_moe.py requires the multiprocess launcher " + "(run_te_ep_moe.sh). Skipping.", + allow_module_level=True, + ) + +from transformer_engine_jax import get_device_compute_capability + +# Grouped GEMM in the MoE custom_vjp requires Blackwell (sm_100+). The +# TE EP NCCL primitives themselves need SM>=90, but the FFN body uses +# grouped_gemm, so the file as a whole gates on sm_100+. +if get_device_compute_capability(0) < 100: + pytest.skip( + "MoE TE EP tests require Blackwell (sm_100+) for grouped GEMM", + allow_module_level=True, + ) + +from transformer_engine.jax.flax import _MoEBlock as MoEBlock +from transformer_engine.jax.moe import moe, record_ep_bootstrap_signature_for_moe +from transformer_engine.jax.ep import ep_bootstrap +from transformer_engine.jax.sharding import MeshResource, global_shard_guard + + +# ----------------------------------------------------------------------------- +# Mesh / shape config +# ----------------------------------------------------------------------------- + +EP_AXIS = "ep" +FSDP_AXIS = "fsdp" +EP_SIZE = 2 +assert ( + jax.device_count() % EP_SIZE == 0 +), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" +FSDP_SIZE = jax.device_count() // EP_SIZE +NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE + +LOGICAL_AXIS_RULES = ( + ("exp", EP_AXIS), + ("embed", FSDP_AXIS), + ("mlp", None), + ("batch", (EP_AXIS, FSDP_AXIS)), +) + +# Small shapes so the parity tests stay tight on bf16. The block still +# has all four ranks participating in dispatch/combine. +DTYPE = jnp.bfloat16 +BATCH = EP_SIZE * FSDP_SIZE * 2 # 8 on 4-GPU, 16 on 8-GPU +SEQ = 32 +HIDDEN = 64 +INTER = 128 +NUM_EXPERTS = 8 +TOPK = 2 + +# bf16 grouped_gemm + softmax-topk + ep all-to-all stack drifts ~1e-1 vs a +# fp32 numpy reference. Keep these tight enough to catch real bugs but +# loose enough to absorb expected bf16 rounding. +FWD_ATOL = 5e-2 +FWD_RTOL = 5e-2 +GRAD_FFN_ATOL = 1e-1 +GRAD_FFN_RTOL = 1e-1 +GRAD_GATE_ATOL = 5e-1 +GRAD_GATE_RTOL = 5e-1 + +# Two TE EP runs that should be bitwise-equal modulo XLA fusion order +# (align_size rounding, etc.). +TE_TO_TE_ATOL = 5e-3 +TE_TO_TE_RTOL = 5e-3 + +# Aux loss is computed in float32 from the SAME logits as the routing +# path. Numerical drift between TE-EP and the reference is dominated by +# the bf16-rounded softmax inside the topk kernel. +AUX_ATOL = 1e-3 +AUX_RTOL = 1e-3 + + +# ----------------------------------------------------------------------------- +# Fixtures +# ----------------------------------------------------------------------------- + + +def _compute_worst_case_recv_pr(): + """Per-rank recv buffer the bootstrap must reserve. + + NCCL EP's HT path lays out the per-rank receive buffer as + ``[num_local_experts, ep_size * max_tokens_per_rank, hidden]`` + (per the LL combine assertion at ``nccl_ep.cc:2185`` and the + HT IPC buffer sizing at ``nccl_ep.cc:415``). We must mirror that + flattened total or ``ncclEpDispatch`` aborts with + ``invalid argument`` at ``ep_backend.cpp:414``. The moe block + computes ``recv_pr`` the same way (see ``moe.py``'s + ``natural_spe = num_ep * max_tokens_per_rank``); keeping the + bootstrap formula in lock-step here. + """ + num_procs = jax.device_count() + num_local_experts = NUM_EXPERTS // EP_SIZE + max_tokens_per_rank = (BATCH // num_procs) * SEQ + natural_spe = EP_SIZE * max_tokens_per_rank + return num_local_experts * natural_spe + + +@pytest.fixture(scope="module") +def mesh(): + if jax.device_count() < NUM_DEVICES_REQUIRED: + pytest.skip( + f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" + f" have {jax.device_count()}" + ) + # ``ep`` must be the inner axis: ``ep_bootstrap`` forms NCCL EP groups + # from consecutive global ranks via ``dp_color = rank // ep_size``, so + # only an (outer_fsdp, inner_ep) device layout groups ranks correctly. + devices = mesh_utils.create_device_mesh((FSDP_SIZE, EP_SIZE)) + mesh_obj = Mesh(devices, axis_names=(FSDP_AXIS, EP_AXIS)) + + num_procs = jax.process_count() + max_tokens_per_rank = (BATCH // num_procs) * SEQ + recv_capacity_per_rank = _compute_worst_case_recv_pr() + + # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather + # and cannot run from inside jax.jit. Sized to the worst-case recv_pr + # across _CONFIGS so every parametrized config is bootstrap-compatible. + with mesh_obj, global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ): + ep_bootstrap( + world_size=num_procs, + rank=jax.process_index(), + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_capacity_per_rank, + hidden_dim=HIDDEN, + max_token_dtype=DTYPE, + ) + record_ep_bootstrap_signature_for_moe( + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_capacity_per_rank, + hidden_dim=HIDDEN, + ep_size=EP_SIZE, + ) + return mesh_obj + + +# ----------------------------------------------------------------------------- +# Pure-JAX reference MoE (no EP). Mirrors the exact math of TE's fused +# router primitive (see tests/jax/test_fused_router.py for the same +# reference applied to the standalone router kernel): +# +# softmax + post-softmax (use_pre_softmax=False, the default): +# 1. top_k by raw logits +# 2. softmax over just the K selected logits (so weights sum to 1) +# +# sigmoid + optional expert_bias: +# 1. scores = sigmoid(logits) +# 2. top_k by (scores + expert_bias) [bias only steers selection] +# 3. weights = scores at top_k positions, normalized when K > 1 +# +# Then for both: +# * weights *= scaling_factor (we leave scaling_factor=1.0 in this +# suite, matching _make_block's default). +# * per-expert FFN: silu(layer_w0) * layer_w1 → wo. +# ----------------------------------------------------------------------------- + + +@partial( + jax.jit, + static_argnames=( + "num_experts", + "num_experts_per_tok", + "aux_loss_coeff", + "score_function", + ), +) +def _pure_jax_moe_reference( + x, + gate_kernel, + wi_0, + wi_1, + wo, + expert_bias=None, + *, + num_experts, + num_experts_per_tok, + aux_loss_coeff: float = 0.0, + score_function: str = "softmax", +): + B, S, H = x.shape + T = B * S + K = num_experts_per_tok + x_2d = x.reshape(T, H) + + gate_kernel_cast = gate_kernel.astype(x.dtype) + logits = (x_2d @ gate_kernel_cast).astype(jnp.float32) # [T, E] + + if score_function == "softmax": + # use_pre_softmax=False: topk on raw logits, then softmax over K. + top_logits, top_indices = jax.lax.top_k(logits, k=K) + weights = jax.nn.softmax(top_logits, axis=-1) # [T, K], sums to 1 + elif score_function == "sigmoid": + scores = jax.nn.sigmoid(logits) # [T, E] + if expert_bias is not None and expert_bias.shape != (0,): + scores_for_routing = scores + expert_bias.astype(jnp.float32)[None, :] + _, top_indices = jax.lax.top_k(scores_for_routing, k=K) + weights = jnp.take_along_axis(scores, top_indices, axis=-1) + else: + weights, top_indices = jax.lax.top_k(scores, k=K) + # Sigmoid weights are normalized when K > 1 (matches the kernel). + if K > 1: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + else: + raise ValueError(f"Unsupported score_function={score_function!r}") + + routing_weights_full = jnp.zeros((T, num_experts), dtype=jnp.float32) + routing_weights_full = routing_weights_full.at[ + jnp.arange(T)[:, None], top_indices + ].set(weights) + + # FFN. ``apply_topk_weights_early`` is a fusion knob that doesn't + # change the math (wo is linear), so the reference is identical for + # both placements. + layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) + layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) + intermediate = jax.nn.silu(layer_w0.astype(jnp.float32)) * layer_w1.astype(jnp.float32) + intermediate = intermediate.astype(x.dtype) + expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] + output_2d = jnp.einsum( + "te,teh->th", routing_weights_full.astype(x.dtype), expert_out + ) + output = output_2d.reshape(B, S, H).astype(x.dtype) + + if aux_loss_coeff > 0.0: + # tex.fused_moe_aux_loss formula (matches the same + # reference_aux_loss helper from test_fused_router.py). The + # "aux scores" use the same score_function but always with + # K-normalised sigmoid (when sigmoid) / plain softmax (when + # softmax) — see tex.fused_topk_with_score_function_fwd with + # compute_aux_scores=True. + if score_function == "softmax": + aux_scores = jax.nn.softmax(logits, axis=-1) + else: # sigmoid + aux_scores = jax.nn.sigmoid(logits) + if K > 1: + aux_scores = aux_scores / ( + aux_scores.sum(axis=-1, keepdims=True) + 1e-20 + ) + routing_map = (routing_weights_full > 0).astype(jnp.int32) + tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] + sum_probs_per_expert = jnp.sum(aux_scores, axis=0) # [E] + aux_loss = (num_experts * aux_loss_coeff / (K * (T**2))) * jnp.sum( + sum_probs_per_expert * tokens_per_expert.astype(jnp.float32) + ) + aux_loss = aux_loss.astype(x.dtype) + else: + aux_loss = jnp.zeros((), dtype=x.dtype) + return output, aux_loss + + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + + +def _make_block( + *, + apply_topk_weights_early=False, + align_size=0, + aux_loss_coeff=0.0, + use_expert_bias=False, + score_function="softmax", + bias_init=None, +): + kwargs = dict( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + data_parallelism_axes=(FSDP_AXIS,), + apply_topk_weights_early=apply_topk_weights_early, + align_size=align_size, + aux_loss_coeff=aux_loss_coeff, + use_expert_bias=use_expert_bias, + score_function=score_function, + dtype=DTYPE, + ) + # Custom bias_init lets tests inject a non-zero expert_bias without + # poking variables['params'] post-init. + if bias_init is not None: + kwargs["bias_init"] = bias_init + return MoEBlock(**kwargs) + + +def _strong_expert_bias_init(key, shape, dtype): + """Half +5, half -5 — large enough to force topk onto the +ve half.""" + del key + n = shape[0] + return jnp.concatenate( + [ + jnp.full((n // 2,), 5.0, dtype=dtype), + jnp.full((n - n // 2,), -5.0, dtype=dtype), + ] + ) + + +def _shard_inputs(x, mesh): + # Match the layout moe.py re-pins to: outer dp axes, then ep innermost. + return jax.lax.with_sharding_constraint( + x, NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None, None)) + ) + + +def _ctx(mesh): + """Combined mesh + global_shard_guard + axis_rules context.""" + + class _Combo: + def __enter__(self_inner): + self_inner._m = mesh.__enter__() + self_inner._gs = global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ) + self_inner._gs.__enter__() + self_inner._ar = nn_partitioning.axis_rules(LOGICAL_AXIS_RULES) + self_inner._ar.__enter__() + return self_inner._m + + def __exit__(self_inner, *args): + self_inner._ar.__exit__(*args) + self_inner._gs.__exit__(*args) + mesh.__exit__(*args) + + return _Combo() + + +def _init_apply(block, mesh, x, key): + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + variables = jax.jit(block.init)(key, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) + output, aux = jax.jit(block.apply)(variables, x_sh) + jax.block_until_ready(output) + return variables, output, aux + + +def _grad_step(block, variables, mesh, x, *, include_aux=False): + """Run jax.grad of mean(out^2) [+ aux if include_aux] vs params.""" + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + + def loss_fn(variables, x): + output, aux = block.apply(variables, x) + loss = jnp.mean(output.astype(jnp.float32) ** 2) + if include_aux and aux is not None: + loss = loss + aux.astype(jnp.float32) + return loss + + grads = jax.jit(jax.grad(loss_fn))(variables, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) + return grads + + +def _grad_aux_only(block, variables, mesh, x): + """Jit'd grad of just the aux loss scalar — proves it reaches the + gate even when no main-output contribution is present.""" + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + + def aux_only(variables, x): + _, aux = block.apply(variables, x) + return aux.astype(jnp.float32) + + grads = jax.jit(jax.grad(aux_only))(variables, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) + return grads + + +def _unwrap(x): + return x.value if hasattr(x, "value") else x + + +def _to_global_numpy(arr, mesh): + """Replicate a sharded JAX array onto every rank and return as numpy. + + Triggers an all-gather inside JIT. The resulting addressable_data(0) + contains the full global array on every process, so we can run the + pure-JAX reference and compare against it from any process. + """ + rep = NamedSharding(mesh, P()) + with mesh: + full = jax.jit(lambda a: jax.lax.with_sharding_constraint(a, rep))(arr) + full.block_until_ready() + return np.asarray(jax.device_get(full.addressable_data(0))) + + +def _params_global_numpy(variables, mesh): + """Pull every entry of variables['params'] to a replicated numpy array.""" + params = variables["params"] + return {name: _to_global_numpy(_unwrap(p), mesh) for name, p in params.items()} + + +def _make_inputs(key): + """Generate a globally-identical input tensor on every process.""" + return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) + + +# ----------------------------------------------------------------------------- +# Tests +# ----------------------------------------------------------------------------- + + +# ----------------------------------------------------------------------------- +# Parametrize variants exercised by both the forward and the backward +# parity tests. Each config is one MoE-block configuration the suite +# wants covered; the test body checks shape, dtype, finiteness AND +# numerical parity vs the same pure-JAX reference (which understands +# the same set of knobs). +# ----------------------------------------------------------------------------- + +_CONFIGS = [ + pytest.param( + dict(score_function="softmax"), + id="softmax", + ), + # TODO: re-add the apply_topk_weights_early=True config once the + # 0*NaN -> NaN leak from padded recv slots in the early-weighting + # multiply (intermediate * recv_w * mask) is debugged. Late + # weighting (combine-side) is unaffected and stays covered above. + # Note: a dedicated align_size=128 config was previously listed + # here. It is no longer interesting because moe.py now floors + # slots_per_expert at 128 unconditionally (effective_align = + # max(align_size, 128)), so align_size=0 (default) and + # align_size=128 produce identical layouts. Re-add a distinct + # case only if the floor is loosened or a >128 align is needed + # by a recipe (e.g. some FP8 paths want 256-aligned slots). + pytest.param( + dict(score_function="sigmoid"), + id="sigmoid", + ), + pytest.param( + dict(score_function="sigmoid", use_expert_bias=True), + id="sigmoid-bias-zero", + ), + pytest.param( + dict( + score_function="sigmoid", + use_expert_bias=True, + bias_init=_strong_expert_bias_init, + ), + id="sigmoid-bias-strong", + ), +] + + +def _reference_kwargs_from_config(config, params_np): + """Pick out the reference-relevant pieces of a parametrize config.""" + return dict( + score_function=config.get("score_function", "softmax"), + expert_bias=( + jnp.asarray(params_np["expert_bias"]) + if config.get("use_expert_bias", False) + else None + ), + ) + + +class TestTeEpMoeForward: + """Per-config forward correctness in a single run: shape, dtype, + finiteness AND numerical parity vs the pure-JAX reference.""" + + @pytest.mark.parametrize("config", _CONFIGS) + def test_forward(self, mesh, config): + block = _make_block(**config) + x = _make_inputs(jax.random.PRNGKey(0)) + variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) + + # Shape / dtype / finiteness (cheap; on the local shard). + assert output.shape == x.shape + assert output.dtype == x.dtype + out_local = np.asarray(jax.device_get(output.addressable_data(0))) + assert np.all(np.isfinite(out_local)), "output has NaN/Inf" + assert aux is None, "aux_loss should be None when aux_loss_coeff == 0" + + # Numerical parity (replicated global view -> single rank's numpy). + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + out_te_np = _to_global_numpy(output, mesh) + + out_ref, _ = _pure_jax_moe_reference( + jnp.asarray(x_np), + jnp.asarray(params_np["gate_kernel"]), + jnp.asarray(params_np["wi_0"]), + jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wo"]), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + **_reference_kwargs_from_config(config, params_np), + ) + np.testing.assert_allclose( + out_te_np.astype(np.float32), + np.asarray(jax.device_get(out_ref)).astype(np.float32), + atol=FWD_ATOL, + rtol=FWD_RTOL, + err_msg=f"forward parity breach for config={config}", + ) + + +class TestTeEpMoeBackward: + """Per-config backward correctness in a single run: per-tensor + grads finite, non-zero AND parity vs the pure-JAX reference.""" + + @pytest.mark.parametrize("config", _CONFIGS) + def test_backward(self, mesh, config): + block = _make_block(**config) + x = _make_inputs(jax.random.PRNGKey(2)) + variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) + grads_te = _grad_step(block, variables, mesh, x) + + # Reference grads via jax.grad over the pure-JAX MoE with the + # same config. + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + ref_kwargs = _reference_kwargs_from_config(config, params_np) + ref_expert_bias = ref_kwargs.pop("expert_bias") + + def loss_fn(params, x): + out, _ = _pure_jax_moe_reference( + x, + params["gate_kernel"], + params["wi_0"], + params["wi_1"], + params["wo"], + ref_expert_bias, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + **ref_kwargs, + ) + return jnp.mean(out.astype(jnp.float32) ** 2) + + grads_ref = jax.jit(jax.grad(loss_fn))( + {k: jnp.asarray(v) for k, v in params_np.items() if k != "expert_bias"}, + jnp.asarray(x_np), + ) + grads_ref_np = {k: np.asarray(jax.device_get(v)) for k, v in grads_ref.items()} + + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + # Per-tensor: finite + non-zero + parity in one pass. + g_te = _to_global_numpy(_unwrap(grads_te["params"][name]), mesh) + assert np.all(np.isfinite(g_te)), f"{name} grad has NaN/Inf [config={config}]" + assert np.any(g_te != 0.0), f"{name} grad identically zero [config={config}]" + atol, rtol = ( + (GRAD_GATE_ATOL, GRAD_GATE_RTOL) + if name == "gate_kernel" + else (GRAD_FFN_ATOL, GRAD_FFN_RTOL) + ) + np.testing.assert_allclose( + g_te.astype(np.float32), + grads_ref_np[name].astype(np.float32), + atol=atol, + rtol=rtol, + err_msg=f"grad parity breach on {name} [config={config}]", + ) + + +class TestTeEpMoeAuxLoss: + """Aux-loss path. Consolidated into: + * ``test_aux_loss``: one run that checks the returned scalar's + shape / dtype / finiteness / magnitude AND numerical parity vs the + reference AND that the aux-only bwd propagates to gate_kernel. + * ``test_combined_loss_grads``: one run for joint main+aux bwd + finite + non-zero per tensor. + """ + + def test_aux_loss(self, mesh): + coeff = 1e-2 + 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)) + + # Shape / dtype / finiteness / magnitude. + assert aux is not None, "aux_loss should be returned when coeff > 0" + assert aux.shape == (), f"aux_loss must be 0-d scalar, got {aux.shape}" + assert aux.dtype == DTYPE, f"aux_loss dtype {aux.dtype} != {DTYPE}" + aux_np = _to_global_numpy(aux, mesh) + assert np.isfinite(aux_np), "aux_loss is NaN/Inf" + assert abs(float(aux_np)) < 1e2, f"aux_loss looks unreasonable: {aux_np}" + + # Numerical parity vs the reference. + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + _, aux_ref = _pure_jax_moe_reference( + jnp.asarray(x_np), + jnp.asarray(params_np["gate_kernel"]), + jnp.asarray(params_np["wi_0"]), + jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wo"]), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + aux_loss_coeff=coeff, + ) + np.testing.assert_allclose( + float(aux_np), + float(jax.device_get(aux_ref)), + atol=AUX_ATOL, + rtol=AUX_RTOL, + ) + + # Aux-only bwd must propagate to gate_kernel — proves the + # fused_moe_aux_loss_bwd → topk(compute_aux_scores)_bwd chain is + # wired. + aux_grads = _grad_aux_only(block, variables, mesh, x) + g_gate = np.asarray( + jax.device_get( + _unwrap(aux_grads["params"]["gate_kernel"]).addressable_data(0) + ) + ) + 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): + """Joint main + aux loss bwd: per-tensor finite + non-zero in + one pass.""" + 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) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g_local = np.asarray( + jax.device_get(_unwrap(grads["params"][name]).addressable_data(0)) + ) + assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf under main+aux" + assert np.any(g_local != 0.0), f"{name} grad zero under main+aux" + + +class TestTeEpMoEBlockFlax: + """Flax wrapper end-to-end in one run: shape/dtype/finiteness on the + forward, numerical parity vs the same reference, and per-tensor + grad finiteness + non-zeroness.""" + + def test_init_apply_parity(self, mesh): + block = _make_block() + x = _make_inputs(jax.random.PRNGKey(12)) + variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(13)) + + assert aux is None + assert output.shape == x.shape + assert output.dtype == x.dtype + out_local = np.asarray(jax.device_get(output.addressable_data(0))) + assert np.all(np.isfinite(out_local)) + + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + out_te_np = _to_global_numpy(output, mesh) + out_ref, _ = _pure_jax_moe_reference( + jnp.asarray(x_np), + jnp.asarray(params_np["gate_kernel"]), + jnp.asarray(params_np["wi_0"]), + jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wo"]), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + ) + np.testing.assert_allclose( + out_te_np.astype(np.float32), + np.asarray(jax.device_get(out_ref)).astype(np.float32), + atol=FWD_ATOL, + rtol=FWD_RTOL, + ) + + grads = _grad_step(block, variables, mesh, x) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g_local = np.asarray( + jax.device_get(_unwrap(grads["params"][name]).addressable_data(0)) + ) + assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf" + assert np.any(g_local != 0.0), f"{name} grad zero" + + +# Keep the bootstrap-signature test last in the module (the "ZZZ" prefix +# ensures pytest's alphabetic class ordering picks it last): it +# intentionally mismatches the NCCL EP bootstrap signature, which +# permanently taints the per-process bootstrap cache for the rest of +# the file. +class TestZZZTeEpMoeBootstrap: + """Per-process NCCL bootstrap re-bootstrap rejection.""" + + def test_bootstrap_signature_mismatch_raises(self, mesh): + block_a = _make_block() + x_a = _make_inputs(jax.random.PRNGKey(14)) + _init_apply(block_a, mesh, x_a, jax.random.PRNGKey(15)) + + # Different hidden dim → different bootstrap signature. + bigger_hidden = HIDDEN * 2 + x_b = jax.random.normal( + jax.random.PRNGKey(16), (BATCH, SEQ, bigger_hidden), dtype=DTYPE + ) + block_b = MoEBlock( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + data_parallelism_axes=(FSDP_AXIS,), + dtype=DTYPE, + ) + with pytest.raises(ValueError, match="bootstrapped"): + _init_apply(block_b, mesh, x_b, jax.random.PRNGKey(17)) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 2a1c818cb3..08348b0104 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -1,76 +1,52 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Functional Mixture-of-Experts (MoE) entry point with a single fused VJP. - -This module exposes :func:`moe`, the framework-agnostic flat function that -implements an entire MoE block (gate -> top-k routing -> token dispatch -> -per-expert FFN -> token combine, plus optional expert parallelism via a -shard_map / ragged_all_to_all collective) under a *single* -``jax.custom_vjp``. It is the moral analog of -:func:`transformer_engine.jax.layernorm_mlp.layernorm_mlp` for MoE: one -custom_vjp boundary covers the whole block so future fusions (FP8 over the -EP wire, fused ``ragged_all_to_all + grouped_gemm``, gate+route+dispatch -fusion) can land without re-architecting the call site. - -Design rationale ----------------- - -The earlier MoE block (:class:`transformer_engine.jax.flax.moe._MoEBlock`) -composed many narrower custom_vjps -- one per :func:`grouped_dense`, one -per :func:`token_dispatch`, etc. Every nested custom_vjp is a place where -a quantized :class:`ScaledTensor` cannot survive (JAX requires custom_vjp -inputs / outputs to be plain ``jnp.ndarray`` ish pytrees). To enable -end-to-end FP8 flow -- in particular FP8 carried over the EP -ragged_all_to_all -- the dispatch's quantize, the a2a, the per-expert -FFN, the inverse a2a, and the combine all have to live inside the same -VJP. This file collapses them into one. - -Implementation conventions --------------------------- - -* No nested ``custom_vjp``. Every primitive's ``_fwd`` and ``_bwd`` is - called directly (e.g. :func:`tex.fused_topk_with_score_function_fwd` / - ``_bwd``, :func:`unpermute_with_mask_map`, - :func:`unpermute_bwd_with_merging_probs`, - :func:`sort_chunks_by_map(is_forward=False)`, - forward + reverse :func:`jax.lax.ragged_all_to_all`) so the outer - ``_moe_bwd_rule`` controls the bwd graph end-to-end without invoking - ``jax.vjp`` for re-linearization. -* The fwd/bwd context (``ctx``) is a plain ``dict`` whose keys depend on - the static configuration (permutation backend, EP active or not, - presence of biases, aux loss enabled). The ``_moe_fwd_rule`` builds a - matching ``ctx_specs`` dict in lockstep when opening the EP shard_map - so ``out_specs`` structurally matches the body's return. -* :func:`_dispatch` is the helper that wraps - ``permute -> a2a -> local_permute`` (forward); :func:`_combine` is its - inverse. Their ``_bwd`` siblings drive the inverse collectives in the - bwd rule. None of these helpers form a custom_vjp boundary. +"""Mixture-of-Experts (MoE) layer for TransformerEngine JAX. + +This module exposes :func:`moe`, a single fused MoE forward pass + bwd +built on top of TE's NCCL-backed Expert Parallelism primitives +(``tex.ep_dispatch`` / ``tex.ep_combine``). The block runs:: + + gate -> topk -> ep_dispatch -> per-expert FFN (grouped GEMMs) + -> ep_combine -> output + +under a single ``jax.custom_vjp`` so the routing, dispatch, FFN and +combine steps fuse cleanly under XLA without leaking intermediate +residuals into the user-facing autograd graph. + +Sharding model +-------------- +* Inbound activations are 3D ``[B, S, H]`` sharded + ``((*data_parallelism_axes, ep_axis), None, None)``. The public + :func:`moe` soft-repins this on entry and warns when a reshard is + inserted. +* The EP primitives operate at global view (their custom_partitioning + rules handle per-shard execution). The FFN GEMMs run per-shard inside + a small ``shard_map`` whose ``in_specs`` and ``out_specs`` mirror the + same ``((dp, ep), ...)`` layout. + +Out-of-scope (for now) +---------------------- +FP8 / MXFP8 quantizer sets are not yet wired on this path; turning +them on requires recipe-aware residual specs and ``ScaledTensor`` +leaves across the ``shard_map`` boundary. ``aux_loss_coeff`` and +``expert_bias`` are supported (the former forces a per-step +all-gather over the routing-side logits, which lives off the critical +path and overlaps with the dispatch collective). """ -import math from dataclasses import dataclass -from enum import Enum from functools import partial -from typing import Any, NewType, Optional, Tuple, Union +from typing import Any, Optional, Tuple, Union +import warnings import jax import jax.numpy as jnp -from flax import struct as flax_struct -from jax.sharding import PartitionSpec as P +from jax.sharding import NamedSharding, PartitionSpec as P +from jax.tree_util import register_pytree_node_class from . import cpp_extensions as tex -from .permutation import ( - PureJaxPermState, - compute_ragged_all_to_all_params, - compute_reverse_ragged_all_to_all_params, - pure_jax_token_combine, - pure_jax_token_dispatch, - routing_map_to_selected_experts, -) from .quantize import ( - QuantizerSet, - ScaledTensor, TensorUsage, noop_quantizer_set, with_sharding_constraint_by_logical_axes, @@ -79,1070 +55,316 @@ from .router import ScoreFunction, _validate_score_function from .sharding import _get_mesh -# Triton-backed primitives are imported lazily: callers on the PURE_JAX -# permutation backend should not need ``triton`` installed. The TRITON -# branches in this module call ``_require_triton()`` first to raise a -# clear error if the import failed. -try: - from .triton_extensions.permutation import ( - make_chunk_sort_map, - make_row_id_map, - permute_with_mask_map, - permute_with_mask_map_and_pad, - sort_chunks_by_map, - unpermute_bwd_with_merging_probs, - unpermute_bwd_with_merging_probs_and_unpad, - unpermute_with_mask_map, - unpermute_with_mask_map_and_unpad, - ) - - _TRITON_AVAILABLE = True -except ImportError: - _TRITON_AVAILABLE = False - make_chunk_sort_map = None - make_row_id_map = None - permute_with_mask_map = None - permute_with_mask_map_and_pad = None - sort_chunks_by_map = None - unpermute_bwd_with_merging_probs = None - unpermute_bwd_with_merging_probs_and_unpad = None - unpermute_with_mask_map = None - unpermute_with_mask_map_and_unpad = None - - -def _require_triton(): - """Raise a clear error if Triton permutation kernels are unavailable.""" - if not _TRITON_AVAILABLE: - raise ImportError( - "PermutationBackend.TRITON requires" - " ``transformer_engine.jax.triton_extensions`` (and ``triton``)." - " Install Triton or pass PermutationBackend.PURE_JAX." - ) - - -PRNGKey = Any -Shape = Tuple[int, ...] -DType = NewType("DType", jnp.dtype) -Array = NewType("Array", jnp.ndarray) +__all__ = ["moe"] -__all__ = ["moe", "PermutationBackend"] +def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: + """Apply a sharding constraint while keeping bwd cotangents in the primal dtype. - -# ============================================================================= -# Enums -# ============================================================================= - - -class PermutationBackend(Enum): - """Token-dispatch / combine backend used by :func:`moe`. - - * ``TRITON``: TE's fused Triton kernels. Faster than ``PURE_JAX`` - on current hardware and the recommended default. - * ``PURE_JAX``: ``jnp.argsort`` + gather paths compiled as plain - XLA; useful as a numerical reference and on builds without - Triton available. + Plain ``jax.lax.with_sharding_constraint`` propagates cotangents in + whatever dtype the upstream gradient lands in; under mixed precision + that can be wider than the primal, blowing up bandwidth and (for + bf16 primals) breaking downstream kernels that pin a bf16 input + layout. This wrapper re-casts the cotangent back to the primal + dtype and re-asserts the same sharding on the bwd path. """ - PURE_JAX = "pure_jax" - TRITON = "triton" + @jax.custom_vjp + def _constraint(y): + return jax.lax.with_sharding_constraint(y, sharding) + def _constraint_fwd(y): + return jax.lax.with_sharding_constraint(y, sharding), jnp.zeros((), dtype=y.dtype) -# ============================================================================= -# Dispatch-state records (carried _dispatch -> _combine / *_bwd) -# ============================================================================= -# -# Two NamedTuples (one per permutation backend) so we get type -# discrimination at the consumer side via ``isinstance``. The backend- -# specific residuals are required fields; the EP-only residuals are -# Optional and are populated only when the run is EP-active. Each field -# is either an ``ndarray`` or ``None`` -- nothing static, since these -# values cross the shard_map pytree boundary and would otherwise be -# coerced into JitTracers. - - -@flax_struct.dataclass -class _PureJaxDispatchState: - """Residuals saved by :func:`_dispatch` on the PURE_JAX path. - - Registered as a JAX pytree via ``flax.struct.dataclass``: each - annotated field is a leaf, ``None`` is a non-leaf sentinel. The - matching spec built by :func:`_build_dispatch_specs` mirrors this - layout so shard_map's value and spec trees line up. - """ - - group_sizes: jnp.ndarray - sorted_indices: jnp.ndarray - routing_weights: jnp.ndarray - # EP-only: - all_shards_tokens_per_expert: Optional[jnp.ndarray] = None - local_perm_row_id_map: Optional[jnp.ndarray] = None - - -@flax_struct.dataclass -class _TritonDispatchState: - """Residuals saved by :func:`_dispatch` on the TRITON path.""" - - group_sizes: jnp.ndarray - row_id_map: jnp.ndarray - pad_offsets: Optional[jnp.ndarray] # populated only when align_size > 0 - merging_probs: jnp.ndarray - # EP-only: - all_shards_tokens_per_expert: Optional[jnp.ndarray] = None - local_perm_row_id_map: Optional[jnp.ndarray] = None - - -_DispatchState = Union[_PureJaxDispatchState, _TritonDispatchState] - - -@flax_struct.dataclass -class _BodyCtx: - """Residuals carried fwd_rule -> bwd_rule by :func:`_body_fwd`. - - Optional fields (``expert_bias``, ``aux_*``) are ``None`` when the - matching feature is disabled. :func:`_build_ctx_specs` mirrors that - layout so the shard_map spec and value trees match leaf-for-leaf. - """ - - # Always present. - x: Any - gate_kernel: Any - logits_2d: Any - saved_scores: Any - routing_map: Any - dispatch: Any # _DispatchState - casted_sorted_x_lhs_trans: Any - casted_wi_rhs_trans: Any # combined [E, H, 2M] residual for fused wi_0|wi_1 bwd - gate_proj_out: Any - up_proj_out: Any - casted_intermediate_lhs_trans: Any - casted_wo_rhs_trans: Any - expert_outputs: Any - local_group_sizes: Any - # Feature-gated. - expert_bias: Any = None - aux_const_buf: Any = None - aux_tokens_per_expert: Any = None - aux_logits_for_score: Any = None - aux_saved_scores: Any = None - + def _constraint_bwd(dtype_ref, grad): + return (jax.lax.with_sharding_constraint(grad.astype(dtype_ref.dtype), sharding),) -# ============================================================================= -# ctx / dispatch-state key conventions -# ============================================================================= -# -# Both ``ctx`` (carried fwd_rule -> bwd_rule) and the dispatch state -# (carried _dispatch -> _combine / _dispatch_bwd / _combine_bwd) are plain -# python dicts. Using a dict (rather than a flax_struct.dataclass) lets us -# vary the populated keys with the static config without breaking -# ``shard_map``'s ``out_specs`` structural match: the spec dict and the -# value dict are built with the SAME keys via :func:`_build_ctx_specs`. -# -# Below is the key glossary so the rest of the file reads cleanly. -# -# DispatchState (dict): values are jnp.ndarray unless noted -# Always present: -# "group_sizes" [n_groups] per-expert token counts -# (n_groups = E for no-EP, -# E_local for EP) -# "ep_active" bool (carried as a Python flag, -# not in the dict; passed -# alongside) -# PURE_JAX backend: -# "sorted_indices" [num_real + padding] argsort indices -# "routing_weights" [num_tokens, topk] per-token-per-expert weights -# TRITON backend: -# "row_id_map" [num_tokens, 2*E + 1] -# "pad_offsets" [E] or None -# "merging_probs" [num_tokens, E] -# EP-only: -# "all_shards_tokens_per_expert" [num_ep, E] -# "local_perm_row_id_map" [recv_buffer_rows] -# "local_perm_inv_row_id_map" [recv_buffer_rows] -# -# NOTE: per-shard compile-time-constant shapes (num_real_tokens, -# padding_size, pre/post_a2a_buffer_shape) are NOT stored in this -# dict; they are recomputed in _body_fwd/_body_bwd via -# _compute_static_shape_info and passed as Python ints / int tuples to -# the dispatch/combine helpers. Storing them in the dict would cause -# JAX's pytree-flatten across the shard_map boundary to coerce them -# into JitTracer 0-d arrays, which breaks Python-level control flow -# (e.g. ``if padding > 0``) and ``jnp.zeros(shape)`` in the bwd. -# -# See :class:`_BodyCtx` (NamedTuple) for the ctx layout and field -# documentation. :func:`_build_ctx_specs` returns a matching ``_BodyCtx`` -# of ``P(...)`` specs so shard_map's value/spec trees line up -# leaf-for-leaf. + _constraint.defvjp(_constraint_fwd, _constraint_bwd) + return _constraint(x) # ============================================================================= -# Static shape helper +# Process-level NCCL EP bootstrap (must run eagerly, outside jax.jit) # ============================================================================= # -# A set of per-shard shape/size values that the dispatch and combine -# helpers (both fwd and bwd) need. They're all derivable from existing -# static args, so we recompute them in both ``_body_fwd`` and -# ``_body_bwd`` and pass them as Python ints / int-tuples through -# explicit kwargs. We MUST NOT stash them inside the dynamic -# ``state`` / ``ctx`` dict: when the dict crosses the EP shard_map's -# out_specs/in_specs boundary, JAX's pytree-flatten coerces any Python -# int leaves into traced 0-d arrays, which then breaks dependent Python -# code in the bwd (e.g. ``if padding > 0`` and ``jnp.zeros(shape)``). - - -@dataclass(frozen=True) -class _StaticShapeInfo: - """Per-shard compile-time-constant shape info used by dispatch / - combine fwd and bwd. Fields are Python ints / int tuples (NOT jnp - arrays) so they can be passed as ordinary static keyword args. - - Attributes - ---------- - num_real_tokens : int - Per-shard count of real (non-padding) permuted tokens, - i.e. ``per_shard_num_tokens * num_experts_per_tok``. - padding_size : int - Per-shard number of alignment-padding tokens appended to the - sort buffer (``num_experts * (align_size - 1)`` when - ``align_size > 0``, else ``0``). - pre_a2a_buffer_shape : tuple[int, int] - ``(num_real_tokens + padding_size, hidden)`` -- the per-shard - shape of the sorted-inputs buffer sent over the EP - ragged_all_to_all in the fwd direction. - post_a2a_buffer_shape : Optional[tuple[int, int]] - ``(recv_buffer_rows, hidden)`` when EP is active, ``None`` - otherwise. - """ +# ``tex.ep_bootstrap`` does a NCCL UID allgather over the JAX runtime, which +# cannot run from inside a jit-traced function. The caller must bootstrap +# eagerly once per process before any jitted MoE call, then record the +# bootstrap signature via ``record_ep_bootstrap_signature_for_moe``. The +# per-call check below verifies the recorded signature is wide enough for +# the current MoE invocation (smaller per-call usage is fine since the C++ +# backend reserves worst-case buffers at bootstrap time). - num_real_tokens: int - padding_size: int - pre_a2a_buffer_shape: Tuple[int, int] - post_a2a_buffer_shape: Optional[Tuple[int, int]] +_te_ep_bootstrap_signature: Optional[Tuple[int, int, int, int, int]] = None -def _compute_static_shape_info( - *, - batch_size: int, - sequence_length: int, - hidden: int, +def record_ep_bootstrap_signature_for_moe( num_experts: int, - num_experts_per_tok: int, - align_size: int, - ep_active: bool, - num_ep: int = 1, - fsdp_sizes: Tuple[int, ...] = (), - recv_buffer_rows: int = 0, - batch_is_per_shard: bool = True, -) -> _StaticShapeInfo: - """Build a :class:`_StaticShapeInfo` for the current rank. - - ``batch_is_per_shard`` controls whether ``batch_size`` is already - sharded (True -- e.g. when this is called from inside a shard_map - body, where ``x.shape[0]`` reports the per-shard batch size) or - global (False -- e.g. when computing from x.shape outside the - shard_map body). + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + ep_size: int, +) -> None: + """Record the params passed to ``ep_bootstrap`` so the per-call check + in ``_moe_fwd_rule`` can verify compatibility. Call this once per + process immediately after ``ep_bootstrap``. """ - if ep_active and not batch_is_per_shard: - dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 - per_shard_batch = batch_size // (num_ep * dp_size) - else: - per_shard_batch = batch_size - per_shard_num_tokens = per_shard_batch * sequence_length - num_real_tokens = per_shard_num_tokens * num_experts_per_tok - padding_size = num_experts * (align_size - 1) if align_size > 0 else 0 - pre_a2a_buffer_shape = (num_real_tokens + padding_size, hidden) - post_a2a_buffer_shape = (recv_buffer_rows, hidden) if ep_active else None - return _StaticShapeInfo( - num_real_tokens=num_real_tokens, - padding_size=padding_size, - pre_a2a_buffer_shape=pre_a2a_buffer_shape, - post_a2a_buffer_shape=post_a2a_buffer_shape, + global _te_ep_bootstrap_signature + _te_ep_bootstrap_signature = ( + num_experts, + max_tokens_per_rank, + recv_capacity_per_rank, + hidden_dim, + ep_size, ) -# ============================================================================= -# Dispatch / combine helpers (no VJP boundary -- pure Python) -# ============================================================================= - - -def _dispatch( - inputs_2d: jnp.ndarray, - sparse_probs: jnp.ndarray, - routing_map: jnp.ndarray, - *, - backend: PermutationBackend, +def _te_ep_assert_compatible_bootstrap( num_experts: int, - num_experts_per_tok: int, - align_size: int, - # EP-only: - ep_active: bool, - ep_axis: Optional[str], - num_ep: int, - recv_buffer_rows: int, - shard_id: Optional[jnp.ndarray] = None, -) -> Tuple[jnp.ndarray, dict]: - """``permute -> (a2a -> local_permute) iff ep_active``. - - Returns ``(sorted_x, state)`` where ``sorted_x`` has shape - ``[buffer_rows, hidden]`` -- ``E`` groups (no-EP) or ``E_local`` groups - (EP) -- and ``state`` is a dict carrying everything :func:`_combine` - and the bwd helpers need to reverse the operation. - - Bypasses the ``custom_vjp``-wrapped public ``token_dispatch`` / - ``pure_jax_token_dispatch`` wrappers (well, mostly: PURE_JAX still - composes through ``pure_jax_token_dispatch`` because that helper has - no ``custom_vjp`` itself -- only its inner ``_sort_activations`` does, - which is fine since we never auto-diff through it from this layer). - For TRITON we call the underlying ``permute_with_mask_map`` / - ``permute_with_mask_map_and_pad`` primitives directly. - """ - num_tokens, hidden = inputs_2d.shape - topk = num_experts_per_tok - - # Backend-specific residuals collected here, then packaged into the - # appropriate _*DispatchState below. - sorted_indices = None - routing_weights_kept = None - row_id_map = None - pad_offsets = None - merging_probs = None - - # ------------------------------------------------------------------ - # Step 1: global permute (every shard routes its own tokens over the - # full expert axis). Backend-specific. - # ------------------------------------------------------------------ - if backend is PermutationBackend.PURE_JAX: - selected_experts, routing_weights = routing_map_to_selected_experts( - sparse_probs, routing_map, topk + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + ep_size: int, +) -> None: + """Verify a prior eager ``ep_bootstrap`` is wide enough for this call.""" + if _te_ep_bootstrap_signature is None: + raise RuntimeError( + "TE EP was not bootstrapped. Call" + " transformer_engine.jax.ep.ep_bootstrap(...) eagerly (outside" + " any jax.jit) once per process, then" + " transformer_engine.jax.moe.record_ep_bootstrap_signature_for_moe(...)" + " with the same params, before invoking moe()." ) - sorted_inputs, perm_state, group_sizes = pure_jax_token_dispatch( - inputs_2d, - selected_experts, - num_experts=num_experts, - num_experts_per_tok=topk, - align_size=align_size, - ) - # NOTE: ``perm_state.num_real_tokens`` and ``perm_state.padding_size`` - # are compile-time Python ints; intentionally NOT stored in the - # returned state (would be coerced to JitTracer 0-d arrays under - # the EP shard_map's pytree flatten). Recompute via - # ``_compute_static_shape_info`` in the bwd / EP-combine - # call sites that need them. - sorted_indices = perm_state.sorted_indices - routing_weights_kept = routing_weights - else: - # TRITON backend -- inline the underlying primitive sequence - # (mirrors ``_token_dispatch_fwd_rule`` but exposes the residuals - # to our ctx instead of saving them inside another custom_vjp). - num_out_tokens = num_tokens * topk - row_id_map = make_row_id_map(routing_map, num_tokens, num_experts) - tokens_per_expert = jnp.sum(routing_map, axis=0).astype(jnp.int32) - if align_size > 0: - target_tokens_per_expert = ( - jnp.ceil(tokens_per_expert / align_size) * align_size - ).astype(jnp.int32) - pad_lengths = target_tokens_per_expert - tokens_per_expert - cum_pad = jnp.cumsum(pad_lengths) - pad_offsets = jnp.concatenate([jnp.array([0], dtype=cum_pad.dtype), cum_pad[:-1]]) - worst_case_out_tokens = ( - (num_out_tokens + num_experts * (align_size - 1)) // align_size - ) * align_size - sorted_inputs, _ = permute_with_mask_map_and_pad( - inputs_2d, - row_id_map, - None, - pad_offsets, - num_tokens, - num_experts, - worst_case_out_tokens, - hidden, - align_size=align_size, - ) - group_sizes = target_tokens_per_expert - else: - sorted_inputs, _ = permute_with_mask_map( - inputs_2d, - row_id_map, - None, - num_tokens, - num_experts, - num_out_tokens, - hidden, - ) - pad_offsets = None - group_sizes = tokens_per_expert - merging_probs = sparse_probs - - def _build_state(group_sizes_val, ep_all=None, ep_local=None): - if backend is PermutationBackend.PURE_JAX: - return _PureJaxDispatchState( - group_sizes=group_sizes_val, - sorted_indices=sorted_indices, - routing_weights=routing_weights_kept, - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - return _TritonDispatchState( - group_sizes=group_sizes_val, - row_id_map=row_id_map, - pad_offsets=pad_offsets, - merging_probs=merging_probs, - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, + b_num_experts, b_max_tpr, b_recv_pr, b_hidden, b_ep_size = _te_ep_bootstrap_signature + if ( + num_experts != b_num_experts + or hidden_dim != b_hidden + or ep_size != b_ep_size + or max_tokens_per_rank > b_max_tpr + or recv_capacity_per_rank > b_recv_pr + ): + raise ValueError( + "TE EP was already bootstrapped with signature" + f" (num_experts={b_num_experts}, max_tokens_per_rank={b_max_tpr}," + f" recv_capacity_per_rank={b_recv_pr}, hidden_dim={b_hidden}," + f" ep_size={b_ep_size}); this moe() call needs" + f" (num_experts={num_experts}, max_tokens_per_rank={max_tokens_per_rank}," + f" recv_capacity_per_rank={recv_capacity_per_rank}, hidden_dim={hidden_dim}," + f" ep_size={ep_size}). Re-bootstrap with wider params (or matching exact" + " sizes) is required." ) - if not ep_active: - return sorted_inputs, _build_state(group_sizes) - - # ------------------------------------------------------------------ - # Step 2 (EP only): all_gather per-expert counts so every shard knows - # the [num_ep, num_experts] token-count matrix. - # ------------------------------------------------------------------ - all_shards_tokens_per_expert = jax.lax.all_gather( - group_sizes[None, :], - axis_name=ep_axis, - axis=0, - tiled=True, - ) - # ------------------------------------------------------------------ - # Step 3 (EP only): forward ragged_all_to_all over the EP axis. - # ------------------------------------------------------------------ - in_off, send_sz, out_off, recv_sz = compute_ragged_all_to_all_params( - all_shards_tokens_per_expert, shard_id, num_ep - ) - post_a2a_buffer_shape = (recv_buffer_rows, hidden) - recv_buf = jnp.zeros(post_a2a_buffer_shape, dtype=sorted_inputs.dtype) - x_recv = jax.lax.ragged_all_to_all( - sorted_inputs, recv_buf, in_off, send_sz, out_off, recv_sz, axis_name=ep_axis - ) - - # ------------------------------------------------------------------ - # Step 4 (EP only): local permute -- (source_shard, expert) -> - # (expert, shard). Inlined ``local_permute_after_a2a`` so we control - # both the row_id_map and its inverse for the bwd. - # ------------------------------------------------------------------ - num_experts_local = num_experts // num_ep - local_expert_start = shard_id * num_experts_local - local_expert_columns = jax.lax.dynamic_slice( - all_shards_tokens_per_expert, - start_indices=(0, local_expert_start), - slice_sizes=(num_ep, num_experts_local), - ) - split_sizes = local_expert_columns.reshape(-1) # source-major - indices_matrix = jnp.arange(num_ep * num_experts_local, dtype=jnp.int32).reshape( - num_ep, num_experts_local - ) - sorted_chunk_indices = indices_matrix.T.reshape(-1) # source-major -> expert-major - num_chunks = num_ep * num_experts_local - # Build a SINGLE row_id_map. ``is_forward=True`` permutes - # source-major -> expert-major; ``is_forward=False`` is the exact - # inverse (this is exactly what ``_sort_chunks_by_index_bwd_rule`` - # uses on the saved residual). _MoEBlock builds two row_id_maps - # only because it calls ``sort_chunks_by_index`` twice -- once in - # ``local_permute_after_a2a`` and again in ``local_unpermute_before_a2a``; - # each of those wrappers calls ``make_chunk_sort_map`` internally. - # Here we share one map across (fwd permute, fwd inverse-permute, - # bwd permute, bwd inverse-permute). - local_perm_row_id_map = make_chunk_sort_map( - split_sizes, sorted_chunk_indices, recv_buffer_rows, num_chunks - ) - sorted_x, _ = sort_chunks_by_map( - x_recv, local_perm_row_id_map, None, recv_buffer_rows, hidden, is_forward=True - ) - local_group_sizes = jnp.sum(local_expert_columns, axis=0) - - # NOTE: pre_a2a_buffer_shape and post_a2a_buffer_shape are compile- - # time int tuples; intentionally NOT stored in the returned state - # (would be coerced to JitTracer 0-d arrays under the EP shard_map's - # pytree flatten). Recompute via ``_compute_static_shape_info`` in - # the bwd call sites that need them. For EP, ``group_sizes`` here is - # the per-local-expert count (the FFN runs over E_local groups, not - # E). The global ``group_sizes`` lives inside - # ``all_shards_tokens_per_expert`` if anyone needs it for - # diagnostics. - return sorted_x, _build_state( - local_group_sizes, - ep_all=all_shards_tokens_per_expert, - ep_local=local_perm_row_id_map, - ) +# ============================================================================= +# Residual container threaded fwd -> bwd +# ============================================================================= -def _combine( - expert_outputs: jnp.ndarray, - state: _DispatchState, - *, - backend: PermutationBackend, - ep_active: bool, - batch_size: int, - sequence_length: int, - dtype: jnp.dtype, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # Computed by _compute_static_shape_info in the caller, passed here - # rather than stored in ``state`` to survive shard_map crossings. - num_real_tokens: int, - padding_size: int, - pre_a2a_buffer_shape: Tuple[int, int], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Inverse of :func:`_dispatch`. - - Returns ``(output, expert_outputs_post_ep)``. ``output`` is the - ``[B, S, H]`` combined activations. ``expert_outputs_post_ep`` is - the FFN-output tensor in the shape that Step 3 of the combine - actually consumed (i.e. after the reverse ragged_all_to_all on EP - runs, or the original input on non-EP). The caller stashes this as - the bwd residual so that ``_combine_bwd``'s Step-3 inverse sees - the same tensor the forward Step 3 used. - """ - if ep_active: - # Step 1 (EP): inverse local permute. Reuse the SAME row_id_map - # built in _dispatch by setting is_forward=False (this is the - # exact inverse, identical to what - # ``_sort_chunks_by_index_bwd_rule`` does with the saved residual). - recv_buffer_rows, hidden = expert_outputs.shape - x_send_back, _ = sort_chunks_by_map( - expert_outputs, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=False, - ) - # Step 2 (EP): reverse ragged_all_to_all. - in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - send_back_buf = jnp.zeros(pre_a2a_buffer_shape, dtype=expert_outputs.dtype) - expert_outputs = jax.lax.ragged_all_to_all( - x_send_back, - send_back_buf, - in_off_r, - send_sz_r, - out_off_r, - recv_sz_r, - axis_name=ep_axis, - ) +# Registered as a pytree so jax.custom_vjp can flatten/unflatten it across +# the fwd -> bwd boundary. ``cfg`` is the only static field (EpLayerConfig +# is a frozen dataclass of ints); the rest are jnp.ndarray, +# GroupedNoScaleTensor (already a pytree), or None when aux_loss_coeff == 0. +@register_pytree_node_class +@dataclass +class _Ctx: + """Residuals carried from the fwd rule into the bwd rule.""" + + x: jnp.ndarray + gate_kernel: jnp.ndarray + expert_bias: jnp.ndarray + logits_2d: jnp.ndarray + saved_scores: jnp.ndarray + routing_map: jnp.ndarray + cfg: Any + handle_mem: Any + token_counts: jnp.ndarray + recv_topk_weights: jnp.ndarray + casted_sorted_x_lhs_trans: Any + casted_wi_rhs_trans: Any + gate_proj_out: jnp.ndarray + up_proj_out: jnp.ndarray + casted_intermediate_lhs_trans: Any + casted_wo_rhs_trans: Any + expert_outputs: jnp.ndarray + local_group_sizes: jnp.ndarray + # Aux-loss residuals; None when aux_loss_coeff == 0. + aux_const_buf: Any = None + aux_tokens_per_expert: Any = None + aux_saved_scores: Any = None - # Step 3: global combine. ``expert_outputs`` here is the post-A2A - # tensor under EP, or the original input under non-EP -- whichever - # value Step 3 actually consumes. Returned as the second tuple - # element so the caller can stash it as the bwd residual. - if backend is PermutationBackend.PURE_JAX: - # Reuse the reference pure-jax implementation; it has no - # custom_vjp on its outer surface so we can call it freely. - perm_state = PureJaxPermState( - sorted_indices=state.sorted_indices, - num_real_tokens=num_real_tokens, - padding_size=padding_size, - ) - output = pure_jax_token_combine( - expert_outputs, - perm_state, - state.routing_weights, - num_experts_per_tok=num_experts_per_tok, - batch_size=batch_size, - sequence_length=sequence_length, + def tree_flatten(self): + children = ( + self.x, + self.gate_kernel, + self.expert_bias, + self.logits_2d, + self.saved_scores, + self.routing_map, + self.handle_mem, + self.token_counts, + self.recv_topk_weights, + self.casted_sorted_x_lhs_trans, + self.casted_wi_rhs_trans, + self.gate_proj_out, + self.up_proj_out, + self.casted_intermediate_lhs_trans, + self.casted_wo_rhs_trans, + self.expert_outputs, + self.local_group_sizes, + self.aux_const_buf, + self.aux_tokens_per_expert, + self.aux_saved_scores, ) - return output, expert_outputs - # TRITON - num_tokens = state.row_id_map.shape[0] - num_experts = (state.row_id_map.shape[1] - 1) // 2 - hidden = expert_outputs.shape[-1] - if state.pad_offsets is not None: - out_2d, _ = unpermute_with_mask_map_and_unpad( - expert_outputs, - state.row_id_map, - state.merging_probs, - None, - state.pad_offsets, - num_tokens, - num_experts, - hidden, - ) - else: - out_2d, _ = unpermute_with_mask_map( + aux_data = (self.cfg,) + return children, aux_data + + @classmethod + def tree_unflatten(cls, aux_data, children): + (cfg,) = aux_data + ( + x, + gate_kernel, + expert_bias, + logits_2d, + saved_scores, + routing_map, + handle_mem, + token_counts, + recv_topk_weights, + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out, + up_proj_out, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, expert_outputs, - state.row_id_map, - state.merging_probs, - None, - num_tokens, - num_experts, - hidden, + local_group_sizes, + aux_const_buf, + aux_tokens_per_expert, + aux_saved_scores, + ) = children + return cls( + 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, + token_counts=token_counts, + 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, + aux_const_buf=aux_const_buf, + aux_tokens_per_expert=aux_tokens_per_expert, + aux_saved_scores=aux_saved_scores, ) - return out_2d.reshape(batch_size, sequence_length, hidden).astype(dtype), expert_outputs - - -def _combine_bwd( # pylint: disable=unused-argument - d_output: jnp.ndarray, - state: _DispatchState, - expert_outputs: jnp.ndarray, - *, - backend: PermutationBackend, - ep_active: bool, - batch_size: int, - sequence_length: int, - dtype: jnp.dtype, - num_experts: int, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # See ``_compute_static_shape_info`` and the note in ``_dispatch`` - # for why these are kwargs rather than state-dict entries. - num_real_tokens: int, - padding_size: int, - post_a2a_buffer_shape: Optional[Tuple[int, int]], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Inverse of :func:`_combine` on the cotangent. - - Returns ``(d_expert_outputs, d_routing_weights_or_merging_probs)``. - - ``expert_outputs`` is the *forward* output of the FFN (same value the - fwd handed to :func:`_combine`). It's required by the TRITON - combine_bwd kernel; for PURE_JAX we don't need it but accept it for - a symmetric signature. - """ - # Step 3 inverse: global combine bwd. - d_output_2d = d_output.reshape(-1, d_output.shape[-1]) - if backend is PermutationBackend.PURE_JAX: - # The pure-jax combine is: - # unsort = _sort_activations(expert_outputs, argsort(sorted_indices)) - # if pad: unsort = unsort[:num_real] - # reshape -> einsum BKE,BK -> BE -> reshape to BSE - # Hand-derive the bwd in plain JAX (no custom_vjp involved): - unsort_indices = jnp.argsort(state.sorted_indices) - topk = num_experts_per_tok - num_real = num_real_tokens - padding = padding_size - # Recover the unsorted intermediate that the fwd produced (we - # need it for the d_routing_weights pullback). Apply the same - # gather the fwd did. - unsort_intermediate = expert_outputs[unsort_indices] - if padding > 0: - unsort_intermediate = unsort_intermediate[:num_real] - # Bwd of einsum/reshape: - # output[B, E] = sum_K intermediate[B, K, E] * weights[B, K] - # d_intermediate[B, K, E] = d_output[B, E] * weights[B, K] - # d_weights[B, K] = sum_E d_output[B, E] * intermediate[B, K, E] - rw = state.routing_weights.reshape(-1, topk) - intermediate_3d = unsort_intermediate.reshape(rw.shape[0], topk, -1) - rw_cast = rw.astype(intermediate_3d.dtype) - d_intermediate_3d = jnp.einsum("BE,BK -> BKE", d_output_2d, rw_cast) - d_routing_weights = jnp.einsum("BE,BKE -> BK", d_output_2d, intermediate_3d).astype( - state.routing_weights.dtype - ) - d_routing_weights = d_routing_weights.reshape(state.routing_weights.shape) - d_unsort_intermediate = d_intermediate_3d.reshape(num_real, -1) - # Pad back with zeros if the fwd stripped padding. - if padding > 0: - d_unsort_intermediate = jnp.concatenate( - [ - d_unsort_intermediate, - jnp.zeros( - (padding, d_unsort_intermediate.shape[-1]), - dtype=d_unsort_intermediate.dtype, - ), - ], - axis=0, - ) - # Bwd of the gather is gather-by-original-indices: - # sorted = unsort[argsort(sorted_indices)] - # d_sorted = scatter d_unsort via argsort(sorted_indices) - # = d_unsort[sorted_indices] (gather by original sorted_indices, - # which is the inverse of argsort(sorted_indices)). - d_expert_outputs_global = d_unsort_intermediate[state.sorted_indices] - else: - # TRITON combine bwd: requires fwd_input (expert_outputs). - num_tokens = state.row_id_map.shape[0] - n_experts = (state.row_id_map.shape[1] - 1) // 2 - hidden = d_output_2d.shape[-1] - num_out_tokens = expert_outputs.shape[0] - if state.pad_offsets is not None: - d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs_and_unpad( - d_output_2d, - state.row_id_map, - expert_outputs, - state.merging_probs, - state.pad_offsets, - num_tokens, - n_experts, - num_out_tokens, - hidden, - ) - # The kernel only writes positions tokens map to; padded - # positions may contain NaN. Replace with zeros (matches - # ``_token_combine_bwd_rule``). - d_expert_outputs_global = jnp.where( - jnp.isnan(d_expert_outputs_global), 0.0, d_expert_outputs_global - ) - else: - d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs( - d_output_2d, - state.row_id_map, - expert_outputs, - state.merging_probs, - num_tokens, - n_experts, - num_out_tokens, - hidden, - ) - d_routing_weights = d_merging_probs - - if not ep_active: - return d_expert_outputs_global, d_routing_weights - - # Step 2 (EP) inverse: bwd of reverse ragged_all_to_all is a forward - # ragged_all_to_all using the SAME forward parameters (sender / - # receiver roles swap from the reverse direction back to forward). - in_off_f, send_sz_f, out_off_f, recv_sz_f = compute_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - recv_buf_for_bwd = jnp.zeros(post_a2a_buffer_shape, dtype=d_expert_outputs_global.dtype) - d_x_send_back = jax.lax.ragged_all_to_all( - d_expert_outputs_global, - recv_buf_for_bwd, - in_off_f, - send_sz_f, - out_off_f, - recv_sz_f, - axis_name=ep_axis, - ) - # Step 1 (EP) inverse: combine fwd applied is_forward=False; the - # bwd is is_forward=True with the SAME row_id_map. - recv_buffer_rows, hidden = d_x_send_back.shape - d_expert_outputs, _ = sort_chunks_by_map( - d_x_send_back, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=True, - ) - return d_expert_outputs, d_routing_weights - - -def _dispatch_bwd( - d_sorted_x: jnp.ndarray, - state: _DispatchState, - inputs_2d_shape: Tuple[int, ...], - *, - backend: PermutationBackend, - ep_active: bool, - num_experts: int, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # See ``_compute_static_shape_info`` and the note in ``_dispatch`` - # for why these are kwargs rather than state-dict entries. - num_real_tokens: int, - padding_size: int, - pre_a2a_buffer_shape: Tuple[int, int], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> jnp.ndarray: - """Inverse of :func:`_dispatch` on the cotangent. Returns ``d_inputs_2d``. - - The probs path through dispatch is always discarded (PURE_JAX never - threads probs through dispatch; TRITON technically does but the - caller drops ``permuted_probs``, so its cotangent is structurally - zero). The probs gradient instead flows back through - :func:`_combine_bwd`. - """ - if ep_active: - # Step 4 inverse: dispatch fwd applied is_forward=True; bwd is - # is_forward=False with the SAME row_id_map. - recv_buffer_rows, hidden = d_sorted_x.shape - d_x_recv, _ = sort_chunks_by_map( - d_sorted_x, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=False, - ) - # Step 3 inverse: bwd of forward ragged_a2a is the reverse-direction - # ragged_a2a using the SAME params with sender/receiver swapped. - in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - recv_buf_pre = jnp.zeros(pre_a2a_buffer_shape, dtype=d_x_recv.dtype) - d_sorted_x = jax.lax.ragged_all_to_all( - d_x_recv, - recv_buf_pre, - in_off_r, - send_sz_r, - out_off_r, - recv_sz_r, - axis_name=ep_axis, - ) - - # Step 1 inverse: global permute bwd. - if backend is PermutationBackend.PURE_JAX: - # Fwd was: replicated = repeat(inputs_2d, topk, axis=0) - # padded = pad(replicated, (0, padding_size)) - # sorted = padded[sorted_indices] - # Bwd: d_padded = scatter via sorted_indices - # = d_sorted[argsort(sorted_indices)] - # d_replicated = d_padded[:num_real] - # d_inputs_2d = d_replicated.reshape(T, topk, H).sum(axis=1) - sorted_indices = state.sorted_indices - num_real = num_real_tokens - padding = padding_size - topk = num_experts_per_tok - unsort_indices = jnp.argsort(sorted_indices) - d_padded = d_sorted_x[unsort_indices] - if padding > 0: - d_replicated = d_padded[:num_real] - else: - d_replicated = d_padded - num_tokens = inputs_2d_shape[0] - hidden = inputs_2d_shape[-1] - d_inputs_2d = d_replicated.reshape(num_tokens, topk, hidden).sum(axis=1) - return d_inputs_2d - - # TRITON: bwd is unpermute_with_mask_map[_and_unpad]. - num_tokens = inputs_2d_shape[0] - hidden = inputs_2d_shape[-1] - if state.pad_offsets is not None: - d_inputs_2d, _ = unpermute_with_mask_map_and_unpad( - d_sorted_x, - state.row_id_map, - None, - None, - state.pad_offsets, - num_tokens, - num_experts, - hidden, - ) - else: - d_inputs_2d, _ = unpermute_with_mask_map( - d_sorted_x, - state.row_id_map, - None, - None, - num_tokens, - num_experts, - hidden, - ) - return d_inputs_2d # ============================================================================= -# Per-shard body +# Per-shard FFN body (runs inside shard_map) # ============================================================================= -def _body_fwd( # pylint: disable=unused-argument - captured: dict, +def _ffn_fwd_per_shard( + recv_tokens_local: jnp.ndarray, + recv_topk_weights_local: jnp.ndarray, + wi_0: jnp.ndarray, + wi_1: jnp.ndarray, + wo: jnp.ndarray, + wi_0_bias: Optional[jnp.ndarray], + wi_1_bias: Optional[jnp.ndarray], + wo_bias: Optional[jnp.ndarray], *, - # Statics - num_experts: int, - num_experts_per_tok: int, + num_local_experts: int, + slots_per_expert: int, activation_type: str, - score_function: ScoreFunction, - use_pre_softmax: bool, - num_groups: Optional[int], - group_topk: Optional[int], - scaling_factor: float, - aux_loss_coeff: float, - permutation_backend: PermutationBackend, - align_size: int, - gate_inside_vjp: bool, - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], - dtype: jnp.dtype, - # EP-only statics - ep_active: bool, - ep_axis: Optional[str], - data_parallelism_axes: Tuple[str, ...], - fsdp_sizes: Tuple[int, ...], - num_ep: int, - num_experts_local: int, - recv_buffer_rows: int, -) -> Tuple[jnp.ndarray, jnp.ndarray, dict]: - """Per-shard forward body. Returns ``(output, aux_loss, ctx_dict)``. - - ``aux_loss`` is always materialized (zeros scalar when disabled) so - the ``shard_map``'s ``out_specs`` has a static structure. - """ - if not gate_inside_vjp: - raise NotImplementedError( - "gate_inside_vjp=False is deferred to a follow-up PR; for now" - " the gate GEMM lives inside the MoE VJP." - ) - - x = captured["inputs"] - gate_kernel = captured["gate_kernel"] - wi_0 = captured["wi_0"] - wi_1 = captured["wi_1"] - wo = captured["wo"] - wi_0_bias = captured.get("wi_0_bias") - wi_1_bias = captured.get("wi_1_bias") - wo_bias = captured.get("wo_bias") - expert_bias = captured.get("expert_bias") - - batch_size, sequence_length, hidden = x.shape - - # ---------------- Stage 1: gate ---------------- - 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) - inputs_2d = x.reshape(-1, hidden) - - # ---------------- Stage 2: routing ---------------- - # Under EP, expert_bias is sharded P(ep_axis); the router needs the - # full E-dim view, so all_gather it. - if ep_active and expert_bias is not None: - full_expert_bias = jax.lax.all_gather(expert_bias, axis_name=ep_axis, tiled=True) - else: - full_expert_bias = expert_bias - # Pass an empty array sentinel when expert_bias is unused (the - # underlying primitive expects a real ndarray, not None). - eb_arg = ( - full_expert_bias if full_expert_bias is not None else jnp.zeros((0,), dtype=jnp.float32) - ) - sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( - logits_2d, - topk=num_experts_per_tok, - 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) - - # ---------------- Stage 2b: aux loss ---------------- - if aux_loss_coeff > 0.0: - if ep_active: - collective_axes: Any = ( - ep_axis if not data_parallelism_axes else (ep_axis, *data_parallelism_axes) - ) - global_logits_2d = jax.lax.all_gather( - logits_2d, axis_name=collective_axes, axis=0, tiled=True - ) - _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( - global_logits_2d, - topk=num_experts_per_tok, - 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) - aux_logits_for_score = global_logits_2d - else: - aux_tokens_per_expert = jnp.sum(routing_map.astype(jnp.int32), axis=0) - aux_logits_for_score = logits_2d - # Aux-side scores: clean per-expert scores (no grouped routing, - # no bias). compute_aux_scores=True takes a separate path that - # ignores the grouping knobs. - aux_probs, _aux_routing_map, aux_saved_scores = tex.fused_topk_with_score_function_fwd( - aux_logits_for_score.astype(jnp.float32), - topk=num_experts_per_tok, - 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=num_experts_per_tok, - coeff=aux_loss_coeff, - ) - else: - aux_loss = jnp.zeros((), dtype=dtype) - aux_const_buf = None - aux_tokens_per_expert = None - aux_logits_for_score = None - aux_saved_scores = None + apply_topk_weights_early: bool, +): + """Per-shard FFN forward. - # ---------------- Stage 3: dispatch ---------------- - shard_id = jax.lax.axis_index(ep_axis) if ep_active else None - sorted_x, dispatch_state = _dispatch( - inputs_2d, - sparse_probs, - routing_map, - backend=permutation_backend, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - ep_axis=ep_axis, - num_ep=num_ep, - recv_buffer_rows=recv_buffer_rows, - shard_id=shard_id, - ) - local_group_sizes = dispatch_state.group_sizes - - # ---------------- Stage 4: per-expert FFN (inlined) ---------------- - q_set_w0, q_set_w1, q_set_wo = quantizer_sets - if q_set_w0 == noop_quantizer_set: - wi_0 = wi_0.astype(sorted_x.dtype) - if q_set_w1 == noop_quantizer_set: - wi_1 = wi_1.astype(sorted_x.dtype) - if q_set_wo == noop_quantizer_set: - wo = wo.astype(sorted_x.dtype) - - # GEMM 1+2 (fused): up_proj_combined = sorted_x @ wi where - # wi := concat([wi_0, wi_1], axis=-1) -> shape [E, H, 2M] - # combined_out := sorted_x @ wi -> shape [T, 2M] - # Splitting the output back into ``gate_proj_out`` / ``up_proj_out`` - # is free (it's a slicing reshape). This collapses two grouped - # GEMMs and two grouped quantizes of ``sorted_x`` (one per kernel) - # into one of each. Bias is concatenated the same way. - # - # FP8/MXFP8 caveat: per-expert amax is now computed over [H, 2M] - # rather than [H, M] for each of wi_0 / wi_1 separately, so the - # representable range for one of the two halves may shift slightly - # vs. the pre-fusion code. Numerics tests cover this. - inter_M = wi_0.shape[-1] + Operates on the shard-local ``[1, recv_pr, H]`` slice that + ``tex.ep_dispatch`` produces. Returns the expert outputs (shaped + ``[1, recv_pr, H_out]`` so the surrounding ``shard_map`` reassembles + them as ``[num_procs, recv_pr, H_out]``) plus the residuals consumed + by the bwd. + """ + hidden = recv_tokens_local.shape[-1] + sorted_x = recv_tokens_local.reshape(-1, hidden) + recv_w_flat = recv_topk_weights_local.reshape(-1) + local_group_sizes = jnp.full((num_local_experts,), slots_per_expert, dtype=jnp.int32) + + wi_0 = wi_0.astype(sorted_x.dtype) + wi_1 = wi_1.astype(sorted_x.dtype) + wo = wo.astype(sorted_x.dtype) + + # wi GEMM uses ONE fused grouped_gemm with the gate/up weights + # concatenated along the trailing (output) axis: wi_combined has + # shape ``(num_local_experts, hidden, 2*H_inter)`` and the resulting + # combined_out has shape ``(num_rows, 2*H_inter)``, which jnp.split + # cleanly slices back into gate / up halves. tex.grouped_gemm only + # supports the canonical (G, K, N) 3D weight layout with + # contracting_dims=((1,),(1,)) -- see the docstring on + # transformer_engine.jax.dense.grouped_dense ("currently only + # supports ((1,), (1,))") and the CI test + # tests/jax/test_multi_process_distributed_grouped_gemm.py. + # An older fused 4D variant built via jnp.stack([wi_0, wi_1], axis=-2) + # put a non-contracting axis in the middle of the RHS, which the + # kernel walked as if it were 3D and read off the end -> NaN. + # Bisected against a jnp.einsum reference: the stack-axis variant + # produced all-NaN output, while the concat-axis variant (this + # path) produces finite outputs matching the reference. wi_combined = jnp.concatenate([wi_0, wi_1], axis=-1) wi_combined_bias = ( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) - casted_sorted_x = tex.grouped_quantize(sorted_x, q_set_w0.x, local_group_sizes, flatten_axis=-1) - casted_wi = tex.grouped_quantize(wi_combined, q_set_w0.kernel, flatten_axis=-1) + + q_set = noop_quantizer_set + casted_sorted_x = tex.grouped_quantize(sorted_x, q_set.x, local_group_sizes, flatten_axis=-1) + casted_wi = tex.grouped_quantize(wi_combined, q_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, ) - gate_proj_out = combined_out[..., :inter_M] - up_proj_out = combined_out[..., inter_M:] + gate_proj_out, up_proj_out = jnp.split(combined_out, 2, axis=-1) casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS) casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS) - if isinstance(casted_sorted_x_lhs_trans, ScaledTensor): - casted_sorted_x_lhs_trans = casted_sorted_x_lhs_trans.checkpoint(q_set_w0.x) - if isinstance(casted_wi_rhs_trans, ScaledTensor): - casted_wi_rhs_trans = casted_wi_rhs_trans.checkpoint(q_set_w0.kernel) - # Activation: intermediate = act(gate_proj_out) * up_proj_out + # Promote the silu+multiply to fp32 to match the pure-JAX reference + # (and ML common practice). bf16 silu accumulation alone drifts ~1% + # vs fp32 silu, which composes through wo -> combine into the + # ~1.4% per-element parity gap we were seeing on softmax. Cast back + # to the activation dtype before the grouped_quantize so the wo GEMM + # input layout is unchanged. act_fn = _convert_to_activation_function(activation_type) - intermediate = act_fn(gate_proj_out) * up_proj_out + intermediate = ( + act_fn(gate_proj_out.astype(jnp.float32)) + * up_proj_out.astype(jnp.float32) + ).astype(sorted_x.dtype) + + if apply_topk_weights_early: + # 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, modulo elementwise op fusion gains. w_b is + # cast to intermediate.dtype so the multiply doesn't promote + # expert_outputs to f32 (NCCL EP combine hard-asserts bf16). + w_b = recv_w_flat[:, None].astype(intermediate.dtype) + mask_b = (recv_w_flat != 0).astype(intermediate.dtype)[:, None] + intermediate = intermediate * w_b * mask_b - # GEMM 3: expert_outputs = intermediate @ wo casted_intermediate = tex.grouped_quantize( - intermediate, q_set_wo.x, local_group_sizes, flatten_axis=-1 + intermediate, q_set.x, local_group_sizes, flatten_axis=-1 ) - casted_wo = tex.grouped_quantize(wo, q_set_wo.kernel, flatten_axis=-1) + casted_wo = tex.grouped_quantize(wo, q_set.kernel, flatten_axis=-1) expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), casted_wo.get_tensor(usage=TensorUsage.RHS), @@ -1151,524 +373,135 @@ def _body_fwd( # pylint: disable=unused-argument ) casted_intermediate_lhs_trans = casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS) casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) - if isinstance(casted_intermediate_lhs_trans, ScaledTensor): - casted_intermediate_lhs_trans = casted_intermediate_lhs_trans.checkpoint(q_set_wo.x) - if isinstance(casted_wo_rhs_trans, ScaledTensor): - casted_wo_rhs_trans = casted_wo_rhs_trans.checkpoint(q_set_wo.kernel) - - # ---------------- Stage 5: combine ---------------- - # Compute per-shard static shape info once and pass through both - # _combine and (later) the bwd helpers via kwargs -- never via the - # state dict, which gets pytree-flattened across shard_map and would - # coerce Python ints into JitTracer 0-d arrays. - _static_shape = _compute_static_shape_info( - batch_size=batch_size, - sequence_length=sequence_length, - hidden=hidden, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - num_ep=num_ep, - fsdp_sizes=fsdp_sizes, - recv_buffer_rows=recv_buffer_rows, - ) - # ``expert_outputs_residual`` is the post-A2A FFN-output tensor that - # Step 3 of the combine actually consumed. Saving this (rather than - # the pre-A2A shard-local FFN output) is what makes - # ``_combine_bwd``'s Step-3 inverse see the same value the forward - # Step 3 saw -- otherwise EP + TRITON yields wrong d_expert_outputs. - output, expert_outputs_residual = _combine( - expert_outputs, - dispatch_state, - backend=permutation_backend, - ep_active=ep_active, - batch_size=batch_size, - sequence_length=sequence_length, - dtype=dtype, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) - # ---------------- Build ctx ---------------- - aux_enabled = aux_loss_coeff > 0.0 - ctx = _BodyCtx( - x=x, - gate_kernel=gate_kernel, - logits_2d=logits_2d, - saved_scores=saved_scores, - routing_map=routing_map, - dispatch=dispatch_state, - 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_residual, - local_group_sizes=local_group_sizes, - expert_bias=expert_bias if expert_bias is not None else None, - aux_const_buf=aux_const_buf if aux_enabled else None, - aux_tokens_per_expert=aux_tokens_per_expert if aux_enabled else None, - aux_logits_for_score=aux_logits_for_score if aux_enabled else None, - aux_saved_scores=aux_saved_scores if aux_enabled else None, + expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) + residuals = ( + 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, ) - - return output, aux_loss, ctx - - -def _body_bwd( # pylint: disable=unused-argument - ctx: _BodyCtx, - dy_pair: Tuple[jnp.ndarray, jnp.ndarray], + return expert_outputs_3d, residuals + + +def _ffn_bwd_per_shard( + d_expert_outputs_local: jnp.ndarray, + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out: jnp.ndarray, + up_proj_out: jnp.ndarray, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes: jnp.ndarray, + recv_topk_weights_local: jnp.ndarray, *, - num_experts: int, - num_experts_per_tok: int, activation_type: str, - score_function: ScoreFunction, - use_pre_softmax: bool, - num_groups: Optional[int], - group_topk: Optional[int], - scaling_factor: float, - aux_loss_coeff: float, - permutation_backend: PermutationBackend, - align_size: int, - gate_inside_vjp: bool, - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], - dtype: jnp.dtype, - ep_active: bool, - ep_axis: Optional[str], - data_parallelism_axes: Tuple[str, ...], - fsdp_sizes: Tuple[int, ...], - num_ep: int, - num_experts_local: int, - recv_buffer_rows: int, - # Static side info (kept here rather than inside ctx because they're - # python flags / shapes, not array leaves): - has_wi_bias: bool, - has_wo_bias: bool, - has_expert_bias: bool, - x_shape: Tuple[int, ...], -) -> dict: - """Per-shard backward body. Returns a dict of grads keyed identically - to the ``captured`` dict consumed by :func:`_body_fwd`.""" - if not gate_inside_vjp: - raise NotImplementedError("gate_inside_vjp=False is deferred to a follow-up PR.") - - d_output, d_aux_loss = dy_pair - # The fused FFN bwd quantizes via ``q_set_w0`` only (one quantize for - # the [E, H, 2M] fused wi tensor and one for the [T, 2M] fused dgrad), - # so ``q_set_w1`` is intentionally unused here. - q_set_w0, _q_set_w1, q_set_wo = quantizer_sets - batch_size, sequence_length, hidden = x_shape - shard_id = jax.lax.axis_index(ep_axis) if ep_active else None - - # Recompute per-shard static shape info from existing statics - # (Python ints / int tuples). Plumbed via kwargs to _combine_bwd - # and _dispatch_bwd -- NOT through the ctx dict, because the - # dict gets pytree-flattened across the bwd shard_map's in_specs - # and Python ints would be coerced into JitTracer 0-d arrays - # (breaking ``if padding > 0`` and ``jnp.zeros(shape)`` callsites). - # ``batch_size`` here is the GLOBAL batch size (captured in - # ``x_shape`` by the outer fwd rule), hence ``batch_is_per_shard=False``. - _static_shape = _compute_static_shape_info( - batch_size=batch_size, - sequence_length=sequence_length, - hidden=hidden, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - num_ep=num_ep, - fsdp_sizes=fsdp_sizes, - recv_buffer_rows=recv_buffer_rows, - batch_is_per_shard=False, - ) - - # Compute per-shard input shape: under the EP shard_map body, the - # gradient tensors live at per-shard shape, so the dispatch_bwd - # reshape target and ``d_x_from_dispatch.reshape(x_shape)`` below - # must use the per-shard shape rather than the captured global - # ``x_shape``. - if ep_active: - dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 - per_shard_batch = batch_size // (num_ep * dp_size) - per_shard_x_shape: Tuple[int, ...] = (per_shard_batch, sequence_length, hidden) - else: - per_shard_x_shape = x_shape - - # ---------------- Combine bwd ---------------- - d_expert_outputs, d_routing_weights = _combine_bwd( - d_output, - ctx.dispatch, - ctx.expert_outputs, - backend=permutation_backend, - ep_active=ep_active, - batch_size=batch_size, - sequence_length=sequence_length, - dtype=dtype, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - post_a2a_buffer_shape=_static_shape.post_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) + apply_topk_weights_early: bool, + has_bias: bool, +): + """Per-shard FFN backward. - # ---------------- FFN bwd: GEMM 3 (wo) ---------------- - casted_d_eo = tex.grouped_quantize( - d_expert_outputs, q_set_wo.dgrad, ctx.local_group_sizes, flatten_axis=-1 - ) + Mirrors :func:`_ffn_fwd_per_shard`. Returns + ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], d_wi_0, d_wi_1, d_wo, + d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. + """ + d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) + recv_w_flat = recv_topk_weights_local.reshape(-1) + q_set = noop_quantizer_set + + # wo bwd + casted_d_eo = tex.grouped_quantize(d_eo_2d, q_set.dgrad, local_group_sizes, flatten_axis=-1) + _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.get_tensor(usage=TensorUsage.LHS), - ctx.casted_wo_rhs_trans, + _casted_d_eo_lhs, + casted_wo_rhs_trans, contracting_dims=((1,), (2,)), ) d_wo = tex.grouped_gemm( - ctx.casted_intermediate_lhs_trans, - casted_d_eo.get_tensor(usage=TensorUsage.RHS), + casted_intermediate_lhs_trans, + _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) - d_wo_bias = tex.grouped_dbias(d_expert_outputs, ctx.local_group_sizes) if has_wo_bias else None + d_wo_bias = tex.grouped_dbias(d_eo_2d, local_group_sizes) if has_bias else None - # ---------------- Activation bwd ---------------- - # intermediate = act(gate_proj_out) * up_proj_out - # d(gate_proj_out) = vjp(act, gate_proj_out)(d_intermediate * up_proj_out) - # d(up_proj_out) = d_intermediate * act(gate_proj_out) act_fn = _convert_to_activation_function(activation_type) - act_gate_proj_out, dact_gate_proj_pullback = jax.vjp(act_fn, ctx.gate_proj_out) - d_up_proj_out = d_intermediate * act_gate_proj_out - (d_gate_proj_out,) = dact_gate_proj_pullback(d_intermediate * ctx.up_proj_out) - - # ---------------- FFN bwd: GEMM 1+2 fused (wi_0 | wi_1) ---------------- - # Concat the two upstream grads along the output (M) axis, do one - # grouped quantize + one dgrad GEMM + one wgrad GEMM, then split. - # ``ctx.casted_wi_rhs_trans`` has shape [E, H, 2M] from the fwd - # fused quantize, so the dgrad math is: - # d_sorted_x = [d_gate | d_up] @ wi_rhs_trans - # = d_gate @ wi_0^T + d_up @ wi_1^T - inter_M = d_gate_proj_out.shape[-1] + if apply_topk_weights_early: + # intermediate' = intermediate * w * mask. Split the cotangent + # across both factors before the activation bwd consumes it. + # Cast w_b so the multiply stays in d_intermediate.dtype and + # d_sorted_x (downstream into ep_dispatch_bwd) stays bf16. + w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) + mask_b = (recv_w_flat != 0).astype(d_intermediate.dtype)[:, None] + intermediate_unweighted = act_fn(gate_proj_out) * up_proj_out + d_recv_w_from_intermediate = jnp.sum( + d_intermediate * intermediate_unweighted * mask_b, axis=-1 + ).astype(recv_w_flat.dtype) + d_intermediate = d_intermediate * w_b * mask_b + else: + d_recv_w_from_intermediate = jnp.zeros_like(recv_w_flat) + + # Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply + # so the silu derivative composes through the gradient at fp32 too; + # cast back to the bf16 layout the wi grouped_quantize expects. + gp_fp32 = gate_proj_out.astype(jnp.float32) + up_fp32 = up_proj_out.astype(jnp.float32) + d_int_fp32 = d_intermediate.astype(jnp.float32) + act_gp_fp32, dact_pullback_fp32 = jax.vjp(act_fn, gp_fp32) + d_up_proj_out = (d_int_fp32 * act_gp_fp32).astype(up_proj_out.dtype) + (d_gate_proj_fp32,) = dact_pullback_fp32(d_int_fp32 * up_fp32) + d_gate_proj_out = d_gate_proj_fp32.astype(gate_proj_out.dtype) + + # 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 back into d_wi_0 / d_wi_1 halves with jnp.split. d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) casted_d_combined = tex.grouped_quantize( - d_combined, q_set_w0.dgrad, ctx.local_group_sizes, flatten_axis=-1 + d_combined, q_set.dgrad, local_group_sizes, flatten_axis=-1 ) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), - ctx.casted_wi_rhs_trans, + casted_wi_rhs_trans, contracting_dims=((1,), (2,)), ) d_wi_combined = tex.grouped_gemm( - ctx.casted_sorted_x_lhs_trans, + casted_sorted_x_lhs_trans, casted_d_combined.get_tensor(usage=TensorUsage.RHS), contracting_dims=((0,), (0,)), ) - d_wi_0 = d_wi_combined[..., :inter_M] - d_wi_1 = d_wi_combined[..., inter_M:] - if has_wi_bias: - d_wi_combined_bias = tex.grouped_dbias(d_combined, ctx.local_group_sizes) - d_wi_0_bias = d_wi_combined_bias[..., :inter_M] - d_wi_1_bias = d_wi_combined_bias[..., inter_M:] + d_wi_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1) + if has_bias: + d_wi_combined_bias = tex.grouped_dbias(d_combined, local_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 d_wi_1_bias = None - # ---------------- Dispatch bwd ---------------- - inputs_2d_shape = (per_shard_x_shape[0] * per_shard_x_shape[1], hidden) - d_inputs_2d = _dispatch_bwd( - d_sorted_x, - ctx.dispatch, - inputs_2d_shape=inputs_2d_shape, - backend=permutation_backend, - ep_active=ep_active, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) - d_x_from_dispatch = d_inputs_2d.reshape(per_shard_x_shape) - - # ---------------- Routing bwd ---------------- - # The probs cotangent comes from _combine_bwd. For PURE_JAX it's the - # cotangent of routing_weights (post-routing_map_to_selected_experts); - # we need to bridge back to sparse_probs. For TRITON it's already the - # cotangent of merging_probs == sparse_probs. - if d_routing_weights is not None: - if permutation_backend is PermutationBackend.PURE_JAX: - # routing_map_to_selected_experts: - # selected_experts = argsort(routing_map)[..., -topk:] - # weights = take_along_axis(sparse_probs, selected_experts, axis=-1) - # routing_map is bool (non-diff); the gradient of weights - # w.r.t. sparse_probs is a scatter-into-zero along the - # selected_experts indices. - selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -num_experts_per_tok:] - d_sparse_probs = jnp.zeros_like(ctx.saved_scores).astype(d_routing_weights.dtype) - d_sparse_probs = jnp.take_along_axis(d_sparse_probs, selected_experts, axis=-1) - # Actually scatter: build via jnp.zeros + .at[].set - d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_routing_weights.dtype) - d_sparse_probs = d_sparse_probs.at[ - jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts - ].set(d_routing_weights) - else: - d_sparse_probs = d_routing_weights.astype(jnp.float32) - else: - d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=jnp.float32) - - # Topk bwd primitive: returns d_logits (no d_expert_bias). - d_logits_2d_main = tex.fused_topk_with_score_function_bwd( - ctx.routing_map, - ctx.saved_scores, - d_sparse_probs.astype(ctx.saved_scores.dtype), - topk=num_experts_per_tok, - use_pre_softmax=use_pre_softmax, - scaling_factor=scaling_factor, - score_function=score_function, - compute_aux_scores=False, - ) - - # ---------------- Aux loss bwd ---------------- - if aux_loss_coeff > 0.0: - # Step 1: aux_loss bwd -> d_aux_probs - aux_num_tokens = ctx.aux_logits_for_score.shape[0] - d_aux_probs = tex.fused_moe_aux_loss_bwd( - ctx.aux_const_buf, - ctx.aux_tokens_per_expert.astype(jnp.int32), - d_aux_loss.reshape(()), - num_tokens=aux_num_tokens, - ) - # Step 2: aux-side topk bwd (compute_aux_scores=True path). - # The routing_map argument is ignored in this branch (the kernel - # uses saved_scores); pass any shape-correct integer tensor. - d_aux_logits = tex.fused_topk_with_score_function_bwd( - jnp.zeros(ctx.aux_logits_for_score.shape, dtype=jnp.bool_), - ctx.aux_saved_scores, - d_aux_probs.astype(ctx.aux_saved_scores.dtype), - topk=num_experts_per_tok, - use_pre_softmax=False, - scaling_factor=1.0, - score_function=score_function, - compute_aux_scores=True, - ) - # Step 3: under EP the aux logits were all_gathered along - # ``(ep_axis, *data_parallelism_axes)`` (the latter being FSDP - # axes that shard the batch). The bwd is the inverse of that - # multi-axis tiled all_gather: ``dynamic_slice`` to pick out - # this shard's local rows from the global cotangent. - # - # JAX's convention for tiled ``all_gather(axis_name=(a, b, ...))`` - # is row-major over the tuple: the shard at mesh position - # ``(i_a, i_b, ...)`` writes to rows - # ``[(i_a * size_b * ... + i_b * ... + ...) * local_T : - # + local_T)``. We invert that by computing the same flat - # index here and slicing. - if ep_active: - local_T_aux = ctx.logits_2d.shape[0] - flat_shard = shard_id # ep is the outermost axis in the gather tuple - for ax, sz in zip(data_parallelism_axes, fsdp_sizes): - flat_shard = flat_shard * sz + jax.lax.axis_index(ax) - d_aux_logits_local = jax.lax.dynamic_slice( - d_aux_logits.astype(ctx.logits_2d.dtype), - start_indices=(flat_shard * local_T_aux, 0), - slice_sizes=(local_T_aux, num_experts), - ) - else: - d_aux_logits_local = d_aux_logits.astype(d_logits_2d_main.dtype) - d_logits_2d = d_logits_2d_main + d_aux_logits_local.astype(d_logits_2d_main.dtype) - else: - d_logits_2d = d_logits_2d_main - - # ---------------- Gate bwd ---------------- - d_gate_logits = d_logits_2d.reshape(per_shard_x_shape[0], per_shard_x_shape[1], 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 - - # Reduce per-rank partial contributions to match the out_specs - # declared by _build_grads_specs: - # gate_kernel : P() -> psum across (ep, *fsdp) - # wi_0/wi_1/wo : P(ep_axis, ...) -> psum across (*fsdp) only - # inputs : P((ep, fsdp), ...) -> already shard-local, no reduction - if ep_active: - replicate_all = (ep_axis,) + tuple(data_parallelism_axes) - d_gate_kernel = jax.lax.psum(d_gate_kernel, axis_name=replicate_all) - if data_parallelism_axes: - replicate_fsdp = tuple(data_parallelism_axes) - d_wi_0 = jax.lax.psum(d_wi_0, axis_name=replicate_fsdp) - d_wi_1 = jax.lax.psum(d_wi_1, axis_name=replicate_fsdp) - d_wo = jax.lax.psum(d_wo, axis_name=replicate_fsdp) - if has_wi_bias: - d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=replicate_fsdp) - d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=replicate_fsdp) - if has_wo_bias: - d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=replicate_fsdp) - - grads: dict = { - "inputs": d_x, - "gate_kernel": d_gate_kernel, - "wi_0": d_wi_0, - "wi_1": d_wi_1, - "wo": d_wo, - } - if has_wi_bias: - grads["wi_0_bias"] = d_wi_0_bias - grads["wi_1_bias"] = d_wi_1_bias - if has_wo_bias: - grads["wo_bias"] = d_wo_bias - if has_expert_bias: - # expert_bias has no gradient through topk (the topk bwd returns - # None for it). Emit a structural zero so the outer rule has - # something to package. - grads["expert_bias"] = jnp.zeros_like(ctx.expert_bias) - return grads - - -# ============================================================================= -# Spec builders for shard_map (lockstep with ctx_dict / captured_dict) -# ============================================================================= - - -def _build_in_specs( - ep_axis: str, - batch_pspec_axis: Any, - *, - has_bias: bool, - has_expert_bias: bool, -) -> dict: - """Build the ``in_specs`` dict for the EP fwd shard_map.""" - specs: dict = { - "inputs": P(batch_pspec_axis, None, None), - "gate_kernel": P(), - "wi_0": P(ep_axis, None, None), - "wi_1": P(ep_axis, None, None), - "wo": P(ep_axis, None, None), - } - if has_bias: - for name in ("wi_0_bias", "wi_1_bias", "wo_bias"): - specs[name] = P(ep_axis, None) - if has_expert_bias: - specs["expert_bias"] = P(ep_axis) - return specs - - -def _build_dispatch_specs( # pylint: disable=unused-argument - ep_axis: str, - *, - backend: PermutationBackend, - ep_active: bool, - align_size: int, -) -> _DispatchState: - """Build the shard_map ``out_specs`` for the dispatch state. - - Returns a :data:`_DispatchState` (either :class:`_PureJaxDispatchState` - or :class:`_TritonDispatchState`) whose fields are - :class:`PartitionSpec` placeholders. Optional fields are set to - ``P()`` when populated by :func:`_dispatch` and to ``None`` when - intentionally omitted, so the spec's pytree structure mirrors the - value's structure leaf-for-leaf. - """ - ep_all = P() if ep_active else None - ep_local = P() if ep_active else None - if backend is PermutationBackend.PURE_JAX: - return _PureJaxDispatchState( - group_sizes=P(), - sorted_indices=P(), - routing_weights=P(), - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - return _TritonDispatchState( - group_sizes=P(), - row_id_map=P(), - pad_offsets=P() if align_size > 0 else None, - merging_probs=P(), - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - - -def _build_ctx_specs( # pylint: disable=unused-argument - ep_axis: str, - batch_pspec_axis: Any, - *, - backend: PermutationBackend, - ep_active: bool, - has_bias: bool, - has_expert_bias: bool, - aux_loss_enabled: bool, - align_size: int, -) -> _BodyCtx: - """Build the spec :class:`_BodyCtx` mirroring :func:`_body_fwd`'s ctx. - - Fields gated off by the static config (``expert_bias``, ``aux_*``) - are ``None`` here so the spec pytree matches the value pytree - leaf-for-leaf. - """ - return _BodyCtx( - # Per-shard local activations along the batch axis. - x=P(batch_pspec_axis, None, None), - gate_kernel=P(), - logits_2d=P(batch_pspec_axis, None), - saved_scores=P(batch_pspec_axis, None), - routing_map=P(batch_pspec_axis, None), - dispatch=_build_dispatch_specs( - ep_axis, backend=backend, ep_active=ep_active, align_size=align_size - ), - # FFN residuals: the LHS_TRANS / RHS_TRANS variants of - # grouped_quantize have leading "rows"/"experts" dims that are - # already shard-local (post-dispatch). Use P(ep_axis,...) on - # leading dim; that works whether the leaf is a plain ndarray - # or a ScaledTensor (shard_map applies the spec leaf-wise to - # the registered ScaledTensor pytree). - casted_sorted_x_lhs_trans=P(), - casted_wi_rhs_trans=P(ep_axis, None, None), - gate_proj_out=P(), - up_proj_out=P(), - casted_intermediate_lhs_trans=P(), - casted_wo_rhs_trans=P(ep_axis, None, None), - expert_outputs=P(), - local_group_sizes=P(), - expert_bias=P(ep_axis) if has_expert_bias else None, - aux_const_buf=P() if aux_loss_enabled else None, - aux_tokens_per_expert=P() if aux_loss_enabled else None, - aux_logits_for_score=P() if aux_loss_enabled else None, - aux_saved_scores=P() if aux_loss_enabled else None, - ) - - -def _build_grads_specs( - ep_axis: str, - batch_pspec_axis: Any, - *, - has_bias: bool, - has_expert_bias: bool, -) -> dict: - """Spec dict for the grads dict returned by :func:`_body_bwd`.""" - return _build_in_specs( - ep_axis, - batch_pspec_axis, - has_bias=has_bias, - has_expert_bias=has_expert_bias, + d_sorted_x_3d = d_sorted_x.reshape(1, d_sorted_x.shape[0], d_sorted_x.shape[1]) + d_recv_w_3d = d_recv_w_from_intermediate.reshape(1, -1) + return ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, ) # ============================================================================= -# Top-level VJP rules +# Full fwd / bwd rules (custom_vjp halves) # ============================================================================= -def _moe_fwd_rule( # pylint: disable=unused-argument - # Args MUST match the positional order of ``_moe`` (diff first, - # then nondiff). See ``_moe_bwd_rule`` for the opposite convention. +def _moe_fwd_rule( x, gate_kernel, wi_0, @@ -1687,170 +520,350 @@ def _moe_fwd_rule( # pylint: disable=unused-argument group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, + align_size, ): - x = with_sharding_constraint_by_logical_axes(x, input_axes) - ep_active = ep_axis is not None - body_kwargs = { - "num_experts": num_experts, - "num_experts_per_tok": num_experts_per_tok, - "activation_type": activation_type, - "score_function": score_function, - "use_pre_softmax": use_pre_softmax, - "num_groups": num_groups, - "group_topk": group_topk, - "scaling_factor": scaling_factor, - "aux_loss_coeff": aux_loss_coeff, - "permutation_backend": permutation_backend, - "align_size": align_size, - "gate_inside_vjp": gate_inside_vjp, - "quantizer_sets": quantizer_sets, - "dtype": dtype, - "ep_axis": ep_axis, - "data_parallelism_axes": data_parallelism_axes, - } - captured: dict = { - "inputs": x, - "gate_kernel": gate_kernel, - "wi_0": wi_0, - "wi_1": wi_1, - "wo": wo, - } - has_bias = wi_0_bias is not None - has_expert_bias = expert_bias is not None - if has_bias: - captured["wi_0_bias"] = wi_0_bias - captured["wi_1_bias"] = wi_1_bias - captured["wo_bias"] = wo_bias - if has_expert_bias: - captured["expert_bias"] = expert_bias - - if not ep_active: - output, aux_loss, ctx = _body_fwd( - captured, - **body_kwargs, - ep_active=False, - fsdp_sizes=(), - num_ep=1, - num_experts_local=num_experts, - recv_buffer_rows=0, - ) - # Carry static side info to the bwd rule alongside ctx. These - # are Python ints/bools/tuples (NOT pytree leaves), so we - # bundle them as a plain dict rather than putting them on the - # ``_BodyCtx`` NamedTuple where shard_map would try to flatten - # them into JitTracers. - static = { - "has_wi_bias": has_bias, - "has_wo_bias": has_bias, - "has_expert_bias": has_expert_bias, - "x_shape": x.shape, - "num_experts_local": num_experts, - "recv_buffer_rows": 0, - } - return (output, aux_loss), (ctx, static) - - # ---------------- EP path ---------------- + """Forward: gate -> topk -> ep_dispatch -> shard_map(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 when ep_axis is set.") + 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_experts_local = num_experts // num_ep + num_local_experts = num_experts // num_ep - # Reject overlapping EP / FSDP axes. Listing ep_axis in - # data_parallelism_axes would produce a duplicate-axis PartitionSpec - # ((ep, ep, ...)) which JAX rejects, and would also double-count - # num_ep in dp_size (under-sizing recv_buffer_rows by a factor of - # num_ep). Catch it up front with a clear error. + dp_size = 1 for ax in data_parallelism_axes: - if ax not in mesh.shape: - raise ValueError( - f"data_parallelism_axes contains {ax!r} but mesh has" - f" axes {tuple(mesh.shape.keys())}" - ) - if ax == ep_axis: - raise ValueError( - f"data_parallelism_axes={data_parallelism_axes!r} contains the EP" - f" axis {ep_axis!r}; EP is implicit in the batch sharding and must" - " not also be listed as a data-parallel axis." - ) + dp_size *= mesh.shape[ax] + num_procs = num_ep * dp_size + + 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 + # Per-rank receive capacity. NCCL EP HT lays out the per-rank receive + # buffer as ``[num_local_experts, num_ep * max_tokens_per_rank, hidden]`` + # (see nccl_ep.cc::init kernel buffer sizing + the LL combine assertion + # at nccl_ep.cc:2185 which spells out the same layout). The natural + # dropless K-expanded count + # ``ceil((B/dp)*S*K / num_local_experts)`` does NOT match: it ignores + # the worst-case where all of one EP group's tokens land on a single + # local expert. We must size to that worst case or NCCL EP's HT kernel + # rejects the dispatch buffer with ``invalid argument``. + natural_spe = num_ep * max_tokens_per_rank # = (B // dp_size) * S + # NCCL EP requires each expert-major output block to be at least + # 128-token aligned. Keep larger caller-requested alignments, but + # do not emit a smaller natural block size for tiny tests. + effective_align = max(int(align_size), 128) + slots_per_expert = ((natural_spe + effective_align - 1) // effective_align) * effective_align + recv_pr = num_local_experts * slots_per_expert + + _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: - batch_pspec_axis = (ep_axis, *data_parallelism_axes) - dp_size = 1 - for ax in data_parallelism_axes: - dp_size *= mesh.shape[ax] + # 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) - global_batch_size, sequence_length, _hidden = x.shape - topk = num_experts_per_tok - if global_batch_size % (num_ep * dp_size) != 0: - raise ValueError(f"batch={global_batch_size} not divisible by ep*dp={num_ep * dp_size}") - recv_buffer_rows = (global_batch_size // dp_size) * sequence_length * topk - if align_size > 0: - recv_buffer_rows += num_experts * (align_size - 1) + # ---------------- 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, + ) + # Sigmoid + K>1 normalises as `weights / (weights.sum + 1e-20)`; for + # tokens whose top-K sigmoid scores all underflow at bf16/fp32 the + # output is NaN at the selected positions. Those NaNs ride + # ep_dispatch -> recv_topk_weights -> combine and poison the per-token + # weighted sum, leaving entire output rows as NaN. Sanitize at the + # source so neither the fwd combine nor the bwd's manual + # `grad_pre_combine * w` sees them. Padded positions in sparse_probs + # are already zero (routing_map is False there); only the rare + # underflow path emits NaN. + sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, 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 - in_specs = _build_in_specs( - ep_axis, - batch_pspec_axis, - has_bias=has_bias, - has_expert_bias=has_expert_bias, + # ---------------- 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) ) - output_spec = P(batch_pspec_axis, None, None) - aux_spec = P() - ctx_spec = _build_ctx_specs( - ep_axis, - batch_pspec_axis, - backend=permutation_backend, - ep_active=True, - has_bias=has_bias, - has_expert_bias=has_expert_bias, - aux_loss_enabled=(aux_loss_coeff > 0.0), - align_size=align_size, + 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=slots_per_expert, + ) + token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( + cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr + ) + 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) if has_bias else None + ffn_in_specs = (ep3_spec, ep2_spec, kernel_spec, kernel_spec, kernel_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, wi_0, wi_1, wo] + if has_bias: + ffn_in_specs = ffn_in_specs + (bias_spec, bias_spec, bias_spec) + ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) + + # FFN residuals live entirely on the local ep rank, so the leading + # "experts" / "rows" dims map to P() (already shard-local). wi is + # fused via jnp.concatenate along the trailing (output) axis + # (see _ffn_fwd_per_shard for rationale), so the residual is a + # single 3D casted_wi_rhs_trans of shape + # (num_local_experts, hidden, 2*H_inter). + residuals_spec = ( + P(), # casted_sorted_x_lhs_trans + P(ep_axis, None, None), # casted_wi_rhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans + P(ep_axis, None, None), # casted_wo_rhs_trans + P(), # local_group_sizes ) + out_specs = (ep3_spec, residuals_spec) - _fsdp_sizes: Tuple[int, ...] = tuple(mesh.shape[ax] for ax in data_parallelism_axes) - - def _shardmap_body(captured_local): - return _body_fwd( - captured_local, - **body_kwargs, - ep_active=True, - fsdp_sizes=_fsdp_sizes, - num_ep=num_ep, - num_experts_local=num_experts_local, - recv_buffer_rows=recv_buffer_rows, + def _body(*args): + if has_bias: + (r_tok, r_w, w0, w1, w_o, w0b, w1b, wob) = args + else: + (r_tok, r_w, w0, w1, w_o) = args + w0b = w1b = wob = None + # Per-rank conditional zero-init of r_tok. Works around a + # narrowly-scoped tex.ep_dispatch_fwd contract gap: the NCCL EP + # HT dispatch kernel zero-initialises the recv buffer correctly + # on ranks that receive at least one token, but leaves + # uninitialised memory on fully-empty-receiver ranks. ``r_w`` + # (the dispatch's own written-or-not indicator: 0 at padded + # slots, non-zero at real-routed slots) gives us a per-shard + # predicate for free. ``jax.lax.cond`` only executes the + # selected branch, so loaded ranks pay nothing at runtime; + # only empty ranks do the zero-fill. + # TODO: remove once tex.ep_dispatch_fwd zero-inits empty-rank + # recv buffers upstream. + rank_has_tokens = jnp.any(r_w != 0) + r_tok = jax.lax.cond( + rank_has_tokens, + lambda x: x, + lambda x: jnp.zeros_like(x), + r_tok, + ) + return _ffn_fwd_per_shard( + r_tok, + r_w, + w0, + w1, + w_o, + w0b, + w1b, + wob, + num_local_experts=num_local_experts, + slots_per_expert=slots_per_expert, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, ) - output, aux_loss, ctx = shard_map( - _shardmap_body, + expert_outputs, ffn_residuals = shard_map( + _body, mesh=mesh, - in_specs=(in_specs,), - out_specs=(output_spec, aux_spec, ctx_spec), + in_specs=ffn_in_specs, + out_specs=out_specs, check_rep=False, - )(captured) + )(*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: + # IEEE 754: NaN * 0 = NaN, so a multiplicative mask cannot kill + # the NaNs ep_dispatch_fwd leaves at padded slots of recv_tokens + # (they ride through the FFN into expert_outputs at the same + # padded positions): mean=NaN on expert_outputs[padded] then + # propagates into the combine output when the kernel's read + # pattern overlaps the padded region. Use jnp.where to overwrite + # padded positions with a literal 0 before combine. + w = recv_topk_weights[..., None].astype(expert_outputs.dtype) + mask_bool = (recv_topk_weights != 0)[..., None] + weighted = jnp.where(mask_bool, expert_outputs * w, jnp.zeros_like(expert_outputs)) + output = tex.ep_combine_fwd( + cfg, + handle_mem, + weighted, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + + ( + 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, + token_counts=token_counts, + 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, + aux_const_buf=aux_const_buf, + aux_tokens_per_expert=aux_tokens_per_expert, + aux_saved_scores=aux_saved_scores, + ) static = { - "has_wi_bias": has_bias, - "has_wo_bias": has_bias, - "has_expert_bias": has_expert_bias, + "has_bias": has_bias, "x_shape": x.shape, - "num_experts_local": num_experts_local, - "recv_buffer_rows": recv_buffer_rows, + "recv_pr": recv_pr, } return (output, aux_loss), (ctx, static) @@ -1865,128 +878,288 @@ def _moe_bwd_rule( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, - ctx, - dy_pair, + apply_topk_weights_early, + align_size, + residuals, + cotangents, ): - ctx, static = ctx # split tensor residuals from static side info - has_wi_bias = static["has_wi_bias"] - has_wo_bias = static["has_wo_bias"] - has_expert_bias = static["has_expert_bias"] - x_shape = static["x_shape"] - num_experts_local = static["num_experts_local"] - recv_buffer_rows = static["recv_buffer_rows"] + """Backward mirror of :func:`_moe_fwd_rule`.""" + del num_groups, group_topk, dtype, align_size # captured in residuals / unused in bwd + from jax.experimental.shard_map import shard_map - ep_active = ep_axis is not None - mesh = _get_mesh() if ep_active else None - fsdp_sizes: Tuple[int, ...] = ( - tuple(mesh.shape[ax] for ax in data_parallelism_axes) if ep_active else () - ) - body_kwargs = { - "num_experts": num_experts, - "num_experts_per_tok": num_experts_per_tok, - "activation_type": activation_type, - "score_function": score_function, - "use_pre_softmax": use_pre_softmax, - "num_groups": num_groups, - "group_topk": group_topk, - "scaling_factor": scaling_factor, - "aux_loss_coeff": aux_loss_coeff, - "permutation_backend": permutation_backend, - "align_size": align_size, - "gate_inside_vjp": gate_inside_vjp, - "quantizer_sets": quantizer_sets, - "dtype": dtype, - "ep_axis": ep_axis, - "data_parallelism_axes": data_parallelism_axes, - "fsdp_sizes": fsdp_sizes, - "num_ep": 1 if not ep_active else mesh.shape[ep_axis], - "num_experts_local": num_experts_local, - "recv_buffer_rows": recv_buffer_rows, - "has_wi_bias": has_wi_bias, - "has_wo_bias": has_wo_bias, - "has_expert_bias": has_expert_bias, - "x_shape": x_shape, - } + d_output, d_aux_loss = cotangents - if not ep_active: - grads = _body_bwd(ctx, dy_pair, ep_active=False, **body_kwargs) - # Apply sharding constraints on grads. - grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( - grads["gate_kernel"], gate_kernel_axes - ) - grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) - grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) - grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) - grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) - return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + ctx, static = residuals + has_bias = static["has_bias"] + x_shape = static["x_shape"] + recv_pr = static["recv_pr"] - from jax.experimental.shard_map import shard_map + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + num_ep = mesh.shape[ep_axis] + dp_size = 1 + for ax in data_parallelism_axes: + dp_size *= mesh.shape[ax] + B, S, _ = x_shape + K = num_experts_per_tok if not data_parallelism_axes: batch_pspec_axis: Any = ep_axis else: - batch_pspec_axis = (ep_axis, *data_parallelism_axes) - ctx_spec = _build_ctx_specs( - ep_axis, - batch_pspec_axis, - backend=permutation_backend, - ep_active=True, - has_bias=has_wi_bias, - has_expert_bias=has_expert_bias, - aux_loss_enabled=(aux_loss_coeff > 0.0), - align_size=align_size, + 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: + # combine_fwd consumed weighted = expert_out * w * mask; + # split the cotangent across both factors. w is cast to + # grad_pre_combine.dtype so the multiply stays bf16 and + # d_sorted_x (downstream into ep_dispatch_bwd) stays bf16. + # + # ep_dispatch_fwd can land NaN into recv_topk_weights on padded + # slots (the public NCCL EP HT path does not zero-fill unused + # recv buffer entries). Untreated, `(NaN != 0) == True` in IEEE, + # so the multiplicative mask cannot suppress the NaN and it + # propagates through grad_pre_combine * w * mask into d_expert_outputs + # and then into every downstream gradient (gate_kernel ends up + # all-NaN). Sanitize once here. + recv_w_clean = jnp.where(jnp.isnan(ctx.recv_topk_weights), 0, ctx.recv_topk_weights) + # IEEE 754: NaN * 0 = NaN, so multiplying grad_pre_combine by a + # 0/1 mask cannot kill the NaNs tex.ep_combine_bwd leaves at + # padded slots of grad_pre_combine: ctx.recv_topk_weights is + # clean after the sanitize above, but grad_pre_combine[padded] + # is still NaN, so grad_pre_combine * w * mask = NaN. Use + # jnp.where to overwrite padded positions with literal 0 + # instead. + w = recv_w_clean[..., None].astype(grad_pre_combine.dtype) + mask_bool = (recv_w_clean != 0)[..., None] + d_expert_outputs = jnp.where( + mask_bool, grad_pre_combine * w, jnp.zeros_like(grad_pre_combine) + ) + # Same masking strategy for the cotangent on recv_topk_weights: + # grad_pre_combine has NaN at padded slots and ctx.expert_outputs + # may too, so the per-element product must be jnp.where'd before + # the sum reduction. + d_recv_w_from_combine = jnp.where( + mask_bool, + grad_pre_combine * ctx.expert_outputs, + jnp.zeros_like(grad_pre_combine), + ).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) if has_bias else None + + bwd_in_specs = ( + ep3_spec, # d_expert_outputs + P(), # casted_sorted_x_lhs_trans + P(ep_axis, None, None), # casted_wi_rhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans + P(ep_axis, None, None), # casted_wo_rhs_trans + P(), # local_group_sizes + ep2_spec, # recv_topk_weights ) - dy_specs = (P(batch_pspec_axis, None, None), P()) - grads_spec = _build_grads_specs( - ep_axis, batch_pspec_axis, has_bias=has_wi_bias, has_expert_bias=has_expert_bias + 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, + ] + bwd_out_specs = ( + ep3_spec, # d_sorted_x + ep2_spec, # d_recv_w_from_intermediate + kernel_spec, # d_wi_0 + kernel_spec, # d_wi_1 + kernel_spec, # d_wo + bias_spec if has_bias else None, # d_wi_0_bias + bias_spec if has_bias else None, # d_wi_1_bias + bias_spec if has_bias else None, # d_wo_bias ) - def _bwd_body(ctx_local, dy_local): - return _body_bwd(ctx_local, dy_local, ep_active=True, **body_kwargs) + def _bwd_body(*args): + ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = _ffn_bwd_per_shard( + *args, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + has_bias=has_bias, + ) + # Weight grads accumulate per-DP-shard inside the body; psum across + # DP axes so each replica sees the full sum (matches out_specs + # P(ep_axis, ...) which is DP-replicated). + if data_parallelism_axes: + dp = tuple(data_parallelism_axes) + d_wi_0 = jax.lax.psum(d_wi_0, axis_name=dp) + d_wi_1 = jax.lax.psum(d_wi_1, axis_name=dp) + d_wo = jax.lax.psum(d_wo, axis_name=dp) + if has_bias: + d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=dp) + d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=dp) + d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=dp) + return ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) - grads = shard_map( + ( + d_sorted_x, + d_recv_w_from_intermediate, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = shard_map( _bwd_body, mesh=mesh, - in_specs=(ctx_spec, dy_specs), - out_specs=grads_spec, + in_specs=bwd_in_specs, + out_specs=bwd_out_specs, check_rep=False, - )(ctx, dy_pair) + )(*bwd_in_args) + + 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, + ) - grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( - grads["gate_kernel"], gate_kernel_axes + # ---------------- 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, ) - grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) - grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) - grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) - grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) - return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + # ---------------- 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_0 = with_sharding_constraint_by_logical_axes(d_wi_0, wi_kernel_axes) + d_wi_1 = with_sharding_constraint_by_logical_axes(d_wi_1, wi_kernel_axes) + d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_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) -def _grads_dict_to_tuple( - grads: dict, has_wi_bias: bool, has_wo_bias: bool, has_expert_bias: bool -) -> Tuple: - """Pack the body_bwd's grads dict into the positional tuple JAX expects.""" return ( - grads["inputs"], - grads["gate_kernel"], - grads["wi_0"], - grads["wi_1"], - grads["wo"], - grads.get("wi_0_bias") if has_wi_bias else None, - grads.get("wi_1_bias") if has_wi_bias else None, - grads.get("wo_bias") if has_wo_bias else None, - grads.get("expert_bias") if has_expert_bias else None, + d_x, + d_gate_kernel, + d_wi_0, + d_wi_1, + 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, ) @@ -1995,7 +1168,7 @@ def _grads_dict_to_tuple( # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 29))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 27))) def _moe( x, gate_kernel, @@ -2015,23 +1188,17 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, + align_size, ): - # Call in `_moe`'s own signature order to match what JAX will pass - # the fwd rule via ``_argnums_partial``. See the comment block at - # the top of ``_moe_fwd_rule`` for why this differs from - # ``_moe_bwd_rule``'s convention. - output_pair, _ = _moe_fwd_rule( + primal, _ = _moe_fwd_rule( x, gate_kernel, wi_0, @@ -2050,19 +1217,17 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, + align_size, ) - return output_pair + return primal _moe.defvjp(_moe_fwd_rule, _moe_bwd_rule) @@ -2079,56 +1244,90 @@ def moe( wo_bias: Optional[jnp.ndarray] = None, expert_bias: Optional[jnp.ndarray] = None, *, - # Architecture num_experts: int, num_experts_per_tok: int, activation_type: str = "silu", - # Routing score_function: Union[str, ScoreFunction] = "softmax", use_pre_softmax: bool = False, num_groups: Optional[int] = None, group_topk: Optional[int] = None, scaling_factor: float = 1.0, aux_loss_coeff: float = 0.0, - # Permutation - permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX, + apply_topk_weights_early: bool = False, align_size: int = 0, - # Gate placement (Phuong: "perhaps as an option") - gate_inside_vjp: bool = True, - # Parallelism (resolved by caller from MeshResource) - ep_axis: Optional[str] = None, + ep_axis: str, data_parallelism_axes: Tuple[str, ...] = (), - # Logical axes for sharding constraints input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed"), - # Quantization - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet] = ( - noop_quantizer_set, - noop_quantizer_set, - noop_quantizer_set, - ), dtype: jnp.dtype = jnp.float32, ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Run a full MoE block under a single fused custom_vjp. + """Run a full MoE block under a single fused custom_vjp on the TE EP path. - Parameters and return are documented at the call site of - ``_MoEBlock.__call__``. See module docstring for design rationale. + Returns ``(output, aux_loss)``. ``aux_loss`` is ``None`` when + ``aux_loss_coeff == 0`` and a 0-d scalar otherwise. + + Parameters + ---------- + expert_bias : Optional[jnp.ndarray] + ``[num_experts]`` learnable router bias added before the top-k + when ``score_function='sigmoid'``. Pass ``None`` to disable. + The bias has no gradient through the top-k primitive itself (it + only steers expert selection); a zero cotangent is returned for + it. + aux_loss_coeff : float + Per-step expert-load-balance loss coefficient. ``0.0`` (default) + disables the aux loss entirely. When non-zero, an extra + all-gather over the routing-side logits is inserted so the + ``fused_moe_aux_loss`` kernel sees a global ``[T_global, E]`` + view; this lives off the dispatch critical path. + align_size : int + Minimum per-expert slot alignment passed to ``tex.ep_prepare`` + as ``dispatch_output_per_expert_alignment``. ``0`` (default) + means use the NCCL-EP-required natural slot count + ``ep_size * max_tokens_per_rank == (B/dp)*S`` (the per-rank + all-tokens-to-one-expert worst case the HT kernel demands). + Any positive value rounds that count up to the nearest + multiple, growing the per-rank receive buffer accordingly. + Set to ``128`` for FP8 recipes that require 128-aligned + grouped-GEMM tiles. + + See module docstring for the rest of the parameter semantics and the + surrounding design rationale. """ - if not isinstance(permutation_backend, PermutationBackend): - raise TypeError( - f"permutation_backend must be a PermutationBackend, got {permutation_backend!r}" - ) - if permutation_backend is PermutationBackend.TRITON: - _require_triton() - # Normalize string score_function ("softmax" / "sigmoid") to the - # ScoreFunction enum once here. The underlying primitive - # ``tex.fused_topk_with_score_function_fwd`` expects an int-coercible - # value (the enum has integer .value), and the public router wrapper - # we bypass also normalizes here. 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, + ) + 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 + output, aux_loss = _moe( x, gate_kernel, @@ -2138,27 +1337,25 @@ def moe( wi_0_bias, wi_1_bias, wo_bias, - expert_bias, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - activation_type=activation_type, - score_function=score_function, - use_pre_softmax=use_pre_softmax, - num_groups=num_groups, - group_topk=group_topk, - scaling_factor=scaling_factor, - aux_loss_coeff=aux_loss_coeff, - permutation_backend=permutation_backend, - align_size=align_size, - gate_inside_vjp=gate_inside_vjp, - ep_axis=ep_axis, - data_parallelism_axes=data_parallelism_axes, - input_axes=input_axes, - gate_kernel_axes=gate_kernel_axes, - wi_kernel_axes=wi_kernel_axes, - wo_kernel_axes=wo_kernel_axes, - quantizer_sets=quantizer_sets, - dtype=dtype, + expert_bias_arg, + 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, + align_size, ) if aux_loss_coeff <= 0.0: aux_loss = None From 006902e8b6063e56554ed037e226ba06ceab3b8b Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 14:55:13 -0700 Subject: [PATCH 02/23] tests/jax: trim TE-EP MoE suite (drop bootstrap, flax-wrapper, bias-zero) * drop ``TestZZZTeEpMoeBootstrap``: the re-bootstrap mismatch is a one-line guard in ``ep_bootstrap`` and not the MoE block's concern; exercising it from this suite also taints the per-process NCCL bootstrap cache for the rest of the file with no real upside. * drop ``TestTeEpMoEBlockFlax::test_init_apply_parity``: every config in ``_CONFIGS`` already runs ``MoEBlock`` (the Flax wrapper) end-to-end via ``test_forward`` / ``test_backward``, so this was a duplicate of ``softmax`` parity in another wrapper -- leave wrapper refactors to devs without paying for an extra CI run each time. * drop ``sigmoid-bias-zero``: with a zero-init bias buffer the routing math collapses to the no-bias case, so ``sigmoid`` already covers that numerical path. The bias-aware codepath is still exercised by ``sigmoid-bias-strong`` (non-zero bias). * refresh the module-level docstring to list intentional non-coverage so future readers don't re-add these tests. Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 135 ++++++++---------------------------- 1 file changed, 28 insertions(+), 107 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index a5ab1c266b..75326af1c6 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -34,17 +34,24 @@ classes: * ``test_forward`` covers the forward across a curated set of - configurations (apply_topk_weights_early on/off, softmax/sigmoid - scoring, optional expert_bias). Each config asserts shape, dtype, - finiteness and numerical parity vs the reference in one run. + configurations (softmax/sigmoid scoring, optional non-zero + expert_bias). Each config asserts shape, dtype, finiteness and + numerical parity vs the reference in one run. * ``test_backward`` mirrors that for gradients. * ``TestTeEpMoeAuxLoss`` covers the second return value end-to-end (returned + parity + aux-only grad propagates to gate + combined main+aux grads stay finite) in two consolidated tests. -* ``TestTeEpMoEBlockFlax`` exercises the Flax wrapper with the same - parity reference. -* ``TestZZZTeEpMoeBootstrap`` verifies the per-process NCCL bootstrap - rejects a mismatched signature. + +Intentional non-coverage: + +* No dedicated "Flax wrapper init+apply" smoke test: every config above + already calls ``MoEBlock`` (the Flax wrapper) end-to-end, so a + separate wrapper smoke would just duplicate ``test_forward[softmax]`` + + ``test_backward[softmax]``. +* No re-bootstrap-mismatch test: ``ep_bootstrap`` rejects a mismatched + signature unconditionally and is a one-line guard; covering it from + this suite would taint the per-process NCCL bootstrap cache for the + rest of the file with no real upside. FP8 / MXFP8 recipes are deferred — the ``quantizer_sets`` plumbing has not yet been re-wired across the TE-EP ``shard_map`` boundary @@ -112,8 +119,7 @@ def _read_mp_options(): if not _MP_ACTIVE: pytest.skip( - "test_te_ep_moe.py requires the multiprocess launcher " - "(run_te_ep_moe.sh). Skipping.", + "test_te_ep_moe.py requires the multiprocess launcher (run_te_ep_moe.sh). Skipping.", allow_module_level=True, ) @@ -231,9 +237,7 @@ def mesh(): # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather # and cannot run from inside jax.jit. Sized to the worst-case recv_pr # across _CONFIGS so every parametrized config is bootstrap-compatible. - with mesh_obj, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ): + with mesh_obj, global_shard_guard(MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS)): ep_bootstrap( world_size=num_procs, rank=jax.process_index(), @@ -323,9 +327,7 @@ def _pure_jax_moe_reference( raise ValueError(f"Unsupported score_function={score_function!r}") routing_weights_full = jnp.zeros((T, num_experts), dtype=jnp.float32) - routing_weights_full = routing_weights_full.at[ - jnp.arange(T)[:, None], top_indices - ].set(weights) + routing_weights_full = routing_weights_full.at[jnp.arange(T)[:, None], top_indices].set(weights) # FFN. ``apply_topk_weights_early`` is a fusion knob that doesn't # change the math (wo is linear), so the reference is identical for @@ -335,9 +337,7 @@ def _pure_jax_moe_reference( intermediate = jax.nn.silu(layer_w0.astype(jnp.float32)) * layer_w1.astype(jnp.float32) intermediate = intermediate.astype(x.dtype) expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] - output_2d = jnp.einsum( - "te,teh->th", routing_weights_full.astype(x.dtype), expert_out - ) + output_2d = jnp.einsum("te,teh->th", routing_weights_full.astype(x.dtype), expert_out) output = output_2d.reshape(B, S, H).astype(x.dtype) if aux_loss_coeff > 0.0: @@ -352,9 +352,7 @@ def _pure_jax_moe_reference( else: # sigmoid aux_scores = jax.nn.sigmoid(logits) if K > 1: - aux_scores = aux_scores / ( - aux_scores.sum(axis=-1, keepdims=True) + 1e-20 - ) + aux_scores = aux_scores / (aux_scores.sum(axis=-1, keepdims=True) + 1e-20) routing_map = (routing_weights_full > 0).astype(jnp.int32) tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] sum_probs_per_expert = jnp.sum(aux_scores, axis=0) # [E] @@ -545,10 +543,12 @@ def _make_inputs(key): dict(score_function="sigmoid"), id="sigmoid", ), - pytest.param( - dict(score_function="sigmoid", use_expert_bias=True), - id="sigmoid-bias-zero", - ), + # NOTE: a ``sigmoid-bias-zero`` config (use_expert_bias=True with a + # zero-initialised bias buffer) was previously exercised here. It + # was dropped because the routing math collapses to the no-bias + # case when the buffer is zero -- ``sigmoid`` already covers that + # numerical path. The bias-aware codepath is still exercised by + # ``sigmoid-bias-strong`` below, which uses a non-zero bias. pytest.param( dict( score_function="sigmoid", @@ -565,9 +565,7 @@ def _reference_kwargs_from_config(config, params_np): return dict( score_function=config.get("score_function", "softmax"), expert_bias=( - jnp.asarray(params_np["expert_bias"]) - if config.get("use_expert_bias", False) - else None + jnp.asarray(params_np["expert_bias"]) if config.get("use_expert_bias", False) else None ), ) @@ -718,9 +716,7 @@ def test_aux_loss(self, mesh): # wired. aux_grads = _grad_aux_only(block, variables, mesh, x) g_gate = np.asarray( - jax.device_get( - _unwrap(aux_grads["params"]["gate_kernel"]).addressable_data(0) - ) + jax.device_get(_unwrap(aux_grads["params"]["gate_kernel"]).addressable_data(0)) ) 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" @@ -733,81 +729,6 @@ def test_combined_loss_grads(self, mesh): variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(23)) grads = _grad_step(block, variables, mesh, x, include_aux=True) for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_local = np.asarray( - jax.device_get(_unwrap(grads["params"][name]).addressable_data(0)) - ) + g_local = np.asarray(jax.device_get(_unwrap(grads["params"][name]).addressable_data(0))) assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf under main+aux" assert np.any(g_local != 0.0), f"{name} grad zero under main+aux" - - -class TestTeEpMoEBlockFlax: - """Flax wrapper end-to-end in one run: shape/dtype/finiteness on the - forward, numerical parity vs the same reference, and per-tensor - grad finiteness + non-zeroness.""" - - def test_init_apply_parity(self, mesh): - block = _make_block() - x = _make_inputs(jax.random.PRNGKey(12)) - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(13)) - - assert aux is None - assert output.shape == x.shape - assert output.dtype == x.dtype - out_local = np.asarray(jax.device_get(output.addressable_data(0))) - assert np.all(np.isfinite(out_local)) - - params_np = _params_global_numpy(variables, mesh) - x_np = np.asarray(jax.device_get(x)) - out_te_np = _to_global_numpy(output, mesh) - out_ref, _ = _pure_jax_moe_reference( - jnp.asarray(x_np), - jnp.asarray(params_np["gate_kernel"]), - jnp.asarray(params_np["wi_0"]), - jnp.asarray(params_np["wi_1"]), - jnp.asarray(params_np["wo"]), - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - ) - np.testing.assert_allclose( - out_te_np.astype(np.float32), - np.asarray(jax.device_get(out_ref)).astype(np.float32), - atol=FWD_ATOL, - rtol=FWD_RTOL, - ) - - grads = _grad_step(block, variables, mesh, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_local = np.asarray( - jax.device_get(_unwrap(grads["params"][name]).addressable_data(0)) - ) - assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf" - assert np.any(g_local != 0.0), f"{name} grad zero" - - -# Keep the bootstrap-signature test last in the module (the "ZZZ" prefix -# ensures pytest's alphabetic class ordering picks it last): it -# intentionally mismatches the NCCL EP bootstrap signature, which -# permanently taints the per-process bootstrap cache for the rest of -# the file. -class TestZZZTeEpMoeBootstrap: - """Per-process NCCL bootstrap re-bootstrap rejection.""" - - def test_bootstrap_signature_mismatch_raises(self, mesh): - block_a = _make_block() - x_a = _make_inputs(jax.random.PRNGKey(14)) - _init_apply(block_a, mesh, x_a, jax.random.PRNGKey(15)) - - # Different hidden dim → different bootstrap signature. - bigger_hidden = HIDDEN * 2 - x_b = jax.random.normal( - jax.random.PRNGKey(16), (BATCH, SEQ, bigger_hidden), dtype=DTYPE - ) - block_b = MoEBlock( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - data_parallelism_axes=(FSDP_AXIS,), - dtype=DTYPE, - ) - with pytest.raises(ValueError, match="bootstrapped"): - _init_apply(block_b, mesh, x_b, jax.random.PRNGKey(17)) From 51a046ad0016b63e7bcf819fac35b858557ecf68 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 15:15:40 -0700 Subject: [PATCH 03/23] jax/router: fix two bwd custom_partitioning bugs (aux-loss rank, topk closure) Two unrelated one-line bugs in the bwd custom_partitioning machinery that only surface once the MoE block's aux-loss path is lifted out of shard_map (the custom_partitioning_sharding_rule check is skipped under shard_map, which is why these never tripped before). 1. FusedMoEAuxLossBwdPrimitive.shardy_sharding_rule: ``grad_aux_loss`` is the cotangent of a scalar loss and is rank-0; declaring it with a spurious ``grad_one`` factor gave it rank-1 and tripped JAX's custom_partitioning_sharding_rule rank check at global view. Change the rule's third operand entry to empty: "const_buf_one, num_experts, grad_one -> i num_experts" -> "const_buf_one, num_experts, -> i num_experts" 2. FusedTopkWithScoreFunctionBwdPrimitive.partition: ``del result_infos, routing_map_format`` removed ``routing_map_format`` from the enclosing scope before the nested ``sharded_impl`` closure was invoked. Python closures resolve names at call time, not definition time, so when XLA finally invoked ``sharded_impl`` for the bwd partitioned impl it raised ``NameError: cannot access free variable 'routing_map_format'``. Drop ``routing_map_format`` from the ``del`` and leave a NOTE so future cleanups don't reintroduce the bug. Sibling partition methods (fwd topk, both aux-loss directions) already only ``del result_infos`` and need no change. Signed-off-by: Teddy Do --- transformer_engine/jax/cpp_extensions/router.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index 8cc94fcaaf..46f51c9d33 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -412,6 +412,11 @@ def partition( arg_infos, result_infos, ): + # NOTE: do NOT include ``routing_map_format`` in this ``del``: the + # ``sharded_impl`` closure below resolves it by name at call time + # (when XLA invokes the partitioned impl), so deleting it here + # raises ``NameError: cannot access free variable 'routing_map_format'`` + # at execution time of the bwd custom_partitioning. del result_infos grad_spec = get_padded_spec(arg_infos[2]) out_sharding = NamedSharding(mesh, PartitionSpec(*grad_spec)) @@ -645,7 +650,14 @@ def shardy_sharding_rule(*args): # backward reconstructs the full [num_tokens, num_experts] grad_probs from # scalar inputs. Shardy will leave num_tokens unsharded, which matches the # replicated PartitionSpec(None, None) in partition(). - return "const_buf_one, num_experts, grad_one -> i num_experts" + # + # grad_aux_loss is the cotangent of a scalar loss and is therefore + # rank-0; the third operand entry is empty (no factor labels). Declaring + # it with the spurious "grad_one" factor gave it rank-1 and tripped + # JAX's custom_partitioning_sharding_rule check once the MoE block + # lifted its aux-loss path out of shard_map (the rule is skipped under + # shard_map, which is why this surfaces only at global view). + return "const_buf_one, num_experts, -> i num_experts" register_primitive(FusedMoEAuxLossBwdPrimitive) From 375d08020548e74476c7707455330f901cf0b206 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 15:15:47 -0700 Subject: [PATCH 04/23] jax/ep: skip size-1 dp/fsdp axis in _ep_outer_axis A dp_resource or fsdp_resource that exists in the active mesh resource config but is sized 1 in the actual mesh would still be returned by ``_ep_outer_axis()``, pinning EP-output PartitionSpecs to a degenerate axis. JAX collapses size-1 mesh axes during lowering, which made the EP-output specs reference an axis that no longer exists at runtime -- breaking shard_map output stitching on configs where DP or FSDP is optional. Treat a size-1 axis as absent: prefer dp -> fsdp, but only when the candidate axis is actually sized > 1 in the current mesh. Falls back to the previous behaviour when no axis is configured at all. Signed-off-by: Teddy Do --- transformer_engine/jax/cpp_extensions/ep.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 2a4e17991a..ad49aaa0d1 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -23,7 +23,7 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive -from ..sharding import global_mesh_resource +from ..sharding import global_mesh_resource, get_mesh_axis_size __all__ = [ "EpConfig", @@ -125,8 +125,15 @@ def _ep_outer_axis(): 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 From f3f77f7cd7f2d266a1807ee20271e62feb51d827 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 15:16:15 -0700 Subject: [PATCH 05/23] jax/flax: realign _MoEBlock with post-resync moe() signature After the upstream PR #3036 resync the moe() API surface lost PermutationBackend (TE-EP is the only backend now), gate_inside_vjp (always True), and the per-call quantizer_sets knob (quantization flows through the standard TE autocast / with_quantizer_set context). It also gained apply_topk_weights_early and renamed the wrapper's private _align_size to the public align_size the test suite already uses. The Flax _MoEBlock wrapper was still passing the old kwargs, which broke every test that touched the wrapper. Wrapper changes: * drop "from ..moe import PermutationBackend" plus the dataclass field, the isinstance(..., PermutationBackend) validation in __post_init__, and the pass-through to moe(). * drop "from ..quantize import noop_quantizer_set" and the quantizer_sets=(noop, noop, noop) pass-through. * drop gate_inside_vjp=True. * rename _align_size: int = 0 -> align_size: int = 0 (matches what tests/jax/test_te_ep_moe.py already passes). * add apply_topk_weights_early: bool = False and pass it through to moe(). * refresh class docstring: drop permutation_backend / _align_size / quantizer_sets descriptions, add apply_topk_weights_early / align_size, note that quantization currently flows only through fp8_autocast. Signed-off-by: Teddy Do --- transformer_engine/jax/flax/moe.py | 41 +++++++++++++----------------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 91346a7a48..b98d5a9549 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -37,8 +37,7 @@ # import P`` without a second jax.sharding import. from jax.sharding import PartitionSpec as P # noqa: F401 # pylint: disable=unused-import -from ..moe import PermutationBackend, moe -from ..quantize import noop_quantizer_set +from ..moe import moe from ..router import ScoreFunction from ..sharding import get_active_resource_axis from .module import TransformerEngineBase @@ -50,7 +49,7 @@ Initializer = Callable[[PRNGKey, Shape, DType], Array] -__all__ = ["PermutationBackend", "_MoEBlock"] +__all__ = ["_MoEBlock"] class _MoEBlock(TransformerEngineBase): @@ -100,12 +99,15 @@ class _MoEBlock(TransformerEngineBase): replicated across non-EP axes within an EP group; set e.g. ``("fsdp",)`` for true FSDP-of-batch where each device owns a unique slice of the batch. - permutation_backend : PermutationBackend - ``PURE_JAX`` (default) or ``TRITON``. - _align_size : int + apply_topk_weights_early : bool + If ``True``, multiply expert outputs by their top-k weights + *inside* each shard before ``ep_combine`` (saves one global + reduction at the cost of an extra broadcast). Default ``False``. + align_size : int Per-expert group-size alignment (``0`` disables; required > 0 - for quantized grouped GEMM). Internal knob; will be inferred - from the active quantization recipe in a follow-up PR. + for quantized grouped GEMM). Forwarded to ``tex.ep_prepare`` as + ``dispatch_output_per_expert_alignment``; will be inferred from + the active quantization recipe in a follow-up PR. dtype : jnp.dtype Compute / parameter dtype. @@ -114,9 +116,9 @@ class _MoEBlock(TransformerEngineBase): Register per-expert FFN biases. Quantization is currently configured via the standard TE autocast - context (``fp8_autocast``/``with_quantizer_set``); per-call - quantizer sets can also be passed through ``__call__``'s - ``quantizer_sets`` keyword once we stabilise the recipe pipeline. + context (``fp8_autocast``/``with_quantizer_set``) and threaded + through ``moe()`` internally; this wrapper does not expose a + per-call ``quantizer_sets`` knob yet. """ # Architecture @@ -143,9 +145,9 @@ class _MoEBlock(TransformerEngineBase): # Parallelism data_parallelism_axes: Tuple[str, ...] = () - # Permutation - permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX - _align_size: int = 0 + # MoE knobs forwarded to ``moe()`` + apply_topk_weights_early: bool = False + align_size: int = 0 # Dtypes / init / misc dtype: DType = jnp.float32 @@ -163,11 +165,6 @@ def __post_init__(self): 1.0, "fan_in", "truncated_normal", dtype=self.dtype ), ) - if not isinstance(self.permutation_backend, PermutationBackend): - raise TypeError( - "permutation_backend must be a PermutationBackend, got" - f" {self.permutation_backend!r}" - ) super().__post_init__() @nn.compact @@ -270,15 +267,13 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: group_topk=self.group_topk, scaling_factor=self.scaling_factor, aux_loss_coeff=self.aux_loss_coeff, - permutation_backend=self.permutation_backend, - align_size=self._align_size, - gate_inside_vjp=True, + apply_topk_weights_early=self.apply_topk_weights_early, + align_size=self.align_size, 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, - quantizer_sets=(noop_quantizer_set, noop_quantizer_set, noop_quantizer_set), dtype=self.dtype, ) From 49058ff508d0c1a8a89b0802c8292d7373d5e54a Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 15:16:37 -0700 Subject: [PATCH 06/23] jax/moe: plumb token_counts to grouped_gemm and zero 0-token wgrad slices Two correctness fixes for the TE-EP MoE custom_vjp that together let the bwd parity tests pass on 0-token-globally experts, and drop a workaround that is no longer needed. (1) Plumb per-expert padded token_counts into grouped_gemm group_sizes. NCCL EP HT dispatch lays out recv_tokens expert-major as [expert_0_padded | expert_1_padded | ... | overalloc_tail] where each per-expert block already includes the dispatch_output_per_expert_alignment zero-padding and only the trailing overalloc tail (slack between sum(token_counts) and the worst-case recv_pr) is unused. Previously _ffn_fwd_per_shard built a static local_group_sizes = jnp.full((num_local_experts,), slots_per_expert), which over-counted by the overalloc tail and forced cuBLAS to run the GEMM for every group including 0-token-routed experts. Pipe the real per-shard token_counts (1, num_local_experts) from ep_prepare through _moe_fwd_rule (added to ffn_in_specs/ffn_in_args with ep2_spec), into _ffn_fwd_per_shard as token_counts_local, and reshape into local_group_sizes for both grouped_quantize and grouped_gemm. cuBLAS now skips both 0-token experts and the trailing overalloc tail. Mirror the residual spec change on the bwd (local_group_sizes residual moves from P() to ep2_spec). (2) Per-group jnp.where zero-fill on wgrad outputs. cuBLAS grouped_gemm skips groups with size_g == 0 without zero-filling the corresponding out[g, :, :] slice (cublaslt_grouped_gemm.cu lines 2086/2096). For a shard hosting an expert that received zero tokens globally, d_wo / d_wi_combined for that expert is left uninit, which propagates NaN straight into the user's optimizer state. Add wgrad_group_active = (local_group_sizes > 0)[:, None, None] in _ffn_bwd_per_shard and apply via jnp.where on d_wo (right after the wo wgrad) and d_wi_combined (right after the fused wi_0+wi_1 wgrad). Mask shape is (num_local_experts, 1, 1) so cost is negligible. (3) Drop the lax.cond zero-init guard on r_tok in _moe_fwd_rule._body. Previously a jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like) wrapper around recv_tokens worked around tex.ep_dispatch_fwd leaving the recv buffer uninit on fully-empty-receiver ranks. With (1) in place, cuBLAS skips experts whose group_sizes == 0 and the per-row trailing tail of dispatched recv_tokens is unread by every downstream consumer (subsequent grouped_gemms read only sum(group_sizes) rows; ep_combine and ep_dispatch_bwd are handle_mem-aware). The only per-row consumer that would propagate the tail is grouped_dbias (per-row segment_sum), which only runs when has_bias=True, and that FFN bias path is currently gated upstream (cuBLAS grouped_gemm has no fused bias on Hopper yet; PR 3083 adds the pure-JAX bias add). With (2) handling the user-visible wgrad-NaN risk on 0-token experts, the lax.cond is now redundant. Replace with a NOTE pointing at the two follow-ups that would force its reintroduction: - a future caller that reads the full recv tile non-group-aware (e.g. an inspect probe), or - the FFN bias path landing, which would resurrect grouped_dbias. Also rewrite the _ffn_fwd_per_shard and _ffn_bwd_per_shard docstrings to spell out the per-row vs per-group uninit semantics so the next person debugging a NaN here has the invariants written down. Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 219 +++++++++++++++++++++++----------- 1 file changed, 147 insertions(+), 72 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 08348b0104..554ee628b8 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -276,6 +276,7 @@ def tree_unflatten(cls, aux_data, children): def _ffn_fwd_per_shard( recv_tokens_local: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, + token_counts_local: jnp.ndarray, wi_0: jnp.ndarray, wi_1: jnp.ndarray, wo: jnp.ndarray, @@ -295,11 +296,50 @@ def _ffn_fwd_per_shard( ``[1, recv_pr, H_out]`` so the surrounding ``shard_map`` reassembles them as ``[num_procs, recv_pr, H_out]``) plus the residuals consumed by the bwd. + + ``token_counts_local`` is the per-expert padded token count (shape + ``[1, num_local_experts]``) from ``tex.ep_prepare``. With NCCL EP's + HT expert-major layout, the dispatch lays out experts contiguously + in ``recv_tokens`` as ``[expert_0_padded | expert_1_padded | ... | + overalloc_tail]``, where each per-expert block already includes the + ``dispatch_output_per_expert_alignment`` zero-padding and only the + trailing overalloc tail (slack between ``sum(token_counts)`` and the + worst-case ``recv_pr``) is unused. Plumbing ``token_counts`` straight + into ``grouped_gemm`` as ``group_sizes`` makes cuBLAS skip both the + overalloc tail (saving FMAs on partially-loaded shards) and any + expert whose per-shard routed count is zero (saving the GEMM + altogether, not just the rows). + + cuBLAS leaves the trailing of each grouped_gemm *per-row* output + (``combined_out``, ``expert_outputs`` in fwd; ``d_intermediate``, + ``d_sorted_x`` in bwd) uninitialised past ``sum(group_sizes)``, and + fully uninitialised on a shard whose every local expert has count 0. + That per-row tail is harmless for everything in this block: + subsequent ``grouped_gemm`` / ``ep_combine`` / ``ep_dispatch_bwd`` + only read valid rows per ``local_group_sizes`` / ``handle_mem``, and + ``act_fn``'s NaN tail only contaminates positions that no + group-aware consumer reads. The one per-row exception is + ``grouped_dbias`` (a ``segment_sum`` that walks every row), which + is only reached when ``has_bias=True``. That FFN bias path is + currently gated upstream (cuBLAS grouped_gemm has no fused bias on + Hopper yet; PR 3083 adds a pure-JAX bias add), so we don't pay for + the tail-zeroing masks needed to keep ``segment_sum`` well-defined. + If the bias path ever lands, re-add ``jnp.where`` masks on + ``combined_out`` / ``d_eo_2d`` / ``d_intermediate``. + + Separately, the grouped_gemm *wgrad* outputs (``d_wo``, + ``d_wi_combined`` in bwd) are per-group ``(num_local_experts, K, N)`` + and are the *user-visible* weight gradients. cuBLAS skips groups + with ``size_g == 0`` without zero-filling, so for 0-token-globally + experts the slice would be NaN and leak into the optimizer. This + is handled by per-group ``jnp.where`` masks (``wgrad_group_active``) + in ``_ffn_bwd_per_shard``; see that docstring. """ hidden = recv_tokens_local.shape[-1] sorted_x = recv_tokens_local.reshape(-1, hidden) recv_w_flat = recv_topk_weights_local.reshape(-1) - local_group_sizes = jnp.full((num_local_experts,), slots_per_expert, dtype=jnp.int32) + local_group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) + del slots_per_expert # not used since group_sizes is plumbed in dynamically wi_0 = wi_0.astype(sorted_x.dtype) wi_1 = wi_1.astype(sorted_x.dtype) @@ -347,8 +387,7 @@ def _ffn_fwd_per_shard( # input layout is unchanged. act_fn = _convert_to_activation_function(activation_type) intermediate = ( - act_fn(gate_proj_out.astype(jnp.float32)) - * up_proj_out.astype(jnp.float32) + act_fn(gate_proj_out.astype(jnp.float32)) * up_proj_out.astype(jnp.float32) ).astype(sorted_x.dtype) if apply_topk_weights_early: @@ -375,6 +414,10 @@ def _ffn_fwd_per_shard( casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) + # Reshape local_group_sizes to (1, num_local_experts) so the + # surrounding shard_map can stitch per-shard counts back into the + # global (num_procs, num_local_experts) layout matching token_counts. + local_group_sizes_3d = local_group_sizes.reshape(1, num_local_experts) residuals = ( casted_sorted_x_lhs_trans, casted_wi_rhs_trans, @@ -382,7 +425,7 @@ def _ffn_fwd_per_shard( up_proj_out, casted_intermediate_lhs_trans, casted_wo_rhs_trans, - local_group_sizes, + local_group_sizes_3d, ) return expert_outputs_3d, residuals @@ -407,10 +450,43 @@ def _ffn_bwd_per_shard( Mirrors :func:`_ffn_fwd_per_shard`. Returns ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], d_wi_0, d_wi_1, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. + + ``local_group_sizes`` arrives as ``(1, num_local_experts)`` (the + fwd-side shard residual), with the same per-expert padded counts the + fwd used as ``grouped_gemm`` ``group_sizes``. cuBLAS leaves rows + past ``sum(group_sizes)`` uninit in the bwd grouped_gemm *dgrad* + outputs (``d_intermediate``, ``d_sorted_x``), but every downstream + consumer of those per-row outputs is group-aware (wi/wo wgrads + contract only over valid rows; ``ep_dispatch_bwd`` reads only + valid positions per ``handle_mem``), so the per-row trailing tail + sits unread. ``grouped_dbias`` (per-row ``segment_sum``) is the + only per-row consumer that would propagate the tail, and it is + only invoked when ``has_bias=True``; that FFN bias path is gated + upstream (see ``_ffn_fwd_per_shard`` docstring) so we skip the + per-row tail-zeroing masks until cuBLAS gains fused-bias grouped + GEMM (or PR 3083's pure-JAX bias add lands). + + The *wgrad* outputs (``d_wo``, ``d_wi_combined``) are different. + They're per-group ``(num_local_experts, K, N)``, and cuBLAS + skips groups with ``size_g == 0`` without zero-filling the + corresponding ``out[g, :, :]`` slice (see + ``cublaslt_grouped_gemm.cu`` lines 2086/2096). For shards hosting + an expert that received zero tokens globally, that expert's + ``d_wo`` / ``d_wi`` slice would be uninit → NaN propagates to the + user's optimizer. We zero those slices via a per-group ``jnp.where`` + immediately after each wgrad. The mask is shape ``(num_groups, 1, 1)`` + and ``num_groups == num_local_experts`` is tiny, so this is cheap. """ + local_group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) recv_w_flat = recv_topk_weights_local.reshape(-1) q_set = noop_quantizer_set + # Per-group active mask for wgrad outputs. cuBLAS grouped_gemm skips + # groups with size_g == 0 and leaves the corresponding output slice + # uninit; without this, ``d_wo[g] / d_wi_combined[g]`` for any expert + # that received zero tokens globally would be NaN and propagate to + # the user's optimizer. + wgrad_group_active = (local_group_sizes > 0)[:, None, None] # wo bwd casted_d_eo = tex.grouped_quantize(d_eo_2d, q_set.dgrad, local_group_sizes, flatten_axis=-1) @@ -426,6 +502,7 @@ def _ffn_bwd_per_shard( _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) + d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) d_wo_bias = tex.grouped_dbias(d_eo_2d, local_group_sizes) if has_bias else None act_fn = _convert_to_activation_function(activation_type) @@ -474,6 +551,9 @@ def _ffn_bwd_per_shard( casted_d_combined.get_tensor(usage=TensorUsage.RHS), contracting_dims=((0,), (0,)), ) + d_wi_combined = jnp.where( + wgrad_group_active, d_wi_combined, jnp.zeros_like(d_wi_combined) + ) d_wi_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1) if has_bias: d_wi_combined_bias = tex.grouped_dbias(d_combined, local_group_sizes) @@ -645,9 +725,7 @@ def _moe_fwd_rule( # 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_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, @@ -698,12 +776,8 @@ def _moe_fwd_rule( # 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) - ) + 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( @@ -723,8 +797,12 @@ def _moe_fwd_rule( has_bias = wi_0_bias is not None kernel_spec = P(ep_axis, None, None) bias_spec = P(ep_axis, None) if has_bias else None - ffn_in_specs = (ep3_spec, ep2_spec, kernel_spec, kernel_spec, kernel_spec) - ffn_in_args = [recv_tokens, recv_topk_weights, wi_0, wi_1, wo] + # token_counts is the per-shard (1, num_local_experts) padded + # per-expert count from ep_prepare; piped into _ffn_fwd_per_shard + # as the grouped_gemm group_sizes so cuBLAS skips both 0-token + # experts and the trailing overalloc tail. + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec, kernel_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi_0, wi_1, wo] if has_bias: ffn_in_specs = ffn_in_specs + (bias_spec, bias_spec, bias_spec) ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) @@ -734,46 +812,49 @@ def _moe_fwd_rule( # fused via jnp.concatenate along the trailing (output) axis # (see _ffn_fwd_per_shard for rationale), so the residual is a # single 3D casted_wi_rhs_trans of shape - # (num_local_experts, hidden, 2*H_inter). + # (num_local_experts, hidden, 2*H_inter). local_group_sizes is + # now per-shard dynamic (= per-shard token_counts), so its + # residual spec mirrors ep2_spec (one row per ep rank). residuals_spec = ( - P(), # casted_sorted_x_lhs_trans - P(ep_axis, None, None), # casted_wi_rhs_trans - P(), # gate_proj_out - P(), # up_proj_out - P(), # casted_intermediate_lhs_trans - P(ep_axis, None, None), # casted_wo_rhs_trans - P(), # local_group_sizes + P(), # casted_sorted_x_lhs_trans + P(ep_axis, None, None), # casted_wi_rhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans + P(ep_axis, None, None), # casted_wo_rhs_trans + ep2_spec, # local_group_sizes (1, num_local_experts) per shard ) out_specs = (ep3_spec, residuals_spec) def _body(*args): if has_bias: - (r_tok, r_w, w0, w1, w_o, w0b, w1b, wob) = args + (r_tok, r_w, tc, w0, w1, w_o, w0b, w1b, wob) = args else: - (r_tok, r_w, w0, w1, w_o) = args + (r_tok, r_w, tc, w0, w1, w_o) = args w0b = w1b = wob = None - # Per-rank conditional zero-init of r_tok. Works around a - # narrowly-scoped tex.ep_dispatch_fwd contract gap: the NCCL EP - # HT dispatch kernel zero-initialises the recv buffer correctly - # on ranks that receive at least one token, but leaves - # uninitialised memory on fully-empty-receiver ranks. ``r_w`` - # (the dispatch's own written-or-not indicator: 0 at padded - # slots, non-zero at real-routed slots) gives us a per-shard - # predicate for free. ``jax.lax.cond`` only executes the - # selected branch, so loaded ranks pay nothing at runtime; - # only empty ranks do the zero-fill. - # TODO: remove once tex.ep_dispatch_fwd zero-inits empty-rank - # recv buffers upstream. - rank_has_tokens = jnp.any(r_w != 0) - r_tok = jax.lax.cond( - rank_has_tokens, - lambda x: x, - lambda x: jnp.zeros_like(x), - r_tok, - ) + # NOTE: tex.ep_dispatch_fwd's NCCL EP HT path leaves the recv + # buffer uninitialised on fully-empty-receiver ranks (and at + # padded slots on partially-loaded ranks). We don't need a + # zero-init guard here anymore because: + # 1. ``tc`` (per-expert padded counts) is plumbed into + # grouped_gemm as group_sizes, so cuBLAS skips both + # 0-token experts and the trailing overalloc tail. + # 2. The per-group wgrad masks in _ffn_bwd_per_shard zero + # ``d_wo`` / ``d_wi_combined`` slices for 0-token-globally + # experts (cuBLAS skips size_g==0 groups without + # zero-filling, which would otherwise leak NaN into the + # user's optimizer). + # 3. All other downstream consumers (ep_combine, + # ep_dispatch_bwd) are handle_mem-aware and read only + # valid positions. + # If a future caller adds a non-group-aware reader of r_tok + # (e.g. an inspect probe over the full recv tile), re-add the + # ``jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like)`` + # guard here. return _ffn_fwd_per_shard( r_tok, r_w, + tc, w0, w1, w_o, @@ -793,9 +874,7 @@ def _body(*args): out_specs=out_specs, check_rep=False, )(*ffn_in_args) - expert_outputs = jax.lax.with_sharding_constraint( - expert_outputs, NamedSharding(mesh, ep3_spec) - ) + 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) @@ -973,15 +1052,15 @@ def _moe_bwd_rule( bias_spec = P(ep_axis, None) if has_bias else None bwd_in_specs = ( - ep3_spec, # d_expert_outputs - P(), # casted_sorted_x_lhs_trans + ep3_spec, # d_expert_outputs + P(), # casted_sorted_x_lhs_trans P(ep_axis, None, None), # casted_wi_rhs_trans - P(), # gate_proj_out - P(), # up_proj_out - P(), # casted_intermediate_lhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans P(ep_axis, None, None), # casted_wo_rhs_trans - P(), # local_group_sizes - ep2_spec, # recv_topk_weights + ep2_spec, # local_group_sizes (1, num_local_experts) per shard + ep2_spec, # recv_topk_weights ) bwd_in_args = [ d_expert_outputs, @@ -995,14 +1074,14 @@ def _moe_bwd_rule( ctx.recv_topk_weights, ] bwd_out_specs = ( - ep3_spec, # d_sorted_x - ep2_spec, # d_recv_w_from_intermediate - kernel_spec, # d_wi_0 - kernel_spec, # d_wi_1 - kernel_spec, # d_wo - bias_spec if has_bias else None, # d_wi_0_bias - bias_spec if has_bias else None, # d_wi_1_bias - bias_spec if has_bias else None, # d_wo_bias + ep3_spec, # d_sorted_x + ep2_spec, # d_recv_w_from_intermediate + kernel_spec, # d_wi_0 + kernel_spec, # d_wi_1 + kernel_spec, # d_wo + bias_spec if has_bias else None, # d_wi_0_bias + bias_spec if has_bias else None, # d_wi_1_bias + bias_spec if has_bias else None, # d_wo_bias ) def _bwd_body(*args): @@ -1059,15 +1138,15 @@ def _bwd_body(*args): in_specs=bwd_in_specs, out_specs=bwd_out_specs, check_rep=False, - )(*bwd_in_args) + )( + *bwd_in_args + ) 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_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, @@ -1114,9 +1193,7 @@ def _bwd_body(*args): ) # 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 - ) + 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, @@ -1305,9 +1382,7 @@ def moe( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - expected_leading: Any = ( - (*data_parallelism_axes, ep_axis) if data_parallelism_axes else ep_axis - ) + 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): From 3e779572ca9f2c38661243eb0c17eae71a7896ac Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 15:31:42 -0700 Subject: [PATCH 07/23] jax/flax,tests: rename use_bias/use_expert_bias for symmetry (PR #3116) Address jberchtold-nvidia's PR #3116 nit "rename use_bias -> use_ffn_bias and use_expert_bias -> use_expert_routing_bias". The two flags are siblings (they enable two different bias buffers) but the old names suggested ``use_bias`` was the general fallback, which wasn't the intent. The new names make the FFN-vs-routing distinction obvious from the call site. * transformer_engine/jax/flax/moe.py use_bias -> use_ffn_bias (dataclass field + branch in __call__ + docstring entry) use_expert_bias -> use_expert_routing_bias (same) * tests/jax/test_te_ep_moe.py _make_block(use_expert_bias=...) -> use_expert_routing_bias sigmoid-bias-strong config key updated _reference_kwargs_from_config now reads use_expert_routing_bias ``_MoEBlock`` is still the experimental underscore-prefixed alias (no public ``MoEBlock`` export yet), so the rename is API-safe. The pre-resync legacy tests (``test_moe_vjp.py``, ``test_multiprocess_moe_vjp.py``) are intentionally not updated -- they already reference removed APIs like ``PermutationBackend`` and need a separate post-resync cleanup pass. Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 39 +++++++++++++++--------------- transformer_engine/jax/flax/moe.py | 37 +++++++++++++++------------- 2 files changed, 39 insertions(+), 37 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 75326af1c6..3a5cbff51a 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -181,7 +181,7 @@ def _read_mp_options(): GRAD_GATE_RTOL = 5e-1 # Two TE EP runs that should be bitwise-equal modulo XLA fusion order -# (align_size rounding, etc.). +# (slot alignment rounding, etc.). TE_TO_TE_ATOL = 5e-3 TE_TO_TE_RTOL = 5e-3 @@ -373,9 +373,8 @@ def _pure_jax_moe_reference( def _make_block( *, apply_topk_weights_early=False, - align_size=0, aux_loss_coeff=0.0, - use_expert_bias=False, + use_expert_routing_bias=False, score_function="softmax", bias_init=None, ): @@ -385,9 +384,8 @@ def _make_block( intermediate_size=INTER, data_parallelism_axes=(FSDP_AXIS,), apply_topk_weights_early=apply_topk_weights_early, - align_size=align_size, aux_loss_coeff=aux_loss_coeff, - use_expert_bias=use_expert_bias, + use_expert_routing_bias=use_expert_routing_bias, score_function=score_function, dtype=DTYPE, ) @@ -532,27 +530,26 @@ def _make_inputs(key): # 0*NaN -> NaN leak from padded recv slots in the early-weighting # multiply (intermediate * recv_w * mask) is debugged. Late # weighting (combine-side) is unaffected and stays covered above. - # Note: a dedicated align_size=128 config was previously listed - # here. It is no longer interesting because moe.py now floors - # slots_per_expert at 128 unconditionally (effective_align = - # max(align_size, 128)), so align_size=0 (default) and - # align_size=128 produce identical layouts. Re-add a distinct - # case only if the floor is loosened or a >128 align is needed - # by a recipe (e.g. some FP8 paths want 256-aligned slots). + # Note: align_size is no longer a user-facing parameter; it is + # hard-coded to _ALIGN_SIZE = 128 in moe.py (per PR #3116 + # review). Re-add a distinct align-size config only if the + # constant is loosened, or a recipe-driven inference is added + # that selects a >128 alignment. pytest.param( dict(score_function="sigmoid"), id="sigmoid", ), - # NOTE: a ``sigmoid-bias-zero`` config (use_expert_bias=True with a - # zero-initialised bias buffer) was previously exercised here. It - # was dropped because the routing math collapses to the no-bias - # case when the buffer is zero -- ``sigmoid`` already covers that - # numerical path. The bias-aware codepath is still exercised by - # ``sigmoid-bias-strong`` below, which uses a non-zero bias. + # NOTE: a ``sigmoid-bias-zero`` config (use_expert_routing_bias=True + # with a zero-initialised bias buffer) was previously exercised + # here. It was dropped because the routing math collapses to the + # no-bias case when the buffer is zero -- ``sigmoid`` already + # covers that numerical path. The bias-aware codepath is still + # exercised by ``sigmoid-bias-strong`` below, which uses a + # non-zero bias. pytest.param( dict( score_function="sigmoid", - use_expert_bias=True, + use_expert_routing_bias=True, bias_init=_strong_expert_bias_init, ), id="sigmoid-bias-strong", @@ -565,7 +562,9 @@ def _reference_kwargs_from_config(config, params_np): return dict( score_function=config.get("score_function", "softmax"), expert_bias=( - jnp.asarray(params_np["expert_bias"]) if config.get("use_expert_bias", False) else None + jnp.asarray(params_np["expert_bias"]) + if config.get("use_expert_routing_bias", False) + else None ), ) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index b98d5a9549..ed36ab835f 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -81,10 +81,12 @@ class _MoEBlock(TransformerEngineBase): Grouped top-k knobs (DeepSeek-style). ``None`` disables grouping. scaling_factor : float Multiplier on the routing weights. - use_expert_bias : bool - If ``True``, registers a per-expert routing bias (shape ``[E]``). - Only meaningful with ``score_function="sigmoid"``; the underlying - primitive validates the pairing. + use_expert_routing_bias : bool + If ``True``, registers a per-expert routing bias (shape ``[E]``) + used by the topk selection. Only meaningful with + ``score_function="sigmoid"``; the underlying primitive validates + the pairing. (Renamed from ``use_expert_bias`` per PR #3116 + review for symmetry with ``use_ffn_bias``.) aux_loss_coeff : float If ``> 0``, return the MoE auxiliary load-balancing loss scalar in addition to the main output. @@ -103,17 +105,20 @@ class _MoEBlock(TransformerEngineBase): If ``True``, multiply expert outputs by their top-k weights *inside* each shard before ``ep_combine`` (saves one global reduction at the cost of an extra broadcast). Default ``False``. - align_size : int - Per-expert group-size alignment (``0`` disables; required > 0 - for quantized grouped GEMM). Forwarded to ``tex.ep_prepare`` as - ``dispatch_output_per_expert_alignment``; will be inferred from - the active quantization recipe in a follow-up PR. + + Note that the per-expert dispatch-slot alignment is fixed internally + at 128 tokens (see ``moe._ALIGN_SIZE``). Per PR #3116 review there's + no current model that wants a >128 alignment, so this is not exposed + as a parameter; re-introduce a knob (or recipe-driven inference) if + a future FP8 recipe needs >128. dtype : jnp.dtype Compute / parameter dtype. kernel_init, bias_init, expert_bias_init : Initializers. - use_bias : bool - Register per-expert FFN biases. + use_ffn_bias : bool + Register per-expert FFN biases (``wi_0_bias``, ``wi_1_bias``, + ``wo_bias``). (Renamed from ``use_bias`` per PR #3116 review + for symmetry with ``use_expert_routing_bias``.) Quantization is currently configured via the standard TE autocast context (``fp8_autocast``/``with_quantizer_set``) and threaded @@ -133,7 +138,7 @@ class _MoEBlock(TransformerEngineBase): num_groups: Optional[int] = None group_topk: Optional[int] = None scaling_factor: float = 1.0 - use_expert_bias: bool = False + use_expert_routing_bias: bool = False aux_loss_coeff: float = 0.0 # Sharding (logical axes) @@ -147,14 +152,13 @@ class _MoEBlock(TransformerEngineBase): # MoE knobs forwarded to ``moe()`` apply_topk_weights_early: bool = False - align_size: int = 0 # Dtypes / init / misc dtype: DType = jnp.float32 kernel_init: Optional[Initializer] = None bias_init: Initializer = nn.initializers.zeros expert_bias_init: Initializer = nn.initializers.zeros - use_bias: bool = False + use_ffn_bias: bool = False def __post_init__(self): if self.kernel_init is None: @@ -218,7 +222,7 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: self.dtype, ) wi_0_bias = wi_1_bias = wo_bias = None - if self.use_bias: + if self.use_ffn_bias: wi_0_bias = self.param( "wi_0_bias", nn.with_logical_partitioning(self.bias_init, ("exp", "mlp")), @@ -238,7 +242,7 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: self.dtype, ) expert_bias = None - if self.use_expert_bias: + if self.use_expert_routing_bias: expert_bias = self.param( "expert_bias", nn.with_logical_partitioning(self.expert_bias_init, ("exp",)), @@ -268,7 +272,6 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: scaling_factor=self.scaling_factor, aux_loss_coeff=self.aux_loss_coeff, apply_topk_weights_early=self.apply_topk_weights_early, - align_size=self.align_size, ep_axis=ep_axis, data_parallelism_axes=self.data_parallelism_axes, input_axes=self.input_axes, From 2f54d83d2d0f8506891fac150bfcd747226348a9 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 15:32:58 -0700 Subject: [PATCH 08/23] jax/moe: address PR #3116 review feedback (hardcode align + expand inline justifications) Responds to jberchtold-nvidia's PR #3116 review threads on ``transformer_engine/jax/moe.py``. All changes are confined to a single file because each review thread targets a localized region and splitting mid-file would risk reordering bugs. Per review thread: 1. "Why do we need _with_sharding_constraint_cast_bwd? I haven't seen something like this required for our other VJPs." -- Expand the helper's docstring to spell out exactly why MoE needs it: unlike LN+MLP, the MoE bwd composes a bf16 cotangent from ep_dispatch_bwd with an fp32 cotangent from fused_topk_with_score_function_bwd (which the fwd's logits_2d -> fp32 promotion forces). Without the cast, ``d_x`` surfaces at fp32 even when ``x`` is bf16, doubling activation grad bandwidth and breaking any downstream LN bwd that pins a bf16 layout. (Review thread "Why do we need this utility function?".) 2. "Why is this dtype casting required? I don't recall us needing it for the non-MoE LNMLP block." -- Expand the comment above the bwd activation fp32 promotion to explain the MoE-specific math: LN+MLP's silu sits behind a downstream LN that absorbs the bf16 rounding error, while MoE's silu sits on the *expert* side of routing -- the bf16 rounding rides directly into expert_outputs and is summed across topk experts by ep_combine. Bf16 silu alone drifts ~1% vs fp32 silu and compounds through wo->combine into the ~1.4% per-element parity gap we measured against the pure-JAX softmax reference. Mirroring the fwd's fp32 promotion in the bwd keeps silu' in lock-step with silu. (Review thread on "# Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply".) 3. "Do we have a use-case for user-specified alignments beyond 128 currently? ... it'd make sense to instead hardcode _ALIGN_SIZE = 128 as a constant at the top of the file for now to simplify this MoEBlock API. We can always expand the API to support a user-specified align size in the future." -- Implement the suggestion. Drop ``align_size`` from ``_moe_fwd_rule`` / ``_moe_bwd_rule`` / ``_moe`` / public ``moe()``; shift the ``custom_vjp`` ``nondiff_argnums`` from ``range(9, 27)`` -> ``range(9, 26)``; replace ``effective_align = max(int(align_size), 128)`` with the new module-level ``_ALIGN_SIZE = 128`` constant. Trim the ``moe()`` docstring accordingly. (Review thread on "natural_spe = num_ep * max_tokens_per_rank".) 4. "Which axis name inputs are physical mesh axes and why can be logical axes? ... No need to make any changes for now, I just want to assess which are which and then we can discuss if it makes sense to support logical on some/all or if some are required to be physical axes." -- Add an "Axis-name parameters" section to ``moe()``'s docstring listing which kwargs are physical mesh axes (``ep_axis``, ``data_parallelism_axes`` -- they index ``Mesh.shape`` directly to compute ``num_ep`` / ``dp_size`` and to construct the ``P((dp..., ep), None, None)`` for ``jax.lax.with_sharding_constraint``) vs logical axes (``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` -- resolved via the Flax logical-axis rules). Also document why ``ep_axis`` / ``data_parallelism_axes`` are intentionally non-logical: the EP comm-group construction (``dp_color = rank // ep_size``) and the bootstrap signature check both require concrete integer sizes. (Review thread on "batch_pspec_axis = (*data_parallelism_axes, ep_axis)".) 5. "Is this NaN filtering a debugging artifact or something we need in the final version?" -- Strengthen the inline comment above ``sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, ...)`` to explicitly call this out as a CORRECTNESS REQUIREMENT, not a debugging artifact: it covers the sigmoid+K>1 underflow path where top-K sigmoid scores all round to zero and the ``weights / (weights.sum + 1e-20)`` normalisation emits NaN. Observationally the filter is a no-op on the dense unit-test distributions, but it must stay in for sparse / production routing. (Review thread on "sparse_probs = jnp.where(jnp.isnan(sparse_probs), ...).") Not addressed in this commit (intentional): * Review thread on the ``align_size: int = 0`` placeholder in ``flax/moe.py`` ("Placeholder comment for me to fix this so align_size is inferred automatically based on the recipe and doesn't need to be specified by the user"). That's jberchtold's own follow-up. * Review thread on the explicit ``tree_flatten`` / ``tree_unflatten`` on ``_Ctx`` ("better to use the ``@flax_struct.dataclass``"). Deferred to a separate, testable commit because changing a ``custom_vjp`` residual's pytree registration touches subtle ordering / None-handling semantics that warrant their own bisect surface. * Review thread on ``use_bias`` / ``use_expert_bias`` renames -- handled in the immediately preceding commit ``jax/flax,tests: rename use_bias/use_expert_bias for symmetry``. * Review thread on the ``expert_bias`` fp32 init -- already resolved during the Phuong PR #3036 resync (the redundant ``jnp.float32`` second-dtype argument on ``self.param`` was dropped; ``expert_bias`` now lives at ``self.dtype``). Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 115 +++++++++++++++++++++++++--------- 1 file changed, 86 insertions(+), 29 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 554ee628b8..fe5c1f576f 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -58,15 +58,44 @@ __all__ = ["moe"] +# 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 a +# 128-token tile, so a single hard-coded constant suffices. +# +# We deliberately omit a user-facing knob: per PR #3116 review there's no +# current model that wants a >128 alignment, and exposing it widens the +# MoEBlock API surface without buying anything. Re-introduce a parameter +# (or recipe-driven inference, see jberchtold's follow-up) if a future +# recipe needs >128. +_ALIGN_SIZE = 128 + + def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: """Apply a sharding constraint while keeping bwd cotangents in the primal dtype. Plain ``jax.lax.with_sharding_constraint`` propagates cotangents in - whatever dtype the upstream gradient lands in; under mixed precision + whatever dtype the upstream gradient lands in. Under mixed precision that can be wider than the primal, blowing up bandwidth and (for bf16 primals) breaking downstream kernels that pin a bf16 input layout. This wrapper re-casts the cotangent back to the primal dtype and re-asserts the same sharding on the bwd path. + + Why MoE specifically needs this (per PR #3116 review): unlike a + plain LN+MLP block, the MoE bwd composes two cotangent paths into + ``d_x`` -- one through ``ep_dispatch_bwd`` (bf16) and one through + ``d_logits_2d @ gate_kernel.T``. The latter starts from + ``fused_topk_with_score_function_bwd``, which returns ``d_logits_2d`` + in fp32 because the fwd promoted ``logits_2d`` to fp32 (the topk / + softmax / sigmoid kernels are only validated at fp32; see + ``tests/pytorch/test_fused_router.py``). The fp32 ``d_logits_2d`` + then composes with ``gate_kernel.T`` and adds into the bf16 + ``d_x_from_dispatch``, yielding an fp32 sum even though the user's + ``x`` is bf16. Without this cast, the user-visible ``d_x`` flows + back into the optimizer at fp32 -- silently doubling the activation + grad bandwidth and tripping any downstream kernel that pins a bf16 + input layout (e.g. an LN bwd that fuses into our ``d_x``). """ @jax.custom_vjp @@ -524,6 +553,19 @@ def _ffn_bwd_per_shard( # Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply # so the silu derivative composes through the gradient at fp32 too; # cast back to the bf16 layout the wi grouped_quantize expects. + # + # Why MoE specifically needs this (per PR #3116 review): the + # non-MoE LN+MLP block can stay in the activation dtype because + # its silu accumulates over a single per-row dot product whose + # numerical drift is absorbed by the downstream LN. The MoE + # silu, by contrast, sits on the *expert* side of the routing, + # so its bf16 rounding error rides directly into ``expert_outputs`` + # and is summed (weighted by routing probs) across topk experts + # by ep_combine -- bf16 silu alone drifts ~1% vs fp32 silu, which + # compounds through wo->combine into the ~1.4% per-element parity + # gap we measured against the pure-JAX softmax reference. Mirroring + # the fwd fp32 promotion keeps the bwd's silu' derivative in lock- + # step with the fwd's silu and preserves grad parity. gp_fp32 = gate_proj_out.astype(jnp.float32) up_fp32 = up_proj_out.astype(jnp.float32) d_int_fp32 = d_intermediate.astype(jnp.float32) @@ -608,7 +650,6 @@ def _moe_fwd_rule( wo_kernel_axes, dtype, apply_topk_weights_early, - align_size, ): """Forward: gate -> topk -> ep_dispatch -> shard_map(FFN) -> ep_combine. @@ -653,10 +694,8 @@ def _moe_fwd_rule( # rejects the dispatch buffer with ``invalid argument``. natural_spe = num_ep * max_tokens_per_rank # = (B // dp_size) * S # NCCL EP requires each expert-major output block to be at least - # 128-token aligned. Keep larger caller-requested alignments, but - # do not emit a smaller natural block size for tiny tests. - effective_align = max(int(align_size), 128) - slots_per_expert = ((natural_spe + effective_align - 1) // effective_align) * effective_align + # ``_ALIGN_SIZE`` (=128) tokens; see the constant's docstring. + slots_per_expert = ((natural_spe + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE recv_pr = num_local_experts * slots_per_expert _te_ep_assert_compatible_bootstrap( @@ -706,15 +745,19 @@ def _moe_fwd_rule( expert_bias=eb_arg, compute_aux_scores=False, ) - # Sigmoid + K>1 normalises as `weights / (weights.sum + 1e-20)`; for - # tokens whose top-K sigmoid scores all underflow at bf16/fp32 the - # output is NaN at the selected positions. Those NaNs ride + # NOTE (PR #3116 review): this NaN filter is a *correctness + # requirement*, NOT a debugging artifact. Sigmoid + K>1 normalises + # as ``weights / (weights.sum + 1e-20)``; for tokens whose top-K + # sigmoid scores all underflow at bf16/fp32, the output is NaN + # at the selected positions. Those NaNs ride # ep_dispatch -> recv_topk_weights -> combine and poison the per-token # weighted sum, leaving entire output rows as NaN. Sanitize at the # source so neither the fwd combine nor the bwd's manual - # `grad_pre_combine * w` sees them. Padded positions in sparse_probs - # are already zero (routing_map is False there); only the rare - # underflow path emits NaN. + # ``grad_pre_combine * w`` sees them. Padded positions in + # sparse_probs are already zero (routing_map is False there); only + # the rare sigmoid-underflow path emits NaN, which is why the + # filter is observationally a no-op in dense unit tests but must + # stay in for sparse / production routing distributions. sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, sparse_probs).astype(dtype) # ---------------- Aux loss (global view, replicated) ---------------- @@ -965,12 +1008,11 @@ def _moe_bwd_rule( wo_kernel_axes, dtype, apply_topk_weights_early, - align_size, residuals, cotangents, ): """Backward mirror of :func:`_moe_fwd_rule`.""" - del num_groups, group_topk, dtype, align_size # captured in residuals / unused in bwd + del num_groups, group_topk, dtype # captured in residuals / unused in bwd from jax.experimental.shard_map import shard_map d_output, d_aux_loss = cotangents @@ -1245,7 +1287,7 @@ def _bwd_body(*args): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 27))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 26))) def _moe( x, gate_kernel, @@ -1273,7 +1315,6 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, - align_size, ): primal, _ = _moe_fwd_rule( x, @@ -1302,7 +1343,6 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, - align_size, ) return primal @@ -1331,7 +1371,6 @@ def moe( scaling_factor: float = 1.0, aux_loss_coeff: float = 0.0, apply_topk_weights_early: bool = False, - align_size: int = 0, ep_axis: str, data_parallelism_axes: Tuple[str, ...] = (), input_axes: Tuple[Optional[str], ...] = (), @@ -1359,16 +1398,35 @@ def moe( all-gather over the routing-side logits is inserted so the ``fused_moe_aux_loss`` kernel sees a global ``[T_global, E]`` view; this lives off the dispatch critical path. - align_size : int - Minimum per-expert slot alignment passed to ``tex.ep_prepare`` - as ``dispatch_output_per_expert_alignment``. ``0`` (default) - means use the NCCL-EP-required natural slot count - ``ep_size * max_tokens_per_rank == (B/dp)*S`` (the per-rank - all-tokens-to-one-expert worst case the HT kernel demands). - Any positive value rounds that count up to the nearest - multiple, growing the per-rank receive buffer accordingly. - Set to ``128`` for FP8 recipes that require 128-aligned - grouped-GEMM tiles. + + 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 (per PR #3116 review): + + * ``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 per-shard + ``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. See module docstring for the rest of the parameter semantics and the surrounding design rationale. @@ -1430,7 +1488,6 @@ def moe( wo_kernel_axes, dtype, apply_topk_weights_early, - align_size, ) if aux_loss_coeff <= 0.0: aux_loss = None From 09d5f78127f8f1d2250ca2a02830bd7a6f2c66ff Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 11 Jun 2026 17:09:25 -0700 Subject: [PATCH 09/23] jax/moe: strip PR-response framing from comments; drop sparse_probs NaN sanitizer * Rewrite the inline justifications added in 078a7d80 so each one reads as standalone code documentation, not as a reply to a reviewer: drop "per PR #3116 review", "review feedback", "Renamed from ... per PR ..." and similar PR/thread references from moe.py, flax/moe.py, and tests/jax/test_te_ep_moe.py. Technical content (why the fp32 promotion is needed for the MoE silu+multiply, why _with_sharding_constraint_cast_bwd exists, physical-vs-logical axis split in moe() docstring, the 128 alignment rationale) is preserved and reframed to be useful to a reader who has no PR context. * Drop the jnp.where(jnp.isnan(sparse_probs), 0, sparse_probs) guard. Tracing fused_topk_with_score_function.cu shows the kernel divides by sum_scores + 1e-20, so finite non-negative sigmoid scores cannot produce NaN here; the filter was only defense against upstream NaNs, which would mask a real regression if anything ever did start producing them. Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 7 +-- transformer_engine/jax/flax/moe.py | 15 ++--- transformer_engine/jax/moe.py | 94 ++++++++++-------------------- 3 files changed, 41 insertions(+), 75 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 3a5cbff51a..ecc3192b13 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -531,10 +531,9 @@ def _make_inputs(key): # multiply (intermediate * recv_w * mask) is debugged. Late # weighting (combine-side) is unaffected and stays covered above. # Note: align_size is no longer a user-facing parameter; it is - # hard-coded to _ALIGN_SIZE = 128 in moe.py (per PR #3116 - # review). Re-add a distinct align-size config only if the - # constant is loosened, or a recipe-driven inference is added - # that selects a >128 alignment. + # hard-coded to _ALIGN_SIZE = 128 in moe.py. Re-add a distinct + # align-size config only if the constant is loosened, or a + # recipe-driven inference is added that selects a >128 alignment. pytest.param( dict(score_function="sigmoid"), id="sigmoid", diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index ed36ab835f..640db29534 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -85,8 +85,7 @@ class _MoEBlock(TransformerEngineBase): If ``True``, registers a per-expert routing bias (shape ``[E]``) used by the topk selection. Only meaningful with ``score_function="sigmoid"``; the underlying primitive validates - the pairing. (Renamed from ``use_expert_bias`` per PR #3116 - review for symmetry with ``use_ffn_bias``.) + the pairing. aux_loss_coeff : float If ``> 0``, return the MoE auxiliary load-balancing loss scalar in addition to the main output. @@ -106,19 +105,17 @@ class _MoEBlock(TransformerEngineBase): *inside* each shard before ``ep_combine`` (saves one global reduction at the cost of an extra broadcast). Default ``False``. - Note that the per-expert dispatch-slot alignment is fixed internally - at 128 tokens (see ``moe._ALIGN_SIZE``). Per PR #3116 review there's - no current model that wants a >128 alignment, so this is not exposed - as a parameter; re-introduce a knob (or recipe-driven inference) if - a future FP8 recipe needs >128. + 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. dtype : jnp.dtype Compute / parameter dtype. kernel_init, bias_init, expert_bias_init : Initializers. use_ffn_bias : bool Register per-expert FFN biases (``wi_0_bias``, ``wi_1_bias``, - ``wo_bias``). (Renamed from ``use_bias`` per PR #3116 review - for symmetry with ``use_expert_routing_bias``.) + ``wo_bias``). Quantization is currently configured via the standard TE autocast context (``fp8_autocast``/``with_quantizer_set``) and threaded diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index fe5c1f576f..18c9343cf9 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -61,41 +61,32 @@ # 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 a -# 128-token tile, so a single hard-coded constant suffices. -# -# We deliberately omit a user-facing knob: per PR #3116 review there's no -# current model that wants a >128 alignment, and exposing it widens the -# MoEBlock API surface without buying anything. Re-introduce a parameter -# (or recipe-driven inference, see jberchtold's follow-up) if a future -# recipe needs >128. +# 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 def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: - """Apply a sharding constraint while keeping bwd cotangents in the primal dtype. - - Plain ``jax.lax.with_sharding_constraint`` propagates cotangents in - whatever dtype the upstream gradient lands in. Under mixed precision - that can be wider than the primal, blowing up bandwidth and (for - bf16 primals) breaking downstream kernels that pin a bf16 input - layout. This wrapper re-casts the cotangent back to the primal - dtype and re-asserts the same sharding on the bwd path. - - Why MoE specifically needs this (per PR #3116 review): unlike a - plain LN+MLP block, the MoE bwd composes two cotangent paths into - ``d_x`` -- one through ``ep_dispatch_bwd`` (bf16) and one through - ``d_logits_2d @ gate_kernel.T``. The latter starts from - ``fused_topk_with_score_function_bwd``, which returns ``d_logits_2d`` - in fp32 because the fwd promoted ``logits_2d`` to fp32 (the topk / - softmax / sigmoid kernels are only validated at fp32; see - ``tests/pytorch/test_fused_router.py``). The fp32 ``d_logits_2d`` - then composes with ``gate_kernel.T`` and adds into the bf16 - ``d_x_from_dispatch``, yielding an fp32 sum even though the user's - ``x`` is bf16. Without this cast, the user-visible ``d_x`` flows - back into the optimizer at fp32 -- silently doubling the activation - grad bandwidth and tripping any downstream kernel that pins a bf16 - input layout (e.g. an LN bwd that fuses into our ``d_x``). + """Sharding constraint that keeps bwd cotangents in the primal dtype. + + Plain ``jax.lax.with_sharding_constraint`` is identity on the fwd + but does not constrain the dtype of the cotangent that flows back + through it. In this MoE bwd, ``d_x`` is built from two paths: + + * ``d_x_from_dispatch`` from ``ep_dispatch_bwd`` -- primal dtype + (bf16 in mixed precision). + * ``d_x_from_gate = d_logits_2d @ gate_kernel.T`` where + ``d_logits_2d`` is produced by + ``fused_topk_with_score_function_bwd``. That primitive runs at + fp32 because the fwd promoted ``logits_2d`` to fp32 (the fused + topk/softmax/sigmoid kernels are only validated at fp32). + + JAX's type promotion then makes ``d_x_from_gate + d_x_from_dispatch`` + fp32, so the user-visible ``d_x`` ends up wider than ``x``. That + doubles activation-grad bandwidth and breaks any downstream kernel + that pins a bf16 input layout. This wrapper inserts an explicit + cast back to the primal dtype on the bwd side and re-asserts the + same sharding there as well. """ @jax.custom_vjp @@ -550,22 +541,14 @@ def _ffn_bwd_per_shard( else: d_recv_w_from_intermediate = jnp.zeros_like(recv_w_flat) - # Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply - # so the silu derivative composes through the gradient at fp32 too; - # cast back to the bf16 layout the wi grouped_quantize expects. - # - # Why MoE specifically needs this (per PR #3116 review): the - # non-MoE LN+MLP block can stay in the activation dtype because - # its silu accumulates over a single per-row dot product whose - # numerical drift is absorbed by the downstream LN. The MoE - # silu, by contrast, sits on the *expert* side of the routing, - # so its bf16 rounding error rides directly into ``expert_outputs`` - # and is summed (weighted by routing probs) across topk experts - # by ep_combine -- bf16 silu alone drifts ~1% vs fp32 silu, which - # compounds through wo->combine into the ~1.4% per-element parity - # gap we measured against the pure-JAX softmax reference. Mirroring - # the fwd fp32 promotion keeps the bwd's silu' derivative in lock- - # step with the fwd's silu and preserves grad parity. + # Activation bwd. The fwd already computes silu+multiply at fp32 + # because the MoE silu sits on the expert side of routing: its + # output rides into ``expert_outputs`` and is then summed -- weighted + # by routing probabilities -- across topk experts by ep_combine. + # Doing silu/silu' in bf16 drifts by ~1% per element vs fp32 and + # that drift compounds through wo->combine. Mirror the fwd's fp32 + # promotion here so silu' lines up with silu, then cast back to the + # bf16 layout the wi grouped_quantize expects. gp_fp32 = gate_proj_out.astype(jnp.float32) up_fp32 = up_proj_out.astype(jnp.float32) d_int_fp32 = d_intermediate.astype(jnp.float32) @@ -745,20 +728,7 @@ def _moe_fwd_rule( expert_bias=eb_arg, compute_aux_scores=False, ) - # NOTE (PR #3116 review): this NaN filter is a *correctness - # requirement*, NOT a debugging artifact. Sigmoid + K>1 normalises - # as ``weights / (weights.sum + 1e-20)``; for tokens whose top-K - # sigmoid scores all underflow at bf16/fp32, the output is NaN - # at the selected positions. Those NaNs ride - # ep_dispatch -> recv_topk_weights -> combine and poison the per-token - # weighted sum, leaving entire output rows as NaN. Sanitize at the - # source so neither the fwd combine nor the bwd's manual - # ``grad_pre_combine * w`` sees them. Padded positions in - # sparse_probs are already zero (routing_map is False there); only - # the rare sigmoid-underflow path emits NaN, which is why the - # filter is observationally a no-op in dense unit tests but must - # stay in for sparse / production routing distributions. - sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, sparse_probs).astype(dtype) + sparse_probs = sparse_probs.astype(dtype) # ---------------- Aux loss (global view, replicated) ---------------- # ``fused_moe_aux_loss_fwd`` sums probs and tokens_per_expert across @@ -1403,7 +1373,7 @@ def moe( 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 (per PR #3116 review): + Axis-name parameters: * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh axis names* -- they index ``jax.sharding.Mesh.shape`` directly From 9d959f29116c1b4f0e4a35fb240e31c720fbc865 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 12 Jun 2026 11:31:40 -0700 Subject: [PATCH 10/23] jax/moe: drop fp32 island around silu+multiply (fwd, bwd, reference) The SwiGLU intermediate (activation inputs gate_proj_out/up_proj_out, silu+multiply, and activation output) was previously promoted to fp32 in _ffn_fwd_per_shard and again in _ffn_bwd_per_shard, then cast back to the wi/wo GEMM dtype. The promotion bought nothing: the activation inputs come out of the wi grouped_gemm in bf16, the activation output is consumed by the wo GEMM (or wo's quantizer for FP8/FP4) in the same dtype, and storing higher precision than either consumer is wasted bandwidth. * _ffn_fwd_per_shard: drop the .astype(jnp.float32) on gate_proj_out and up_proj_out and the trailing .astype(sorted_x.dtype). The multiply now stays in the wi GEMM output dtype end-to-end. * _ffn_bwd_per_shard: symmetric simplification. jax.vjp(act_fn, ...) runs at bf16, both d_intermediate * silu' and d_intermediate * up stay at bf16, no casts. silu' is now consistent with silu (both bf16) so the chain rule composes cleanly without the prior fp32 detour. * tests/jax/test_te_ep_moe.py::_pure_jax_moe_reference: drop the matching fp32 silu in the parity reference so the test compares bf16-vs-bf16. Parity tolerance was not loosened; expect the comparison to tighten now that both sides round silu identically. Also fix an inaccurate inline comment at the apply_topk_weights_early fwd branch: the bf16 requirement on expert_outputs is enforced by ep_bootstrap (which rejects max_token_dtype != bf16 and sizes the NCCL EP HT mega-buffer for 2-byte slots accordingly), not by a runtime assert in the combine FFI. Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 6 +++-- transformer_engine/jax/moe.py | 42 ++++++++++++++--------------------- 2 files changed, 21 insertions(+), 27 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index ecc3192b13..428379d3bd 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -334,8 +334,10 @@ def _pure_jax_moe_reference( # both placements. layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) - intermediate = jax.nn.silu(layer_w0.astype(jnp.float32)) * layer_w1.astype(jnp.float32) - intermediate = intermediate.astype(x.dtype) + # Activation runs in x.dtype (typically bf16) to mirror the impl -- + # the impl keeps silu+multiply in the wi GEMM output dtype because + # storing higher precision than the consumer (wo) GEMM buys nothing. + intermediate = jax.nn.silu(layer_w0) * layer_w1 expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] output_2d = jnp.einsum("te,teh->th", routing_weights_full.astype(x.dtype), expert_out) output = output_2d.reshape(B, S, H).astype(x.dtype) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 18c9343cf9..c7fe5b5f7c 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -399,23 +399,23 @@ def _ffn_fwd_per_shard( casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS) casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS) - # Promote the silu+multiply to fp32 to match the pure-JAX reference - # (and ML common practice). bf16 silu accumulation alone drifts ~1% - # vs fp32 silu, which composes through wo -> combine into the - # ~1.4% per-element parity gap we were seeing on softmax. Cast back - # to the activation dtype before the grouped_quantize so the wo GEMM - # input layout is unchanged. + # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM + # output dtype; the activation output (`intermediate`) stays in the + # dtype the wo GEMM / wo's quantized input consumes. For bf16 compute + # that's all bf16; for FP8/FP4 the downstream grouped_quantize is what + # transitions to the target precision. Storing a higher precision than + # the consumer GEMM buys nothing. act_fn = _convert_to_activation_function(activation_type) - intermediate = ( - act_fn(gate_proj_out.astype(jnp.float32)) * up_proj_out.astype(jnp.float32) - ).astype(sorted_x.dtype) + intermediate = act_fn(gate_proj_out) * up_proj_out if apply_topk_weights_early: # 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, modulo elementwise op fusion gains. w_b is # cast to intermediate.dtype so the multiply doesn't promote - # expert_outputs to f32 (NCCL EP combine hard-asserts bf16). + # expert_outputs above the EP buffer's element width + # (ep_bootstrap rejects max_token_dtype != bf16, and the NCCL EP + # HT mega-buffer is sized for 2-byte slots accordingly). w_b = recv_w_flat[:, None].astype(intermediate.dtype) mask_b = (recv_w_flat != 0).astype(intermediate.dtype)[:, None] intermediate = intermediate * w_b * mask_b @@ -541,21 +541,13 @@ def _ffn_bwd_per_shard( else: d_recv_w_from_intermediate = jnp.zeros_like(recv_w_flat) - # Activation bwd. The fwd already computes silu+multiply at fp32 - # because the MoE silu sits on the expert side of routing: its - # output rides into ``expert_outputs`` and is then summed -- weighted - # by routing probabilities -- across topk experts by ep_combine. - # Doing silu/silu' in bf16 drifts by ~1% per element vs fp32 and - # that drift compounds through wo->combine. Mirror the fwd's fp32 - # promotion here so silu' lines up with silu, then cast back to the - # bf16 layout the wi grouped_quantize expects. - gp_fp32 = gate_proj_out.astype(jnp.float32) - up_fp32 = up_proj_out.astype(jnp.float32) - d_int_fp32 = d_intermediate.astype(jnp.float32) - act_gp_fp32, dact_pullback_fp32 = jax.vjp(act_fn, gp_fp32) - d_up_proj_out = (d_int_fp32 * act_gp_fp32).astype(up_proj_out.dtype) - (d_gate_proj_fp32,) = dact_pullback_fp32(d_int_fp32 * up_fp32) - d_gate_proj_out = d_gate_proj_fp32.astype(gate_proj_out.dtype) + # 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_out) + d_up_proj_out = d_intermediate * act_gp + (d_gate_proj_out,) = dact_pullback(d_intermediate * up_proj_out) # wi bwd (fused gate/up via concat). Mirror the fused fwd: pack the # gate/up cotangents along the trailing axis, run a single From b7d3a85f07f75cc2241448f423a0465faaa78440 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 12 Jun 2026 15:15:05 -0700 Subject: [PATCH 11/23] remove useless comments Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 121 +++++----------------------------- 1 file changed, 16 insertions(+), 105 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index c7fe5b5f7c..ee61540801 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -317,43 +317,10 @@ def _ffn_fwd_per_shard( them as ``[num_procs, recv_pr, H_out]``) plus the residuals consumed by the bwd. - ``token_counts_local`` is the per-expert padded token count (shape - ``[1, num_local_experts]``) from ``tex.ep_prepare``. With NCCL EP's - HT expert-major layout, the dispatch lays out experts contiguously - in ``recv_tokens`` as ``[expert_0_padded | expert_1_padded | ... | - overalloc_tail]``, where each per-expert block already includes the - ``dispatch_output_per_expert_alignment`` zero-padding and only the - trailing overalloc tail (slack between ``sum(token_counts)`` and the - worst-case ``recv_pr``) is unused. Plumbing ``token_counts`` straight - into ``grouped_gemm`` as ``group_sizes`` makes cuBLAS skip both the - overalloc tail (saving FMAs on partially-loaded shards) and any - expert whose per-shard routed count is zero (saving the GEMM - altogether, not just the rows). - - cuBLAS leaves the trailing of each grouped_gemm *per-row* output - (``combined_out``, ``expert_outputs`` in fwd; ``d_intermediate``, - ``d_sorted_x`` in bwd) uninitialised past ``sum(group_sizes)``, and - fully uninitialised on a shard whose every local expert has count 0. - That per-row tail is harmless for everything in this block: - subsequent ``grouped_gemm`` / ``ep_combine`` / ``ep_dispatch_bwd`` - only read valid rows per ``local_group_sizes`` / ``handle_mem``, and - ``act_fn``'s NaN tail only contaminates positions that no - group-aware consumer reads. The one per-row exception is - ``grouped_dbias`` (a ``segment_sum`` that walks every row), which - is only reached when ``has_bias=True``. That FFN bias path is - currently gated upstream (cuBLAS grouped_gemm has no fused bias on - Hopper yet; PR 3083 adds a pure-JAX bias add), so we don't pay for - the tail-zeroing masks needed to keep ``segment_sum`` well-defined. - If the bias path ever lands, re-add ``jnp.where`` masks on - ``combined_out`` / ``d_eo_2d`` / ``d_intermediate``. - - Separately, the grouped_gemm *wgrad* outputs (``d_wo``, - ``d_wi_combined`` in bwd) are per-group ``(num_local_experts, K, N)`` - and are the *user-visible* weight gradients. cuBLAS skips groups - with ``size_g == 0`` without zero-filling, so for 0-token-globally - experts the slice would be NaN and leak into the optimizer. This - is handled by per-group ``jnp.where`` masks (``wgrad_group_active``) - in ``_ffn_bwd_per_shard``; see that docstring. + ``token_counts_local`` (``[1, num_local_experts]``, from + ``tex.ep_prepare``) is passed to ``grouped_gemm`` as ``group_sizes`` + so cuBLAS skips both 0-token-routed experts and the dispatch + overalloc tail. """ hidden = recv_tokens_local.shape[-1] sorted_x = recv_tokens_local.reshape(-1, hidden) @@ -365,22 +332,10 @@ def _ffn_fwd_per_shard( wi_1 = wi_1.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) - # wi GEMM uses ONE fused grouped_gemm with the gate/up weights - # concatenated along the trailing (output) axis: wi_combined has - # shape ``(num_local_experts, hidden, 2*H_inter)`` and the resulting - # combined_out has shape ``(num_rows, 2*H_inter)``, which jnp.split - # cleanly slices back into gate / up halves. tex.grouped_gemm only - # supports the canonical (G, K, N) 3D weight layout with - # contracting_dims=((1,),(1,)) -- see the docstring on - # transformer_engine.jax.dense.grouped_dense ("currently only - # supports ((1,), (1,))") and the CI test - # tests/jax/test_multi_process_distributed_grouped_gemm.py. - # An older fused 4D variant built via jnp.stack([wi_0, wi_1], axis=-2) - # put a non-contracting axis in the middle of the RHS, which the - # kernel walked as if it were 3D and read off the end -> NaN. - # Bisected against a jnp.einsum reference: the stack-axis variant - # produced all-NaN output, while the concat-axis variant (this - # path) produces finite outputs matching the reference. + # Concat wi_0/wi_1 along the trailing axis (NOT stack on a new + # axis). grouped_gemm requires the 3D (G, K, N) weight layout with + # contracting_dims=((1,), (1,)); a 4D stack variant walks off the + # end of the RHS and returns NaN. wi_combined = jnp.concatenate([wi_0, wi_1], axis=-1) wi_combined_bias = ( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None @@ -403,8 +358,7 @@ def _ffn_fwd_per_shard( # output dtype; the activation output (`intermediate`) stays in the # dtype the wo GEMM / wo's quantized input consumes. For bf16 compute # that's all bf16; for FP8/FP4 the downstream grouped_quantize is what - # transitions to the target precision. Storing a higher precision than - # the consumer GEMM buys nothing. + # transitions to the target precision. act_fn = _convert_to_activation_function(activation_type) intermediate = act_fn(gate_proj_out) * up_proj_out @@ -468,44 +422,16 @@ def _ffn_bwd_per_shard( """Per-shard FFN backward. Mirrors :func:`_ffn_fwd_per_shard`. Returns - ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], d_wi_0, d_wi_1, d_wo, - d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. - - ``local_group_sizes`` arrives as ``(1, num_local_experts)`` (the - fwd-side shard residual), with the same per-expert padded counts the - fwd used as ``grouped_gemm`` ``group_sizes``. cuBLAS leaves rows - past ``sum(group_sizes)`` uninit in the bwd grouped_gemm *dgrad* - outputs (``d_intermediate``, ``d_sorted_x``), but every downstream - consumer of those per-row outputs is group-aware (wi/wo wgrads - contract only over valid rows; ``ep_dispatch_bwd`` reads only - valid positions per ``handle_mem``), so the per-row trailing tail - sits unread. ``grouped_dbias`` (per-row ``segment_sum``) is the - only per-row consumer that would propagate the tail, and it is - only invoked when ``has_bias=True``; that FFN bias path is gated - upstream (see ``_ffn_fwd_per_shard`` docstring) so we skip the - per-row tail-zeroing masks until cuBLAS gains fused-bias grouped - GEMM (or PR 3083's pure-JAX bias add lands). - - The *wgrad* outputs (``d_wo``, ``d_wi_combined``) are different. - They're per-group ``(num_local_experts, K, N)``, and cuBLAS - skips groups with ``size_g == 0`` without zero-filling the - corresponding ``out[g, :, :]`` slice (see - ``cublaslt_grouped_gemm.cu`` lines 2086/2096). For shards hosting - an expert that received zero tokens globally, that expert's - ``d_wo`` / ``d_wi`` slice would be uninit → NaN propagates to the - user's optimizer. We zero those slices via a per-group ``jnp.where`` - immediately after each wgrad. The mask is shape ``(num_groups, 1, 1)`` - and ``num_groups == num_local_experts`` is tiny, so this is cheap. + ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], + d_wi_0, d_wi_1, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. """ local_group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) recv_w_flat = recv_topk_weights_local.reshape(-1) q_set = noop_quantizer_set - # Per-group active mask for wgrad outputs. cuBLAS grouped_gemm skips - # groups with size_g == 0 and leaves the corresponding output slice - # uninit; without this, ``d_wo[g] / d_wi_combined[g]`` for any expert - # that received zero tokens globally would be NaN and propagate to - # the user's optimizer. + # cuBLAS grouped_gemm skips size_g == 0 groups without zero-filling + # the output slice; mask 0-token-expert wgrads to zero so the + # optimizer never sees uninit memory. wgrad_group_active = (local_group_sizes > 0)[:, None, None] # wo bwd @@ -568,9 +494,7 @@ def _ffn_bwd_per_shard( casted_d_combined.get_tensor(usage=TensorUsage.RHS), contracting_dims=((0,), (0,)), ) - d_wi_combined = jnp.where( - wgrad_group_active, d_wi_combined, jnp.zeros_like(d_wi_combined) - ) + d_wi_combined = jnp.where(wgrad_group_active, d_wi_combined, jnp.zeros_like(d_wi_combined)) d_wi_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1) if has_bias: d_wi_combined_bias = tex.grouped_dbias(d_combined, local_group_sizes) @@ -1015,26 +939,13 @@ def _moe_bwd_rule( d_expert_outputs = grad_pre_combine d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) else: - # combine_fwd consumed weighted = expert_out * w * mask; - # split the cotangent across both factors. w is cast to - # grad_pre_combine.dtype so the multiply stays bf16 and - # d_sorted_x (downstream into ep_dispatch_bwd) stays bf16. - # # ep_dispatch_fwd can land NaN into recv_topk_weights on padded - # slots (the public NCCL EP HT path does not zero-fill unused - # recv buffer entries). Untreated, `(NaN != 0) == True` in IEEE, + # slots. Untreated, `(NaN != 0) == True` in IEEE, # so the multiplicative mask cannot suppress the NaN and it # propagates through grad_pre_combine * w * mask into d_expert_outputs # and then into every downstream gradient (gate_kernel ends up # all-NaN). Sanitize once here. recv_w_clean = jnp.where(jnp.isnan(ctx.recv_topk_weights), 0, ctx.recv_topk_weights) - # IEEE 754: NaN * 0 = NaN, so multiplying grad_pre_combine by a - # 0/1 mask cannot kill the NaNs tex.ep_combine_bwd leaves at - # padded slots of grad_pre_combine: ctx.recv_topk_weights is - # clean after the sanitize above, but grad_pre_combine[padded] - # is still NaN, so grad_pre_combine * w * mask = NaN. Use - # jnp.where to overwrite padded positions with literal 0 - # instead. w = recv_w_clean[..., None].astype(grad_pre_combine.dtype) mask_bool = (recv_w_clean != 0)[..., None] d_expert_outputs = jnp.where( From ff50f447b94bfe6a5b65f39c1a9bb3a9a4b5a8e3 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 12 Jun 2026 15:32:03 -0700 Subject: [PATCH 12/23] tests/jax: remove legacy MoE VJP tests + launcher; point CI at TE-EP successor test_moe_vjp.py and test_multiprocess_moe_vjp.py both import PermutationBackend from transformer_engine.jax.moe -- an API that was removed during the Phuong PR #3036 resync. Both files have been dead-on-import ever since; the multiprocess launcher run_multiprocess_moe_vjp.sh only points at the dead test. test_te_ep_moe.py (the TE-EP-only custom_vjp suite) already covers everything the legacy files exercised that is still meaningful: fwd, bwd parity vs the pure-JAX reference, aux loss, both score functions, multi-process. The legacy parametrize axis (PermutationBackend.PURE_JAX vs TRITON) no longer exists. * Delete tests/jax/test_moe_vjp.py * Delete tests/jax/test_multiprocess_moe_vjp.py * Delete tests/jax/run_multiprocess_moe_vjp.sh * qa/L0_jax_distributed_unittest/test.sh: switch the MoE VJP distributed suite invocation from run_multiprocess_moe_vjp.sh / test_multiprocess_moe_vjp.py to run_te_ep_moe.sh / test_te_ep_moe.py. * tests/jax/conftest.py: docstring reference updated. * tests/jax/test_te_ep_moe.py: drop stale "successor to ..." aside and the "mirroring run_multiprocess_moe_vjp.sh" parenthetical. Net: -981 / +9. Signed-off-by: Teddy Do --- qa/L0_jax_distributed_unittest/test.sh | 8 +- tests/jax/conftest.py | 4 +- tests/jax/run_multiprocess_moe_vjp.sh | 132 -------- tests/jax/test_moe_vjp.py | 443 ------------------------- tests/jax/test_multiprocess_moe_vjp.py | 406 ---------------------- tests/jax/test_te_ep_moe.py | 8 +- 6 files changed, 9 insertions(+), 992 deletions(-) delete mode 100755 tests/jax/run_multiprocess_moe_vjp.sh delete mode 100644 tests/jax/test_moe_vjp.py delete mode 100644 tests/jax/test_multiprocess_moe_vjp.py diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index f86cea284e..a563e6908d 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -41,12 +41,12 @@ wait TE_PATH=$TE_PATH bash $TE_PATH/examples/jax/ep/run_test_ep.sh || test_fail "run_test_ep.sh" wait -# MoE custom_vjp distributed suite. Runs one Python process per GPU -# via tests/jax/run_multiprocess_moe_vjp.sh (mirrors the pattern in +# TE-EP MoE custom_vjp distributed suite. Runs one Python process per +# GPU via tests/jax/run_te_ep_moe.sh (mirrors the pattern in # examples/jax/encoder/run_test_multiprocessing_encoder.sh). Requires # >=4 visible GPUs. -TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_multiprocess_moe_vjp.sh \ - || test_fail "test_multiprocess_moe_vjp.py" +TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_te_ep_moe.sh \ + || test_fail "test_te_ep_moe.py" # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs; # auto-skips otherwise). CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)" diff --git a/tests/jax/conftest.py b/tests/jax/conftest.py index 74cb91202c..d729bfd1c7 100644 --- a/tests/jax/conftest.py +++ b/tests/jax/conftest.py @@ -90,8 +90,8 @@ def pytest_addoption(parser): """CLI options used by multiprocess JAX tests. ``--num-process`` and ``--process-id`` let a multiprocess launcher - (see ``tests/jax/run_multiprocess_moe_vjp.sh``) fork one pytest - process per GPU and tell each child its rank, so the test module + (see ``tests/jax/run_te_ep_moe.sh``) fork one pytest process per + GPU and tell each child its rank, so the test module can call ``jax.distributed.initialize(...)`` with the right ``local_device_ids``. Both default to 0; non-multiprocess tests ignore them. diff --git a/tests/jax/run_multiprocess_moe_vjp.sh b/tests/jax/run_multiprocess_moe_vjp.sh deleted file mode 100755 index 8dc1d2eb04..0000000000 --- a/tests/jax/run_multiprocess_moe_vjp.sh +++ /dev/null @@ -1,132 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. -# -# Multiprocess (one-GPU-per-process) launcher for the unified MoE VJP -# test suite. Forks one pytest invocation per visible GPU, passing each -# its own --num-process=N --process-id=i, and waits for all of them. -# Each child calls jax.distributed.initialize(..., local_device_ids= -# process_id) so each Python process only sees its one GPU as a local -# device and the participating processes form a global mesh. - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" -TEST_FILE="$TE_ROOT/tests/jax/test_multiprocess_moe_vjp.py" -PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" - -NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" -if [ "$NUM_GPUS" -lt 4 ]; then - echo "[run_multiprocess_moe_vjp.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2 - exit 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 MOE_VJP_COORDINATOR_ADDRESS="${MOE_VJP_COORDINATOR_ADDRESS:-127.0.0.1:13456}" - -echo "============================================================" -echo "MoE VJP MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" -echo " test file : $TEST_FILE" -echo " coordinator : $MOE_VJP_COORDINATOR_ADDRESS" -echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" -echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" -echo "============================================================" - -# Per-process logs. MOE_VJP_MP_LOG_DIR can be set to a host-mounted dir -# (e.g. when running inside a container that throws away /tmp on exit) -# so logs survive for postmortem inspection. Defaults to a fresh /tmp. -if [ -n "${MOE_VJP_MP_LOG_DIR:-}" ]; then - LOG_DIR="$MOE_VJP_MP_LOG_DIR" - mkdir -p "$LOG_DIR" -else - LOG_DIR=$(mktemp -d -t moe_vjp_mp_XXXXXX) -fi -echo "Per-process logs: $LOG_DIR" - -PIDS=() - -cleanup() { - for pid in "${PIDS[@]:-}"; do - if kill -0 "$pid" 2>/dev/null; then - kill -TERM "$pid" 2>/dev/null || true - fi - done - sleep 1 - for pid in "${PIDS[@]:-}"; do - if kill -0 "$pid" 2>/dev/null; then - kill -KILL "$pid" 2>/dev/null || true - fi - done -} -trap cleanup EXIT INT TERM - -# Launch one pytest per GPU. Process 0 streams to stdout; others log -# only to file so the live output isn't a mosaic. -for i in $(seq 0 $((NUM_GPUS - 1))); do - 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" & - else - "${PYTEST_CMD[@]}" > "$LOG_FILE" 2>&1 & - fi - PIDS+=("$!") -done - -# Wait for all and collect exit codes. -EXITS=() -for pid in "${PIDS[@]}"; do - if wait "$pid"; then - EXITS+=("0") - else - EXITS+=("$?") - fi -done - -# Summary. -echo -echo "============================================================" -echo "Per-process exit codes:" -for i in "${!EXITS[@]}"; do - echo " proc $i -> ${EXITS[$i]}" -done - -# Final pass/fail. Any non-zero in any process fails the suite, but -# we tolerate non-zero on the non-zero processes only if proc 0 -# reports PASS (this matches the encoder launcher's logic). Simplest -# 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. Anything else is a failure. -FAILED=0 -for e in "${EXITS[@]}"; do - if [ "$e" != "0" ] && [ "$e" != "5" ]; then - FAILED=1 - break - fi -done - -echo -if [ "$FAILED" -eq 0 ]; then - echo "[run_multiprocess_moe_vjp.sh] all processes PASSED" - if [ -z "${MOE_VJP_MP_LOG_DIR:-}" ]; then - rm -rf "$LOG_DIR" - fi - exit 0 -fi - -echo "[run_multiprocess_moe_vjp.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_moe_vjp.py b/tests/jax/test_moe_vjp.py deleted file mode 100644 index cc458d039e..0000000000 --- a/tests/jax/test_moe_vjp.py +++ /dev/null @@ -1,443 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Single-device tests for the unified MoE custom_vjp at -``transformer_engine.jax.moe.moe`` (and its Flax wrapper -``transformer_engine.jax.flax._MoEBlock``). - -Strategy --------- - -Rather than reproducing every internal kernel residual, we rely on a -single end-to-end pure-JAX *reference* implementation of the whole -MoE block (``_pure_jax_moe_reference`` below) and compare the TE -``moe(...)`` forward output AND parameter gradients against it. This -gives us coverage of: - -* the gate GEMM, -* the fused top-k routing primitive (and its bwd), -* the dispatch / per-expert FFN / combine pipeline (and their bwds - threaded through the absorbed primitives), -* the optional aux-loss path (and its bwd). - -The reference uses only ``jnp`` ops + ``jax.vjp``, so we get a -"definitive" pullback to compare against without needing the TE -primitive bwd kernels. - -Distributed (EP + FSDP) testing is intentionally NOT in this file -- -that needs a multi-device setup and lives in -``tests/jax/test_distributed_moe_vjp.py`` (follow-up). -""" - -from functools import partial -from typing import Optional, Tuple - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -from transformer_engine_jax import get_device_compute_capability -from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import PermutationBackend, moe - -# The MoE custom_vjp uses grouped GEMM, which is currently -# Blackwell-only (sm_100+). Skip the whole file on older arches. -if get_device_compute_capability(0) < 100: - pytest.skip( - "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", - allow_module_level=True, - ) - -# Parametrize values for the dispatch / combine backend. Only the -# ``triton`` variant is gated by the ``triton`` marker (so the -# ``pure_jax`` variant still runs on environments without Triton). -BACKEND_PARAMS = [ - pytest.param("pure_jax", id="pure_jax"), - pytest.param("triton", id="triton", marks=pytest.mark.triton), -] - - -# ----------------------------------------------------------------------------- -# Test config -# ----------------------------------------------------------------------------- - -DTYPE = jnp.float32 # use fp32 for tighter parity assertions -BATCH_SIZE = 2 -SEQUENCE_LENGTH = 16 -HIDDEN_SIZE = 32 -INTERMEDIATE_SIZE = 64 -NUM_EXPERTS = 8 -NUM_EXPERTS_PER_TOK = 2 - - -def _make_inputs(key: jax.Array, *, batch=BATCH_SIZE, seq=SEQUENCE_LENGTH) -> jax.Array: - return jax.random.normal(key, (batch, seq, HIDDEN_SIZE), dtype=DTYPE) - - -# ----------------------------------------------------------------------------- -# Pure-JAX reference MoE -# ----------------------------------------------------------------------------- -# -# Implements EXACTLY the same math as ``moe(...)`` for the no-EP, -# softmax-routing, no-bias, silu activation, no-quantization path. -# Returns ``(output, aux_loss_or_zero)``. Used as ground truth for both -# fwd and bwd parity. - - -@partial( - jax.jit, - static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff"), -) -def _pure_jax_moe_reference( - x: jnp.ndarray, - gate_kernel: jnp.ndarray, - wi_0: jnp.ndarray, - wi_1: jnp.ndarray, - wo: jnp.ndarray, - *, - num_experts: int, - num_experts_per_tok: int, - aux_loss_coeff: float = 0.0, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Reference no-EP MoE forward (pure JAX, no TE primitives). - - Mirrors :func:`transformer_engine.jax.moe._body_fwd` for the - PURE_JAX backend, no biases, softmax routing, silu activation, - no quantization. Linear ops only -- ``jax.vjp`` over this gives - the canonical bwd to compare against. - """ - B, S, H = x.shape - T = B * S - x_2d = x.reshape(T, H) - - # Gate - logits = x_2d @ gate_kernel # [T, E] - - # Softmax + topk (no expert_bias, no grouping, scale=1.0) - probs_full = jax.nn.softmax(logits, axis=-1) # [T, E] - # top-k by probability: - sorted_idx = jnp.argsort(probs_full, axis=-1) # ascending - selected = sorted_idx[:, -num_experts_per_tok:] # [T, K] - weights = jnp.take_along_axis(probs_full, selected, axis=-1) # [T, K] - # Normalize topk weights to sum to 1 (matches softmax->topk semantics - # of fused_topk_with_score_function with use_pre_softmax=False): - weights = weights / jnp.sum(weights, axis=-1, keepdims=True) - - # Build a sparse routing_map [T, E] with weights at selected positions - routing_weights_full = jnp.zeros_like(probs_full) - routing_weights_full = routing_weights_full.at[jnp.arange(T)[:, None], selected].set(weights) - - # Per-expert FFN: replicate each token K times, gather by expert, - # run through wi_0 / wi_1 / wo, gather back, weighted-sum. - # - # Vectorize the gather without sorting: for each (token, slot k), - # multiply the corresponding expert's FFN by routing_weights[t, k] - # and sum over experts. - # x_2d: [T, H], wi_0: [E, H, M], wi_1: [E, H, M], wo: [E, M, H] - # For each expert e: layer_w0_e = x_2d @ wi_0[e]; layer_w1_e = x_2d @ wi_1[e] - # intermediate_e = silu(layer_w0_e) * layer_w1_e - # expert_out_e = intermediate_e @ wo[e] - # output[t, h] = sum_e routing_weights_full[t, e] * expert_out_e[t, h] - layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) # [T, E, M] - layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) # [T, E, M] - intermediate = jax.nn.silu(layer_w0) * layer_w1 # [T, E, M] - expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] - output_2d = jnp.einsum("te,teh->th", routing_weights_full, expert_out) # [T, H] - output = output_2d.reshape(B, S, H) - - if aux_loss_coeff > 0.0: - # aux scores: clean per-expert softmax (compute_aux_scores=True - # kernel uses a clean softmax, no bias, scale=1, no grouping). - aux_probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1) - # tokens_per_expert from REAL routing_map (post-grouping); here - # there's no grouping so == count of non-zero positions per expert. - routing_map = (routing_weights_full > 0).astype(jnp.int32) - tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] - # aux_loss formula: (E * coeff / (k * T^2)) * sum_e - # (sum_t aux_probs[t, e]) * tokens_per_expert[e] - sum_probs_per_expert = jnp.sum(aux_probs, axis=0) # [E] - aux_loss = (num_experts * aux_loss_coeff / (num_experts_per_tok * (T**2))) * jnp.sum( - sum_probs_per_expert * tokens_per_expert.astype(jnp.float32) - ) - else: - aux_loss = jnp.zeros((), dtype=DTYPE) - - return output, aux_loss - - -# ----------------------------------------------------------------------------- -# Helpers -# ----------------------------------------------------------------------------- - - -def _init_params(key: jax.Array) -> dict: - k_g, k_w0, k_w1, k_wo = jax.random.split(key, 4) - init = jax.nn.initializers.variance_scaling(1.0, "fan_in", "truncated_normal") - return dict( - gate_kernel=init(k_g, (HIDDEN_SIZE, NUM_EXPERTS), DTYPE), - wi_0=init(k_w0, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), - wi_1=init(k_w1, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), - wo=init(k_wo, (NUM_EXPERTS, INTERMEDIATE_SIZE, HIDDEN_SIZE), DTYPE), - ) - - -@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) -def _run_te_moe( - x: jnp.ndarray, - params: dict, - *, - permutation_backend, - aux_loss_coeff: float = 0.0, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - return moe( - x, - params["gate_kernel"], - params["wi_0"], - params["wi_1"], - params["wo"], - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - activation_type="silu", - score_function="softmax", - use_pre_softmax=False, - scaling_factor=1.0, - aux_loss_coeff=aux_loss_coeff, - permutation_backend=permutation_backend, - align_size=0, - dtype=DTYPE, - ) - - -@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) -def _grads_te_main_loss(params, x, *, permutation_backend, aux_loss_coeff: float = 0.0): - """jit'd grad of ``mean(out**2)`` w.r.t. params (no aux contribution).""" - - def loss(params, x): - out, _ = _run_te_moe( - x, params, permutation_backend=permutation_backend, aux_loss_coeff=aux_loss_coeff - ) - return jnp.mean(out**2) - - return jax.grad(loss)(params, x) - - -@partial(jax.jit, static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff")) -def _grads_ref_main_loss(params, x, *, num_experts, num_experts_per_tok, aux_loss_coeff=0.0): - """jit'd grad of ``mean(out**2)`` w.r.t. params on the pure-JAX ref.""" - - def loss(params, x): - out, _ = _pure_jax_moe_reference( - x, - **params, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - aux_loss_coeff=aux_loss_coeff, - ) - return jnp.mean(out**2) - - return jax.grad(loss)(params, x) - - -@partial(jax.jit, static_argnames=("permutation_backend",)) -def _grad_te_aux_only(params, x, *, permutation_backend): - """jit'd grad of just the aux loss scalar (no main contribution).""" - - def aux_only(params, x): - _, aux = _run_te_moe( - x, params, permutation_backend=permutation_backend, aux_loss_coeff=1e-2 - ) - return aux.astype(jnp.float32) - - return jax.grad(aux_only)(params, x) - - -# ----------------------------------------------------------------------------- -# Tests -# ----------------------------------------------------------------------------- - - -class TestMoeVjpForward: - """Forward shape / finiteness / parity vs pure-JAX reference.""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_forward_shape_and_finite(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(0) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out, aux = _run_te_moe(x, params, permutation_backend=backend) - assert out.shape == x.shape - assert out.dtype == x.dtype - assert jnp.all(jnp.isfinite(out)) - assert aux is None - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_forward_parity_vs_pure_jax_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(1) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out_te, _ = _run_te_moe(x, params, permutation_backend=backend) - out_ref, _ = _pure_jax_moe_reference( - x, - **params, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - ) - # FP32, small shapes -> tight tolerance - np.testing.assert_allclose(np.array(out_te), np.array(out_ref), atol=2e-5, rtol=2e-5) - - def test_pure_jax_triton_equivalence(self): - key = jax.random.PRNGKey(2) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out_pj, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.PURE_JAX) - out_tr, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.TRITON) - np.testing.assert_allclose(np.array(out_pj), np.array(out_tr), atol=2e-5, rtol=2e-5) - - -class TestMoeVjpBackward: - """Backward parity vs pure-JAX reference (which uses ``jax.vjp`` over - plain JAX ops, giving us the canonical pullback).""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_grads_finite_and_nonzero(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(3) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - grads = _grads_te_main_loss(params, x, permutation_backend=backend) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g = grads[name] - assert jnp.all(jnp.isfinite(g)), f"{name} grad has NaN/Inf" - assert jnp.any(g != 0.0), f"{name} grad is identically zero" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_grads_match_pure_jax_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(4) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - grads_te = _grads_te_main_loss(params, x, permutation_backend=backend) - grads_ref = _grads_ref_main_loss( - params, - x, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - ) - # Loose-ish tol on grads: routing path has discrete topk so the - # softmax cotangent paths through the non-topk experts diverge - # slightly between TE (which uses the fused topk bwd) and the - # reference (which uses argsort-based take_along_axis). - # Tighter than the bf16 tests. - for name in ("wi_0", "wi_1", "wo"): - np.testing.assert_allclose( - np.array(grads_te[name]), - np.array(grads_ref[name]), - atol=5e-5, - rtol=5e-5, - err_msg=f"grad mismatch on {name}", - ) - # Gate grad has more error budget because it propagates through - # the topk derivative kernel (which differs in zero-pattern - # treatment from a plain take_along_axis). - np.testing.assert_allclose( - np.array(grads_te["gate_kernel"]), - np.array(grads_ref["gate_kernel"]), - atol=5e-4, - rtol=5e-4, - err_msg="grad mismatch on gate_kernel", - ) - - -class TestMoeVjpAuxLoss: - """Aux-loss path: forward + grad parity.""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_returned_and_finite(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(5) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - _, aux = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) - assert aux is not None - assert aux.shape == () - assert jnp.isfinite(aux) - assert jnp.abs(aux) < 1e2 - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_parity_vs_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(6) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - _, aux_te = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) - _, aux_ref = _pure_jax_moe_reference( - x, - **params, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - aux_loss_coeff=1e-2, - ) - np.testing.assert_allclose(float(aux_te), float(aux_ref), atol=1e-5, rtol=1e-5) - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_grads_propagate_to_logits(self, backend_name): - """The aux-loss bwd path must produce non-zero gate-kernel grads - when only the aux-loss scalar is differentiated (no main-output - contribution).""" - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(7) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - g_gate = _grad_te_aux_only(params, x, permutation_backend=backend)["gate_kernel"] - assert jnp.all(jnp.isfinite(g_gate)) - assert jnp.any( - g_gate != 0.0 - ), "aux_loss bwd should propagate to gate_kernel via fused_topk bwd" - - -# ----------------------------------------------------------------------------- -# Flax wrapper smoke test -# ----------------------------------------------------------------------------- - - -class TestMoEBlockFlaxWrapper: - """Sanity-check the thin Flax wrapper: forward + grad on init.""" - - def test_init_and_apply(self): - block = MoEBlock( - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - intermediate_size=INTERMEDIATE_SIZE, - permutation_backend=PermutationBackend.PURE_JAX, - dtype=DTYPE, - ) - key = jax.random.PRNGKey(8) - ki, kx = jax.random.split(key) - x = _make_inputs(kx) - variables = jax.jit(block.init)(ki, x) - out, aux = jax.jit(block.apply)(variables, x) - assert out.shape == x.shape - assert aux is None - - @jax.jit - def grad_fn(variables, x): - return jax.grad(lambda v, x: jnp.mean(block.apply(v, x)[0] ** 2))(variables, x) - - grads = grad_fn(variables, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g = grads["params"][name] - g = g.value if hasattr(g, "value") else g - assert jnp.all(jnp.isfinite(g)), f"{name} grad NaN/Inf" - assert jnp.any(g != 0.0), f"{name} grad zero" diff --git a/tests/jax/test_multiprocess_moe_vjp.py b/tests/jax/test_multiprocess_moe_vjp.py deleted file mode 100644 index 97044780f0..0000000000 --- a/tests/jax/test_multiprocess_moe_vjp.py +++ /dev/null @@ -1,406 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Multi-process (one-GPU-per-process) tests for the unified MoE custom_vjp. - -The launcher ``tests/jax/run_multiprocess_moe_vjp.sh`` forks one pytest -process per visible GPU (mirroring -``examples/jax/encoder/run_test_multiprocessing_encoder.sh``). Each -process binds to exactly one device via -``jax.distributed.initialize(..., local_device_ids=process_id)``; the -participating processes form a global mesh through JAX's distributed -runtime. - -How to run ----------- - -You typically do NOT invoke pytest on this file directly -- use the -launcher, which passes ``--num-process=N --process-id=i`` to each -forked process. Driving it directly with only one process will skip -every test because :func:`jax.distributed.initialize` requires -multiple participants. - - bash tests/jax/run_multiprocess_moe_vjp.sh - -CI invocation lives in ``qa/L0_jax_distributed_unittest/test.sh``. -""" - -import os - -# NCCL needs HBM headroom that JAX's default 90% preallocation does -# not leave. Set before any jax import below. -os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") -os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") - -import sys - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -from jax.experimental import mesh_utils -from jax.sharding import Mesh, NamedSharding, PartitionSpec as P -from flax.linen import partitioning as nn_partitioning - - -# Per-process distributed bootstrap. Each pytest invocation initializes -# JAX with exactly one local device (its assigned GPU). Once -# initialized, the four processes form one global mesh of 4 devices. -def _init_distributed(num_process: int, process_id: int) -> bool: - """Initialize jax.distributed for this pytest process. - - Returns True if initialization succeeded (i.e. this is a real - multi-process launch), False if num_process == 0 / 1 meaning the - file is being collected without a launcher and tests should be - skipped at module level. - """ - if num_process <= 1: - return False - coord = os.environ.get("MOE_VJP_COORDINATOR_ADDRESS", "127.0.0.1:1234") - jax.distributed.initialize( - coordinator_address=coord, - num_processes=num_process, - process_id=process_id, - local_device_ids=process_id, - ) - assert jax.local_device_count() == 1, "one GPU per process is the whole point" - assert ( - jax.device_count() == num_process - ), f"global device_count {jax.device_count()} != num_process {num_process}" - return True - - -# Read --num-process / --process-id BEFORE pytest collects any tests so -# we can fast-skip the whole module when not in a multiprocess launch. -def _read_mp_options(): - # Use pytest's option lookup via the request fixture isn't available - # at module top-level; parse argv ourselves the same way encoder - # test does. CLI form is e.g. "pytest ... --num-process=4 --process-id=0". - num = int(os.environ.get("MP_NUM_PROCESS", "0") or "0") - pid = int(os.environ.get("MP_PROCESS_ID", "0") or "0") - for i, a in enumerate(sys.argv): - if a.startswith("--num-process="): - num = int(a.split("=", 1)[1]) - elif a == "--num-process" and i + 1 < len(sys.argv): - num = int(sys.argv[i + 1]) - elif a.startswith("--process-id="): - pid = int(a.split("=", 1)[1]) - elif a == "--process-id" and i + 1 < len(sys.argv): - pid = int(sys.argv[i + 1]) - return num, pid - - -_MP_NUM_PROCESS, _MP_PROCESS_ID = _read_mp_options() -_MP_ACTIVE = _init_distributed(_MP_NUM_PROCESS, _MP_PROCESS_ID) - -if not _MP_ACTIVE: - # Skip the entire module if not launched via the multiprocess - # runner. Lets `pytest tests/jax/` collect this file harmlessly. - pytest.skip( - "test_multiprocess_moe_vjp.py requires the multiprocess launcher " - "(run_multiprocess_moe_vjp.sh). Skipping.", - allow_module_level=True, - ) - -from transformer_engine_jax import get_device_compute_capability - -# Grouped GEMM in the MoE custom_vjp currently requires Blackwell -# (sm_100+). Skip the whole file on older arches. -if get_device_compute_capability(0) < 100: - pytest.skip( - "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", - allow_module_level=True, - ) - -import transformer_engine.jax as te -from transformer_engine.common import recipe as te_recipe -from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import PermutationBackend -from transformer_engine.jax.sharding import MeshResource, global_shard_guard - -# Parametrize values for the dispatch / combine backend. Only the -# ``triton`` variant carries the ``triton`` marker, so the -# ``pure_jax`` variant still runs on environments without Triton. -BACKEND_PARAMS = [ - pytest.param("pure_jax", id="pure_jax"), - pytest.param("triton", id="triton", marks=pytest.mark.triton), -] - - -EP_AXIS = "ep" -FSDP_AXIS = "fsdp" -EP_SIZE = 2 -# FSDP_SIZE adapts to whatever the launcher gave us: dlcluster GB200 -# gives 4 GPUs (FSDP=2), CI B200 gives 8 GPUs (FSDP=4). Both stay -# 128-aligned for MXFP8 and divide num_experts/topk cleanly. -assert ( - jax.device_count() % EP_SIZE == 0 -), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" -FSDP_SIZE = jax.device_count() // EP_SIZE -NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE - -LOGICAL_AXIS_RULES = ( - ("exp", EP_AXIS), - ("embed", FSDP_AXIS), - ("mlp", None), - ("batch", (EP_AXIS, FSDP_AXIS)), -) - - -@pytest.fixture(scope="module") -def mesh(): - if jax.device_count() < NUM_DEVICES_REQUIRED: - pytest.skip( - f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" - f" have {jax.device_count()}" - ) - devices = mesh_utils.create_device_mesh((EP_SIZE, FSDP_SIZE)) - return Mesh(devices, axis_names=(EP_AXIS, FSDP_AXIS)) - - -# ``recipe`` parametrize values used across all tests below. ``None`` -# = plain bf16; the named recipes route through TE's autocast and -# exercise the FP8/MXFP8 quantization paths in _body_fwd/_body_bwd. -# Only recipes that work on TE Blackwell are included; older GPUs -# skip via the ``hardware_supports`` guard below. -RECIPE_NAMES = ("bf16", "MXFP8BlockScaling") - - -def _resolve_recipe(name): - """Return ``(use_fp8, recipe_instance)`` for the parametrize id.""" - if name == "bf16": - return False, None - if name == "MXFP8BlockScaling": - return True, te_recipe.MXFP8BlockScaling() - raise ValueError(f"unknown recipe name: {name!r}") - - -def _hardware_supports(recipe_name): - """Skip an FP8 recipe on GPUs that don't have the hw for it.""" - if recipe_name == "bf16": - return True - from transformer_engine_jax import get_device_compute_capability - - arch = get_device_compute_capability(0) - if recipe_name == "MXFP8BlockScaling": - return arch >= 100 - return False - - -def _autocast_ctx(recipe_name): - """Context manager that turns FP8 on for non-bf16 recipes.""" - use_fp8, recipe_inst = _resolve_recipe(recipe_name) - return te.autocast(enabled=use_fp8, recipe=recipe_inst) - - -def _tol_finite_grad(recipe_name): - """Per-recipe absolute tolerance for parity grad comparison.""" - if recipe_name == "bf16": - return 5e-2 - # MXFP8 grads carry block-scale quantization noise; loosen accordingly. - return 3e-1 - - -# ----------------------------------------------------------------------------- -# Helpers -# ----------------------------------------------------------------------------- - - -def _make_block( - *, - num_experts, - num_experts_per_tok, - intermediate_size, - permutation_backend, - aux_loss_coeff=0.0, - dtype=jnp.bfloat16, - align_size=0, -): - return MoEBlock( - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - intermediate_size=intermediate_size, - permutation_backend=permutation_backend, - data_parallelism_axes=(FSDP_AXIS,), - aux_loss_coeff=aux_loss_coeff, - dtype=dtype, - _align_size=align_size, - ) - - -def _shard_inputs(x, mesh): - return jax.lax.with_sharding_constraint( - x, NamedSharding(mesh, P((EP_AXIS, FSDP_AXIS), None, None)) - ) - - -def _init_apply(block, mesh, x, key): - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x = _shard_inputs(x, mesh) - variables = jax.jit(block.init)(key, x) - jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) - output, aux = jax.jit(block.apply)(variables, x) - jax.block_until_ready(output) - return variables, output, aux - - -def _grad_step(block, variables, mesh, x): - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x = _shard_inputs(x, mesh) - - def loss_fn(variables, x): - output, aux = block.apply(variables, x) - main = jnp.mean(output.astype(jnp.float32) ** 2) - return main + (aux.astype(jnp.float32) if aux is not None else 0.0) - - grads = jax.jit(jax.grad(loss_fn))(variables, x) - jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) - return grads - - -def _unwrap(x): - return x.value if hasattr(x, "value") else x - - -def _local_shard(x): - """Return the local (this-process) shard of a global JAX Array as numpy. - - Every assertion in this file is structural (finite-ness, non-zero, - parity within tolerance). For all of these, checking the local - shard on each process is sufficient and avoids any cross-process - collective in the test machinery. ``arr.addressable_data(0)`` - returns the local-device view of the sharded array -- with one - GPU per process there is exactly one addressable shard. - """ - return np.asarray(jax.device_get(x.addressable_data(0))) - - -# ----------------------------------------------------------------------------- -# Mixtral-style shapes, sized to fit on a single 4-GPU bf16 box (a -# 4-way data-parallel shard of a Mixtral-8 block). -# ----------------------------------------------------------------------------- - -BATCH = EP_SIZE * FSDP_SIZE * 4 # 16 on 4-GPU, 32 on 8-GPU -SEQ = 2048 -HIDDEN = 1024 -INTER = 4096 -NUM_EXPERTS = 8 -TOPK = 2 - - -class TestMoeVjpMultiprocess: - """Multiprocess (one-GPU-per-process) correctness checks for the - unified MoE custom_vjp. - """ - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_fwd_and_bwd(self, mesh, backend_name, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - backend = PermutationBackend(backend_name) - block = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=backend, - ) - x = jax.random.normal( - jax.random.PRNGKey(0), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - with _autocast_ctx(recipe_name): - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) - # Local-shard checks (see _local_shard docstring for why). - out_local = _local_shard(output) - assert output.dtype == x.dtype - assert np.all(np.isfinite(out_local)), "output has NaN/Inf" - assert aux is None - with _autocast_ctx(recipe_name): - grads = _grad_step(block, variables, mesh, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_local = _local_shard(_unwrap(grads["params"][name])) - assert np.all(np.isfinite(g_local)), f"{name} grad has NaN/Inf" - assert np.any(g_local != 0.0), f"{name} grad is identically zero" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_aux_loss(self, mesh, backend_name, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - backend = PermutationBackend(backend_name) - block = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=backend, - aux_loss_coeff=1e-2, - ) - x = jax.random.normal( - jax.random.PRNGKey(4), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - with _autocast_ctx(recipe_name): - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(5)) - out_local = _local_shard(output) - assert np.all(np.isfinite(out_local)), "output has NaN/Inf under aux" - assert aux is not None - assert aux.shape == () - aux_local = _local_shard(aux) - assert np.isfinite(aux_local), "aux is NaN/Inf" - with _autocast_ctx(recipe_name): - grads = _grad_step(block, variables, mesh, x) - g_gate_local = _local_shard(_unwrap(grads["params"]["gate_kernel"])) - assert np.all(np.isfinite(g_gate_local)), "gate grad NaN/Inf under aux" - - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_pure_jax_triton_parity(self, mesh, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - block_pj = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=PermutationBackend.PURE_JAX, - ) - block_tr = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=PermutationBackend.TRITON, - ) - x = jax.random.normal( - jax.random.PRNGKey(6), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - tol = _tol_finite_grad(recipe_name) - with _autocast_ctx(recipe_name): - variables, out_pj, _ = _init_apply(block_pj, mesh, x, jax.random.PRNGKey(7)) - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x_sh = _shard_inputs(x, mesh) - out_tr, _ = jax.jit(block_tr.apply)(variables, x_sh) - - out_pj_local = _local_shard(out_pj) - out_tr_local = _local_shard(out_tr) - diff = float(np.max(np.abs(out_pj_local - out_tr_local))) - assert diff < tol, f"forward parity breach: max_abs_diff={diff} (tol={tol})" - - with _autocast_ctx(recipe_name): - grads_pj = _grad_step(block_pj, variables, mesh, x) - grads_tr = _grad_step(block_tr, variables, mesh, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_pj = _local_shard(_unwrap(grads_pj["params"][name])) - g_tr = _local_shard(_unwrap(grads_tr["params"][name])) - d = float(np.max(np.abs(g_pj - g_tr))) - assert d < tol, f"grad parity breach on {name}: max_abs_diff={d} (tol={tol})" diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 428379d3bd..6f27ba1d33 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -5,8 +5,7 @@ """Multi-process (one-GPU-per-process) tests for the TE-EP MoE custom_vjp. The launcher ``tests/jax/run_te_ep_moe.sh`` forks one pytest process per -visible GPU (mirroring ``run_multiprocess_moe_vjp.sh``). Each process binds -to exactly one device via +visible GPU. Each process binds to exactly one device via ``jax.distributed.initialize(..., local_device_ids=process_id)``; the participating processes form a global ``(ep, fsdp)`` mesh through JAX's distributed runtime. @@ -26,9 +25,8 @@ What this suite covers ---------------------- -This file is the TE-EP-only successor to ``test_moe_vjp.py`` and -``test_multiprocess_moe_vjp.py``. Each test exercises one MoE-block -run and bundles every check that single run supports — shape, dtype, +Each test exercises one MoE-block run and bundles every check that +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: From 641b6b89548f3828366737358a56539a3ad75b98 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 12 Jun 2026 15:33:56 -0700 Subject: [PATCH 13/23] jax/moe: swap _Ctx to @flax.struct.dataclass, drop manual pytree boilerplate Per reviewer feedback (Jaberchtold on PR #3036): the manual tree_flatten / tree_unflatten on _Ctx duplicate exactly what @flax.struct.dataclass auto-generates, and the permutation dataclasses elsewhere in this module already use flax.struct. Switching to @flax.struct.dataclass: * Removes ~75 lines of mechanical tree_flatten / tree_unflatten that have to be kept in sync with the field list by hand. * Keeps cfg as the single static field via flax.struct.field(pytree_node=False), so the fwd -> bwd boundary behavior under jax.custom_vjp is unchanged. * Drops two now-unused imports (dataclasses.dataclass, jax.tree_util.register_pytree_node_class) and adds flax.struct. Field order and the (children, aux_data) split are byte-equivalent to the previous manual implementation, so the pytree treedef seen by jax.custom_vjp is identical. Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 97 ++++------------------------------- 1 file changed, 10 insertions(+), 87 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index ee61540801..09c9b5dd7c 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -35,15 +35,14 @@ path and overlaps with the dispatch collective). """ -from dataclasses import dataclass from functools import partial from typing import Any, Optional, Tuple, Union import warnings +import flax.struct import jax import jax.numpy as jnp from jax.sharding import NamedSharding, PartitionSpec as P -from jax.tree_util import register_pytree_node_class from . import cpp_extensions as tex from .quantize import ( @@ -180,14 +179,14 @@ def _te_ep_assert_compatible_bootstrap( # ============================================================================= -# Registered as a pytree so jax.custom_vjp can flatten/unflatten it across -# the fwd -> bwd boundary. ``cfg`` is the only static field (EpLayerConfig -# is a frozen dataclass of ints); the rest are jnp.ndarray, -# GroupedNoScaleTensor (already a pytree), or None when aux_loss_coeff == 0. -@register_pytree_node_class -@dataclass +@flax.struct.dataclass class _Ctx: - """Residuals carried from the fwd rule into the bwd rule.""" + """Residuals carried from the fwd rule into the bwd rule. + + Flattened automatically by jax.custom_vjp; ``cfg`` is the only + static field (the rest are jnp.ndarray, GroupedNoScaleTensor, or + None when aux_loss_coeff == 0). + """ x: jnp.ndarray gate_kernel: jnp.ndarray @@ -195,8 +194,8 @@ class _Ctx: logits_2d: jnp.ndarray saved_scores: jnp.ndarray routing_map: jnp.ndarray - cfg: Any - handle_mem: Any + cfg: Any = flax.struct.field(pytree_node=False) + handle_mem: jnp.ndarray token_counts: jnp.ndarray recv_topk_weights: jnp.ndarray casted_sorted_x_lhs_trans: Any @@ -207,86 +206,10 @@ class _Ctx: casted_wo_rhs_trans: Any expert_outputs: jnp.ndarray local_group_sizes: jnp.ndarray - # Aux-loss residuals; None when aux_loss_coeff == 0. aux_const_buf: Any = None aux_tokens_per_expert: Any = None aux_saved_scores: Any = None - def tree_flatten(self): - children = ( - self.x, - self.gate_kernel, - self.expert_bias, - self.logits_2d, - self.saved_scores, - self.routing_map, - self.handle_mem, - self.token_counts, - self.recv_topk_weights, - self.casted_sorted_x_lhs_trans, - self.casted_wi_rhs_trans, - self.gate_proj_out, - self.up_proj_out, - self.casted_intermediate_lhs_trans, - self.casted_wo_rhs_trans, - self.expert_outputs, - self.local_group_sizes, - self.aux_const_buf, - self.aux_tokens_per_expert, - self.aux_saved_scores, - ) - aux_data = (self.cfg,) - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): - (cfg,) = aux_data - ( - x, - gate_kernel, - expert_bias, - logits_2d, - saved_scores, - routing_map, - handle_mem, - token_counts, - recv_topk_weights, - casted_sorted_x_lhs_trans, - casted_wi_rhs_trans, - gate_proj_out, - up_proj_out, - casted_intermediate_lhs_trans, - casted_wo_rhs_trans, - expert_outputs, - local_group_sizes, - aux_const_buf, - aux_tokens_per_expert, - aux_saved_scores, - ) = children - return cls( - 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, - token_counts=token_counts, - 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, - aux_const_buf=aux_const_buf, - aux_tokens_per_expert=aux_tokens_per_expert, - aux_saved_scores=aux_saved_scores, - ) - # ============================================================================= # Per-shard FFN body (runs inside shard_map) From 055b72cc3f6642f70246753e53d938eb9c627ceb Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 12 Jun 2026 16:15:48 -0700 Subject: [PATCH 14/23] jax/moe: drop bwd recv_topk_weights NaN sanitizer; trust the dispatch contract Mirrors the sparse_probs NaN-sanitizer removal in fe446974: we trust ep_dispatch_fwd's contract that recv_topk_weights does not contain NaN, and would rather see NaN propagate (catching a contract violation immediately) than silently sanitize it. The mask_bool dance itself stays: ctx.expert_outputs and grad_pre_combine still carry NaN at padded slots (ep_dispatch_fwd leaves uninit memory in recv_tokens, FFN and combine_bwd propagate it), and IEEE NaN * 0 = NaN means jnp.where is structurally needed to overwrite padded positions with literal zeros before the sum reduction. What changed: * Drop `recv_w_clean = jnp.where(jnp.isnan(...), 0, ...)` and thread ctx.recv_topk_weights directly into w / mask_bool. * Replace the NaN-defensive comment block with a shorter note that explains the structural reason the mask is still needed (NaN in expert_outputs / grad_pre_combine at padded slots), without claiming anything about recv_topk_weights. Addresses Greptile P1 by removing the asymmetry (fwd had no sanitizer, bwd did) -- chosen direction is "remove the bwd sanitizer", matching the project-wide stance of trusting kernel contracts rather than papering over violations. Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 09c9b5dd7c..bd876e47a2 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -862,22 +862,16 @@ def _moe_bwd_rule( d_expert_outputs = grad_pre_combine d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) else: - # ep_dispatch_fwd can land NaN into recv_topk_weights on padded - # slots. Untreated, `(NaN != 0) == True` in IEEE, - # so the multiplicative mask cannot suppress the NaN and it - # propagates through grad_pre_combine * w * mask into d_expert_outputs - # and then into every downstream gradient (gate_kernel ends up - # all-NaN). Sanitize once here. - recv_w_clean = jnp.where(jnp.isnan(ctx.recv_topk_weights), 0, ctx.recv_topk_weights) - w = recv_w_clean[..., None].astype(grad_pre_combine.dtype) - mask_bool = (recv_w_clean != 0)[..., None] + # Bwd mirror of the fwd mask: grad_pre_combine and ctx.expert_outputs + # both carry NaN at padded slots (ep_dispatch_fwd leaves uninit + # memory in recv_tokens, the FFN and combine_bwd propagate), and + # IEEE NaN * 0 = NaN, so jnp.where is needed to overwrite padded + # positions with literal zeros before the sum reduction. + w = ctx.recv_topk_weights[..., None].astype(grad_pre_combine.dtype) + mask_bool = (ctx.recv_topk_weights != 0)[..., None] d_expert_outputs = jnp.where( mask_bool, grad_pre_combine * w, jnp.zeros_like(grad_pre_combine) ) - # Same masking strategy for the cotangent on recv_topk_weights: - # grad_pre_combine has NaN at padded slots and ctx.expert_outputs - # may too, so the per-element product must be jnp.where'd before - # the sum reduction. d_recv_w_from_combine = jnp.where( mask_bool, grad_pre_combine * ctx.expert_outputs, From e6705116936ce15bcc25db4178e7f6f51c66f905 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 16 Jun 2026 16:34:23 -0700 Subject: [PATCH 15/23] jax/moe: assert output dtype; tests cover d_x parity (dtype + values) Two related dtype-contract changes: 1. moe.py: one-line assert at the moe() return path that output.dtype == x.dtype. Cheap structural guard against any future bug that lets the public output drift wider than the user-supplied input dtype. 2. test_te_ep_moe.py: extend test_backward to also check d_x, the gradient propagated back to the previous layer in backprop. _grad_step now uses jax.grad(loss_fn, argnums=(0, 1)) and returns (grads_variables, grad_x); the reference path does the same so we can compare. d_x is checked for: * shape == x.shape * dtype == x.dtype (protects the _with_sharding_constraint_cast_bwd wrapper that casts the fp32-promoted gate path back to the primal dtype on bwd; a regression in that wrapper would silently double activation gradient bandwidth) * finiteness + non-zero * numerical parity vs the pure-JAX reference d_x Addresses jberchtold review comment on test_te_ep_moe.py:650 ("we also need to check the final propagated gradient that will be passed onto the next layer in backprop"). test_combined_loss_grads is adjusted to ``grads, _`` unpacking; it doesn't need d_x for its main+aux finiteness check. Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 46 +++++++++++++++++++++++++++++------ transformer_engine/jax/moe.py | 1 + 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 6f27ba1d33..7e03683c5c 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -448,7 +448,12 @@ def _init_apply(block, mesh, x, key): def _grad_step(block, variables, mesh, x, *, include_aux=False): - """Run jax.grad of mean(out^2) [+ aux if include_aux] vs params.""" + """Run jax.grad of mean(out^2) [+ aux if include_aux] vs (params, x). + + Returns ``(grads_variables, grad_x)`` so callers can check both the + weight gradients and the input-activation gradient that propagates + back to the previous layer. + """ with _ctx(mesh): x_sh = _shard_inputs(x, mesh) @@ -459,9 +464,10 @@ def loss_fn(variables, x): loss = loss + aux.astype(jnp.float32) return loss - grads = jax.jit(jax.grad(loss_fn))(variables, x_sh) - jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) - return grads + grads_v, grad_x = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))(variables, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(grads_v)[0]) + jax.block_until_ready(grad_x) + return grads_v, grad_x def _grad_aux_only(block, variables, mesh, x): @@ -618,10 +624,11 @@ def test_backward(self, mesh, config): block = _make_block(**config) x = _make_inputs(jax.random.PRNGKey(2)) variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) - grads_te = _grad_step(block, variables, mesh, x) + grads_te, grad_x_te = _grad_step(block, variables, mesh, x) # Reference grads via jax.grad over the pure-JAX MoE with the - # same config. + # same config. argnums=(0, 1) so the reference also produces a + # d_x for the propagated-gradient parity check below. params_np = _params_global_numpy(variables, mesh) x_np = np.asarray(jax.device_get(x)) ref_kwargs = _reference_kwargs_from_config(config, params_np) @@ -641,11 +648,12 @@ def loss_fn(params, x): ) return jnp.mean(out.astype(jnp.float32) ** 2) - grads_ref = jax.jit(jax.grad(loss_fn))( + 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), ) grads_ref_np = {k: np.asarray(jax.device_get(v)) for k, v in grads_ref.items()} + grad_x_ref_np = np.asarray(jax.device_get(grad_x_ref)) for name in ("gate_kernel", "wi_0", "wi_1", "wo"): # Per-tensor: finite + non-zero + parity in one pass. @@ -665,6 +673,28 @@ def loss_fn(params, x): err_msg=f"grad parity breach on {name} [config={config}]", ) + # d_x: the gradient propagated back to the previous layer. Checks + # shape, dtype (must match x.dtype — protects the + # _with_sharding_constraint_cast_bwd wrapper that casts the + # fp32-promoted gate path back to bf16), finiteness, non-zero + # AND numerical parity vs the pure-JAX reference d_x. + grad_x_te_np = _to_global_numpy(grad_x_te, mesh) + assert grad_x_te.shape == x.shape, ( + f"d_x shape {grad_x_te.shape} != x.shape {x.shape} [config={config}]" + ) + assert grad_x_te.dtype == x.dtype, ( + f"d_x dtype {grad_x_te.dtype} != x.dtype {x.dtype} [config={config}]" + ) + assert np.all(np.isfinite(grad_x_te_np)), f"d_x has NaN/Inf [config={config}]" + assert np.any(grad_x_te_np != 0.0), f"d_x identically zero [config={config}]" + np.testing.assert_allclose( + grad_x_te_np.astype(np.float32), + grad_x_ref_np.astype(np.float32), + atol=GRAD_FFN_ATOL, + rtol=GRAD_FFN_RTOL, + err_msg=f"d_x parity breach [config={config}]", + ) + class TestTeEpMoeAuxLoss: """Aux-loss path. Consolidated into: @@ -725,7 +755,7 @@ def test_combined_loss_grads(self, mesh): 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) + grads, _ = _grad_step(block, variables, mesh, x, include_aux=True) for name in ("gate_kernel", "wi_0", "wi_1", "wo"): g_local = np.asarray(jax.device_get(_unwrap(grads["params"][name]).addressable_data(0))) assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf under main+aux" diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index bd876e47a2..df88667b9f 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -1281,4 +1281,5 @@ def moe( ) 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 From 518e1773fb2a2c94954415f2013cd203fe00bcca Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 16 Jun 2026 16:34:47 -0700 Subject: [PATCH 16/23] tests/jax/test_te_ep_moe: strip docstring to just "what this suite covers" Drops two paragraphs whose content was agent-flavoured PR-review notes rather than user-facing test docs: * The final "FP8 / MXFP8 deferred" paragraph that referenced an internal review artifact (``.pr3036-review/INTEGRATION_DESIGN.md``) not in the repo. * The "Intentional non-coverage" section that explained which tests deliberately do not exist (no Flax-wrapper smoke, no re-bootstrap-mismatch test) and why -- exactly the kind of defensive / forward-looking justification prose CLAUDE.md says to keep out of the codebase. The remaining docstring covers what readers actually need: how to launch the suite, what each test class exercises, and a short note on the parametrize-vs-class layout. Addresses jberchtold review comment on test_te_ep_moe.py:54. Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 7e03683c5c..8287eff151 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -39,21 +39,6 @@ * ``TestTeEpMoeAuxLoss`` covers the second return value end-to-end (returned + parity + aux-only grad propagates to gate + combined main+aux grads stay finite) in two consolidated tests. - -Intentional non-coverage: - -* No dedicated "Flax wrapper init+apply" smoke test: every config above - already calls ``MoEBlock`` (the Flax wrapper) end-to-end, so a - separate wrapper smoke would just duplicate ``test_forward[softmax]`` - + ``test_backward[softmax]``. -* No re-bootstrap-mismatch test: ``ep_bootstrap`` rejects a mismatched - signature unconditionally and is a one-line guard; covering it from - this suite would taint the per-process NCCL bootstrap cache for the - rest of the file with no real upside. - -FP8 / MXFP8 recipes are deferred — the ``quantizer_sets`` plumbing -has not yet been re-wired across the TE-EP ``shard_map`` boundary -(see ``.pr3036-review/INTEGRATION_DESIGN.md``). """ import os From bd052ceb629adee8b28311bf8d93f27b4cf43d63 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 6 Jul 2026 19:05:02 -0700 Subject: [PATCH 17/23] jax/moe: address TE EP alignment review feedback Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 31 +++++++++++++++--------------- transformer_engine/jax/flax/moe.py | 4 +++- transformer_engine/jax/moe.py | 4 ++-- 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 8287eff151..4e3cdbee8f 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -118,7 +118,7 @@ def _read_mp_options(): ) from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import moe, record_ep_bootstrap_signature_for_moe +from transformer_engine.jax.moe import _ALIGN_SIZE, moe, record_ep_bootstrap_signature_for_moe from transformer_engine.jax.ep import ep_bootstrap from transformer_engine.jax.sharding import MeshResource, global_shard_guard @@ -190,14 +190,15 @@ def _compute_worst_case_recv_pr(): flattened total or ``ncclEpDispatch`` aborts with ``invalid argument`` at ``ep_backend.cpp:414``. The moe block computes ``recv_pr`` the same way (see ``moe.py``'s - ``natural_spe = num_ep * max_tokens_per_rank``); keeping the - bootstrap formula in lock-step here. + ``natural_spe = num_ep * max_tokens_per_rank`` rounded up to + ``_ALIGN_SIZE``); keeping the bootstrap formula in lock-step here. """ num_procs = jax.device_count() num_local_experts = NUM_EXPERTS // EP_SIZE max_tokens_per_rank = (BATCH // num_procs) * SEQ natural_spe = EP_SIZE * max_tokens_per_rank - return num_local_experts * natural_spe + slots_per_expert = ((natural_spe + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + return num_local_experts * slots_per_expert @pytest.fixture(scope="module") @@ -361,7 +362,7 @@ def _make_block( aux_loss_coeff=0.0, use_expert_routing_bias=False, score_function="softmax", - bias_init=None, + expert_bias_init=None, ): kwargs = dict( num_experts=NUM_EXPERTS, @@ -374,10 +375,10 @@ def _make_block( score_function=score_function, dtype=DTYPE, ) - # Custom bias_init lets tests inject a non-zero expert_bias without + # Custom expert_bias_init lets tests inject a non-zero expert_bias without # poking variables['params'] post-init. - if bias_init is not None: - kwargs["bias_init"] = bias_init + if expert_bias_init is not None: + kwargs["expert_bias_init"] = expert_bias_init return MoEBlock(**kwargs) @@ -540,7 +541,7 @@ def _make_inputs(key): dict( score_function="sigmoid", use_expert_routing_bias=True, - bias_init=_strong_expert_bias_init, + expert_bias_init=_strong_expert_bias_init, ), id="sigmoid-bias-strong", ), @@ -664,12 +665,12 @@ def loss_fn(params, x): # fp32-promoted gate path back to bf16), finiteness, non-zero # AND numerical parity vs the pure-JAX reference d_x. grad_x_te_np = _to_global_numpy(grad_x_te, mesh) - assert grad_x_te.shape == x.shape, ( - f"d_x shape {grad_x_te.shape} != x.shape {x.shape} [config={config}]" - ) - assert grad_x_te.dtype == x.dtype, ( - f"d_x dtype {grad_x_te.dtype} != x.dtype {x.dtype} [config={config}]" - ) + assert ( + grad_x_te.shape == x.shape + ), f"d_x shape {grad_x_te.shape} != x.shape {x.shape} [config={config}]" + assert ( + grad_x_te.dtype == x.dtype + ), f"d_x dtype {grad_x_te.dtype} != x.dtype {x.dtype} [config={config}]" assert np.all(np.isfinite(grad_x_te_np)), f"d_x has NaN/Inf [config={config}]" assert np.any(grad_x_te_np != 0.0), f"d_x identically zero [config={config}]" np.testing.assert_allclose( diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 640db29534..3629346e33 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -240,11 +240,13 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: ) expert_bias = None if self.use_expert_routing_bias: + # The router logits are promoted to fp32 before fused top-k; keep + # the routing bias in the same dtype so it only affects selection. expert_bias = self.param( "expert_bias", nn.with_logical_partitioning(self.expert_bias_init, ("exp",)), (self.num_experts,), - self.dtype, + jnp.float32, ) ep_axis = get_active_resource_axis("ep_resource") diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index df88667b9f..daa69ab07e 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -634,7 +634,7 @@ def _moe_fwd_rule( # ---------------- TE EP dispatch (global view) ---------------- cfg = tex.EpLayerConfig( top_k=K, - dispatch_output_per_expert_alignment=slots_per_expert, + dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( @@ -1249,7 +1249,7 @@ def moe( if expert_bias is None: expert_bias_arg = jnp.zeros((0,), dtype=jnp.float32) else: - expert_bias_arg = expert_bias + expert_bias_arg = expert_bias.astype(jnp.float32) output, aux_loss = _moe( x, From c2869ede04fea8c58613c182f3b5c79da5c363c2 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 6 Jul 2026 20:17:03 -0700 Subject: [PATCH 18/23] jax/moe: fix early topk weighting padded-slot masking Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 12 ++++------ transformer_engine/jax/moe.py | 41 +++++++++++++++++++++-------------- 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 4e3cdbee8f..26bd070ce1 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -518,14 +518,10 @@ def _make_inputs(key): dict(score_function="softmax"), id="softmax", ), - # TODO: re-add the apply_topk_weights_early=True config once the - # 0*NaN -> NaN leak from padded recv slots in the early-weighting - # multiply (intermediate * recv_w * mask) is debugged. Late - # weighting (combine-side) is unaffected and stays covered above. - # Note: align_size is no longer a user-facing parameter; it is - # hard-coded to _ALIGN_SIZE = 128 in moe.py. Re-add a distinct - # align-size config only if the constant is loosened, or a - # recipe-driven inference is added that selects a >128 alignment. + pytest.param( + dict(score_function="softmax", apply_topk_weights_early=True), + id="softmax-early-weighting", + ), pytest.param( dict(score_function="sigmoid"), id="sigmoid", diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index daa69ab07e..934eb5f32e 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -288,14 +288,14 @@ def _ffn_fwd_per_shard( if apply_topk_weights_early: # 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, modulo elementwise op fusion gains. w_b is - # cast to intermediate.dtype so the multiply doesn't promote - # expert_outputs above the EP buffer's element width - # (ep_bootstrap rejects max_token_dtype != bf16, and the NCCL EP - # HT mega-buffer is sized for 2-byte slots accordingly). + # late-weighting path. Padded recv slots can contain uninitialized + # data, so overwrite inactive rows with literal zeros instead of + # relying on multiplication by a zero mask (IEEE NaN * 0 = NaN). + # ``w_b`` is cast to ``intermediate.dtype`` so the multiply doesn't + # promote expert_outputs above the EP buffer's element width. w_b = recv_w_flat[:, None].astype(intermediate.dtype) - mask_b = (recv_w_flat != 0).astype(intermediate.dtype)[:, None] - intermediate = intermediate * w_b * mask_b + active = (recv_w_flat != 0)[:, None] + intermediate = jnp.where(active, intermediate * w_b, jnp.zeros_like(intermediate)) casted_intermediate = tex.grouped_quantize( intermediate, q_set.x, local_group_sizes, flatten_axis=-1 @@ -377,26 +377,35 @@ def _ffn_bwd_per_shard( act_fn = _convert_to_activation_function(activation_type) if apply_topk_weights_early: # intermediate' = intermediate * w * mask. Split the cotangent - # across both factors before the activation bwd consumes it. - # Cast w_b so the multiply stays in d_intermediate.dtype and - # d_sorted_x (downstream into ep_dispatch_bwd) stays bf16. + # across both factors before the activation bwd consumes it. Padded + # recv slots may still be NaN in the saved activation residuals, so + # use zero-filled residuals on inactive rows before the activation VJP. w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) - mask_b = (recv_w_flat != 0).astype(d_intermediate.dtype)[:, None] - intermediate_unweighted = act_fn(gate_proj_out) * up_proj_out + active = (recv_w_flat != 0)[:, None] + gate_proj_for_bwd = jnp.where(active, gate_proj_out, jnp.zeros_like(gate_proj_out)) + up_proj_for_bwd = jnp.where(active, up_proj_out, jnp.zeros_like(up_proj_out)) + intermediate_unweighted = act_fn(gate_proj_for_bwd) * up_proj_for_bwd d_recv_w_from_intermediate = jnp.sum( - d_intermediate * intermediate_unweighted * mask_b, axis=-1 + jnp.where( + active, + d_intermediate * intermediate_unweighted, + jnp.zeros_like(d_intermediate), + ), + axis=-1, ).astype(recv_w_flat.dtype) - d_intermediate = d_intermediate * w_b * mask_b + d_intermediate = jnp.where(active, d_intermediate * w_b, jnp.zeros_like(d_intermediate)) 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_out) + 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_out) + (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 From 2dd0c6c60e86d76180189463976eee92d2308784 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 6 Jul 2026 21:40:01 -0700 Subject: [PATCH 19/23] jax/moe: remove unused EP mesh size Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 934eb5f32e..b1211c1bf2 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -843,7 +843,6 @@ def _moe_bwd_rule( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - num_ep = mesh.shape[ep_axis] dp_size = 1 for ax in data_parallelism_axes: dp_size *= mesh.shape[ax] From 9a13ef8512b5537fd2487934526d7acb4c65b12b Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 7 Jul 2026 17:15:42 -0700 Subject: [PATCH 20/23] jax/moe: tighten TE EP recv capacity bound Signed-off-by: Teddy Do --- tests/jax/test_te_ep_moe.py | 27 +++++++++++++++------------ transformer_engine/jax/moe.py | 31 ++++++++++++++----------------- 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 26bd070ce1..e627296b06 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -183,22 +183,25 @@ def _read_mp_options(): def _compute_worst_case_recv_pr(): """Per-rank recv buffer the bootstrap must reserve. - NCCL EP's HT path lays out the per-rank receive buffer as - ``[num_local_experts, ep_size * max_tokens_per_rank, hidden]`` - (per the LL combine assertion at ``nccl_ep.cc:2185`` and the - HT IPC buffer sizing at ``nccl_ep.cc:415``). We must mirror that - flattened total or ``ncclEpDispatch`` aborts with - ``invalid argument`` at ``ep_backend.cpp:414``. The moe block - computes ``recv_pr`` the same way (see ``moe.py``'s - ``natural_spe = num_ep * max_tokens_per_rank`` rounded up to - ``_ALIGN_SIZE``); keeping the bootstrap formula in lock-step here. + NCCL EP HT expert-major uses one flat recv buffer with variable + per-expert zones. Each non-empty expert zone is padded to + ``_ALIGN_SIZE`` slots, so the reserve must cover the worst-case + total assignments plus independent per-zone padding. """ num_procs = jax.device_count() num_local_experts = NUM_EXPERTS // EP_SIZE max_tokens_per_rank = (BATCH // num_procs) * SEQ - natural_spe = EP_SIZE * max_tokens_per_rank - slots_per_expert = ((natural_spe + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE - return num_local_experts * slots_per_expert + tokens_per_ep_group = EP_SIZE * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(TOPK, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts + aligned_total_bound = ( + (padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE + ) * _ALIGN_SIZE + per_expert_bound = num_local_experts * ( + (tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE + ) * _ALIGN_SIZE + return min(per_expert_bound, aligned_total_bound) @pytest.fixture(scope="module") diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index b1211c1bf2..5399fbb2ee 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -228,7 +228,6 @@ def _ffn_fwd_per_shard( wo_bias: Optional[jnp.ndarray], *, num_local_experts: int, - slots_per_expert: int, activation_type: str, apply_topk_weights_early: bool, ): @@ -249,7 +248,6 @@ def _ffn_fwd_per_shard( sorted_x = recv_tokens_local.reshape(-1, hidden) recv_w_flat = recv_topk_weights_local.reshape(-1) local_group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) - del slots_per_expert # not used since group_sizes is plumbed in dynamically wi_0 = wi_0.astype(sorted_x.dtype) wi_1 = wi_1.astype(sorted_x.dtype) @@ -514,20 +512,20 @@ def _moe_fwd_rule( # Per-rank send capacity: B/num_procs rows x S tokens per rank. max_tokens_per_rank = (B // num_procs) * S - # Per-rank receive capacity. NCCL EP HT lays out the per-rank receive - # buffer as ``[num_local_experts, num_ep * max_tokens_per_rank, hidden]`` - # (see nccl_ep.cc::init kernel buffer sizing + the LL combine assertion - # at nccl_ep.cc:2185 which spells out the same layout). The natural - # dropless K-expanded count - # ``ceil((B/dp)*S*K / num_local_experts)`` does NOT match: it ignores - # the worst-case where all of one EP group's tokens land on a single - # local expert. We must size to that worst case or NCCL EP's HT kernel - # rejects the dispatch buffer with ``invalid argument``. - natural_spe = num_ep * max_tokens_per_rank # = (B // dp_size) * S - # NCCL EP requires each expert-major output block to be at least - # ``_ALIGN_SIZE`` (=128) tokens; see the constant's docstring. - slots_per_expert = ((natural_spe + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE - recv_pr = num_local_experts * slots_per_expert + # Per-rank receive capacity. NCCL EP HT expert-major lays out variable + # per-expert zones in one flat recv buffer, with each non-empty zone padded + # to ``dispatch_output_per_expert_alignment``. + tokens_per_ep_group = num_ep * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(K, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts + aligned_total_bound = ( + (padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE + ) * _ALIGN_SIZE + per_expert_bound = num_local_experts * ( + (tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE + ) * _ALIGN_SIZE + recv_pr = min(per_expert_bound, aligned_total_bound) _te_ep_assert_compatible_bootstrap( num_experts=num_experts, @@ -723,7 +721,6 @@ def _body(*args): w1b, wob, num_local_experts=num_local_experts, - slots_per_expert=slots_per_expert, activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, ) From a49dc6bfeadb04398298a7ece354c66eb1c41240 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 7 Jul 2026 18:03:10 -0700 Subject: [PATCH 21/23] jax/moe: simplify late TE EP weighting Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 5399fbb2ee..9bb9873a49 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -746,16 +746,10 @@ def _body(*args): out_partition_spec=out_partition_spec, ) else: - # IEEE 754: NaN * 0 = NaN, so a multiplicative mask cannot kill - # the NaNs ep_dispatch_fwd leaves at padded slots of recv_tokens - # (they ride through the FFN into expert_outputs at the same - # padded positions): mean=NaN on expert_outputs[padded] then - # propagates into the combine output when the kernel's read - # pattern overlaps the padded region. Use jnp.where to overwrite - # padded positions with a literal 0 before combine. + # 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) - mask_bool = (recv_topk_weights != 0)[..., None] - weighted = jnp.where(mask_bool, expert_outputs * w, jnp.zeros_like(expert_outputs)) + weighted = expert_outputs * w output = tex.ep_combine_fwd( cfg, handle_mem, @@ -867,11 +861,9 @@ def _moe_bwd_rule( d_expert_outputs = grad_pre_combine d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) else: - # Bwd mirror of the fwd mask: grad_pre_combine and ctx.expert_outputs - # both carry NaN at padded slots (ep_dispatch_fwd leaves uninit - # memory in recv_tokens, the FFN and combine_bwd propagate), and - # IEEE NaN * 0 = NaN, so jnp.where is needed to overwrite padded - # positions with literal zeros before the sum reduction. + # Reverse the late-weighting multiply. Padded expert-major rows are + # part of the physical grouped-GEMM ranges, so write literal zero + # cotangents for inactive rows instead of relying on NaN * 0. w = ctx.recv_topk_weights[..., None].astype(grad_pre_combine.dtype) mask_bool = (ctx.recv_topk_weights != 0)[..., None] d_expert_outputs = jnp.where( From bd5d8ba23859cbe4aedd7e120d0d9fb9b0d78174 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 01:38:41 +0000 Subject: [PATCH 22/23] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_te_ep_moe.py | 10 ++++------ transformer_engine/jax/moe.py | 10 ++++------ 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index e627296b06..d08765e184 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -195,12 +195,10 @@ def _compute_worst_case_recv_pr(): max_local_assignments = tokens_per_ep_group * min(TOPK, num_local_experts) max_nonempty_experts = min(num_local_experts, max_local_assignments) padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts - aligned_total_bound = ( - (padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE - ) * _ALIGN_SIZE - per_expert_bound = num_local_experts * ( - (tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE - ) * _ALIGN_SIZE + aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + ) return min(per_expert_bound, aligned_total_bound) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 9bb9873a49..a721940d22 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -519,12 +519,10 @@ def _moe_fwd_rule( max_local_assignments = tokens_per_ep_group * min(K, num_local_experts) max_nonempty_experts = min(num_local_experts, max_local_assignments) padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts - aligned_total_bound = ( - (padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE - ) * _ALIGN_SIZE - per_expert_bound = num_local_experts * ( - (tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE - ) * _ALIGN_SIZE + aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + ) recv_pr = min(per_expert_bound, aligned_total_bound) _te_ep_assert_compatible_bootstrap( From e626a7f64fe2467e229b77c4f6fa0686fc01c9a1 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 7 Jul 2026 18:57:40 -0700 Subject: [PATCH 23/23] jax/moe: reduce padded-slot recv weight masking Signed-off-by: Teddy Do --- transformer_engine/jax/moe.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index a721940d22..887e005de6 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -384,11 +384,7 @@ def _ffn_bwd_per_shard( up_proj_for_bwd = jnp.where(active, up_proj_out, jnp.zeros_like(up_proj_out)) intermediate_unweighted = act_fn(gate_proj_for_bwd) * up_proj_for_bwd d_recv_w_from_intermediate = jnp.sum( - jnp.where( - active, - d_intermediate * intermediate_unweighted, - jnp.zeros_like(d_intermediate), - ), + d_intermediate * intermediate_unweighted, axis=-1, ).astype(recv_w_flat.dtype) d_intermediate = jnp.where(active, d_intermediate * w_b, jnp.zeros_like(d_intermediate)) @@ -867,11 +863,7 @@ def _moe_bwd_rule( d_expert_outputs = jnp.where( mask_bool, grad_pre_combine * w, jnp.zeros_like(grad_pre_combine) ) - d_recv_w_from_combine = jnp.where( - mask_bool, - grad_pre_combine * ctx.expert_outputs, - jnp.zeros_like(grad_pre_combine), - ).sum(axis=-1) + 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) ----------------