diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh index 32d5f21956..33352d3b3a 100755 --- a/tests/jax/run_te_ep_moe.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -14,7 +14,7 @@ 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" +TEST_FILE="${TEST_FILE:-$TE_ROOT/tests/jax/test_te_ep_moe.py}" PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 0e9cb6598d..f92519fe6b 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -31,11 +31,12 @@ on the block are pytest parametrize values rather than separate test classes: -* ``test_forward`` covers the forward across a curated set of - configurations (softmax/sigmoid scoring, optional non-zero - expert_bias). Each config asserts shape, dtype, finiteness and - numerical parity vs the reference in one run. -* ``test_backward`` mirrors that for gradients. +* ``test_forward`` covers BF16 and MXFP8 forward execution across a + curated set of configurations (softmax/sigmoid scoring, optional + non-zero expert_bias). Each config asserts shape, dtype, finiteness + and numerical parity vs the same BF16 reference in one run. +* ``test_backward`` mirrors that for gradients. BF16 and MXFP8 share + the full test body and differ only in the grouped-GEMM quantizer sets. * ``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. @@ -108,18 +109,21 @@ def _read_mp_options(): from transformer_engine_jax import get_device_compute_capability -# Grouped GEMM in the MoE custom_vjp requires Blackwell (sm_100+). The -# TE EP NCCL primitives themselves need SM>=90, but the FFN body uses -# grouped_gemm, so the file as a whole gates on sm_100+. -if get_device_compute_capability(0) < 100: +# TE EP NCCL primitives need SM>=90 +if get_device_compute_capability(0) < 90: pytest.skip( - "MoE TE EP tests require Blackwell (sm_100+) for grouped GEMM", + "MoE TE EP tests require Hopper (sm_90+) or newer for TE EP", allow_module_level=True, ) from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import _ALIGN_SIZE, 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.common.recipe import MXFP8BlockScaling from transformer_engine.jax.sharding import MeshResource, global_shard_guard @@ -148,7 +152,7 @@ def _read_mp_options(): DTYPE = jnp.bfloat16 BATCH = EP_SIZE * FSDP_SIZE * 2 # 8 on 4-GPU, 16 on 8-GPU SEQ = 32 -HIDDEN = 64 +HIDDEN = 128 INTER = 128 NUM_EXPERTS = 8 TOPK = 2 @@ -275,8 +279,7 @@ def mesh(): def _pure_jax_moe_reference( x, gate_kernel, - wi_0, - wi_1, + wi, wo, expert_bias=None, *, @@ -317,6 +320,7 @@ def _pure_jax_moe_reference( # 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. + wi_0, wi_1 = jnp.split(wi, 2, axis=-1) layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) # Activation runs in x.dtype (typically bf16) to mirror the impl -- @@ -365,6 +369,7 @@ def _make_block( score_function="softmax", expert_bias_init=None, input_axes=("batch", None, None), + quantization_recipe=None, ): kwargs = dict( num_experts=NUM_EXPERTS, @@ -377,6 +382,7 @@ def _make_block( score_function=score_function, dtype=DTYPE, input_axes=input_axes, + quantization_recipe=quantization_recipe, ) # Custom expert_bias_init lets tests inject a non-zero expert_bias without # poking variables['params'] post-init. @@ -436,7 +442,14 @@ def _init_apply(block, mesh, x, key): return variables, output, aux -def _grad_step(block, variables, mesh, x, *, include_aux=False): +def _grad_step( + block, + variables, + mesh, + x, + *, + include_aux=False, +): """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 @@ -503,6 +516,13 @@ def _make_inputs(key): return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) +def _quantization_recipe(quantization): + if quantization == "bf16": + return None + assert quantization == "mxfp8" + return MXFP8BlockScaling() + + # ----------------------------------------------------------------------------- # Tests # ----------------------------------------------------------------------------- @@ -546,6 +566,13 @@ def _make_inputs(key): ), ] +_QUANTIZATION_CASES = [ + pytest.param("bf16", id="bf16"), +] + +if get_device_compute_capability(0) >= 100: + _QUANTIZATION_CASES.append(pytest.param("mxfp8", id="mxfp8")) + def _reference_kwargs_from_config(config, params_np): """Pick out the reference-relevant pieces of a parametrize config.""" @@ -564,8 +591,9 @@ class TestTeEpMoeForward: finiteness AND numerical parity vs the pure-JAX reference.""" @pytest.mark.parametrize("config", _CONFIGS) - def test_forward(self, mesh, config): - block = _make_block(**config) + @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) + def test_forward(self, mesh, config, quantization): + block = _make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(0)) variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) @@ -584,8 +612,7 @@ def test_forward(self, mesh, config): 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["wi"]), jnp.asarray(params_np["wo"]), num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -596,7 +623,7 @@ def test_forward(self, mesh, config): 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}", + err_msg=f"forward parity breach for config={config}, quantization={quantization}", ) @@ -605,8 +632,9 @@ class TestTeEpMoeBackward: 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) + @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) + def test_backward(self, mesh, config, quantization): + block = _make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(2)) variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) grads_te, grad_x_te = _grad_step(block, variables, mesh, x) @@ -623,8 +651,7 @@ def loss_fn(params, x): out, _ = _pure_jax_moe_reference( x, params["gate_kernel"], - params["wi_0"], - params["wi_1"], + params["wi"], params["wo"], ref_expert_bias, num_experts=NUM_EXPERTS, @@ -640,7 +667,7 @@ def loss_fn(params, x): 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"): + for name in ("gate_kernel", "wi", "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}]" @@ -655,7 +682,9 @@ def loss_fn(params, x): grads_ref_np[name].astype(np.float32), atol=atol, rtol=rtol, - err_msg=f"grad parity breach on {name} [config={config}]", + err_msg=( + f"grad parity breach on {name} [config={config}, quantization={quantization}]" + ), ) # d_x: the gradient propagated back to the previous layer. Checks @@ -677,7 +706,7 @@ def loss_fn(params, x): grad_x_ref_np.astype(np.float32), atol=GRAD_FFN_ATOL, rtol=GRAD_FFN_RTOL, - err_msg=f"d_x parity breach [config={config}]", + err_msg=f"d_x parity breach [config={config}, quantization={quantization}]", ) @@ -710,8 +739,7 @@ def test_aux_loss(self, mesh): _, 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["wi"]), jnp.asarray(params_np["wo"]), num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -741,7 +769,7 @@ def test_combined_loss_grads(self, mesh): 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"): + for name in ("gate_kernel", "wi", "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" diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 7138cfcf40..2d56eb4769 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -1293,9 +1293,18 @@ def grouped_quantize( return quantizer.quantize(x, flatten_axis=flatten_axis, group_sizes=group_sizes) n_groups = group_sizes.size original_shape = x.shape - assert n_groups == len( - quantizer.quantizers - ), f"n_groups={n_groups} != n_quantizers = {len(quantizer.quantizers)}" + n_quantizers = len(quantizer.quantizers) + if quantizer.scaling_mode.is_mxfp8_scaling: + # Stateless MXFP8 quantizers may describe a global grouped operation + # while this primitive is traced inside shard_map on only the local + # groups. The recipe is identical for every group, and no per-group + # state is selected here, so the global descriptor only needs to cover + # the local operation. + assert ( + n_groups <= n_quantizers + ), f"local n_groups={n_groups} exceeds global n_quantizers={n_quantizers}" + else: + assert n_groups == n_quantizers, f"n_groups={n_groups} != n_quantizers={n_quantizers}" scale = jnp.ones((n_groups,), jnp.float32) if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index c73bad33be..460714098d 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -32,14 +32,18 @@ import jax.numpy as jnp from flax import linen as nn +from transformer_engine.common.recipe import Recipe # Re-exported so downstream users can ``from transformer_engine.jax.flax.moe # import P`` without a second jax.sharding import. -from jax.sharding import PartitionSpec as P # noqa: F401 # pylint: disable=unused-import +from jax.sharding import ( + PartitionSpec as P, +) # noqa: F401 # pylint: disable=unused-import from ..moe import moe +from ..quantize import QuantizerSet from ..router import ScoreFunction -from ..sharding import get_active_resource_axis +from ..sharding import _get_mesh, get_active_resource_axis from .module import TransformerEngineBase PRNGKey = Any @@ -104,6 +108,9 @@ 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``. + recv_capacity_per_rank : Optional[int] + Exact aligned receive capacity per EP rank. ``None`` reserves the + dropless worst case. The per-expert dispatch-slot alignment is fixed internally at 128 tokens (see ``moe._ALIGN_SIZE``) -- the value required by NCCL EP @@ -117,10 +124,10 @@ class _MoEBlock(TransformerEngineBase): Register per-expert FFN biases (``wi_0_bias``, ``wi_1_bias``, ``wo_bias``). - Quantization is currently configured via the standard TE autocast - context (``fp8_autocast``/``with_quantizer_set``) and threaded - through ``moe()`` internally; this wrapper does not expose a - per-call ``quantizer_sets`` knob yet. + quantization_recipe : Optional[Recipe] + Recipe used to construct the FC1 and FC2 grouped-GEMM quantizer + sets. ``None`` uses the recipe from the active TE autocast context, + or no-op quantizers when autocast is disabled. """ # Architecture @@ -149,6 +156,7 @@ class _MoEBlock(TransformerEngineBase): # MoE knobs forwarded to ``moe()`` apply_topk_weights_early: bool = False + recv_capacity_per_rank: Optional[int] = None # Dtypes / init / misc dtype: DType = jnp.float32 @@ -156,6 +164,7 @@ class _MoEBlock(TransformerEngineBase): bias_init: Initializer = nn.initializers.zeros expert_bias_init: Initializer = nn.initializers.zeros use_ffn_bias: bool = False + quantization_recipe: Optional[Recipe] = None def __post_init__(self): if self.kernel_init is None: @@ -176,7 +185,6 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: ---------- inputs : jnp.ndarray ``[batch, sequence, hidden]``. - Returns ------- output : jnp.ndarray @@ -203,16 +211,14 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: (hidden_size, self.num_experts), self.dtype, ) - wi_0 = self.param( - "wi_0", + # FC1 is stored as one gated-SwiGLU kernel. Keeping its two + # projections contiguous lets the functional MoE path quantize and + # all-gather one FP8 data buffer (and one scale buffer), rather than + # materializing a concatenate inside the custom-VJP. + wi = self.param( + "wi", nn.with_logical_partitioning(self.kernel_init, self.wi_kernel_axes), - (self.num_experts, hidden_size, self.intermediate_size), - self.dtype, - ) - wi_1 = self.param( - "wi_1", - nn.with_logical_partitioning(self.kernel_init, self.wi_kernel_axes), - (self.num_experts, hidden_size, self.intermediate_size), + (self.num_experts, hidden_size, 2 * self.intermediate_size), self.dtype, ) wo = self.param( @@ -253,12 +259,39 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: ) ep_axis = get_active_resource_axis("ep_resource") + mesh = _get_mesh() + data_parallel_size = 1 + for axis in self.data_parallelism_axes: + data_parallel_size *= mesh.shape[axis] + + def make_grouped_quantizer_set(postfix): + # Dispatched token groups span every data-parallel replica, + # whereas expert kernels have one group per global expert. + token_set = self.generate_quantizer_set( + f"{postfix}_token", + fp8_recipe=self.quantization_recipe, + n_groups=data_parallel_size * self.num_experts, + ) + expert_set = self.generate_quantizer_set( + f"{postfix}_expert", + fp8_recipe=self.quantization_recipe, + n_groups=self.num_experts, + ) + return QuantizerSet( + x=token_set.x, + kernel=expert_set.kernel, + dgrad=token_set.dgrad, + ) + + quantizer_sets = ( + make_grouped_quantizer_set("_fc1"), + make_grouped_quantizer_set("_fc2"), + ) return moe( inputs, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, @@ -274,6 +307,8 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: scaling_factor=self.scaling_factor, aux_loss_coeff=self.aux_loss_coeff, apply_topk_weights_early=self.apply_topk_weights_early, + quantizer_sets=quantizer_sets, + recv_capacity_per_rank=self.recv_capacity_per_rank, ep_axis=ep_axis, data_parallelism_axes=self.data_parallelism_axes, input_axes=self.input_axes, diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 21ef59ba5a..eae56872e8 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -20,24 +20,20 @@ ``((*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). +* The EP, grouped-quantize, and grouped-GEMM primitives operate at global + view. Their custom partitioning rules handle per-shard execution, + including EP placement and DP/FSDP gathers and reductions. + +FC1 and FC2 use independent quantizer sets. The sets are differentiable +``custom_vjp`` arguments and are returned by the backward rule so +stateful recipes follow the same update semantics as the other TE MLPs. +``aux_loss_coeff`` and ``expert_bias`` are also supported. """ +import math +import warnings from functools import partial from typing import Any, Optional, Tuple, Union -import warnings import flax.struct import jax @@ -46,6 +42,8 @@ from . import cpp_extensions as tex from .quantize import ( + GroupedQuantizer, + QuantizerSet, TensorUsage, noop_quantizer_set, with_sharding_constraint_by_logical_axes, @@ -54,7 +52,7 @@ from .router import ScoreFunction, _validate_score_function from .sharding import _get_mesh -__all__ = ["moe"] +__all__ = ["get_moe_recv_capacity_per_rank", "moe"] # Per-expert dispatch-slot alignment fed to ``tex.ep_prepare`` as @@ -65,6 +63,62 @@ _ALIGN_SIZE = 128 +def get_moe_recv_capacity_per_rank( + *, + num_experts: int, + num_experts_per_tok: int, + max_tokens_per_rank: int, + ep_size: int, + recv_capacity_factor: Optional[float] = None, + alignment: int = _ALIGN_SIZE, +) -> int: + """Return the aligned receive capacity for one EP rank. + + ``recv_capacity_factor=None`` reserves the dropless worst case. A finite + factor >= 1 scales the capacity needed by perfectly balanced routing and + is capped at the worst case. The balanced baseline includes the independent + per-local-expert alignment required by NCCL EP. + """ + if num_experts <= 0 or num_experts_per_tok <= 0 or max_tokens_per_rank <= 0: + raise ValueError( + "num_experts, num_experts_per_tok, and max_tokens_per_rank must be positive" + ) + if ep_size <= 0 or num_experts % ep_size != 0: + raise ValueError(f"num_experts={num_experts} must be divisible by ep_size={ep_size}") + if alignment <= 0: + raise ValueError(f"alignment must be positive, got {alignment}") + if recv_capacity_factor is not None: + recv_capacity_factor = float(recv_capacity_factor) + if not math.isfinite(recv_capacity_factor) or recv_capacity_factor < 1.0: + raise ValueError( + "recv_capacity_factor must be finite and >= 1.0, or None for worst-case capacity; " + f"got {recv_capacity_factor}" + ) + + num_local_experts = num_experts // ep_size + tokens_per_ep_group = ep_size * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(num_experts_per_tok, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (alignment - 1) * max_nonempty_experts + aligned_total_bound = ((padded_total_bound + alignment - 1) // alignment) * alignment + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + alignment - 1) // alignment) * alignment + ) + worst_case = min(per_expert_bound, aligned_total_bound) + if recv_capacity_factor is None: + return worst_case + + balanced_per_expert = ( + max_tokens_per_rank * num_experts_per_tok + num_local_experts - 1 + ) // num_local_experts + balanced_aligned = ( + num_local_experts * ((balanced_per_expert + alignment - 1) // alignment) * alignment + ) + requested = math.ceil(balanced_aligned * recv_capacity_factor) + requested = ((requested + alignment - 1) // alignment) * alignment + return min(requested, worst_case) + + def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: """Sharding constraint that keeps bwd cotangents in the primal dtype. @@ -110,9 +164,9 @@ def _constraint_bwd(dtype_ref, grad): # 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). +# per-call check below verifies the recorded signature matches the current +# MoE invocation. NCCL EP permits a smaller token count than the bootstrap +# maximum, but the dispatch receive capacity itself must match exactly. _te_ep_bootstrap_signature: Optional[Tuple[int, int, int, int, int]] = None @@ -145,7 +199,7 @@ def _te_ep_assert_compatible_bootstrap( hidden_dim: int, ep_size: int, ) -> None: - """Verify a prior eager ``ep_bootstrap`` is wide enough for this call.""" + """Verify a prior eager ``ep_bootstrap`` is compatible with this call.""" if _te_ep_bootstrap_signature is None: raise RuntimeError( "TE EP was not bootstrapped. Call" @@ -160,7 +214,7 @@ def _te_ep_assert_compatible_bootstrap( 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 + or recv_capacity_per_rank != b_recv_pr ): raise ValueError( "TE EP was already bootstrapped with signature" @@ -170,7 +224,7 @@ def _te_ep_assert_compatible_bootstrap( 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." + " sizes) is required. NCCL EP dispatch capacity must exactly match bootstrap." ) @@ -196,7 +250,6 @@ class _Ctx: routing_map: jnp.ndarray 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 casted_wi_rhs_trans: Any @@ -206,65 +259,120 @@ class _Ctx: casted_wo_rhs_trans: Any expert_outputs: jnp.ndarray local_group_sizes: jnp.ndarray + quantizer_sets: Any aux_const_buf: Any = None aux_tokens_per_expert: Any = None aux_saved_scores: Any = None # ============================================================================= -# Per-shard FFN body (runs inside shard_map) +# Per-shard FFN body # ============================================================================= +def _validate_moe_quantizer_sets( + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], + *, + num_token_groups: int, + num_expert_groups: int, +) -> None: + """Validate the current global-view MoE quantizer contract. + + Quantizers passed to the public MoE API always describe the global logical + operation. The shard-mapped FFN consumes only its local group count, but it + must not rewrite that public metadata into a shard-local representation. + + Stateful grouped recipes will eventually require sharded leading group + dimensions on their internal state. Until that representation exists, MoE + supports only no-op quantizers and stateless MXFP8 grouped quantizers. + """ + if not isinstance(quantizer_sets, tuple) or len(quantizer_sets) != 2: + raise TypeError("MoE quantizer_sets must be a tuple of FC1 and FC2 QuantizerSet objects.") + + expected_groups = { + "x": num_token_groups, + "kernel": num_expert_groups, + "dgrad": num_token_groups, + } + for set_name, quantizer_set in zip(("FC1", "FC2"), quantizer_sets): + if not isinstance(quantizer_set, QuantizerSet): + raise TypeError(f"MoE {set_name} quantizer must be a QuantizerSet.") + quantizers = { + "x": quantizer_set.x, + "kernel": quantizer_set.kernel, + "dgrad": quantizer_set.dgrad, + } + if all(quantizer is None for quantizer in quantizers.values()): + continue + if any(quantizer is None for quantizer in quantizers.values()): + raise TypeError( + f"MoE {set_name} must use either all no-op quantizers or all grouped MXFP8 " + "quantizers." + ) + + for source, quantizer in quantizers.items(): + if not isinstance(quantizer, GroupedQuantizer): + raise TypeError( + f"MoE {set_name} {source} quantizer must be a GroupedQuantizer; " + f"got {type(quantizer).__name__}." + ) + if not quantizer.scaling_mode.is_mxfp8_scaling: + raise NotImplementedError( + "TE MoE currently supports only BF16/no-op and stateless MXFP8 grouped " + f"quantizers; {set_name} {source} uses {quantizer.scaling_mode}." + ) + if jax.tree_util.tree_leaves(quantizer): + raise NotImplementedError( + "TE MoE does not yet support stateful grouped quantizers. Quantizer state " + "must first be represented with a sharded global group dimension." + ) + expected = expected_groups[source] + if quantizer.n_groups != expected or len(quantizer.quantizers) != expected: + raise ValueError( + f"MoE {set_name} {source} quantizer must describe the global logical " + f"group count {expected}; got n_groups={quantizer.n_groups} and " + f"{len(quantizer.quantizers)} child quantizers." + ) + + 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, + wi: jnp.ndarray, wo: jnp.ndarray, wi_0_bias: Optional[jnp.ndarray], wi_1_bias: Optional[jnp.ndarray], wo_bias: Optional[jnp.ndarray], + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, num_local_experts: int, activation_type: str, apply_topk_weights_early: bool, ): - """Per-shard FFN forward. - - 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. - - ``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. - """ + """Run the grouped FFN on one shard's EP receive buffer.""" 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 = token_counts_local.reshape(-1).astype(jnp.int32) + group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) - wi_0 = wi_0.astype(sorted_x.dtype) - wi_1 = wi_1.astype(sorted_x.dtype) + wi = wi.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) - # 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`` is stored in its gated-SwiGLU layout [expert, hidden, 2*mlp]. + # Keeping it contiguous lets grouped quantize/GEMM consume it directly. wi_combined_bias = ( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) - 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) + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets + casted_sorted_x = tex.grouped_quantize( + sorted_x, + fc1_quantizer_set.x, + group_sizes, + flatten_axis=-1, + ) + casted_wi = tex.grouped_quantize(wi, fc1_quantizer_set.kernel, flatten_axis=-1) combined_out = tex.grouped_gemm( casted_sorted_x.get_tensor(usage=TensorUsage.LHS), casted_wi.get_tensor(usage=TensorUsage.RHS), @@ -272,8 +380,6 @@ def _ffn_fwd_per_shard( bias=wi_combined_bias, ) 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) # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM # output dtype; the activation output (`intermediate`) stays in the @@ -282,45 +388,35 @@ def _ffn_fwd_per_shard( # transitions to the target precision. act_fn = _convert_to_activation_function(activation_type) intermediate = act_fn(gate_proj_out) * up_proj_out - if apply_topk_weights_early: # 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. 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) - active = (recv_w_flat != 0)[:, None] - intermediate = jnp.where(active, intermediate * w_b, jnp.zeros_like(intermediate)) + # late-weighting path. Grouped GEMM skips padding automatically. + intermediate = intermediate * recv_w_flat[:, None].astype(intermediate.dtype) casted_intermediate = tex.grouped_quantize( - intermediate, q_set.x, local_group_sizes, flatten_axis=-1 + intermediate, + fc2_quantizer_set.x, + group_sizes, + flatten_axis=-1, ) - casted_wo = tex.grouped_quantize(wo, q_set.kernel, flatten_axis=-1) + casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), casted_wo.get_tensor(usage=TensorUsage.RHS), contracting_dims=((1,), (1,)), bias=wo_bias, ) - casted_intermediate_lhs_trans = casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS) - 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) + group_sizes_2d = group_sizes.reshape(1, num_local_experts) residuals = ( - casted_sorted_x_lhs_trans, - casted_wi_rhs_trans, + casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(fc1_quantizer_set.x), + casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(fc1_quantizer_set.kernel), gate_proj_out, up_proj_out, - casted_intermediate_lhs_trans, - casted_wo_rhs_trans, - local_group_sizes_3d, + casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(fc2_quantizer_set.x), + casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(fc2_quantizer_set.kernel), + group_sizes_2d, ) return expert_outputs_3d, residuals @@ -335,28 +431,25 @@ def _ffn_bwd_per_shard( casted_wo_rhs_trans, local_group_sizes: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, activation_type: str, apply_topk_weights_early: bool, has_bias: bool, ): - """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 = local_group_sizes.reshape(-1).astype(jnp.int32) + """Backward mirror of :func:`_ffn_fwd_per_shard`.""" + 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 - # 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] + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets # wo bwd - casted_d_eo = tex.grouped_quantize(d_eo_2d, q_set.dgrad, local_group_sizes, flatten_axis=-1) + casted_d_eo = tex.grouped_quantize( + d_eo_2d, + fc2_quantizer_set.dgrad, + 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( @@ -369,25 +462,21 @@ 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 + d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None act_fn = _convert_to_activation_function(activation_type) if apply_topk_weights_early: - # intermediate' = intermediate * w * mask. Split the cotangent - # 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. + # intermediate' = intermediate * w. The subsequent dgrad and EP + # operations consume group_sizes, so they skip padded ragged rows. w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) - 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 + gate_proj_for_bwd = gate_proj_out + up_proj_for_bwd = up_proj_out + intermediate_unweighted = act_fn(gate_proj_out) * up_proj_out d_recv_w_from_intermediate = jnp.sum( d_intermediate * intermediate_unweighted, axis=-1, ).astype(recv_w_flat.dtype) - d_intermediate = jnp.where(active, d_intermediate * w_b, jnp.zeros_like(d_intermediate)) + d_intermediate = d_intermediate * w_b else: gate_proj_for_bwd = gate_proj_out up_proj_for_bwd = up_proj_out @@ -405,10 +494,13 @@ def _ffn_bwd_per_shard( # 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. + # wgrad result remains in the contiguous gated-SwiGLU ``wi`` layout. d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) casted_d_combined = tex.grouped_quantize( - d_combined, q_set.dgrad, local_group_sizes, flatten_axis=-1 + d_combined, + fc1_quantizer_set.dgrad, + group_sizes, + flatten_axis=-1, ) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), @@ -420,10 +512,8 @@ 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) + d_wi_combined_bias = tex.grouped_dbias(d_combined, 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 @@ -434,8 +524,7 @@ def _ffn_bwd_per_shard( return ( d_sorted_x_3d, d_recv_w_3d, - d_wi_0, - d_wi_1, + d_wi_combined, d_wo, d_wi_0_bias, d_wi_1_bias, @@ -451,13 +540,13 @@ def _ffn_bwd_per_shard( def _moe_fwd_rule( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -475,8 +564,9 @@ def _moe_fwd_rule( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ): - """Forward: gate -> topk -> ep_dispatch -> shard_map(FFN) -> ep_combine. + """Forward: gate -> topk -> ep_dispatch -> FFN -> ep_combine. Returns ``(output, aux_loss)``. ``aux_loss`` is a zero scalar when ``aux_loss_coeff == 0``. @@ -500,6 +590,11 @@ def _moe_fwd_rule( for ax in data_parallelism_axes: dp_size *= mesh.shape[ax] num_procs = num_ep * dp_size + _validate_moe_quantizer_sets( + quantizer_sets, + num_token_groups=dp_size * num_experts, + num_expert_groups=num_experts, + ) B, S, H = x.shape K = num_experts_per_tok @@ -508,18 +603,21 @@ 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 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 + worst_case_recv_pr = get_moe_recv_capacity_per_rank( + num_experts=num_experts, + num_experts_per_tok=K, + max_tokens_per_rank=max_tokens_per_rank, + ep_size=num_ep, ) - recv_pr = min(per_expert_bound, aligned_total_bound) + if recv_capacity_per_rank is None: + recv_pr = worst_case_recv_pr + else: + recv_pr = int(recv_capacity_per_rank) + if recv_pr <= 0 or recv_pr % _ALIGN_SIZE != 0: + raise ValueError( + f"recv_capacity_per_rank must be a positive multiple of {_ALIGN_SIZE}, got" + f" {recv_pr}" + ) _te_ep_assert_compatible_bootstrap( num_experts=num_experts, @@ -638,6 +736,7 @@ def _moe_fwd_rule( dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + token_counts = jax.lax.with_sharding_constraint(token_counts, NamedSharding(mesh, ep2_spec)) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr ) @@ -649,81 +748,57 @@ def _moe_fwd_rule( # ---------------- FFN (per-shard via shard_map) ---------------- has_bias = wi_0_bias is not None kernel_spec = P(ep_axis, None, None) - bias_spec = P(ep_axis, None) if has_bias else None - # 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] + bias_spec = P(ep_axis, None) + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi, wo] if has_bias: - ffn_in_specs = ffn_in_specs + (bias_spec, bias_spec, bias_spec) + 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). local_group_sizes is - # now per-shard dynamic (= per-shard token_counts), so its - # residual spec mirrors ep2_spec (one row per ep rank). + # Quantized grouped tensors store their data, scales, and group metadata + # as physical buffers rather than in the source tensor's logical shape. + # A PartitionSpec used as a pytree prefix applies the same ownership to + # every array leaf of the grouped tensor: dispatched-token buffers belong + # to the compound batch shard, while expert-weight buffers belong to EP. + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) residuals_spec = ( - 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 + token_buffer_spec, + expert_buffer_spec, + token_matrix_spec, + token_matrix_spec, + token_buffer_spec, + expert_buffer_spec, + ep2_spec, ) - out_specs = (ep3_spec, residuals_spec) - def _body(*args): + def _ffn_fwd_body(*args): if has_bias: - (r_tok, r_w, tc, w0, w1, w_o, w0b, w1b, wob) = args + r_tok, r_w, tc, local_wi, local_wo, w0b, w1b, wob = args else: - (r_tok, r_w, tc, w0, w1, w_o) = args + r_tok, r_w, tc, local_wi, local_wo = args w0b = w1b = wob = None - # 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, + local_wi, + local_wo, w0b, w1b, wob, + quantizer_sets, num_local_experts=num_local_experts, activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, ) expert_outputs, ffn_residuals = shard_map( - _body, + _ffn_fwd_body, mesh=mesh, in_specs=ffn_in_specs, - out_specs=out_specs, + out_specs=(ep3_spec, residuals_spec), check_rep=False, )(*ffn_in_args) expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) @@ -773,7 +848,6 @@ def _body(*args): 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, @@ -783,6 +857,7 @@ def _body(*args): casted_wo_rhs_trans=casted_wo_rhs_trans, expert_outputs=expert_outputs, local_group_sizes=local_group_sizes, + quantizer_sets=quantizer_sets, aux_const_buf=aux_const_buf, aux_tokens_per_expert=aux_tokens_per_expert, aux_saved_scores=aux_saved_scores, @@ -814,11 +889,12 @@ def _moe_bwd_rule( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, residuals, cotangents, ): """Backward mirror of :func:`_moe_fwd_rule`.""" - del num_groups, group_topk, dtype # captured in residuals / unused in bwd + del num_groups, group_topk, dtype, recv_capacity_per_rank # captured / unused in bwd from jax.experimental.shard_map import shard_map # total_recv_tokens is a non-differentiable output; its cotangent is unused. @@ -832,10 +908,6 @@ def _moe_bwd_rule( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - 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: @@ -852,39 +924,35 @@ def _moe_bwd_rule( 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: - # 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. + # Reverse the late-weighting multiply. The subsequent dgrad and EP + # operations consume group_sizes, so they skip padded ragged rows. 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) - ) + d_expert_outputs = grad_pre_combine * w d_recv_w_from_combine = (grad_pre_combine * ctx.expert_outputs).sum(axis=-1) d_recv_w_from_combine = d_recv_w_from_combine.astype(ctx.recv_topk_weights.dtype) # ---------------- FFN bwd (per-shard via shard_map) ---------------- kernel_spec = P(ep_axis, None, None) - bias_spec = P(ep_axis, None) 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 - ep2_spec, # local_group_sizes (1, num_local_experts) per shard - ep2_spec, # recv_topk_weights + bias_spec = P(ep_axis, None) + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) + residuals_specs = ( + token_buffer_spec, + expert_buffer_spec, + token_matrix_spec, + token_matrix_spec, + token_buffer_spec, + expert_buffer_spec, + ep2_spec, ) + bwd_in_specs = (ep3_spec, *residuals_specs, ep2_spec) bwd_in_args = [ d_expert_outputs, ctx.casted_sorted_x_lhs_trans, @@ -896,67 +964,65 @@ def _moe_bwd_rule( 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(*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( + def _ffn_bwd_body(*args): + grads = _ffn_bwd_per_shard( *args, + ctx.quantizer_sets, activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, has_bias=has_bias, ) - # 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). + ( + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, + ) = grads if data_parallelism_axes: - dp = 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) + dp_axes = tuple(data_parallelism_axes) + d_wi_local = jax.lax.psum(d_wi_local, axis_name=dp_axes) + d_wo_local = jax.lax.psum(d_wo_local, axis_name=dp_axes) if has_bias: - d_wi_0_bias = 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) + d_wi_0_bias_local = jax.lax.psum(d_wi_0_bias_local, axis_name=dp_axes) + d_wi_1_bias_local = jax.lax.psum(d_wi_1_bias_local, axis_name=dp_axes) + d_wo_bias_local = jax.lax.psum(d_wo_bias_local, axis_name=dp_axes) return ( - d_sorted_x_3d, - d_recv_w_3d, - d_wi_0, - d_wi_1, - d_wo, - d_wi_0_bias, - d_wi_1_bias, - d_wo_bias, + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, ) + if has_bias: + bwd_out_specs = ( + ep3_spec, + ep2_spec, + kernel_spec, + kernel_spec, + bias_spec, + bias_spec, + bias_spec, + ) + else: + bwd_out_specs = (ep3_spec, ep2_spec, kernel_spec, kernel_spec, None, None, None) + ( d_sorted_x, d_recv_w_from_intermediate, - d_wi_0, - d_wi_1, + d_wi, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias, ) = shard_map( - _bwd_body, + _ffn_bwd_body, mesh=mesh, in_specs=bwd_in_specs, out_specs=bwd_out_specs, @@ -1040,9 +1106,14 @@ def _bwd_body(*args): # 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_wi = with_sharding_constraint_by_logical_axes(d_wi, wi_kernel_axes) d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_axes) + if has_bias: + wi_bias_axes = (wi_kernel_axes[0], *wi_kernel_axes[2:]) + wo_bias_axes = (wo_kernel_axes[0], *wo_kernel_axes[2:]) + d_wi_0_bias = with_sharding_constraint_by_logical_axes(d_wi_0_bias, wi_bias_axes) + d_wi_1_bias = with_sharding_constraint_by_logical_axes(d_wi_1_bias, wi_bias_axes) + d_wo_bias = with_sharding_constraint_by_logical_axes(d_wo_bias, wo_bias_axes) # expert_bias has no learnable bwd path through fused_topk: the # primitive's bwd returns None for the bias slot. Match that with a @@ -1053,13 +1124,13 @@ def _bwd_body(*args): return ( d_x, d_gate_kernel, - d_wi_0, - d_wi_1, + d_wi, d_wo, d_wi_0_bias if has_bias else None, d_wi_1_bias if has_bias else None, d_wo_bias if has_bias else None, d_expert_bias, + ctx.quantizer_sets, ) @@ -1068,17 +1139,17 @@ def _bwd_body(*args): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 26))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 27))) def _moe( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1096,17 +1167,18 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ): primal, _ = _moe_fwd_rule( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1124,6 +1196,7 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ) return primal @@ -1134,8 +1207,7 @@ def _moe( def moe( x: jnp.ndarray, gate_kernel: jnp.ndarray, - wi_0: jnp.ndarray, - wi_1: jnp.ndarray, + wi: jnp.ndarray, wo: jnp.ndarray, wi_0_bias: Optional[jnp.ndarray] = None, wi_1_bias: Optional[jnp.ndarray] = None, @@ -1152,6 +1224,10 @@ def moe( scaling_factor: float = 1.0, aux_loss_coeff: float = 0.0, apply_topk_weights_early: bool = False, + quantizer_sets: Tuple[QuantizerSet, QuantizerSet] = ( + noop_quantizer_set, + noop_quantizer_set, + ), ep_axis: str, data_parallelism_axes: Tuple[str, ...] = (), input_axes: Tuple[Optional[str], ...] = (), @@ -1159,6 +1235,7 @@ def moe( wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed"), dtype: jnp.dtype = jnp.float32, + recv_capacity_per_rank: Optional[int] = None, ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: """Run a full MoE block under a single fused custom_vjp on the TE EP path. @@ -1181,6 +1258,18 @@ 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. + quantizer_sets : Tuple[QuantizerSet, QuantizerSet] + Independent FC1 and FC2 quantizer sets describing the global logical + operation. Token quantizers have ``dp_size * num_experts`` groups and + kernel quantizers have ``num_experts`` groups; shard-local FFN calls use + this global descriptor unchanged. Currently only no-op (BF16) and + stateless grouped MXFP8 quantizers are supported. They are differentiable + custom-VJP arguments so recipe state is threaded through backward. + recv_capacity_per_rank : Optional[int] + Exact aligned receive-buffer capacity for each EP rank. ``None`` + (default) reserves the dropless aligned worst case. The value must match + the capacity used by ``ep_bootstrap``. Overflow is reported through + ``total_recv_tokens`` when bootstrap used ``drop_on_overflow=True``. Note that the per-expert dispatch-slot alignment is fixed internally at 128 tokens (``_ALIGN_SIZE``); see that constant's docstring for @@ -1191,7 +1280,7 @@ def moe( * ``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 + ``P((dp..., ep), None, None)`` for the physical ``jax.lax.with_sharding_constraint`` calls that JAX requires to refer to real mesh axes). * ``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, @@ -1247,13 +1336,13 @@ def moe( output, aux_loss, total_recv_tokens = _moe( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias_arg, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1271,6 +1360,7 @@ def moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ) if aux_loss_coeff <= 0.0: aux_loss = None diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index c5ad0451fd..edcec01924 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -8,9 +8,10 @@ both single-scale (1x) and double-scale (2x) quantization schemes. It supports rowwise and colwise quantization modes with proper scaling and dequantization. """ +import math +from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Callable, Optional, Tuple -from abc import ABC, abstractmethod import jax.numpy as jnp from jax.tree_util import register_pytree_node_class @@ -436,6 +437,15 @@ def __post_init__(self): 0 < self.flatten_axis < data_ndim ), f"flatten_axis {self.flatten_axis} is out of bounds for data.ndim = {data_ndim}" + # A grouped tensor used as a shard_map residual temporarily has global + # physical data/scale buffers while ``original_shape`` intentionally + # continues to describe the shard-local logical tensor. The matching + # shard_map input restores local leaves before grouped GEMM consumes + # the wrapper. Validate scale layout only when this is a genuine local + # tensor view; the global transport view is not itself a GEMM operand. + if self.data.size != math.prod(self.original_shape): + return + active_dims = ( self.first_dims if self.first_dims is not None and self.first_dims.size > 0