From b14df81d33f3522bfecdaa9ad8570430642d0764 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 26 May 2026 16:12:16 -0700 Subject: [PATCH 01/44] Add grouped GMM custom partitioning rules --- tests/jax/test_grouped_gemm_partitioning.py | 248 +++++++++++ ..._multi_process_distributed_grouped_gemm.py | 146 +++++-- transformer_engine/jax/cpp_extensions/gemm.py | 392 ++++++++++++++++++ .../jax/cpp_extensions/quantization.py | 192 +++++++++ transformer_engine/jax/dense.py | 56 ++- transformer_engine/jax/flax/module.py | 2 +- transformer_engine/jax/sharding.py | 2 + 7 files changed, 989 insertions(+), 49 deletions(-) create mode 100644 tests/jax/test_grouped_gemm_partitioning.py diff --git a/tests/jax/test_grouped_gemm_partitioning.py b/tests/jax/test_grouped_gemm_partitioning.py new file mode 100644 index 0000000000..5fb93e30ae --- /dev/null +++ b/tests/jax/test_grouped_gemm_partitioning.py @@ -0,0 +1,248 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Partitioning tests for grouped quantize and grouped GEMM.""" + +from types import SimpleNamespace + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, NamedSharding, PartitionSpec + +from transformer_engine.jax.cpp_extensions.gemm import GroupedGemmPrimitive +from transformer_engine.jax.cpp_extensions.quantization import GroupedQuantizePrimitive +from transformer_engine.jax.dense import grouped_dense +from transformer_engine.jax.quantize import QuantizeLayout, QuantizerFactory, ScalingMode +from transformer_engine.jax.sharding import MeshResource, global_shard_guard + + +def _mesh(): + return Mesh(np.asarray(jax.devices()[:1]).reshape(1, 1), ("expert", "fsdp")) + + +def _arg_info(mesh, shape, spec): + return SimpleNamespace( + shape=shape, + ndim=len(shape), + size=int(np.prod(shape)), + sharding=NamedSharding(mesh, PartitionSpec(*spec)), + ) + + +def _normalize_spec(spec): + if isinstance(spec, PartitionSpec): + return tuple(spec) + return spec + + +def _mxfp8_grouped_quantizer_set(n_groups): + return QuantizerFactory.create_set( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + fwd_dtype=jnp.float8_e4m3fn, + bwd_dtype=jnp.float8_e4m3fn, + is_2x2x=True, + n_groups=n_groups, + ) + + +def test_grouped_quantize_specs_preserve_ep_and_fsdp_for_block_scales(): + mesh = _mesh() + with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): + _, _, out_shardings, _ = GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 64), ("expert", None, "fsdp")), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (8,), ("expert",)), + ), + (), + ) + + specs = tuple(tuple(sharding.spec) for sharding in out_shardings) + assert _normalize_spec(specs[0]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[2]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[4]) == ("expert",) + + +def test_grouped_quantize_mxfp8_colwise_specs_preserve_ep_and_fsdp(): + mesh = _mesh() + with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): + _, _, out_shardings, _ = GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE_COLWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 128), ("expert", None, "fsdp")), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (8,), ("expert",)), + ), + (), + ) + + specs = tuple(tuple(sharding.spec) for sharding in out_shardings) + assert _normalize_spec(specs[0]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[1]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[2]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[3]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[4]) == ("expert",) + + +def test_grouped_gemm_rhs_weight_specs_gather_fsdp_but_preserve_ep(): + mesh = _mesh() + arg_infos = ( + _arg_info(mesh, (8192,), (None,)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (65536,), (("expert", "fsdp"),)), + _arg_info(mesh, (2048,), (("expert", "fsdp"),)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (1,), (None,)), + _arg_info(mesh, (0,), (None,)), + ) + with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): + _, _, out_sharding, arg_shardings = GroupedGemmPrimitive.partition( + False, + False, + ScalingMode.NO_SCALING.value, + jnp.bfloat16, + False, + False, + False, + 1, + 1, + (1, 128, 64), + 128, + 64, + 128, + 64, + mesh, + arg_infos, + (), + ) + + assert tuple(arg_shardings[2].spec) == ("expert",) + assert tuple(arg_shardings[3].spec) == ("expert",) + assert tuple(out_sharding[0].spec) == (None, None, None) + + +def test_grouped_partitioning_shardy_rules_smoke(): + mesh = _mesh() + quantize_rule = GroupedQuantizePrimitive.shardy_sharding_rule( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + SimpleNamespace(shape=(8, 128, 64)), + SimpleNamespace(shape=(8,)), + SimpleNamespace(shape=(8,)), + ), + ( + SimpleNamespace(shape=(8 * 128 * 64,)), + SimpleNamespace(shape=(1,)), + SimpleNamespace(shape=(8 * 128 * 64,)), + SimpleNamespace(shape=(1,)), + SimpleNamespace(shape=(8,)), + ), + ) + gemm_rule = GroupedGemmPrimitive.shardy_sharding_rule( + False, + False, + ScalingMode.NO_SCALING.value, + jnp.bfloat16, + False, + False, + False, + 1, + 2, + (128, 64), + 128, + 64, + 128, + 64, + mesh, + tuple(SimpleNamespace(shape=(1,)) for _ in range(13)), + (SimpleNamespace(shape=(128, 64)),), + ) + + assert quantize_rule is not None + assert gemm_rule is not None + + +def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process(): + mesh = _mesh() + n_groups = 2 + group_tokens = 128 + hidden = 128 + out_hidden = 128 + x_shape = (n_groups * group_tokens, hidden) + w_shape = (n_groups, hidden, out_hidden) + + x_sharding = NamedSharding(mesh, PartitionSpec("expert", None)) + w_sharding = NamedSharding(mesh, PartitionSpec("expert", "fsdp", None)) + group_sharding = NamedSharding(mesh, PartitionSpec("expert")) + out_sharding = NamedSharding(mesh, PartitionSpec("expert", None)) + + quantizer_set = _mxfp8_grouped_quantizer_set(n_groups) + + with mesh, global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): + x = jax.device_put( + jax.random.normal(jax.random.PRNGKey(20), x_shape, dtype=jnp.bfloat16) + * jnp.asarray(0.01, dtype=jnp.bfloat16), + x_sharding, + ) + w = jax.device_put( + jax.random.normal(jax.random.PRNGKey(21), w_shape, dtype=jnp.bfloat16) + * jnp.asarray(0.01, dtype=jnp.bfloat16), + w_sharding, + ) + group_sizes = jax.device_put( + jnp.full((n_groups,), group_tokens, dtype=jnp.int32), + group_sharding, + ) + + def apply_with_vjp(x, w, group_sizes): + def apply(x, w): + return grouped_dense( + x, + w, + group_sizes, + contracting_dims=((1,), (1,)), + quantizer_set=quantizer_set, + kernel_fsdp_info=("fsdp", 1), + ) + + out, vjp_fn = jax.vjp(apply, x, w) + dx, dw = vjp_fn(out) + return out, dx, dw + + out, dx, dw = jax.jit( + apply_with_vjp, + in_shardings=(x_sharding, w_sharding, group_sharding), + out_shardings=(out_sharding, x_sharding, w_sharding), + )(x, w, group_sizes) + out, dx, dw = jax.block_until_ready((out, dx, dw)) + + assert tuple(out.sharding.spec) == ("expert", None) + assert tuple(dx.sharding.spec) == ("expert", None) + assert tuple(dw.sharding.spec) == ("expert", "fsdp", None) + for value in (out, dx, dw): + local_value = np.asarray(jax.device_get(value.addressable_data(0))) + assert np.all(np.isfinite(local_value)) + assert np.any(local_value != 0.0) diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py index 94fed0859f..30a1452a07 100644 --- a/tests/jax/test_multi_process_distributed_grouped_gemm.py +++ b/tests/jax/test_multi_process_distributed_grouped_gemm.py @@ -7,18 +7,34 @@ import jax import jax.numpy as jnp import jax.experimental.multihost_utils as jem +import numpy as np +from jax.experimental import shard_map +from jax.sharding import NamedSharding, PartitionSpec from transformer_engine.jax.dense import grouped_dense as te_grouped_dense from transformer_engine.jax.quantize import ( QuantizerFactory, ScalingMode, ) +from transformer_engine.jax.sharding import MeshResource, global_shard_guard from utils import assert_allclose, dtype_tols N_GROUP = 8 -MESH_AXIS_NAME = "fsdp" +EP_AXIS_NAME = "ep" +FSDP_AXIS_NAME = "fsdp" +MESH_AXIS_NAME = FSDP_AXIS_NAME + + +def _mxfp8_grouped_quantizer_set(n_groups): + return QuantizerFactory.create_set( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + fwd_dtype=jnp.float8_e4m3fn, + bwd_dtype=jnp.float8_e4m3fn, + is_2x2x=True, + n_groups=n_groups, + ) def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis): @@ -36,13 +52,18 @@ def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis): def init_data(): x_key = jax.random.PRNGKey(0) w_key = jax.random.PRNGKey(1) - x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16) - w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16) - w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim)) - return x, w, w, w_amax + x = ( + jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16) + * jnp.asarray(0.01, dtype=jnp.bfloat16) + ) + w = ( + jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16) + * jnp.asarray(0.01, dtype=jnp.bfloat16) + ) + return x, w, w - def test_func(outter_x, outter_w, outter_w_amax): - in_specs = (x_sharding.spec, w_sharding.spec, None) + def test_func(outter_x, outter_w): + in_specs = (x_sharding.spec, w_sharding.spec) out_specs = x_sharding.spec @partial( @@ -52,36 +73,29 @@ def test_func(outter_x, outter_w, outter_w_amax): out_specs=out_specs, check_rep=False, ) - def sharded_group_gemm(x, w, w_amax): + def sharded_group_gemm(x, w): group_size = x.shape[0] x_reshaped = x.reshape(-1, x.shape[-1]) n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) - quantizer_set = QuantizerFactory.create_set( - scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING, - fwd_dtype=jnp.float8_e4m3fn, - bwd_dtype=jnp.float8_e5m2, - is_2x2x=True, - n_groups=group_size, - ) + quantizer_set = _mxfp8_grouped_quantizer_set(group_size) output = te_grouped_dense( x_reshaped, w, n_groups, - kernel_amax=w_amax, quantizer_set=quantizer_set, kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis), ) output = output.reshape(*x.shape[:-1], -1) return output - def run(x, w, w_amax): - output = sharded_group_gemm(x, w, w_amax) + def run(x, w): + output = sharded_group_gemm(x, w) return output - output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax) - dx, dw, _ = vjp_fn(output) + output, vjp_fn = jax.vjp(run, outter_x, outter_w) + dx, dw = vjp_fn(output) return output, dx, dw def ref_func(outter_x, outter_w): @@ -101,13 +115,7 @@ def sharded_group_gemm(x, w): x_reshaped = x.reshape(-1, x.shape[-1]) n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) - quantizer_set = QuantizerFactory.create_set( - scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING, - fwd_dtype=jnp.float8_e4m3fn, - bwd_dtype=jnp.float8_e5m2, - is_2x2x=True, - n_groups=group_size, - ) + quantizer_set = _mxfp8_grouped_quantizer_set(group_size) output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set) output = output.reshape(*x.shape[:-1], -1) return output @@ -120,13 +128,13 @@ def run(x, w): dx, dw = vjp_fn(output) return output, dx, dw - init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None)) - x, w, w_global, w_amax = init_func() + init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding)) + x, w, w_global = init_func() o_sharding = x_sharding test_func_jitted = jax.jit( test_func, - in_shardings=(x_sharding, w_sharding, None), + in_shardings=(x_sharding, w_sharding), out_shardings=(o_sharding, x_sharding, w_sharding), ) ref_func_jitted = jax.jit( @@ -135,24 +143,76 @@ def run(x, w): out_shardings=(o_sharding, x_sharding, w_no_sharding), ) - out, dx, dw = test_func_jitted(x, w, w_amax) + out, dx, dw = test_func_jitted(x, w) ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global) e4m3_tols = dtype_tols(jnp.float8_e4m3fn) - e5m2_tols = dtype_tols(jnp.float8_e5m2) - out, ref_out = jem.process_allgather((out, ref_out)) - dx, ref_dx = jem.process_allgather((dx, ref_dx)) - dw, ref_dw = jem.process_allgather((dw, ref_dw)) + out, ref_out = jem.process_allgather((out, ref_out), tiled=True) + dx, ref_dx = jem.process_allgather((dx, ref_dx), tiled=True) + dw, ref_dw = jem.process_allgather((dw, ref_dw), tiled=True) + + assert_allclose(out, ref_out, **e4m3_tols) + assert_allclose(dx, ref_dx, **e4m3_tols) + assert_allclose(dw, ref_dw, **e4m3_tols) + + +def run_grouped_dense_mxfp8_ep_fsdp_outside_shard_map(): + n_groups = 4 + group_tokens = 128 + hidden = 256 + out_hidden = 128 + x_shape = (n_groups * group_tokens, hidden) + w_shape = (n_groups, hidden, out_hidden) + quantizer_set = _mxfp8_grouped_quantizer_set(n_groups) + + x_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME, None)) + w_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME, FSDP_AXIS_NAME, None)) + group_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME)) + out_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME, None)) + + with mesh, global_shard_guard( + MeshResource(ep_resource=EP_AXIS_NAME, fsdp_resource=FSDP_AXIS_NAME) + ): + x = jax.device_put( + jax.random.normal(jax.random.PRNGKey(20), x_shape, dtype=jnp.bfloat16) + * jnp.asarray(0.01, dtype=jnp.bfloat16), + x_sharding, + ) + w = jax.device_put( + jax.random.normal(jax.random.PRNGKey(21), w_shape, dtype=jnp.bfloat16) + * jnp.asarray(0.01, dtype=jnp.bfloat16), + w_sharding, + ) + group_sizes = jax.device_put( + jnp.full((n_groups,), group_tokens, dtype=jnp.int32), + group_sharding, + ) - jnp.allclose(out, ref_out, **e4m3_tols) - jnp.allclose(dx, ref_dx, **e5m2_tols) - jnp.allclose(dw, ref_dw, **e5m2_tols) + def apply(x, w, group_sizes): + return te_grouped_dense( + x, + w, + group_sizes, + contracting_dims=((1,), (1,)), + quantizer_set=quantizer_set, + kernel_fsdp_info=(FSDP_AXIS_NAME, 1), + ) + + out = jax.jit( + apply, + in_shardings=(x_sharding, w_sharding, group_sharding), + out_shardings=out_sharding, + )(x, w, group_sizes) + jax.block_until_ready(out) + + local_out = np.asarray(jax.device_get(out.addressable_data(0))) + assert tuple(out.sharding.spec) == (EP_AXIS_NAME, None) + assert np.all(np.isfinite(local_out)) + assert np.any(local_out != 0.0) if __name__ == "__main__": - from jax.sharding import NamedSharding, PartitionSpec - from jax.experimental import shard_map import sys coord_addr = sys.argv[1] @@ -163,10 +223,14 @@ def run(x, w): coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id ) - mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,)) + mesh = jax.make_mesh((num_procs,), (FSDP_AXIS_NAME,)) with mesh: data_shapes = [((4, 16, 128, 7168), (7168, 2048))] for data_shape in data_shapes: for kernel_fsdp_axis in [1, 2]: test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis) + + if num_procs == 4: + mesh = jax.make_mesh((2, 2), (EP_AXIS_NAME, FSDP_AXIS_NAME)) + run_grouped_dense_mxfp8_ep_fsdp_outside_shard_map() diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 4ff6d07986..b0d971ee25 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -211,6 +211,115 @@ def _get_nvfp4_tensor_scale_inv(amax): return amax / (DATA_DTYPE_MAX * SCALE_DTYPE_MAX) +def _axis_spec_contains(axis_spec, axis): + if axis is None or axis_spec is None: + return False + if isinstance(axis_spec, tuple): + return axis in axis_spec + return axis_spec == axis + + +def _spec_contains_axis(spec, axis): + return any(_axis_spec_contains(axis_spec, axis) for axis_spec in spec) + + +def _strip_axis_from_axis_spec(axis_spec, axis): + if axis is None or axis_spec is None: + return axis_spec + if isinstance(axis_spec, tuple): + stripped = tuple(a for a in axis_spec if a != axis) + if len(stripped) == 0: + return None + return stripped[0] if len(stripped) == 1 else stripped + return None if axis_spec == axis else axis_spec + + +def _strip_axis_from_spec(spec, axis): + return tuple(_strip_axis_from_axis_spec(axis_spec, axis) for axis_spec in spec) + + +def _common_axis(spec_a, spec_b): + axes = [] + for spec in (spec_a, spec_b): + for axis_spec in spec: + if axis_spec is None: + continue + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + for axis in axis_tuple: + if axis is not None and axis not in axes: + axes.append(axis) + for axis in axes: + if _spec_contains_axis(spec_a, axis) and _spec_contains_axis(spec_b, axis): + return axis + return None + + +def _merge_axis_spec(axis_spec_a, axis_spec_b): + axes = [] + for axis_spec in (axis_spec_a, axis_spec_b): + if axis_spec is None: + continue + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + for axis in axis_tuple: + if axis is not None and axis not in axes: + axes.append(axis) + if len(axes) == 0: + return None + return axes[0] if len(axes) == 1 else tuple(axes) + + +def _partition_spec_from_result(mesh, result_info, fallback_spec): + if result_info is not None and result_info.sharding is not None: + return result_info.sharding + return NamedSharding(mesh, PartitionSpec(*fallback_spec)) + + +def _local_shape_from_spec(global_shape, spec, mesh): + local_shape = [] + for dim, axis_spec in zip(global_shape, spec): + axis_size = _axis_spec_size(axis_spec, mesh) + local_shape.append(dim // axis_size) + return tuple(local_shape) + + +def _axis_spec_size(axis_spec, mesh): + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + axis_size = 1 + for axis in axis_tuple: + if axis is not None: + axis_size *= mesh.shape[axis] + return axis_size + + +def _spec_size(spec, mesh): + axis_size = 1 + for axis_spec in spec: + axis_size *= _axis_spec_size(axis_spec, mesh) + return axis_size + + +def _local_2d_sizes_from_spec(shape, spec, axis_boundary, left_size, right_size, mesh): + if len(shape) == len(spec) and len(shape) > 1: + local_shape = _local_shape_from_spec(shape, spec, mesh) + return ( + math.prod(local_shape[:axis_boundary]), + math.prod(local_shape[axis_boundary:]), + ) + + spec_size = _spec_size(spec, mesh) + if spec_size == 1: + return left_size, right_size + if left_size % spec_size == 0: + return left_size // spec_size, right_size + if right_size % spec_size == 0: + return left_size, right_size // spec_size + raise ValueError( + "Cannot derive local grouped GEMM 2D sizes from sharding spec. " + f"shape={shape}, spec={spec}, axis_boundary={axis_boundary}, " + f"left_size={left_size}, right_size={right_size}, spec_size={spec_size}" + ) + + def collective_gemm_bootstrap( num_total_devices, num_devices_per_process, @@ -1738,6 +1847,289 @@ def impl( ) return (out,) + @staticmethod + def _parse_partition_specs( + mesh, + arg_infos, + result_infos, + out_shape=None, + lhs_is_trans=None, + lhs_axis_boundary=None, + ): + del mesh + gsr = global_mesh_resource() + fsdp_axis = gsr.fsdp_resource + + lhs_data_spec = get_padded_spec(arg_infos[0]) + lhs_scale_spec = get_padded_spec(arg_infos[1]) + rhs_data_spec = get_padded_spec(arg_infos[2]) + rhs_scale_spec = get_padded_spec(arg_infos[3]) + bias_spec = get_padded_spec(arg_infos[4]) + + lhs_first_dims_spec = get_padded_spec(arg_infos[5]) + lhs_last_dims_spec = get_padded_spec(arg_infos[6]) + rhs_first_dims_spec = get_padded_spec(arg_infos[7]) + rhs_last_dims_spec = get_padded_spec(arg_infos[8]) + out_first_dims_spec = get_padded_spec(arg_infos[9]) + out_last_dims_spec = get_padded_spec(arg_infos[10]) + additional_arg_0_spec = get_padded_spec(arg_infos[11]) + additional_arg_1_spec = get_padded_spec(arg_infos[12]) + + grouped_dim_specs = ( + lhs_first_dims_spec, + lhs_last_dims_spec, + rhs_first_dims_spec, + rhs_last_dims_spec, + out_first_dims_spec, + out_last_dims_spec, + ) + grouped_dim_infos = arg_infos[5:11] + active_group_spec = next( + (spec for spec, info in zip(grouped_dim_specs, grouped_dim_infos) if info.size > 0), + (None,), + ) + if arg_infos[11].size > 1: + additional_arg_0_spec = active_group_spec + if arg_infos[12].size > 1: + additional_arg_1_spec = active_group_spec + + rhs_is_ragged = arg_infos[7].size > 0 or arg_infos[8].size > 0 + ep_axis = gsr.ep_resource + if ep_axis is not None and not rhs_is_ragged and _spec_contains_axis(active_group_spec, ep_axis): + if len(rhs_data_spec) > 0 and not _spec_contains_axis(rhs_data_spec, ep_axis): + rhs_data_spec = ( + _merge_axis_spec(rhs_data_spec[0], ep_axis), + *rhs_data_spec[1:], + ) + if len(rhs_scale_spec) > 0 and not _spec_contains_axis(rhs_scale_spec, ep_axis): + rhs_scale_spec = ( + _merge_axis_spec(rhs_scale_spec[0], ep_axis), + *rhs_scale_spec[1:], + ) + if len(bias_spec) > 0 and not _spec_contains_axis(bias_spec, ep_axis): + bias_spec = (_merge_axis_spec(bias_spec[0], ep_axis), *bias_spec[1:]) + + gather_rhs_fsdp = ( + fsdp_axis is not None + and not rhs_is_ragged + and ( + _spec_contains_axis(rhs_data_spec, fsdp_axis) + or _spec_contains_axis(rhs_scale_spec, fsdp_axis) + or _spec_contains_axis(bias_spec, fsdp_axis) + ) + ) + + if gather_rhs_fsdp: + rhs_data_spec = _strip_axis_from_spec(rhs_data_spec, fsdp_axis) + rhs_scale_spec = _strip_axis_from_spec(rhs_scale_spec, fsdp_axis) + bias_spec = _strip_axis_from_spec(bias_spec, fsdp_axis) + + reduce_axis = _common_axis(lhs_data_spec, rhs_data_spec) + if reduce_axis not in (gsr.dp_resource, gsr.fsdp_resource): + reduce_axis = None + if reduce_axis is not None and gather_rhs_fsdp: + reduce_axis = None + + if result_infos: + out_spec = get_padded_spec(result_infos[0]) + else: + out_spec = (None,) * (len(out_shape) if out_shape is not None else 1) + + if rhs_is_ragged and lhs_is_trans is not None and lhs_axis_boundary is not None: + lhs_non_contracting_dims = ( + range(lhs_axis_boundary, len(lhs_data_spec)) + if lhs_is_trans + else range(0, lhs_axis_boundary) + ) + lhs_data_spec = list(lhs_data_spec) + for out_idx, lhs_dim in enumerate(lhs_non_contracting_dims, start=1): + if out_idx < len(out_spec): + lhs_data_spec[lhs_dim] = _merge_axis_spec( + lhs_data_spec[lhs_dim], out_spec[out_idx] + ) + lhs_data_spec = tuple(lhs_data_spec) + + return ( + ( + lhs_data_spec, + lhs_scale_spec, + rhs_data_spec, + rhs_scale_spec, + bias_spec, + lhs_first_dims_spec, + lhs_last_dims_spec, + rhs_first_dims_spec, + rhs_last_dims_spec, + out_first_dims_spec, + out_last_dims_spec, + additional_arg_0_spec, + additional_arg_1_spec, + ), + out_spec, + reduce_axis, + ) + + @staticmethod + def partition( + lhs_is_trans, + rhs_is_trans, + scaling_mode, + out_dtype, + has_bias, + use_async_d2h_group_sizes, + use_v2_ffi, + lhs_axis_boundary, + rhs_axis_boundary, + out_shape, + lhs_left_size, + lhs_right_size, + rhs_left_size, + rhs_right_size, + mesh, + arg_infos, + result_infos, + ): + arg_specs, out_spec, reduce_axis = GroupedGemmPrimitive._parse_partition_specs( + mesh, + arg_infos, + result_infos, + out_shape, + lhs_is_trans=lhs_is_trans, + lhs_axis_boundary=lhs_axis_boundary, + ) + arg_shardings = tuple(NamedSharding(mesh, PartitionSpec(*spec)) for spec in arg_specs) + result_info = result_infos[0] if result_infos else None + out_sharding = (_partition_spec_from_result(mesh, result_info, out_spec),) + local_out_shape = _local_shape_from_spec(out_shape, out_spec, mesh) + local_lhs_left_size, local_lhs_right_size = _local_2d_sizes_from_spec( + arg_infos[0].shape, + arg_specs[0], + lhs_axis_boundary, + lhs_left_size, + lhs_right_size, + mesh, + ) + local_rhs_left_size, local_rhs_right_size = _local_2d_sizes_from_spec( + arg_infos[2].shape, + arg_specs[2], + rhs_axis_boundary, + rhs_left_size, + rhs_right_size, + mesh, + ) + + def sharded_impl( + lhs_data, + lhs_scale_inv, + rhs_data, + rhs_scale_inv, + bias, + lhs_first_dims, + lhs_last_dims, + rhs_first_dims, + rhs_last_dims, + out_first_dims, + out_last_dims, + additional_arg_0, + additional_arg_1, + ): + (out,) = GroupedGemmPrimitive.impl( + lhs_data, + lhs_scale_inv, + rhs_data, + rhs_scale_inv, + bias, + lhs_first_dims, + lhs_last_dims, + rhs_first_dims, + rhs_last_dims, + out_first_dims, + out_last_dims, + additional_arg_0, + additional_arg_1, + lhs_is_trans=lhs_is_trans, + rhs_is_trans=rhs_is_trans, + scaling_mode=scaling_mode, + out_dtype=out_dtype, + has_bias=has_bias, + use_async_d2h_group_sizes=use_async_d2h_group_sizes, + use_v2_ffi=use_v2_ffi, + lhs_axis_boundary=lhs_axis_boundary, + rhs_axis_boundary=rhs_axis_boundary, + out_shape=local_out_shape, + lhs_left_size=local_lhs_left_size, + lhs_right_size=local_lhs_right_size, + rhs_left_size=local_rhs_left_size, + rhs_right_size=local_rhs_right_size, + ) + + if reduce_axis is not None: + if is_all_reduce_in_float32(): + out = jax.lax.psum(out.astype(jnp.float32), reduce_axis).astype(out_dtype) + else: + out = jax.lax.psum(out, reduce_axis) + return (out,) + + return mesh, sharded_impl, out_sharding, arg_shardings + + @staticmethod + def shardy_sharding_rule( + lhs_is_trans, + rhs_is_trans, + scaling_mode, + out_dtype, + has_bias, + use_async_d2h_group_sizes, + use_v2_ffi, + lhs_axis_boundary, + rhs_axis_boundary, + out_shape, + lhs_left_size, + lhs_right_size, + rhs_left_size, + rhs_right_size, + mesh, + operand_types, + result_types, + ): + del ( + lhs_is_trans, + rhs_is_trans, + scaling_mode, + out_dtype, + has_bias, + use_async_d2h_group_sizes, + use_v2_ffi, + lhs_axis_boundary, + rhs_axis_boundary, + out_shape, + lhs_left_size, + lhs_right_size, + rhs_left_size, + rhs_right_size, + mesh, + ) + + prefix = "GroupedGemm" + + def spec_for(name, rank): + if rank == 0: + return () + return tuple(f"{prefix}_{name}_{i}" for i in range(rank)) + + operand_mappings = tuple( + spec_for(f"arg{i}", len(operand_type.shape)) + for i, operand_type in enumerate(operand_types) + ) + result_mappings = tuple( + spec_for(f"out{i}", len(result_type.shape)) + for i, result_type in enumerate(result_types) + ) + return SdyShardingRule( + operand_mappings=operand_mappings, + result_mappings=result_mappings, + ) + register_primitive(GroupedGemmPrimitive) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 7138cfcf40..761933b3f9 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -32,6 +32,8 @@ all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp, get_num_devices_in_mesh, + global_mesh_resource, + lax_paral_op, ) from ..quantize import ( ScaledTensor2x, @@ -52,6 +54,59 @@ __all__ = ["quantize", "quantize_dbias", "grouped_quantize", "grouped_dbias"] +def _merge_axis_specs(axis_specs): + axes = [] + for axis_spec in axis_specs: + if axis_spec is None: + continue + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + for axis in axis_tuple: + if axis is not None and axis not in axes: + axes.append(axis) + if len(axes) == 0: + return None + return axes[0] if len(axes) == 1 else tuple(axes) + + +def _flat_data_spec(input_spec): + return (_merge_axis_specs(input_spec),) + + +def _axis_spec_size(axis_spec, mesh): + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + axis_size = 1 + for axis in axis_tuple: + if axis is not None: + axis_size *= mesh.shape[axis] + return axis_size + + +def _local_shape_from_spec(global_shape, spec, mesh): + local_shape = [] + for dim, axis_spec in zip(global_shape, spec): + local_shape.append(dim // _axis_spec_size(axis_spec, mesh)) + return tuple(local_shape) + + +def _pad_or_slice_to_shape(x, target_shape): + if target_shape is None or x.shape == target_shape: + return x + target_size = math.prod(target_shape) + current_size = math.prod(x.shape) + x = x.reshape(-1) + if current_size > target_size: + return x[:target_size].reshape(target_shape) + return jnp.pad(x, (0, target_size - current_size)).reshape(target_shape) + + +def _all_reduce_grouped_amax_along_dp_fsdp(amax, mesh): + gsr = global_mesh_resource() + for axis in (gsr.dp_resource, gsr.fsdp_resource): + if axis is not None and axis in mesh.axis_names: + amax = lax_paral_op(amax, jax.lax.pmax, axis, mesh) + return amax + + class BaseDBiasQuantizePrimitive(BasePrimitive): """ Cast Primitive wrapping nvte_quantize and nvte_quantize_dbias @@ -1236,6 +1291,143 @@ def impl( ) return rowwise_out, colwise_out, rowwise_scale_inv, colwise_scale_inv, updated_amax + @staticmethod + def _parse_partition_specs(scaling_mode, q_layout, mesh, arg_infos): + del mesh + x_spec = get_padded_spec(arg_infos[0]) + group_spec = get_padded_spec(arg_infos[2]) + if group_spec == (None,) and len(x_spec) > 0: + group_spec = (x_spec[0],) + flat_spec = _flat_data_spec(x_spec) + replicated_spec = (None,) + + rowwise_out_spec = flat_spec if q_layout.has_rowwise else replicated_spec + colwise_out_spec = flat_spec if q_layout.has_colwise else replicated_spec + + rowwise_scale_inv_spec = replicated_spec + colwise_scale_inv_spec = replicated_spec + if ScalingMode(scaling_mode).is_block_scaling: + rowwise_scale_inv_spec = flat_spec if q_layout.has_rowwise else replicated_spec + colwise_scale_inv_spec = flat_spec if q_layout.has_colwise else replicated_spec + elif ScalingMode(scaling_mode).is_tensor_scaling(): + rowwise_scale_inv_spec = group_spec if q_layout.has_rowwise else replicated_spec + colwise_scale_inv_spec = group_spec if q_layout.has_colwise else replicated_spec + + updated_amax_spec = group_spec + return ( + x_spec, + group_spec, + ( + rowwise_out_spec, + colwise_out_spec, + rowwise_scale_inv_spec, + colwise_scale_inv_spec, + updated_amax_spec, + ), + ) + + @staticmethod + def partition( + out_dtype, + scaling_mode, + q_layout, + flatten_axis, + scale_dtype, + mesh, + arg_infos, + result_infos, + ): + x_spec, group_spec, out_specs = GroupedQuantizePrimitive._parse_partition_specs( + scaling_mode, q_layout, mesh, arg_infos + ) + local_out_shapes = ( + tuple(_local_shape_from_spec(info.shape, spec, mesh) for info, spec in zip(result_infos, out_specs)) + if result_infos + else (None,) * len(out_specs) + ) + + arg_shardings = ( + NamedSharding(mesh, PartitionSpec(*x_spec)), + NamedSharding(mesh, PartitionSpec(*group_spec)), + NamedSharding(mesh, PartitionSpec(*group_spec)), + ) + out_shardings = tuple(NamedSharding(mesh, PartitionSpec(*spec)) for spec in out_specs) + + def sharded_impl(x, scale, group_sizes): + ( + rowwise_out, + colwise_out, + rowwise_scale_inv, + colwise_scale_inv, + updated_amax, + ) = GroupedQuantizePrimitive.impl( + x, + scale, + group_sizes, + out_dtype=out_dtype, + scaling_mode=scaling_mode, + q_layout=q_layout, + flatten_axis=flatten_axis, + scale_dtype=scale_dtype, + ) + if ScalingMode(scaling_mode).is_block_scaling: + rowwise_scale_inv = _pad_or_slice_to_shape(rowwise_scale_inv, local_out_shapes[2]) + colwise_scale_inv = _pad_or_slice_to_shape(colwise_scale_inv, local_out_shapes[3]) + if ScalingMode(scaling_mode).is_tensor_scaling(): + updated_amax = _all_reduce_grouped_amax_along_dp_fsdp(updated_amax, mesh) + return ( + rowwise_out, + colwise_out, + rowwise_scale_inv, + colwise_scale_inv, + updated_amax, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings + + @staticmethod + def shardy_sharding_rule( + out_dtype, + scaling_mode, + q_layout, + flatten_axis, + scale_dtype, + mesh, + value_types, + result_types, + ): + del out_dtype, scale_dtype, mesh, result_types, flatten_axis + + prefix = "GroupedQuantize" + input_spec = tuple(f"{prefix}_x_{i}" for i in range(len(value_types[0].shape))) + flat_spec = (f"{prefix}_flat",) + group_spec = (BATCHING + f"{prefix}_group",) + scalar_spec = (BATCHING + f"{prefix}_scalar",) + + rowwise_out_spec = flat_spec if q_layout.has_rowwise else scalar_spec + colwise_out_spec = flat_spec if q_layout.has_colwise else scalar_spec + + if ScalingMode(scaling_mode).is_block_scaling: + rowwise_scale_spec = flat_spec if q_layout.has_rowwise else scalar_spec + colwise_scale_spec = flat_spec if q_layout.has_colwise else scalar_spec + elif ScalingMode(scaling_mode).is_tensor_scaling(): + rowwise_scale_spec = group_spec if q_layout.has_rowwise else scalar_spec + colwise_scale_spec = group_spec if q_layout.has_colwise else scalar_spec + else: + rowwise_scale_spec = scalar_spec + colwise_scale_spec = scalar_spec + + return SdyShardingRule( + operand_mappings=(input_spec, group_spec, group_spec), + result_mappings=( + rowwise_out_spec, + colwise_out_spec, + rowwise_scale_spec, + colwise_scale_spec, + group_spec, + ), + ) + register_primitive(GroupedQuantizePrimitive) diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index f8c30ffccb..70151a44b7 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -14,9 +14,11 @@ import warnings import jax import jax.numpy as jnp +from jax.sharding import PartitionSpec from . import cpp_extensions as tex from .cpp_extensions.amax import AmaxScope +from .sharding import global_mesh_resource, with_sharding_constraint from .quantize import ( ScaledTensor, QuantizerSet, @@ -54,6 +56,10 @@ def _psum_scatter_kernel(kernel, scattered_kernel_shape, mesh_axis, axis_idx): return kernel +def _is_manual_mesh_axis(mesh_axis): + return mesh_axis is not None and mesh_axis in jax.sharding.get_abstract_mesh().manual_axes + + def dense( x: jnp.ndarray, kernel: jnp.ndarray, @@ -349,6 +355,18 @@ def grouped_dense( Returns: A jnp.ndarray containing the result of the grouped linear operation """ + x_contracting_dims, kernel_contracting_dims = contracting_dims + x_contracting_dims = tex.sanitize_dims(x.ndim, x_contracting_dims) + kernel_contracting_dims = tex.sanitize_dims(kernel.ndim, kernel_contracting_dims) + contracting_dims = (x_contracting_dims, kernel_contracting_dims) + + restore_leading_ep_axis = False + if x.ndim == 3 and x.shape[0] == 1: + if x_contracting_dims == (x.ndim - 1,): + restore_leading_ep_axis = True + x = x.reshape(*x.shape[1:]) + contracting_dims = ((x.ndim - 1,), kernel_contracting_dims) + output = _grouped_dense( x, kernel, @@ -361,6 +379,8 @@ def grouped_dense( quantizer_set, kernel_fsdp_info, ) + if restore_leading_ep_axis: + output = output.reshape(1, *output.shape) return output @@ -406,10 +426,7 @@ def _grouped_dense_fwd_rule( ): use_bias = bias is not None - kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info - kernel_fsdp_enabled = kernel_fsdp_mesh_axis is not None - assert not kernel_fsdp_enabled, "FSDP sharding for grouped_dense is not supported yet." - del kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx, kernel_fsdp_info, kernel_fsdp_enabled + del kernel_fsdp_info x_contracting_dims, k_contracting_dims = contracting_dims flatten_axis_x = -len(x_contracting_dims) @@ -478,9 +495,7 @@ def _grouped_dense_fwd_rule( def _grouped_dense_bwd_rule( contracting_dims, precision, preferred_element_type, group_offset, kernel_fsdp_info, ctx, grad ): - kernel_fsdp_mesh_axis, _ = kernel_fsdp_info - kernel_fsdp_enabled = kernel_fsdp_mesh_axis is not None - assert not kernel_fsdp_enabled, "FSDP sharding for grouped_dense is not supported yet." + kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info fwd_x_contracting_dims, fwd_k_contracting_dims = contracting_dims @@ -530,6 +545,14 @@ def _grouped_dense_bwd_rule( preferred_element_type=preferred_element_type, group_offset=group_offset, ) + if _is_manual_mesh_axis(kernel_fsdp_mesh_axis): + if kernel_fsdp_axis_idx in fwd_k_contracting_dims: + dgrad_axis_idx = fwd_x_contracting_dims[ + fwd_k_contracting_dims.index(kernel_fsdp_axis_idx) + ] + dgrad = _all_gather_kernel(dgrad, kernel_fsdp_mesh_axis, dgrad_axis_idx) + else: + dgrad = jax.lax.psum(dgrad, kernel_fsdp_mesh_axis) wgrad = tex.grouped_gemm( wgrad_x_T, @@ -539,6 +562,25 @@ def _grouped_dense_bwd_rule( preferred_element_type=preferred_element_type, group_offset=group_offset, ) + if _is_manual_mesh_axis(kernel_fsdp_mesh_axis): + if kernel_fsdp_axis_idx in fwd_k_contracting_dims: + wgrad = _psum_scatter_kernel( + wgrad, kernel_shape, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx + ) + else: + wgrad = jax.lax.psum(wgrad, kernel_fsdp_mesh_axis) + if kernel_fsdp_mesh_axis is not None: + wgrad_spec = [None] * len(kernel_shape) + ep_resource = None + try: + ep_resource = global_mesh_resource().ep_resource + except AssertionError: + pass + if len(wgrad_spec) > 0: + wgrad_spec[0] = ep_resource + if 0 <= kernel_fsdp_axis_idx < len(wgrad_spec): + wgrad_spec[kernel_fsdp_axis_idx] = kernel_fsdp_mesh_axis + wgrad = with_sharding_constraint(wgrad, PartitionSpec(*wgrad_spec)) group_sizes_grad = None dbias = tex.grouped_dbias(grad, group_sizes) if use_bias else None diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 17c9a242f0..14783ecbe2 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -1471,7 +1471,7 @@ def te_grouped_dot_general(generate_quantizer_set, x, kernel, group_sizes, **kwa x, kernel, group_sizes=group_sizes, - contracting_dims=((1,), (1,)), + contracting_dims=((-1,), (1,)), quantizer_set=quantizer_set, ) return out diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 9b13412c14..16a10c860d 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -330,6 +330,7 @@ class MeshResource: tp_resource: Axis name for tensor parallelism (hidden dimension sharding), default is None tpsp_resource: Axis name for tensor sequence parallelism (hidden and sequence sharding), default is None fsdp_resource: Axis name for full-sharded data parallelism, default is None + ep_resource: Axis name for expert parallelism (expert sharding), default is None pp_resource: Axis name for pipeline parallelism (layer sharding), default is None cp_resource: Axis name for context parallelism (sequence sharding), default is None """ @@ -338,6 +339,7 @@ class MeshResource: tp_resource: str = None tpsp_resource: str = None fsdp_resource: str = None + ep_resource: str = None pp_resource: str = None cp_resource: str = None From 60a0b50400b2309c27bd8b3d8566115004176bcc Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 26 May 2026 16:13:53 -0700 Subject: [PATCH 02/44] Add outside shard_map grouped GMM backward test --- ..._multi_process_distributed_grouped_gemm.py | 42 +++++++++++-------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py index 30a1452a07..de5e031566 100644 --- a/tests/jax/test_multi_process_distributed_grouped_gemm.py +++ b/tests/jax/test_multi_process_distributed_grouped_gemm.py @@ -189,27 +189,35 @@ def run_grouped_dense_mxfp8_ep_fsdp_outside_shard_map(): group_sharding, ) - def apply(x, w, group_sizes): - return te_grouped_dense( - x, - w, - group_sizes, - contracting_dims=((1,), (1,)), - quantizer_set=quantizer_set, - kernel_fsdp_info=(FSDP_AXIS_NAME, 1), - ) - - out = jax.jit( - apply, + def apply_with_vjp(x, w, group_sizes): + def apply(x, w): + return te_grouped_dense( + x, + w, + group_sizes, + contracting_dims=((1,), (1,)), + quantizer_set=quantizer_set, + kernel_fsdp_info=(FSDP_AXIS_NAME, 1), + ) + + out, vjp_fn = jax.vjp(apply, x, w) + dx, dw = vjp_fn(out) + return out, dx, dw + + out, dx, dw = jax.jit( + apply_with_vjp, in_shardings=(x_sharding, w_sharding, group_sharding), - out_shardings=out_sharding, + out_shardings=(out_sharding, x_sharding, w_sharding), )(x, w, group_sizes) - jax.block_until_ready(out) + out, dx, dw = jax.block_until_ready((out, dx, dw)) - local_out = np.asarray(jax.device_get(out.addressable_data(0))) assert tuple(out.sharding.spec) == (EP_AXIS_NAME, None) - assert np.all(np.isfinite(local_out)) - assert np.any(local_out != 0.0) + assert tuple(dx.sharding.spec) == (EP_AXIS_NAME, None) + assert tuple(dw.sharding.spec) == (EP_AXIS_NAME, FSDP_AXIS_NAME, None) + for value in (out, dx, dw): + local_value = np.asarray(jax.device_get(value.addressable_data(0))) + assert np.all(np.isfinite(local_value)) + assert np.any(local_value != 0.0) if __name__ == "__main__": From 786fa1d961de8d5ae7acac09b544241adda9dc7d Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 28 May 2026 16:07:07 -0700 Subject: [PATCH 03/44] Use 4 GPU mesh for grouped GEMM partitioning test --- tests/jax/test_grouped_gemm_partitioning.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/jax/test_grouped_gemm_partitioning.py b/tests/jax/test_grouped_gemm_partitioning.py index 5fb93e30ae..d69f0c29b8 100644 --- a/tests/jax/test_grouped_gemm_partitioning.py +++ b/tests/jax/test_grouped_gemm_partitioning.py @@ -8,6 +8,7 @@ import jax import jax.numpy as jnp import numpy as np +import pytest from jax.sharding import Mesh, NamedSharding, PartitionSpec from transformer_engine.jax.cpp_extensions.gemm import GroupedGemmPrimitive @@ -18,7 +19,10 @@ def _mesh(): - return Mesh(np.asarray(jax.devices()[:1]).reshape(1, 1), ("expert", "fsdp")) + devices = jax.devices() + if len(devices) < 4: + pytest.skip("Grouped GEMM partitioning tests require at least 4 visible GPUs.") + return Mesh(np.asarray(devices[:4]).reshape(2, 2), ("expert", "fsdp")) def _arg_info(mesh, shape, spec): @@ -187,9 +191,9 @@ def test_grouped_partitioning_shardy_rules_smoke(): def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process(): mesh = _mesh() - n_groups = 2 + n_groups = 4 group_tokens = 128 - hidden = 128 + hidden = 256 out_hidden = 128 x_shape = (n_groups * group_tokens, hidden) w_shape = (n_groups, hidden, out_hidden) From ff0407dada98fed561347146162720fb94d4c51f Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 1 Jun 2026 08:46:14 -0700 Subject: [PATCH 04/44] progress Signed-off-by: Jeremy Berchtold --- tests/jax/test_grouped_gemm_partitioning.py | 242 +++++++++++++++++- ..._multi_process_distributed_grouped_gemm.py | 78 ++++-- transformer_engine/jax/cpp_extensions/gemm.py | 73 ++++-- .../jax/cpp_extensions/quantization.py | 45 +++- transformer_engine/jax/dense.py | 69 ++++- transformer_engine/jax/sharding.py | 26 +- 6 files changed, 448 insertions(+), 85 deletions(-) diff --git a/tests/jax/test_grouped_gemm_partitioning.py b/tests/jax/test_grouped_gemm_partitioning.py index d69f0c29b8..831e245c0e 100644 --- a/tests/jax/test_grouped_gemm_partitioning.py +++ b/tests/jax/test_grouped_gemm_partitioning.py @@ -25,6 +25,23 @@ def _mesh(): return Mesh(np.asarray(devices[:4]).reshape(2, 2), ("expert", "fsdp")) +def _mesh_with_dp_tp(): + devices = jax.devices() + if len(devices) < 4: + pytest.skip("Grouped GEMM partitioning tests require at least 4 visible GPUs.") + return Mesh(np.asarray(devices[:4]).reshape(2, 1, 2, 1), ("expert", "dp", "fsdp", "tp")) + + +def _mesh_with_arbitrary_axis(): + devices = jax.devices() + if len(devices) < 4: + pytest.skip("Grouped GEMM partitioning tests require at least 4 visible GPUs.") + return Mesh( + np.asarray(devices[:4]).reshape(2, 1, 2, 1), + ("expert", "dp", "fsdp", "myaxis123"), + ) + + def _arg_info(mesh, shape, spec): return SimpleNamespace( shape=shape, @@ -40,6 +57,14 @@ def _normalize_spec(spec): return spec +def _spec_contains_axis(spec, axis): + for axis_spec in spec: + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + if axis in axis_tuple: + return True + return False + + def _mxfp8_grouped_quantizer_set(n_groups): return QuantizerFactory.create_set( scaling_mode=ScalingMode.MXFP8_1D_SCALING, @@ -50,10 +75,10 @@ def _mxfp8_grouped_quantizer_set(n_groups): ) -def test_grouped_quantize_specs_preserve_ep_and_fsdp_for_block_scales(): +def test_grouped_quantize_gathers_hidden_axis_for_block_scales(): mesh = _mesh() with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - _, _, out_shardings, _ = GroupedQuantizePrimitive.partition( + _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( jnp.float8_e4m3fn, ScalingMode.MXFP8_1D_SCALING.value, QuantizeLayout.ROWWISE, @@ -68,16 +93,17 @@ def test_grouped_quantize_specs_preserve_ep_and_fsdp_for_block_scales(): (), ) + assert tuple(arg_shardings[0].spec) == ("expert", None, None) specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == (("expert", "fsdp"),) - assert _normalize_spec(specs[2]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[0]) == ("expert",) + assert _normalize_spec(specs[2]) == ("expert",) assert _normalize_spec(specs[4]) == ("expert",) -def test_grouped_quantize_mxfp8_colwise_specs_preserve_ep_and_fsdp(): +def test_grouped_quantize_mxfp8_colwise_specs_gather_hidden_axis(): mesh = _mesh() with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - _, _, out_shardings, _ = GroupedQuantizePrimitive.partition( + _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( jnp.float8_e4m3fn, ScalingMode.MXFP8_1D_SCALING.value, QuantizeLayout.ROWWISE_COLWISE, @@ -92,14 +118,47 @@ def test_grouped_quantize_mxfp8_colwise_specs_preserve_ep_and_fsdp(): (), ) + assert tuple(arg_shardings[0].spec) == ("expert", None, None) specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == (("expert", "fsdp"),) - assert _normalize_spec(specs[1]) == (("expert", "fsdp"),) - assert _normalize_spec(specs[2]) == (("expert", "fsdp"),) - assert _normalize_spec(specs[3]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[0]) == ("expert",) + assert _normalize_spec(specs[1]) == ("expert",) + assert _normalize_spec(specs[2]) == ("expert",) + assert _normalize_spec(specs[3]) == ("expert",) assert _normalize_spec(specs[4]) == ("expert",) +def test_grouped_quantize_strips_unsupported_axes_and_gathers_hidden_axes(): + mesh = _mesh_with_dp_tp() + with jax.set_mesh(mesh), global_shard_guard( + MeshResource(dp_resource="dp", tp_resource="tp", fsdp_resource="fsdp", ep_resource="expert") + ): + _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 128), ("expert", "dp", ("fsdp", "tp"))), + _arg_info(mesh, (8,), (("expert", "tp"),)), + _arg_info(mesh, (8,), (("expert", "tp"),)), + ), + (), + ) + + assert tuple(arg_shardings[0].spec) == ("expert", None, None) + assert tuple(arg_shardings[1].spec) == ("expert",) + assert tuple(arg_shardings[2].spec) == ("expert",) + + out_specs = tuple(tuple(sharding.spec) for sharding in out_shardings) + assert _normalize_spec(out_specs[0]) == ("expert",) + assert _normalize_spec(out_specs[2]) == ("expert",) + assert _normalize_spec(out_specs[4]) == ("expert",) + for spec in (*out_specs, *(tuple(sharding.spec) for sharding in arg_shardings)): + assert not _spec_contains_axis(spec, "tp") + + def test_grouped_gemm_rhs_weight_specs_gather_fsdp_but_preserve_ep(): mesh = _mesh() arg_infos = ( @@ -143,6 +202,169 @@ def test_grouped_gemm_rhs_weight_specs_gather_fsdp_but_preserve_ep(): assert tuple(out_sharding[0].spec) == (None, None, None) +def test_grouped_gemm_strips_unsupported_axes_preserves_dp_and_gathers_rhs_fsdp(): + mesh = _mesh_with_dp_tp() + arg_infos = ( + _arg_info(mesh, (8192,), (("dp", "tp"),)), + _arg_info(mesh, (0,), (("tp",),)), + _arg_info(mesh, (65536,), (("expert", "fsdp", "tp"),)), + _arg_info(mesh, (2048,), (("expert", "fsdp", "tp"),)), + _arg_info(mesh, (0,), (("fsdp", "tp"),)), + _arg_info(mesh, (8,), (("expert", "tp"),)), + _arg_info(mesh, (0,), (("tp",),)), + _arg_info(mesh, (0,), (("tp",),)), + _arg_info(mesh, (0,), (("tp",),)), + _arg_info(mesh, (8,), (("expert", "tp"),)), + _arg_info(mesh, (0,), (("tp",),)), + _arg_info(mesh, (1,), (("tp",),)), + _arg_info(mesh, (0,), (("tp",),)), + ) + result_infos = (_arg_info(mesh, (1, 128, 64), ("expert", "tp", None)),) + with jax.set_mesh(mesh), global_shard_guard( + MeshResource(dp_resource="dp", tp_resource="tp", fsdp_resource="fsdp", ep_resource="expert") + ): + _, _, out_sharding, arg_shardings = GroupedGemmPrimitive.partition( + False, + False, + ScalingMode.NO_SCALING.value, + jnp.bfloat16, + False, + False, + False, + 1, + 1, + (1, 128, 64), + 128, + 64, + 128, + 64, + mesh, + arg_infos, + result_infos, + ) + + assert tuple(arg_shardings[0].spec) == ("dp",) + assert tuple(arg_shardings[2].spec) == ("expert",) + assert tuple(arg_shardings[3].spec) == ("expert",) + assert tuple(arg_shardings[5].spec) == ("expert",) + assert tuple(out_sharding[0].spec) == ("expert", None, None) + for spec in ( + *(tuple(sharding.spec) for sharding in arg_shardings), + tuple(out_sharding[0].spec), + ): + assert not _spec_contains_axis(spec, "tp") + + +def test_grouped_gemm_reduce_axis_skips_ep_and_uses_dp(): + mesh = _mesh_with_dp_tp() + arg_infos = ( + _arg_info(mesh, (8192,), (("expert", "dp"),)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (8192,), (("expert", "dp"),)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (1,), (None,)), + _arg_info(mesh, (0,), (None,)), + ) + + with jax.set_mesh(mesh), global_shard_guard( + MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="expert") + ): + _, _, reduce_axis = GroupedGemmPrimitive._parse_partition_specs( + mesh, + arg_infos, + (), + out_shape=(1, 128, 64), + lhs_is_trans=False, + lhs_axis_boundary=1, + ) + + assert reduce_axis == "dp" + + +def test_grouped_partitioning_strips_arbitrary_unsupported_axis(): + mesh = _mesh_with_arbitrary_axis() + mesh_resource = MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="expert") + + with jax.set_mesh(mesh), global_shard_guard(mesh_resource): + _, _, quantize_out_shardings, quantize_arg_shardings = GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 128), ("expert", "myaxis123", ("dp", "fsdp"))), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + ), + (), + ) + + gemm_arg_infos = ( + _arg_info(mesh, (8192,), (("dp", "myaxis123"),)), + _arg_info(mesh, (0,), (("myaxis123",),)), + _arg_info(mesh, (65536,), (("expert", "fsdp", "myaxis123"),)), + _arg_info(mesh, (2048,), (("expert", "fsdp", "myaxis123"),)), + _arg_info(mesh, (0,), (("fsdp", "myaxis123"),)), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + _arg_info(mesh, (0,), (("myaxis123",),)), + _arg_info(mesh, (0,), (("myaxis123",),)), + _arg_info(mesh, (0,), (("myaxis123",),)), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + _arg_info(mesh, (0,), (("myaxis123",),)), + _arg_info(mesh, (1,), (("myaxis123",),)), + _arg_info(mesh, (0,), (("myaxis123",),)), + ) + gemm_result_infos = (_arg_info(mesh, (1, 128, 64), ("expert", "myaxis123", None)),) + _, _, gemm_out_sharding, gemm_arg_shardings = GroupedGemmPrimitive.partition( + False, + False, + ScalingMode.NO_SCALING.value, + jnp.bfloat16, + False, + False, + False, + 1, + 1, + (1, 128, 64), + 128, + 64, + 128, + 64, + mesh, + gemm_arg_infos, + gemm_result_infos, + ) + + assert tuple(quantize_arg_shardings[0].spec) == ("expert", None, None) + assert tuple(quantize_arg_shardings[1].spec) == ("expert",) + quantize_out_specs = tuple(tuple(sharding.spec) for sharding in quantize_out_shardings) + assert _normalize_spec(quantize_out_specs[0]) == ("expert",) + assert _normalize_spec(quantize_out_specs[2]) == ("expert",) + + assert tuple(gemm_arg_shardings[0].spec) == ("dp",) + assert tuple(gemm_arg_shardings[2].spec) == ("expert",) + assert tuple(gemm_arg_shardings[3].spec) == ("expert",) + assert tuple(gemm_out_sharding[0].spec) == ("expert", None, None) + + all_specs = ( + *quantize_out_specs, + *(tuple(sharding.spec) for sharding in quantize_arg_shardings), + *(tuple(sharding.spec) for sharding in gemm_arg_shardings), + tuple(gemm_out_sharding[0].spec), + ) + for spec in all_specs: + assert not _spec_contains_axis(spec, "myaxis123") + + def test_grouped_partitioning_shardy_rules_smoke(): mesh = _mesh() quantize_rule = GroupedQuantizePrimitive.shardy_sharding_rule( diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py index de5e031566..ce52126dea 100644 --- a/tests/jax/test_multi_process_distributed_grouped_gemm.py +++ b/tests/jax/test_multi_process_distributed_grouped_gemm.py @@ -47,11 +47,18 @@ def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis): if kernel_fsdp_axis == 2 else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None)) ) + b_sharding = ( + NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME)) + if kernel_fsdp_axis == 2 + else NamedSharding(mesh, PartitionSpec(None, None)) + ) w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None)) + b_no_sharding = NamedSharding(mesh, PartitionSpec(None, None)) def init_data(): x_key = jax.random.PRNGKey(0) w_key = jax.random.PRNGKey(1) + b_key = jax.random.PRNGKey(2) x = ( jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16) * jnp.asarray(0.01, dtype=jnp.bfloat16) @@ -60,10 +67,14 @@ def init_data(): jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16) * jnp.asarray(0.01, dtype=jnp.bfloat16) ) - return x, w, w + b = ( + jax.random.normal(b_key, shape=(N_GROUP, w_shape[-1]), dtype=jnp.bfloat16) + * jnp.asarray(0.01, dtype=jnp.bfloat16) + ) + return x, w, w, b, b - def test_func(outter_x, outter_w): - in_specs = (x_sharding.spec, w_sharding.spec) + def test_func(outter_x, outter_w, outter_b): + in_specs = (x_sharding.spec, w_sharding.spec, b_sharding.spec) out_specs = x_sharding.spec @partial( @@ -73,7 +84,7 @@ def test_func(outter_x, outter_w): out_specs=out_specs, check_rep=False, ) - def sharded_group_gemm(x, w): + def sharded_group_gemm(x, w, b): group_size = x.shape[0] x_reshaped = x.reshape(-1, x.shape[-1]) n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) @@ -84,23 +95,24 @@ def sharded_group_gemm(x, w): x_reshaped, w, n_groups, + bias=b, quantizer_set=quantizer_set, kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis), ) output = output.reshape(*x.shape[:-1], -1) return output - def run(x, w): - output = sharded_group_gemm(x, w) + def run(x, w, b): + output = sharded_group_gemm(x, w, b) return output - output, vjp_fn = jax.vjp(run, outter_x, outter_w) - dx, dw = vjp_fn(output) - return output, dx, dw + output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_b) + dx, dw, db = vjp_fn(output) + return output, dx, dw, db - def ref_func(outter_x, outter_w): + def ref_func(outter_x, outter_w, outter_b): - in_specs = (x_sharding.spec, w_no_sharding.spec) + in_specs = (x_sharding.spec, w_no_sharding.spec, b_no_sharding.spec) out_specs = x_sharding.spec @partial( @@ -110,51 +122,63 @@ def ref_func(outter_x, outter_w): out_specs=out_specs, check_rep=False, ) - def sharded_group_gemm(x, w): + def sharded_group_gemm(x, w, b): group_size = x.shape[0] x_reshaped = x.reshape(-1, x.shape[-1]) n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) quantizer_set = _mxfp8_grouped_quantizer_set(group_size) - output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set) + output = te_grouped_dense( + x_reshaped, + w, + n_groups, + bias=b, + quantizer_set=quantizer_set, + ) output = output.reshape(*x.shape[:-1], -1) return output - def run(x, w): - output = sharded_group_gemm(x, w) + def run(x, w, b): + output = sharded_group_gemm(x, w, b) return output - output, vjp_fn = jax.vjp(run, outter_x, outter_w) - dx, dw = vjp_fn(output) - return output, dx, dw + output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_b) + dx, dw, db = vjp_fn(output) + return output, dx, dw, db - init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding)) - x, w, w_global = init_func() + init_func = jax.jit( + init_data, + out_shardings=(x_sharding, w_sharding, w_no_sharding, b_sharding, b_no_sharding), + ) + x, w, w_global, b, b_global = init_func() o_sharding = x_sharding test_func_jitted = jax.jit( test_func, - in_shardings=(x_sharding, w_sharding), - out_shardings=(o_sharding, x_sharding, w_sharding), + in_shardings=(x_sharding, w_sharding, b_sharding), + out_shardings=(o_sharding, x_sharding, w_sharding, b_sharding), ) ref_func_jitted = jax.jit( ref_func, - in_shardings=(x_sharding, w_no_sharding), - out_shardings=(o_sharding, x_sharding, w_no_sharding), + in_shardings=(x_sharding, w_no_sharding, b_no_sharding), + out_shardings=(o_sharding, x_sharding, w_no_sharding, b_no_sharding), ) - out, dx, dw = test_func_jitted(x, w) - ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global) + out, dx, dw, db = test_func_jitted(x, w, b) + ref_out, ref_dx, ref_dw, ref_db = ref_func_jitted(x, w_global, b_global) - e4m3_tols = dtype_tols(jnp.float8_e4m3fn) + # Avoid creating a host scalar JAX array under the multi-process mesh in dtype_tols. + e4m3_tols = dtype_tols(jnp.float8_e4m3fn, rtol=0.25, atol=0.25) out, ref_out = jem.process_allgather((out, ref_out), tiled=True) dx, ref_dx = jem.process_allgather((dx, ref_dx), tiled=True) dw, ref_dw = jem.process_allgather((dw, ref_dw), tiled=True) + db, ref_db = jem.process_allgather((db, ref_db), tiled=True) assert_allclose(out, ref_out, **e4m3_tols) assert_allclose(dx, ref_dx, **e4m3_tols) assert_allclose(dw, ref_dw, **e4m3_tols) + assert_allclose(db, ref_db, **e4m3_tols) def run_grouped_dense_mxfp8_ep_fsdp_outside_shard_map(): diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index b0d971ee25..61521873d2 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -238,7 +238,30 @@ def _strip_axis_from_spec(spec, axis): return tuple(_strip_axis_from_axis_spec(axis_spec, axis) for axis_spec in spec) -def _common_axis(spec_a, spec_b): +def _filter_axis_spec(axis_spec, allowed_axes): + if axis_spec is None: + return None + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + axes = tuple(axis for axis in axis_tuple if axis in allowed_axes) + if len(axes) == 0: + return None + return axes[0] if len(axes) == 1 else axes + + +def _filter_spec_axes(spec, allowed_axes): + return tuple(_filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) + + +def _supported_grouped_gemm_axes(mesh): + gsr = global_mesh_resource(validate=False) + return { + axis + for axis in (gsr.ep_resource, gsr.dp_resource, gsr.fsdp_resource) + if axis is not None and axis in mesh.axis_names + } + + +def _common_axis(spec_a, spec_b, allowed_axes=None): axes = [] for spec in (spec_a, spec_b): for axis_spec in spec: @@ -249,6 +272,8 @@ def _common_axis(spec_a, spec_b): if axis is not None and axis not in axes: axes.append(axis) for axis in axes: + if allowed_axes is not None and axis not in allowed_axes: + continue if _spec_contains_axis(spec_a, axis) and _spec_contains_axis(spec_b, axis): return axis return None @@ -1856,24 +1881,24 @@ def _parse_partition_specs( lhs_is_trans=None, lhs_axis_boundary=None, ): - del mesh - gsr = global_mesh_resource() + gsr = global_mesh_resource(validate=False) fsdp_axis = gsr.fsdp_resource - - lhs_data_spec = get_padded_spec(arg_infos[0]) - lhs_scale_spec = get_padded_spec(arg_infos[1]) - rhs_data_spec = get_padded_spec(arg_infos[2]) - rhs_scale_spec = get_padded_spec(arg_infos[3]) - bias_spec = get_padded_spec(arg_infos[4]) - - lhs_first_dims_spec = get_padded_spec(arg_infos[5]) - lhs_last_dims_spec = get_padded_spec(arg_infos[6]) - rhs_first_dims_spec = get_padded_spec(arg_infos[7]) - rhs_last_dims_spec = get_padded_spec(arg_infos[8]) - out_first_dims_spec = get_padded_spec(arg_infos[9]) - out_last_dims_spec = get_padded_spec(arg_infos[10]) - additional_arg_0_spec = get_padded_spec(arg_infos[11]) - additional_arg_1_spec = get_padded_spec(arg_infos[12]) + allowed_axes = _supported_grouped_gemm_axes(mesh) + + lhs_data_spec = _filter_spec_axes(get_padded_spec(arg_infos[0]), allowed_axes) + lhs_scale_spec = _filter_spec_axes(get_padded_spec(arg_infos[1]), allowed_axes) + rhs_data_spec = _filter_spec_axes(get_padded_spec(arg_infos[2]), allowed_axes) + rhs_scale_spec = _filter_spec_axes(get_padded_spec(arg_infos[3]), allowed_axes) + bias_spec = _filter_spec_axes(get_padded_spec(arg_infos[4]), allowed_axes) + + lhs_first_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[5]), allowed_axes) + lhs_last_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[6]), allowed_axes) + rhs_first_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[7]), allowed_axes) + rhs_last_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[8]), allowed_axes) + out_first_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[9]), allowed_axes) + out_last_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[10]), allowed_axes) + additional_arg_0_spec = _filter_spec_axes(get_padded_spec(arg_infos[11]), allowed_axes) + additional_arg_1_spec = _filter_spec_axes(get_padded_spec(arg_infos[12]), allowed_axes) grouped_dim_specs = ( lhs_first_dims_spec, @@ -1924,14 +1949,15 @@ def _parse_partition_specs( rhs_scale_spec = _strip_axis_from_spec(rhs_scale_spec, fsdp_axis) bias_spec = _strip_axis_from_spec(bias_spec, fsdp_axis) - reduce_axis = _common_axis(lhs_data_spec, rhs_data_spec) - if reduce_axis not in (gsr.dp_resource, gsr.fsdp_resource): - reduce_axis = None + reducible_axes = tuple( + axis for axis in (gsr.dp_resource, gsr.fsdp_resource) if axis is not None + ) + reduce_axis = _common_axis(lhs_data_spec, rhs_data_spec, reducible_axes) if reduce_axis is not None and gather_rhs_fsdp: reduce_axis = None if result_infos: - out_spec = get_padded_spec(result_infos[0]) + out_spec = _filter_spec_axes(get_padded_spec(result_infos[0]), allowed_axes) else: out_spec = (None,) * (len(out_shape) if out_shape is not None else 1) @@ -1998,8 +2024,7 @@ def partition( lhs_axis_boundary=lhs_axis_boundary, ) arg_shardings = tuple(NamedSharding(mesh, PartitionSpec(*spec)) for spec in arg_specs) - result_info = result_infos[0] if result_infos else None - out_sharding = (_partition_spec_from_result(mesh, result_info, out_spec),) + out_sharding = (NamedSharding(mesh, PartitionSpec(*out_spec)),) local_out_shape = _local_shape_from_spec(out_shape, out_spec, mesh) local_lhs_left_size, local_lhs_right_size = _local_2d_sizes_from_spec( arg_infos[0].shape, diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 761933b3f9..f74ca21a38 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -72,6 +72,40 @@ def _flat_data_spec(input_spec): return (_merge_axis_specs(input_spec),) +def _normalize_flatten_axis(flatten_axis, ndim): + return flatten_axis + ndim if flatten_axis < 0 else flatten_axis + + +def _contiguous_flat_input_spec(input_spec, flatten_axis): + flatten_axis = _normalize_flatten_axis(flatten_axis, len(input_spec)) + if flatten_axis <= 0 or len(input_spec) == 0: + return (None,) * len(input_spec) + return (input_spec[0], *((None,) * (len(input_spec) - 1))) + + +def _filter_axis_spec(axis_spec, allowed_axes): + if axis_spec is None: + return None + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + axes = tuple(axis for axis in axis_tuple if axis in allowed_axes) + if len(axes) == 0: + return None + return axes[0] if len(axes) == 1 else axes + + +def _filter_spec_axes(spec, allowed_axes): + return tuple(_filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) + + +def _supported_grouped_quantize_axes(mesh): + gsr = global_mesh_resource(validate=False) + return { + axis + for axis in (gsr.ep_resource, gsr.dp_resource, gsr.fsdp_resource) + if axis is not None and axis in mesh.axis_names + } + + def _axis_spec_size(axis_spec, mesh): axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) axis_size = 1 @@ -1292,10 +1326,11 @@ def impl( return rowwise_out, colwise_out, rowwise_scale_inv, colwise_scale_inv, updated_amax @staticmethod - def _parse_partition_specs(scaling_mode, q_layout, mesh, arg_infos): - del mesh - x_spec = get_padded_spec(arg_infos[0]) - group_spec = get_padded_spec(arg_infos[2]) + def _parse_partition_specs(scaling_mode, q_layout, flatten_axis, mesh, arg_infos): + allowed_axes = _supported_grouped_quantize_axes(mesh) + x_spec = _filter_spec_axes(get_padded_spec(arg_infos[0]), allowed_axes) + x_spec = _contiguous_flat_input_spec(x_spec, flatten_axis) + group_spec = _filter_spec_axes(get_padded_spec(arg_infos[2]), allowed_axes) if group_spec == (None,) and len(x_spec) > 0: group_spec = (x_spec[0],) flat_spec = _flat_data_spec(x_spec) @@ -1338,7 +1373,7 @@ def partition( result_infos, ): x_spec, group_spec, out_specs = GroupedQuantizePrimitive._parse_partition_specs( - scaling_mode, q_layout, mesh, arg_infos + scaling_mode, q_layout, flatten_axis, mesh, arg_infos ) local_out_shapes = ( tuple(_local_shape_from_spec(info.shape, spec, mesh) for info, spec in zip(result_infos, out_specs)) diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index 70151a44b7..35eed4d6cc 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -18,7 +18,7 @@ from . import cpp_extensions as tex from .cpp_extensions.amax import AmaxScope -from .sharding import global_mesh_resource, with_sharding_constraint +from .sharding import global_mesh_resource, get_mesh_axis_size, with_sharding_constraint from .quantize import ( ScaledTensor, QuantizerSet, @@ -60,6 +60,16 @@ def _is_manual_mesh_axis(mesh_axis): return mesh_axis is not None and mesh_axis in jax.sharding.get_abstract_mesh().manual_axes +def _kernel_non_contracting_axis_to_bias_axis(kernel_axis_idx, kernel_contracting_dims): + if kernel_axis_idx in kernel_contracting_dims: + return None + bias_axis_idx = 1 + for dim in range(1, kernel_axis_idx): + if dim not in kernel_contracting_dims: + bias_axis_idx += 1 + return bias_axis_idx + + def dense( x: jnp.ndarray, kernel: jnp.ndarray, @@ -426,9 +436,36 @@ def _grouped_dense_fwd_rule( ): use_bias = bias is not None - del kernel_fsdp_info - x_contracting_dims, k_contracting_dims = contracting_dims + local_kernel_shape = kernel.shape + kernel_was_gathered = False + bias_shape = bias.shape if use_bias else None + bias_fsdp_axis_idx = -1 + bias_was_gathered = False + + kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info + if ( + _is_manual_mesh_axis(kernel_fsdp_mesh_axis) + and 0 < kernel_fsdp_axis_idx < kernel.ndim + ): + kernel = _all_gather_kernel(kernel, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx) + kernel_was_gathered = True + + if use_bias and kernel_fsdp_axis_idx not in k_contracting_dims: + bias_fsdp_axis_idx = _kernel_non_contracting_axis_to_bias_axis( + kernel_fsdp_axis_idx, k_contracting_dims + ) + mesh_axis_size = get_mesh_axis_size(kernel_fsdp_mesh_axis) + if ( + bias_fsdp_axis_idx is not None + and 0 < bias_fsdp_axis_idx < bias.ndim + and mesh_axis_size > 1 + and bias.shape[bias_fsdp_axis_idx] * mesh_axis_size + == kernel.shape[kernel_fsdp_axis_idx] + ): + bias = _all_gather_kernel(bias, kernel_fsdp_mesh_axis, bias_fsdp_axis_idx) + bias_was_gathered = True + flatten_axis_x = -len(x_contracting_dims) flatten_axis_k = len(k_contracting_dims) - len(kernel.shape) + 1 # +1 for G axis @@ -484,10 +521,14 @@ def _grouped_dense_fwd_rule( else ctx_kernel ), x.shape, - kernel.shape, + local_kernel_shape, use_bias, quantizer_set, flatten_axis_k, + kernel_was_gathered, + bias_shape, + bias_fsdp_axis_idx, + bias_was_gathered, ) return output, ctx @@ -508,6 +549,10 @@ def _grouped_dense_bwd_rule( use_bias, quantizer_set, flatten_axis_k, + kernel_was_gathered, + bias_shape, + bias_fsdp_axis_idx, + bias_was_gathered, ) = ctx # The 1 in range is for excluding the group dimension (shall we use the hardcoded results below?) @@ -545,7 +590,7 @@ def _grouped_dense_bwd_rule( preferred_element_type=preferred_element_type, group_offset=group_offset, ) - if _is_manual_mesh_axis(kernel_fsdp_mesh_axis): + if _is_manual_mesh_axis(kernel_fsdp_mesh_axis) and not kernel_was_gathered: if kernel_fsdp_axis_idx in fwd_k_contracting_dims: dgrad_axis_idx = fwd_x_contracting_dims[ fwd_k_contracting_dims.index(kernel_fsdp_axis_idx) @@ -563,7 +608,11 @@ def _grouped_dense_bwd_rule( group_offset=group_offset, ) if _is_manual_mesh_axis(kernel_fsdp_mesh_axis): - if kernel_fsdp_axis_idx in fwd_k_contracting_dims: + if kernel_was_gathered: + wgrad = _psum_scatter_kernel( + wgrad, kernel_shape, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx + ) + elif kernel_fsdp_axis_idx in fwd_k_contracting_dims: wgrad = _psum_scatter_kernel( wgrad, kernel_shape, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx ) @@ -584,6 +633,14 @@ def _grouped_dense_bwd_rule( group_sizes_grad = None dbias = tex.grouped_dbias(grad, group_sizes) if use_bias else None + if ( + dbias is not None + and _is_manual_mesh_axis(kernel_fsdp_mesh_axis) + and bias_was_gathered + ): + dbias = _psum_scatter_kernel( + dbias, bias_shape, kernel_fsdp_mesh_axis, bias_fsdp_axis_idx + ) return dgrad, wgrad, group_sizes_grad, dbias, quantizer_set diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 16a10c860d..8dffb71196 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -133,8 +133,8 @@ def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec): """ A wrapper function to jax.lax.with_sharding_constraint 1. Does nothing if mesh is empty. - 2. If all mesh axes are manual axes, replaces pspec with all Nones. - 3. Otherwise, strips only the manual axes. + 2. Keeps only auto axes in pspec. + 3. Returns x unchanged if no auto axes remain. """ if pspec is None: return x @@ -143,22 +143,21 @@ def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec): if mesh.empty: return x - # We want to exclude the axes that already used by shard_map and shard_map - # only sets those in the abstract_mesh, not the physical one - manual_axis_names = get_abstract_mesh().manual_axes + # with_sharding_constraint can only refer to auto axes. Explicit axes are + # already fixed by the active mesh, and manual axes are managed by shard_map. + abstract_mesh = get_abstract_mesh() + auto_axis_names = set(abstract_mesh.auto_axes) # Multiple mesh axes can be mapped to a single shape axis, so we need to unpack and process tuples here too - def filter_manual_axes(name_or_tuple): + def filter_non_auto_axes(name_or_tuple): if isinstance(name_or_tuple, tuple): - out = tuple(n for n in name_or_tuple if n not in manual_axis_names) + out = tuple(n for n in name_or_tuple if n in auto_axis_names) if len(out) == 0: return None return out - if name_or_tuple in manual_axis_names: - return None - return name_or_tuple + return name_or_tuple if name_or_tuple in auto_axis_names else None - cleaned_axis_names = tuple(filter_manual_axes(name_or_tuple) for name_or_tuple in pspec) + cleaned_axis_names = tuple(filter_non_auto_axes(name_or_tuple) for name_or_tuple in pspec) if cleaned_axis_names == (None,) * len(cleaned_axis_names): return x @@ -366,7 +365,7 @@ def global_shard_guard(resource: MeshResource): _GLOBAL_MESH_RESOURCE = old_resources -def global_mesh_resource() -> MeshResource: +def global_mesh_resource(validate: bool = True) -> MeshResource: """Get the current global mesh resource configuration. Returns: @@ -377,7 +376,8 @@ def global_mesh_resource() -> MeshResource: " context. If you are not using multiple GPUs, you can use an empty MeshResource by" " wrapping your program in 'with global_shard_guard(MeshResource()):'" ) - _validate_mesh_resource_configuration(_GLOBAL_MESH_RESOURCE) + if validate: + _validate_mesh_resource_configuration(_GLOBAL_MESH_RESOURCE) return _GLOBAL_MESH_RESOURCE From 1bd6b54db95fbd7b6f02d855d6854c94a70adf82 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 2 Jun 2026 13:10:53 -0700 Subject: [PATCH 05/44] Add warnings Signed-off-by: Jeremy Berchtold --- tests/jax/test_grouped_gemm_partitioning.py | 136 +++++++++--------- transformer_engine/jax/cpp_extensions/gemm.py | 109 ++++++++++---- .../jax/cpp_extensions/quantization.py | 34 ++++- 3 files changed, 182 insertions(+), 97 deletions(-) diff --git a/tests/jax/test_grouped_gemm_partitioning.py b/tests/jax/test_grouped_gemm_partitioning.py index 831e245c0e..e62e33a616 100644 --- a/tests/jax/test_grouped_gemm_partitioning.py +++ b/tests/jax/test_grouped_gemm_partitioning.py @@ -132,20 +132,21 @@ def test_grouped_quantize_strips_unsupported_axes_and_gathers_hidden_axes(): with jax.set_mesh(mesh), global_shard_guard( MeshResource(dp_resource="dp", tp_resource="tp", fsdp_resource="fsdp", ep_resource="expert") ): - _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - mesh, - ( - _arg_info(mesh, (8, 128, 128), ("expert", "dp", ("fsdp", "tp"))), - _arg_info(mesh, (8,), (("expert", "tp"),)), - _arg_info(mesh, (8,), (("expert", "tp"),)), - ), - (), - ) + with pytest.warns(RuntimeWarning, match="Grouped quantize.*tp"): + _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 128), ("expert", "dp", ("fsdp", "tp"))), + _arg_info(mesh, (8,), (("expert", "tp"),)), + _arg_info(mesh, (8,), (("expert", "tp"),)), + ), + (), + ) assert tuple(arg_shardings[0].spec) == ("expert", None, None) assert tuple(arg_shardings[1].spec) == ("expert",) @@ -223,25 +224,26 @@ def test_grouped_gemm_strips_unsupported_axes_preserves_dp_and_gathers_rhs_fsdp( with jax.set_mesh(mesh), global_shard_guard( MeshResource(dp_resource="dp", tp_resource="tp", fsdp_resource="fsdp", ep_resource="expert") ): - _, _, out_sharding, arg_shardings = GroupedGemmPrimitive.partition( - False, - False, - ScalingMode.NO_SCALING.value, - jnp.bfloat16, - False, - False, - False, - 1, - 1, - (1, 128, 64), - 128, - 64, - 128, - 64, - mesh, - arg_infos, - result_infos, - ) + with pytest.warns(RuntimeWarning, match="Grouped GEMM.*tp"): + _, _, out_sharding, arg_shardings = GroupedGemmPrimitive.partition( + False, + False, + ScalingMode.NO_SCALING.value, + jnp.bfloat16, + False, + False, + False, + 1, + 1, + (1, 128, 64), + 128, + 64, + 128, + 64, + mesh, + arg_infos, + result_infos, + ) assert tuple(arg_shardings[0].spec) == ("dp",) assert tuple(arg_shardings[2].spec) == ("expert",) @@ -293,20 +295,21 @@ def test_grouped_partitioning_strips_arbitrary_unsupported_axis(): mesh_resource = MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="expert") with jax.set_mesh(mesh), global_shard_guard(mesh_resource): - _, _, quantize_out_shardings, quantize_arg_shardings = GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - mesh, - ( - _arg_info(mesh, (8, 128, 128), ("expert", "myaxis123", ("dp", "fsdp"))), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - ), - (), - ) + with pytest.warns(RuntimeWarning, match="Grouped quantize.*myaxis123"): + _, _, quantize_out_shardings, quantize_arg_shardings = GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 128), ("expert", "myaxis123", ("dp", "fsdp"))), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + ), + (), + ) gemm_arg_infos = ( _arg_info(mesh, (8192,), (("dp", "myaxis123"),)), @@ -324,25 +327,26 @@ def test_grouped_partitioning_strips_arbitrary_unsupported_axis(): _arg_info(mesh, (0,), (("myaxis123",),)), ) gemm_result_infos = (_arg_info(mesh, (1, 128, 64), ("expert", "myaxis123", None)),) - _, _, gemm_out_sharding, gemm_arg_shardings = GroupedGemmPrimitive.partition( - False, - False, - ScalingMode.NO_SCALING.value, - jnp.bfloat16, - False, - False, - False, - 1, - 1, - (1, 128, 64), - 128, - 64, - 128, - 64, - mesh, - gemm_arg_infos, - gemm_result_infos, - ) + with pytest.warns(RuntimeWarning, match="Grouped GEMM.*myaxis123"): + _, _, gemm_out_sharding, gemm_arg_shardings = GroupedGemmPrimitive.partition( + False, + False, + ScalingMode.NO_SCALING.value, + jnp.bfloat16, + False, + False, + False, + 1, + 1, + (1, 128, 64), + 128, + 64, + 128, + 64, + mesh, + gemm_arg_infos, + gemm_result_infos, + ) assert tuple(quantize_arg_shardings[0].spec) == ("expert", None, None) assert tuple(quantize_arg_shardings[1].spec) == ("expert",) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 61521873d2..ce364597c6 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -252,6 +252,30 @@ def _filter_spec_axes(spec, allowed_axes): return tuple(_filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) +def _spec_axes(spec): + axes = [] + for axis_spec in spec: + if axis_spec is None: + continue + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + for axis in axis_tuple: + if axis is not None and axis not in axes: + axes.append(axis) + return axes + + +def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): + ignored_axes = tuple(axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec)) + if ignored_axes: + warnings.warn( + "Grouped GEMM custom partitioning will ignore/replicate sharding " + f"axes {ignored_axes} from {arg_name}; only DP/FSDP/EP grouped " + "partitioning axes are preserved.", + RuntimeWarning, + stacklevel=3, + ) + + def _supported_grouped_gemm_axes(mesh): gsr = global_mesh_resource(validate=False) return { @@ -1885,20 +1909,21 @@ def _parse_partition_specs( fsdp_axis = gsr.fsdp_resource allowed_axes = _supported_grouped_gemm_axes(mesh) - lhs_data_spec = _filter_spec_axes(get_padded_spec(arg_infos[0]), allowed_axes) - lhs_scale_spec = _filter_spec_axes(get_padded_spec(arg_infos[1]), allowed_axes) - rhs_data_spec = _filter_spec_axes(get_padded_spec(arg_infos[2]), allowed_axes) - rhs_scale_spec = _filter_spec_axes(get_padded_spec(arg_infos[3]), allowed_axes) - bias_spec = _filter_spec_axes(get_padded_spec(arg_infos[4]), allowed_axes) - - lhs_first_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[5]), allowed_axes) - lhs_last_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[6]), allowed_axes) - rhs_first_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[7]), allowed_axes) - rhs_last_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[8]), allowed_axes) - out_first_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[9]), allowed_axes) - out_last_dims_spec = _filter_spec_axes(get_padded_spec(arg_infos[10]), allowed_axes) - additional_arg_0_spec = _filter_spec_axes(get_padded_spec(arg_infos[11]), allowed_axes) - additional_arg_1_spec = _filter_spec_axes(get_padded_spec(arg_infos[12]), allowed_axes) + original_arg_specs = tuple(get_padded_spec(arg_info) for arg_info in arg_infos) + lhs_data_spec = _filter_spec_axes(original_arg_specs[0], allowed_axes) + lhs_scale_spec = _filter_spec_axes(original_arg_specs[1], allowed_axes) + rhs_data_spec = _filter_spec_axes(original_arg_specs[2], allowed_axes) + rhs_scale_spec = _filter_spec_axes(original_arg_specs[3], allowed_axes) + bias_spec = _filter_spec_axes(original_arg_specs[4], allowed_axes) + + lhs_first_dims_spec = _filter_spec_axes(original_arg_specs[5], allowed_axes) + lhs_last_dims_spec = _filter_spec_axes(original_arg_specs[6], allowed_axes) + rhs_first_dims_spec = _filter_spec_axes(original_arg_specs[7], allowed_axes) + rhs_last_dims_spec = _filter_spec_axes(original_arg_specs[8], allowed_axes) + out_first_dims_spec = _filter_spec_axes(original_arg_specs[9], allowed_axes) + out_last_dims_spec = _filter_spec_axes(original_arg_specs[10], allowed_axes) + additional_arg_0_spec = _filter_spec_axes(original_arg_specs[11], allowed_axes) + additional_arg_1_spec = _filter_spec_axes(original_arg_specs[12], allowed_axes) grouped_dim_specs = ( lhs_first_dims_spec, @@ -1957,8 +1982,10 @@ def _parse_partition_specs( reduce_axis = None if result_infos: - out_spec = _filter_spec_axes(get_padded_spec(result_infos[0]), allowed_axes) + original_out_spec = get_padded_spec(result_infos[0]) + out_spec = _filter_spec_axes(original_out_spec, allowed_axes) else: + original_out_spec = None out_spec = (None,) * (len(out_shape) if out_shape is not None else 1) if rhs_is_ragged and lhs_is_trans is not None and lhs_axis_boundary is not None: @@ -1975,22 +2002,46 @@ def _parse_partition_specs( ) lhs_data_spec = tuple(lhs_data_spec) - return ( + final_arg_specs = ( + lhs_data_spec, + lhs_scale_spec, + rhs_data_spec, + rhs_scale_spec, + bias_spec, + lhs_first_dims_spec, + lhs_last_dims_spec, + rhs_first_dims_spec, + rhs_last_dims_spec, + out_first_dims_spec, + out_last_dims_spec, + additional_arg_0_spec, + additional_arg_1_spec, + ) + for arg_name, original_spec, partition_spec in zip( ( - lhs_data_spec, - lhs_scale_spec, - rhs_data_spec, - rhs_scale_spec, - bias_spec, - lhs_first_dims_spec, - lhs_last_dims_spec, - rhs_first_dims_spec, - rhs_last_dims_spec, - out_first_dims_spec, - out_last_dims_spec, - additional_arg_0_spec, - additional_arg_1_spec, + "lhs_data", + "lhs_scale_inv", + "rhs_data", + "rhs_scale_inv", + "bias", + "lhs_first_dims", + "lhs_last_dims", + "rhs_first_dims", + "rhs_last_dims", + "out_first_dims", + "out_last_dims", + "additional_arg_0", + "additional_arg_1", ), + original_arg_specs, + final_arg_specs, + ): + _warn_if_axes_ignored(arg_name, original_spec, partition_spec) + if original_out_spec is not None: + _warn_if_axes_ignored("output", original_out_spec, out_spec) + + return ( + final_arg_specs, out_spec, reduce_axis, ) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index f74ca21a38..31884fe354 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -6,6 +6,7 @@ from functools import reduce from typing import Tuple, Optional, Union import math +import warnings import jax @@ -97,6 +98,30 @@ def _filter_spec_axes(spec, allowed_axes): return tuple(_filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) +def _spec_axes(spec): + axes = [] + for axis_spec in spec: + if axis_spec is None: + continue + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + for axis in axis_tuple: + if axis is not None and axis not in axes: + axes.append(axis) + return axes + + +def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): + ignored_axes = tuple(axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec)) + if ignored_axes: + warnings.warn( + "Grouped quantize custom partitioning will ignore/replicate sharding " + f"axes {ignored_axes} from {arg_name}; only supported packed grouped " + "data axes are preserved.", + RuntimeWarning, + stacklevel=3, + ) + + def _supported_grouped_quantize_axes(mesh): gsr = global_mesh_resource(validate=False) return { @@ -1328,11 +1353,16 @@ def impl( @staticmethod def _parse_partition_specs(scaling_mode, q_layout, flatten_axis, mesh, arg_infos): allowed_axes = _supported_grouped_quantize_axes(mesh) - x_spec = _filter_spec_axes(get_padded_spec(arg_infos[0]), allowed_axes) + original_x_spec = get_padded_spec(arg_infos[0]) + x_spec = _filter_spec_axes(original_x_spec, allowed_axes) x_spec = _contiguous_flat_input_spec(x_spec, flatten_axis) - group_spec = _filter_spec_axes(get_padded_spec(arg_infos[2]), allowed_axes) + _warn_if_axes_ignored("x", original_x_spec, x_spec) + + original_group_spec = get_padded_spec(arg_infos[2]) + group_spec = _filter_spec_axes(original_group_spec, allowed_axes) if group_spec == (None,) and len(x_spec) > 0: group_spec = (x_spec[0],) + _warn_if_axes_ignored("group_sizes", original_group_spec, group_spec) flat_spec = _flat_data_spec(x_spec) replicated_spec = (None,) From 3c30c9b3912f0a4b092874e9a5916e64c48a75ec Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 2 Jun 2026 13:25:35 -0700 Subject: [PATCH 06/44] Remove kernel_fsdp_info Signed-off-by: Jeremy Berchtold --- tests/jax/test_grouped_gemm_partitioning.py | 31 +++- ..._multi_process_distributed_grouped_gemm.py | 2 - .../jax/cpp_extensions/quantization.py | 2 +- transformer_engine/jax/dense.py | 136 +----------------- 4 files changed, 31 insertions(+), 140 deletions(-) diff --git a/tests/jax/test_grouped_gemm_partitioning.py b/tests/jax/test_grouped_gemm_partitioning.py index e62e33a616..9fa34a8f9f 100644 --- a/tests/jax/test_grouped_gemm_partitioning.py +++ b/tests/jax/test_grouped_gemm_partitioning.py @@ -127,6 +127,30 @@ def test_grouped_quantize_mxfp8_colwise_specs_gather_hidden_axis(): assert _normalize_spec(specs[4]) == ("expert",) +def test_grouped_quantize_preserves_row_side_fsdp_for_kernel(): + mesh = _mesh() + with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): + _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 64), ("expert", "fsdp", None)), + _arg_info(mesh, (8,), ("expert",)), + _arg_info(mesh, (8,), ("expert",)), + ), + (), + ) + + assert tuple(arg_shardings[0].spec) == ("expert", "fsdp", None) + specs = tuple(tuple(sharding.spec) for sharding in out_shardings) + assert _normalize_spec(specs[0]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[2]) == (("expert", "fsdp"),) + + def test_grouped_quantize_strips_unsupported_axes_and_gathers_hidden_axes(): mesh = _mesh_with_dp_tp() with jax.set_mesh(mesh), global_shard_guard( @@ -148,13 +172,13 @@ def test_grouped_quantize_strips_unsupported_axes_and_gathers_hidden_axes(): (), ) - assert tuple(arg_shardings[0].spec) == ("expert", None, None) + assert tuple(arg_shardings[0].spec) == ("expert", "dp", None) assert tuple(arg_shardings[1].spec) == ("expert",) assert tuple(arg_shardings[2].spec) == ("expert",) out_specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(out_specs[0]) == ("expert",) - assert _normalize_spec(out_specs[2]) == ("expert",) + assert _normalize_spec(out_specs[0]) == (("expert", "dp"),) + assert _normalize_spec(out_specs[2]) == (("expert", "dp"),) assert _normalize_spec(out_specs[4]) == ("expert",) for spec in (*out_specs, *(tuple(sharding.spec) for sharding in arg_shardings)): assert not _spec_contains_axis(spec, "tp") @@ -455,7 +479,6 @@ def apply(x, w): group_sizes, contracting_dims=((1,), (1,)), quantizer_set=quantizer_set, - kernel_fsdp_info=("fsdp", 1), ) out, vjp_fn = jax.vjp(apply, x, w) diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py index ce52126dea..cb7ea2cd60 100644 --- a/tests/jax/test_multi_process_distributed_grouped_gemm.py +++ b/tests/jax/test_multi_process_distributed_grouped_gemm.py @@ -97,7 +97,6 @@ def sharded_group_gemm(x, w, b): n_groups, bias=b, quantizer_set=quantizer_set, - kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis), ) output = output.reshape(*x.shape[:-1], -1) return output @@ -221,7 +220,6 @@ def apply(x, w): group_sizes, contracting_dims=((1,), (1,)), quantizer_set=quantizer_set, - kernel_fsdp_info=(FSDP_AXIS_NAME, 1), ) out, vjp_fn = jax.vjp(apply, x, w) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 31884fe354..34b2cadfab 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -81,7 +81,7 @@ def _contiguous_flat_input_spec(input_spec, flatten_axis): flatten_axis = _normalize_flatten_axis(flatten_axis, len(input_spec)) if flatten_axis <= 0 or len(input_spec) == 0: return (None,) * len(input_spec) - return (input_spec[0], *((None,) * (len(input_spec) - 1))) + return (*input_spec[:flatten_axis], *((None,) * (len(input_spec) - flatten_axis))) def _filter_axis_spec(axis_spec, allowed_axes): diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index 35eed4d6cc..13ee446fbb 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -14,11 +14,9 @@ import warnings import jax import jax.numpy as jnp -from jax.sharding import PartitionSpec from . import cpp_extensions as tex from .cpp_extensions.amax import AmaxScope -from .sharding import global_mesh_resource, get_mesh_axis_size, with_sharding_constraint from .quantize import ( ScaledTensor, QuantizerSet, @@ -28,48 +26,6 @@ ) -def _all_gather_kernel(kernel, mesh_axis, axis_idx): - assert mesh_axis is not None - assert 0 < axis_idx < len(kernel.shape) - - # TODO(Ming Hunag): Add a condition branch for with/without shmap. - kernel_shape = kernel.shape - kernel_whole_shape = (*kernel_shape[:axis_idx], -1, *kernel_shape[axis_idx + 1 :]) - global_kernel = jax.lax.all_gather(kernel, mesh_axis, axis=axis_idx) - global_kernel = global_kernel.reshape(*kernel_whole_shape) - return global_kernel - - -def _psum_scatter_kernel(kernel, scattered_kernel_shape, mesh_axis, axis_idx): - assert mesh_axis is not None - assert 0 < axis_idx < len(scattered_kernel_shape) - - # TODO(Ming Hunag): Add a condition branch for with/without shmap. - kernel = kernel.reshape( - *scattered_kernel_shape[:axis_idx], - -1, - scattered_kernel_shape[axis_idx], - *scattered_kernel_shape[axis_idx + 1 :], - ) - kernel = jax.lax.psum_scatter(kernel, mesh_axis, scatter_dimension=axis_idx) - kernel = kernel.reshape(scattered_kernel_shape) - return kernel - - -def _is_manual_mesh_axis(mesh_axis): - return mesh_axis is not None and mesh_axis in jax.sharding.get_abstract_mesh().manual_axes - - -def _kernel_non_contracting_axis_to_bias_axis(kernel_axis_idx, kernel_contracting_dims): - if kernel_axis_idx in kernel_contracting_dims: - return None - bias_axis_idx = 1 - for dim in range(1, kernel_axis_idx): - if dim not in kernel_contracting_dims: - bias_axis_idx += 1 - return bias_axis_idx - - def dense( x: jnp.ndarray, kernel: jnp.ndarray, @@ -341,7 +297,6 @@ def grouped_dense( preferred_element_type: jnp.dtype = None, group_offset: jnp.array = None, quantizer_set: QuantizerSet = noop_quantizer_set, - kernel_fsdp_info: Tuple[str, int] = (None, -1), ): """ Perform grouped dense (linear) layer transformation with optional quantization. @@ -357,10 +312,6 @@ def grouped_dense( preferred_element_type: Preferred data type for the output tensor group_offset: 1D array containing offsets for each group (not yet implemented) quantizer_set: Set of quantizers for FP8 quantization of the input and output - kernel_fsdp_info: A tuple containing FSDP-related information for a weight matrix - represented in the format (str, int). The first element is the - FSDP mesh axis, and the second element is the dimension along - which the weight is sharded. Returns: A jnp.ndarray containing the result of the grouped linear operation @@ -387,14 +338,13 @@ def grouped_dense( preferred_element_type, group_offset, quantizer_set, - kernel_fsdp_info, ) if restore_leading_ep_axis: output = output.reshape(1, *output.shape) return output -@partial(jax.custom_vjp, nondiff_argnums=(3, 5, 6, 7, 9)) +@partial(jax.custom_vjp, nondiff_argnums=(3, 5, 6, 7)) def _grouped_dense( x, kernel, @@ -405,7 +355,6 @@ def _grouped_dense( preferred_element_type, group_offset, quantizer_set, - kernel_fsdp_info, ): output, _ = _grouped_dense_fwd_rule( x, @@ -417,7 +366,6 @@ def _grouped_dense( preferred_element_type, group_offset, quantizer_set, - kernel_fsdp_info, ) return output @@ -432,39 +380,10 @@ def _grouped_dense_fwd_rule( preferred_element_type, group_offset, quantizer_set, - kernel_fsdp_info, ): use_bias = bias is not None x_contracting_dims, k_contracting_dims = contracting_dims - local_kernel_shape = kernel.shape - kernel_was_gathered = False - bias_shape = bias.shape if use_bias else None - bias_fsdp_axis_idx = -1 - bias_was_gathered = False - - kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info - if ( - _is_manual_mesh_axis(kernel_fsdp_mesh_axis) - and 0 < kernel_fsdp_axis_idx < kernel.ndim - ): - kernel = _all_gather_kernel(kernel, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx) - kernel_was_gathered = True - - if use_bias and kernel_fsdp_axis_idx not in k_contracting_dims: - bias_fsdp_axis_idx = _kernel_non_contracting_axis_to_bias_axis( - kernel_fsdp_axis_idx, k_contracting_dims - ) - mesh_axis_size = get_mesh_axis_size(kernel_fsdp_mesh_axis) - if ( - bias_fsdp_axis_idx is not None - and 0 < bias_fsdp_axis_idx < bias.ndim - and mesh_axis_size > 1 - and bias.shape[bias_fsdp_axis_idx] * mesh_axis_size - == kernel.shape[kernel_fsdp_axis_idx] - ): - bias = _all_gather_kernel(bias, kernel_fsdp_mesh_axis, bias_fsdp_axis_idx) - bias_was_gathered = True flatten_axis_x = -len(x_contracting_dims) flatten_axis_k = len(k_contracting_dims) - len(kernel.shape) + 1 # +1 for G axis @@ -521,23 +440,17 @@ def _grouped_dense_fwd_rule( else ctx_kernel ), x.shape, - local_kernel_shape, + kernel.shape, use_bias, quantizer_set, flatten_axis_k, - kernel_was_gathered, - bias_shape, - bias_fsdp_axis_idx, - bias_was_gathered, ) return output, ctx def _grouped_dense_bwd_rule( - contracting_dims, precision, preferred_element_type, group_offset, kernel_fsdp_info, ctx, grad + contracting_dims, precision, preferred_element_type, group_offset, ctx, grad ): - kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info - fwd_x_contracting_dims, fwd_k_contracting_dims = contracting_dims ( @@ -549,10 +462,6 @@ def _grouped_dense_bwd_rule( use_bias, quantizer_set, flatten_axis_k, - kernel_was_gathered, - bias_shape, - bias_fsdp_axis_idx, - bias_was_gathered, ) = ctx # The 1 in range is for excluding the group dimension (shall we use the hardcoded results below?) @@ -590,14 +499,6 @@ def _grouped_dense_bwd_rule( preferred_element_type=preferred_element_type, group_offset=group_offset, ) - if _is_manual_mesh_axis(kernel_fsdp_mesh_axis) and not kernel_was_gathered: - if kernel_fsdp_axis_idx in fwd_k_contracting_dims: - dgrad_axis_idx = fwd_x_contracting_dims[ - fwd_k_contracting_dims.index(kernel_fsdp_axis_idx) - ] - dgrad = _all_gather_kernel(dgrad, kernel_fsdp_mesh_axis, dgrad_axis_idx) - else: - dgrad = jax.lax.psum(dgrad, kernel_fsdp_mesh_axis) wgrad = tex.grouped_gemm( wgrad_x_T, @@ -607,40 +508,9 @@ def _grouped_dense_bwd_rule( preferred_element_type=preferred_element_type, group_offset=group_offset, ) - if _is_manual_mesh_axis(kernel_fsdp_mesh_axis): - if kernel_was_gathered: - wgrad = _psum_scatter_kernel( - wgrad, kernel_shape, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx - ) - elif kernel_fsdp_axis_idx in fwd_k_contracting_dims: - wgrad = _psum_scatter_kernel( - wgrad, kernel_shape, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx - ) - else: - wgrad = jax.lax.psum(wgrad, kernel_fsdp_mesh_axis) - if kernel_fsdp_mesh_axis is not None: - wgrad_spec = [None] * len(kernel_shape) - ep_resource = None - try: - ep_resource = global_mesh_resource().ep_resource - except AssertionError: - pass - if len(wgrad_spec) > 0: - wgrad_spec[0] = ep_resource - if 0 <= kernel_fsdp_axis_idx < len(wgrad_spec): - wgrad_spec[kernel_fsdp_axis_idx] = kernel_fsdp_mesh_axis - wgrad = with_sharding_constraint(wgrad, PartitionSpec(*wgrad_spec)) group_sizes_grad = None dbias = tex.grouped_dbias(grad, group_sizes) if use_bias else None - if ( - dbias is not None - and _is_manual_mesh_axis(kernel_fsdp_mesh_axis) - and bias_was_gathered - ): - dbias = _psum_scatter_kernel( - dbias, bias_shape, kernel_fsdp_mesh_axis, bias_fsdp_axis_idx - ) return dgrad, wgrad, group_sizes_grad, dbias, quantizer_set From 273e066e5be8341b329988defe331c3e5f24e2c3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 2 Jun 2026 20:26:43 +0000 Subject: [PATCH 07/44] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_grouped_gemm_partitioning.py | 28 ++++++++++--------- ..._multi_process_distributed_grouped_gemm.py | 17 +++++------ transformer_engine/jax/cpp_extensions/gemm.py | 10 +++++-- .../jax/cpp_extensions/quantization.py | 9 ++++-- 4 files changed, 37 insertions(+), 27 deletions(-) diff --git a/tests/jax/test_grouped_gemm_partitioning.py b/tests/jax/test_grouped_gemm_partitioning.py index 9fa34a8f9f..be2487a8df 100644 --- a/tests/jax/test_grouped_gemm_partitioning.py +++ b/tests/jax/test_grouped_gemm_partitioning.py @@ -320,19 +320,21 @@ def test_grouped_partitioning_strips_arbitrary_unsupported_axis(): with jax.set_mesh(mesh), global_shard_guard(mesh_resource): with pytest.warns(RuntimeWarning, match="Grouped quantize.*myaxis123"): - _, _, quantize_out_shardings, quantize_arg_shardings = GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - mesh, - ( - _arg_info(mesh, (8, 128, 128), ("expert", "myaxis123", ("dp", "fsdp"))), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - ), - (), + _, _, quantize_out_shardings, quantize_arg_shardings = ( + GroupedQuantizePrimitive.partition( + jnp.float8_e4m3fn, + ScalingMode.MXFP8_1D_SCALING.value, + QuantizeLayout.ROWWISE, + -1, + jnp.float8_e8m0fnu, + mesh, + ( + _arg_info(mesh, (8, 128, 128), ("expert", "myaxis123", ("dp", "fsdp"))), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + _arg_info(mesh, (8,), (("expert", "myaxis123"),)), + ), + (), + ) ) gemm_arg_infos = ( diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py index cb7ea2cd60..b2978922e1 100644 --- a/tests/jax/test_multi_process_distributed_grouped_gemm.py +++ b/tests/jax/test_multi_process_distributed_grouped_gemm.py @@ -59,18 +59,15 @@ def init_data(): x_key = jax.random.PRNGKey(0) w_key = jax.random.PRNGKey(1) b_key = jax.random.PRNGKey(2) - x = ( - jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16) - * jnp.asarray(0.01, dtype=jnp.bfloat16) + x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16) * jnp.asarray( + 0.01, dtype=jnp.bfloat16 ) - w = ( - jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16) - * jnp.asarray(0.01, dtype=jnp.bfloat16) - ) - b = ( - jax.random.normal(b_key, shape=(N_GROUP, w_shape[-1]), dtype=jnp.bfloat16) - * jnp.asarray(0.01, dtype=jnp.bfloat16) + w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16) * jnp.asarray( + 0.01, dtype=jnp.bfloat16 ) + b = jax.random.normal( + b_key, shape=(N_GROUP, w_shape[-1]), dtype=jnp.bfloat16 + ) * jnp.asarray(0.01, dtype=jnp.bfloat16) return x, w, w, b, b def test_func(outter_x, outter_w, outter_b): diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index ce364597c6..9bd0a1700e 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -265,7 +265,9 @@ def _spec_axes(spec): def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): - ignored_axes = tuple(axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec)) + ignored_axes = tuple( + axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec) + ) if ignored_axes: warnings.warn( "Grouped GEMM custom partitioning will ignore/replicate sharding " @@ -1945,7 +1947,11 @@ def _parse_partition_specs( rhs_is_ragged = arg_infos[7].size > 0 or arg_infos[8].size > 0 ep_axis = gsr.ep_resource - if ep_axis is not None and not rhs_is_ragged and _spec_contains_axis(active_group_spec, ep_axis): + if ( + ep_axis is not None + and not rhs_is_ragged + and _spec_contains_axis(active_group_spec, ep_axis) + ): if len(rhs_data_spec) > 0 and not _spec_contains_axis(rhs_data_spec, ep_axis): rhs_data_spec = ( _merge_axis_spec(rhs_data_spec[0], ep_axis), diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 34b2cadfab..828bdb6067 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -111,7 +111,9 @@ def _spec_axes(spec): def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): - ignored_axes = tuple(axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec)) + ignored_axes = tuple( + axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec) + ) if ignored_axes: warnings.warn( "Grouped quantize custom partitioning will ignore/replicate sharding " @@ -1406,7 +1408,10 @@ def partition( scaling_mode, q_layout, flatten_axis, mesh, arg_infos ) local_out_shapes = ( - tuple(_local_shape_from_spec(info.shape, spec, mesh) for info, spec in zip(result_infos, out_specs)) + tuple( + _local_shape_from_spec(info.shape, spec, mesh) + for info, spec in zip(result_infos, out_specs) + ) if result_infos else (None,) * len(out_specs) ) From fe906fda11d93486e4e9fbd35d648c844a51865a Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 2 Jun 2026 13:42:33 -0700 Subject: [PATCH 08/44] Remove unnecessary reshape Signed-off-by: Jeremy Berchtold --- transformer_engine/jax/dense.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index 13ee446fbb..b60810f5eb 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -321,13 +321,6 @@ def grouped_dense( kernel_contracting_dims = tex.sanitize_dims(kernel.ndim, kernel_contracting_dims) contracting_dims = (x_contracting_dims, kernel_contracting_dims) - restore_leading_ep_axis = False - if x.ndim == 3 and x.shape[0] == 1: - if x_contracting_dims == (x.ndim - 1,): - restore_leading_ep_axis = True - x = x.reshape(*x.shape[1:]) - contracting_dims = ((x.ndim - 1,), kernel_contracting_dims) - output = _grouped_dense( x, kernel, @@ -339,8 +332,6 @@ def grouped_dense( group_offset, quantizer_set, ) - if restore_leading_ep_axis: - output = output.reshape(1, *output.shape) return output From e76d20aff1ae112f89fd1094ad4e7806320495d4 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 2 Jun 2026 14:06:13 -0700 Subject: [PATCH 09/44] Remove multi-process grouped GEMM tests Signed-off-by: Jeremy Berchtold --- qa/L1_jax_distributed_unittest/test.sh | 1 - ..._multi_process_distributed_grouped_gemm.py | 263 ------------------ 2 files changed, 264 deletions(-) delete mode 100644 tests/jax/test_multi_process_distributed_grouped_gemm.py diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 4f92d1c783..031bb72995 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -38,7 +38,6 @@ XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false" python3 -m pyt python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py" # TODO(Phuong): add this test back after it is verified -# SCRIPT_NAME=$TE_PATH/tests/jax/test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh || test_fail "test_multi_process_distributed_grouped_gemm.py" if [ $RET -ne 0 ]; then echo "Error: some sub-tests failed: $FAILED_CASES" diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py deleted file mode 100644 index b2978922e1..0000000000 --- a/tests/jax/test_multi_process_distributed_grouped_gemm.py +++ /dev/null @@ -1,263 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -from functools import partial - -import jax -import jax.numpy as jnp -import jax.experimental.multihost_utils as jem -import numpy as np -from jax.experimental import shard_map -from jax.sharding import NamedSharding, PartitionSpec - -from transformer_engine.jax.dense import grouped_dense as te_grouped_dense -from transformer_engine.jax.quantize import ( - QuantizerFactory, - ScalingMode, -) -from transformer_engine.jax.sharding import MeshResource, global_shard_guard - -from utils import assert_allclose, dtype_tols - - -N_GROUP = 8 -EP_AXIS_NAME = "ep" -FSDP_AXIS_NAME = "fsdp" -MESH_AXIS_NAME = FSDP_AXIS_NAME - - -def _mxfp8_grouped_quantizer_set(n_groups): - return QuantizerFactory.create_set( - scaling_mode=ScalingMode.MXFP8_1D_SCALING, - fwd_dtype=jnp.float8_e4m3fn, - bwd_dtype=jnp.float8_e4m3fn, - is_2x2x=True, - n_groups=n_groups, - ) - - -def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis): - assert kernel_fsdp_axis in [1, 2] - x_shape, w_shape = data_shapes - - x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None)) - w_sharding = ( - NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME)) - if kernel_fsdp_axis == 2 - else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None)) - ) - b_sharding = ( - NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME)) - if kernel_fsdp_axis == 2 - else NamedSharding(mesh, PartitionSpec(None, None)) - ) - w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None)) - b_no_sharding = NamedSharding(mesh, PartitionSpec(None, None)) - - def init_data(): - x_key = jax.random.PRNGKey(0) - w_key = jax.random.PRNGKey(1) - b_key = jax.random.PRNGKey(2) - x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16) * jnp.asarray( - 0.01, dtype=jnp.bfloat16 - ) - w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16) * jnp.asarray( - 0.01, dtype=jnp.bfloat16 - ) - b = jax.random.normal( - b_key, shape=(N_GROUP, w_shape[-1]), dtype=jnp.bfloat16 - ) * jnp.asarray(0.01, dtype=jnp.bfloat16) - return x, w, w, b, b - - def test_func(outter_x, outter_w, outter_b): - in_specs = (x_sharding.spec, w_sharding.spec, b_sharding.spec) - out_specs = x_sharding.spec - - @partial( - shard_map.shard_map, - mesh=mesh, - in_specs=in_specs, - out_specs=out_specs, - check_rep=False, - ) - def sharded_group_gemm(x, w, b): - group_size = x.shape[0] - x_reshaped = x.reshape(-1, x.shape[-1]) - n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) - - quantizer_set = _mxfp8_grouped_quantizer_set(group_size) - - output = te_grouped_dense( - x_reshaped, - w, - n_groups, - bias=b, - quantizer_set=quantizer_set, - ) - output = output.reshape(*x.shape[:-1], -1) - return output - - def run(x, w, b): - output = sharded_group_gemm(x, w, b) - return output - - output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_b) - dx, dw, db = vjp_fn(output) - return output, dx, dw, db - - def ref_func(outter_x, outter_w, outter_b): - - in_specs = (x_sharding.spec, w_no_sharding.spec, b_no_sharding.spec) - out_specs = x_sharding.spec - - @partial( - shard_map.shard_map, - mesh=mesh, - in_specs=in_specs, - out_specs=out_specs, - check_rep=False, - ) - def sharded_group_gemm(x, w, b): - group_size = x.shape[0] - x_reshaped = x.reshape(-1, x.shape[-1]) - n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) - - quantizer_set = _mxfp8_grouped_quantizer_set(group_size) - output = te_grouped_dense( - x_reshaped, - w, - n_groups, - bias=b, - quantizer_set=quantizer_set, - ) - output = output.reshape(*x.shape[:-1], -1) - return output - - def run(x, w, b): - output = sharded_group_gemm(x, w, b) - return output - - output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_b) - dx, dw, db = vjp_fn(output) - return output, dx, dw, db - - init_func = jax.jit( - init_data, - out_shardings=(x_sharding, w_sharding, w_no_sharding, b_sharding, b_no_sharding), - ) - x, w, w_global, b, b_global = init_func() - - o_sharding = x_sharding - test_func_jitted = jax.jit( - test_func, - in_shardings=(x_sharding, w_sharding, b_sharding), - out_shardings=(o_sharding, x_sharding, w_sharding, b_sharding), - ) - ref_func_jitted = jax.jit( - ref_func, - in_shardings=(x_sharding, w_no_sharding, b_no_sharding), - out_shardings=(o_sharding, x_sharding, w_no_sharding, b_no_sharding), - ) - - out, dx, dw, db = test_func_jitted(x, w, b) - ref_out, ref_dx, ref_dw, ref_db = ref_func_jitted(x, w_global, b_global) - - # Avoid creating a host scalar JAX array under the multi-process mesh in dtype_tols. - e4m3_tols = dtype_tols(jnp.float8_e4m3fn, rtol=0.25, atol=0.25) - - out, ref_out = jem.process_allgather((out, ref_out), tiled=True) - dx, ref_dx = jem.process_allgather((dx, ref_dx), tiled=True) - dw, ref_dw = jem.process_allgather((dw, ref_dw), tiled=True) - db, ref_db = jem.process_allgather((db, ref_db), tiled=True) - - assert_allclose(out, ref_out, **e4m3_tols) - assert_allclose(dx, ref_dx, **e4m3_tols) - assert_allclose(dw, ref_dw, **e4m3_tols) - assert_allclose(db, ref_db, **e4m3_tols) - - -def run_grouped_dense_mxfp8_ep_fsdp_outside_shard_map(): - n_groups = 4 - group_tokens = 128 - hidden = 256 - out_hidden = 128 - x_shape = (n_groups * group_tokens, hidden) - w_shape = (n_groups, hidden, out_hidden) - quantizer_set = _mxfp8_grouped_quantizer_set(n_groups) - - x_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME, None)) - w_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME, FSDP_AXIS_NAME, None)) - group_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME)) - out_sharding = NamedSharding(mesh, PartitionSpec(EP_AXIS_NAME, None)) - - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS_NAME, fsdp_resource=FSDP_AXIS_NAME) - ): - x = jax.device_put( - jax.random.normal(jax.random.PRNGKey(20), x_shape, dtype=jnp.bfloat16) - * jnp.asarray(0.01, dtype=jnp.bfloat16), - x_sharding, - ) - w = jax.device_put( - jax.random.normal(jax.random.PRNGKey(21), w_shape, dtype=jnp.bfloat16) - * jnp.asarray(0.01, dtype=jnp.bfloat16), - w_sharding, - ) - group_sizes = jax.device_put( - jnp.full((n_groups,), group_tokens, dtype=jnp.int32), - group_sharding, - ) - - def apply_with_vjp(x, w, group_sizes): - def apply(x, w): - return te_grouped_dense( - x, - w, - group_sizes, - contracting_dims=((1,), (1,)), - quantizer_set=quantizer_set, - ) - - out, vjp_fn = jax.vjp(apply, x, w) - dx, dw = vjp_fn(out) - return out, dx, dw - - out, dx, dw = jax.jit( - apply_with_vjp, - in_shardings=(x_sharding, w_sharding, group_sharding), - out_shardings=(out_sharding, x_sharding, w_sharding), - )(x, w, group_sizes) - out, dx, dw = jax.block_until_ready((out, dx, dw)) - - assert tuple(out.sharding.spec) == (EP_AXIS_NAME, None) - assert tuple(dx.sharding.spec) == (EP_AXIS_NAME, None) - assert tuple(dw.sharding.spec) == (EP_AXIS_NAME, FSDP_AXIS_NAME, None) - for value in (out, dx, dw): - local_value = np.asarray(jax.device_get(value.addressable_data(0))) - assert np.all(np.isfinite(local_value)) - assert np.any(local_value != 0.0) - - -if __name__ == "__main__": - import sys - - coord_addr = sys.argv[1] - proc_id = int(sys.argv[2]) - num_procs = int(sys.argv[3]) - - jax.distributed.initialize( - coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id - ) - - mesh = jax.make_mesh((num_procs,), (FSDP_AXIS_NAME,)) - - with mesh: - data_shapes = [((4, 16, 128, 7168), (7168, 2048))] - for data_shape in data_shapes: - for kernel_fsdp_axis in [1, 2]: - test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis) - - if num_procs == 4: - mesh = jax.make_mesh((2, 2), (EP_AXIS_NAME, FSDP_AXIS_NAME)) - run_grouped_dense_mxfp8_ep_fsdp_outside_shard_map() From 1aa1e827597535cbd13dfafc9da618cfe421d0fe Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 2 Jun 2026 14:24:09 -0700 Subject: [PATCH 10/44] Refactor helpers into sharding.py Signed-off-by: Jeremy Berchtold --- transformer_engine/jax/cpp_extensions/gemm.py | 225 ++++-------------- .../jax/cpp_extensions/quantization.py | 82 +------ transformer_engine/jax/sharding.py | 150 ++++++++++++ 3 files changed, 205 insertions(+), 252 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 9bd0a1700e..655aa6e3f9 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -49,7 +49,16 @@ ) from .misc import get_padded_spec, is_all_reduce_in_float32, get_min_device_compute_capability from ..sharding import ( + common_spec_axis, + filter_spec_axes, global_mesh_resource, + local_2d_sizes_from_spec, + local_shape_from_spec, + merge_axis_specs, + spec_axes, + spec_contains_axis, + strip_axis_from_spec, + supported_grouped_partition_axes, tpsp_axis_size, dp_or_fsdp_axis_size, ) @@ -211,62 +220,9 @@ def _get_nvfp4_tensor_scale_inv(amax): return amax / (DATA_DTYPE_MAX * SCALE_DTYPE_MAX) -def _axis_spec_contains(axis_spec, axis): - if axis is None or axis_spec is None: - return False - if isinstance(axis_spec, tuple): - return axis in axis_spec - return axis_spec == axis - - -def _spec_contains_axis(spec, axis): - return any(_axis_spec_contains(axis_spec, axis) for axis_spec in spec) - - -def _strip_axis_from_axis_spec(axis_spec, axis): - if axis is None or axis_spec is None: - return axis_spec - if isinstance(axis_spec, tuple): - stripped = tuple(a for a in axis_spec if a != axis) - if len(stripped) == 0: - return None - return stripped[0] if len(stripped) == 1 else stripped - return None if axis_spec == axis else axis_spec - - -def _strip_axis_from_spec(spec, axis): - return tuple(_strip_axis_from_axis_spec(axis_spec, axis) for axis_spec in spec) - - -def _filter_axis_spec(axis_spec, allowed_axes): - if axis_spec is None: - return None - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - axes = tuple(axis for axis in axis_tuple if axis in allowed_axes) - if len(axes) == 0: - return None - return axes[0] if len(axes) == 1 else axes - - -def _filter_spec_axes(spec, allowed_axes): - return tuple(_filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) - - -def _spec_axes(spec): - axes = [] - for axis_spec in spec: - if axis_spec is None: - continue - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - for axis in axis_tuple: - if axis is not None and axis not in axes: - axes.append(axis) - return axes - - def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): ignored_axes = tuple( - axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec) + axis for axis in spec_axes(original_spec) if axis not in spec_axes(partition_spec) ) if ignored_axes: warnings.warn( @@ -278,99 +234,6 @@ def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): ) -def _supported_grouped_gemm_axes(mesh): - gsr = global_mesh_resource(validate=False) - return { - axis - for axis in (gsr.ep_resource, gsr.dp_resource, gsr.fsdp_resource) - if axis is not None and axis in mesh.axis_names - } - - -def _common_axis(spec_a, spec_b, allowed_axes=None): - axes = [] - for spec in (spec_a, spec_b): - for axis_spec in spec: - if axis_spec is None: - continue - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - for axis in axis_tuple: - if axis is not None and axis not in axes: - axes.append(axis) - for axis in axes: - if allowed_axes is not None and axis not in allowed_axes: - continue - if _spec_contains_axis(spec_a, axis) and _spec_contains_axis(spec_b, axis): - return axis - return None - - -def _merge_axis_spec(axis_spec_a, axis_spec_b): - axes = [] - for axis_spec in (axis_spec_a, axis_spec_b): - if axis_spec is None: - continue - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - for axis in axis_tuple: - if axis is not None and axis not in axes: - axes.append(axis) - if len(axes) == 0: - return None - return axes[0] if len(axes) == 1 else tuple(axes) - - -def _partition_spec_from_result(mesh, result_info, fallback_spec): - if result_info is not None and result_info.sharding is not None: - return result_info.sharding - return NamedSharding(mesh, PartitionSpec(*fallback_spec)) - - -def _local_shape_from_spec(global_shape, spec, mesh): - local_shape = [] - for dim, axis_spec in zip(global_shape, spec): - axis_size = _axis_spec_size(axis_spec, mesh) - local_shape.append(dim // axis_size) - return tuple(local_shape) - - -def _axis_spec_size(axis_spec, mesh): - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - axis_size = 1 - for axis in axis_tuple: - if axis is not None: - axis_size *= mesh.shape[axis] - return axis_size - - -def _spec_size(spec, mesh): - axis_size = 1 - for axis_spec in spec: - axis_size *= _axis_spec_size(axis_spec, mesh) - return axis_size - - -def _local_2d_sizes_from_spec(shape, spec, axis_boundary, left_size, right_size, mesh): - if len(shape) == len(spec) and len(shape) > 1: - local_shape = _local_shape_from_spec(shape, spec, mesh) - return ( - math.prod(local_shape[:axis_boundary]), - math.prod(local_shape[axis_boundary:]), - ) - - spec_size = _spec_size(spec, mesh) - if spec_size == 1: - return left_size, right_size - if left_size % spec_size == 0: - return left_size // spec_size, right_size - if right_size % spec_size == 0: - return left_size, right_size // spec_size - raise ValueError( - "Cannot derive local grouped GEMM 2D sizes from sharding spec. " - f"shape={shape}, spec={spec}, axis_boundary={axis_boundary}, " - f"left_size={left_size}, right_size={right_size}, spec_size={spec_size}" - ) - - def collective_gemm_bootstrap( num_total_devices, num_devices_per_process, @@ -1909,23 +1772,23 @@ def _parse_partition_specs( ): gsr = global_mesh_resource(validate=False) fsdp_axis = gsr.fsdp_resource - allowed_axes = _supported_grouped_gemm_axes(mesh) + allowed_axes = supported_grouped_partition_axes(mesh) original_arg_specs = tuple(get_padded_spec(arg_info) for arg_info in arg_infos) - lhs_data_spec = _filter_spec_axes(original_arg_specs[0], allowed_axes) - lhs_scale_spec = _filter_spec_axes(original_arg_specs[1], allowed_axes) - rhs_data_spec = _filter_spec_axes(original_arg_specs[2], allowed_axes) - rhs_scale_spec = _filter_spec_axes(original_arg_specs[3], allowed_axes) - bias_spec = _filter_spec_axes(original_arg_specs[4], allowed_axes) - - lhs_first_dims_spec = _filter_spec_axes(original_arg_specs[5], allowed_axes) - lhs_last_dims_spec = _filter_spec_axes(original_arg_specs[6], allowed_axes) - rhs_first_dims_spec = _filter_spec_axes(original_arg_specs[7], allowed_axes) - rhs_last_dims_spec = _filter_spec_axes(original_arg_specs[8], allowed_axes) - out_first_dims_spec = _filter_spec_axes(original_arg_specs[9], allowed_axes) - out_last_dims_spec = _filter_spec_axes(original_arg_specs[10], allowed_axes) - additional_arg_0_spec = _filter_spec_axes(original_arg_specs[11], allowed_axes) - additional_arg_1_spec = _filter_spec_axes(original_arg_specs[12], allowed_axes) + lhs_data_spec = filter_spec_axes(original_arg_specs[0], allowed_axes) + lhs_scale_spec = filter_spec_axes(original_arg_specs[1], allowed_axes) + rhs_data_spec = filter_spec_axes(original_arg_specs[2], allowed_axes) + rhs_scale_spec = filter_spec_axes(original_arg_specs[3], allowed_axes) + bias_spec = filter_spec_axes(original_arg_specs[4], allowed_axes) + + lhs_first_dims_spec = filter_spec_axes(original_arg_specs[5], allowed_axes) + lhs_last_dims_spec = filter_spec_axes(original_arg_specs[6], allowed_axes) + rhs_first_dims_spec = filter_spec_axes(original_arg_specs[7], allowed_axes) + rhs_last_dims_spec = filter_spec_axes(original_arg_specs[8], allowed_axes) + out_first_dims_spec = filter_spec_axes(original_arg_specs[9], allowed_axes) + out_last_dims_spec = filter_spec_axes(original_arg_specs[10], allowed_axes) + additional_arg_0_spec = filter_spec_axes(original_arg_specs[11], allowed_axes) + additional_arg_1_spec = filter_spec_axes(original_arg_specs[12], allowed_axes) grouped_dim_specs = ( lhs_first_dims_spec, @@ -1950,46 +1813,46 @@ def _parse_partition_specs( if ( ep_axis is not None and not rhs_is_ragged - and _spec_contains_axis(active_group_spec, ep_axis) + and spec_contains_axis(active_group_spec, ep_axis) ): - if len(rhs_data_spec) > 0 and not _spec_contains_axis(rhs_data_spec, ep_axis): + if len(rhs_data_spec) > 0 and not spec_contains_axis(rhs_data_spec, ep_axis): rhs_data_spec = ( - _merge_axis_spec(rhs_data_spec[0], ep_axis), + merge_axis_specs(rhs_data_spec[0], ep_axis), *rhs_data_spec[1:], ) - if len(rhs_scale_spec) > 0 and not _spec_contains_axis(rhs_scale_spec, ep_axis): + if len(rhs_scale_spec) > 0 and not spec_contains_axis(rhs_scale_spec, ep_axis): rhs_scale_spec = ( - _merge_axis_spec(rhs_scale_spec[0], ep_axis), + merge_axis_specs(rhs_scale_spec[0], ep_axis), *rhs_scale_spec[1:], ) - if len(bias_spec) > 0 and not _spec_contains_axis(bias_spec, ep_axis): - bias_spec = (_merge_axis_spec(bias_spec[0], ep_axis), *bias_spec[1:]) + if len(bias_spec) > 0 and not spec_contains_axis(bias_spec, ep_axis): + bias_spec = (merge_axis_specs(bias_spec[0], ep_axis), *bias_spec[1:]) gather_rhs_fsdp = ( fsdp_axis is not None and not rhs_is_ragged and ( - _spec_contains_axis(rhs_data_spec, fsdp_axis) - or _spec_contains_axis(rhs_scale_spec, fsdp_axis) - or _spec_contains_axis(bias_spec, fsdp_axis) + spec_contains_axis(rhs_data_spec, fsdp_axis) + or spec_contains_axis(rhs_scale_spec, fsdp_axis) + or spec_contains_axis(bias_spec, fsdp_axis) ) ) if gather_rhs_fsdp: - rhs_data_spec = _strip_axis_from_spec(rhs_data_spec, fsdp_axis) - rhs_scale_spec = _strip_axis_from_spec(rhs_scale_spec, fsdp_axis) - bias_spec = _strip_axis_from_spec(bias_spec, fsdp_axis) + rhs_data_spec = strip_axis_from_spec(rhs_data_spec, fsdp_axis) + rhs_scale_spec = strip_axis_from_spec(rhs_scale_spec, fsdp_axis) + bias_spec = strip_axis_from_spec(bias_spec, fsdp_axis) reducible_axes = tuple( axis for axis in (gsr.dp_resource, gsr.fsdp_resource) if axis is not None ) - reduce_axis = _common_axis(lhs_data_spec, rhs_data_spec, reducible_axes) + reduce_axis = common_spec_axis(lhs_data_spec, rhs_data_spec, reducible_axes) if reduce_axis is not None and gather_rhs_fsdp: reduce_axis = None if result_infos: original_out_spec = get_padded_spec(result_infos[0]) - out_spec = _filter_spec_axes(original_out_spec, allowed_axes) + out_spec = filter_spec_axes(original_out_spec, allowed_axes) else: original_out_spec = None out_spec = (None,) * (len(out_shape) if out_shape is not None else 1) @@ -2003,7 +1866,7 @@ def _parse_partition_specs( lhs_data_spec = list(lhs_data_spec) for out_idx, lhs_dim in enumerate(lhs_non_contracting_dims, start=1): if out_idx < len(out_spec): - lhs_data_spec[lhs_dim] = _merge_axis_spec( + lhs_data_spec[lhs_dim] = merge_axis_specs( lhs_data_spec[lhs_dim], out_spec[out_idx] ) lhs_data_spec = tuple(lhs_data_spec) @@ -2082,8 +1945,8 @@ def partition( ) arg_shardings = tuple(NamedSharding(mesh, PartitionSpec(*spec)) for spec in arg_specs) out_sharding = (NamedSharding(mesh, PartitionSpec(*out_spec)),) - local_out_shape = _local_shape_from_spec(out_shape, out_spec, mesh) - local_lhs_left_size, local_lhs_right_size = _local_2d_sizes_from_spec( + local_out_shape = local_shape_from_spec(out_shape, out_spec, mesh) + local_lhs_left_size, local_lhs_right_size = local_2d_sizes_from_spec( arg_infos[0].shape, arg_specs[0], lhs_axis_boundary, @@ -2091,7 +1954,7 @@ def partition( lhs_right_size, mesh, ) - local_rhs_left_size, local_rhs_right_size = _local_2d_sizes_from_spec( + local_rhs_left_size, local_rhs_right_size = local_2d_sizes_from_spec( arg_infos[2].shape, arg_specs[2], rhs_axis_boundary, diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 828bdb6067..9266ab08f0 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -32,9 +32,14 @@ from ..sharding import ( all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp, + filter_spec_axes, get_num_devices_in_mesh, global_mesh_resource, lax_paral_op, + local_shape_from_spec, + merge_axis_specs, + spec_axes, + supported_grouped_partition_axes, ) from ..quantize import ( ScaledTensor2x, @@ -55,22 +60,8 @@ __all__ = ["quantize", "quantize_dbias", "grouped_quantize", "grouped_dbias"] -def _merge_axis_specs(axis_specs): - axes = [] - for axis_spec in axis_specs: - if axis_spec is None: - continue - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - for axis in axis_tuple: - if axis is not None and axis not in axes: - axes.append(axis) - if len(axes) == 0: - return None - return axes[0] if len(axes) == 1 else tuple(axes) - - def _flat_data_spec(input_spec): - return (_merge_axis_specs(input_spec),) + return (merge_axis_specs(*input_spec),) def _normalize_flatten_axis(flatten_axis, ndim): @@ -84,35 +75,9 @@ def _contiguous_flat_input_spec(input_spec, flatten_axis): return (*input_spec[:flatten_axis], *((None,) * (len(input_spec) - flatten_axis))) -def _filter_axis_spec(axis_spec, allowed_axes): - if axis_spec is None: - return None - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - axes = tuple(axis for axis in axis_tuple if axis in allowed_axes) - if len(axes) == 0: - return None - return axes[0] if len(axes) == 1 else axes - - -def _filter_spec_axes(spec, allowed_axes): - return tuple(_filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) - - -def _spec_axes(spec): - axes = [] - for axis_spec in spec: - if axis_spec is None: - continue - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - for axis in axis_tuple: - if axis is not None and axis not in axes: - axes.append(axis) - return axes - - def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): ignored_axes = tuple( - axis for axis in _spec_axes(original_spec) if axis not in _spec_axes(partition_spec) + axis for axis in spec_axes(original_spec) if axis not in spec_axes(partition_spec) ) if ignored_axes: warnings.warn( @@ -124,31 +89,6 @@ def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): ) -def _supported_grouped_quantize_axes(mesh): - gsr = global_mesh_resource(validate=False) - return { - axis - for axis in (gsr.ep_resource, gsr.dp_resource, gsr.fsdp_resource) - if axis is not None and axis in mesh.axis_names - } - - -def _axis_spec_size(axis_spec, mesh): - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - axis_size = 1 - for axis in axis_tuple: - if axis is not None: - axis_size *= mesh.shape[axis] - return axis_size - - -def _local_shape_from_spec(global_shape, spec, mesh): - local_shape = [] - for dim, axis_spec in zip(global_shape, spec): - local_shape.append(dim // _axis_spec_size(axis_spec, mesh)) - return tuple(local_shape) - - def _pad_or_slice_to_shape(x, target_shape): if target_shape is None or x.shape == target_shape: return x @@ -1354,14 +1294,14 @@ def impl( @staticmethod def _parse_partition_specs(scaling_mode, q_layout, flatten_axis, mesh, arg_infos): - allowed_axes = _supported_grouped_quantize_axes(mesh) + allowed_axes = supported_grouped_partition_axes(mesh) original_x_spec = get_padded_spec(arg_infos[0]) - x_spec = _filter_spec_axes(original_x_spec, allowed_axes) + x_spec = filter_spec_axes(original_x_spec, allowed_axes) x_spec = _contiguous_flat_input_spec(x_spec, flatten_axis) _warn_if_axes_ignored("x", original_x_spec, x_spec) original_group_spec = get_padded_spec(arg_infos[2]) - group_spec = _filter_spec_axes(original_group_spec, allowed_axes) + group_spec = filter_spec_axes(original_group_spec, allowed_axes) if group_spec == (None,) and len(x_spec) > 0: group_spec = (x_spec[0],) _warn_if_axes_ignored("group_sizes", original_group_spec, group_spec) @@ -1409,7 +1349,7 @@ def partition( ) local_out_shapes = ( tuple( - _local_shape_from_spec(info.shape, spec, mesh) + local_shape_from_spec(info.shape, spec, mesh) for info, spec in zip(result_infos, out_specs) ) if result_infos diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 8dffb71196..86ddf72ff7 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -12,6 +12,7 @@ from contextlib import contextmanager from dataclasses import dataclass from typing import Callable, Optional +import math import warnings import jax @@ -233,6 +234,155 @@ def get_padded_spec(spec, ndim): return spec + (None,) * (ndim - len(spec)) +def spec_axes(spec): + """Return unique non-None mesh axes used by a PartitionSpec-like tuple.""" + axes = [] + for axis_spec in spec: + if axis_spec is None: + continue + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + for axis in axis_tuple: + if axis is not None and axis not in axes: + axes.append(axis) + return axes + + +def axis_spec_contains(axis_spec, axis): + """Return whether one dimension's axis spec contains a mesh axis.""" + if axis is None or axis_spec is None: + return False + if isinstance(axis_spec, tuple): + return axis in axis_spec + return axis_spec == axis + + +def spec_contains_axis(spec, axis): + """Return whether a PartitionSpec-like tuple contains a mesh axis.""" + return any(axis_spec_contains(axis_spec, axis) for axis_spec in spec) + + +def filter_axis_spec(axis_spec, allowed_axes): + """Keep only allowed axes in one dimension's axis spec.""" + if axis_spec is None: + return None + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + axes = tuple(axis for axis in axis_tuple if axis in allowed_axes) + if len(axes) == 0: + return None + return axes[0] if len(axes) == 1 else axes + + +def filter_spec_axes(spec, allowed_axes): + """Keep only allowed axes in a PartitionSpec-like tuple.""" + return tuple(filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) + + +def supported_grouped_partition_axes(mesh): + """Return mesh axes supported by grouped quantize/GEMM custom partitioning.""" + gsr = global_mesh_resource(validate=False) + return { + axis + for axis in (gsr.ep_resource, gsr.dp_resource, gsr.fsdp_resource) + if axis is not None and axis in mesh.axis_names + } + + +def strip_axis_from_axis_spec(axis_spec, axis): + """Remove one mesh axis from one dimension's axis spec.""" + if axis is None or axis_spec is None: + return axis_spec + if isinstance(axis_spec, tuple): + stripped = tuple(a for a in axis_spec if a != axis) + if len(stripped) == 0: + return None + return stripped[0] if len(stripped) == 1 else stripped + return None if axis_spec == axis else axis_spec + + +def strip_axis_from_spec(spec, axis): + """Remove one mesh axis from a PartitionSpec-like tuple.""" + return tuple(strip_axis_from_axis_spec(axis_spec, axis) for axis_spec in spec) + + +def merge_axis_specs(*axis_specs): + """Merge dimension axis specs while preserving first-seen axis order.""" + axes = [] + for axis_spec in axis_specs: + if axis_spec is None: + continue + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + for axis in axis_tuple: + if axis is not None and axis not in axes: + axes.append(axis) + if len(axes) == 0: + return None + return axes[0] if len(axes) == 1 else tuple(axes) + + +def common_spec_axis(spec_a, spec_b, allowed_axes=None): + """Return the first mesh axis that appears in both specs.""" + axes = [] + for spec in (spec_a, spec_b): + for axis in spec_axes(spec): + if axis not in axes: + axes.append(axis) + for axis in axes: + if allowed_axes is not None and axis not in allowed_axes: + continue + if spec_contains_axis(spec_a, axis) and spec_contains_axis(spec_b, axis): + return axis + return None + + +def axis_spec_size(axis_spec, mesh): + """Return the device count represented by one dimension's axis spec.""" + axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) + axis_size = 1 + for axis in axis_tuple: + if axis is not None: + axis_size *= mesh.shape[axis] + return axis_size + + +def spec_size(spec, mesh): + """Return the total device count represented by a PartitionSpec-like tuple.""" + axis_size = 1 + for axis_spec in spec: + axis_size *= axis_spec_size(axis_spec, mesh) + return axis_size + + +def local_shape_from_spec(global_shape, spec, mesh): + """Derive a local shape from a global shape and PartitionSpec-like tuple.""" + local_shape = [] + for dim, axis_spec in zip(global_shape, spec): + local_shape.append(dim // axis_spec_size(axis_spec, mesh)) + return tuple(local_shape) + + +def local_2d_sizes_from_spec(shape, spec, axis_boundary, left_size, right_size, mesh): + """Derive local collapsed 2D dimensions from a global shape and sharding spec.""" + if len(shape) == len(spec) and len(shape) > 1: + local_shape = local_shape_from_spec(shape, spec, mesh) + return ( + math.prod(local_shape[:axis_boundary]), + math.prod(local_shape[axis_boundary:]), + ) + + size = spec_size(spec, mesh) + if size == 1: + return left_size, right_size + if left_size % size == 0: + return left_size // size, right_size + if right_size % size == 0: + return left_size, right_size // size + raise ValueError( + "Cannot derive local 2D sizes from sharding spec. " + f"shape={shape}, spec={spec}, axis_boundary={axis_boundary}, " + f"left_size={left_size}, right_size={right_size}, spec_size={size}" + ) + + def lax_paral_op( x: jnp.array, ops: Callable, mesh_resource: str, mesh: jax.sharding.Mesh, **kwargs ): From 6fd04b573446e27e0ac9d1d7666f24a423a971bf Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 2 Jun 2026 15:00:31 -0700 Subject: [PATCH 11/44] Rename distributed grouped GEMM tests Signed-off-by: Jeremy Berchtold --- qa/L1_jax_distributed_unittest/test.sh | 4 ++-- ..._gemm_partitioning.py => test_distributed_grouped_gemm.py} | 0 2 files changed, 2 insertions(+), 2 deletions(-) rename tests/jax/{test_grouped_gemm_partitioning.py => test_distributed_grouped_gemm.py} (100%) diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 031bb72995..9d481d4332 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -25,6 +25,8 @@ export XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_dense.xml $TE_PATH/tests/jax/test_distributed_dense.py || test_fail "test_distributed_dense.py" +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_grouped_gemm.xml $TE_PATH/tests/jax/test_distributed_grouped_gemm.py || test_fail "test_distributed_grouped_gemm.py" + python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_helper.xml $TE_PATH/tests/jax/test_distributed_helper.py || test_fail "test_distributed_helper.py" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_layernorm.xml $TE_PATH/tests/jax/test_distributed_layernorm.py || test_fail "test_distributed_layernorm.py" @@ -37,8 +39,6 @@ XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false" python3 -m pyt python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py" -# TODO(Phuong): add this test back after it is verified - if [ $RET -ne 0 ]; then echo "Error: some sub-tests failed: $FAILED_CASES" exit 1 diff --git a/tests/jax/test_grouped_gemm_partitioning.py b/tests/jax/test_distributed_grouped_gemm.py similarity index 100% rename from tests/jax/test_grouped_gemm_partitioning.py rename to tests/jax/test_distributed_grouped_gemm.py From 81d45affb0d768c2d6fecb2fbbc4ed2c06dcb2be Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 16 Jul 2026 16:10:28 -0700 Subject: [PATCH 12/44] Run JAX TE-EP MoE FFN without shard_map --- transformer_engine/jax/cpp_extensions/gemm.py | 12 + transformer_engine/jax/moe.py | 387 ++++++++---------- 2 files changed, 179 insertions(+), 220 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index f4964f9d5a..54e4b510cf 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -1863,9 +1863,16 @@ def _parse_partition_specs( if len(bias_spec) > 0 and not spec_contains_axis(bias_spec, ep_axis): bias_spec = (merge_axis_specs(bias_spec[0], ep_axis), *bias_spec[1:]) + # A compound leading group dimension may use FSDP as the outer + # data-parallel axis (for example MoE groups ordered + # (fsdp, ep, local_expert)). In that case FSDP describes distinct + # groups, not a sharded RHS contracting dimension, and must remain + # on the group axis. Otherwise gather the FSDP-sharded RHS as usual. + fsdp_is_group_axis = spec_contains_axis(active_group_spec, fsdp_axis) gather_rhs_fsdp = ( fsdp_axis is not None and not rhs_is_ragged + and not fsdp_is_group_axis and ( spec_contains_axis(rhs_data_spec, fsdp_axis) or spec_contains_axis(rhs_scale_spec, fsdp_axis) @@ -1882,6 +1889,11 @@ def _parse_partition_specs( axis for axis in (gsr.dp_resource, gsr.fsdp_resource) if axis is not None ) reduce_axis = common_spec_axis(lhs_data_spec, rhs_data_spec, reducible_axes) + # A common DP/FSDP group axis represents independent groups, not a + # partitioned contraction. Keep reductions for genuinely sharded + # contracting dimensions (including grouped wgrad), but not here. + if reduce_axis is not None and spec_contains_axis(active_group_spec, reduce_axis): + reduce_axis = None if reduce_axis is not None and gather_rhs_fsdp: reduce_axis = None diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 887e005de6..85a0ed94f6 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -20,19 +20,17 @@ ``((*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. +* 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. 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). +them on requires recipe-aware residual handling for ``ScaledTensor`` +leaves. ``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). """ from functools import partial @@ -212,14 +210,14 @@ class _Ctx: # ============================================================================= -# Per-shard FFN body (runs inside shard_map) +# Global-view FFN body # ============================================================================= -def _ffn_fwd_per_shard( - recv_tokens_local: jnp.ndarray, - recv_topk_weights_local: jnp.ndarray, - token_counts_local: jnp.ndarray, +def _ffn_fwd_global( + recv_tokens: jnp.ndarray, + recv_topk_weights: jnp.ndarray, + token_counts: jnp.ndarray, wi_0: jnp.ndarray, wi_1: jnp.ndarray, wo: jnp.ndarray, @@ -227,32 +225,55 @@ def _ffn_fwd_per_shard( wi_1_bias: Optional[jnp.ndarray], wo_bias: Optional[jnp.ndarray], *, + dp_size: int, + num_ep: int, num_local_experts: int, activation_type: str, apply_topk_weights_early: bool, + flat_token_sharding: NamedSharding, + flat_group_sharding: NamedSharding, + grouped_weight_sharding: NamedSharding, + grouped_bias_sharding: NamedSharding, ): - """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 FFN on global EP-dispatch buffers. + + Grouped-operation custom partitioning lowers the global operands to + the per-device problem. ``token_counts`` from ``tex.ep_prepare`` is + passed through as the dynamic grouped-GEMM group sizes, so cuBLAS + skips both 0-token experts and dispatch-buffer over-allocation. """ - 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) + hidden = recv_tokens.shape[-1] + sorted_x = recv_tokens.reshape(-1, hidden) + recv_w_flat = recv_topk_weights.reshape(-1) + group_sizes = token_counts.reshape(-1).astype(jnp.int32) + sorted_x = jax.lax.with_sharding_constraint(sorted_x, flat_token_sharding) + recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) + group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) wi_0 = wi_0.astype(sorted_x.dtype) wi_1 = wi_1.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) + # Dispatch groups flatten in (dp, ep, local_expert) order. Broadcast + # each global expert parameter set over the outer DP dimension before + # flattening so the grouped weights have exactly the same ordering. + num_groups = dp_size * num_ep * num_local_experts + + def _broadcast_experts(value, sharding): + value = jnp.broadcast_to( + value.reshape(1, num_ep, num_local_experts, *value.shape[1:]), + (dp_size, num_ep, num_local_experts, *value.shape[1:]), + ).reshape(num_groups, *value.shape[1:]) + return jax.lax.with_sharding_constraint(value, sharding) + + wi_0 = _broadcast_experts(wi_0, grouped_weight_sharding) + wi_1 = _broadcast_experts(wi_1, grouped_weight_sharding) + wo = _broadcast_experts(wo, grouped_weight_sharding) + if wi_0_bias is not None: + wi_0_bias = _broadcast_experts(wi_0_bias, grouped_bias_sharding) + wi_1_bias = _broadcast_experts(wi_1_bias, grouped_bias_sharding) + wo_bias = _broadcast_experts(wo_bias, grouped_bias_sharding) + # 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 @@ -263,7 +284,7 @@ def _ffn_fwd_per_shard( ) q_set = noop_quantizer_set - casted_sorted_x = tex.grouped_quantize(sorted_x, q_set.x, local_group_sizes, flatten_axis=-1) + casted_sorted_x = tex.grouped_quantize(sorted_x, q_set.x, 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), @@ -271,6 +292,7 @@ def _ffn_fwd_per_shard( contracting_dims=((1,), (1,)), bias=wi_combined_bias, ) + combined_out = jax.lax.with_sharding_constraint(combined_out, flat_token_sharding) 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) @@ -282,6 +304,7 @@ 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 + intermediate = jax.lax.with_sharding_constraint(intermediate, flat_token_sharding) if apply_topk_weights_early: # Fold the per-token combine weights into the FFN intermediate; @@ -295,9 +318,7 @@ def _ffn_fwd_per_shard( 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 - ) + casted_intermediate = tex.grouped_quantize(intermediate, q_set.x, group_sizes, 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), @@ -305,14 +326,12 @@ def _ffn_fwd_per_shard( contracting_dims=((1,), (1,)), bias=wo_bias, ) + expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, flat_token_sharding) 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) + expert_outputs_3d = expert_outputs.reshape(*recv_tokens.shape[:-1], expert_outputs.shape[-1]) + group_sizes_nd = group_sizes.reshape(token_counts.shape) residuals = ( casted_sorted_x_lhs_trans, casted_wi_rhs_trans, @@ -320,13 +339,13 @@ def _ffn_fwd_per_shard( up_proj_out, casted_intermediate_lhs_trans, casted_wo_rhs_trans, - local_group_sizes_3d, + group_sizes_nd, ) return expert_outputs_3d, residuals -def _ffn_bwd_per_shard( - d_expert_outputs_local: jnp.ndarray, +def _ffn_bwd_global( + d_expert_outputs: jnp.ndarray, casted_sorted_x_lhs_trans, casted_wi_rhs_trans, gate_proj_out: jnp.ndarray, @@ -334,29 +353,36 @@ def _ffn_bwd_per_shard( casted_intermediate_lhs_trans, casted_wo_rhs_trans, local_group_sizes: jnp.ndarray, - recv_topk_weights_local: jnp.ndarray, + recv_topk_weights: jnp.ndarray, *, activation_type: str, apply_topk_weights_early: bool, has_bias: bool, + flat_token_sharding: NamedSharding, + flat_group_sharding: NamedSharding, + grouped_weight_sharding: NamedSharding, + grouped_bias_sharding: NamedSharding, ): - """Per-shard FFN backward. + """Run the FFN backward on global residuals. - Mirrors :func:`_ffn_fwd_per_shard`. Returns - ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], + Mirrors :func:`_ffn_fwd_global`. Returns + ``(d_sorted_x [num_procs, recv_pr, H], d_recv_w [num_procs, 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) + group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) + d_eo_2d = d_expert_outputs.reshape(-1, d_expert_outputs.shape[-1]) + recv_w_flat = recv_topk_weights.reshape(-1) + d_eo_2d = jax.lax.with_sharding_constraint(d_eo_2d, flat_token_sharding) + recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) + group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) 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] + wgrad_group_active = (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) + casted_d_eo = tex.grouped_quantize(d_eo_2d, q_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( @@ -364,13 +390,17 @@ def _ffn_bwd_per_shard( casted_wo_rhs_trans, contracting_dims=((1,), (2,)), ) + d_intermediate = jax.lax.with_sharding_constraint(d_intermediate, flat_token_sharding) d_wo = tex.grouped_gemm( casted_intermediate_lhs_trans, _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 = jax.lax.with_sharding_constraint(d_wo, grouped_weight_sharding) + d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None + if has_bias: + d_wo_bias = jax.lax.with_sharding_constraint(d_wo_bias, grouped_bias_sharding) act_fn = _convert_to_activation_function(activation_type) if apply_topk_weights_early: @@ -407,30 +437,34 @@ def _ffn_bwd_per_shard( # 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.dgrad, local_group_sizes, flatten_axis=-1 - ) + d_combined = jax.lax.with_sharding_constraint(d_combined, flat_token_sharding) + casted_d_combined = tex.grouped_quantize(d_combined, q_set.dgrad, group_sizes, flatten_axis=-1) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), casted_wi_rhs_trans, contracting_dims=((1,), (2,)), ) + d_sorted_x = jax.lax.with_sharding_constraint(d_sorted_x, flat_token_sharding) d_wi_combined = tex.grouped_gemm( casted_sorted_x_lhs_trans, 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 = jax.lax.with_sharding_constraint(d_wi_combined, grouped_weight_sharding) 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_combined_bias = jax.lax.with_sharding_constraint( + d_wi_combined_bias, grouped_bias_sharding + ) 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 - 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) + d_sorted_x_3d = d_sorted_x.reshape(*d_expert_outputs.shape[:-1], d_sorted_x.shape[-1]) + d_recv_w_3d = d_recv_w_from_intermediate.reshape(recv_topk_weights.shape) return ( d_sorted_x_3d, d_recv_w_3d, @@ -476,13 +510,12 @@ def _moe_fwd_rule( dtype, apply_topk_weights_early, ): - """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``. """ 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) @@ -538,6 +571,10 @@ def _moe_fwd_rule( batch_pspec_axis = (*data_parallelism_axes, ep_axis) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) + flat_token_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) + flat_group_sharding = NamedSharding(mesh, P(batch_pspec_axis)) + grouped_weight_sharding = NamedSharding(mesh, P(batch_pspec_axis, None, None)) + grouped_bias_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, ep3_spec)) # ---------------- Gate (global view) ---------------- @@ -638,6 +675,7 @@ def _moe_fwd_rule( dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) token_counts, 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 ) @@ -646,86 +684,32 @@ def _moe_fwd_rule( recv_topk_weights, NamedSharding(mesh, ep2_spec) ) - # ---------------- FFN (per-shard via shard_map) ---------------- + # ---------------- FFN (global view, custom-partitioned primitives) ---------------- 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] - 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). 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 - ep2_spec, # local_group_sizes (1, num_local_experts) per shard + # The NCCL EP receive buffer may contain uninitialized padded slots. + # Dynamic group sizes keep grouped GEMMs from reading those rows; the + # backward masks skipped wgrad groups, while EP combine/dispatch bwd + # consume only positions described by handle_mem. + expert_outputs, ffn_residuals = _ffn_fwd_global( + recv_tokens, + recv_topk_weights, + token_counts, + wi_0, + wi_1, + wo, + wi_0_bias if has_bias else None, + wi_1_bias if has_bias else None, + wo_bias if has_bias else None, + dp_size=dp_size, + num_ep=num_ep, + num_local_experts=num_local_experts, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + flat_token_sharding=flat_token_sharding, + flat_group_sharding=flat_group_sharding, + grouped_weight_sharding=grouped_weight_sharding, + grouped_bias_sharding=grouped_bias_sharding, ) - out_specs = (ep3_spec, residuals_spec) - - def _body(*args): - if has_bias: - (r_tok, r_w, tc, w0, w1, w_o, w0b, w1b, wob) = args - else: - (r_tok, r_w, tc, w0, w1, w_o) = 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, - w0b, - w1b, - wob, - 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, - mesh=mesh, - in_specs=ffn_in_specs, - out_specs=out_specs, - check_rep=False, - )(*ffn_in_args) expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) # ---------------- TE EP combine (global view) ---------------- @@ -816,7 +800,6 @@ def _moe_bwd_rule( ): """Backward mirror of :func:`_moe_fwd_rule`.""" 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 @@ -828,6 +811,8 @@ 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] + num_local_experts = num_experts // num_ep dp_size = 1 for ax in data_parallelism_axes: dp_size *= mesh.shape[ax] @@ -840,6 +825,10 @@ def _moe_bwd_rule( batch_pspec_axis = (*data_parallelism_axes, ep_axis) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) + flat_token_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) + flat_group_sharding = NamedSharding(mesh, P(batch_pspec_axis)) + grouped_weight_sharding = NamedSharding(mesh, P(batch_pspec_axis, None, None)) + grouped_bias_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) out_partition_spec = (batch_pspec_axis, None, None) # ---------------- Combine bwd (global view) ---------------- @@ -866,22 +855,17 @@ def _moe_bwd_rule( 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 - ) - bwd_in_args = [ + # ---------------- FFN bwd (global view, custom-partitioned primitives) ---------------- + ( + 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, + ) = _ffn_bwd_global( d_expert_outputs, ctx.casted_sorted_x_lhs_trans, ctx.casted_wi_rhs_trans, @@ -891,75 +875,32 @@ def _moe_bwd_rule( 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 + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + has_bias=has_bias, + flat_token_sharding=flat_token_sharding, + flat_group_sharding=flat_group_sharding, + grouped_weight_sharding=grouped_weight_sharding, + grouped_bias_sharding=grouped_bias_sharding, ) - 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) + # The forward broadcast introduced one expert-gradient group per DP + # replica, ordered (dp, ep, local_expert). Sum that outer dimension + # to recover the public parameter shapes [num_experts, ...]. + def _fold_dp_groups(grad): 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, + grad.reshape(dp_size, num_ep, num_local_experts, *grad.shape[1:]) + .sum(axis=0) + .reshape(num_experts, *grad.shape[1:]) ) - ( - 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=bwd_in_specs, - out_specs=bwd_out_specs, - check_rep=False, - )( - *bwd_in_args - ) + d_wi_0 = _fold_dp_groups(d_wi_0) + d_wi_1 = _fold_dp_groups(d_wi_1) + d_wo = _fold_dp_groups(d_wo) + if has_bias: + d_wi_0_bias = _fold_dp_groups(d_wi_0_bias) + d_wi_1_bias = _fold_dp_groups(d_wi_1_bias) + d_wo_bias = _fold_dp_groups(d_wo_bias) d_recv_w_total = d_recv_w_from_combine + d_recv_w_from_intermediate @@ -1039,6 +980,12 @@ def _bwd_body(*args): 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) + 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 @@ -1185,7 +1132,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``, From 3672316632fab7f640016c1b5888ca188a676446 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 16 Jul 2026 18:15:02 -0700 Subject: [PATCH 13/44] Support MXFP8 quantization in global-view JAX MoE --- tests/jax/test_distributed_grouped_gemm.py | 109 +++++--- transformer_engine/jax/cpp_extensions/gemm.py | 8 + .../jax/cpp_extensions/quantization.py | 239 ++++++++++++++++-- transformer_engine/jax/moe.py | 74 ++++-- .../jax/quantize/dequantizer.py | 9 +- transformer_engine/jax/quantize/tensor.py | 37 ++- 6 files changed, 390 insertions(+), 86 deletions(-) diff --git a/tests/jax/test_distributed_grouped_gemm.py b/tests/jax/test_distributed_grouped_gemm.py index be2487a8df..28241966fb 100644 --- a/tests/jax/test_distributed_grouped_gemm.py +++ b/tests/jax/test_distributed_grouped_gemm.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Partitioning tests for grouped quantize and grouped GEMM.""" +import math from types import SimpleNamespace import jax @@ -75,7 +76,27 @@ def _mxfp8_grouped_quantizer_set(n_groups): ) -def test_grouped_quantize_gathers_hidden_axis_for_block_scales(): +def test_grouped_quantize_large_abstract_shape_preserves_inferred_hidden_dim(): + input_shape = (786432, 4096) + group_count = 64 + outputs = GroupedQuantizePrimitive.abstract( + jax.core.ShapedArray(input_shape, jnp.bfloat16), + jax.core.ShapedArray((group_count,), jnp.float32), + jax.core.ShapedArray((group_count,), jnp.int32), + out_dtype=jnp.float8_e4m3fn, + scaling_mode=ScalingMode.MXFP8_1D_SCALING.value, + q_layout=QuantizeLayout.ROWWISE, + flatten_axis=-1, + scale_dtype=jnp.float8_e8m0fnu, + uniform_groups=False, + ) + + assert outputs[0].shape == input_shape + assert max(outputs[0].shape) < 2**31 + assert math.prod(outputs[0].shape) == 3221225472 + + +def test_grouped_quantize_preserves_output_side_fsdp_for_uniform_kernel(): mesh = _mesh() with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( @@ -84,23 +105,24 @@ def test_grouped_quantize_gathers_hidden_axis_for_block_scales(): QuantizeLayout.ROWWISE, -1, jnp.float8_e8m0fnu, + True, mesh, ( - _arg_info(mesh, (8, 128, 64), ("expert", None, "fsdp")), + _arg_info(mesh, (8, 128, 256), ("expert", None, "fsdp")), _arg_info(mesh, (8,), ("expert",)), _arg_info(mesh, (8,), ("expert",)), ), (), ) - assert tuple(arg_shardings[0].spec) == ("expert", None, None) + assert tuple(arg_shardings[0].spec) == ("expert", None, "fsdp") specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == ("expert",) - assert _normalize_spec(specs[2]) == ("expert",) + assert _normalize_spec(specs[0]) == ("expert", None, "fsdp") + assert _normalize_spec(specs[2]) == ("expert", None, "fsdp") assert _normalize_spec(specs[4]) == ("expert",) -def test_grouped_quantize_mxfp8_colwise_specs_gather_hidden_axis(): +def test_grouped_quantize_mxfp8_colwise_scale_tracks_output_side_fsdp(): mesh = _mesh() with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( @@ -109,21 +131,22 @@ def test_grouped_quantize_mxfp8_colwise_specs_gather_hidden_axis(): QuantizeLayout.ROWWISE_COLWISE, -1, jnp.float8_e8m0fnu, + True, mesh, ( - _arg_info(mesh, (8, 128, 128), ("expert", None, "fsdp")), + _arg_info(mesh, (8, 128, 256), ("expert", None, "fsdp")), _arg_info(mesh, (8,), ("expert",)), _arg_info(mesh, (8,), ("expert",)), ), (), ) - assert tuple(arg_shardings[0].spec) == ("expert", None, None) + assert tuple(arg_shardings[0].spec) == ("expert", None, "fsdp") specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == ("expert",) - assert _normalize_spec(specs[1]) == ("expert",) - assert _normalize_spec(specs[2]) == ("expert",) - assert _normalize_spec(specs[3]) == ("expert",) + assert _normalize_spec(specs[0]) == ("expert", None, "fsdp") + assert _normalize_spec(specs[1]) == ("expert", None, "fsdp") + assert _normalize_spec(specs[2]) == ("expert", None, "fsdp") + assert _normalize_spec(specs[3]) == ("expert", "fsdp", None) assert _normalize_spec(specs[4]) == ("expert",) @@ -136,9 +159,10 @@ def test_grouped_quantize_preserves_row_side_fsdp_for_kernel(): QuantizeLayout.ROWWISE, -1, jnp.float8_e8m0fnu, + True, mesh, ( - _arg_info(mesh, (8, 128, 64), ("expert", "fsdp", None)), + _arg_info(mesh, (8, 256, 128), ("expert", "fsdp", None)), _arg_info(mesh, (8,), ("expert",)), _arg_info(mesh, (8,), ("expert",)), ), @@ -147,11 +171,11 @@ def test_grouped_quantize_preserves_row_side_fsdp_for_kernel(): assert tuple(arg_shardings[0].spec) == ("expert", "fsdp", None) specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == (("expert", "fsdp"),) - assert _normalize_spec(specs[2]) == (("expert", "fsdp"),) + assert _normalize_spec(specs[0]) == ("expert", "fsdp", None) + assert _normalize_spec(specs[2]) == ("expert", "fsdp", None) -def test_grouped_quantize_strips_unsupported_axes_and_gathers_hidden_axes(): +def test_grouped_quantize_strips_unsupported_axes_and_preserves_supported_axes(): mesh = _mesh_with_dp_tp() with jax.set_mesh(mesh), global_shard_guard( MeshResource(dp_resource="dp", tp_resource="tp", fsdp_resource="fsdp", ep_resource="expert") @@ -163,22 +187,23 @@ def test_grouped_quantize_strips_unsupported_axes_and_gathers_hidden_axes(): QuantizeLayout.ROWWISE, -1, jnp.float8_e8m0fnu, + True, mesh, ( - _arg_info(mesh, (8, 128, 128), ("expert", "dp", ("fsdp", "tp"))), + _arg_info(mesh, (8, 128, 256), ("expert", "dp", ("fsdp", "tp"))), _arg_info(mesh, (8,), (("expert", "tp"),)), _arg_info(mesh, (8,), (("expert", "tp"),)), ), (), ) - assert tuple(arg_shardings[0].spec) == ("expert", "dp", None) + assert tuple(arg_shardings[0].spec) == ("expert", "dp", "fsdp") assert tuple(arg_shardings[1].spec) == ("expert",) assert tuple(arg_shardings[2].spec) == ("expert",) out_specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(out_specs[0]) == (("expert", "dp"),) - assert _normalize_spec(out_specs[2]) == (("expert", "dp"),) + assert _normalize_spec(out_specs[0]) == ("expert", "dp", "fsdp") + assert _normalize_spec(out_specs[2]) == ("expert", "dp", "fsdp") assert _normalize_spec(out_specs[4]) == ("expert",) for spec in (*out_specs, *(tuple(sharding.spec) for sharding in arg_shardings)): assert not _spec_contains_axis(spec, "tp") @@ -327,9 +352,10 @@ def test_grouped_partitioning_strips_arbitrary_unsupported_axis(): QuantizeLayout.ROWWISE, -1, jnp.float8_e8m0fnu, + True, mesh, ( - _arg_info(mesh, (8, 128, 128), ("expert", "myaxis123", ("dp", "fsdp"))), + _arg_info(mesh, (8, 128, 256), ("expert", "myaxis123", ("dp", "fsdp"))), _arg_info(mesh, (8,), (("expert", "myaxis123"),)), _arg_info(mesh, (8,), (("expert", "myaxis123"),)), ), @@ -374,11 +400,11 @@ def test_grouped_partitioning_strips_arbitrary_unsupported_axis(): gemm_result_infos, ) - assert tuple(quantize_arg_shardings[0].spec) == ("expert", None, None) + assert tuple(quantize_arg_shardings[0].spec) == ("expert", None, ("dp", "fsdp")) assert tuple(quantize_arg_shardings[1].spec) == ("expert",) quantize_out_specs = tuple(tuple(sharding.spec) for sharding in quantize_out_shardings) - assert _normalize_spec(quantize_out_specs[0]) == ("expert",) - assert _normalize_spec(quantize_out_specs[2]) == ("expert",) + assert _normalize_spec(quantize_out_specs[0]) == ("expert", None, ("dp", "fsdp")) + assert _normalize_spec(quantize_out_specs[2]) == ("expert", None, ("dp", "fsdp")) assert tuple(gemm_arg_shardings[0].spec) == ("dp",) assert tuple(gemm_arg_shardings[2].spec) == ("expert",) @@ -403,16 +429,17 @@ def test_grouped_partitioning_shardy_rules_smoke(): QuantizeLayout.ROWWISE, -1, jnp.float8_e8m0fnu, + True, mesh, ( - SimpleNamespace(shape=(8, 128, 64)), + SimpleNamespace(shape=(8, 128, 128)), SimpleNamespace(shape=(8,)), SimpleNamespace(shape=(8,)), ), ( - SimpleNamespace(shape=(8 * 128 * 64,)), + SimpleNamespace(shape=(8, 128, 128)), SimpleNamespace(shape=(1,)), - SimpleNamespace(shape=(8 * 128 * 64,)), + SimpleNamespace(shape=(8, 1, 512)), SimpleNamespace(shape=(1,)), SimpleNamespace(shape=(8,)), ), @@ -441,17 +468,25 @@ def test_grouped_partitioning_shardy_rules_smoke(): assert gemm_rule is not None -def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process(): +@pytest.mark.parametrize( + "weight_spec", + [ + ("expert", "fsdp", None), + ("expert", None, "fsdp"), + ], + ids=("contracting-fsdp", "output-fsdp"), +) +def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process(weight_spec): mesh = _mesh() n_groups = 4 group_tokens = 128 hidden = 256 - out_hidden = 128 + out_hidden = 256 x_shape = (n_groups * group_tokens, hidden) w_shape = (n_groups, hidden, out_hidden) x_sharding = NamedSharding(mesh, PartitionSpec("expert", None)) - w_sharding = NamedSharding(mesh, PartitionSpec("expert", "fsdp", None)) + w_sharding = NamedSharding(mesh, PartitionSpec(*weight_spec)) group_sharding = NamedSharding(mesh, PartitionSpec("expert")) out_sharding = NamedSharding(mesh, PartitionSpec("expert", None)) @@ -496,8 +531,20 @@ def apply(x, w): assert tuple(out.sharding.spec) == ("expert", None) assert tuple(dx.sharding.spec) == ("expert", None) - assert tuple(dw.sharding.spec) == ("expert", "fsdp", None) + assert tuple(dw.sharding.spec) == weight_spec for value in (out, dx, dw): local_value = np.asarray(jax.device_get(value.addressable_data(0))) assert np.all(np.isfinite(local_value)) assert np.any(local_value != 0.0) + + x_global = np.asarray(jax.device_get(x)).reshape(n_groups, group_tokens, hidden) + w_global = np.asarray(jax.device_get(w)) + reference = np.einsum( + "gth,gho->gto", x_global.astype(np.float32), w_global.astype(np.float32) + ).reshape(x_shape[0], out_hidden) + np.testing.assert_allclose( + np.asarray(jax.device_get(out)).astype(np.float32), + reference.astype(np.float32), + atol=5e-3, + rtol=5e-2, + ) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 54e4b510cf..0326641471 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -2025,6 +2025,14 @@ def sharded_impl( additional_arg_0, additional_arg_1, ): + # Grouped quantization may expose rank-3 scale carriers so Shardy + # can gather expert and FSDP axes independently. At this point the + # requested local shardings/collectives have already been applied; + # the FFI consumes the same contiguous pre-swizzled bytes as 1D. + if lhs_scale_inv.ndim > 2: + lhs_scale_inv = lhs_scale_inv.reshape(-1) + if rhs_scale_inv.ndim > 2: + rhs_scale_inv = rhs_scale_inv.reshape(-1) (out,) = GroupedGemmPrimitive.impl( lhs_data, lhs_scale_inv, diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 9266ab08f0..44c71fdecd 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -12,7 +12,7 @@ import jax import jax.numpy as jnp from jax import dtypes, ffi -from jax.experimental.custom_partitioning import SdyShardingRule, BATCHING +from jax.experimental.custom_partitioning import SdyShardingRule, BATCHING, CompoundFactor from jax.sharding import PartitionSpec import transformer_engine_jax @@ -32,6 +32,7 @@ from ..sharding import ( all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp, + axis_spec_size, filter_spec_axes, get_num_devices_in_mesh, global_mesh_resource, @@ -64,6 +65,41 @@ def _flat_data_spec(input_spec): return (merge_axis_specs(*input_spec),) +def _grouped_data_spec(input_spec, flatten_axis): + return _contiguous_flat_input_spec(input_spec, flatten_axis) + + +def _uniform_mxfp8_scale_carrier_shapes(x_shape, n_groups, flatten_axis): + """Return sharding-friendly scale carriers for uniform 3D MXFP8 kernels.""" + flatten_axis = _normalize_flatten_axis(flatten_axis, len(x_shape)) + if len(x_shape) != 3 or flatten_axis != 2 or x_shape[0] != n_groups: + return None + + group_rows = x_shape[1] + columns = x_shape[2] + if group_rows % 128 != 0 or columns % 128 != 0: + return None + + # The fused MXFP8 swizzle groups scales in 128-row/column tiles. These + # carriers expose group and column ownership as separate dimensions while + # preserving the exact contiguous byte order consumed by grouped GEMM. + rowwise_shape = (n_groups, group_rows // 128, columns * 4) + colwise_shape = (n_groups, columns // 128, group_rows * 4) + return rowwise_shape, colwise_shape + + +def _uniform_mxfp8_scale_carrier_specs(x_spec, flatten_axis): + """Map a uniform 3D kernel's data sharding onto its scale carriers.""" + flatten_axis = _normalize_flatten_axis(flatten_axis, len(x_spec)) + assert len(x_spec) == 3 and flatten_axis == 2 + group_spec = x_spec[0] + row_spec = merge_axis_specs(*x_spec[1:flatten_axis]) + column_spec = merge_axis_specs(*x_spec[flatten_axis:]) + return ( + (group_spec, row_spec, column_spec), + (group_spec, column_spec, row_spec), + ) + def _normalize_flatten_axis(flatten_axis, ndim): return flatten_axis + ndim if flatten_axis < 0 else flatten_axis @@ -1059,7 +1095,8 @@ class GroupedQuantizePrimitive(BasePrimitive): 5, 6, 7, - ) # out_dtype, scaling_mode, q_layout, flatten_axis, scale_dtype + 8, + ) # out_dtype, scaling_mode, q_layout, flatten_axis, scale_dtype, uniform_groups inner_primitive = None outer_primitive = None @@ -1122,13 +1159,16 @@ def abstract( q_layout, flatten_axis, scale_dtype, + uniform_groups, ): """ te_dbias_quantize_p abstract """ dtype = dtypes.canonicalize_dtype(x_aval.dtype) assert dtype in [jnp.float32, jnp.float16, jnp.bfloat16] - out_shape = math.prod(x_aval.shape) + # Preserve logical rank for custom partitioning. The FFI still consumes + # the same contiguous buffer, while avoiding oversized flat Shardy dims. + out_shape = x_aval.shape # TODO(Phuong): can scale_aval be None? assert scale_aval is None or scale_aval.dtype == jnp.float32 @@ -1137,14 +1177,22 @@ def abstract( f" be one of {ScalingMode(scaling_mode).get_compatible_q_dtypes()}" ) - rowwise_scale_inv_shape, colwise_scale_inv_shape = ScalingMode( - scaling_mode - ).get_grouped_scale_shape_2x( - x_aval.shape, - group_sizes_aval.size, - is_padded=True, - flatten_axis=flatten_axis, - ) + scale_carrier_shapes = None + if uniform_groups and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING: + scale_carrier_shapes = _uniform_mxfp8_scale_carrier_shapes( + x_aval.shape, group_sizes_aval.size, flatten_axis + ) + if scale_carrier_shapes is None: + rowwise_scale_inv_shape, colwise_scale_inv_shape = ScalingMode( + scaling_mode + ).get_grouped_scale_shape_2x( + x_aval.shape, + group_sizes_aval.size, + is_padded=True, + flatten_axis=flatten_axis, + ) + else: + rowwise_scale_inv_shape, colwise_scale_inv_shape = scale_carrier_shapes if q_layout.has_rowwise: rowwise_out_shape = out_shape @@ -1224,11 +1272,12 @@ def lowering( q_layout, flatten_axis, scale_dtype, + uniform_groups, ): """ te_dbias_quantize_p lowering rules """ - del out_dtype, scale_dtype + del out_dtype, scale_dtype, uniform_groups x_aval, scale_aval, group_sizes_aval = ctx.avals_in assert x_aval.dtype in [jnp.float32, jnp.float16, jnp.bfloat16] assert scale_aval.dtype == jnp.float32 @@ -1268,6 +1317,7 @@ def impl( q_layout, flatten_axis, scale_dtype, + uniform_groups, ): """ te_dbias_quantize_p implementation @@ -1289,15 +1339,34 @@ def impl( q_layout=q_layout, flatten_axis=flatten_axis, scale_dtype=scale_dtype, + uniform_groups=uniform_groups, ) return rowwise_out, colwise_out, rowwise_scale_inv, colwise_scale_inv, updated_amax @staticmethod - def _parse_partition_specs(scaling_mode, q_layout, flatten_axis, mesh, arg_infos): + def _parse_partition_specs( + scaling_mode, q_layout, flatten_axis, uniform_groups, mesh, arg_infos + ): allowed_axes = supported_grouped_partition_axes(mesh) original_x_spec = get_padded_spec(arg_infos[0]) x_spec = filter_spec_axes(original_x_spec, allowed_axes) - x_spec = _contiguous_flat_input_spec(x_spec, flatten_axis) + use_scale_carrier = ( + uniform_groups + and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING + and _uniform_mxfp8_scale_carrier_shapes( + arg_infos[0].shape, arg_infos[2].size, flatten_axis + ) + is not None + ) + if use_scale_carrier: + x_spec = list(x_spec) + for dim in (1, 2): + local_dim = arg_infos[0].shape[dim] // axis_spec_size(x_spec[dim], mesh) + if local_dim % 128 != 0: + x_spec[dim] = None + x_spec = tuple(x_spec) + else: + x_spec = _contiguous_flat_input_spec(x_spec, flatten_axis) _warn_if_axes_ignored("x", original_x_spec, x_spec) original_group_spec = get_padded_spec(arg_infos[2]) @@ -1306,16 +1375,28 @@ def _parse_partition_specs(scaling_mode, q_layout, flatten_axis, mesh, arg_infos group_spec = (x_spec[0],) _warn_if_axes_ignored("group_sizes", original_group_spec, group_spec) flat_spec = _flat_data_spec(x_spec) + data_spec = tuple(x_spec) if use_scale_carrier else _grouped_data_spec(x_spec, flatten_axis) replicated_spec = (None,) - rowwise_out_spec = flat_spec if q_layout.has_rowwise else replicated_spec - colwise_out_spec = flat_spec if q_layout.has_colwise else replicated_spec + rowwise_out_spec = data_spec if q_layout.has_rowwise else replicated_spec + colwise_out_spec = data_spec if q_layout.has_colwise else replicated_spec rowwise_scale_inv_spec = replicated_spec colwise_scale_inv_spec = replicated_spec if ScalingMode(scaling_mode).is_block_scaling: - rowwise_scale_inv_spec = flat_spec if q_layout.has_rowwise else replicated_spec - colwise_scale_inv_spec = flat_spec if q_layout.has_colwise else replicated_spec + if use_scale_carrier: + rowwise_carrier_spec, colwise_carrier_spec = _uniform_mxfp8_scale_carrier_specs( + x_spec, flatten_axis + ) + rowwise_scale_inv_spec = ( + rowwise_carrier_spec if q_layout.has_rowwise else replicated_spec + ) + colwise_scale_inv_spec = ( + colwise_carrier_spec if q_layout.has_colwise else replicated_spec + ) + else: + rowwise_scale_inv_spec = flat_spec if q_layout.has_rowwise else replicated_spec + colwise_scale_inv_spec = flat_spec if q_layout.has_colwise else replicated_spec elif ScalingMode(scaling_mode).is_tensor_scaling(): rowwise_scale_inv_spec = group_spec if q_layout.has_rowwise else replicated_spec colwise_scale_inv_spec = group_spec if q_layout.has_colwise else replicated_spec @@ -1340,12 +1421,21 @@ def partition( q_layout, flatten_axis, scale_dtype, + uniform_groups, mesh, arg_infos, result_infos, ): x_spec, group_spec, out_specs = GroupedQuantizePrimitive._parse_partition_specs( - scaling_mode, q_layout, flatten_axis, mesh, arg_infos + scaling_mode, q_layout, flatten_axis, uniform_groups, mesh, arg_infos + ) + use_scale_carrier = ( + uniform_groups + and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING + and _uniform_mxfp8_scale_carrier_shapes( + arg_infos[0].shape, arg_infos[2].size, flatten_axis + ) + is not None ) local_out_shapes = ( tuple( @@ -1379,6 +1469,10 @@ def sharded_impl(x, scale, group_sizes): q_layout=q_layout, flatten_axis=flatten_axis, scale_dtype=scale_dtype, + # The rank-3 carrier is a partitioning-only representation. Keep + # the inner FFI contract flat and reshape its local scale buffers + # after quantization. + uniform_groups=False if use_scale_carrier else uniform_groups, ) if ScalingMode(scaling_mode).is_block_scaling: rowwise_scale_inv = _pad_or_slice_to_shape(rowwise_scale_inv, local_out_shapes[2]) @@ -1402,20 +1496,117 @@ def shardy_sharding_rule( q_layout, flatten_axis, scale_dtype, + uniform_groups, mesh, value_types, result_types, ): - del out_dtype, scale_dtype, mesh, result_types, flatten_axis + del out_dtype, scale_dtype, mesh prefix = "GroupedQuantize" - input_spec = tuple(f"{prefix}_x_{i}" for i in range(len(value_types[0].shape))) + input_shape = value_types[0].shape + input_spec = tuple(f"{prefix}_x_{i}" for i in range(len(input_shape))) + normalized_flatten_axis = _normalize_flatten_axis(flatten_axis, len(input_spec)) + use_scale_carrier = ( + uniform_groups + and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING + and _uniform_mxfp8_scale_carrier_shapes( + input_shape, value_types[2].shape[0], flatten_axis + ) + is not None + ) + if use_scale_carrier: + group_factor = f"{prefix}_kernel_group" + row_factor = f"{prefix}_kernel_row_tiles" + column_factor = f"{prefix}_kernel_column_tiles" + row_block_factor = f"{prefix}_kernel_row_block" + column_block_factor = f"{prefix}_kernel_column_block" + row_pack_factor = f"{prefix}_kernel_row_pack" + column_pack_factor = f"{prefix}_kernel_column_pack" + row_unit_factor = f"{prefix}_kernel_row_unit" + column_unit_factor = f"{prefix}_kernel_column_unit" + + row_tiles = input_shape[1] // 128 + column_tiles = input_shape[2] // 128 + factor_sizes = {} + if row_tiles == 1: + row_input_factor = f"{prefix}_kernel_row" + row_scale_factor = row_unit_factor + row_packed_factors = (row_input_factor,) + else: + row_input_factor = CompoundFactor(row_factor, row_block_factor) + row_scale_factor = row_factor + row_packed_factors = (row_factor, row_block_factor) + factor_sizes[row_block_factor] = 128 + if column_tiles == 1: + column_input_factor = f"{prefix}_kernel_column" + column_scale_factor = column_unit_factor + column_packed_factors = (column_input_factor,) + else: + column_input_factor = CompoundFactor(column_factor, column_block_factor) + column_scale_factor = column_factor + column_packed_factors = (column_factor, column_block_factor) + factor_sizes[column_block_factor] = 128 + + input_spec = ( + group_factor, + row_input_factor, + column_input_factor, + ) + data_spec = input_spec + rowwise_scale_spec = ( + ( + group_factor, + row_scale_factor, + CompoundFactor(*column_packed_factors, row_pack_factor), + ) + if q_layout.has_rowwise + else (BATCHING + f"{prefix}_scalar",) + ) + colwise_scale_spec = ( + ( + group_factor, + column_scale_factor, + CompoundFactor(*row_packed_factors, column_pack_factor), + ) + if q_layout.has_colwise + else (BATCHING + f"{prefix}_scalar",) + ) + if q_layout.has_rowwise: + factor_sizes[row_pack_factor] = 4 + elif row_tiles > 1: + factor_sizes[row_factor] = input_shape[1] // 128 + if q_layout.has_colwise: + factor_sizes[column_pack_factor] = 4 + elif column_tiles > 1: + factor_sizes[column_factor] = input_shape[2] // 128 + + scalar_spec = (BATCHING + f"{prefix}_scalar",) + rowwise_out_spec = data_spec if q_layout.has_rowwise else scalar_spec + colwise_out_spec = data_spec if q_layout.has_colwise else scalar_spec + group_spec = (BATCHING + f"{prefix}_group",) + return SdyShardingRule( + operand_mappings=(input_spec, group_spec, group_spec), + result_mappings=( + rowwise_out_spec, + colwise_out_spec, + rowwise_scale_spec, + colwise_scale_spec, + group_spec, + ), + **factor_sizes, + ) + + data_spec = tuple( + input_spec[i] if i < normalized_flatten_axis else f"{prefix}_data_{i}" + for i in range(len(input_spec)) + ) flat_spec = (f"{prefix}_flat",) group_spec = (BATCHING + f"{prefix}_group",) scalar_spec = (BATCHING + f"{prefix}_scalar",) - rowwise_out_spec = flat_spec if q_layout.has_rowwise else scalar_spec - colwise_out_spec = flat_spec if q_layout.has_colwise else scalar_spec + rowwise_out_spec = data_spec if q_layout.has_rowwise else scalar_spec + colwise_out_spec = data_spec if q_layout.has_colwise else scalar_spec if ScalingMode(scaling_mode).is_block_scaling: rowwise_scale_spec = flat_spec if q_layout.has_rowwise else scalar_spec @@ -1488,6 +1679,7 @@ def grouped_quantize( ), f"Only flatten_axis = -1 is supported for now, got {flatten_axis}" ragged_first_dims = group_sizes # None if no explicit group_sizes (kernel case) + uniform_groups = group_sizes is None if group_sizes is None: group_sizes = jnp.ones(x.shape[0], dtype=jnp.int32) @@ -1535,6 +1727,7 @@ def grouped_quantize( q_layout=q_layout, flatten_axis=flatten_axis, scale_dtype=quantizer.get_scale_dtype(), + uniform_groups=uniform_groups, ) # For DelayedScaling2x and CurrentScaling2x, the scale buffer diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 85a0ed94f6..55dbb63b55 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -24,13 +24,10 @@ view. Their custom partitioning rules handle per-shard execution, including EP placement and DP/FSDP gathers and reductions. -Out-of-scope (for now) ----------------------- -FP8 / MXFP8 quantizer sets are not yet wired on this path; turning -them on requires recipe-aware residual handling for ``ScaledTensor`` -leaves. ``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). +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. """ from functools import partial @@ -44,6 +41,7 @@ from . import cpp_extensions as tex from .quantize import ( + QuantizerSet, TensorUsage, noop_quantizer_set, with_sharding_constraint_by_logical_axes, @@ -204,6 +202,7 @@ 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 @@ -224,6 +223,7 @@ def _ffn_fwd_global( wi_0_bias: Optional[jnp.ndarray], wi_1_bias: Optional[jnp.ndarray], wo_bias: Optional[jnp.ndarray], + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, dp_size: int, num_ep: int, @@ -283,9 +283,13 @@ def _broadcast_experts(value, sharding): 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, 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_combined, 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), @@ -294,8 +298,12 @@ def _broadcast_experts(value, sharding): ) combined_out = jax.lax.with_sharding_constraint(combined_out, flat_token_sharding) 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) + casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor( + usage=TensorUsage.LHS_TRANS + ).checkpoint(fc1_quantizer_set.x) + casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint( + fc1_quantizer_set.kernel + ) # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM # output dtype; the activation output (`intermediate`) stays in the @@ -318,8 +326,10 @@ def _broadcast_experts(value, sharding): 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, group_sizes, flatten_axis=-1) - casted_wo = tex.grouped_quantize(wo, q_set.kernel, flatten_axis=-1) + casted_intermediate = tex.grouped_quantize( + intermediate, fc2_quantizer_set.x, group_sizes, flatten_axis=-1 + ) + casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), casted_wo.get_tensor(usage=TensorUsage.RHS), @@ -327,8 +337,12 @@ def _broadcast_experts(value, sharding): bias=wo_bias, ) expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, flat_token_sharding) - casted_intermediate_lhs_trans = casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS) - casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) + casted_intermediate_lhs_trans = casted_intermediate.get_tensor( + usage=TensorUsage.LHS_TRANS + ).checkpoint(fc2_quantizer_set.x) + casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint( + fc2_quantizer_set.kernel + ) expert_outputs_3d = expert_outputs.reshape(*recv_tokens.shape[:-1], expert_outputs.shape[-1]) group_sizes_nd = group_sizes.reshape(token_counts.shape) @@ -354,6 +368,7 @@ def _ffn_bwd_global( casted_wo_rhs_trans, local_group_sizes: jnp.ndarray, recv_topk_weights: jnp.ndarray, + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, activation_type: str, apply_topk_weights_early: bool, @@ -375,14 +390,16 @@ def _ffn_bwd_global( d_eo_2d = jax.lax.with_sharding_constraint(d_eo_2d, flat_token_sharding) recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) - q_set = noop_quantizer_set + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets # 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 = (group_sizes > 0)[:, None, None] # wo bwd - casted_d_eo = tex.grouped_quantize(d_eo_2d, q_set.dgrad, 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( @@ -438,7 +455,9 @@ def _ffn_bwd_global( # 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) d_combined = jax.lax.with_sharding_constraint(d_combined, flat_token_sharding) - casted_d_combined = tex.grouped_quantize(d_combined, q_set.dgrad, group_sizes, flatten_axis=-1) + casted_d_combined = tex.grouped_quantize( + 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), casted_wi_rhs_trans, @@ -492,6 +511,7 @@ def _moe_fwd_rule( wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -700,6 +720,7 @@ def _moe_fwd_rule( wi_0_bias if has_bias else None, wi_1_bias if has_bias else None, wo_bias if has_bias else None, + quantizer_sets, dp_size=dp_size, num_ep=num_ep, num_local_experts=num_local_experts, @@ -765,6 +786,7 @@ def _moe_fwd_rule( 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, @@ -875,6 +897,7 @@ def _moe_bwd_rule( ctx.casted_wo_rhs_trans, ctx.local_group_sizes, ctx.recv_topk_weights, + ctx.quantizer_sets, activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, has_bias=has_bias, @@ -1003,6 +1026,7 @@ def _fold_dp_groups(grad): d_wi_1_bias if has_bias else None, d_wo_bias if has_bias else None, d_expert_bias, + ctx.quantizer_sets, ) @@ -1011,7 +1035,7 @@ def _fold_dp_groups(grad): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 26))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(10, 27))) def _moe( x, gate_kernel, @@ -1022,6 +1046,7 @@ def _moe( wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1050,6 +1075,7 @@ def _moe( wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1095,6 +1121,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], ...] = (), @@ -1122,6 +1152,9 @@ 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. They are differentiable + custom-VJP arguments so recipe state is threaded through backward. Note that the per-expert dispatch-slot alignment is fixed internally at 128 tokens (``_ALIGN_SIZE``); see that constant's docstring for @@ -1195,6 +1228,7 @@ def moe( wi_1_bias, wo_bias, expert_bias_arg, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, diff --git a/transformer_engine/jax/quantize/dequantizer.py b/transformer_engine/jax/quantize/dequantizer.py index ca44c2e4af..269c824089 100644 --- a/transformer_engine/jax/quantize/dequantizer.py +++ b/transformer_engine/jax/quantize/dequantizer.py @@ -303,8 +303,13 @@ def _grouped_dequantize(grouped_scaled_tensor): Returns: List of dequantized tensors for each group """ - data = grouped_scaled_tensor.data - scale_inv = grouped_scaled_tensor.scale_inv + # Group offsets are scalar-element offsets even though grouped quantization + # preserves the logical N-D carrier shape to avoid oversized 1-D dimensions. + data = grouped_scaled_tensor.data.reshape(-1) + # Uniform grouped kernels may use a multidimensional carrier so expert and + # FSDP ownership survive custom partitioning. Group offsets still address + # the same contiguous pre-swizzled byte stream. + scale_inv = grouped_scaled_tensor.scale_inv.reshape(-1) group_sizes = ( grouped_scaled_tensor.first_dims if grouped_scaled_tensor.first_dims is not None diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index c5ad0451fd..95bb8ae51c 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. """ +from abc import ABC, abstractmethod from dataclasses import dataclass +import math 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 @@ -427,8 +428,14 @@ def group_sizes(self) -> jnp.ndarray: return jnp.ones((self.original_shape[0],), dtype=jnp.int32) def __post_init__(self): - assert self.scale_inv.ndim == 1, "Only support flattened scale_inv" - assert self.data.ndim == 1, "Only support flattened data" + assert self.scale_inv.ndim in (1, 3), ( + "Grouped scale_inv must be flat or use the uniform-kernel 3D carrier, " + f"got shape {self.scale_inv.shape}" + ) + assert self.data.size == math.prod(self.original_shape), ( + f"Quantized data has {self.data.size} elements, expected " + f"{math.prod(self.original_shape)} for original shape {self.original_shape}" + ) assert self.flatten_axis > 0 data_ndim = len(self.original_shape) @@ -446,13 +453,23 @@ def __post_init__(self): else: num_groups = self.original_shape[0] - expected_scale_shape = self.scaling_mode.get_grouped_scale_shape( - self.original_shape, - num_groups, - self.is_colwise, - is_padded=True, - flatten_axis=self.flatten_axis, - ) + if self.scale_inv.ndim == 3: + assert self.scaling_mode == ScalingMode.MXFP8_1D_SCALING + assert len(self.original_shape) == 3 and self.flatten_axis == 2 + groups, rows, columns = self.original_shape + expected_scale_shape = ( + (groups, columns // 128, rows * 4) + if self.is_colwise + else (groups, rows // 128, columns * 4) + ) + else: + expected_scale_shape = self.scaling_mode.get_grouped_scale_shape( + self.original_shape, + num_groups, + self.is_colwise, + is_padded=True, + flatten_axis=self.flatten_axis, + ) assert self.scale_inv.shape == expected_scale_shape, ( f"Unexpected scale_inv shape! \nExpect {expected_scale_shape} for padded" From 330b2ccbc7cfdfc54030d8bb51aa19943fbdac73 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Fri, 17 Jul 2026 09:42:02 -0700 Subject: [PATCH 14/44] Quantize JAX MoE weights before FSDP gather --- tests/jax/test_distributed_grouped_gemm.py | 43 +++++++++++++++++++ transformer_engine/jax/cpp_extensions/gemm.py | 21 +++++++-- transformer_engine/jax/moe.py | 32 +++++++------- 3 files changed, 76 insertions(+), 20 deletions(-) diff --git a/tests/jax/test_distributed_grouped_gemm.py b/tests/jax/test_distributed_grouped_gemm.py index 28241966fb..5e850094ba 100644 --- a/tests/jax/test_distributed_grouped_gemm.py +++ b/tests/jax/test_distributed_grouped_gemm.py @@ -252,6 +252,49 @@ def test_grouped_gemm_rhs_weight_specs_gather_fsdp_but_preserve_ep(): assert tuple(out_sharding[0].spec) == (None, None, None) +def test_grouped_gemm_gathers_smaller_moe_rhs_across_fsdp_group_axis(): + """A global MoE RHS has E groups while token counts have dp * E groups.""" + mesh = _mesh() + arg_infos = ( + _arg_info(mesh, (8192,), (None,)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (32, 128, 64), (("fsdp", "expert"), None, None)), + _arg_info(mesh, (2048,), (("fsdp", "expert"),)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (64,), (("fsdp", "expert"),)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (64,), (("fsdp", "expert"),)), + _arg_info(mesh, (0,), (None,)), + _arg_info(mesh, (1,), (None,)), + _arg_info(mesh, (0,), (None,)), + ) + with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): + _, _, _, arg_shardings = GroupedGemmPrimitive.partition( + False, + False, + ScalingMode.NO_SCALING.value, + jnp.bfloat16, + False, + False, + False, + 1, + 1, + (64, 128, 64), + 128, + 64, + 128, + 64, + mesh, + arg_infos, + (), + ) + + assert tuple(arg_shardings[2].spec) == ("expert", None, None) + assert tuple(arg_shardings[3].spec) == ("expert",) + + def test_grouped_gemm_strips_unsupported_axes_preserves_dp_and_gathers_rhs_fsdp(): mesh = _mesh_with_dp_tp() arg_infos = ( diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 0326641471..a9805bf282 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -1865,14 +1865,27 @@ def _parse_partition_specs( # A compound leading group dimension may use FSDP as the outer # data-parallel axis (for example MoE groups ordered - # (fsdp, ep, local_expert)). In that case FSDP describes distinct - # groups, not a sharded RHS contracting dimension, and must remain - # on the group axis. Otherwise gather the FSDP-sharded RHS as usual. + # (fsdp, ep, local_expert)). Normally that means FSDP describes + # distinct groups, not a sharded RHS contracting dimension. MoE + # model weights are the exception: their global expert-group axis + # is smaller than the active token-group array, so FSDP shards a + # shared RHS that must be gathered locally. Keep that gather inside + # this custom partitioning boundary, after grouped quantization. fsdp_is_group_axis = spec_contains_axis(active_group_spec, fsdp_axis) + active_group_count = next( + (info.shape[0] for info in grouped_dim_infos if info.size > 0), + None, + ) + rhs_group_count = arg_infos[2].shape[0] if len(arg_infos[2].shape) > 0 else None + rhs_has_fewer_groups = ( + active_group_count is not None + and rhs_group_count is not None + and rhs_group_count < active_group_count + ) gather_rhs_fsdp = ( fsdp_axis is not None and not rhs_is_ragged - and not fsdp_is_group_axis + and (not fsdp_is_group_axis or rhs_has_fewer_groups) and ( spec_contains_axis(rhs_data_spec, fsdp_axis) or spec_contains_axis(rhs_scale_spec, fsdp_axis) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 55dbb63b55..fb627718d7 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -232,7 +232,6 @@ def _ffn_fwd_global( apply_topk_weights_early: bool, flat_token_sharding: NamedSharding, flat_group_sharding: NamedSharding, - grouped_weight_sharding: NamedSharding, grouped_bias_sharding: NamedSharding, ): """Run the FFN on global EP-dispatch buffers. @@ -254,25 +253,27 @@ def _ffn_fwd_global( wi_1 = wi_1.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) - # Dispatch groups flatten in (dp, ep, local_expert) order. Broadcast - # each global expert parameter set over the outer DP dimension before - # flattening so the grouped weights have exactly the same ordering. + # Dispatch groups flatten in (dp, ep, local_expert) order. Keep the + # model weights in their original global [expert, ...] layout: grouped + # quantize must see the FSDP shard before grouped_gemm's custom + # partitioning gathers it. The grouped-GEMM partitioner maps that + # smaller RHS group axis onto the local dispatch groups after quantizing. + # + # Bias is the one exception. grouped_gemm's public bias contract is one + # row per global GEMM, so retain its inexpensive logical expansion here. num_groups = dp_size * num_ep * num_local_experts - def _broadcast_experts(value, sharding): + def _broadcast_bias(value): value = jnp.broadcast_to( value.reshape(1, num_ep, num_local_experts, *value.shape[1:]), (dp_size, num_ep, num_local_experts, *value.shape[1:]), ).reshape(num_groups, *value.shape[1:]) - return jax.lax.with_sharding_constraint(value, sharding) + return jax.lax.with_sharding_constraint(value, grouped_bias_sharding) - wi_0 = _broadcast_experts(wi_0, grouped_weight_sharding) - wi_1 = _broadcast_experts(wi_1, grouped_weight_sharding) - wo = _broadcast_experts(wo, grouped_weight_sharding) if wi_0_bias is not None: - wi_0_bias = _broadcast_experts(wi_0_bias, grouped_bias_sharding) - wi_1_bias = _broadcast_experts(wi_1_bias, grouped_bias_sharding) - wo_bias = _broadcast_experts(wo_bias, grouped_bias_sharding) + wi_0_bias = _broadcast_bias(wi_0_bias) + wi_1_bias = _broadcast_bias(wi_1_bias) + wo_bias = _broadcast_bias(wo_bias) # 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 @@ -728,7 +729,6 @@ def _moe_fwd_rule( apply_topk_weights_early=apply_topk_weights_early, flat_token_sharding=flat_token_sharding, flat_group_sharding=flat_group_sharding, - grouped_weight_sharding=grouped_weight_sharding, grouped_bias_sharding=grouped_bias_sharding, ) expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) @@ -907,9 +907,9 @@ def _moe_bwd_rule( grouped_bias_sharding=grouped_bias_sharding, ) - # The forward broadcast introduced one expert-gradient group per DP - # replica, ordered (dp, ep, local_expert). Sum that outer dimension - # to recover the public parameter shapes [num_experts, ...]. + # Wgrad has one expert-gradient group per outer data replica, ordered + # (dp, ep, local_expert). Sum that dimension to recover the public + # parameter shapes [num_experts, ...]. def _fold_dp_groups(grad): return ( grad.reshape(dp_size, num_ep, num_local_experts, *grad.shape[1:]) From d88cbc10191a6043ec3862e947626424ffe4f2ae Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 20 Jul 2026 15:36:41 -0700 Subject: [PATCH 15/44] Simplify TE MoE early weighting padding handling --- transformer_engine/jax/moe.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index fb627718d7..a56c43ec56 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -318,14 +318,8 @@ def _broadcast_bias(value): 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, fc2_quantizer_set.x, group_sizes, flatten_axis=-1 From 1b783b98ba93a1e27854bf3872730925e1b7777d Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 20 Jul 2026 15:55:59 -0700 Subject: [PATCH 16/44] Remove redundant masks from TE MoE ragged paths --- transformer_engine/jax/moe.py | 32 +++++++++----------------------- 1 file changed, 9 insertions(+), 23 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index a56c43ec56..e6c65663da 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -386,11 +386,6 @@ def _ffn_bwd_global( recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) fc1_quantizer_set, fc2_quantizer_set = quantizer_sets - # 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 = (group_sizes > 0)[:, None, None] - # wo bwd casted_d_eo = tex.grouped_quantize( d_eo_2d, fc2_quantizer_set.dgrad, group_sizes, flatten_axis=-1 @@ -408,7 +403,6 @@ def _ffn_bwd_global( _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) - d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) d_wo = jax.lax.with_sharding_constraint(d_wo, grouped_weight_sharding) d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None if has_bias: @@ -416,20 +410,17 @@ def _ffn_bwd_global( 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 @@ -464,7 +455,6 @@ def _ffn_bwd_global( 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 = jax.lax.with_sharding_constraint(d_wi_combined, grouped_weight_sharding) d_wi_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1) if has_bias: @@ -860,14 +850,10 @@ def _moe_bwd_rule( 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) From 35cef8059875d887b6da6ac3e1cd622b421bee7e Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 21 Jul 2026 09:49:45 -0700 Subject: [PATCH 17/44] Keep ragged MoE scales on local shards --- .../jax/cpp_extensions/quantization.py | 24 +++++++++++++++++++ transformer_engine/jax/moe.py | 24 +++++++++++++++---- 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 44c71fdecd..e758600527 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -1638,6 +1638,7 @@ def grouped_quantize( quantizer: GroupedQuantizer, group_sizes: jnp.ndarray = None, flatten_axis: int = -1, + ragged_scale_sharding: NamedSharding | None = None, ) -> Union[GroupedScaledTensor1x, GroupedNoScaleTensor]: """Quantize a tensor in grouped manner. @@ -1650,6 +1651,8 @@ def grouped_quantize( quantizer: The quantizer to use for quantization group_sizes: Array of ints containing the size of each group (default: None) flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1) + ragged_scale_sharding: Data sharding to preserve for flat ragged scale buffers. + Its first partition-spec entry is applied to each scale buffer. Returns: A GroupedScaledTensor1x containing the quantized data @@ -1735,6 +1738,27 @@ def grouped_quantize( if is_tensor_scaling and quantizer.q_layout.is_rowwise_colwise or apply_colwise_war: colwise_scale_inv = rowwise_scale_inv + if ragged_scale_sharding is not None: + if ragged_first_dims is None: + raise ValueError("ragged_scale_sharding requires explicit group_sizes") + if not quantizer.scaling_mode.is_block_scaling: + raise ValueError("ragged_scale_sharding requires block scaling") + data_spec = ragged_scale_sharding.spec + if len(data_spec) == 0: + raise ValueError("ragged_scale_sharding must have a token/group dimension") + scale_sharding = NamedSharding( + ragged_scale_sharding.mesh, + PartitionSpec(data_spec[0]), + ) + if q_layout.has_rowwise: + rowwise_scale_inv = jax.lax.with_sharding_constraint( + rowwise_scale_inv, scale_sharding + ) + if q_layout.has_colwise: + colwise_scale_inv = jax.lax.with_sharding_constraint( + colwise_scale_inv, scale_sharding + ) + # TODO(Phuong): store the whole updated_amax in the grouped_quantize instead? if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: for i, quantizer_i in enumerate(quantizer.quantizers): diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index e6c65663da..a49963d25d 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -286,7 +286,11 @@ def _broadcast_bias(value): 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 + sorted_x, + fc1_quantizer_set.x, + group_sizes, + flatten_axis=-1, + ragged_scale_sharding=flat_token_sharding, ) casted_wi = tex.grouped_quantize( wi_combined, fc1_quantizer_set.kernel, flatten_axis=-1 @@ -322,7 +326,11 @@ def _broadcast_bias(value): intermediate = intermediate * recv_w_flat[:, None].astype(intermediate.dtype) casted_intermediate = tex.grouped_quantize( - intermediate, fc2_quantizer_set.x, group_sizes, flatten_axis=-1 + intermediate, + fc2_quantizer_set.x, + group_sizes, + flatten_axis=-1, + ragged_scale_sharding=flat_token_sharding, ) casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) expert_outputs = tex.grouped_gemm( @@ -388,7 +396,11 @@ def _ffn_bwd_global( fc1_quantizer_set, fc2_quantizer_set = quantizer_sets # wo bwd casted_d_eo = tex.grouped_quantize( - d_eo_2d, fc2_quantizer_set.dgrad, group_sizes, flatten_axis=-1 + d_eo_2d, + fc2_quantizer_set.dgrad, + group_sizes, + flatten_axis=-1, + ragged_scale_sharding=flat_token_sharding, ) _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) @@ -442,7 +454,11 @@ def _ffn_bwd_global( d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) d_combined = jax.lax.with_sharding_constraint(d_combined, flat_token_sharding) casted_d_combined = tex.grouped_quantize( - d_combined, fc1_quantizer_set.dgrad, group_sizes, flatten_axis=-1 + d_combined, + fc1_quantizer_set.dgrad, + group_sizes, + flatten_axis=-1, + ragged_scale_sharding=flat_token_sharding, ) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), From f4beba5f8213eb37a3ee683a8341852176a0334f Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Tue, 21 Jul 2026 10:08:55 -0700 Subject: [PATCH 18/44] Use contiguous wi in JAX MoE VJP --- tests/jax/test_te_ep_moe.py | 16 ++++----- transformer_engine/jax/flax/moe.py | 19 +++++------ transformer_engine/jax/moe.py | 55 +++++++++++------------------- 3 files changed, 35 insertions(+), 55 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index d08765e184..e2e192e341 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -582,8 +582,8 @@ 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"])[..., :INTER], + jnp.asarray(params_np["wi"])[..., INTER:], jnp.asarray(params_np["wo"]), num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -621,8 +621,8 @@ def loss_fn(params, x): out, _ = _pure_jax_moe_reference( x, params["gate_kernel"], - params["wi_0"], - params["wi_1"], + params["wi"][..., :INTER], + params["wi"][..., INTER:], params["wo"], ref_expert_bias, num_experts=NUM_EXPERTS, @@ -638,7 +638,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}]" @@ -708,8 +708,8 @@ 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"])[..., :INTER], + jnp.asarray(params_np["wi"])[..., INTER:], jnp.asarray(params_np["wo"]), num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -739,7 +739,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/flax/moe.py b/transformer_engine/jax/flax/moe.py index 3629346e33..064db7e695 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -200,16 +200,14 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[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( @@ -254,8 +252,7 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: return moe( inputs, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index a49963d25d..4c8c013e1d 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -217,8 +217,7 @@ def _ffn_fwd_global( recv_tokens: jnp.ndarray, recv_topk_weights: jnp.ndarray, token_counts: 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], @@ -249,8 +248,7 @@ def _ffn_fwd_global( recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) - 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) # Dispatch groups flatten in (dp, ep, local_expert) order. Keep the @@ -275,11 +273,8 @@ def _broadcast_bias(value): wi_1_bias = _broadcast_bias(wi_1_bias) wo_bias = _broadcast_bias(wo_bias) - # 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 ) @@ -293,7 +288,7 @@ def _broadcast_bias(value): ragged_scale_sharding=flat_token_sharding, ) casted_wi = tex.grouped_quantize( - wi_combined, fc1_quantizer_set.kernel, flatten_axis=-1 + wi, fc1_quantizer_set.kernel, flatten_axis=-1 ) combined_out = tex.grouped_gemm( casted_sorted_x.get_tensor(usage=TensorUsage.LHS), @@ -385,7 +380,7 @@ def _ffn_bwd_global( Mirrors :func:`_ffn_fwd_global`. Returns ``(d_sorted_x [num_procs, recv_pr, H], d_recv_w [num_procs, recv_pr], - d_wi_0, d_wi_1, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. + d_wi, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. """ group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) d_eo_2d = d_expert_outputs.reshape(-1, d_expert_outputs.shape[-1]) @@ -450,7 +445,7 @@ def _ffn_bwd_global( # 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) d_combined = jax.lax.with_sharding_constraint(d_combined, flat_token_sharding) casted_d_combined = tex.grouped_quantize( @@ -472,7 +467,6 @@ def _ffn_bwd_global( contracting_dims=((0,), (0,)), ) d_wi_combined = jax.lax.with_sharding_constraint(d_wi_combined, grouped_weight_sharding) - 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, group_sizes) d_wi_combined_bias = jax.lax.with_sharding_constraint( @@ -488,8 +482,7 @@ def _ffn_bwd_global( 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, @@ -505,8 +498,7 @@ def _ffn_bwd_global( def _moe_fwd_rule( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, @@ -715,8 +707,7 @@ def _moe_fwd_rule( recv_tokens, recv_topk_weights, token_counts, - wi_0, - wi_1, + wi, wo, wi_0_bias if has_bias else None, wi_1_bias if has_bias else None, @@ -877,8 +868,7 @@ def _moe_bwd_rule( ( 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, @@ -913,8 +903,7 @@ def _fold_dp_groups(grad): .reshape(num_experts, *grad.shape[1:]) ) - d_wi_0 = _fold_dp_groups(d_wi_0) - d_wi_1 = _fold_dp_groups(d_wi_1) + d_wi = _fold_dp_groups(d_wi) d_wo = _fold_dp_groups(d_wo) if has_bias: d_wi_0_bias = _fold_dp_groups(d_wi_0_bias) @@ -996,8 +985,7 @@ def _fold_dp_groups(grad): # 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:]) @@ -1015,8 +1003,7 @@ def _fold_dp_groups(grad): 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, @@ -1031,12 +1018,11 @@ def _fold_dp_groups(grad): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(10, 27))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 26))) def _moe( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, @@ -1064,8 +1050,7 @@ def _moe( primal, _ = _moe_fwd_rule( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, @@ -1099,8 +1084,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, @@ -1217,8 +1201,7 @@ def moe( output, aux_loss = _moe( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, From 2b9079b6a9e351df6c66eb853a35b2404b86bdcb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:10:37 +0000 Subject: [PATCH 19/44] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/jax/cpp_extensions/quantization.py | 9 +++------ transformer_engine/jax/moe.py | 10 ++++------ 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index e758600527..a888c43ded 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -100,6 +100,7 @@ def _uniform_mxfp8_scale_carrier_specs(x_spec, flatten_axis): (group_spec, column_spec, row_spec), ) + def _normalize_flatten_axis(flatten_axis, ndim): return flatten_axis + ndim if flatten_axis < 0 else flatten_axis @@ -1751,13 +1752,9 @@ def grouped_quantize( PartitionSpec(data_spec[0]), ) if q_layout.has_rowwise: - rowwise_scale_inv = jax.lax.with_sharding_constraint( - rowwise_scale_inv, scale_sharding - ) + rowwise_scale_inv = jax.lax.with_sharding_constraint(rowwise_scale_inv, scale_sharding) if q_layout.has_colwise: - colwise_scale_inv = jax.lax.with_sharding_constraint( - colwise_scale_inv, scale_sharding - ) + colwise_scale_inv = jax.lax.with_sharding_constraint(colwise_scale_inv, scale_sharding) # TODO(Phuong): store the whole updated_amax in the grouped_quantize instead? if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 4c8c013e1d..cee5a28c60 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -287,9 +287,7 @@ def _broadcast_bias(value): flatten_axis=-1, ragged_scale_sharding=flat_token_sharding, ) - casted_wi = tex.grouped_quantize( - wi, fc1_quantizer_set.kernel, 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), @@ -298,9 +296,9 @@ def _broadcast_bias(value): ) combined_out = jax.lax.with_sharding_constraint(combined_out, flat_token_sharding) 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 - ).checkpoint(fc1_quantizer_set.x) + casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint( + fc1_quantizer_set.x + ) casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint( fc1_quantizer_set.kernel ) From 294a8ec37cf17d054dc24535c75dab336608155b Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 27 Jul 2026 10:08:00 -0700 Subject: [PATCH 20/44] Revert "Remove redundant masks from TE MoE ragged paths" This reverts commit 1b783b98ba93a1e27854bf3872730925e1b7777d. --- transformer_engine/jax/moe.py | 32 +++++++++++++++++++++++--------- 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index cee5a28c60..f7fd0024dc 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -387,6 +387,11 @@ def _ffn_bwd_global( recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) fc1_quantizer_set, fc2_quantizer_set = quantizer_sets + # 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 = (group_sizes > 0)[:, None, None] + # wo bwd casted_d_eo = tex.grouped_quantize( d_eo_2d, @@ -408,6 +413,7 @@ def _ffn_bwd_global( _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) + d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) d_wo = jax.lax.with_sharding_constraint(d_wo, grouped_weight_sharding) d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None if has_bias: @@ -415,17 +421,20 @@ def _ffn_bwd_global( act_fn = _convert_to_activation_function(activation_type) if apply_topk_weights_early: - # intermediate' = intermediate * w. The subsequent dgrad and EP - # operations consume group_sizes, so they skip padded ragged rows. + # 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. w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) - gate_proj_for_bwd = gate_proj_out - up_proj_for_bwd = up_proj_out - intermediate_unweighted = act_fn(gate_proj_out) * up_proj_out + 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, axis=-1, ).astype(recv_w_flat.dtype) - d_intermediate = d_intermediate * w_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 @@ -464,6 +473,7 @@ def _ffn_bwd_global( 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 = jax.lax.with_sharding_constraint(d_wi_combined, grouped_weight_sharding) if has_bias: d_wi_combined_bias = tex.grouped_dbias(d_combined, group_sizes) @@ -855,10 +865,14 @@ def _moe_bwd_rule( d_expert_outputs = grad_pre_combine d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) else: - # Reverse the late-weighting multiply. The subsequent dgrad and EP - # operations consume group_sizes, so they skip padded ragged rows. + # 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) - d_expert_outputs = grad_pre_combine * w + 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_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) From 720cb77b836f1d12a028509f0e46962e341ff394 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 29 Jul 2026 07:03:28 -0700 Subject: [PATCH 21/44] Add handle_mem_ptr logging to ep_api.cpp --- transformer_engine/common/ep/ep_api.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index 0981289ffe..fff90dd8f5 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -75,6 +75,7 @@ size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg) { void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv_tokens_per_expert, NVTETensor total_recv_tokens_per_rank, const NVTEEpLayerConfig* layer_cfg, cudaStream_t stream) { + printf("nvte_ep_prepare handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); NVTEEpLayerConfig cfg = normalize_ep_config(layer_cfg, kLayerConfigMinSize, "layer_cfg"); EPBackend::get().prepare(handle_mem_ptr(handle_mem), topk_idx, recv_tokens_per_expert, total_recv_tokens_per_rank, cfg, stream); @@ -85,6 +86,7 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok NVTECommWindow topk_weights_win, NVTETensor recv_tokens, NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, NVTECommWindow recv_topk_weights_win, cudaStream_t stream) { + printf("nvte_ep_dispatch handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); EPBackend::get().dispatch(handle_mem_ptr(handle_mem), topk_idx, tokens, tokens_win, topk_weights, topk_weights_win, recv_tokens, recv_tokens_win, recv_topk_weights, recv_topk_weights_win, stream); @@ -92,6 +94,7 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok void nvte_ep_combine(NVTETensor handle_mem, NVTETensor expert_out, NVTECommWindow expert_out_win, NVTETensor result, cudaStream_t stream) { + printf("nvte_ep_combine handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); EPBackend::get().combine(handle_mem_ptr(handle_mem), expert_out, expert_out_win, result, stream); } @@ -99,6 +102,7 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow NVTETensor g_recv_topk_weights, NVTECommWindow g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream) { + printf("nvte_ep_dispatch_bwd handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); EPBackend::get().dispatch_bwd(handle_mem_ptr(handle_mem), grad, grad_win, g_recv_topk_weights, g_recv_topk_weights_win, grad_tokens, grad_topk_weights, stream); } @@ -106,6 +110,7 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, cudaStream_t stream) { + printf("nvte_combine_bwd handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); EPBackend::get().combine_bwd(handle_mem_ptr(handle_mem), grad, grad_win, grad_expert_out, grad_expert_out_win, stream); } From 3af791dae6bfef10bc883e76b330caf3e182aab7 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 29 Jul 2026 08:12:12 -0700 Subject: [PATCH 22/44] Additional handle_mem and cache logging --- transformer_engine/common/ep/ep_backend.cpp | 97 +++++++++++++++++++-- transformer_engine/common/ep/ep_backend.h | 7 +- 2 files changed, 94 insertions(+), 10 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index a82ec1c98d..e4704ec9d4 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include "../common.h" @@ -29,6 +30,63 @@ namespace ep { namespace { +std::atomic ep_trace_sequence{0}; +std::atomic ep_prepare_update_sequence{0}; + +bool trace_ep_handles() { + static const bool enabled = [] { + const char* value = std::getenv("NVTE_EP_TRACE_HANDLES"); + return value != nullptr && std::strcmp(value, "0") != 0; + }(); + return enabled; +} + +uint64_t fingerprint_routing_descriptor(const ncclEpTensor_t& routing) { + // This intentionally fingerprints only host-visible descriptor metadata. + // Reading routing values would require extra device work or host completion + // and could perturb CUDA graph capture or EP stream ordering. + constexpr uint64_t kFnvOffset = 1469598103934665603ULL; + constexpr uint64_t kFnvPrime = 1099511628211ULL; + uint64_t hash = kFnvOffset; + const auto mix = [&hash](uint64_t value) { + hash ^= value; + hash *= kFnvPrime; + }; + mix(static_cast(reinterpret_cast(routing.data))); + mix(static_cast(routing.ndim)); + mix(static_cast(routing.datatype)); + for (int i = 0; i < routing.ndim; ++i) { + mix(static_cast(routing.sizes[i])); + } + return hash; +} + +void trace_ep_handle(const char* op, void* handle_mem_ptr, ncclEpHandle_t handle, + cudaStream_t stream, const char* cache_state = "", + const ncclEpTensor_t* routing = nullptr, uint64_t update_id = 0) { + if (!trace_ep_handles()) return; + const uint64_t sequence = ep_trace_sequence.fetch_add(1, std::memory_order_relaxed); + const void* routing_ptr = routing == nullptr ? nullptr : routing->data; + const int routing_ndim = routing == nullptr ? 0 : routing->ndim; + const int64_t routing_dim0 = + routing == nullptr || routing->ndim < 1 ? 0 : routing->sizes[0]; + const int64_t routing_dim1 = + routing == nullptr || routing->ndim < 2 ? 0 : routing->sizes[1]; + const uint64_t routing_fingerprint = + routing == nullptr ? 0 : fingerprint_routing_descriptor(*routing); + std::fprintf(stderr, + "NVTE_EP_HANDLE_TRACE pid=%d seq=%lu op=%s handle_mem_ptr=%p " + "nccl_ep_handle=%p cache=%s stream=%p update_id=%lu " + "routing_ptr=%p routing_shape=[%ld,%ld] routing_ndim=%d " + "routing_descriptor_fingerprint=0x%016lx\n", + static_cast(getpid()), static_cast(sequence), op, handle_mem_ptr, + static_cast(handle), cache_state, static_cast(stream), + static_cast(update_id), routing_ptr, + static_cast(routing_dim0), static_cast(routing_dim1), routing_ndim, + static_cast(routing_fingerprint)); + std::fflush(stderr); +} + ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) { switch (dtype) { case kNVTEFloat32: @@ -247,7 +305,8 @@ size_t EPBackend::cache_cap_locked() { return handle_cache_cap_; } -ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg) { +ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg, + cudaStream_t stream) { // Update the program-wide fallback cfg so dispatch/combine/_bwd can // reconstruct the handle on a pointer-cache miss (WAR for XLA buffer reloc // between runs; one cfg per process). Remove this once XLA preserves the @@ -267,6 +326,8 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon auto it = index_.find(handle_mem); if (it != index_.end()) { lru_.splice(lru_.begin(), lru_, it->second); + trace_ep_handle("prepare_handle", handle_mem, it->second->handle, stream, "hit", nullptr, + it->second->last_update_id); return it->second->handle; } ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; @@ -276,7 +337,8 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon layer_cfg.top_k)); ncclEpHandle_t h = open_handle(handle_mem, hm_size, layer_cfg.top_k, layer_cfg.dispatch_output_per_expert_alignment); - lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size}); + trace_ep_handle("prepare_handle", handle_mem, h, stream, "miss"); + lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size, 0}); index_.emplace(handle_mem, lru_.begin()); while (lru_.size() > cache_cap_locked()) { HandleEntry& victim = lru_.back(); @@ -287,10 +349,14 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon return h; } -ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem) { +ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem, cudaStream_t stream, + uint64_t* last_update_id) { auto it = index_.find(handle_mem); if (it != index_.end()) { lru_.splice(lru_.begin(), lru_, it->second); + *last_update_id = it->second->last_update_id; + trace_ep_handle("lookup_handle", handle_mem, it->second->handle, stream, "hit", nullptr, + *last_update_id); return it->second->handle; } // Miss: reconstruct from the process-wide cached cfg. XLA may relocate @@ -300,7 +366,9 @@ ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem) { const uintptr_t hm_addr = reinterpret_cast(handle_mem); NVTE_CHECK(fallback_layer_cfg_.has_value(), "ep op on handle_mem=0x", hm_addr, " with no cached entry and no prior nvte_ep_prepare; call prepare first."); - return prepare_handle_locked(handle_mem, *fallback_layer_cfg_); + ncclEpHandle_t handle = prepare_handle_locked(handle_mem, *fallback_layer_cfg_, stream); + *last_update_id = 0; + return handle; } // --------------------------------------------------------------------------- @@ -344,7 +412,13 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); + ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg, stream); + const uint64_t update_id = + ep_prepare_update_sequence.fetch_add(1, std::memory_order_relaxed) + 1; + auto entry = index_.find(handle_mem); + NVTE_CHECK(entry != index_.end(), "EP handle cache entry disappeared during prepare"); + entry->second->last_update_id = update_id; + trace_ep_handle("prepare_update", handle_mem, h, stream, "", &nccl_topk_idx, update_id); NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } @@ -407,7 +481,10 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - ncclEpHandle_t h = lookup_handle_locked(handle_mem); + uint64_t update_id = 0; + ncclEpHandle_t h = lookup_handle_locked(handle_mem, stream, &update_id); + trace_ep_handle(is_forward ? "dispatch_fwd" : "combine_bwd", handle_mem, h, stream, "", nullptr, + update_id); NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, /*layout_info=*/nullptr, &dispatch_cfg, stream)); } @@ -431,7 +508,9 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - ncclEpHandle_t h = lookup_handle_locked(handle_mem); + uint64_t update_id = 0; + ncclEpHandle_t h = lookup_handle_locked(handle_mem, stream, &update_id); + trace_ep_handle("combine_fwd", handle_mem, h, stream, "", nullptr, update_id); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } @@ -469,7 +548,9 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - ncclEpHandle_t h = lookup_handle_locked(handle_mem); + uint64_t update_id = 0; + ncclEpHandle_t h = lookup_handle_locked(handle_mem, stream, &update_id); + trace_ep_handle("dispatch_bwd", handle_mem, h, stream, "", nullptr, update_id); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); } diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index 80c9b9cea3..b6eb4101b2 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -93,6 +93,7 @@ class EPBackend { ncclEpHandle_t handle; NVTEEpLayerConfig layer_cfg; size_t handle_mem_size; + uint64_t last_update_id; }; ncclEpGroup_t ep_group_{nullptr}; @@ -106,8 +107,10 @@ class EPBackend { std::optional fallback_layer_cfg_; // Caller must hold mutex_. - ncclEpHandle_t prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg); - ncclEpHandle_t lookup_handle_locked(void* handle_mem); + ncclEpHandle_t prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg, + cudaStream_t stream); + ncclEpHandle_t lookup_handle_locked(void* handle_mem, cudaStream_t stream, + uint64_t* last_update_id); size_t cache_cap_locked(); }; From c378de935d486d4e3ef11d0372a5b46a8eddf7a2 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 30 Jul 2026 07:57:57 -0700 Subject: [PATCH 23/44] Add custom-partitioned MoE diagnostics and tests --- tests/jax/test_distributed_grouped_gemm.py | 31 +- tests/jax/test_multi_process_ep.py | 302 +++++ tests/jax/test_te_ep_moe.py | 657 +++++++++- transformer_engine/jax/cpp_extensions/gemm.py | 52 +- transformer_engine/jax/moe.py | 1052 ++++++++++++++++- 5 files changed, 2022 insertions(+), 72 deletions(-) diff --git a/tests/jax/test_distributed_grouped_gemm.py b/tests/jax/test_distributed_grouped_gemm.py index 5e850094ba..a9e9d159da 100644 --- a/tests/jax/test_distributed_grouped_gemm.py +++ b/tests/jax/test_distributed_grouped_gemm.py @@ -252,14 +252,14 @@ def test_grouped_gemm_rhs_weight_specs_gather_fsdp_but_preserve_ep(): assert tuple(out_sharding[0].spec) == (None, None, None) -def test_grouped_gemm_gathers_smaller_moe_rhs_across_fsdp_group_axis(): - """A global MoE RHS has E groups while token counts have dp * E groups.""" +def test_grouped_gemm_gathers_fsdp_shared_moe_rhs_with_exact_group_ratio(): + """A global MoE RHS has E groups while token counts have fsdp * E groups.""" mesh = _mesh() arg_infos = ( _arg_info(mesh, (8192,), (None,)), _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (32, 128, 64), (("fsdp", "expert"), None, None)), - _arg_info(mesh, (2048,), (("fsdp", "expert"),)), + _arg_info(mesh, (32, 128, 64), (("expert", "fsdp"), None, None)), + _arg_info(mesh, (2048,), (("expert", "fsdp"),)), _arg_info(mesh, (0,), (None,)), _arg_info(mesh, (64,), (("fsdp", "expert"),)), _arg_info(mesh, (0,), (None,)), @@ -512,14 +512,17 @@ def test_grouped_partitioning_shardy_rules_smoke(): @pytest.mark.parametrize( - "weight_spec", + ("group_spec", "weight_spec"), [ - ("expert", "fsdp", None), - ("expert", None, "fsdp"), + ("expert", ("expert", "fsdp", None)), + ("expert", ("expert", None, "fsdp")), + (("fsdp", "expert"), (("fsdp", "expert"), None, None)), ], - ids=("contracting-fsdp", "output-fsdp"), + ids=("contracting-fsdp", "output-fsdp", "compound-fsdp-expert-groups"), ) -def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process(weight_spec): +def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process( + group_spec, weight_spec +): mesh = _mesh() n_groups = 4 group_tokens = 128 @@ -528,10 +531,10 @@ def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process(weight_spe x_shape = (n_groups * group_tokens, hidden) w_shape = (n_groups, hidden, out_hidden) - x_sharding = NamedSharding(mesh, PartitionSpec("expert", None)) + x_sharding = NamedSharding(mesh, PartitionSpec(group_spec, None)) w_sharding = NamedSharding(mesh, PartitionSpec(*weight_spec)) - group_sharding = NamedSharding(mesh, PartitionSpec("expert")) - out_sharding = NamedSharding(mesh, PartitionSpec("expert", None)) + group_sharding = NamedSharding(mesh, PartitionSpec(group_spec)) + out_sharding = NamedSharding(mesh, PartitionSpec(group_spec, None)) quantizer_set = _mxfp8_grouped_quantizer_set(n_groups) @@ -572,8 +575,8 @@ def apply(x, w): )(x, w, group_sizes) out, dx, dw = jax.block_until_ready((out, dx, dw)) - assert tuple(out.sharding.spec) == ("expert", None) - assert tuple(dx.sharding.spec) == ("expert", None) + assert tuple(out.sharding.spec) == (group_spec, None) + assert tuple(dx.sharding.spec) == (group_spec, None) assert tuple(dw.sharding.spec) == weight_spec for value in (out, dx, dw): local_value = np.asarray(jax.device_get(value.addressable_data(0))) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 0b8bb25f3f..a5190264d1 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -13,6 +13,8 @@ - ``ep_combine`` custom_vjp: ``max|grad_eo| ≈ eo_const / TOP_K`` (closed form). - ``ep_dispatch`` custom_vjp: exact per-(t, k) ``grad_topk_weights`` under skewed upstream gradients (no k-axis averaging). + - ``jax.lax.scan`` over distinct top-1 routing maps matches an unrolled + dispatch/combine reference in both forward and backward. - HLO reshard guard: compile-only, no XLA collectives outside the EP FFI. Launch via tests/jax/multi_process_launch_ep.sh (one process per GPU). @@ -295,6 +297,306 @@ def run(idx, ta_, tb_, w_): rtol=5e-2, ) + def _make_scan_inputs(self, num_layers=4): + """Top-1 inputs with a distinct random routing map for every layer.""" + num_layers = int(os.environ.get("NVTE_TEST_EP_SCAN_LAYERS", num_layers)) + num_tokens = TOKENS_PER_DP_SHARD * self.dp + rng = np.random.default_rng(seed=20260728) + routes = rng.integers( + 0, + self.num_experts, + size=(num_layers, num_tokens, 1), + dtype=np.int32, + ) + # A random draw could theoretically repeat a complete map. Make the + # invariant explicit so this test never becomes probabilistic. + seen = set() + for layer in range(num_layers): + while routes[layer].tobytes() in seen: + routes[layer] = rng.integers( + 0, self.num_experts, size=(num_tokens, 1), dtype=np.int32 + ) + seen.add(routes[layer].tobytes()) + tokens = jnp.asarray( + rng.standard_normal((num_tokens, HIDDEN_DIM), dtype=np.float32) * 0.25, + dtype=jnp.bfloat16, + ) + weights = jnp.ones((num_tokens, 1), dtype=jnp.float32) + return jnp.asarray(routes), tokens, weights + + def _scan_ep_layer(self, cfg, route, tokens, weights, slot_dependent_scale): + """One dispatch/combine-only layer used by the scan regression tests.""" + token_spec = PartitionSpec(("dp", "ep"), None) + ep_token_spec = PartitionSpec(("dp", "ep"), None, None) + ep_weight_spec = PartitionSpec(("dp", "ep"), None) + + recv_tokens, recv_weights, handle_mem, token_counts = ep_dispatch( + cfg, route, tokens, weights, self.recv_capacity_per_rank + ) + recv_tokens = jax.lax.with_sharding_constraint( + recv_tokens, NamedSharding(self.mesh, ep_token_spec) + ) + recv_weights = jax.lax.with_sharding_constraint( + recv_weights, NamedSharding(self.mesh, ep_weight_spec) + ) + + # A plain top-1 identity round-trip cannot expose a stale handle if + # dispatch and combine both use the same stale routing map: the two + # wrong permutations cancel. Scaling by packed dispatch slot makes the + # result (and its token gradient) depend on the map that was prepared. + if slot_dependent_scale: + slot = jnp.arange(recv_tokens.shape[-2], dtype=jnp.float32) + scale = 0.5 + (slot % 7) * 0.125 + expert_out = recv_tokens.astype(jnp.float32) * scale[None, :, None] + else: + expert_out = recv_tokens.astype(jnp.float32) + expert_out = jnp.where( + recv_weights[..., None] != 0, + expert_out * recv_weights[..., None], + 0.0, + ).astype(recv_tokens.dtype) + expert_out = jax.lax.with_sharding_constraint( + expert_out, NamedSharding(self.mesh, ep_token_spec) + ) + out = ep_combine( + cfg, + handle_mem, + token_counts, + expert_out, + tokens.shape[0], + out_sharding=(("dp", "ep"), None), + ) + return jax.lax.with_sharding_constraint( + out, NamedSharding(self.mesh, token_spec) + ) + + def test_scan_dispatch_combine_top1_identity(self): + """Top-1 dispatch/combine remains an exact identity across scan layers.""" + routes, tokens, weights = self._make_scan_inputs() + cfg = EpLayerConfig(top_k=1, dispatch_output_per_expert_alignment=16) + route_spec = PartitionSpec(None, ("dp", "ep"), None) + token_spec = PartitionSpec(("dp", "ep"), None) + + with self.mesh, global_shard_guard(self.mr): + routes = jax.lax.with_sharding_constraint( + routes, NamedSharding(self.mesh, route_spec) + ) + tokens_s = jax.lax.with_sharding_constraint( + tokens, NamedSharding(self.mesh, token_spec) + ) + weights_s = jax.lax.with_sharding_constraint( + weights, NamedSharding(self.mesh, token_spec) + ) + + @jax.jit + def run(route_maps, initial_tokens, topk_weights): + def body(current_tokens, route): + out = self._scan_ep_layer( + cfg, + route, + current_tokens, + topk_weights, + slot_dependent_scale=False, + ) + return out, None + + return jax.lax.scan(body, initial_tokens, route_maps)[0] + + out = run(routes, tokens_s, weights_s) + out.block_until_ready() + out_global = jmu.process_allgather(out, tiled=True) + + if self.rank == 0: + np.testing.assert_array_equal(np.asarray(out_global), np.asarray(tokens)) + + def test_scan_dispatch_only_distinct_routing_maps(self): + """Every scan iteration's packed dispatch matches an isolated dispatch. + + This is the direct stale-routing oracle: there is no combine operation + whose inverse permutation could hide a reused routing-map handle. + """ + routes, tokens, weights = self._make_scan_inputs() + cfg = EpLayerConfig(top_k=1, dispatch_output_per_expert_alignment=16) + route_spec = PartitionSpec(None, ("dp", "ep"), None) + token_spec = PartitionSpec(("dp", "ep"), None) + ep_token_spec = PartitionSpec(("dp", "ep"), None, None) + ep_weight_spec = PartitionSpec(("dp", "ep"), None) + + with self.mesh, global_shard_guard(self.mr): + routes_s = jax.lax.with_sharding_constraint( + routes, NamedSharding(self.mesh, route_spec) + ) + tokens_s = jax.lax.with_sharding_constraint( + tokens, NamedSharding(self.mesh, token_spec) + ) + weights_s = jax.lax.with_sharding_constraint( + weights, NamedSharding(self.mesh, token_spec) + ) + + def dispatch_one(route, input_tokens, topk_weights): + recv_tokens, recv_weights, _handle_mem, token_counts = ep_dispatch( + cfg, + route, + input_tokens, + topk_weights, + self.recv_capacity_per_rank, + ) + recv_tokens = jax.lax.with_sharding_constraint( + recv_tokens, NamedSharding(self.mesh, ep_token_spec) + ) + recv_weights = jax.lax.with_sharding_constraint( + recv_weights, NamedSharding(self.mesh, ep_weight_spec) + ) + # NCCL EP leaves capacity padding unspecified, including the + # weight buffer. Derive the valid prefix from token_counts. + align = max(int(cfg.dispatch_output_per_expert_alignment), 1) + padded_counts = ((token_counts + align - 1) // align) * align + valid_count = jnp.sum(padded_counts, axis=-1, keepdims=True) + slot = jnp.arange(self.recv_capacity_per_rank, dtype=jnp.int32) + valid = slot[None, :] < valid_count + recv_tokens = jnp.where( + valid[..., None], recv_tokens, jnp.zeros_like(recv_tokens) + ) + recv_weights = jnp.where( + valid, recv_weights, jnp.zeros_like(recv_weights) + ) + return recv_tokens, recv_weights + + @jax.jit + def scan_dispatch(route_maps, input_tokens, topk_weights): + def body(carry, route): + return carry, dispatch_one(route, input_tokens, topk_weights) + + return jax.lax.scan(body, (), route_maps)[1] + + scan_tokens, scan_weights = scan_dispatch(routes_s, tokens_s, weights_s) + scan_tokens.block_until_ready() + scan_tokens_global = jmu.process_allgather(scan_tokens, tiled=True) + scan_weights_global = jmu.process_allgather(scan_weights, tiled=True) + + # Run each route in a separate executable invocation so its + # expected packed layout cannot be affected by another layer's + # live handle. + isolated_dispatch = jax.jit(dispatch_one) + ref_tokens = [] + ref_weights = [] + for layer in range(routes.shape[0]): + layer_tokens, layer_weights = isolated_dispatch( + routes_s[layer], tokens_s, weights_s + ) + layer_tokens.block_until_ready() + ref_tokens.append( + np.asarray(jmu.process_allgather(layer_tokens, tiled=True)) + ) + ref_weights.append( + np.asarray(jmu.process_allgather(layer_weights, tiled=True)) + ) + + if self.rank == 0: + np.testing.assert_array_equal( + np.asarray(scan_tokens_global), np.stack(ref_tokens) + ) + np.testing.assert_array_equal( + np.asarray(scan_weights_global), np.stack(ref_weights) + ) + + def test_scan_dispatch_combine_routing_handle_fwd_bwd(self): + """Scan must match unrolled layers when every iteration has a new map. + + Slot-dependent scaling prevents a stale dispatch/combine handle from + cancelling as it does for an identity expert. Comparing VJPs also + exercises ep_combine_bwd and ep_dispatch_bwd through the scan transpose. + """ + routes, tokens, weights = self._make_scan_inputs() + num_layers = routes.shape[0] + cfg = EpLayerConfig(top_k=1, dispatch_output_per_expert_alignment=16) + route_spec = PartitionSpec(None, ("dp", "ep"), None) + token_spec = PartitionSpec(("dp", "ep"), None) + + with self.mesh, global_shard_guard(self.mr): + routes_s = jax.lax.with_sharding_constraint( + routes, NamedSharding(self.mesh, route_spec) + ) + tokens_s = jax.lax.with_sharding_constraint( + tokens, NamedSharding(self.mesh, token_spec) + ) + weights_s = jax.lax.with_sharding_constraint( + weights, NamedSharding(self.mesh, token_spec) + ) + cotangent = jnp.asarray( + np.linspace( + 0.25, + 1.25, + tokens.size, + dtype=np.float32, + ).reshape(tokens.shape), + dtype=tokens.dtype, + ) + cotangent = jax.lax.with_sharding_constraint( + cotangent, NamedSharding(self.mesh, token_spec) + ) + + def scan_fwd(route_maps, initial_tokens, topk_weights): + def body(current_tokens, route): + out = self._scan_ep_layer( + cfg, + route, + current_tokens, + topk_weights, + slot_dependent_scale=True, + ) + return out, None + + return jax.lax.scan(body, initial_tokens, route_maps)[0] + + def unrolled_fwd(route_maps, initial_tokens, topk_weights): + out = initial_tokens + for layer in range(num_layers): + out = self._scan_ep_layer( + cfg, + route_maps[layer], + out, + topk_weights, + slot_dependent_scale=True, + ) + return out + + def value_and_token_vjp( + fn, route_maps, initial_tokens, topk_weights, out_cotangent + ): + out, pullback = jax.vjp( + lambda x: fn(route_maps, x, topk_weights), initial_tokens + ) + return out, pullback(out_cotangent)[0] + + scan_run = jax.jit( + lambda r, t, w, g: value_and_token_vjp(scan_fwd, r, t, w, g) + ) + unrolled_run = jax.jit( + lambda r, t, w, g: value_and_token_vjp(unrolled_fwd, r, t, w, g) + ) + scan_out, scan_grad = scan_run(routes_s, tokens_s, weights_s, cotangent) + ref_out, ref_grad = unrolled_run(routes_s, tokens_s, weights_s, cotangent) + scan_grad.block_until_ready() + ref_grad.block_until_ready() + + scan_out_global = jmu.process_allgather(scan_out, tiled=True) + ref_out_global = jmu.process_allgather(ref_out, tiled=True) + scan_grad_global = jmu.process_allgather(scan_grad, tiled=True) + ref_grad_global = jmu.process_allgather(ref_grad, tiled=True) + + if self.rank == 0: + self.assertFalse( + np.array_equal(np.asarray(ref_out_global), np.asarray(tokens)), + "slot-sensitive reference unexpectedly collapsed to an identity", + ) + np.testing.assert_array_equal( + np.asarray(scan_out_global), np.asarray(ref_out_global) + ) + np.testing.assert_array_equal( + np.asarray(scan_grad_global), np.asarray(ref_grad_global) + ) + def test_primitive_prepare(self): """ep_prepare returns token_counts and handle_mem of the expected shapes.""" T_global, topk_idx, _tokens, _w = self._make_identity_inputs() diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index e2e192e341..fc3d20b07e 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -118,7 +118,11 @@ def _read_mp_options(): ) 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.jax.sharding import MeshResource, global_shard_guard @@ -129,7 +133,8 @@ def _read_mp_options(): EP_AXIS = "ep" FSDP_AXIS = "fsdp" -EP_SIZE = 2 +EP_SIZE = int(os.environ.get("TE_EP_MOE_EP_SIZE", "2")) +assert EP_SIZE in (2, 4), f"TE_EP_MOE_EP_SIZE must be 2 or 4, got {EP_SIZE}" assert ( jax.device_count() % EP_SIZE == 0 ), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" @@ -138,7 +143,11 @@ def _read_mp_options(): LOGICAL_AXIS_RULES = ( ("exp", EP_AXIS), + # Match MaxText's converging layout: FSDP is the outer component of + # the compound expert dimension and EP is inner. + ("exp_fsdp", (FSDP_AXIS, EP_AXIS)), ("embed", FSDP_AXIS), + ("embed_replicated", None), ("mlp", None), ("batch", (EP_AXIS, FSDP_AXIS)), ) @@ -167,6 +176,25 @@ def _read_mp_options(): # (slot alignment rounding, etc.). TE_TO_TE_ATOL = 5e-3 TE_TO_TE_RTOL = 5e-3 +TE_TO_TE_GRAD_NORM_RATIO = (0.98, 1.02) +TE_TO_TE_GRAD_COSINE = 0.999 + + +def _assert_gradient_direction_and_scale(actual, expected, *, name): + """Catch permutations/reduction factors hidden by bf16 absolute tolerances.""" + actual = np.asarray(actual, dtype=np.float64).reshape(-1) + expected = np.asarray(expected, dtype=np.float64).reshape(-1) + actual_norm = np.linalg.norm(actual) + expected_norm = np.linalg.norm(expected) + assert actual_norm > 0.0 and expected_norm > 0.0, f"{name}: zero gradient norm" + norm_ratio = actual_norm / expected_norm + cosine = np.dot(actual, expected) / (actual_norm * expected_norm) + assert 0.8 <= norm_ratio <= 1.2, ( + f"{name}: gradient norm ratio {norm_ratio:.6f} outside [0.8, 1.2] " + f"(cosine={cosine:.6f})" + ) + assert cosine >= 0.98, f"{name}: gradient cosine similarity {cosine:.6f} < 0.98" + # 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 @@ -194,10 +222,16 @@ def _compute_worst_case_recv_pr(): 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 + 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 + num_local_experts + * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) + * _ALIGN_SIZE ) return min(per_expert_bound, aligned_total_bound) @@ -222,7 +256,9 @@ 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(), @@ -312,7 +348,9 @@ 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 @@ -324,7 +362,9 @@ def _pure_jax_moe_reference( # 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_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: @@ -339,7 +379,9 @@ 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] @@ -364,6 +406,7 @@ def _make_block( use_expert_routing_bias=False, score_function="softmax", expert_bias_init=None, + compound_expert_sharding=False, ): kwargs = dict( num_experts=NUM_EXPERTS, @@ -376,6 +419,11 @@ def _make_block( score_function=score_function, dtype=DTYPE, ) + if compound_expert_sharding: + # Match MaxText shard_exp_on_fsdp=True: FSDP and EP both shard the + # expert group axis while the hidden dimensions remain replicated. + kwargs["wi_kernel_axes"] = ("exp_fsdp", "embed_replicated", "mlp") + kwargs["wo_kernel_axes"] = ("exp_fsdp", "mlp", "embed_replicated") # Custom expert_bias_init lets tests inject a non-zero expert_bias without # poking variables['params'] post-init. if expert_bias_init is not None: @@ -501,6 +549,279 @@ def _make_inputs(key): return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) +def _make_stacked_moe_params(mesh, num_layers): + """Create distinct production-shaped MoE parameters with a scan axis. + + The expert parameters use the same compound ``(fsdp, ep)`` expert + sharding as MaxText's ``shard_exp_on_fsdp=True`` integration. A rotating + strong routing bias makes each layer's routing map distinct while also + leaving half of the experts empty in every layer. + """ + gate_key, wi_key, wo_key = jax.random.split(jax.random.PRNGKey(101), 3) + + with _ctx(mesh): + + @jax.jit + def initialize(): + gate_kernel = jax.random.normal( + gate_key, + (num_layers, HIDDEN, NUM_EXPERTS), + dtype=DTYPE, + ) / np.sqrt(HIDDEN) + wi = jax.random.normal( + wi_key, + (num_layers, NUM_EXPERTS, HIDDEN, 2 * INTER), + dtype=DTYPE, + ) / np.sqrt(HIDDEN) + wo = jax.random.normal( + wo_key, + (num_layers, NUM_EXPERTS, INTER, HIDDEN), + dtype=DTYPE, + ) / np.sqrt(INTER) + base_bias = jnp.concatenate( + ( + jnp.full((NUM_EXPERTS // 2,), 10.0, dtype=jnp.float32), + jnp.full( + (NUM_EXPERTS - NUM_EXPERTS // 2,), + -10.0, + dtype=jnp.float32, + ), + ) + ) + expert_bias = jnp.stack( + [ + jnp.roll(base_bias, (2 * layer) % NUM_EXPERTS) + for layer in range(num_layers) + ] + ) + + gate_kernel = jax.lax.with_sharding_constraint( + gate_kernel, NamedSharding(mesh, P(None, None, None)) + ) + wi = jax.lax.with_sharding_constraint( + wi, + NamedSharding( + mesh, + P(None, (FSDP_AXIS, EP_AXIS), None, None), + ), + ) + wo = jax.lax.with_sharding_constraint( + wo, + NamedSharding( + mesh, + P(None, (FSDP_AXIS, EP_AXIS), None, None), + ), + ) + expert_bias = jax.lax.with_sharding_constraint( + expert_bias, NamedSharding(mesh, P(None, None)) + ) + return { + "gate_kernel": gate_kernel, + "wi": wi, + "wo": wo, + "expert_bias": expert_bias, + } + + params = initialize() + jax.block_until_ready(params["wi"]) + return params + + +def _functional_production_moe_layer(params, x): + """One residual TE MoE layer matching the failing integration knobs.""" + normed_x = ( + x.astype(jnp.float32) + * jax.lax.rsqrt( + jnp.mean(x.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + 1.0e-6 + ) + ).astype(x.dtype) + branch, _ = moe( + normed_x, + params["gate_kernel"], + params["wi"], + params["wo"], + expert_bias=params["expert_bias"], + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + score_function="sigmoid", + apply_topk_weights_early=False, + ep_axis=EP_AXIS, + data_parallelism_axes=(FSDP_AXIS,), + wi_kernel_axes=("exp_fsdp", "embed_replicated", "mlp"), + wo_kernel_axes=("exp_fsdp", "mlp", "embed_replicated"), + dtype=DTYPE, + ) + return x + branch + + +def _run_stacked_te_moe(params, x, *, use_scan, remat): + """Run identical per-layer parameters scanned or Python-unrolled.""" + layer_fn = _functional_production_moe_layer + if remat: + # MaxText applies nn.remat to its layer before passing the layer to + # nn.scan. Checkpointing the functional body gives the same important + # lowering: forward recomputation occurs inside reverse scan. + layer_fn = jax.checkpoint(layer_fn, prevent_cse=True) + + if use_scan: + + def scan_body(value, layer_params): + return layer_fn(layer_params, value), None + + return jax.lax.scan(scan_body, x, params)[0] + + value = x + for layer in range(params["gate_kernel"].shape[0]): + layer_params = jax.tree_util.tree_map(lambda p: p[layer], params) + value = layer_fn(layer_params, value) + return value + + +def _run_stacked_jax_reference(params, x): + """Pure-JAX residual stack using the same distinct layer parameters.""" + value = x + for layer in range(params["gate_kernel"].shape[0]): + layer_params = jax.tree_util.tree_map(lambda p: p[layer], params) + normed_value = ( + value.astype(jnp.float32) + * jax.lax.rsqrt( + jnp.mean(value.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + + 1.0e-6 + ) + ).astype(value.dtype) + branch, _ = _pure_jax_moe_reference( + normed_value, + layer_params["gate_kernel"], + layer_params["wi"][..., :INTER], + layer_params["wi"][..., INTER:], + layer_params["wo"], + layer_params["expert_bias"], + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + score_function="sigmoid", + ) + value = value + branch + return value + + +def _stack_value_and_grad(params, x, *, use_scan, remat): + def loss_fn(p, value): + output = _run_stacked_te_moe( + p, + value, + use_scan=use_scan, + remat=remat, + ) + return jnp.mean(output.astype(jnp.float32) ** 2), output + + return jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True)(params, x) + + +def _gradient_similarity(actual, expected): + actual = np.asarray(actual, dtype=np.float64).reshape(-1) + expected = np.asarray(expected, dtype=np.float64).reshape(-1) + actual_norm = np.linalg.norm(actual) + expected_norm = np.linalg.norm(expected) + norm_ratio = actual_norm / expected_norm + cosine = np.dot(actual, expected) / (actual_norm * expected_norm) + return actual_norm, expected_norm, norm_ratio, cosine + + +def _assert_te_to_te_gradient(actual, expected, *, name): + """Strict scan-vs-unrolled oracle with useful failure diagnostics.""" + actual = np.asarray(actual) + expected = np.asarray(expected) + assert np.all(np.isfinite(actual)), f"{name}: scan gradient has NaN/Inf" + assert np.all(np.isfinite(expected)), f"{name}: unrolled gradient has NaN/Inf" + actual_norm, expected_norm, norm_ratio, cosine = _gradient_similarity( + actual, expected + ) + assert actual_norm > 0.0 and expected_norm > 0.0, ( + f"{name}: zero gradient norm " + f"(scan={actual_norm:.9e}, unrolled={expected_norm:.9e})" + ) + lo, hi = TE_TO_TE_GRAD_NORM_RATIO + assert lo <= norm_ratio <= hi and cosine >= TE_TO_TE_GRAD_COSINE, ( + f"{name}: scan/unrolled gradient direction or scale mismatch: " + f"scan_norm={actual_norm:.9e}, unrolled_norm={expected_norm:.9e}, " + f"ratio={norm_ratio:.9f}, cosine={cosine:.9f}, " + f"scan_mean={actual.astype(np.float64).mean():.9e}, " + f"scan_std={actual.astype(np.float64).std():.9e}, " + f"scan_absmax={np.abs(actual.astype(np.float64)).max():.9e}, " + f"unrolled_mean={expected.astype(np.float64).mean():.9e}, " + f"unrolled_std={expected.astype(np.float64).std():.9e}, " + f"unrolled_absmax={np.abs(expected.astype(np.float64)).max():.9e}" + ) + np.testing.assert_allclose( + actual.astype(np.float32), + expected.astype(np.float32), + atol=TE_TO_TE_ATOL, + rtol=TE_TO_TE_RTOL, + err_msg=f"{name}: scan/unrolled elementwise mismatch", + ) + + +def _assert_stacked_reference_gradient(actual, expected, *, name): + """Broad TE-vs-JAX stack control; strict parity is tested per block.""" + actual = np.asarray(actual) + expected = np.asarray(expected) + assert np.all(np.isfinite(actual)), f"{name}: TE gradient has NaN/Inf" + assert np.all(np.isfinite(expected)), f"{name}: reference gradient has NaN/Inf" + actual_norm, expected_norm, norm_ratio, cosine = _gradient_similarity( + actual, expected + ) + assert actual_norm > 0.0 and expected_norm > 0.0, ( + f"{name}: zero gradient norm " + f"(TE={actual_norm:.9e}, reference={expected_norm:.9e})" + ) + assert 0.5 <= norm_ratio <= 1.5 and cosine >= 0.85, ( + f"{name}: gross TE/reference stack mismatch: " + f"TE_norm={actual_norm:.9e}, reference_norm={expected_norm:.9e}, " + f"ratio={norm_ratio:.9f}, cosine={cosine:.9f}" + ) + + +def _assert_distinct_layer_routes(params_np, x_np): + """Prove the test does not accidentally reuse one routing map.""" + value = jnp.asarray(x_np) + signatures = [] + for layer in range(params_np["gate_kernel"].shape[0]): + gate = jnp.asarray(params_np["gate_kernel"][layer]) + bias = jnp.asarray(params_np["expert_bias"][layer]) + normed_value = ( + value.astype(jnp.float32) + * jax.lax.rsqrt( + jnp.mean(value.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + + 1.0e-6 + ) + ).astype(value.dtype) + logits = jnp.einsum("bsh,he->bse", normed_value, gate).astype(jnp.float32) + _, indices = jax.lax.top_k(jax.nn.sigmoid(logits) + bias, TOPK) + indices_np = np.asarray(jax.device_get(indices), dtype=np.int64) + position = np.arange(indices_np.size, dtype=np.int64).reshape(indices_np.shape) + signatures.append( + ( + int(indices_np.sum()), + int((indices_np * (position + 1)).sum()), + ) + ) + branch, _ = _pure_jax_moe_reference( + normed_value, + gate, + jnp.asarray(params_np["wi"][layer])[..., :INTER], + jnp.asarray(params_np["wi"][layer])[..., INTER:], + jnp.asarray(params_np["wo"][layer]), + bias, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + score_function="sigmoid", + ) + value = value + branch + assert ( + len(set(signatures)) > 1 + ), f"all layers used one routing signature: {signatures}" + + # ----------------------------------------------------------------------------- # Tests # ----------------------------------------------------------------------------- @@ -523,6 +844,18 @@ def _make_inputs(key): dict(score_function="softmax", apply_topk_weights_early=True), id="softmax-early-weighting", ), + pytest.param( + dict(score_function="softmax", compound_expert_sharding=True), + id="softmax-compound-fsdp-expert", + ), + pytest.param( + dict( + score_function="sigmoid", + apply_topk_weights_early=True, + compound_expert_sharding=True, + ), + id="sigmoid-early-weighting-compound-fsdp-expert", + ), pytest.param( dict(score_function="sigmoid"), id="sigmoid", @@ -641,8 +974,12 @@ def loss_fn(params, x): 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}]" - assert np.any(g_te != 0.0), f"{name} grad identically zero [config={config}]" + 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" @@ -655,6 +992,11 @@ def loss_fn(params, x): rtol=rtol, err_msg=f"grad parity breach on {name} [config={config}]", ) + _assert_gradient_direction_and_scale( + g_te, + grads_ref_np[name], + name=f"{name} [config={config}]", + ) # d_x: the gradient propagated back to the previous layer. Checks # shape, dtype (must match x.dtype — protects the @@ -677,6 +1019,291 @@ def loss_fn(params, x): rtol=GRAD_FFN_RTOL, err_msg=f"d_x parity breach [config={config}]", ) + _assert_gradient_direction_and_scale( + grad_x_te_np, + grad_x_ref_np, + name=f"d_x [config={config}]", + ) + + def test_repeated_production_block_backward(self, mesh): + """Catch small d_x errors that amplify across a stack of MoE blocks.""" + repeats = 8 + config = dict( + score_function="sigmoid", + apply_topk_weights_early=True, + compound_expert_sharding=True, + use_expert_routing_bias=True, + ) + block = _make_block(**config) + x = _make_inputs(jax.random.PRNGKey(22)) + variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(23)) + + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + + def te_loss_fn(variables, value): + for _ in range(repeats): + branch, _ = block.apply(variables, value) + value = value + branch + return jnp.mean(value.astype(jnp.float32) ** 2) + + grads_te, grad_x_te = jax.jit(jax.grad(te_loss_fn, argnums=(0, 1)))( + variables, + x_sh, + ) + jax.block_until_ready(grad_x_te) + + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + expert_bias = jnp.asarray(params_np["expert_bias"]) + + def reference_loss_fn(params, value): + for _ in range(repeats): + branch, _ = _pure_jax_moe_reference( + value, + params["gate_kernel"], + params["wi"][..., :INTER], + params["wi"][..., INTER:], + params["wo"], + expert_bias, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + score_function="sigmoid", + ) + value = value + branch + return jnp.mean(value.astype(jnp.float32) ** 2) + + grads_ref, grad_x_ref = jax.jit(jax.grad(reference_loss_fn, argnums=(0, 1)))( + {k: jnp.asarray(v) for k, v in params_np.items() if k != "expert_bias"}, + jnp.asarray(x_np), + ) + + for name in ("gate_kernel", "wi", "wo"): + _assert_gradient_direction_and_scale( + _to_global_numpy(_unwrap(grads_te["params"][name]), mesh), + np.asarray(jax.device_get(grads_ref[name])), + name=f"repeated {name}", + ) + _assert_gradient_direction_and_scale( + _to_global_numpy(grad_x_te, mesh), + np.asarray(jax.device_get(grad_x_ref)), + name="repeated d_x", + ) + + @pytest.mark.parametrize("num_layers", (4, 8)) + def test_scanned_production_block_backward(self, mesh, num_layers): + """Full MoE scan/remat must match the identical unrolled stack. + + Unlike the primitive scan tests, this exercises top-2 routing, + compound FSDP/EP expert parameters, grouped-GEMM weight gradients, + the full ``moe`` custom VJP residual, and forward rematerialization + inside the reverse scan. + """ + params = _make_stacked_moe_params(mesh, num_layers) + x = _make_inputs(jax.random.PRNGKey(102)) + + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + unrolled_run = jax.jit( + partial( + _stack_value_and_grad, + use_scan=False, + remat=False, + ) + ) + scan_run = jax.jit( + partial( + _stack_value_and_grad, + use_scan=True, + remat=False, + ) + ) + scan_remat_run = jax.jit( + partial( + _stack_value_and_grad, + use_scan=True, + remat=True, + ) + ) + + unrolled = unrolled_run(params, x_sh) + jax.block_until_ready(unrolled) + scanned = scan_run(params, x_sh) + jax.block_until_ready(scanned) + scanned_remat = scan_remat_run(params, x_sh) + jax.block_until_ready(scanned_remat) + + (unrolled_loss, unrolled_output), (unrolled_grads, unrolled_grad_x) = unrolled + modes = ( + ("scan", scanned), + ("scan-remat", scanned_remat), + ) + unrolled_output_np = _to_global_numpy(unrolled_output, mesh) + unrolled_grad_x_np = _to_global_numpy(unrolled_grad_x, mesh) + unrolled_grads_np = { + name: _to_global_numpy(grad, mesh) for name, grad in unrolled_grads.items() + } + + params_np = { + name: _to_global_numpy(param, mesh) for name, param in params.items() + } + x_np = np.asarray(jax.device_get(x)) + if jax.process_index() == 0: + _assert_distinct_layer_routes(params_np, x_np) + + mismatches = [] + + def record_check(check): + try: + check() + except AssertionError as error: + mismatches.append(str(error)) + + for mode_name, mode_result in modes: + (mode_loss, mode_output), (mode_grads, mode_grad_x) = mode_result + record_check( + lambda: np.testing.assert_allclose( + np.asarray(jax.device_get(mode_loss), dtype=np.float32), + np.asarray(jax.device_get(unrolled_loss), dtype=np.float32), + atol=TE_TO_TE_ATOL, + rtol=TE_TO_TE_RTOL, + err_msg=f"{mode_name}: loss mismatch", + ) + ) + mode_output_np = _to_global_numpy(mode_output, mesh) + record_check( + lambda: np.testing.assert_allclose( + mode_output_np.astype(np.float32), + unrolled_output_np.astype(np.float32), + atol=TE_TO_TE_ATOL, + rtol=TE_TO_TE_RTOL, + err_msg=f"{mode_name}: output mismatch", + ) + ) + mode_grad_x_np = _to_global_numpy(mode_grad_x, mesh) + record_check( + lambda: _assert_te_to_te_gradient( + mode_grad_x_np, + unrolled_grad_x_np, + name=f"{mode_name} d_x", + ) + ) + for param_name in ("gate_kernel", "wi", "wo"): + mode_grad_np = _to_global_numpy(mode_grads[param_name], mesh) + record_check( + lambda mode_grad_np=mode_grad_np, param_name=param_name: ( + _assert_te_to_te_gradient( + mode_grad_np, + unrolled_grads_np[param_name], + name=f"{mode_name} {param_name}", + ) + ) + ) + for layer in range(num_layers): + record_check( + lambda layer=layer, mode_grad_np=mode_grad_np, param_name=param_name: ( + _assert_te_to_te_gradient( + mode_grad_np[layer], + unrolled_grads_np[param_name][layer], + name=f"{mode_name} layer={layer} {param_name}", + ) + ) + ) + + # The unrolled TE execution is the control. Check it against a + # communication-free JAX reference so a scan/unrolled agreement cannot + # hide a bug shared by both TE executions. + reference_params = { + name: jnp.asarray(value) for name, value in params_np.items() + } + + def reference_loss_fn(p, value): + output = _run_stacked_jax_reference(p, value) + return jnp.mean(output.astype(jnp.float32) ** 2) + + reference_grads, reference_grad_x = jax.jit( + jax.grad(reference_loss_fn, argnums=(0, 1)) + )(reference_params, jnp.asarray(x_np)) + for param_name in ("gate_kernel", "wi", "wo"): + record_check( + lambda param_name=param_name: _assert_stacked_reference_gradient( + unrolled_grads_np[param_name], + np.asarray(jax.device_get(reference_grads[param_name])), + name=f"unrolled-reference {param_name}", + ) + ) + record_check( + lambda: _assert_stacked_reference_gradient( + unrolled_grad_x_np, + np.asarray(jax.device_get(reference_grad_x)), + name="unrolled-reference d_x", + ) + ) + if mismatches: + pytest.fail("\n\n".join(mismatches)) + + def test_scanned_training_trajectory(self, mesh): + """Three SGD steps must retain scan/remat vs unrolled parity.""" + num_layers = 4 + learning_rate = 1.0e-3 + params = _make_stacked_moe_params(mesh, num_layers) + x = _make_inputs(jax.random.PRNGKey(103)) + + def make_step(*, use_scan, remat): + def loss_fn(p, value): + output = _run_stacked_te_moe( + p, + value, + use_scan=use_scan, + remat=remat, + ) + return jnp.mean(output.astype(jnp.float32) ** 2) + + def step(p, value): + loss, grads = jax.value_and_grad(loss_fn)(p, value) + updated = jax.tree_util.tree_map( + lambda weight, grad: weight - learning_rate * grad, + p, + grads, + ) + return updated, loss + + return jax.jit(step) + + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + unrolled_step = make_step(use_scan=False, remat=False) + scan_remat_step = make_step(use_scan=True, remat=True) + unrolled_params = params + scan_remat_params = params + unrolled_losses = [] + scan_remat_losses = [] + for _ in range(3): + unrolled_params, unrolled_loss = unrolled_step(unrolled_params, x_sh) + scan_remat_params, scan_remat_loss = scan_remat_step( + scan_remat_params, x_sh + ) + jax.block_until_ready(unrolled_params) + jax.block_until_ready(scan_remat_params) + unrolled_losses.append(float(jax.device_get(unrolled_loss))) + scan_remat_losses.append(float(jax.device_get(scan_remat_loss))) + + np.testing.assert_allclose( + np.asarray(scan_remat_losses, dtype=np.float32), + np.asarray(unrolled_losses, dtype=np.float32), + atol=TE_TO_TE_ATOL, + rtol=TE_TO_TE_RTOL, + err_msg=( + "scan-remat training trajectory differs from unrolled: " + f"scan={scan_remat_losses}, unrolled={unrolled_losses}" + ), + ) + for param_name in ("gate_kernel", "wi", "wo"): + _assert_te_to_te_gradient( + _to_global_numpy(scan_remat_params[param_name], mesh), + _to_global_numpy(unrolled_params[param_name], mesh), + name=f"three-step parameter {param_name}", + ) class TestTeEpMoeAuxLoss: @@ -727,7 +1354,9 @@ 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" @@ -740,6 +1369,8 @@ 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", "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" diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index a9805bf282..18a39d4e41 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -6,6 +6,7 @@ import math import operator import os +import sys from collections.abc import Iterable from dataclasses import dataclass from functools import partial, reduce, cache @@ -83,6 +84,16 @@ num_cublas_streams = get_num_compute_streams() +_debug_python_patch = os.getenv("NVTE_DEBUG_PYTHON_PATCH", "0") == "1" +_debug_grouped_gemm_partition_count = 0 +if _debug_python_patch: + print( + "[TE patch debug] imported transformer_engine.jax.cpp_extensions.gemm " + f"from {__file__} (rank={os.getenv('SLURM_PROCID', 'unknown')})", + file=sys.stderr, + flush=True, + ) + # Cache whether the CUDA-graphable grouped GEMM implementation is available at import time. # Calling get_grouped_gemm_setup_workspace_size raises a RuntimeError mentioning "cublas" when # compiled against cuBLAS < 13.2, in which case the cuda-graphable path is unavailable. @@ -1865,27 +1876,29 @@ def _parse_partition_specs( # A compound leading group dimension may use FSDP as the outer # data-parallel axis (for example MoE groups ordered - # (fsdp, ep, local_expert)). Normally that means FSDP describes - # distinct groups, not a sharded RHS contracting dimension. MoE - # model weights are the exception: their global expert-group axis - # is smaller than the active token-group array, so FSDP shards a - # shared RHS that must be gathered locally. Keep that gather inside - # this custom partitioning boundary, after grouped quantization. + # (fsdp, ep, local_expert)). Normally FSDP then describes distinct + # groups. MoE model weights are the exception: they retain one + # global expert set while the active group array has one copy per + # FSDP rank. Gather that shared, already-quantized RHS only when the + # group-count ratio exactly matches the FSDP mesh size. The exact + # ratio prevents unrelated unequal group shapes from selecting this + # specialized mapping. fsdp_is_group_axis = spec_contains_axis(active_group_spec, fsdp_axis) active_group_count = next( (info.shape[0] for info in grouped_dim_infos if info.size > 0), None, ) rhs_group_count = arg_infos[2].shape[0] if len(arg_infos[2].shape) > 0 else None - rhs_has_fewer_groups = ( - active_group_count is not None + rhs_is_fsdp_shared_group_set = ( + fsdp_axis is not None + and active_group_count is not None and rhs_group_count is not None - and rhs_group_count < active_group_count + and rhs_group_count * mesh.shape[fsdp_axis] == active_group_count ) gather_rhs_fsdp = ( fsdp_axis is not None and not rhs_is_ragged - and (not fsdp_is_group_axis or rhs_has_fewer_groups) + and (not fsdp_is_group_axis or rhs_is_fsdp_shared_group_set) and ( spec_contains_axis(rhs_data_spec, fsdp_axis) or spec_contains_axis(rhs_scale_spec, fsdp_axis) @@ -1893,6 +1906,25 @@ def _parse_partition_specs( ) ) + global _debug_grouped_gemm_partition_count + if _debug_python_patch and _debug_grouped_gemm_partition_count < 20: + _debug_grouped_gemm_partition_count += 1 + print( + "[TE patch debug] GroupedGemmPrimitive._parse_partition_specs " + f"call={_debug_grouped_gemm_partition_count} " + f"active_group_spec={active_group_spec} " + f"active_group_count={active_group_count} " + f"rhs_group_count={rhs_group_count} " + f"fsdp_axis={fsdp_axis!r} " + f"fsdp_mesh_size={mesh.shape.get(fsdp_axis) if fsdp_axis is not None else None} " + f"rhs_data_spec_before={rhs_data_spec} " + f"rhs_scale_spec_before={rhs_scale_spec} " + f"rhs_is_fsdp_shared_group_set={rhs_is_fsdp_shared_group_set} " + f"gather_rhs_fsdp={gather_rhs_fsdp}", + file=sys.stderr, + flush=True, + ) + if gather_rhs_fsdp: rhs_data_spec = strip_axis_from_spec(rhs_data_spec, fsdp_axis) rhs_scale_spec = strip_axis_from_spec(rhs_scale_spec, fsdp_axis) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index f7fd0024dc..53726abf63 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -30,9 +30,11 @@ ``aux_loss_coeff`` and ``expert_bias`` are also supported. """ +import os +import sys +import warnings from functools import partial from typing import Any, Optional, Tuple, Union -import warnings import flax.struct import jax @@ -60,6 +62,347 @@ # same 128-token tile, so a single constant covers every supported path. _ALIGN_SIZE = 128 +_debug_python_patch = os.getenv("NVTE_DEBUG_PYTHON_PATCH", "0") == "1" +_debug_moe_numerics = os.getenv("NVTE_DEBUG_MOE_NUMERICS", "0") == "1" +_debug_moe_input_grad = os.getenv("NVTE_DEBUG_MOE_INPUT_GRAD", "0") == "1" +_use_reference_fwd = os.getenv("NVTE_MOE_REFERENCE_FWD", "0") == "1" +_use_reference_dgrad = os.getenv("NVTE_MOE_REFERENCE_DGRAD", "0") == "1" +_zero_dispatch_weight_grad = os.getenv("NVTE_MOE_ZERO_DISPATCH_WEIGHT_GRAD", "0") == "1" +_refresh_ep_handle_in_bwd = os.getenv("NVTE_MOE_REFRESH_EP_HANDLE_IN_BWD", "0") == "1" +_refresh_ep_handle_before_combine = ( + os.getenv("NVTE_MOE_REFRESH_EP_HANDLE_BEFORE_COMBINE", "0") == "1" +) +_validate_ep_routing = os.getenv("NVTE_MOE_VALIDATE_EP_ROUTING", "0") == "1" +_debug_moe_fwd_recompute = os.getenv("NVTE_MOE_DEBUG_FWD_RECOMPUTE", "0") == "1" +_validate_ep_token_roundtrip = ( + os.getenv("NVTE_MOE_VALIDATE_EP_TOKEN_ROUNDTRIP", "0") == "1" +) +_zero_moe_input_grad = os.getenv("NVTE_MOE_ZERO_INPUT_GRAD", "0") == "1" +_skip_moe_backward = os.getenv("NVTE_MOE_SKIP_BACKWARD", "0") == "1" +_debug_handle_mem = os.getenv("NVTE_MOE_DEBUG_HANDLE_MEM", "0") == "1" +_validate_ep_forward_roundtrip = ( + os.getenv("NVTE_MOE_VALIDATE_EP_FORWARD_ROUNDTRIP", "0") == "1" +) +_zero_moe_output = os.getenv("NVTE_MOE_ZERO_OUTPUT", "0") == "1" +_skip_moe_forward = os.getenv("NVTE_MOE_SKIP_FORWARD", "0") == "1" +_EP_ROUTING_PROBE_WIDTH = 16 +_debug_ffn_fwd_global_count = 0 +_debug_reference_dgrad_count = 0 +if _debug_python_patch: + print( + "[TE patch debug] imported transformer_engine.jax.moe " + f"from {__file__} (rank={os.getenv('SLURM_PROCID', 'unknown')}, " + f"reference_fwd={_use_reference_fwd}, reference_dgrad={_use_reference_dgrad})", + file=sys.stderr, + flush=True, + ) +if _use_reference_dgrad: + print( + "[TE reference dgrad] enabled: activation dgrad uses jax.lax.ragged_dot; " + "forward and weight gradients remain on the TE grouped-GEMM path", + file=sys.stderr, + flush=True, + ) +if _use_reference_fwd: + print( + "[TE reference fwd] enabled: FFN forward uses shard-local " + "jax.lax.ragged_dot; weight gradients remain on the TE grouped-GEMM path", + file=sys.stderr, + flush=True, + ) +if _zero_dispatch_weight_grad: + print( + "[TE dispatch diagnostic] routing-weight cotangent into ep_dispatch_bwd " + "is forced to zero; token dgrad remains enabled", + file=sys.stderr, + flush=True, + ) +if _refresh_ep_handle_in_bwd: + print( + "[TE EP handle diagnostic] backward refreshes the NCCL EP routing " + "handle from the saved routing map before combine/dispatch", + file=sys.stderr, + flush=True, + ) +if _refresh_ep_handle_before_combine: + print( + "[TE EP handle diagnostic] forward refreshes the NCCL EP routing " + "handle immediately before combine", + file=sys.stderr, + flush=True, + ) +if _validate_ep_routing: + print( + "[TE EP routing probe] enabled: validates effective combine-fwd and " + "dispatch-bwd handle mappings with deterministic expert codes", + file=sys.stderr, + flush=True, + ) +if _debug_moe_fwd_recompute: + print( + "[TE MoE fwd recompute debug] enabled: logs order-sensitive input/output " + "statistics keyed by the routing signature", + file=sys.stderr, + flush=True, + ) +if _validate_ep_token_roundtrip: + print( + "[TE EP token roundtrip] enabled: validates intra-expert token ordering " + "across dispatch-fwd and dispatch-bwd", + file=sys.stderr, + flush=True, + ) +if _zero_moe_input_grad: + print( + "[TE MoE input-grad diagnostic] the complete MoE input cotangent is " + "forced to zero after routing and gate gradients are combined", + file=sys.stderr, + flush=True, + ) +if _skip_moe_backward: + print( + "[TE MoE backward diagnostic] bypassing the complete TE MoE backward " + "body and returning zero activation/parameter cotangents", + file=sys.stderr, + flush=True, + ) +if _debug_handle_mem: + print( + "[TE EP handle-value diagnostic] logging each handle_mem's complete-byte " + "signatures and exact head/tail byte samples after ep_prepare", + file=sys.stderr, + flush=True, + ) +if _validate_ep_forward_roundtrip: + print( + "[TE EP forward roundtrip] validating real dispatched token values and " + "ordering through combine_fwd", + file=sys.stderr, + flush=True, + ) +if _zero_moe_output: + print( + "[TE MoE output diagnostic] forcing the routed-MoE forward output to " + "zero after all TE forward operations and validation probes", + file=sys.stderr, + flush=True, + ) +if _skip_moe_forward: + print( + "[TE MoE forward diagnostic] bypassing the complete TE routed-MoE " + "forward body and returning a zero output", + file=sys.stderr, + flush=True, + ) + + +def _debug_stable_stats(value, row_active=None): + """Return bounded sampled stats without materializing a full float32 copy.""" + value = jnp.asarray(value) + # A dispatch tensor is commonly [num_procs, recv_capacity, hidden] while + # its activity mask is [num_procs, recv_capacity]. Treat every mask entry + # as one logical row; using only value.shape[0] here would accidentally + # apply the first few token-mask entries to whole process-sized slabs. + if row_active is None: + matrix = value.reshape(value.shape[0], -1) + else: + row_active = jnp.asarray(row_active, jnp.bool_).reshape(-1) + if value.size % row_active.size != 0: + raise ValueError( + "Debug row mask must divide the sampled tensor size, but got " + f"value.shape={value.shape} and row_active.shape={row_active.shape}." + ) + matrix = value.reshape(row_active.size, -1) + max_samples = 65536 + stride = max((matrix.size + max_samples - 1) // max_samples, 1) + sample = matrix.reshape(-1)[::stride][:max_samples].astype(jnp.float32) + # Avoid forming flattened indices that can exceed int32 for production + # dispatch buffers (>180B logical elements). + stride_rows, stride_remainder = divmod(stride, matrix.shape[1]) + sample_indices = jnp.arange(sample.size, dtype=jnp.int32) + sampled_rows = ( + sample_indices * stride_rows + + (sample_indices * stride_remainder) // matrix.shape[1] + ) + if row_active is None: + active = jnp.ones(sample.shape, dtype=jnp.bool_) + else: + active = row_active[sampled_rows] + finite = jnp.isfinite(sample) & active + finite_value = jnp.where(finite, sample, 0.0) + absmax = jnp.max(jnp.abs(finite_value)) + safe_scale = jnp.where(absmax > 0, absmax, 1.0) + scaled = finite_value / safe_scale + element_count = jnp.maximum(jnp.sum(active, dtype=jnp.float32), 1.0) + scaled_mean = jnp.sum(scaled) / element_count + scaled_square_mean = jnp.sum(jnp.square(scaled)) / element_count + mean = absmax * scaled_mean + abs_mean = absmax * (jnp.sum(jnp.abs(scaled)) / element_count) + stddev = absmax * jnp.sqrt(jnp.maximum(scaled_square_mean - jnp.square(scaled_mean), 0.0)) + finite_fraction = jnp.sum(finite, dtype=jnp.float32) / element_count + return mean, abs_mean, stddev, absmax, finite_fraction + + +def _ep_routing_probe_code_table(num_experts, dtype): + """Return deterministic ±1 codes that identify experts across 16 channels.""" + expert = jnp.arange(num_experts, dtype=jnp.uint32)[:, None] + channel = jnp.arange(_EP_ROUTING_PROBE_WIDTH, dtype=jnp.uint32)[None, :] + value = (expert + jnp.uint32(1)) * jnp.uint32(0x9E3779B1) + value ^= (channel + jnp.uint32(1)) * jnp.uint32(0x85EBCA77) + value ^= value >> jnp.uint32(16) + value *= jnp.uint32(0xC2B2AE3D) + value ^= value >> jnp.uint32(13) + return jnp.where((value & jnp.uint32(1)) != 0, 1, -1).astype(dtype) + + +def _ep_routing_probe_packed_codes( + token_counts, + recv_capacity_per_rank, + num_ep, + num_local_experts, + dtype, +): + """Build expert-major probe rows matching the native EP receive layout.""" + leading_size = token_counts.shape[0] + ep_rank = jnp.arange(leading_size, dtype=jnp.int32) % num_ep + local_expert = jnp.arange(num_local_experts, dtype=jnp.int32) + global_expert = ep_rank[:, None] * num_local_experts + local_expert[None, :] + code_table = _ep_routing_probe_code_table(num_ep * num_local_experts, dtype) + codes_by_group = code_table[global_expert] + + def _repeat_one_leading_group(group_codes, group_counts): + return jnp.repeat( + group_codes, + group_counts.astype(jnp.int32), + axis=0, + total_repeat_length=recv_capacity_per_rank, + ) + + packed = jax.vmap(_repeat_one_leading_group)(codes_by_group, token_counts) + active_rows = ( + jnp.arange(recv_capacity_per_rank, dtype=jnp.int32)[None, :] + < jnp.sum(token_counts, axis=-1, dtype=jnp.int32)[:, None] + ) + return jnp.where(active_rows[..., None], packed, jnp.zeros_like(packed)) + + +def _ep_routing_probe_signature(topk_idx): + """Return two order-sensitive uint32 signatures for the intended map.""" + flat = topk_idx.reshape(-1).astype(jnp.uint32) + position = jnp.arange(flat.size, dtype=jnp.uint32) + signature_0 = jnp.sum( + (flat + jnp.uint32(1)) + * (position * jnp.uint32(0x9E3779B1) + jnp.uint32(0x85EBCA77)), + dtype=jnp.uint32, + ) + signature_1 = jnp.sum( + (flat + jnp.uint32(17)) + * (position * jnp.uint32(0xC2B2AE3D) + jnp.uint32(0x27D4EB2F)), + dtype=jnp.uint32, + ) + return signature_0, signature_1 + + +def _print_handle_mem_value(label, handle_mem, topk_idx): + """Log exact handle bytes plus full-buffer fingerprints after ep_prepare. + + ``handle_mem`` is an opaque uint8 tensor with one row per global EP/DP + rank. Printing every byte for every scanned layer would make the + multi-host log impractically large, so the exact first/last 64 bytes are + printed and two order-sensitive signatures cover every byte in each row. + The routing-map signature in the same record makes repeated-forward + comparisons unambiguous. + """ + if not _debug_handle_mem: + return + rows = handle_mem.reshape(-1, handle_mem.shape[-1]).astype(jnp.uint32) + position = jnp.arange(rows.shape[-1], dtype=jnp.uint32) + signature_0 = jnp.sum( + (rows + jnp.uint32(1)) + * (position * jnp.uint32(0x9E3779B1) + jnp.uint32(0x85EBCA77)), + axis=-1, + dtype=jnp.uint32, + ) + signature_1 = jnp.sum( + (rows + jnp.uint32(17)) + * (position * jnp.uint32(0xC2B2AE3D) + jnp.uint32(0x27D4EB2F)), + axis=-1, + dtype=jnp.uint32, + ) + byte_sum = jnp.sum(rows, axis=-1, dtype=jnp.uint32) + route_signature_0, route_signature_1 = _ep_routing_probe_signature(topk_idx) + sample_width = min(64, handle_mem.shape[-1]) + jax.debug.print( + f"[TE EP handle value] label={label} " + f"shape={handle_mem.shape} " + "route_sig=({route_sig0},{route_sig1}) " + "byte_sum={byte_sum} handle_sig0={handle_sig0} " + "handle_sig1={handle_sig1} head={head} tail={tail}", + route_sig0=route_signature_0, + route_sig1=route_signature_1, + byte_sum=byte_sum, + handle_sig0=signature_0, + handle_sig1=signature_1, + head=handle_mem[..., :sample_width], + tail=handle_mem[..., -sample_width:], + ordered=False, + ) + + +def _debug_ordered_tensor_stats(value): + """Return stable scalar stats plus two order-sensitive sampled projections.""" + value = jnp.asarray(value) + max_samples = 65536 + stride = max((value.size + max_samples - 1) // max_samples, 1) + sample = value.reshape(-1)[::stride][:max_samples].astype(jnp.float32) + finite = jnp.isfinite(sample) + finite_value = jnp.where(finite, sample, 0.0) + absmax = jnp.max(jnp.abs(finite_value)) + safe_scale = jnp.where(absmax > 0, absmax, 1.0) + scaled = finite_value / safe_scale + count = jnp.maximum(jnp.sum(finite, dtype=jnp.float32), 1.0) + mean = absmax * jnp.sum(scaled) / count + square_mean = jnp.sum(jnp.square(scaled)) / count + stddev = absmax * jnp.sqrt( + jnp.maximum(square_mean - jnp.square(jnp.sum(scaled) / count), 0.0) + ) + position = jnp.arange(sample.size, dtype=jnp.uint32) + sign_0 = jnp.where( + ((position * jnp.uint32(0x9E3779B1) + jnp.uint32(0x85EBCA77)) >> 31) != 0, + 1.0, + -1.0, + ) + sign_1 = jnp.where( + ((position * jnp.uint32(0xC2B2AE3D) + jnp.uint32(0x27D4EB2F)) >> 31) != 0, + 1.0, + -1.0, + ) + projection_0 = jnp.sum(scaled * sign_0) / count + projection_1 = jnp.sum(scaled * sign_1) / count + return mean, stddev, absmax, projection_0, projection_1, jnp.mean(finite) + + +def _print_ep_routing_probe_result(label, topk_idx, actual, expected, tolerance): + """Print an elementwise comparison for an expert-code routing probe.""" + difference = actual.astype(jnp.float32) - expected.astype(jnp.float32) + abs_difference = jnp.abs(difference) + mismatch = abs_difference > tolerance + signature_0, signature_1 = _ep_routing_probe_signature(topk_idx) + jax.debug.print( + "[TE EP routing probe] {label} route_sig=({signature_0},{signature_1}) " + "match={match} " + "mismatch_fraction={mismatch_fraction:.6e} " + "absmean={absmean:.6e} absmax={absmax:.6e}", + label=label, + signature_0=signature_0, + signature_1=signature_1, + match=jnp.all(~mismatch), + mismatch_fraction=jnp.mean(mismatch.astype(jnp.float32)), + absmean=jnp.mean(abs_difference), + absmax=jnp.max(abs_difference), + ordered=False, + ) + def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: """Sharding constraint that keeps bwd cotangents in the primal dtype. @@ -194,12 +537,18 @@ class _Ctx: handle_mem: jnp.ndarray token_counts: jnp.ndarray recv_topk_weights: jnp.ndarray + recv_token_probe: Any 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 + wi: Any + wo: Any + wi_0_bias: Any + wi_1_bias: Any + wo_bias: Any expert_outputs: jnp.ndarray local_group_sizes: jnp.ndarray quantizer_sets: Any @@ -231,6 +580,7 @@ def _ffn_fwd_global( apply_topk_weights_early: bool, flat_token_sharding: NamedSharding, flat_group_sharding: NamedSharding, + grouped_weight_sharding: NamedSharding, grouped_bias_sharding: NamedSharding, ): """Run the FFN on global EP-dispatch buffers. @@ -240,6 +590,21 @@ def _ffn_fwd_global( passed through as the dynamic grouped-GEMM group sizes, so cuBLAS skips both 0-token experts and dispatch-buffer over-allocation. """ + global _debug_ffn_fwd_global_count + if _debug_python_patch and _debug_ffn_fwd_global_count < 10: + _debug_ffn_fwd_global_count += 1 + print( + "[TE patch debug] _ffn_fwd_global " + f"call={_debug_ffn_fwd_global_count} " + f"recv_tokens.shape={recv_tokens.shape} " + f"token_counts.shape={token_counts.shape} " + f"wi.shape={wi.shape} wo.shape={wo.shape} " + f"dp_size={dp_size} num_ep={num_ep} " + f"num_local_experts={num_local_experts} " + f"flat_group_sharding={flat_group_sharding}", + file=sys.stderr, + flush=True, + ) hidden = recv_tokens.shape[-1] sorted_x = recv_tokens.reshape(-1, hidden) recv_w_flat = recv_topk_weights.reshape(-1) @@ -254,13 +619,42 @@ def _ffn_fwd_global( # Dispatch groups flatten in (dp, ep, local_expert) order. Keep the # model weights in their original global [expert, ...] layout: grouped # quantize must see the FSDP shard before grouped_gemm's custom - # partitioning gathers it. The grouped-GEMM partitioner maps that - # smaller RHS group axis onto the local dispatch groups after quantizing. + # partitioning gathers it. This avoids materializing a global dp*expert + # weight tensor and keeps the FSDP collective inside the grouped-GEMM + # custom partitioning boundary. # # Bias is the one exception. grouped_gemm's public bias contract is one # row per global GEMM, so retain its inexpensive logical expansion here. num_groups = dp_size * num_ep * num_local_experts + def _weights_in_dispatch_group_order(weight): + expanded = jnp.broadcast_to( + weight[None, ...], + (dp_size, *weight.shape), + ).reshape(num_groups, *weight.shape[1:]) + return jax.lax.with_sharding_constraint(expanded, grouped_weight_sharding) + + def _reference_ragged_dot(lhs, rhs): + def _per_shard(local_lhs, local_rhs, local_group_sizes): + return jax.lax.ragged_dot(local_lhs, local_rhs, local_group_sizes) + + return jax.shard_map( + _per_shard, + mesh=flat_token_sharding.mesh, + in_specs=( + flat_token_sharding.spec, + grouped_weight_sharding.spec, + flat_group_sharding.spec, + ), + out_specs=flat_token_sharding.spec, + )(lhs, rhs, group_sizes) + + if _use_reference_fwd: + if wi_0_bias is not None: + raise ValueError("NVTE_MOE_REFERENCE_FWD does not support expert biases.") + wi_for_fwd = _weights_in_dispatch_group_order(wi) + wo_for_fwd = _weights_in_dispatch_group_order(wo) + def _broadcast_bias(value): value = jnp.broadcast_to( value.reshape(1, num_ep, num_local_experts, *value.shape[1:]), @@ -288,12 +682,15 @@ def _broadcast_bias(value): ragged_scale_sharding=flat_token_sharding, ) casted_wi = tex.grouped_quantize(wi, fc1_quantizer_set.kernel, flatten_axis=-1) - combined_out = tex.grouped_gemm( - casted_sorted_x.get_tensor(usage=TensorUsage.LHS), - casted_wi.get_tensor(usage=TensorUsage.RHS), - contracting_dims=((1,), (1,)), - bias=wi_combined_bias, - ) + if _use_reference_fwd: + combined_out = _reference_ragged_dot(sorted_x, wi_for_fwd) + else: + combined_out = tex.grouped_gemm( + casted_sorted_x.get_tensor(usage=TensorUsage.LHS), + casted_wi.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((1,), (1,)), + bias=wi_combined_bias, + ) combined_out = jax.lax.with_sharding_constraint(combined_out, flat_token_sharding) 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).checkpoint( @@ -326,12 +723,15 @@ def _broadcast_bias(value): ragged_scale_sharding=flat_token_sharding, ) 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, - ) + if _use_reference_fwd: + expert_outputs = _reference_ragged_dot(intermediate, wo_for_fwd) + else: + 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, + ) expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, flat_token_sharding) casted_intermediate_lhs_trans = casted_intermediate.get_tensor( usage=TensorUsage.LHS_TRANS @@ -362,6 +762,8 @@ def _ffn_bwd_global( up_proj_out: jnp.ndarray, casted_intermediate_lhs_trans, casted_wo_rhs_trans, + wi: jnp.ndarray, + wo: jnp.ndarray, local_group_sizes: jnp.ndarray, recv_topk_weights: jnp.ndarray, quantizer_sets: Tuple[QuantizerSet, QuantizerSet], @@ -380,6 +782,13 @@ def _ffn_bwd_global( ``(d_sorted_x [num_procs, recv_pr, H], d_recv_w [num_procs, recv_pr], d_wi, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. """ + # Each leading process row has its own packed expert prefix and trailing + # over-allocation. Preserve that boundary for diagnostics before flattening + # the group sizes for the shard-local ragged operations. + active_rows_2d = ( + jnp.arange(recv_topk_weights.shape[-1])[None, :] + < jnp.sum(local_group_sizes, axis=-1, dtype=jnp.int32)[:, None] + ) group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) d_eo_2d = d_expert_outputs.reshape(-1, d_expert_outputs.shape[-1]) recv_w_flat = recv_topk_weights.reshape(-1) @@ -387,6 +796,54 @@ def _ffn_bwd_global( recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) fc1_quantizer_set, fc2_quantizer_set = quantizer_sets + + def _weights_in_dispatch_group_order(weight): + """Repeat one global expert set for each outer FSDP token replica.""" + if group_sizes.shape[0] % weight.shape[0] != 0: + raise ValueError( + "Reference MoE dgrad requires dispatch groups to contain an " + "integer number of global expert sets, but got " + f"{group_sizes.shape[0]} groups and {weight.shape[0]} experts." + ) + num_replicas = group_sizes.shape[0] // weight.shape[0] + expanded = jnp.broadcast_to( + weight[None, ...], + (num_replicas, *weight.shape), + ).reshape(group_sizes.shape[0], *weight.shape[1:]) + return jax.lax.with_sharding_constraint(expanded, grouped_weight_sharding) + + def _reference_ragged_dot(lhs, rhs): + """Run the JAX reference on matching shard-local rows and expert groups.""" + def _per_shard(local_lhs, local_rhs, local_group_sizes): + return jax.lax.ragged_dot(local_lhs, local_rhs, local_group_sizes) + + return jax.shard_map( + _per_shard, + mesh=flat_token_sharding.mesh, + in_specs=( + flat_token_sharding.spec, + grouped_weight_sharding.spec, + flat_group_sharding.spec, + ), + out_specs=flat_token_sharding.spec, + )(lhs, rhs, group_sizes) + + if _use_reference_dgrad: + global _debug_reference_dgrad_count + if _debug_python_patch and _debug_reference_dgrad_count < 10: + _debug_reference_dgrad_count += 1 + print( + "[TE reference dgrad] tracing " + f"call={_debug_reference_dgrad_count} " + f"group_sizes={group_sizes.shape} wi={wi.shape} wo={wo.shape} " + f"group_order=(fsdp,ep,local_expert) " + f"output_constraint={flat_token_sharding}", + file=sys.stderr, + flush=True, + ) + wi_for_dgrad = _weights_in_dispatch_group_order(wi) + wo_for_dgrad = _weights_in_dispatch_group_order(wo) + # 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. @@ -402,11 +859,17 @@ def _ffn_bwd_global( ) _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) - d_intermediate = tex.grouped_gemm( - _casted_d_eo_lhs, - casted_wo_rhs_trans, - contracting_dims=((1,), (2,)), - ) + if _use_reference_dgrad: + d_intermediate = _reference_ragged_dot( + d_eo_2d, + jnp.swapaxes(wo_for_dgrad, -1, -2), + ) + else: + d_intermediate = tex.grouped_gemm( + _casted_d_eo_lhs, + casted_wo_rhs_trans, + contracting_dims=((1,), (2,)), + ) d_intermediate = jax.lax.with_sharding_constraint(d_intermediate, flat_token_sharding) d_wo = tex.grouped_gemm( casted_intermediate_lhs_trans, @@ -462,11 +925,17 @@ def _ffn_bwd_global( flatten_axis=-1, ragged_scale_sharding=flat_token_sharding, ) - d_sorted_x = tex.grouped_gemm( - casted_d_combined.get_tensor(usage=TensorUsage.LHS), - casted_wi_rhs_trans, - contracting_dims=((1,), (2,)), - ) + if _use_reference_dgrad: + d_sorted_x = _reference_ragged_dot( + d_combined, + jnp.swapaxes(wi_for_dgrad, -1, -2), + ) + else: + d_sorted_x = tex.grouped_gemm( + casted_d_combined.get_tensor(usage=TensorUsage.LHS), + casted_wi_rhs_trans, + contracting_dims=((1,), (2,)), + ) d_sorted_x = jax.lax.with_sharding_constraint(d_sorted_x, flat_token_sharding) d_wi_combined = tex.grouped_gemm( casted_sorted_x_lhs_trans, @@ -485,6 +954,46 @@ def _ffn_bwd_global( d_wi_0_bias = None d_wi_1_bias = None + if _debug_moe_numerics: + # ragged_dot/grouped_gemm consume one packed prefix per process row. + # recv_w padding is not guaranteed initialized, so recv_w != 0 is not + # a valid activity mask. + active_rows = active_rows_2d.reshape(-1) + deo_stats = _debug_stable_stats(d_eo_2d, active_rows) + dint_stats = _debug_stable_stats(d_intermediate, active_rows) + dcombined_stats = _debug_stable_stats(d_combined, active_rows) + dx_stats = _debug_stable_stats(d_sorted_x, active_rows) + drw_stats = _debug_stable_stats(d_recv_w_from_intermediate, active_rows) + jax.debug.print( + "[TE bwd stats] " + "dout(mean={deo_mean:.3e},absmean={deo_absmean:.3e},std={deo_std:.3e}," + "absmax={deo_absmax:.3e},finite={deo_finite:.6f}) " + "dfc2(absmean={dint_absmean:.3e},std={dint_std:.3e},absmax={dint_absmax:.3e}) " + "dact(absmean={dc_absmean:.3e},std={dc_std:.3e},absmax={dc_absmax:.3e}) " + "din(absmean={dx_absmean:.3e},std={dx_std:.3e},absmax={dx_absmax:.3e}," + "finite={dx_finite:.6f}) " + "dweight(absmean={drw_absmean:.3e},std={drw_std:.3e},absmax={drw_absmax:.3e})", + deo_mean=deo_stats[0], + deo_absmean=deo_stats[1], + deo_std=deo_stats[2], + deo_absmax=deo_stats[3], + deo_finite=deo_stats[4], + dint_absmean=dint_stats[1], + dint_std=dint_stats[2], + dint_absmax=dint_stats[3], + dc_absmean=dcombined_stats[1], + dc_std=dcombined_stats[2], + dc_absmax=dcombined_stats[3], + dx_absmean=dx_stats[1], + dx_std=dx_stats[2], + dx_absmax=dx_stats[3], + dx_finite=dx_stats[4], + drw_absmean=drw_stats[1], + drw_std=drw_stats[2], + drw_absmax=drw_stats[3], + ordered=False, + ) + d_sorted_x_3d = d_sorted_x.reshape(*d_expert_outputs.shape[:-1], d_sorted_x.shape[-1]) d_recv_w_3d = d_recv_w_from_intermediate.reshape(recv_topk_weights.shape) return ( @@ -540,6 +1049,49 @@ def _moe_fwd_rule( x = with_sharding_constraint_by_logical_axes(x, input_axes) + if _skip_moe_forward: + if not _skip_moe_backward: + raise RuntimeError( + "NVTE_MOE_SKIP_FORWARD requires NVTE_MOE_SKIP_BACKWARD=1." + ) + has_bias = wi_0_bias is not None + ctx = _Ctx( + x=x, + gate_kernel=gate_kernel, + expert_bias=expert_bias, + logits_2d=None, + saved_scores=None, + routing_map=None, + cfg=None, + handle_mem=None, + token_counts=None, + recv_topk_weights=None, + recv_token_probe=None, + casted_sorted_x_lhs_trans=None, + casted_wi_rhs_trans=None, + gate_proj_out=None, + up_proj_out=None, + casted_intermediate_lhs_trans=None, + casted_wo_rhs_trans=None, + wi=wi, + wo=wo, + wi_0_bias=wi_0_bias if has_bias else None, + wi_1_bias=wi_1_bias if has_bias else None, + wo_bias=wo_bias if has_bias else None, + expert_outputs=None, + local_group_sizes=None, + quantizer_sets=quantizer_sets, + ) + static = { + "has_bias": has_bias, + "x_shape": x.shape, + "recv_pr": 0, + } + return ( + jnp.zeros_like(x), + jnp.zeros((), dtype=x.dtype), + ), (ctx, static) + mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") @@ -696,6 +1248,7 @@ def _moe_fwd_rule( dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + _print_handle_mem_value("forward_prepare", handle_mem, 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 @@ -728,17 +1281,38 @@ def _moe_fwd_rule( apply_topk_weights_early=apply_topk_weights_early, flat_token_sharding=flat_token_sharding, flat_group_sharding=flat_group_sharding, + grouped_weight_sharding=grouped_weight_sharding, grouped_bias_sharding=grouped_bias_sharding, ) 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) + combine_handle_mem = handle_mem + if _refresh_ep_handle_before_combine: + # A scanned forward can reuse one physical handle_mem allocation across + # loop iterations. Re-prepare the current routing map at the combine + # consumption point so combine cannot rely on routing state left by a + # different iteration. Keep the original handle in the VJP residual: + # this diagnostic intentionally isolates forward combine. + refreshed_token_counts, combine_handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + _print_handle_mem_value( + "forward_pre_combine_refresh", combine_handle_mem, topk_idx_3d + ) + refreshed_token_counts = jax.lax.with_sharding_constraint( + refreshed_token_counts, NamedSharding(mesh, ep2_spec) + ) + if _debug_moe_numerics: + jax.debug.print( + "[TE EP pre-combine refresh] token_counts_match={match}", + match=jnp.all(refreshed_token_counts == token_counts), + ordered=False, + ) if apply_topk_weights_early: # expert_outputs is already weighted upstream. output = tex.ep_combine_fwd( cfg, - handle_mem, + combine_handle_mem, expert_outputs, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, @@ -750,12 +1324,108 @@ def _moe_fwd_rule( weighted = expert_outputs * w output = tex.ep_combine_fwd( cfg, - handle_mem, + combine_handle_mem, weighted, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, ) + if _validate_ep_forward_roundtrip: + # This validates more than the synthetic expert-code probe below: + # dispatch the actual input token values, combine them immediately, + # and compare against the weighted identity analytically. It detects + # a mutually consistent but wrong dispatch/combine token permutation. + probe_width = min(_EP_ROUTING_PROBE_WIDTH, H) + roundtrip_weighted = ( + recv_tokens[..., :probe_width].astype(jnp.float32) + * recv_topk_weights[..., None] + ).astype(x.dtype) + roundtrip_output = tex.ep_combine_fwd( + cfg, + combine_handle_mem, + roundtrip_weighted, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + expected_roundtrip = ( + x[..., :probe_width].astype(jnp.float32) + * jnp.sum(topk_w_3d, axis=-1, keepdims=True) + ).astype(x.dtype) + _print_ep_routing_probe_result( + "forward_token_roundtrip", + topk_idx_3d, + roundtrip_output, + expected_roundtrip, + tolerance=6.25e-2, + ) + + if _validate_ep_routing: + probe_code_table = _ep_routing_probe_code_table(num_experts, x.dtype) + probe_packed = _ep_routing_probe_packed_codes( + token_counts, + recv_pr, + num_ep, + num_local_experts, + x.dtype, + ) + probe_packed = jax.lax.with_sharding_constraint( + probe_packed, NamedSharding(mesh, ep3_spec) + ) + probe_weighted = ( + probe_packed.astype(jnp.float32) * recv_topk_weights[..., None] + ).astype(x.dtype) + probe_combined = tex.ep_combine_fwd( + cfg, + combine_handle_mem, + probe_weighted, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + expected_probe_terms = ( + probe_code_table[topk_idx_3d].astype(jnp.float32) + * topk_w_3d[..., None] + ).astype(x.dtype) + expected_probe_combined = jnp.sum( + expected_probe_terms.astype(jnp.float32), axis=-2 + ).astype(x.dtype) + _print_ep_routing_probe_result( + "combine_fwd", + topk_idx_3d, + probe_combined, + expected_probe_combined, + tolerance=6.25e-2, + ) + + if _zero_moe_output: + output = jnp.zeros_like(output) + + if _debug_moe_fwd_recompute: + signature_0, signature_1 = _ep_routing_probe_signature(topk_idx_3d) + x_stats = _debug_ordered_tensor_stats(x) + output_stats = _debug_ordered_tensor_stats(output) + jax.debug.print( + "[TE MoE fwd recompute] route_sig=({signature_0},{signature_1}) " + "input(mean={x_mean:.6e},std={x_std:.6e},absmax={x_absmax:.6e}," + "proj=({x_proj0:.9e},{x_proj1:.9e}),finite={x_finite:.6f}) " + "output(mean={out_mean:.6e},std={out_std:.6e},absmax={out_absmax:.6e}," + "proj=({out_proj0:.9e},{out_proj1:.9e}),finite={out_finite:.6f})", + signature_0=signature_0, + signature_1=signature_1, + x_mean=x_stats[0], + x_std=x_stats[1], + x_absmax=x_stats[2], + x_proj0=x_stats[3], + x_proj1=x_stats[4], + x_finite=x_stats[5], + out_mean=output_stats[0], + out_std=output_stats[1], + out_absmax=output_stats[2], + out_proj0=output_stats[3], + out_proj1=output_stats[4], + out_finite=output_stats[5], + ordered=False, + ) + ( casted_sorted_x_lhs_trans, casted_wi_rhs_trans, @@ -777,12 +1447,25 @@ def _moe_fwd_rule( handle_mem=handle_mem, token_counts=token_counts, recv_topk_weights=recv_topk_weights, + recv_token_probe=( + jax.lax.with_sharding_constraint( + recv_tokens[..., :_EP_ROUTING_PROBE_WIDTH], + NamedSharding(mesh, ep3_spec), + ) + if _validate_ep_token_roundtrip + else None + ), 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, + wi=wi if (_use_reference_dgrad or _skip_moe_backward) else None, + wo=wo if (_use_reference_dgrad or _skip_moe_backward) else None, + wi_0_bias=wi_0_bias if (has_bias and _skip_moe_backward) else None, + wi_1_bias=wi_1_bias if (has_bias and _skip_moe_backward) else None, + wo_bias=wo_bias if (has_bias and _skip_moe_backward) else None, expert_outputs=expert_outputs, local_group_sizes=local_group_sizes, quantizer_sets=quantizer_sets, @@ -829,6 +1512,56 @@ def _moe_bwd_rule( x_shape = static["x_shape"] recv_pr = static["recv_pr"] + if _skip_moe_backward: + # Strong isolation diagnostic: do not execute combine_bwd, grouped + # GEMM backward, dispatch_bwd, or router backward. This differs from + # NVTE_MOE_ZERO_INPUT_GRAD, which discards d_x only after all of those + # operations have already run and therefore cannot rule out an + # asynchronous side effect from a backward custom call. + if ctx.wi is None or ctx.wo is None: + raise RuntimeError( + "NVTE_MOE_SKIP_BACKWARD requires wi/wo in the VJP residual." + ) + d_x = with_sharding_constraint_by_logical_axes( + jnp.zeros_like(ctx.x), input_axes + ) + d_gate_kernel = with_sharding_constraint_by_logical_axes( + jnp.zeros_like(ctx.gate_kernel), gate_kernel_axes + ) + d_wi = with_sharding_constraint_by_logical_axes( + jnp.zeros_like(ctx.wi), wi_kernel_axes + ) + d_wo = with_sharding_constraint_by_logical_axes( + jnp.zeros_like(ctx.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( + jnp.zeros_like(ctx.wi_0_bias), wi_bias_axes + ) + d_wi_1_bias = with_sharding_constraint_by_logical_axes( + jnp.zeros_like(ctx.wi_1_bias), wi_bias_axes + ) + d_wo_bias = with_sharding_constraint_by_logical_axes( + jnp.zeros_like(ctx.wo_bias), wo_bias_axes + ) + else: + d_wi_0_bias = None + d_wi_1_bias = None + d_wo_bias = None + return ( + d_x, + d_gate_kernel, + d_wi, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + jnp.zeros_like(ctx.expert_bias), + ctx.quantizer_sets, + ) + mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") @@ -852,12 +1585,53 @@ def _moe_bwd_rule( grouped_bias_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) out_partition_spec = (batch_pspec_axis, None, None) + # A scanned layer can reuse the same physical handle_mem buffer for + # different loop iterations. The native NCCL EP cache is keyed by that + # buffer pointer and retains routing state established by ep_prepare. + # Refreshing here updates the cached handle to the routing plan for the + # current reverse-scan iteration before either backward EP operation. + bwd_handle_mem = ctx.handle_mem + if _refresh_ep_handle_in_bwd: + bwd_selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] + bwd_topk_idx = bwd_selected_experts.reshape(B, S, K).astype(jnp.int32) + bwd_topk_idx = jax.lax.with_sharding_constraint( + bwd_topk_idx, NamedSharding(mesh, ep3_spec) + ) + refreshed_token_counts, bwd_handle_mem = tex.ep_prepare(ctx.cfg, bwd_topk_idx) + _print_handle_mem_value( + "backward_refresh", bwd_handle_mem, bwd_topk_idx + ) + refreshed_token_counts = jax.lax.with_sharding_constraint( + refreshed_token_counts, NamedSharding(mesh, ep2_spec) + ) + if _debug_moe_numerics: + token_count_match = jnp.all(refreshed_token_counts == ctx.token_counts) + jax.debug.print( + "[TE EP handle refresh] token_counts_match={match}", + match=token_count_match, + ordered=False, + ) + # ---------------- Combine bwd (global view) ---------------- d_output = jax.lax.with_sharding_constraint(d_output, NamedSharding(mesh, ep3_spec)) - grad_pre_combine = tex.ep_combine_bwd(ctx.cfg, ctx.handle_mem, d_output, recv_pr) + grad_pre_combine = tex.ep_combine_bwd(ctx.cfg, bwd_handle_mem, d_output, recv_pr) grad_pre_combine = jax.lax.with_sharding_constraint( grad_pre_combine, NamedSharding(mesh, ep3_spec) ) + # The EP kernel writes only the per-process packed expert prefix. Its + # over-allocation tail is intentionally left uninitialized, which is safe + # only while every downstream consumer is perfectly handle/group aware. + # Materialize the contract here so padding cannot leak through elementwise + # weighting, compiler fusion, or a later dispatch backward. + active_recv_rows = ( + jnp.arange(ctx.recv_topk_weights.shape[-1])[None, :] + < jnp.sum(ctx.token_counts, axis=-1, dtype=jnp.int32)[:, None] + ) + grad_pre_combine = jnp.where( + active_recv_rows[..., None], + grad_pre_combine, + jnp.zeros_like(grad_pre_combine), + ) if apply_topk_weights_early: # combine_fwd consumed already-weighted expert_outputs; the recv_w @@ -868,14 +1642,64 @@ def _moe_bwd_rule( # 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( - mask_bool, grad_pre_combine * w, jnp.zeros_like(grad_pre_combine) + active_recv_rows_3d = active_recv_rows[..., None] + safe_recv_weights = jnp.where( + active_recv_rows, + ctx.recv_topk_weights, + jnp.zeros_like(ctx.recv_topk_weights), + ) + safe_expert_outputs = jnp.where( + active_recv_rows_3d, + ctx.expert_outputs, + jnp.zeros_like(ctx.expert_outputs), ) - d_recv_w_from_combine = (grad_pre_combine * ctx.expert_outputs).sum(axis=-1) + d_expert_outputs = ( + grad_pre_combine + * safe_recv_weights[..., None].astype(grad_pre_combine.dtype) + ) + d_recv_w_from_combine = (grad_pre_combine * safe_expert_outputs).sum(axis=-1) d_recv_w_from_combine = d_recv_w_from_combine.astype(ctx.recv_topk_weights.dtype) + if _debug_moe_numerics: + d_output_stats = _debug_stable_stats(d_output) + grad_pre_combine_stats = _debug_stable_stats( + grad_pre_combine, active_recv_rows + ) + d_expert_output_stats = _debug_stable_stats( + d_expert_outputs, active_recv_rows + ) + d_recv_w_stats = _debug_stable_stats( + d_recv_w_from_combine, active_recv_rows + ) + jax.debug.print( + "[TE combine-bwd stats] " + "upstream(absmean={up_absmean:.3e},std={up_std:.3e}," + "absmax={up_absmax:.3e},finite={up_finite:.6f}) " + "combine(absmean={combine_absmean:.3e},std={combine_std:.3e}," + "absmax={combine_absmax:.3e},finite={combine_finite:.6f}) " + "weighted(absmean={weighted_absmean:.3e},std={weighted_std:.3e}," + "absmax={weighted_absmax:.3e},finite={weighted_finite:.6f}) " + "dweight(absmean={dw_absmean:.3e},std={dw_std:.3e}," + "absmax={dw_absmax:.3e},finite={dw_finite:.6f})", + up_absmean=d_output_stats[1], + up_std=d_output_stats[2], + up_absmax=d_output_stats[3], + up_finite=d_output_stats[4], + combine_absmean=grad_pre_combine_stats[1], + combine_std=grad_pre_combine_stats[2], + combine_absmax=grad_pre_combine_stats[3], + combine_finite=grad_pre_combine_stats[4], + weighted_absmean=d_expert_output_stats[1], + weighted_std=d_expert_output_stats[2], + weighted_absmax=d_expert_output_stats[3], + weighted_finite=d_expert_output_stats[4], + dw_absmean=d_recv_w_stats[1], + dw_std=d_recv_w_stats[2], + dw_absmax=d_recv_w_stats[3], + dw_finite=d_recv_w_stats[4], + ordered=False, + ) + # ---------------- FFN bwd (global view, custom-partitioned primitives) ---------------- ( d_sorted_x, @@ -893,6 +1717,8 @@ def _moe_bwd_rule( ctx.up_proj_out, ctx.casted_intermediate_lhs_trans, ctx.casted_wo_rhs_trans, + ctx.wi, + ctx.wo, ctx.local_group_sizes, ctx.recv_topk_weights, ctx.quantizer_sets, @@ -927,15 +1753,125 @@ def _fold_dp_groups(grad): # ---------------- 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)) + dispatch_weight_cotangent = ( + jnp.zeros_like(d_recv_w_total) if _zero_dispatch_weight_grad else d_recv_w_total + ) d_x_from_dispatch, d_topk_w = tex.ep_dispatch_bwd( ctx.cfg, - ctx.handle_mem, + bwd_handle_mem, d_sorted_x, - d_recv_w_total, + dispatch_weight_cotangent, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, ) + if _validate_ep_token_roundtrip: + roundtrip_topk_idx = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] + roundtrip_topk_idx_3d = roundtrip_topk_idx.reshape(B, S, K).astype(jnp.int32) + roundtrip_topk_idx_3d = jax.lax.with_sharding_constraint( + roundtrip_topk_idx_3d, NamedSharding(mesh, ep3_spec) + ) + roundtrip_tokens, _ = tex.ep_dispatch_bwd( + ctx.cfg, + bwd_handle_mem, + ctx.recv_token_probe, + jnp.zeros_like(ctx.recv_topk_weights), + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + expected_roundtrip = ( + ctx.x[..., :_EP_ROUTING_PROBE_WIDTH].astype(jnp.float32) * float(K) + ).astype(roundtrip_tokens.dtype) + _print_ep_routing_probe_result( + "dispatch_token_roundtrip", + roundtrip_topk_idx_3d, + roundtrip_tokens, + expected_roundtrip, + tolerance=6.25e-2, + ) + + if _validate_ep_routing: + probe_topk_idx = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] + probe_topk_idx_3d = probe_topk_idx.reshape(B, S, K).astype(jnp.int32) + probe_topk_idx_3d = jax.lax.with_sharding_constraint( + probe_topk_idx_3d, NamedSharding(mesh, ep3_spec) + ) + probe_code_table = _ep_routing_probe_code_table(num_experts, d_sorted_x.dtype) + probe_packed = _ep_routing_probe_packed_codes( + ctx.token_counts, + recv_pr, + num_ep, + num_local_experts, + d_sorted_x.dtype, + ) + probe_packed = jax.lax.with_sharding_constraint( + probe_packed, NamedSharding(mesh, ep3_spec) + ) + probe_dispatch, _ = tex.ep_dispatch_bwd( + ctx.cfg, + bwd_handle_mem, + probe_packed, + jnp.zeros_like(ctx.recv_topk_weights), + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + expected_probe_dispatch = jnp.sum( + probe_code_table[probe_topk_idx_3d].astype(jnp.float32), + axis=-2, + ).astype(d_sorted_x.dtype) + _print_ep_routing_probe_result( + "dispatch_bwd", + probe_topk_idx_3d, + probe_dispatch, + expected_probe_dispatch, + tolerance=1.0e-3, + ) + + if _debug_moe_numerics or _debug_moe_input_grad: + d_sorted_x_stats = _debug_stable_stats(d_sorted_x, active_recv_rows) + d_recv_combine_stats = _debug_stable_stats( + d_recv_w_from_combine, active_recv_rows + ) + d_recv_intermediate_stats = _debug_stable_stats( + d_recv_w_from_intermediate, active_recv_rows + ) + d_recv_input_stats = _debug_stable_stats( + dispatch_weight_cotangent, active_recv_rows + ) + d_dispatch_stats = _debug_stable_stats(d_x_from_dispatch) + d_topk_stats = _debug_stable_stats(d_topk_w) + jax.debug.print( + "[TE dispatch-bwd stats] " + "token_in(absmean={token_in_absmean:.3e},std={token_in_std:.3e}," + "absmax={token_in_absmax:.3e},finite={token_in_finite:.6f}) " + "weight_combine(absmean={wc_absmean:.3e},absmax={wc_absmax:.3e}) " + "weight_ffn(absmean={wf_absmean:.3e},absmax={wf_absmax:.3e}) " + "weight_input(absmean={wi_absmean:.3e},absmax={wi_absmax:.3e}) " + "token_out(absmean={token_out_absmean:.3e},std={token_out_std:.3e}," + "absmax={token_out_absmax:.3e},finite={token_out_finite:.6f}) " + "weight_out(absmean={weight_out_absmean:.3e},std={weight_out_std:.3e}," + "absmax={weight_out_absmax:.3e},finite={weight_out_finite:.6f})", + token_in_absmean=d_sorted_x_stats[1], + token_in_std=d_sorted_x_stats[2], + token_in_absmax=d_sorted_x_stats[3], + token_in_finite=d_sorted_x_stats[4], + wc_absmean=d_recv_combine_stats[1], + wc_absmax=d_recv_combine_stats[3], + wf_absmean=d_recv_intermediate_stats[1], + wf_absmax=d_recv_intermediate_stats[3], + wi_absmean=d_recv_input_stats[1], + wi_absmax=d_recv_input_stats[3], + token_out_absmean=d_dispatch_stats[1], + token_out_std=d_dispatch_stats[2], + token_out_absmax=d_dispatch_stats[3], + token_out_finite=d_dispatch_stats[4], + weight_out_absmean=d_topk_stats[1], + weight_out_std=d_topk_stats[2], + weight_out_absmax=d_topk_stats[3], + weight_out_finite=d_topk_stats[4], + ordered=False, + ) + # ---------------- Routing bwd (global view) ---------------- # The cotangent on routing_weights is a sparse scatter into sparse_probs # at the selected_experts indices. @@ -957,6 +1893,26 @@ def _fold_dp_groups(grad): compute_aux_scores=False, ) + if _debug_moe_numerics or _debug_moe_input_grad: + sparse_prob_stats = _debug_stable_stats(d_sparse_probs) + logits_stats = _debug_stable_stats(d_logits_2d) + jax.debug.print( + "[TE router-bwd stats] " + "sparse_in(absmean={sparse_absmean:.3e},std={sparse_std:.3e}," + "absmax={sparse_absmax:.3e},finite={sparse_finite:.6f}) " + "logits_out(absmean={logits_absmean:.3e},std={logits_std:.3e}," + "absmax={logits_absmax:.3e},finite={logits_finite:.6f})", + sparse_absmean=sparse_prob_stats[1], + sparse_std=sparse_prob_stats[2], + sparse_absmax=sparse_prob_stats[3], + sparse_finite=sparse_prob_stats[4], + logits_absmean=logits_stats[1], + logits_std=logits_stats[2], + logits_absmax=logits_stats[3], + logits_finite=logits_stats[4], + ordered=False, + ) + # ---------------- 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 @@ -992,6 +1948,32 @@ def _fold_dp_groups(grad): 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 + if _zero_moe_input_grad: + d_x = jnp.zeros_like(d_x) + + if _debug_moe_numerics or _debug_moe_input_grad: + dispatch_stats = _debug_stable_stats(d_x_from_dispatch) + gate_stats = _debug_stable_stats(d_x_from_gate) + total_stats = _debug_stable_stats(d_x) + jax.debug.print( + "[TE input-grad stats] " + "dispatch(absmean={dispatch_absmean:.3e},std={dispatch_std:.3e}," + "absmax={dispatch_absmax:.3e}) " + "gate(absmean={gate_absmean:.3e},std={gate_std:.3e},absmax={gate_absmax:.3e}) " + "total(absmean={total_absmean:.3e},std={total_std:.3e}," + "absmax={total_absmax:.3e},finite={total_finite:.6f})", + dispatch_absmean=dispatch_stats[1], + dispatch_std=dispatch_stats[2], + dispatch_absmax=dispatch_stats[3], + gate_absmean=gate_stats[1], + gate_std=gate_stats[2], + gate_absmax=gate_stats[3], + total_absmean=total_stats[1], + total_std=total_stats[2], + total_absmax=total_stats[3], + total_finite=total_stats[4], + ordered=False, + ) # Pin output grads to the declared logical axes so downstream # optimizers see consistent shardings. From 6c5cea4ad1d38a3d3ad0808bf445851ce5242598 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 30 Jul 2026 15:01:26 -0700 Subject: [PATCH 24/44] Add synchronous MoE FFI diagnostics --- tests/jax/test_te_ep_moe.py | 2 ++ transformer_engine/jax/csrc/extensions/ep.cpp | 25 +++++++---------- .../jax/csrc/extensions/ffi.cpp | 27 +++++++++++++++++++ transformer_engine/jax/csrc/extensions/ffi.h | 2 ++ .../jax/csrc/extensions/gemm.cpp | 7 +++-- .../jax/csrc/extensions/router.cpp | 12 ++++----- 6 files changed, 50 insertions(+), 25 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index fc3d20b07e..176275a9bb 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -1090,6 +1090,7 @@ def reference_loss_fn(params, value): name="repeated d_x", ) + @pytest.mark.skip(reason="Experimental scan/remat parity coverage is currently disabled.") @pytest.mark.parametrize("num_layers", (4, 8)) def test_scanned_production_block_backward(self, mesh, num_layers): """Full MoE scan/remat must match the identical unrolled stack. @@ -1242,6 +1243,7 @@ def reference_loss_fn(p, value): if mismatches: pytest.fail("\n\n".join(mismatches)) + @pytest.mark.skip(reason="Experimental scan/remat parity coverage is currently disabled.") def test_scanned_training_trajectory(self, mesh): """Three SGD steps must retain scan/remat vs unrolled parity.""" num_layers = 4 diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index cdd7730204..b28a84355b 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -224,7 +224,7 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T static_cast(config.dispatch_output_per_expert_alignment)}; nvte_ep_prepare(handle_mem_.data(), topk_idx_.data(), recv_tokens_per_expert_.data(), /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check(stream, "EpPrepareFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, @@ -234,8 +234,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, .Arg() // topk_idx .Ret() // recv_tokens_per_expert .Ret() // handle_mem - .Attrs(), - FFI_CudaGraph_Traits); + .Attrs()); // ── ep_dispatch ─────────────────────────────────────────────────────────────── @@ -299,7 +298,7 @@ Error_Type EpDispatchFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_ topk_weights_.data(), no_win, recv_tokens_.data(), no_win, recv_topk_weights_.data(), no_win, stream); - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check(stream, "EpDispatchFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchHandler, EpDispatchFFI, @@ -312,8 +311,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchHandler, EpDispatchFFI, .Arg() // topk_weights .Ret() // recv_tokens .Ret() // recv_topk_weights - .Attrs(), - FFI_CudaGraph_Traits); + .Attrs()); // ── ep_combine ──────────────────────────────────────────────────────────────── @@ -343,7 +341,7 @@ Error_Type EpCombineFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T NVTECommWindow no_win{nullptr, 0}; nvte_ep_combine(handle_mem_.data(), expert_out_.data(), no_win, result_.data(), stream); - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check(stream, "EpCombineFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineHandler, EpCombineFFI, @@ -353,8 +351,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineHandler, EpCombineFFI, .Arg() // handle_mem .Arg() // expert_out .Ret() // result - .Attrs(), - FFI_CudaGraph_Traits); + .Attrs()); // ── ep_dispatch_bwd ─────────────────────────────────────────────────────────── @@ -411,7 +408,7 @@ Error_Type EpDispatchBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buff nvte_ep_dispatch_bwd(handle_mem_.data(), grad_.data(), no_win, g_recv_topk_weights_.data(), no_win, grad_tokens_.data(), grad_topk_weights_.data(), stream); - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check(stream, "EpDispatchBwdFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchBwdHandler, EpDispatchBwdFFI, @@ -423,8 +420,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchBwdHandler, EpDispatchBwdFFI, .Arg() // g_recv_topk_weights .Ret() // grad_tokens .Ret() // grad_topk_weights - .Attrs(), - FFI_CudaGraph_Traits); + .Attrs()); // ── ep_combine_bwd ──────────────────────────────────────────────────────────── @@ -457,7 +453,7 @@ Error_Type EpCombineBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffe nvte_ep_combine_bwd(handle_mem_.data(), grad_.data(), no_win, grad_expert_out_.data(), no_win, stream); - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check(stream, "EpCombineBwdFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineBwdHandler, EpCombineBwdFFI, @@ -467,8 +463,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineBwdHandler, EpCombineBwdFFI, .Arg() // handle_mem .Arg() // grad (w.r.t. result) .Ret() // grad_expert_out - .Attrs(), - FFI_CudaGraph_Traits); + .Attrs()); } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/ffi.cpp b/transformer_engine/jax/csrc/extensions/ffi.cpp index 6bb2f18234..fde21db5e7 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.cpp +++ b/transformer_engine/jax/csrc/extensions/ffi.cpp @@ -61,5 +61,32 @@ Error_Type ffi_with_cuda_error_check() { return Error_Type::Success(); } +Error_Type ffi_with_cuda_stream_sync_and_error_check(cudaStream_t stream, + const char* operation_name) { + cudaError_t launch_error = cudaGetLastError(); + if (launch_error != cudaSuccess) { + return Error_Type(XLA_FFI_Error_Code_INTERNAL, + std::string(operation_name) + + " CUDA launch error before stream synchronization: " + + cudaGetErrorString(launch_error)); + } + + cudaError_t sync_error = cudaStreamSynchronize(stream); + if (sync_error != cudaSuccess) { + return Error_Type(XLA_FFI_Error_Code_INTERNAL, + std::string(operation_name) + " CUDA stream synchronization error: " + + cudaGetErrorString(sync_error)); + } + + cudaError_t last_error = cudaGetLastError(); + if (last_error != cudaSuccess) { + return Error_Type(XLA_FFI_Error_Code_INTERNAL, + std::string(operation_name) + + " CUDA error after stream synchronization: " + + cudaGetErrorString(last_error)); + } + return Error_Type::Success(); +} + } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/ffi.h b/transformer_engine/jax/csrc/extensions/ffi.h index f9d327102b..7626419929 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.h +++ b/transformer_engine/jax/csrc/extensions/ffi.h @@ -30,6 +30,8 @@ constexpr auto FFI_CudaGraph_Traits = {xla::ffi::Traits::kCmdBufferCompatible}; DType convert_ffi_datatype_to_te_dtype(const xla::ffi::DataType& type); Error_Type ffi_with_cuda_error_check(); +Error_Type ffi_with_cuda_stream_sync_and_error_check(cudaStream_t stream, + const char* operation_name); // source_location is not available in C++17, so we implement it ourselves #if defined(__GNUC__) || defined(__clang__) diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index a36bf6cd22..e246910922 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -1017,7 +1017,7 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), workspace_cublas.data(), gemmConfig, stream); - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check(stream, "GroupedGemmV2FFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmV2Handler, GroupedGemmV2FFI, @@ -1040,8 +1040,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmV2Handler, GroupedGemmV2FFI, .Ret() // cublas_workspace .Ret() // setup_workspace .Ret() // int64_workspace - .Attrs(), - FFI_CudaGraph_Traits); + .Attrs()); Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, Buffer_Type rhs_data, Buffer_Type rhs_sinv, Buffer_Type bias, @@ -1425,7 +1424,7 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type grad, workspace_list.data(), accumulate, use_split_accumulator, num_math_sm, stream); - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check(stream, "GroupedGemmFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmHandler, GroupedGemmFFI, diff --git a/transformer_engine/jax/csrc/extensions/router.cpp b/transformer_engine/jax/csrc/extensions/router.cpp index a1d0731b6d..3682f0d40e 100644 --- a/transformer_engine/jax/csrc/extensions/router.cpp +++ b/transformer_engine/jax/csrc/extensions/router.cpp @@ -80,7 +80,8 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( routing_map_tensor.data(), routing_map_format_nvte, intermediate_tensor.data(), stream); } - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check( + stream, "FusedTopkWithScoreFunctionForwardFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, @@ -99,8 +100,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, .Attr("scaling_factor") .Attr("score_function") .Attr("compute_aux_scores") - .Attr("routing_map_format"), - FFI_CudaGraph_Traits); + .Attr("routing_map_format")); // ============================================================================ // Fused Top-K with Score Function - Backward @@ -157,7 +157,8 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( static_cast(score_function), grad_logits_tensor.data(), stream); } - return ffi_with_cuda_error_check(); + return ffi_with_cuda_stream_sync_and_error_check( + stream, "FusedTopkWithScoreFunctionBackwardFFI"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler, @@ -173,8 +174,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler, .Attr("scaling_factor") .Attr("score_function") .Attr("compute_aux_scores") - .Attr("routing_map_format"), - FFI_CudaGraph_Traits); + .Attr("routing_map_format")); // ============================================================================ // Fused MoE Aux Loss - Forward From 0a34e1e5a3c913c84290e2e927a785935db9845c Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 30 Jul 2026 16:30:36 -0700 Subject: [PATCH 25/44] Add device sync diagnostics to MoE FFI ops --- transformer_engine/jax/csrc/extensions/ep.cpp | 20 +++++++++---- .../jax/csrc/extensions/ffi.cpp | 29 +++++++++---------- transformer_engine/jax/csrc/extensions/ffi.h | 4 +-- .../jax/csrc/extensions/gemm.cpp | 8 +++-- .../jax/csrc/extensions/router.cpp | 14 ++++++--- 5 files changed, 46 insertions(+), 29 deletions(-) diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index b28a84355b..9058d8cfd7 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -198,6 +198,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpInstantiateHandler, EpInstantiateImpl, FFI::Bind Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type topk_idx, Result_Type recv_tokens_per_expert, Result_Type handle_mem, EpConfig config) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpPrepareFFI", "begin"); + if (sync_error.failure()) return sync_error; (void)ep_state; // lifetime only. auto topk_dims = topk_idx.dimensions(); NVTE_CHECK(topk_dims.size() >= 2, @@ -224,7 +226,7 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T static_cast(config.dispatch_output_per_expert_alignment)}; nvte_ep_prepare(handle_mem_.data(), topk_idx_.data(), recv_tokens_per_expert_.data(), /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); - return ffi_with_cuda_stream_sync_and_error_check(stream, "EpPrepareFFI"); + return ffi_with_cuda_device_sync_and_error_check("EpPrepareFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, @@ -241,6 +243,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, Error_Type EpDispatchFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type handle_mem, Buffer_Type topk_idx, Buffer_Type tokens, Buffer_Type topk_weights, Result_Type recv_tokens, Result_Type recv_topk_weights, EpConfig config) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpDispatchFFI", "begin"); + if (sync_error.failure()) return sync_error; (void)ep_state; auto token_dims = tokens.dimensions(); NVTE_CHECK(token_dims.size() >= 2, @@ -298,7 +302,7 @@ Error_Type EpDispatchFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_ topk_weights_.data(), no_win, recv_tokens_.data(), no_win, recv_topk_weights_.data(), no_win, stream); - return ffi_with_cuda_stream_sync_and_error_check(stream, "EpDispatchFFI"); + return ffi_with_cuda_device_sync_and_error_check("EpDispatchFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchHandler, EpDispatchFFI, @@ -317,6 +321,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchHandler, EpDispatchFFI, Error_Type EpCombineFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type handle_mem, Buffer_Type expert_out, Result_Type result, EpConfig config) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpCombineFFI", "begin"); + if (sync_error.failure()) return sync_error; (void)ep_state; auto eo_dims = expert_out.dimensions(); NVTE_CHECK(eo_dims.size() >= 2, @@ -341,7 +347,7 @@ Error_Type EpCombineFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T NVTECommWindow no_win{nullptr, 0}; nvte_ep_combine(handle_mem_.data(), expert_out_.data(), no_win, result_.data(), stream); - return ffi_with_cuda_stream_sync_and_error_check(stream, "EpCombineFFI"); + return ffi_with_cuda_device_sync_and_error_check("EpCombineFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineHandler, EpCombineFFI, @@ -359,6 +365,8 @@ Error_Type EpDispatchBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buff Buffer_Type grad, Buffer_Type g_recv_topk_weights, Result_Type grad_tokens, Result_Type grad_topk_weights, EpConfig config) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpDispatchBwdFFI", "begin"); + if (sync_error.failure()) return sync_error; (void)ep_state; auto grad_dims = grad.dimensions(); NVTE_CHECK(grad_dims.size() >= 2, @@ -408,7 +416,7 @@ Error_Type EpDispatchBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buff nvte_ep_dispatch_bwd(handle_mem_.data(), grad_.data(), no_win, g_recv_topk_weights_.data(), no_win, grad_tokens_.data(), grad_topk_weights_.data(), stream); - return ffi_with_cuda_stream_sync_and_error_check(stream, "EpDispatchBwdFFI"); + return ffi_with_cuda_device_sync_and_error_check("EpDispatchBwdFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchBwdHandler, EpDispatchBwdFFI, @@ -426,6 +434,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchBwdHandler, EpDispatchBwdFFI, Error_Type EpCombineBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type handle_mem, Buffer_Type grad, Result_Type grad_expert_out, EpConfig config) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpCombineBwdFFI", "begin"); + if (sync_error.failure()) return sync_error; (void)ep_state; auto grad_dims = grad.dimensions(); NVTE_CHECK(grad_dims.size() >= 2, @@ -453,7 +463,7 @@ Error_Type EpCombineBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffe nvte_ep_combine_bwd(handle_mem_.data(), grad_.data(), no_win, grad_expert_out_.data(), no_win, stream); - return ffi_with_cuda_stream_sync_and_error_check(stream, "EpCombineBwdFFI"); + return ffi_with_cuda_device_sync_and_error_check("EpCombineBwdFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineBwdHandler, EpCombineBwdFFI, diff --git a/transformer_engine/jax/csrc/extensions/ffi.cpp b/transformer_engine/jax/csrc/extensions/ffi.cpp index fde21db5e7..fd9ea4d39b 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.cpp +++ b/transformer_engine/jax/csrc/extensions/ffi.cpp @@ -5,6 +5,7 @@ ************************************************************************/ #include "extensions/ffi.h" +#include #include namespace transformer_engine { @@ -61,29 +62,25 @@ Error_Type ffi_with_cuda_error_check() { return Error_Type::Success(); } -Error_Type ffi_with_cuda_stream_sync_and_error_check(cudaStream_t stream, - const char* operation_name) { - cudaError_t launch_error = cudaGetLastError(); - if (launch_error != cudaSuccess) { - return Error_Type(XLA_FFI_Error_Code_INTERNAL, - std::string(operation_name) + - " CUDA launch error before stream synchronization: " + - cudaGetErrorString(launch_error)); - } - - cudaError_t sync_error = cudaStreamSynchronize(stream); +Error_Type ffi_with_cuda_device_sync_and_error_check(const char* operation_name, + const char* synchronization_phase) { + cudaError_t sync_error = cudaDeviceSynchronize(); + std::printf( + "[TE FFI device sync] operation=%s phase=%s code=%d name=%s message=%s\n", + operation_name, synchronization_phase, static_cast(sync_error), + cudaGetErrorName(sync_error), cudaGetErrorString(sync_error)); + std::fflush(stdout); if (sync_error != cudaSuccess) { return Error_Type(XLA_FFI_Error_Code_INTERNAL, - std::string(operation_name) + " CUDA stream synchronization error: " + - cudaGetErrorString(sync_error)); + std::string(operation_name) + " CUDA device synchronization error at " + + synchronization_phase + ": " + cudaGetErrorString(sync_error)); } cudaError_t last_error = cudaGetLastError(); if (last_error != cudaSuccess) { return Error_Type(XLA_FFI_Error_Code_INTERNAL, - std::string(operation_name) + - " CUDA error after stream synchronization: " + - cudaGetErrorString(last_error)); + std::string(operation_name) + " CUDA error after device synchronization at " + + synchronization_phase + ": " + cudaGetErrorString(last_error)); } return Error_Type::Success(); } diff --git a/transformer_engine/jax/csrc/extensions/ffi.h b/transformer_engine/jax/csrc/extensions/ffi.h index 7626419929..cdf90de8c5 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.h +++ b/transformer_engine/jax/csrc/extensions/ffi.h @@ -30,8 +30,8 @@ constexpr auto FFI_CudaGraph_Traits = {xla::ffi::Traits::kCmdBufferCompatible}; DType convert_ffi_datatype_to_te_dtype(const xla::ffi::DataType& type); Error_Type ffi_with_cuda_error_check(); -Error_Type ffi_with_cuda_stream_sync_and_error_check(cudaStream_t stream, - const char* operation_name); +Error_Type ffi_with_cuda_device_sync_and_error_check(const char* operation_name, + const char* synchronization_phase); // source_location is not available in C++17, so we implement it ourselves #if defined(__GNUC__) || defined(__clang__) diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index e246910922..5be03e3f46 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -926,6 +926,8 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty Buffer_Type alpha, Buffer_Type beta, Result_Type output, Result_Type cublas_workspace, Result_Type setup_workspace, Result_Type int64_workspace, GroupedGemmV2Config config) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check("GroupedGemmV2FFI", "begin"); + if (sync_error.failure()) return sync_error; auto [lhs_is_trans, rhs_is_trans, scaling_mode, lhs_axis_boundary, rhs_axis_boundary, lhs_left_size, lhs_right_size, rhs_left_size, rhs_right_size] = config; @@ -1017,7 +1019,7 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), workspace_cublas.data(), gemmConfig, stream); - return ffi_with_cuda_stream_sync_and_error_check(stream, "GroupedGemmV2FFI"); + return ffi_with_cuda_device_sync_and_error_check("GroupedGemmV2FFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmV2Handler, GroupedGemmV2FFI, @@ -1049,6 +1051,8 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type Buffer_Type out_first_dims, Buffer_Type out_last_dims, Buffer_Type group_offset, Result_Type output, Result_Type workspace, GroupedGemmConfig config) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check("GroupedGemmFFI", "begin"); + if (sync_error.failure()) return sync_error; auto [lhs_is_trans, rhs_is_trans, scaling_mode, has_bias, use_async_d2h_group_sizes, lhs_axis_boundary, rhs_axis_boundary, lhs_left_size, lhs_right_size, rhs_left_size, rhs_right_size] = config; @@ -1424,7 +1428,7 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type grad, workspace_list.data(), accumulate, use_split_accumulator, num_math_sm, stream); - return ffi_with_cuda_stream_sync_and_error_check(stream, "GroupedGemmFFI"); + return ffi_with_cuda_device_sync_and_error_check("GroupedGemmFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmHandler, GroupedGemmFFI, diff --git a/transformer_engine/jax/csrc/extensions/router.cpp b/transformer_engine/jax/csrc/extensions/router.cpp index 3682f0d40e..b591e7a7a4 100644 --- a/transformer_engine/jax/csrc/extensions/router.cpp +++ b/transformer_engine/jax/csrc/extensions/router.cpp @@ -26,6 +26,9 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( int64_t topk, int64_t use_pre_softmax, int64_t num_groups, int64_t group_topk, double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores, JAXX_Routing_Map_Format routing_map_format) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check( + "FusedTopkWithScoreFunctionForwardFFI", "begin"); + if (sync_error.failure()) return sync_error; auto dtype = convert_ffi_datatype_to_te_dtype(logits_buf.element_type()); auto dims = logits_buf.dimensions(); auto num_tokens = static_cast(product(dims, 0, dims.size() - 1)); @@ -80,8 +83,8 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( routing_map_tensor.data(), routing_map_format_nvte, intermediate_tensor.data(), stream); } - return ffi_with_cuda_stream_sync_and_error_check( - stream, "FusedTopkWithScoreFunctionForwardFFI"); + return ffi_with_cuda_device_sync_and_error_check( + "FusedTopkWithScoreFunctionForwardFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, @@ -115,6 +118,9 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( int64_t topk, int64_t use_pre_softmax, double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores, JAXX_Routing_Map_Format routing_map_format) { + auto sync_error = ffi_with_cuda_device_sync_and_error_check( + "FusedTopkWithScoreFunctionBackwardFFI", "begin"); + if (sync_error.failure()) return sync_error; // intermediate is always float32 (CompType) regardless of logits dtype. auto intermediate_dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf.element_type()); NVTE_CHECK( @@ -157,8 +163,8 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( static_cast(score_function), grad_logits_tensor.data(), stream); } - return ffi_with_cuda_stream_sync_and_error_check( - stream, "FusedTopkWithScoreFunctionBackwardFFI"); + return ffi_with_cuda_device_sync_and_error_check( + "FusedTopkWithScoreFunctionBackwardFFI", "end"); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler, From 307aa1ed17a354f9b09c696d2484b9906a91e8db Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Fri, 31 Jul 2026 08:13:33 -0700 Subject: [PATCH 26/44] Restore graph-compatible MoE FFI execution --- transformer_engine/common/ep/ep_api.cpp | 23 +++++++++--- transformer_engine/jax/csrc/extensions/ep.cpp | 35 ++++++++----------- .../jax/csrc/extensions/ffi.cpp | 24 ------------- transformer_engine/jax/csrc/extensions/ffi.h | 2 -- .../jax/csrc/extensions/gemm.cpp | 11 +++--- .../jax/csrc/extensions/router.cpp | 18 ++++------ 6 files changed, 43 insertions(+), 70 deletions(-) diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index fff90dd8f5..f29b9990f2 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -15,6 +15,8 @@ #include #include +#include +#include #include #include "../util/logging.h" @@ -57,6 +59,17 @@ inline void* handle_mem_ptr(NVTETensor mem) { NVTE_CHECK(p != nullptr, "handle_mem tensor data must not be null"); return p; } + +bool trace_handle_mem_ptrs() { + const char* value = std::getenv("NVTE_EP_TRACE_HANDLE_MEM_PTRS"); + return value != nullptr && std::strcmp(value, "0") != 0; +} + +void trace_handle_mem_ptr(const char* operation, NVTETensor handle_mem) { + if (!trace_handle_mem_ptrs()) return; + std::printf("%s handle_mem_ptr: %p\n", operation, handle_mem_ptr(handle_mem)); + std::fflush(stdout); +} } // namespace void nvte_ep_initialize(void* ep_comm, const NVTEEpGroupConfig* group_config) { @@ -75,7 +88,7 @@ size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg) { void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv_tokens_per_expert, NVTETensor total_recv_tokens_per_rank, const NVTEEpLayerConfig* layer_cfg, cudaStream_t stream) { - printf("nvte_ep_prepare handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); + trace_handle_mem_ptr("nvte_ep_prepare", handle_mem); NVTEEpLayerConfig cfg = normalize_ep_config(layer_cfg, kLayerConfigMinSize, "layer_cfg"); EPBackend::get().prepare(handle_mem_ptr(handle_mem), topk_idx, recv_tokens_per_expert, total_recv_tokens_per_rank, cfg, stream); @@ -86,7 +99,7 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok NVTECommWindow topk_weights_win, NVTETensor recv_tokens, NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, NVTECommWindow recv_topk_weights_win, cudaStream_t stream) { - printf("nvte_ep_dispatch handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); + trace_handle_mem_ptr("nvte_ep_dispatch", handle_mem); EPBackend::get().dispatch(handle_mem_ptr(handle_mem), topk_idx, tokens, tokens_win, topk_weights, topk_weights_win, recv_tokens, recv_tokens_win, recv_topk_weights, recv_topk_weights_win, stream); @@ -94,7 +107,7 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok void nvte_ep_combine(NVTETensor handle_mem, NVTETensor expert_out, NVTECommWindow expert_out_win, NVTETensor result, cudaStream_t stream) { - printf("nvte_ep_combine handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); + trace_handle_mem_ptr("nvte_ep_combine", handle_mem); EPBackend::get().combine(handle_mem_ptr(handle_mem), expert_out, expert_out_win, result, stream); } @@ -102,7 +115,7 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow NVTETensor g_recv_topk_weights, NVTECommWindow g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream) { - printf("nvte_ep_dispatch_bwd handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); + trace_handle_mem_ptr("nvte_ep_dispatch_bwd", handle_mem); EPBackend::get().dispatch_bwd(handle_mem_ptr(handle_mem), grad, grad_win, g_recv_topk_weights, g_recv_topk_weights_win, grad_tokens, grad_topk_weights, stream); } @@ -110,7 +123,7 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, cudaStream_t stream) { - printf("nvte_combine_bwd handle_mem_ptr: 0x%lx\n", handle_mem_ptr(handle_mem)); + trace_handle_mem_ptr("nvte_ep_combine_bwd", handle_mem); EPBackend::get().combine_bwd(handle_mem_ptr(handle_mem), grad, grad_win, grad_expert_out, grad_expert_out_win, stream); } diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index 9058d8cfd7..cdd7730204 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -198,8 +198,6 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpInstantiateHandler, EpInstantiateImpl, FFI::Bind Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type topk_idx, Result_Type recv_tokens_per_expert, Result_Type handle_mem, EpConfig config) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpPrepareFFI", "begin"); - if (sync_error.failure()) return sync_error; (void)ep_state; // lifetime only. auto topk_dims = topk_idx.dimensions(); NVTE_CHECK(topk_dims.size() >= 2, @@ -226,7 +224,7 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T static_cast(config.dispatch_output_per_expert_alignment)}; nvte_ep_prepare(handle_mem_.data(), topk_idx_.data(), recv_tokens_per_expert_.data(), /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); - return ffi_with_cuda_device_sync_and_error_check("EpPrepareFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, @@ -236,15 +234,14 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, .Arg() // topk_idx .Ret() // recv_tokens_per_expert .Ret() // handle_mem - .Attrs()); + .Attrs(), + FFI_CudaGraph_Traits); // ── ep_dispatch ─────────────────────────────────────────────────────────────── Error_Type EpDispatchFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type handle_mem, Buffer_Type topk_idx, Buffer_Type tokens, Buffer_Type topk_weights, Result_Type recv_tokens, Result_Type recv_topk_weights, EpConfig config) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpDispatchFFI", "begin"); - if (sync_error.failure()) return sync_error; (void)ep_state; auto token_dims = tokens.dimensions(); NVTE_CHECK(token_dims.size() >= 2, @@ -302,7 +299,7 @@ Error_Type EpDispatchFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_ topk_weights_.data(), no_win, recv_tokens_.data(), no_win, recv_topk_weights_.data(), no_win, stream); - return ffi_with_cuda_device_sync_and_error_check("EpDispatchFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchHandler, EpDispatchFFI, @@ -315,14 +312,13 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchHandler, EpDispatchFFI, .Arg() // topk_weights .Ret() // recv_tokens .Ret() // recv_topk_weights - .Attrs()); + .Attrs(), + FFI_CudaGraph_Traits); // ── ep_combine ──────────────────────────────────────────────────────────────── Error_Type EpCombineFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type handle_mem, Buffer_Type expert_out, Result_Type result, EpConfig config) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpCombineFFI", "begin"); - if (sync_error.failure()) return sync_error; (void)ep_state; auto eo_dims = expert_out.dimensions(); NVTE_CHECK(eo_dims.size() >= 2, @@ -347,7 +343,7 @@ Error_Type EpCombineFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T NVTECommWindow no_win{nullptr, 0}; nvte_ep_combine(handle_mem_.data(), expert_out_.data(), no_win, result_.data(), stream); - return ffi_with_cuda_device_sync_and_error_check("EpCombineFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineHandler, EpCombineFFI, @@ -357,7 +353,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineHandler, EpCombineFFI, .Arg() // handle_mem .Arg() // expert_out .Ret() // result - .Attrs()); + .Attrs(), + FFI_CudaGraph_Traits); // ── ep_dispatch_bwd ─────────────────────────────────────────────────────────── @@ -365,8 +362,6 @@ Error_Type EpDispatchBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buff Buffer_Type grad, Buffer_Type g_recv_topk_weights, Result_Type grad_tokens, Result_Type grad_topk_weights, EpConfig config) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpDispatchBwdFFI", "begin"); - if (sync_error.failure()) return sync_error; (void)ep_state; auto grad_dims = grad.dimensions(); NVTE_CHECK(grad_dims.size() >= 2, @@ -416,7 +411,7 @@ Error_Type EpDispatchBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buff nvte_ep_dispatch_bwd(handle_mem_.data(), grad_.data(), no_win, g_recv_topk_weights_.data(), no_win, grad_tokens_.data(), grad_topk_weights_.data(), stream); - return ffi_with_cuda_device_sync_and_error_check("EpDispatchBwdFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchBwdHandler, EpDispatchBwdFFI, @@ -428,14 +423,13 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchBwdHandler, EpDispatchBwdFFI, .Arg() // g_recv_topk_weights .Ret() // grad_tokens .Ret() // grad_topk_weights - .Attrs()); + .Attrs(), + FFI_CudaGraph_Traits); // ── ep_combine_bwd ──────────────────────────────────────────────────────────── Error_Type EpCombineBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type handle_mem, Buffer_Type grad, Result_Type grad_expert_out, EpConfig config) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check("EpCombineBwdFFI", "begin"); - if (sync_error.failure()) return sync_error; (void)ep_state; auto grad_dims = grad.dimensions(); NVTE_CHECK(grad_dims.size() >= 2, @@ -463,7 +457,7 @@ Error_Type EpCombineBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffe nvte_ep_combine_bwd(handle_mem_.data(), grad_.data(), no_win, grad_expert_out_.data(), no_win, stream); - return ffi_with_cuda_device_sync_and_error_check("EpCombineBwdFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineBwdHandler, EpCombineBwdFFI, @@ -473,7 +467,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineBwdHandler, EpCombineBwdFFI, .Arg() // handle_mem .Arg() // grad (w.r.t. result) .Ret() // grad_expert_out - .Attrs()); + .Attrs(), + FFI_CudaGraph_Traits); } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/ffi.cpp b/transformer_engine/jax/csrc/extensions/ffi.cpp index fd9ea4d39b..6bb2f18234 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.cpp +++ b/transformer_engine/jax/csrc/extensions/ffi.cpp @@ -5,7 +5,6 @@ ************************************************************************/ #include "extensions/ffi.h" -#include #include namespace transformer_engine { @@ -62,28 +61,5 @@ Error_Type ffi_with_cuda_error_check() { return Error_Type::Success(); } -Error_Type ffi_with_cuda_device_sync_and_error_check(const char* operation_name, - const char* synchronization_phase) { - cudaError_t sync_error = cudaDeviceSynchronize(); - std::printf( - "[TE FFI device sync] operation=%s phase=%s code=%d name=%s message=%s\n", - operation_name, synchronization_phase, static_cast(sync_error), - cudaGetErrorName(sync_error), cudaGetErrorString(sync_error)); - std::fflush(stdout); - if (sync_error != cudaSuccess) { - return Error_Type(XLA_FFI_Error_Code_INTERNAL, - std::string(operation_name) + " CUDA device synchronization error at " + - synchronization_phase + ": " + cudaGetErrorString(sync_error)); - } - - cudaError_t last_error = cudaGetLastError(); - if (last_error != cudaSuccess) { - return Error_Type(XLA_FFI_Error_Code_INTERNAL, - std::string(operation_name) + " CUDA error after device synchronization at " + - synchronization_phase + ": " + cudaGetErrorString(last_error)); - } - return Error_Type::Success(); -} - } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/ffi.h b/transformer_engine/jax/csrc/extensions/ffi.h index cdf90de8c5..f9d327102b 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.h +++ b/transformer_engine/jax/csrc/extensions/ffi.h @@ -30,8 +30,6 @@ constexpr auto FFI_CudaGraph_Traits = {xla::ffi::Traits::kCmdBufferCompatible}; DType convert_ffi_datatype_to_te_dtype(const xla::ffi::DataType& type); Error_Type ffi_with_cuda_error_check(); -Error_Type ffi_with_cuda_device_sync_and_error_check(const char* operation_name, - const char* synchronization_phase); // source_location is not available in C++17, so we implement it ourselves #if defined(__GNUC__) || defined(__clang__) diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 5be03e3f46..a36bf6cd22 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -926,8 +926,6 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty Buffer_Type alpha, Buffer_Type beta, Result_Type output, Result_Type cublas_workspace, Result_Type setup_workspace, Result_Type int64_workspace, GroupedGemmV2Config config) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check("GroupedGemmV2FFI", "begin"); - if (sync_error.failure()) return sync_error; auto [lhs_is_trans, rhs_is_trans, scaling_mode, lhs_axis_boundary, rhs_axis_boundary, lhs_left_size, lhs_right_size, rhs_left_size, rhs_right_size] = config; @@ -1019,7 +1017,7 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), workspace_cublas.data(), gemmConfig, stream); - return ffi_with_cuda_device_sync_and_error_check("GroupedGemmV2FFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmV2Handler, GroupedGemmV2FFI, @@ -1042,7 +1040,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmV2Handler, GroupedGemmV2FFI, .Ret() // cublas_workspace .Ret() // setup_workspace .Ret() // int64_workspace - .Attrs()); + .Attrs(), + FFI_CudaGraph_Traits); Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, Buffer_Type rhs_data, Buffer_Type rhs_sinv, Buffer_Type bias, @@ -1051,8 +1050,6 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type Buffer_Type out_first_dims, Buffer_Type out_last_dims, Buffer_Type group_offset, Result_Type output, Result_Type workspace, GroupedGemmConfig config) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check("GroupedGemmFFI", "begin"); - if (sync_error.failure()) return sync_error; auto [lhs_is_trans, rhs_is_trans, scaling_mode, has_bias, use_async_d2h_group_sizes, lhs_axis_boundary, rhs_axis_boundary, lhs_left_size, lhs_right_size, rhs_left_size, rhs_right_size] = config; @@ -1428,7 +1425,7 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type grad, workspace_list.data(), accumulate, use_split_accumulator, num_math_sm, stream); - return ffi_with_cuda_device_sync_and_error_check("GroupedGemmFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmHandler, GroupedGemmFFI, diff --git a/transformer_engine/jax/csrc/extensions/router.cpp b/transformer_engine/jax/csrc/extensions/router.cpp index b591e7a7a4..a1d0731b6d 100644 --- a/transformer_engine/jax/csrc/extensions/router.cpp +++ b/transformer_engine/jax/csrc/extensions/router.cpp @@ -26,9 +26,6 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( int64_t topk, int64_t use_pre_softmax, int64_t num_groups, int64_t group_topk, double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores, JAXX_Routing_Map_Format routing_map_format) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check( - "FusedTopkWithScoreFunctionForwardFFI", "begin"); - if (sync_error.failure()) return sync_error; auto dtype = convert_ffi_datatype_to_te_dtype(logits_buf.element_type()); auto dims = logits_buf.dimensions(); auto num_tokens = static_cast(product(dims, 0, dims.size() - 1)); @@ -83,8 +80,7 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( routing_map_tensor.data(), routing_map_format_nvte, intermediate_tensor.data(), stream); } - return ffi_with_cuda_device_sync_and_error_check( - "FusedTopkWithScoreFunctionForwardFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, @@ -103,7 +99,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, .Attr("scaling_factor") .Attr("score_function") .Attr("compute_aux_scores") - .Attr("routing_map_format")); + .Attr("routing_map_format"), + FFI_CudaGraph_Traits); // ============================================================================ // Fused Top-K with Score Function - Backward @@ -118,9 +115,6 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( int64_t topk, int64_t use_pre_softmax, double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores, JAXX_Routing_Map_Format routing_map_format) { - auto sync_error = ffi_with_cuda_device_sync_and_error_check( - "FusedTopkWithScoreFunctionBackwardFFI", "begin"); - if (sync_error.failure()) return sync_error; // intermediate is always float32 (CompType) regardless of logits dtype. auto intermediate_dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf.element_type()); NVTE_CHECK( @@ -163,8 +157,7 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( static_cast(score_function), grad_logits_tensor.data(), stream); } - return ffi_with_cuda_device_sync_and_error_check( - "FusedTopkWithScoreFunctionBackwardFFI", "end"); + return ffi_with_cuda_error_check(); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler, @@ -180,7 +173,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler, .Attr("scaling_factor") .Attr("score_function") .Attr("compute_aux_scores") - .Attr("routing_map_format")); + .Attr("routing_map_format"), + FFI_CudaGraph_Traits); // ============================================================================ // Fused MoE Aux Loss - Forward From 9776f651fe6b0aec458444ec96f43b6038a51677 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Fri, 31 Jul 2026 09:13:32 -0700 Subject: [PATCH 27/44] Reapply "Remove redundant masks from TE MoE ragged paths" This reverts commit 294a8ec37cf17d054dc24535c75dab336608155b. --- transformer_engine/jax/moe.py | 46 +++++++++-------------------------- 1 file changed, 11 insertions(+), 35 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 53726abf63..31e18a0870 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -844,11 +844,6 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): wi_for_dgrad = _weights_in_dispatch_group_order(wi) wo_for_dgrad = _weights_in_dispatch_group_order(wo) - # 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 = (group_sizes > 0)[:, None, None] - # wo bwd casted_d_eo = tex.grouped_quantize( d_eo_2d, @@ -876,7 +871,6 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) - d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) d_wo = jax.lax.with_sharding_constraint(d_wo, grouped_weight_sharding) d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None if has_bias: @@ -884,20 +878,17 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): 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 @@ -942,7 +933,6 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): 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 = jax.lax.with_sharding_constraint(d_wi_combined, grouped_weight_sharding) if has_bias: d_wi_combined_bias = tex.grouped_dbias(d_combined, group_sizes) @@ -1639,25 +1629,11 @@ def _moe_bwd_rule( 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. - active_recv_rows_3d = active_recv_rows[..., None] - safe_recv_weights = jnp.where( - active_recv_rows, - ctx.recv_topk_weights, - jnp.zeros_like(ctx.recv_topk_weights), - ) - safe_expert_outputs = jnp.where( - active_recv_rows_3d, - ctx.expert_outputs, - jnp.zeros_like(ctx.expert_outputs), - ) - d_expert_outputs = ( - grad_pre_combine - * safe_recv_weights[..., None].astype(grad_pre_combine.dtype) - ) - d_recv_w_from_combine = (grad_pre_combine * safe_expert_outputs).sum(axis=-1) + # 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) + 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) if _debug_moe_numerics: From 97ba3293750a8aeba95f5c1dc4f917e3d2602be3 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 10 Aug 2026 11:30:40 -0700 Subject: [PATCH 28/44] Optimize JAX MoE gated activation backward --- transformer_engine/jax/moe.py | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 31e18a0870..e465a21a22 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -902,12 +902,25 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): d_up_proj_out = d_intermediate * act_gp (d_gate_proj_out,) = dact_pullback(d_intermediate * up_proj_for_bwd) - # wi bwd (fused gate/up via concat). Mirror the fused fwd: pack the - # gate/up cotangents along the trailing axis, run a single - # grouped_quantize + two grouped_gemm pair (one dgrad, one wgrad) - # against the fused casted_wi_rhs_trans residual, then split the - # wgrad result remains in the contiguous gated-SwiGLU ``wi`` layout. - d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) + # wi bwd (fused gate/up). 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. Padding + # the two halves into disjoint regions and adding them is equivalent + # to concatenate, but lowers substantially faster for this large + # activation-gradient buffer on GPU. + intermediate_dim = d_gate_proj_out.shape[-1] + d_gate_padded = jax.lax.pad( + d_gate_proj_out, + jnp.array(0, dtype=d_gate_proj_out.dtype), + ((0, 0, 0), (0, intermediate_dim, 0)), + ) + d_up_padded = jax.lax.pad( + d_up_proj_out, + jnp.array(0, dtype=d_up_proj_out.dtype), + ((0, 0, 0), (intermediate_dim, 0, 0)), + ) + d_combined = d_gate_padded + d_up_padded d_combined = jax.lax.with_sharding_constraint(d_combined, flat_token_sharding) casted_d_combined = tex.grouped_quantize( d_combined, From 0170c88122a812f24af17da56cc1fda199b7b6c1 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 10 Aug 2026 12:00:51 -0700 Subject: [PATCH 29/44] Revert "Optimize JAX MoE gated activation backward" This reverts commit 97ba3293750a8aeba95f5c1dc4f917e3d2602be3. This is no longer necessary with "--xla_gpu_experimental_max_unroll_factor=8". The performance of the previous approach is fixed with the previous JAX code and this XLA flag --- transformer_engine/jax/moe.py | 25 ++++++------------------- 1 file changed, 6 insertions(+), 19 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index e465a21a22..31e18a0870 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -902,25 +902,12 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): d_up_proj_out = d_intermediate * act_gp (d_gate_proj_out,) = dact_pullback(d_intermediate * up_proj_for_bwd) - # wi bwd (fused gate/up). 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. Padding - # the two halves into disjoint regions and adding them is equivalent - # to concatenate, but lowers substantially faster for this large - # activation-gradient buffer on GPU. - intermediate_dim = d_gate_proj_out.shape[-1] - d_gate_padded = jax.lax.pad( - d_gate_proj_out, - jnp.array(0, dtype=d_gate_proj_out.dtype), - ((0, 0, 0), (0, intermediate_dim, 0)), - ) - d_up_padded = jax.lax.pad( - d_up_proj_out, - jnp.array(0, dtype=d_up_proj_out.dtype), - ((0, 0, 0), (intermediate_dim, 0, 0)), - ) - d_combined = d_gate_padded + d_up_padded + # wi bwd (fused gate/up via concat). Mirror the fused fwd: pack the + # gate/up cotangents along the trailing axis, run a single + # grouped_quantize + two grouped_gemm pair (one dgrad, one wgrad) + # against the fused casted_wi_rhs_trans residual, then split the + # wgrad result remains in the contiguous gated-SwiGLU ``wi`` layout. + d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) d_combined = jax.lax.with_sharding_constraint(d_combined, flat_token_sharding) casted_d_combined = tex.grouped_quantize( d_combined, From 4fb46473b213fa54db3c3297730bfec361a83005 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 10 Aug 2026 12:37:23 -0700 Subject: [PATCH 30/44] Reset cached NCCL EP library path --- transformer_engine/common/CMakeLists.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 9734e24fbe..9b810b6cef 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -475,6 +475,10 @@ message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") # imports stay unresolved (and harmless) under default ELF lazy binding when # the gate trips. LD_BIND_NOW environments lose this property. set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib") +# Editable builds reuse their CMake directory. Clear a library path cached by +# an older checkout so headers and the static archive always come from the +# same nccl-extensions build. +unset(NCCL_EP_LIB CACHE) find_file(NCCL_EP_LIB NAMES libnccl_ep.a HINTS ${NCCL_EP_LIB_DIR} From a59b738573e6a872b4950d71c769ed36a03f1cac Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 10 Aug 2026 13:04:16 -0700 Subject: [PATCH 31/44] Fix EP backward sharding after upstream merge --- tests/jax/test_multi_process_ep.py | 3 +++ tests/jax/test_te_ep_moe.py | 2 +- transformer_engine/jax/cpp_extensions/ep.py | 7 ++++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 04afd0bdbc..99864b014c 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -385,6 +385,7 @@ def _scan_ep_layer(self, cfg, route, tokens, weights, slot_dependent_scale): out, NamedSharding(self.mesh, token_spec) ) + @unittest.skip("Branch-only scan experiment; superseded by MaxText integration coverage.") def test_scan_dispatch_combine_top1_identity(self): """Top-1 dispatch/combine remains an exact identity across scan layers.""" routes, tokens, weights = self._make_scan_inputs() @@ -424,6 +425,7 @@ def body(current_tokens, route): if self.rank == 0: np.testing.assert_array_equal(np.asarray(out_global), np.asarray(tokens)) + @unittest.skip("Branch-only scan experiment; superseded by MaxText integration coverage.") def test_scan_dispatch_only_distinct_routing_maps(self): """Every scan iteration's packed dispatch matches an isolated dispatch. @@ -515,6 +517,7 @@ def body(carry, route): np.asarray(scan_weights_global), np.stack(ref_weights) ) + @unittest.skip("Branch-only scan experiment; superseded by MaxText integration coverage.") def test_scan_dispatch_combine_routing_handle_fwd_bwd(self): """Scan must match unrolled layers when every iteration has a new map. diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 2a86a9495c..d95f9fe9a4 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -1045,7 +1045,7 @@ def test_repeated_production_block_backward(self, mesh): def te_loss_fn(variables, value): for _ in range(repeats): - branch, _ = block.apply(variables, value) + branch, _, _ = block.apply(variables, value) value = value + branch return jnp.mean(value.astype(jnp.float32) ** 2) diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index ca70ea145c..54ec369b90 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -900,7 +900,12 @@ def partition( result_infos, ): del is_outer, result_infos - arg_shardings = tuple(a.sharding for a in arg_infos) + # The combine cotangent must have the same token sharding as the + # forward combine output. Transpose propagation can otherwise infer a + # replicated grad and pass the global token count to a handle prepared + # for only the rank-local tokens. + grad_sharding = NamedSharding(mesh, _ep_output_spec()) + arg_shardings = (arg_infos[0].sharding, grad_sharding) # EP-output leading (trailing dims auto-pad to None). out_sharding = NamedSharding(mesh, _ep_output_spec()) From 94c8495e1f80b20a4b141b048a2b86c517f49841 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 10 Aug 2026 13:59:05 -0700 Subject: [PATCH 32/44] Add configurable EP MoE receive capacity --- tests/jax/test_te_ep_moe.py | 169 ++++++++++++++++++++++++----- transformer_engine/jax/flax/moe.py | 5 + transformer_engine/jax/moe.py | 118 ++++++++++++++++---- 3 files changed, 242 insertions(+), 50 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index d95f9fe9a4..06cbce9d11 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -120,10 +120,11 @@ def _read_mp_options(): from transformer_engine.jax.flax import _MoEBlock as MoEBlock from transformer_engine.jax.moe import ( _ALIGN_SIZE, + get_moe_recv_capacity_per_rank, moe, record_ep_bootstrap_signature_for_moe, ) -from transformer_engine.jax.ep import ep_bootstrap +from transformer_engine.jax.ep import ep_bootstrap, ep_finalize from transformer_engine.jax.sharding import MeshResource, global_shard_guard @@ -208,34 +209,6 @@ def _assert_gradient_direction_and_scale(actual, expected, *, name): # ----------------------------------------------------------------------------- -def _compute_worst_case_recv_pr(): - """Per-rank recv buffer the bootstrap must reserve. - - 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 - 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") def mesh(): if jax.device_count() < NUM_DEVICES_REQUIRED: @@ -251,7 +224,13 @@ def mesh(): num_procs = jax.process_count() max_tokens_per_rank = (BATCH // num_procs) * SEQ - recv_capacity_per_rank = _compute_worst_case_recv_pr() + recv_capacity_per_rank = get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=max_tokens_per_rank, + ep_size=EP_SIZE, + recv_capacity_factor=2.0, + ) # 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 @@ -267,6 +246,7 @@ def mesh(): recv_capacity_per_rank=recv_capacity_per_rank, hidden_dim=HIDDEN, max_token_dtype=DTYPE, + drop_on_overflow=True, ) record_ep_bootstrap_signature_for_moe( num_experts=NUM_EXPERTS, @@ -408,7 +388,16 @@ def _make_block( expert_bias_init=None, compound_expert_sharding=False, input_axes=("batch", None, None), + recv_capacity_per_rank=None, ): + if recv_capacity_per_rank is None: + recv_capacity_per_rank = get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=(BATCH // jax.process_count()) * SEQ, + ep_size=EP_SIZE, + recv_capacity_factor=2.0, + ) kwargs = dict( num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -420,6 +409,7 @@ def _make_block( score_function=score_function, dtype=DTYPE, input_axes=input_axes, + recv_capacity_per_rank=recv_capacity_per_rank, ) if compound_expert_sharding: # Match MaxText shard_exp_on_fsdp=True: FSDP and EP both shard the @@ -445,6 +435,13 @@ def _strong_expert_bias_init(key, shape, dtype): ) +def _cross_rank_expert_bias_init(key, shape, dtype): + """Select one expert on each EP rank for a balanced reduced-capacity test.""" + del key + bias = jnp.full(shape, -5.0, dtype=dtype) + return bias.at[jnp.asarray((0, shape[0] // EP_SIZE))].set(5.0) + + 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( @@ -892,6 +889,45 @@ def _reference_kwargs_from_config(config, params_np): ) +class TestTeEpMoeReceiveCapacity: + """Reduced receive buffers preserve valid results and report overflow.""" + + def test_capacity_helper(self): + # Use a production-like token count where alignment does not collapse + # balanced and worst-case capacities to the same small test buffer. + max_tpr = 256 + worst = get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=max_tpr, + ep_size=EP_SIZE, + ) + balanced = get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=max_tpr, + ep_size=EP_SIZE, + recv_capacity_factor=1.0, + ) + headroom = get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=max_tpr, + ep_size=EP_SIZE, + recv_capacity_factor=2.0, + ) + assert balanced < headroom < worst + assert balanced % _ALIGN_SIZE == 0 + with pytest.raises(ValueError, match="finite and >= 1.0"): + get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=max_tpr, + ep_size=EP_SIZE, + recv_capacity_factor=0.5, + ) + + class TestTeEpMoeForward: """Per-config forward correctness in a single run: shape, dtype, finiteness AND numerical parity vs the pure-JAX reference.""" @@ -1378,3 +1414,76 @@ def test_combined_loss_grads(self, mesh): ) 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 TestZZTeEpMoeOverflow: + """Run last with a small exact capacity: valid routing then overflow.""" + + @pytest.fixture(scope="class", autouse=True) + @classmethod + def reduced_capacity_bootstrap(cls, mesh): + del cls + max_tpr = (BATCH // jax.process_count()) * SEQ + capacity = _ALIGN_SIZE + ep_finalize() + with mesh, global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ): + ep_bootstrap( + world_size=jax.process_count(), + rank=jax.process_index(), + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tpr, + recv_capacity_per_rank=capacity, + hidden_dim=HIDDEN, + max_token_dtype=DTYPE, + drop_on_overflow=True, + ) + record_ep_bootstrap_signature_for_moe( + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tpr, + recv_capacity_per_rank=capacity, + hidden_dim=HIDDEN, + ep_size=EP_SIZE, + ) + yield capacity + + @pytest.mark.parametrize( + ("expert_bias_init", "expect_overflow"), + ( + pytest.param(_cross_rank_expert_bias_init, False, id="within-capacity"), + pytest.param(_strong_expert_bias_init, True, id="overflow"), + ), + ) + def test_reduced_capacity_vjp( + self, mesh, reduced_capacity_bootstrap, expert_bias_init, expect_overflow + ): + capacity = reduced_capacity_bootstrap + + x = jax.random.normal(jax.random.PRNGKey(72), (BATCH, SEQ, HIDDEN), dtype=DTYPE) + block = _make_block( + score_function="sigmoid", + use_expert_routing_bias=True, + expert_bias_init=expert_bias_init, + recv_capacity_per_rank=capacity, + ) + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + variables = jax.jit(block.init)(jax.random.PRNGKey(73), x_sh) + + def loss_fn(variables, inputs): + output, _aux, totals = block.apply(variables, inputs) + return jnp.mean(output.astype(jnp.float32) ** 2), totals + + (loss, totals), grads = jax.jit( + jax.value_and_grad(loss_fn, has_aux=True) + )(variables, x_sh) + jax.block_until_ready((loss, totals, grads)) + + observed = int(_to_global_numpy(totals, mesh).max()) + assert (observed > capacity) is expect_overflow + assert np.isfinite(float(jax.device_get(loss))) + assert all( + np.all(np.isfinite(np.asarray(jax.device_get(leaf.addressable_data(0))))) + for leaf in jax.tree_util.tree_leaves(grads) + ) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 80d19097f8..3941a00b04 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -104,6 +104,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 @@ -149,6 +152,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 @@ -271,6 +275,7 @@ 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, + 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 897a88d03e..3a70589c1a 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -30,6 +30,7 @@ ``aux_loss_coeff`` and ``expert_bias`` are also supported. """ +import math import os import sys import warnings @@ -52,7 +53,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 @@ -62,6 +63,69 @@ # same 128-token tile, so a single constant covers every supported path. _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) + + _debug_python_patch = os.getenv("NVTE_DEBUG_PYTHON_PATCH", "0") == "1" _debug_moe_numerics = os.getenv("NVTE_DEBUG_MOE_NUMERICS", "0") == "1" _debug_moe_input_grad = os.getenv("NVTE_DEBUG_MOE_INPUT_GRAD", "0") == "1" @@ -449,9 +513,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 @@ -484,7 +548,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" @@ -499,7 +563,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" @@ -509,7 +573,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." ) @@ -1029,6 +1093,7 @@ def _moe_fwd_rule( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ): """Forward: gate -> topk -> ep_dispatch -> FFN -> ep_combine. @@ -1080,6 +1145,7 @@ def _moe_fwd_rule( return ( jnp.zeros_like(x), jnp.zeros((), dtype=x.dtype), + jnp.zeros((1,), dtype=jnp.int32), ), (ctx, static) mesh = _get_mesh() @@ -1104,18 +1170,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 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 {recv_pr}" + ) _te_ep_assert_compatible_bootstrap( num_experts=num_experts, @@ -1492,11 +1560,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 # total_recv_tokens is a non-differentiable output; its cotangent is unused. d_output, d_aux_loss, _d_total_recv_tokens = cotangents @@ -1992,7 +2061,7 @@ def _fold_dp_groups(grad): # ============================================================================= -@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, @@ -2020,6 +2089,7 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ): primal, _ = _moe_fwd_rule( x, @@ -2048,6 +2118,7 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ) return primal @@ -2086,6 +2157,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. @@ -2111,6 +2183,11 @@ def moe( quantizer_sets : Tuple[QuantizerSet, QuantizerSet] Independent FC1 and FC2 quantizer sets. 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 @@ -2201,6 +2278,7 @@ def moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ) if aux_loss_coeff <= 0.0: aux_loss = None From d89f5fc5bb306e126aaf283c70796091e142661b Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 12 Aug 2026 08:40:43 -0700 Subject: [PATCH 33/44] Revert branch changes outside JAX MoE --- qa/L1_jax_distributed_unittest/test.sh | 5 +- tests/cpp/operator/test_cast_mxfp8.cu | 238 +++++ tests/jax/test_distributed_grouped_gemm.py | 596 ------------- ..._multi_process_distributed_grouped_gemm.py | 172 ++++ tests/jax/test_multi_process_ep.py | 305 ------- tests/jax/test_te_ep_moe.py | 826 +----------------- .../test_mxfp8_quantize_swizzle_fusion.py | 33 + tests/pytorch/test_grouped_mlp.py | 6 +- transformer_engine/common/CMakeLists.txt | 4 - .../common/cast/fp8/gated_fp8.cuh | 4 +- .../common/cast/mxfp8/gated_mxfp8.cuh | 4 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 83 +- .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 590 ++++++++++--- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 4 +- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 8 +- transformer_engine/common/common.h | 20 + transformer_engine/common/ep/ep_api.cpp | 18 - transformer_engine/common/ep/ep_backend.cpp | 97 +- transformer_engine/common/ep/ep_backend.h | 7 +- transformer_engine/jax/cpp_extensions/ep.py | 7 +- transformer_engine/jax/cpp_extensions/gemm.py | 402 --------- .../jax/cpp_extensions/quantization.py | 440 +--------- transformer_engine/jax/dense.py | 56 +- transformer_engine/jax/flax/module.py | 2 +- .../jax/quantize/dequantizer.py | 9 +- transformer_engine/jax/quantize/tensor.py | 37 +- transformer_engine/jax/sharding.py | 178 +--- .../pytorch/ops/fused/grouped_mlp.py | 16 +- 28 files changed, 1153 insertions(+), 3014 deletions(-) delete mode 100644 tests/jax/test_distributed_grouped_gemm.py create mode 100644 tests/jax/test_multi_process_distributed_grouped_gemm.py diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 6d5265ac9e..c734567e1d 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -25,8 +25,6 @@ export XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_dense.xml $TE_PATH/tests/jax/test_distributed_dense.py || test_fail "test_distributed_dense.py" -python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_grouped_gemm.xml $TE_PATH/tests/jax/test_distributed_grouped_gemm.py || test_fail "test_distributed_grouped_gemm.py" - python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_helper.xml $TE_PATH/tests/jax/test_distributed_helper.py || test_fail "test_distributed_helper.py" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_layernorm.xml $TE_PATH/tests/jax/test_distributed_layernorm.py || test_fail "test_distributed_layernorm.py" @@ -42,6 +40,9 @@ XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false" python3 -m pyt # NCCL EP multi-process suite. Self-skips on <4 GPUs. TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh || test_fail "test_multi_process_ep.py" +# TODO(Phuong): add this test back after it is verified +# SCRIPT_NAME=$TE_PATH/tests/jax/test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh || test_fail "test_multi_process_distributed_grouped_gemm.py" + if [ $RET -ne 0 ]; then echo "Error: some sub-tests failed: $FAILED_CASES" exit 1 diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index c7c778ce1e..edeb87ebe1 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -7,10 +7,12 @@ #include #include #include +#include #include #include #include +#include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -735,3 +737,239 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::ValuesIn(input_scenarios)), test_name_generator); + +// ============================================================================ +// Swizzled-scales cast-only tests +// +// Validate the WITH_GEMM_SWIZZLED_SCALES=true code path added by the +// CastTraitsSwizzle port. The specialized kernel dispatches to +// CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> whenever the +// output tensor has set_with_gemm_swizzled_scales(true), producing scales +// directly in GEMM-swizzled layout. +// +// Reference construction: run the well-tested linear-scale path, then apply +// nvte_swizzle_scaling_factors (independently tested by SwizzleTestSuite) to +// transform the linear scales into the swizzled layout. Byte-compare the +// direct swizzled cast output against this reference for both scale tensors +// and the FP8 output data itself. +// +// Pass criteria per test: +// 1. FP8 rowwise data byte-identical between linear and swizzled paths. +// 2. FP8 colwise data byte-identical between linear and swizzled paths. +// 3. Swizzled rowwise scale bytes match linear->swizzle reference. +// 4. Swizzled colwise scale bytes match linear->swizzle reference. +// 5. Rowwise-only FP8 data and swizzled scales match the same reference. +// 6. No CUDA errors from any launch. +// ============================================================================ + +class SwizzledScalesFusedCastMXFP8TestSuite : public ::testing::TestWithParam< + std::tuple, + transformer_engine::DType, + transformer_engine::DType>> {}; + +TEST_P(SwizzledScalesFusedCastMXFP8TestSuite, TestSwizzledCastMXFP8) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const auto shape = std::get<0>(GetParam()); + const DType itype = std::get<1>(GetParam()); + const DType otype = std::get<2>(GetParam()); + + // BIDIMENSIONAL scaling needs a 2D+ input. + if (shape.size() < 2) { + GTEST_SKIP(); + } + + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + const size_t out_bytes = rows * cols; // fp8 = 1 byte/elem + + // Input filled with the same values for all casts. + Tensor input("input", shape, itype); + fillUniform(&input); + + // Target: swizzled-scale cast. Dispatcher routes to + // CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> on kernel #3. + Tensor output_swizzled("output_swizzled", shape, otype, true, true, NVTE_MXFP8_1D_SCALING); + output_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_swizzled.data(), 0); + + // Rowwise-only activations take the pointer-based cast-only kernel. This + // directly exercises its WITH_GEMM_SWIZZLED_SCALES trait specialization + // for shapes whose column count is a multiple of 128. + Tensor output_rowwise_swizzled("output_rowwise_swizzled", shape, otype, + /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + output_rowwise_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_rowwise_swizzled.data(), 0); + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "swizzled-scale nvte_quantize failed"; + + // Reference construction: nvte_swizzle_scaling_factors accepts tensors with + // exactly one scale direction, so we build the rowwise and colwise references + // independently. Each is a single-direction linear cast followed by a + // single-direction swizzle transform, using the same input as the target. + // MXFP8 is a deterministic per-direction operation, so a rowwise-only cast + // produces the same rowwise fp8 bytes and scales as a rowwise+colwise cast. + + // Rowwise reference. + Tensor linear_row("linear_row", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_row.data(), 0); + Tensor ref_row_swz("ref_row_swz", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + ref_row_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_row_swz.rowwise_dptr(), linear_row.rowwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_row.data(), ref_row_swz.data(), 0); + } + + // Colwise reference. + Tensor linear_col("linear_col", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_col.data(), 0); + Tensor ref_col_swz("ref_col_swz", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + ref_col_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_col_swz.columnwise_dptr(), linear_col.columnwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_col.data(), ref_col_swz.data(), 0); + } + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "reference construction failed"; + + // ---- Comparisons ---- + + // (1) & (2): FP8 output data byte-identical to single-direction linear casts. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test_row = reinterpret_cast( + output_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref_row = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_row[i], ref_row[i]) + << "rowwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_row[i]) + << " linear=" << static_cast(ref_row[i]) << ")"; + } + const uint8_t *test_col = reinterpret_cast( + output_swizzled.columnwise_cpu_dptr()); + const uint8_t *ref_col = reinterpret_cast( + linear_col.columnwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_col[i], ref_col[i]) + << "colwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_col[i]) + << " linear=" << static_cast(ref_col[i]) << ")"; + } + } + ); + + // (3): Swizzled rowwise scale bytes — directly-written vs linear-then-transform. + { + const size_t n = product(output_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled rowwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (4): Swizzled colwise scale bytes — exercises the uint16-pair-packed flush + // path added by CACHE_COLWISE_SCALE_IN_SMEM + WITH_SWIZZLED_SCALES. + { + const size_t n = product(output_swizzled.columnwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.columnwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_col_swz.columnwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled colwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (5) & (6): The rowwise-only swizzled path matches the same independently + // constructed linear-then-swizzle reference. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test = reinterpret_cast( + output_rowwise_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only fp8 data mismatch at index " << i; + } + } + ); + { + const size_t n = product(output_rowwise_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = + output_rowwise_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only swizzled scale mismatch at offset " << i; + } + } + +} + +std::string swizzled_test_name_generator( + const testing::TestParamInfo& info) { + std::string name; + const auto &shape = std::get<0>(info.param); + for (size_t i = 0; i < shape.size(); ++i) { + if (i > 0) name += "x"; + name += std::to_string(shape[i]); + } + name += "X" + test::typeName(std::get<1>(info.param)) + + "X" + test::typeName(std::get<2>(info.param)); + return name; +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest_FusedCastMXFP8_SwizzledCastOnly, + SwizzledScalesFusedCastMXFP8TestSuite, + ::testing::Values( + // 1. Aligned, small — sanity. + std::make_tuple(std::vector{128, 128}, DType::kBFloat16, DType::kFloat8E4M3), + // 2. Aligned, small — second dtype pair. + std::make_tuple(std::vector{128, 128}, DType::kFloat32, DType::kFloat8E5M2), + // 3. Aligned, medium. + std::make_tuple(std::vector{256, 384}, DType::kBFloat16, DType::kFloat8E4M3), + // 4. Odd number of scale rows (96/32 = 3) - exercises colwise flush + // odd-row scalar tail. + std::make_tuple(std::vector{96, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 5. cols/32 not a multiple of 4 - exercises rowwise flush scalar tail + // (need cols multiple of 8 for TMA-alignment on 16-bit input). + std::make_tuple(std::vector{256, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 6. Odd colwise scale rows + rowwise tail cols - combined tail paths. + std::make_tuple(std::vector{96, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 7. Small but tile-aligned. + std::make_tuple(std::vector{32, 64}, DType::kBFloat16, DType::kFloat8E4M3), + // 8. Minimum aligned size - single 32x32 tile. + std::make_tuple(std::vector{32, 32}, DType::kBFloat16, DType::kFloat8E4M3), + // 9. Multi-CTA at scale — stresses gmem writes across many CTAs. + std::make_tuple(std::vector{4096, 32768}, DType::kBFloat16, DType::kFloat8E4M3), + // 10. 4D input — matches the existing suite's rank coverage. + std::make_tuple(std::vector{16, 8, 4, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 11. Third input dtype. + std::make_tuple(std::vector{128, 128}, DType::kFloat16, DType::kFloat8E4M3), + // 12. Larger fp32. + std::make_tuple(std::vector{1024, 1024}, DType::kFloat32, DType::kFloat8E4M3) + ), + swizzled_test_name_generator); diff --git a/tests/jax/test_distributed_grouped_gemm.py b/tests/jax/test_distributed_grouped_gemm.py deleted file mode 100644 index a9e9d159da..0000000000 --- a/tests/jax/test_distributed_grouped_gemm.py +++ /dev/null @@ -1,596 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. -"""Partitioning tests for grouped quantize and grouped GEMM.""" - -import math -from types import SimpleNamespace - -import jax -import jax.numpy as jnp -import numpy as np -import pytest -from jax.sharding import Mesh, NamedSharding, PartitionSpec - -from transformer_engine.jax.cpp_extensions.gemm import GroupedGemmPrimitive -from transformer_engine.jax.cpp_extensions.quantization import GroupedQuantizePrimitive -from transformer_engine.jax.dense import grouped_dense -from transformer_engine.jax.quantize import QuantizeLayout, QuantizerFactory, ScalingMode -from transformer_engine.jax.sharding import MeshResource, global_shard_guard - - -def _mesh(): - devices = jax.devices() - if len(devices) < 4: - pytest.skip("Grouped GEMM partitioning tests require at least 4 visible GPUs.") - return Mesh(np.asarray(devices[:4]).reshape(2, 2), ("expert", "fsdp")) - - -def _mesh_with_dp_tp(): - devices = jax.devices() - if len(devices) < 4: - pytest.skip("Grouped GEMM partitioning tests require at least 4 visible GPUs.") - return Mesh(np.asarray(devices[:4]).reshape(2, 1, 2, 1), ("expert", "dp", "fsdp", "tp")) - - -def _mesh_with_arbitrary_axis(): - devices = jax.devices() - if len(devices) < 4: - pytest.skip("Grouped GEMM partitioning tests require at least 4 visible GPUs.") - return Mesh( - np.asarray(devices[:4]).reshape(2, 1, 2, 1), - ("expert", "dp", "fsdp", "myaxis123"), - ) - - -def _arg_info(mesh, shape, spec): - return SimpleNamespace( - shape=shape, - ndim=len(shape), - size=int(np.prod(shape)), - sharding=NamedSharding(mesh, PartitionSpec(*spec)), - ) - - -def _normalize_spec(spec): - if isinstance(spec, PartitionSpec): - return tuple(spec) - return spec - - -def _spec_contains_axis(spec, axis): - for axis_spec in spec: - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - if axis in axis_tuple: - return True - return False - - -def _mxfp8_grouped_quantizer_set(n_groups): - return QuantizerFactory.create_set( - scaling_mode=ScalingMode.MXFP8_1D_SCALING, - fwd_dtype=jnp.float8_e4m3fn, - bwd_dtype=jnp.float8_e4m3fn, - is_2x2x=True, - n_groups=n_groups, - ) - - -def test_grouped_quantize_large_abstract_shape_preserves_inferred_hidden_dim(): - input_shape = (786432, 4096) - group_count = 64 - outputs = GroupedQuantizePrimitive.abstract( - jax.core.ShapedArray(input_shape, jnp.bfloat16), - jax.core.ShapedArray((group_count,), jnp.float32), - jax.core.ShapedArray((group_count,), jnp.int32), - out_dtype=jnp.float8_e4m3fn, - scaling_mode=ScalingMode.MXFP8_1D_SCALING.value, - q_layout=QuantizeLayout.ROWWISE, - flatten_axis=-1, - scale_dtype=jnp.float8_e8m0fnu, - uniform_groups=False, - ) - - assert outputs[0].shape == input_shape - assert max(outputs[0].shape) < 2**31 - assert math.prod(outputs[0].shape) == 3221225472 - - -def test_grouped_quantize_preserves_output_side_fsdp_for_uniform_kernel(): - mesh = _mesh() - with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - True, - mesh, - ( - _arg_info(mesh, (8, 128, 256), ("expert", None, "fsdp")), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (8,), ("expert",)), - ), - (), - ) - - assert tuple(arg_shardings[0].spec) == ("expert", None, "fsdp") - specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == ("expert", None, "fsdp") - assert _normalize_spec(specs[2]) == ("expert", None, "fsdp") - assert _normalize_spec(specs[4]) == ("expert",) - - -def test_grouped_quantize_mxfp8_colwise_scale_tracks_output_side_fsdp(): - mesh = _mesh() - with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE_COLWISE, - -1, - jnp.float8_e8m0fnu, - True, - mesh, - ( - _arg_info(mesh, (8, 128, 256), ("expert", None, "fsdp")), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (8,), ("expert",)), - ), - (), - ) - - assert tuple(arg_shardings[0].spec) == ("expert", None, "fsdp") - specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == ("expert", None, "fsdp") - assert _normalize_spec(specs[1]) == ("expert", None, "fsdp") - assert _normalize_spec(specs[2]) == ("expert", None, "fsdp") - assert _normalize_spec(specs[3]) == ("expert", "fsdp", None) - assert _normalize_spec(specs[4]) == ("expert",) - - -def test_grouped_quantize_preserves_row_side_fsdp_for_kernel(): - mesh = _mesh() - with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - True, - mesh, - ( - _arg_info(mesh, (8, 256, 128), ("expert", "fsdp", None)), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (8,), ("expert",)), - ), - (), - ) - - assert tuple(arg_shardings[0].spec) == ("expert", "fsdp", None) - specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(specs[0]) == ("expert", "fsdp", None) - assert _normalize_spec(specs[2]) == ("expert", "fsdp", None) - - -def test_grouped_quantize_strips_unsupported_axes_and_preserves_supported_axes(): - mesh = _mesh_with_dp_tp() - with jax.set_mesh(mesh), global_shard_guard( - MeshResource(dp_resource="dp", tp_resource="tp", fsdp_resource="fsdp", ep_resource="expert") - ): - with pytest.warns(RuntimeWarning, match="Grouped quantize.*tp"): - _, _, out_shardings, arg_shardings = GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - True, - mesh, - ( - _arg_info(mesh, (8, 128, 256), ("expert", "dp", ("fsdp", "tp"))), - _arg_info(mesh, (8,), (("expert", "tp"),)), - _arg_info(mesh, (8,), (("expert", "tp"),)), - ), - (), - ) - - assert tuple(arg_shardings[0].spec) == ("expert", "dp", "fsdp") - assert tuple(arg_shardings[1].spec) == ("expert",) - assert tuple(arg_shardings[2].spec) == ("expert",) - - out_specs = tuple(tuple(sharding.spec) for sharding in out_shardings) - assert _normalize_spec(out_specs[0]) == ("expert", "dp", "fsdp") - assert _normalize_spec(out_specs[2]) == ("expert", "dp", "fsdp") - assert _normalize_spec(out_specs[4]) == ("expert",) - for spec in (*out_specs, *(tuple(sharding.spec) for sharding in arg_shardings)): - assert not _spec_contains_axis(spec, "tp") - - -def test_grouped_gemm_rhs_weight_specs_gather_fsdp_but_preserve_ep(): - mesh = _mesh() - arg_infos = ( - _arg_info(mesh, (8192,), (None,)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (65536,), (("expert", "fsdp"),)), - _arg_info(mesh, (2048,), (("expert", "fsdp"),)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (1,), (None,)), - _arg_info(mesh, (0,), (None,)), - ) - with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - _, _, out_sharding, arg_shardings = GroupedGemmPrimitive.partition( - False, - False, - ScalingMode.NO_SCALING.value, - jnp.bfloat16, - False, - False, - False, - 1, - 1, - (1, 128, 64), - 128, - 64, - 128, - 64, - mesh, - arg_infos, - (), - ) - - assert tuple(arg_shardings[2].spec) == ("expert",) - assert tuple(arg_shardings[3].spec) == ("expert",) - assert tuple(out_sharding[0].spec) == (None, None, None) - - -def test_grouped_gemm_gathers_fsdp_shared_moe_rhs_with_exact_group_ratio(): - """A global MoE RHS has E groups while token counts have fsdp * E groups.""" - mesh = _mesh() - arg_infos = ( - _arg_info(mesh, (8192,), (None,)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (32, 128, 64), (("expert", "fsdp"), None, None)), - _arg_info(mesh, (2048,), (("expert", "fsdp"),)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (64,), (("fsdp", "expert"),)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (64,), (("fsdp", "expert"),)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (1,), (None,)), - _arg_info(mesh, (0,), (None,)), - ) - with global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - _, _, _, arg_shardings = GroupedGemmPrimitive.partition( - False, - False, - ScalingMode.NO_SCALING.value, - jnp.bfloat16, - False, - False, - False, - 1, - 1, - (64, 128, 64), - 128, - 64, - 128, - 64, - mesh, - arg_infos, - (), - ) - - assert tuple(arg_shardings[2].spec) == ("expert", None, None) - assert tuple(arg_shardings[3].spec) == ("expert",) - - -def test_grouped_gemm_strips_unsupported_axes_preserves_dp_and_gathers_rhs_fsdp(): - mesh = _mesh_with_dp_tp() - arg_infos = ( - _arg_info(mesh, (8192,), (("dp", "tp"),)), - _arg_info(mesh, (0,), (("tp",),)), - _arg_info(mesh, (65536,), (("expert", "fsdp", "tp"),)), - _arg_info(mesh, (2048,), (("expert", "fsdp", "tp"),)), - _arg_info(mesh, (0,), (("fsdp", "tp"),)), - _arg_info(mesh, (8,), (("expert", "tp"),)), - _arg_info(mesh, (0,), (("tp",),)), - _arg_info(mesh, (0,), (("tp",),)), - _arg_info(mesh, (0,), (("tp",),)), - _arg_info(mesh, (8,), (("expert", "tp"),)), - _arg_info(mesh, (0,), (("tp",),)), - _arg_info(mesh, (1,), (("tp",),)), - _arg_info(mesh, (0,), (("tp",),)), - ) - result_infos = (_arg_info(mesh, (1, 128, 64), ("expert", "tp", None)),) - with jax.set_mesh(mesh), global_shard_guard( - MeshResource(dp_resource="dp", tp_resource="tp", fsdp_resource="fsdp", ep_resource="expert") - ): - with pytest.warns(RuntimeWarning, match="Grouped GEMM.*tp"): - _, _, out_sharding, arg_shardings = GroupedGemmPrimitive.partition( - False, - False, - ScalingMode.NO_SCALING.value, - jnp.bfloat16, - False, - False, - False, - 1, - 1, - (1, 128, 64), - 128, - 64, - 128, - 64, - mesh, - arg_infos, - result_infos, - ) - - assert tuple(arg_shardings[0].spec) == ("dp",) - assert tuple(arg_shardings[2].spec) == ("expert",) - assert tuple(arg_shardings[3].spec) == ("expert",) - assert tuple(arg_shardings[5].spec) == ("expert",) - assert tuple(out_sharding[0].spec) == ("expert", None, None) - for spec in ( - *(tuple(sharding.spec) for sharding in arg_shardings), - tuple(out_sharding[0].spec), - ): - assert not _spec_contains_axis(spec, "tp") - - -def test_grouped_gemm_reduce_axis_skips_ep_and_uses_dp(): - mesh = _mesh_with_dp_tp() - arg_infos = ( - _arg_info(mesh, (8192,), (("expert", "dp"),)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (8192,), (("expert", "dp"),)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (8,), ("expert",)), - _arg_info(mesh, (0,), (None,)), - _arg_info(mesh, (1,), (None,)), - _arg_info(mesh, (0,), (None,)), - ) - - with jax.set_mesh(mesh), global_shard_guard( - MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="expert") - ): - _, _, reduce_axis = GroupedGemmPrimitive._parse_partition_specs( - mesh, - arg_infos, - (), - out_shape=(1, 128, 64), - lhs_is_trans=False, - lhs_axis_boundary=1, - ) - - assert reduce_axis == "dp" - - -def test_grouped_partitioning_strips_arbitrary_unsupported_axis(): - mesh = _mesh_with_arbitrary_axis() - mesh_resource = MeshResource(dp_resource="dp", fsdp_resource="fsdp", ep_resource="expert") - - with jax.set_mesh(mesh), global_shard_guard(mesh_resource): - with pytest.warns(RuntimeWarning, match="Grouped quantize.*myaxis123"): - _, _, quantize_out_shardings, quantize_arg_shardings = ( - GroupedQuantizePrimitive.partition( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - True, - mesh, - ( - _arg_info(mesh, (8, 128, 256), ("expert", "myaxis123", ("dp", "fsdp"))), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - ), - (), - ) - ) - - gemm_arg_infos = ( - _arg_info(mesh, (8192,), (("dp", "myaxis123"),)), - _arg_info(mesh, (0,), (("myaxis123",),)), - _arg_info(mesh, (65536,), (("expert", "fsdp", "myaxis123"),)), - _arg_info(mesh, (2048,), (("expert", "fsdp", "myaxis123"),)), - _arg_info(mesh, (0,), (("fsdp", "myaxis123"),)), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - _arg_info(mesh, (0,), (("myaxis123",),)), - _arg_info(mesh, (0,), (("myaxis123",),)), - _arg_info(mesh, (0,), (("myaxis123",),)), - _arg_info(mesh, (8,), (("expert", "myaxis123"),)), - _arg_info(mesh, (0,), (("myaxis123",),)), - _arg_info(mesh, (1,), (("myaxis123",),)), - _arg_info(mesh, (0,), (("myaxis123",),)), - ) - gemm_result_infos = (_arg_info(mesh, (1, 128, 64), ("expert", "myaxis123", None)),) - with pytest.warns(RuntimeWarning, match="Grouped GEMM.*myaxis123"): - _, _, gemm_out_sharding, gemm_arg_shardings = GroupedGemmPrimitive.partition( - False, - False, - ScalingMode.NO_SCALING.value, - jnp.bfloat16, - False, - False, - False, - 1, - 1, - (1, 128, 64), - 128, - 64, - 128, - 64, - mesh, - gemm_arg_infos, - gemm_result_infos, - ) - - assert tuple(quantize_arg_shardings[0].spec) == ("expert", None, ("dp", "fsdp")) - assert tuple(quantize_arg_shardings[1].spec) == ("expert",) - quantize_out_specs = tuple(tuple(sharding.spec) for sharding in quantize_out_shardings) - assert _normalize_spec(quantize_out_specs[0]) == ("expert", None, ("dp", "fsdp")) - assert _normalize_spec(quantize_out_specs[2]) == ("expert", None, ("dp", "fsdp")) - - assert tuple(gemm_arg_shardings[0].spec) == ("dp",) - assert tuple(gemm_arg_shardings[2].spec) == ("expert",) - assert tuple(gemm_arg_shardings[3].spec) == ("expert",) - assert tuple(gemm_out_sharding[0].spec) == ("expert", None, None) - - all_specs = ( - *quantize_out_specs, - *(tuple(sharding.spec) for sharding in quantize_arg_shardings), - *(tuple(sharding.spec) for sharding in gemm_arg_shardings), - tuple(gemm_out_sharding[0].spec), - ) - for spec in all_specs: - assert not _spec_contains_axis(spec, "myaxis123") - - -def test_grouped_partitioning_shardy_rules_smoke(): - mesh = _mesh() - quantize_rule = GroupedQuantizePrimitive.shardy_sharding_rule( - jnp.float8_e4m3fn, - ScalingMode.MXFP8_1D_SCALING.value, - QuantizeLayout.ROWWISE, - -1, - jnp.float8_e8m0fnu, - True, - mesh, - ( - SimpleNamespace(shape=(8, 128, 128)), - SimpleNamespace(shape=(8,)), - SimpleNamespace(shape=(8,)), - ), - ( - SimpleNamespace(shape=(8, 128, 128)), - SimpleNamespace(shape=(1,)), - SimpleNamespace(shape=(8, 1, 512)), - SimpleNamespace(shape=(1,)), - SimpleNamespace(shape=(8,)), - ), - ) - gemm_rule = GroupedGemmPrimitive.shardy_sharding_rule( - False, - False, - ScalingMode.NO_SCALING.value, - jnp.bfloat16, - False, - False, - False, - 1, - 2, - (128, 64), - 128, - 64, - 128, - 64, - mesh, - tuple(SimpleNamespace(shape=(1,)) for _ in range(13)), - (SimpleNamespace(shape=(128, 64)),), - ) - - assert quantize_rule is not None - assert gemm_rule is not None - - -@pytest.mark.parametrize( - ("group_spec", "weight_spec"), - [ - ("expert", ("expert", "fsdp", None)), - ("expert", ("expert", None, "fsdp")), - (("fsdp", "expert"), (("fsdp", "expert"), None, None)), - ], - ids=("contracting-fsdp", "output-fsdp", "compound-fsdp-expert-groups"), -) -def test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process( - group_spec, weight_spec -): - mesh = _mesh() - n_groups = 4 - group_tokens = 128 - hidden = 256 - out_hidden = 256 - x_shape = (n_groups * group_tokens, hidden) - w_shape = (n_groups, hidden, out_hidden) - - x_sharding = NamedSharding(mesh, PartitionSpec(group_spec, None)) - w_sharding = NamedSharding(mesh, PartitionSpec(*weight_spec)) - group_sharding = NamedSharding(mesh, PartitionSpec(group_spec)) - out_sharding = NamedSharding(mesh, PartitionSpec(group_spec, None)) - - quantizer_set = _mxfp8_grouped_quantizer_set(n_groups) - - with mesh, global_shard_guard(MeshResource(fsdp_resource="fsdp", ep_resource="expert")): - x = jax.device_put( - jax.random.normal(jax.random.PRNGKey(20), x_shape, dtype=jnp.bfloat16) - * jnp.asarray(0.01, dtype=jnp.bfloat16), - x_sharding, - ) - w = jax.device_put( - jax.random.normal(jax.random.PRNGKey(21), w_shape, dtype=jnp.bfloat16) - * jnp.asarray(0.01, dtype=jnp.bfloat16), - w_sharding, - ) - group_sizes = jax.device_put( - jnp.full((n_groups,), group_tokens, dtype=jnp.int32), - group_sharding, - ) - - def apply_with_vjp(x, w, group_sizes): - def apply(x, w): - return grouped_dense( - x, - w, - group_sizes, - contracting_dims=((1,), (1,)), - quantizer_set=quantizer_set, - ) - - out, vjp_fn = jax.vjp(apply, x, w) - dx, dw = vjp_fn(out) - return out, dx, dw - - out, dx, dw = jax.jit( - apply_with_vjp, - in_shardings=(x_sharding, w_sharding, group_sharding), - out_shardings=(out_sharding, x_sharding, w_sharding), - )(x, w, group_sizes) - out, dx, dw = jax.block_until_ready((out, dx, dw)) - - assert tuple(out.sharding.spec) == (group_spec, None) - assert tuple(dx.sharding.spec) == (group_spec, None) - assert tuple(dw.sharding.spec) == weight_spec - for value in (out, dx, dw): - local_value = np.asarray(jax.device_get(value.addressable_data(0))) - assert np.all(np.isfinite(local_value)) - assert np.any(local_value != 0.0) - - x_global = np.asarray(jax.device_get(x)).reshape(n_groups, group_tokens, hidden) - w_global = np.asarray(jax.device_get(w)) - reference = np.einsum( - "gth,gho->gto", x_global.astype(np.float32), w_global.astype(np.float32) - ).reshape(x_shape[0], out_hidden) - np.testing.assert_allclose( - np.asarray(jax.device_get(out)).astype(np.float32), - reference.astype(np.float32), - atol=5e-3, - rtol=5e-2, - ) diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py new file mode 100644 index 0000000000..94fed0859f --- /dev/null +++ b/tests/jax/test_multi_process_distributed_grouped_gemm.py @@ -0,0 +1,172 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +from functools import partial + +import jax +import jax.numpy as jnp +import jax.experimental.multihost_utils as jem + +from transformer_engine.jax.dense import grouped_dense as te_grouped_dense +from transformer_engine.jax.quantize import ( + QuantizerFactory, + ScalingMode, +) + +from utils import assert_allclose, dtype_tols + + +N_GROUP = 8 +MESH_AXIS_NAME = "fsdp" + + +def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis): + assert kernel_fsdp_axis in [1, 2] + x_shape, w_shape = data_shapes + + x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None)) + w_sharding = ( + NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME)) + if kernel_fsdp_axis == 2 + else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None)) + ) + w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None)) + + def init_data(): + x_key = jax.random.PRNGKey(0) + w_key = jax.random.PRNGKey(1) + x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16) + w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16) + w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim)) + return x, w, w, w_amax + + def test_func(outter_x, outter_w, outter_w_amax): + in_specs = (x_sharding.spec, w_sharding.spec, None) + out_specs = x_sharding.spec + + @partial( + shard_map.shard_map, + mesh=mesh, + in_specs=in_specs, + out_specs=out_specs, + check_rep=False, + ) + def sharded_group_gemm(x, w, w_amax): + group_size = x.shape[0] + x_reshaped = x.reshape(-1, x.shape[-1]) + n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) + + quantizer_set = QuantizerFactory.create_set( + scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING, + fwd_dtype=jnp.float8_e4m3fn, + bwd_dtype=jnp.float8_e5m2, + is_2x2x=True, + n_groups=group_size, + ) + + output = te_grouped_dense( + x_reshaped, + w, + n_groups, + kernel_amax=w_amax, + quantizer_set=quantizer_set, + kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis), + ) + output = output.reshape(*x.shape[:-1], -1) + return output + + def run(x, w, w_amax): + output = sharded_group_gemm(x, w, w_amax) + return output + + output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax) + dx, dw, _ = vjp_fn(output) + return output, dx, dw + + def ref_func(outter_x, outter_w): + + in_specs = (x_sharding.spec, w_no_sharding.spec) + out_specs = x_sharding.spec + + @partial( + shard_map.shard_map, + mesh=mesh, + in_specs=in_specs, + out_specs=out_specs, + check_rep=False, + ) + def sharded_group_gemm(x, w): + group_size = x.shape[0] + x_reshaped = x.reshape(-1, x.shape[-1]) + n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size) + + quantizer_set = QuantizerFactory.create_set( + scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING, + fwd_dtype=jnp.float8_e4m3fn, + bwd_dtype=jnp.float8_e5m2, + is_2x2x=True, + n_groups=group_size, + ) + output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set) + output = output.reshape(*x.shape[:-1], -1) + return output + + def run(x, w): + output = sharded_group_gemm(x, w) + return output + + output, vjp_fn = jax.vjp(run, outter_x, outter_w) + dx, dw = vjp_fn(output) + return output, dx, dw + + init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None)) + x, w, w_global, w_amax = init_func() + + o_sharding = x_sharding + test_func_jitted = jax.jit( + test_func, + in_shardings=(x_sharding, w_sharding, None), + out_shardings=(o_sharding, x_sharding, w_sharding), + ) + ref_func_jitted = jax.jit( + ref_func, + in_shardings=(x_sharding, w_no_sharding), + out_shardings=(o_sharding, x_sharding, w_no_sharding), + ) + + out, dx, dw = test_func_jitted(x, w, w_amax) + ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global) + + e4m3_tols = dtype_tols(jnp.float8_e4m3fn) + e5m2_tols = dtype_tols(jnp.float8_e5m2) + + out, ref_out = jem.process_allgather((out, ref_out)) + dx, ref_dx = jem.process_allgather((dx, ref_dx)) + dw, ref_dw = jem.process_allgather((dw, ref_dw)) + + jnp.allclose(out, ref_out, **e4m3_tols) + jnp.allclose(dx, ref_dx, **e5m2_tols) + jnp.allclose(dw, ref_dw, **e5m2_tols) + + +if __name__ == "__main__": + from jax.sharding import NamedSharding, PartitionSpec + from jax.experimental import shard_map + import sys + + coord_addr = sys.argv[1] + proc_id = int(sys.argv[2]) + num_procs = int(sys.argv[3]) + + jax.distributed.initialize( + coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id + ) + + mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,)) + + with mesh: + data_shapes = [((4, 16, 128, 7168), (7168, 2048))] + for data_shape in data_shapes: + for kernel_fsdp_axis in [1, 2]: + test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 99864b014c..47af0b0c39 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -13,8 +13,6 @@ - ``ep_combine`` custom_vjp: ``max|grad_eo| ≈ eo_const / TOP_K`` (closed form). - ``ep_dispatch`` custom_vjp: exact per-(t, k) ``grad_topk_weights`` under skewed upstream gradients (no k-axis averaging). - - ``jax.lax.scan`` over distinct top-1 routing maps matches an unrolled - dispatch/combine reference in both forward and backward. - HLO reshard guard: compile-only, no XLA collectives outside the EP FFI. - Drop-on-overflow: ``ep_finalize`` + re-bootstrap with a small recv capacity drops the overflow instead of trapping; ``total_recv_tokens`` reports the @@ -312,309 +310,6 @@ def run(idx, ta_, tb_, w_): rtol=5e-2, ) - def _make_scan_inputs(self, num_layers=4): - """Top-1 inputs with a distinct random routing map for every layer.""" - num_layers = int(os.environ.get("NVTE_TEST_EP_SCAN_LAYERS", num_layers)) - num_tokens = TOKENS_PER_DP_SHARD * self.dp - rng = np.random.default_rng(seed=20260728) - routes = rng.integers( - 0, - self.num_experts, - size=(num_layers, num_tokens, 1), - dtype=np.int32, - ) - # A random draw could theoretically repeat a complete map. Make the - # invariant explicit so this test never becomes probabilistic. - seen = set() - for layer in range(num_layers): - while routes[layer].tobytes() in seen: - routes[layer] = rng.integers( - 0, self.num_experts, size=(num_tokens, 1), dtype=np.int32 - ) - seen.add(routes[layer].tobytes()) - tokens = jnp.asarray( - rng.standard_normal((num_tokens, HIDDEN_DIM), dtype=np.float32) * 0.25, - dtype=jnp.bfloat16, - ) - weights = jnp.ones((num_tokens, 1), dtype=jnp.float32) - return jnp.asarray(routes), tokens, weights - - def _scan_ep_layer(self, cfg, route, tokens, weights, slot_dependent_scale): - """One dispatch/combine-only layer used by the scan regression tests.""" - token_spec = PartitionSpec(("dp", "ep"), None) - ep_token_spec = PartitionSpec(("dp", "ep"), None, None) - ep_weight_spec = PartitionSpec(("dp", "ep"), None) - - recv_tokens, recv_weights, handle_mem, token_counts = ep_dispatch( - cfg, route, tokens, weights, self.recv_capacity_per_rank - ) - recv_tokens = jax.lax.with_sharding_constraint( - recv_tokens, NamedSharding(self.mesh, ep_token_spec) - ) - recv_weights = jax.lax.with_sharding_constraint( - recv_weights, NamedSharding(self.mesh, ep_weight_spec) - ) - - # A plain top-1 identity round-trip cannot expose a stale handle if - # dispatch and combine both use the same stale routing map: the two - # wrong permutations cancel. Scaling by packed dispatch slot makes the - # result (and its token gradient) depend on the map that was prepared. - if slot_dependent_scale: - slot = jnp.arange(recv_tokens.shape[-2], dtype=jnp.float32) - scale = 0.5 + (slot % 7) * 0.125 - expert_out = recv_tokens.astype(jnp.float32) * scale[None, :, None] - else: - expert_out = recv_tokens.astype(jnp.float32) - expert_out = jnp.where( - recv_weights[..., None] != 0, - expert_out * recv_weights[..., None], - 0.0, - ).astype(recv_tokens.dtype) - expert_out = jax.lax.with_sharding_constraint( - expert_out, NamedSharding(self.mesh, ep_token_spec) - ) - out = ep_combine( - cfg, - handle_mem, - token_counts, - expert_out, - tokens.shape[0], - out_sharding=(("dp", "ep"), None), - ) - return jax.lax.with_sharding_constraint( - out, NamedSharding(self.mesh, token_spec) - ) - - @unittest.skip("Branch-only scan experiment; superseded by MaxText integration coverage.") - def test_scan_dispatch_combine_top1_identity(self): - """Top-1 dispatch/combine remains an exact identity across scan layers.""" - routes, tokens, weights = self._make_scan_inputs() - cfg = EpLayerConfig(top_k=1, dispatch_output_per_expert_alignment=16) - route_spec = PartitionSpec(None, ("dp", "ep"), None) - token_spec = PartitionSpec(("dp", "ep"), None) - - with self.mesh, global_shard_guard(self.mr): - routes = jax.lax.with_sharding_constraint( - routes, NamedSharding(self.mesh, route_spec) - ) - tokens_s = jax.lax.with_sharding_constraint( - tokens, NamedSharding(self.mesh, token_spec) - ) - weights_s = jax.lax.with_sharding_constraint( - weights, NamedSharding(self.mesh, token_spec) - ) - - @jax.jit - def run(route_maps, initial_tokens, topk_weights): - def body(current_tokens, route): - out = self._scan_ep_layer( - cfg, - route, - current_tokens, - topk_weights, - slot_dependent_scale=False, - ) - return out, None - - return jax.lax.scan(body, initial_tokens, route_maps)[0] - - out = run(routes, tokens_s, weights_s) - out.block_until_ready() - out_global = jmu.process_allgather(out, tiled=True) - - if self.rank == 0: - np.testing.assert_array_equal(np.asarray(out_global), np.asarray(tokens)) - - @unittest.skip("Branch-only scan experiment; superseded by MaxText integration coverage.") - def test_scan_dispatch_only_distinct_routing_maps(self): - """Every scan iteration's packed dispatch matches an isolated dispatch. - - This is the direct stale-routing oracle: there is no combine operation - whose inverse permutation could hide a reused routing-map handle. - """ - routes, tokens, weights = self._make_scan_inputs() - cfg = EpLayerConfig(top_k=1, dispatch_output_per_expert_alignment=16) - route_spec = PartitionSpec(None, ("dp", "ep"), None) - token_spec = PartitionSpec(("dp", "ep"), None) - ep_token_spec = PartitionSpec(("dp", "ep"), None, None) - ep_weight_spec = PartitionSpec(("dp", "ep"), None) - - with self.mesh, global_shard_guard(self.mr): - routes_s = jax.lax.with_sharding_constraint( - routes, NamedSharding(self.mesh, route_spec) - ) - tokens_s = jax.lax.with_sharding_constraint( - tokens, NamedSharding(self.mesh, token_spec) - ) - weights_s = jax.lax.with_sharding_constraint( - weights, NamedSharding(self.mesh, token_spec) - ) - - def dispatch_one(route, input_tokens, topk_weights): - recv_tokens, recv_weights, _handle_mem, token_counts = ep_dispatch( - cfg, - route, - input_tokens, - topk_weights, - self.recv_capacity_per_rank, - ) - recv_tokens = jax.lax.with_sharding_constraint( - recv_tokens, NamedSharding(self.mesh, ep_token_spec) - ) - recv_weights = jax.lax.with_sharding_constraint( - recv_weights, NamedSharding(self.mesh, ep_weight_spec) - ) - # NCCL EP leaves capacity padding unspecified, including the - # weight buffer. Derive the valid prefix from token_counts. - align = max(int(cfg.dispatch_output_per_expert_alignment), 1) - padded_counts = ((token_counts + align - 1) // align) * align - valid_count = jnp.sum(padded_counts, axis=-1, keepdims=True) - slot = jnp.arange(self.recv_capacity_per_rank, dtype=jnp.int32) - valid = slot[None, :] < valid_count - recv_tokens = jnp.where( - valid[..., None], recv_tokens, jnp.zeros_like(recv_tokens) - ) - recv_weights = jnp.where( - valid, recv_weights, jnp.zeros_like(recv_weights) - ) - return recv_tokens, recv_weights - - @jax.jit - def scan_dispatch(route_maps, input_tokens, topk_weights): - def body(carry, route): - return carry, dispatch_one(route, input_tokens, topk_weights) - - return jax.lax.scan(body, (), route_maps)[1] - - scan_tokens, scan_weights = scan_dispatch(routes_s, tokens_s, weights_s) - scan_tokens.block_until_ready() - scan_tokens_global = jmu.process_allgather(scan_tokens, tiled=True) - scan_weights_global = jmu.process_allgather(scan_weights, tiled=True) - - # Run each route in a separate executable invocation so its - # expected packed layout cannot be affected by another layer's - # live handle. - isolated_dispatch = jax.jit(dispatch_one) - ref_tokens = [] - ref_weights = [] - for layer in range(routes.shape[0]): - layer_tokens, layer_weights = isolated_dispatch( - routes_s[layer], tokens_s, weights_s - ) - layer_tokens.block_until_ready() - ref_tokens.append( - np.asarray(jmu.process_allgather(layer_tokens, tiled=True)) - ) - ref_weights.append( - np.asarray(jmu.process_allgather(layer_weights, tiled=True)) - ) - - if self.rank == 0: - np.testing.assert_array_equal( - np.asarray(scan_tokens_global), np.stack(ref_tokens) - ) - np.testing.assert_array_equal( - np.asarray(scan_weights_global), np.stack(ref_weights) - ) - - @unittest.skip("Branch-only scan experiment; superseded by MaxText integration coverage.") - def test_scan_dispatch_combine_routing_handle_fwd_bwd(self): - """Scan must match unrolled layers when every iteration has a new map. - - Slot-dependent scaling prevents a stale dispatch/combine handle from - cancelling as it does for an identity expert. Comparing VJPs also - exercises ep_combine_bwd and ep_dispatch_bwd through the scan transpose. - """ - routes, tokens, weights = self._make_scan_inputs() - num_layers = routes.shape[0] - cfg = EpLayerConfig(top_k=1, dispatch_output_per_expert_alignment=16) - route_spec = PartitionSpec(None, ("dp", "ep"), None) - token_spec = PartitionSpec(("dp", "ep"), None) - - with self.mesh, global_shard_guard(self.mr): - routes_s = jax.lax.with_sharding_constraint( - routes, NamedSharding(self.mesh, route_spec) - ) - tokens_s = jax.lax.with_sharding_constraint( - tokens, NamedSharding(self.mesh, token_spec) - ) - weights_s = jax.lax.with_sharding_constraint( - weights, NamedSharding(self.mesh, token_spec) - ) - cotangent = jnp.asarray( - np.linspace( - 0.25, - 1.25, - tokens.size, - dtype=np.float32, - ).reshape(tokens.shape), - dtype=tokens.dtype, - ) - cotangent = jax.lax.with_sharding_constraint( - cotangent, NamedSharding(self.mesh, token_spec) - ) - - def scan_fwd(route_maps, initial_tokens, topk_weights): - def body(current_tokens, route): - out = self._scan_ep_layer( - cfg, - route, - current_tokens, - topk_weights, - slot_dependent_scale=True, - ) - return out, None - - return jax.lax.scan(body, initial_tokens, route_maps)[0] - - def unrolled_fwd(route_maps, initial_tokens, topk_weights): - out = initial_tokens - for layer in range(num_layers): - out = self._scan_ep_layer( - cfg, - route_maps[layer], - out, - topk_weights, - slot_dependent_scale=True, - ) - return out - - def value_and_token_vjp( - fn, route_maps, initial_tokens, topk_weights, out_cotangent - ): - out, pullback = jax.vjp( - lambda x: fn(route_maps, x, topk_weights), initial_tokens - ) - return out, pullback(out_cotangent)[0] - - scan_run = jax.jit( - lambda r, t, w, g: value_and_token_vjp(scan_fwd, r, t, w, g) - ) - unrolled_run = jax.jit( - lambda r, t, w, g: value_and_token_vjp(unrolled_fwd, r, t, w, g) - ) - scan_out, scan_grad = scan_run(routes_s, tokens_s, weights_s, cotangent) - ref_out, ref_grad = unrolled_run(routes_s, tokens_s, weights_s, cotangent) - scan_grad.block_until_ready() - ref_grad.block_until_ready() - - scan_out_global = jmu.process_allgather(scan_out, tiled=True) - ref_out_global = jmu.process_allgather(ref_out, tiled=True) - scan_grad_global = jmu.process_allgather(scan_grad, tiled=True) - ref_grad_global = jmu.process_allgather(ref_grad, tiled=True) - - if self.rank == 0: - self.assertFalse( - np.array_equal(np.asarray(ref_out_global), np.asarray(tokens)), - "slot-sensitive reference unexpectedly collapsed to an identity", - ) - np.testing.assert_array_equal( - np.asarray(scan_out_global), np.asarray(ref_out_global) - ) - np.testing.assert_array_equal( - np.asarray(scan_grad_global), np.asarray(ref_grad_global) - ) - def test_primitive_prepare(self): """ep_prepare returns token_counts, total_recv_tokens and handle_mem. diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 06cbce9d11..0e9cb6598d 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -118,13 +118,8 @@ def _read_mp_options(): ) from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import ( - _ALIGN_SIZE, - get_moe_recv_capacity_per_rank, - moe, - record_ep_bootstrap_signature_for_moe, -) -from transformer_engine.jax.ep import ep_bootstrap, ep_finalize +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 @@ -134,8 +129,7 @@ def _read_mp_options(): EP_AXIS = "ep" FSDP_AXIS = "fsdp" -EP_SIZE = int(os.environ.get("TE_EP_MOE_EP_SIZE", "2")) -assert EP_SIZE in (2, 4), f"TE_EP_MOE_EP_SIZE must be 2 or 4, got {EP_SIZE}" +EP_SIZE = 2 assert ( jax.device_count() % EP_SIZE == 0 ), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" @@ -144,11 +138,7 @@ def _read_mp_options(): LOGICAL_AXIS_RULES = ( ("exp", EP_AXIS), - # Match MaxText's converging layout: FSDP is the outer component of - # the compound expert dimension and EP is inner. - ("exp_fsdp", (FSDP_AXIS, EP_AXIS)), ("embed", FSDP_AXIS), - ("embed_replicated", None), ("mlp", None), ("batch", (FSDP_AXIS, EP_AXIS)), ) @@ -177,25 +167,6 @@ def _read_mp_options(): # (slot alignment rounding, etc.). TE_TO_TE_ATOL = 5e-3 TE_TO_TE_RTOL = 5e-3 -TE_TO_TE_GRAD_NORM_RATIO = (0.98, 1.02) -TE_TO_TE_GRAD_COSINE = 0.999 - - -def _assert_gradient_direction_and_scale(actual, expected, *, name): - """Catch permutations/reduction factors hidden by bf16 absolute tolerances.""" - actual = np.asarray(actual, dtype=np.float64).reshape(-1) - expected = np.asarray(expected, dtype=np.float64).reshape(-1) - actual_norm = np.linalg.norm(actual) - expected_norm = np.linalg.norm(expected) - assert actual_norm > 0.0 and expected_norm > 0.0, f"{name}: zero gradient norm" - norm_ratio = actual_norm / expected_norm - cosine = np.dot(actual, expected) / (actual_norm * expected_norm) - assert 0.8 <= norm_ratio <= 1.2, ( - f"{name}: gradient norm ratio {norm_ratio:.6f} outside [0.8, 1.2] " - f"(cosine={cosine:.6f})" - ) - assert cosine >= 0.98, f"{name}: gradient cosine similarity {cosine:.6f} < 0.98" - # 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 @@ -209,6 +180,28 @@ def _assert_gradient_direction_and_scale(actual, expected, *, name): # ----------------------------------------------------------------------------- +def _compute_worst_case_recv_pr(): + """Per-rank recv buffer the bootstrap must reserve. + + 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 + 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") def mesh(): if jax.device_count() < NUM_DEVICES_REQUIRED: @@ -224,20 +217,12 @@ def mesh(): num_procs = jax.process_count() max_tokens_per_rank = (BATCH // num_procs) * SEQ - recv_capacity_per_rank = get_moe_recv_capacity_per_rank( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - max_tokens_per_rank=max_tokens_per_rank, - ep_size=EP_SIZE, - recv_capacity_factor=2.0, - ) + 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) - ): + 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(), @@ -246,7 +231,6 @@ def mesh(): recv_capacity_per_rank=recv_capacity_per_rank, hidden_dim=HIDDEN, max_token_dtype=DTYPE, - drop_on_overflow=True, ) record_ep_bootstrap_signature_for_moe( num_experts=NUM_EXPERTS, @@ -328,9 +312,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 @@ -342,9 +324,7 @@ def _pure_jax_moe_reference( # 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_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: @@ -359,9 +339,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] @@ -386,18 +364,8 @@ def _make_block( use_expert_routing_bias=False, score_function="softmax", expert_bias_init=None, - compound_expert_sharding=False, input_axes=("batch", None, None), - recv_capacity_per_rank=None, ): - if recv_capacity_per_rank is None: - recv_capacity_per_rank = get_moe_recv_capacity_per_rank( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - max_tokens_per_rank=(BATCH // jax.process_count()) * SEQ, - ep_size=EP_SIZE, - recv_capacity_factor=2.0, - ) kwargs = dict( num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -409,13 +377,7 @@ def _make_block( score_function=score_function, dtype=DTYPE, input_axes=input_axes, - recv_capacity_per_rank=recv_capacity_per_rank, ) - if compound_expert_sharding: - # Match MaxText shard_exp_on_fsdp=True: FSDP and EP both shard the - # expert group axis while the hidden dimensions remain replicated. - kwargs["wi_kernel_axes"] = ("exp_fsdp", "embed_replicated", "mlp") - kwargs["wo_kernel_axes"] = ("exp_fsdp", "mlp", "embed_replicated") # Custom expert_bias_init lets tests inject a non-zero expert_bias without # poking variables['params'] post-init. if expert_bias_init is not None: @@ -435,13 +397,6 @@ def _strong_expert_bias_init(key, shape, dtype): ) -def _cross_rank_expert_bias_init(key, shape, dtype): - """Select one expert on each EP rank for a balanced reduced-capacity test.""" - del key - bias = jnp.full(shape, -5.0, dtype=dtype) - return bias.at[jnp.asarray((0, shape[0] // EP_SIZE))].set(5.0) - - 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( @@ -548,279 +503,6 @@ def _make_inputs(key): return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) -def _make_stacked_moe_params(mesh, num_layers): - """Create distinct production-shaped MoE parameters with a scan axis. - - The expert parameters use the same compound ``(fsdp, ep)`` expert - sharding as MaxText's ``shard_exp_on_fsdp=True`` integration. A rotating - strong routing bias makes each layer's routing map distinct while also - leaving half of the experts empty in every layer. - """ - gate_key, wi_key, wo_key = jax.random.split(jax.random.PRNGKey(101), 3) - - with _ctx(mesh): - - @jax.jit - def initialize(): - gate_kernel = jax.random.normal( - gate_key, - (num_layers, HIDDEN, NUM_EXPERTS), - dtype=DTYPE, - ) / np.sqrt(HIDDEN) - wi = jax.random.normal( - wi_key, - (num_layers, NUM_EXPERTS, HIDDEN, 2 * INTER), - dtype=DTYPE, - ) / np.sqrt(HIDDEN) - wo = jax.random.normal( - wo_key, - (num_layers, NUM_EXPERTS, INTER, HIDDEN), - dtype=DTYPE, - ) / np.sqrt(INTER) - base_bias = jnp.concatenate( - ( - jnp.full((NUM_EXPERTS // 2,), 10.0, dtype=jnp.float32), - jnp.full( - (NUM_EXPERTS - NUM_EXPERTS // 2,), - -10.0, - dtype=jnp.float32, - ), - ) - ) - expert_bias = jnp.stack( - [ - jnp.roll(base_bias, (2 * layer) % NUM_EXPERTS) - for layer in range(num_layers) - ] - ) - - gate_kernel = jax.lax.with_sharding_constraint( - gate_kernel, NamedSharding(mesh, P(None, None, None)) - ) - wi = jax.lax.with_sharding_constraint( - wi, - NamedSharding( - mesh, - P(None, (FSDP_AXIS, EP_AXIS), None, None), - ), - ) - wo = jax.lax.with_sharding_constraint( - wo, - NamedSharding( - mesh, - P(None, (FSDP_AXIS, EP_AXIS), None, None), - ), - ) - expert_bias = jax.lax.with_sharding_constraint( - expert_bias, NamedSharding(mesh, P(None, None)) - ) - return { - "gate_kernel": gate_kernel, - "wi": wi, - "wo": wo, - "expert_bias": expert_bias, - } - - params = initialize() - jax.block_until_ready(params["wi"]) - return params - - -def _functional_production_moe_layer(params, x): - """One residual TE MoE layer matching the failing integration knobs.""" - normed_x = ( - x.astype(jnp.float32) - * jax.lax.rsqrt( - jnp.mean(x.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + 1.0e-6 - ) - ).astype(x.dtype) - branch, _ = moe( - normed_x, - params["gate_kernel"], - params["wi"], - params["wo"], - expert_bias=params["expert_bias"], - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - score_function="sigmoid", - apply_topk_weights_early=False, - ep_axis=EP_AXIS, - data_parallelism_axes=(FSDP_AXIS,), - wi_kernel_axes=("exp_fsdp", "embed_replicated", "mlp"), - wo_kernel_axes=("exp_fsdp", "mlp", "embed_replicated"), - dtype=DTYPE, - ) - return x + branch - - -def _run_stacked_te_moe(params, x, *, use_scan, remat): - """Run identical per-layer parameters scanned or Python-unrolled.""" - layer_fn = _functional_production_moe_layer - if remat: - # MaxText applies nn.remat to its layer before passing the layer to - # nn.scan. Checkpointing the functional body gives the same important - # lowering: forward recomputation occurs inside reverse scan. - layer_fn = jax.checkpoint(layer_fn, prevent_cse=True) - - if use_scan: - - def scan_body(value, layer_params): - return layer_fn(layer_params, value), None - - return jax.lax.scan(scan_body, x, params)[0] - - value = x - for layer in range(params["gate_kernel"].shape[0]): - layer_params = jax.tree_util.tree_map(lambda p: p[layer], params) - value = layer_fn(layer_params, value) - return value - - -def _run_stacked_jax_reference(params, x): - """Pure-JAX residual stack using the same distinct layer parameters.""" - value = x - for layer in range(params["gate_kernel"].shape[0]): - layer_params = jax.tree_util.tree_map(lambda p: p[layer], params) - normed_value = ( - value.astype(jnp.float32) - * jax.lax.rsqrt( - jnp.mean(value.astype(jnp.float32) ** 2, axis=-1, keepdims=True) - + 1.0e-6 - ) - ).astype(value.dtype) - branch, _ = _pure_jax_moe_reference( - normed_value, - layer_params["gate_kernel"], - layer_params["wi"][..., :INTER], - layer_params["wi"][..., INTER:], - layer_params["wo"], - layer_params["expert_bias"], - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - score_function="sigmoid", - ) - value = value + branch - return value - - -def _stack_value_and_grad(params, x, *, use_scan, remat): - def loss_fn(p, value): - output = _run_stacked_te_moe( - p, - value, - use_scan=use_scan, - remat=remat, - ) - return jnp.mean(output.astype(jnp.float32) ** 2), output - - return jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True)(params, x) - - -def _gradient_similarity(actual, expected): - actual = np.asarray(actual, dtype=np.float64).reshape(-1) - expected = np.asarray(expected, dtype=np.float64).reshape(-1) - actual_norm = np.linalg.norm(actual) - expected_norm = np.linalg.norm(expected) - norm_ratio = actual_norm / expected_norm - cosine = np.dot(actual, expected) / (actual_norm * expected_norm) - return actual_norm, expected_norm, norm_ratio, cosine - - -def _assert_te_to_te_gradient(actual, expected, *, name): - """Strict scan-vs-unrolled oracle with useful failure diagnostics.""" - actual = np.asarray(actual) - expected = np.asarray(expected) - assert np.all(np.isfinite(actual)), f"{name}: scan gradient has NaN/Inf" - assert np.all(np.isfinite(expected)), f"{name}: unrolled gradient has NaN/Inf" - actual_norm, expected_norm, norm_ratio, cosine = _gradient_similarity( - actual, expected - ) - assert actual_norm > 0.0 and expected_norm > 0.0, ( - f"{name}: zero gradient norm " - f"(scan={actual_norm:.9e}, unrolled={expected_norm:.9e})" - ) - lo, hi = TE_TO_TE_GRAD_NORM_RATIO - assert lo <= norm_ratio <= hi and cosine >= TE_TO_TE_GRAD_COSINE, ( - f"{name}: scan/unrolled gradient direction or scale mismatch: " - f"scan_norm={actual_norm:.9e}, unrolled_norm={expected_norm:.9e}, " - f"ratio={norm_ratio:.9f}, cosine={cosine:.9f}, " - f"scan_mean={actual.astype(np.float64).mean():.9e}, " - f"scan_std={actual.astype(np.float64).std():.9e}, " - f"scan_absmax={np.abs(actual.astype(np.float64)).max():.9e}, " - f"unrolled_mean={expected.astype(np.float64).mean():.9e}, " - f"unrolled_std={expected.astype(np.float64).std():.9e}, " - f"unrolled_absmax={np.abs(expected.astype(np.float64)).max():.9e}" - ) - np.testing.assert_allclose( - actual.astype(np.float32), - expected.astype(np.float32), - atol=TE_TO_TE_ATOL, - rtol=TE_TO_TE_RTOL, - err_msg=f"{name}: scan/unrolled elementwise mismatch", - ) - - -def _assert_stacked_reference_gradient(actual, expected, *, name): - """Broad TE-vs-JAX stack control; strict parity is tested per block.""" - actual = np.asarray(actual) - expected = np.asarray(expected) - assert np.all(np.isfinite(actual)), f"{name}: TE gradient has NaN/Inf" - assert np.all(np.isfinite(expected)), f"{name}: reference gradient has NaN/Inf" - actual_norm, expected_norm, norm_ratio, cosine = _gradient_similarity( - actual, expected - ) - assert actual_norm > 0.0 and expected_norm > 0.0, ( - f"{name}: zero gradient norm " - f"(TE={actual_norm:.9e}, reference={expected_norm:.9e})" - ) - assert 0.5 <= norm_ratio <= 1.5 and cosine >= 0.85, ( - f"{name}: gross TE/reference stack mismatch: " - f"TE_norm={actual_norm:.9e}, reference_norm={expected_norm:.9e}, " - f"ratio={norm_ratio:.9f}, cosine={cosine:.9f}" - ) - - -def _assert_distinct_layer_routes(params_np, x_np): - """Prove the test does not accidentally reuse one routing map.""" - value = jnp.asarray(x_np) - signatures = [] - for layer in range(params_np["gate_kernel"].shape[0]): - gate = jnp.asarray(params_np["gate_kernel"][layer]) - bias = jnp.asarray(params_np["expert_bias"][layer]) - normed_value = ( - value.astype(jnp.float32) - * jax.lax.rsqrt( - jnp.mean(value.astype(jnp.float32) ** 2, axis=-1, keepdims=True) - + 1.0e-6 - ) - ).astype(value.dtype) - logits = jnp.einsum("bsh,he->bse", normed_value, gate).astype(jnp.float32) - _, indices = jax.lax.top_k(jax.nn.sigmoid(logits) + bias, TOPK) - indices_np = np.asarray(jax.device_get(indices), dtype=np.int64) - position = np.arange(indices_np.size, dtype=np.int64).reshape(indices_np.shape) - signatures.append( - ( - int(indices_np.sum()), - int((indices_np * (position + 1)).sum()), - ) - ) - branch, _ = _pure_jax_moe_reference( - normed_value, - gate, - jnp.asarray(params_np["wi"][layer])[..., :INTER], - jnp.asarray(params_np["wi"][layer])[..., INTER:], - jnp.asarray(params_np["wo"][layer]), - bias, - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - score_function="sigmoid", - ) - value = value + branch - assert ( - len(set(signatures)) > 1 - ), f"all layers used one routing signature: {signatures}" - - # ----------------------------------------------------------------------------- # Tests # ----------------------------------------------------------------------------- @@ -843,18 +525,6 @@ def _assert_distinct_layer_routes(params_np, x_np): dict(score_function="softmax", apply_topk_weights_early=True), id="softmax-early-weighting", ), - pytest.param( - dict(score_function="softmax", compound_expert_sharding=True), - id="softmax-compound-fsdp-expert", - ), - pytest.param( - dict( - score_function="sigmoid", - apply_topk_weights_early=True, - compound_expert_sharding=True, - ), - id="sigmoid-early-weighting-compound-fsdp-expert", - ), pytest.param( dict(score_function="sigmoid"), id="sigmoid", @@ -889,45 +559,6 @@ def _reference_kwargs_from_config(config, params_np): ) -class TestTeEpMoeReceiveCapacity: - """Reduced receive buffers preserve valid results and report overflow.""" - - def test_capacity_helper(self): - # Use a production-like token count where alignment does not collapse - # balanced and worst-case capacities to the same small test buffer. - max_tpr = 256 - worst = get_moe_recv_capacity_per_rank( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - max_tokens_per_rank=max_tpr, - ep_size=EP_SIZE, - ) - balanced = get_moe_recv_capacity_per_rank( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - max_tokens_per_rank=max_tpr, - ep_size=EP_SIZE, - recv_capacity_factor=1.0, - ) - headroom = get_moe_recv_capacity_per_rank( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - max_tokens_per_rank=max_tpr, - ep_size=EP_SIZE, - recv_capacity_factor=2.0, - ) - assert balanced < headroom < worst - assert balanced % _ALIGN_SIZE == 0 - with pytest.raises(ValueError, match="finite and >= 1.0"): - get_moe_recv_capacity_per_rank( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - max_tokens_per_rank=max_tpr, - ep_size=EP_SIZE, - recv_capacity_factor=0.5, - ) - - class TestTeEpMoeForward: """Per-config forward correctness in a single run: shape, dtype, finiteness AND numerical parity vs the pure-JAX reference.""" @@ -953,8 +584,8 @@ 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"])[..., :INTER], - jnp.asarray(params_np["wi"])[..., INTER:], + 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, @@ -992,8 +623,8 @@ def loss_fn(params, x): out, _ = _pure_jax_moe_reference( x, params["gate_kernel"], - params["wi"][..., :INTER], - params["wi"][..., INTER:], + params["wi_0"], + params["wi_1"], params["wo"], ref_expert_bias, num_experts=NUM_EXPERTS, @@ -1009,15 +640,11 @@ 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", "wo"): + 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}]" + 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" @@ -1030,11 +657,6 @@ def loss_fn(params, x): rtol=rtol, err_msg=f"grad parity breach on {name} [config={config}]", ) - _assert_gradient_direction_and_scale( - g_te, - grads_ref_np[name], - name=f"{name} [config={config}]", - ) # d_x: the gradient propagated back to the previous layer. Checks # shape, dtype (must match x.dtype — protects the @@ -1057,293 +679,6 @@ def loss_fn(params, x): rtol=GRAD_FFN_RTOL, err_msg=f"d_x parity breach [config={config}]", ) - _assert_gradient_direction_and_scale( - grad_x_te_np, - grad_x_ref_np, - name=f"d_x [config={config}]", - ) - - def test_repeated_production_block_backward(self, mesh): - """Catch small d_x errors that amplify across a stack of MoE blocks.""" - repeats = 8 - config = dict( - score_function="sigmoid", - apply_topk_weights_early=True, - compound_expert_sharding=True, - use_expert_routing_bias=True, - ) - block = _make_block(**config) - x = _make_inputs(jax.random.PRNGKey(22)) - variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(23)) - - with _ctx(mesh): - x_sh = _shard_inputs(x, mesh) - - def te_loss_fn(variables, value): - for _ in range(repeats): - branch, _, _ = block.apply(variables, value) - value = value + branch - return jnp.mean(value.astype(jnp.float32) ** 2) - - grads_te, grad_x_te = jax.jit(jax.grad(te_loss_fn, argnums=(0, 1)))( - variables, - x_sh, - ) - jax.block_until_ready(grad_x_te) - - params_np = _params_global_numpy(variables, mesh) - x_np = np.asarray(jax.device_get(x)) - expert_bias = jnp.asarray(params_np["expert_bias"]) - - def reference_loss_fn(params, value): - for _ in range(repeats): - branch, _ = _pure_jax_moe_reference( - value, - params["gate_kernel"], - params["wi"][..., :INTER], - params["wi"][..., INTER:], - params["wo"], - expert_bias, - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - score_function="sigmoid", - ) - value = value + branch - return jnp.mean(value.astype(jnp.float32) ** 2) - - grads_ref, grad_x_ref = jax.jit(jax.grad(reference_loss_fn, argnums=(0, 1)))( - {k: jnp.asarray(v) for k, v in params_np.items() if k != "expert_bias"}, - jnp.asarray(x_np), - ) - - for name in ("gate_kernel", "wi", "wo"): - _assert_gradient_direction_and_scale( - _to_global_numpy(_unwrap(grads_te["params"][name]), mesh), - np.asarray(jax.device_get(grads_ref[name])), - name=f"repeated {name}", - ) - _assert_gradient_direction_and_scale( - _to_global_numpy(grad_x_te, mesh), - np.asarray(jax.device_get(grad_x_ref)), - name="repeated d_x", - ) - - @pytest.mark.skip(reason="Experimental scan/remat parity coverage is currently disabled.") - @pytest.mark.parametrize("num_layers", (4, 8)) - def test_scanned_production_block_backward(self, mesh, num_layers): - """Full MoE scan/remat must match the identical unrolled stack. - - Unlike the primitive scan tests, this exercises top-2 routing, - compound FSDP/EP expert parameters, grouped-GEMM weight gradients, - the full ``moe`` custom VJP residual, and forward rematerialization - inside the reverse scan. - """ - params = _make_stacked_moe_params(mesh, num_layers) - x = _make_inputs(jax.random.PRNGKey(102)) - - with _ctx(mesh): - x_sh = _shard_inputs(x, mesh) - unrolled_run = jax.jit( - partial( - _stack_value_and_grad, - use_scan=False, - remat=False, - ) - ) - scan_run = jax.jit( - partial( - _stack_value_and_grad, - use_scan=True, - remat=False, - ) - ) - scan_remat_run = jax.jit( - partial( - _stack_value_and_grad, - use_scan=True, - remat=True, - ) - ) - - unrolled = unrolled_run(params, x_sh) - jax.block_until_ready(unrolled) - scanned = scan_run(params, x_sh) - jax.block_until_ready(scanned) - scanned_remat = scan_remat_run(params, x_sh) - jax.block_until_ready(scanned_remat) - - (unrolled_loss, unrolled_output), (unrolled_grads, unrolled_grad_x) = unrolled - modes = ( - ("scan", scanned), - ("scan-remat", scanned_remat), - ) - unrolled_output_np = _to_global_numpy(unrolled_output, mesh) - unrolled_grad_x_np = _to_global_numpy(unrolled_grad_x, mesh) - unrolled_grads_np = { - name: _to_global_numpy(grad, mesh) for name, grad in unrolled_grads.items() - } - - params_np = { - name: _to_global_numpy(param, mesh) for name, param in params.items() - } - x_np = np.asarray(jax.device_get(x)) - if jax.process_index() == 0: - _assert_distinct_layer_routes(params_np, x_np) - - mismatches = [] - - def record_check(check): - try: - check() - except AssertionError as error: - mismatches.append(str(error)) - - for mode_name, mode_result in modes: - (mode_loss, mode_output), (mode_grads, mode_grad_x) = mode_result - record_check( - lambda: np.testing.assert_allclose( - np.asarray(jax.device_get(mode_loss), dtype=np.float32), - np.asarray(jax.device_get(unrolled_loss), dtype=np.float32), - atol=TE_TO_TE_ATOL, - rtol=TE_TO_TE_RTOL, - err_msg=f"{mode_name}: loss mismatch", - ) - ) - mode_output_np = _to_global_numpy(mode_output, mesh) - record_check( - lambda: np.testing.assert_allclose( - mode_output_np.astype(np.float32), - unrolled_output_np.astype(np.float32), - atol=TE_TO_TE_ATOL, - rtol=TE_TO_TE_RTOL, - err_msg=f"{mode_name}: output mismatch", - ) - ) - mode_grad_x_np = _to_global_numpy(mode_grad_x, mesh) - record_check( - lambda: _assert_te_to_te_gradient( - mode_grad_x_np, - unrolled_grad_x_np, - name=f"{mode_name} d_x", - ) - ) - for param_name in ("gate_kernel", "wi", "wo"): - mode_grad_np = _to_global_numpy(mode_grads[param_name], mesh) - record_check( - lambda mode_grad_np=mode_grad_np, param_name=param_name: ( - _assert_te_to_te_gradient( - mode_grad_np, - unrolled_grads_np[param_name], - name=f"{mode_name} {param_name}", - ) - ) - ) - for layer in range(num_layers): - record_check( - lambda layer=layer, mode_grad_np=mode_grad_np, param_name=param_name: ( - _assert_te_to_te_gradient( - mode_grad_np[layer], - unrolled_grads_np[param_name][layer], - name=f"{mode_name} layer={layer} {param_name}", - ) - ) - ) - - # The unrolled TE execution is the control. Check it against a - # communication-free JAX reference so a scan/unrolled agreement cannot - # hide a bug shared by both TE executions. - reference_params = { - name: jnp.asarray(value) for name, value in params_np.items() - } - - def reference_loss_fn(p, value): - output = _run_stacked_jax_reference(p, value) - return jnp.mean(output.astype(jnp.float32) ** 2) - - reference_grads, reference_grad_x = jax.jit( - jax.grad(reference_loss_fn, argnums=(0, 1)) - )(reference_params, jnp.asarray(x_np)) - for param_name in ("gate_kernel", "wi", "wo"): - record_check( - lambda param_name=param_name: _assert_stacked_reference_gradient( - unrolled_grads_np[param_name], - np.asarray(jax.device_get(reference_grads[param_name])), - name=f"unrolled-reference {param_name}", - ) - ) - record_check( - lambda: _assert_stacked_reference_gradient( - unrolled_grad_x_np, - np.asarray(jax.device_get(reference_grad_x)), - name="unrolled-reference d_x", - ) - ) - if mismatches: - pytest.fail("\n\n".join(mismatches)) - - @pytest.mark.skip(reason="Experimental scan/remat parity coverage is currently disabled.") - def test_scanned_training_trajectory(self, mesh): - """Three SGD steps must retain scan/remat vs unrolled parity.""" - num_layers = 4 - learning_rate = 1.0e-3 - params = _make_stacked_moe_params(mesh, num_layers) - x = _make_inputs(jax.random.PRNGKey(103)) - - def make_step(*, use_scan, remat): - def loss_fn(p, value): - output = _run_stacked_te_moe( - p, - value, - use_scan=use_scan, - remat=remat, - ) - return jnp.mean(output.astype(jnp.float32) ** 2) - - def step(p, value): - loss, grads = jax.value_and_grad(loss_fn)(p, value) - updated = jax.tree_util.tree_map( - lambda weight, grad: weight - learning_rate * grad, - p, - grads, - ) - return updated, loss - - return jax.jit(step) - - with _ctx(mesh): - x_sh = _shard_inputs(x, mesh) - unrolled_step = make_step(use_scan=False, remat=False) - scan_remat_step = make_step(use_scan=True, remat=True) - unrolled_params = params - scan_remat_params = params - unrolled_losses = [] - scan_remat_losses = [] - for _ in range(3): - unrolled_params, unrolled_loss = unrolled_step(unrolled_params, x_sh) - scan_remat_params, scan_remat_loss = scan_remat_step( - scan_remat_params, x_sh - ) - jax.block_until_ready(unrolled_params) - jax.block_until_ready(scan_remat_params) - unrolled_losses.append(float(jax.device_get(unrolled_loss))) - scan_remat_losses.append(float(jax.device_get(scan_remat_loss))) - - np.testing.assert_allclose( - np.asarray(scan_remat_losses, dtype=np.float32), - np.asarray(unrolled_losses, dtype=np.float32), - atol=TE_TO_TE_ATOL, - rtol=TE_TO_TE_RTOL, - err_msg=( - "scan-remat training trajectory differs from unrolled: " - f"scan={scan_remat_losses}, unrolled={unrolled_losses}" - ), - ) - for param_name in ("gate_kernel", "wi", "wo"): - _assert_te_to_te_gradient( - _to_global_numpy(scan_remat_params[param_name], mesh), - _to_global_numpy(unrolled_params[param_name], mesh), - name=f"three-step parameter {param_name}", - ) class TestTeEpMoeAuxLoss: @@ -1375,8 +710,8 @@ 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"])[..., :INTER], - jnp.asarray(params_np["wi"])[..., INTER:], + 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, @@ -1394,9 +729,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" @@ -1408,82 +741,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", "wo"): - g_local = np.asarray( - jax.device_get(_unwrap(grads["params"][name]).addressable_data(0)) - ) + 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 TestZZTeEpMoeOverflow: - """Run last with a small exact capacity: valid routing then overflow.""" - - @pytest.fixture(scope="class", autouse=True) - @classmethod - def reduced_capacity_bootstrap(cls, mesh): - del cls - max_tpr = (BATCH // jax.process_count()) * SEQ - capacity = _ALIGN_SIZE - ep_finalize() - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ): - ep_bootstrap( - world_size=jax.process_count(), - rank=jax.process_index(), - num_experts=NUM_EXPERTS, - max_tokens_per_rank=max_tpr, - recv_capacity_per_rank=capacity, - hidden_dim=HIDDEN, - max_token_dtype=DTYPE, - drop_on_overflow=True, - ) - record_ep_bootstrap_signature_for_moe( - num_experts=NUM_EXPERTS, - max_tokens_per_rank=max_tpr, - recv_capacity_per_rank=capacity, - hidden_dim=HIDDEN, - ep_size=EP_SIZE, - ) - yield capacity - - @pytest.mark.parametrize( - ("expert_bias_init", "expect_overflow"), - ( - pytest.param(_cross_rank_expert_bias_init, False, id="within-capacity"), - pytest.param(_strong_expert_bias_init, True, id="overflow"), - ), - ) - def test_reduced_capacity_vjp( - self, mesh, reduced_capacity_bootstrap, expert_bias_init, expect_overflow - ): - capacity = reduced_capacity_bootstrap - - x = jax.random.normal(jax.random.PRNGKey(72), (BATCH, SEQ, HIDDEN), dtype=DTYPE) - block = _make_block( - score_function="sigmoid", - use_expert_routing_bias=True, - expert_bias_init=expert_bias_init, - recv_capacity_per_rank=capacity, - ) - with _ctx(mesh): - x_sh = _shard_inputs(x, mesh) - variables = jax.jit(block.init)(jax.random.PRNGKey(73), x_sh) - - def loss_fn(variables, inputs): - output, _aux, totals = block.apply(variables, inputs) - return jnp.mean(output.astype(jnp.float32) ** 2), totals - - (loss, totals), grads = jax.jit( - jax.value_and_grad(loss_fn, has_aux=True) - )(variables, x_sh) - jax.block_until_ready((loss, totals, grads)) - - observed = int(_to_global_numpy(totals, mesh).max()) - assert (observed > capacity) is expect_overflow - assert np.isfinite(float(jax.device_get(loss))) - assert all( - np.all(np.isfinite(np.asarray(jax.device_get(leaf.addressable_data(0))))) - for leaf in jax.tree_util.tree_leaves(grads) - ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py index 127b487650..16a2d75de6 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py +++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py @@ -129,3 +129,36 @@ def test_mxfp8_quantize_swizzle_fusion( return_rowwise=return_rowwise, return_transpose=return_transpose, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("M, N", [(96, 160), (4096, 576), (4096, 2112)]) +def test_mxfp8_bidirectional_swizzled_row_scale_padding(M: int, N: int) -> None: + """The specialized bidirectional kernel must not overwrite padded row scales.""" + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + quantizer.optimize_for_gemm = True + scale = quantizer(x)._rowwise_scale_inv.view(torch.uint8) + + scale_rows = torch.arange(M, device=scale.device, dtype=torch.int64).view(-1, 1) + scale_cols = torch.arange(N // 32, device=scale.device, dtype=torch.int64).view(1, -1) + num_tiles_x = math.ceil(N / 128) + scale_indices = ( + ((scale_rows // 128) * num_tiles_x + scale_cols // 4) * (128 * 4) + + (scale_rows % 32) * 16 + + ((scale_rows % 128) // 32) * 4 + + scale_cols % 4 + ) + valid_mask = torch.zeros(scale.numel(), dtype=torch.bool, device=scale.device) + valid_mask[scale_indices.view(-1)] = True + + torch.testing.assert_close( + scale.view(-1)[~valid_mask], + torch.zeros_like(scale.view(-1)[~valid_mask]), + atol=0, + rtol=0, + ) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d09c92ad49..499c22509c 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1559,13 +1559,15 @@ def test_grouped_mlp_single_group_mxfp8( """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" if ( runtime_offsets_supported - and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets() + and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets( + te.ops.ScaledSwiGLU + ) ): pytest.skip("Requires cuDNN frontend >= 1.27.0") monkeypatch.setattr( grouped_mlp_module, "_cudnn_frontend_supports_single_group_runtime_offsets", - lambda: runtime_offsets_supported, + lambda _activation_type: runtime_offsets_supported, ) self.test_grouped_mlp( group_size=1, diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 9b810b6cef..9734e24fbe 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -475,10 +475,6 @@ message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") # imports stay unresolved (and harmless) under default ELF lazy binding when # the gate trips. LD_BIND_NOW environments lose this property. set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib") -# Editable builds reuse their CMake directory. Clear a library path cached by -# an older checkout so headers and the static archive always come from the -# same nccl-extensions build. -unset(NCCL_EP_LIB CACHE) find_file(NCCL_EP_LIB NAMES libnccl_ep.a HINTS ${NCCL_EP_LIB_DIR} diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh index 522a9add8f..52ab6f49e9 100644 --- a/transformer_engine/common/cast/fp8/gated_fp8.cuh +++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh @@ -68,11 +68,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_elems = SHMEM_DIM_Y * SHMEM_DIM_X; constexpr size_t buff_elems_total = BUFFERS_NUM * buff_elems; diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 83b5a49cae..879c6befc3 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -124,11 +124,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) __shared__ float subamax_colwise_buff[SUBAMAX_BUFF_DIM_Y][CHUNK_DIM_X]; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 8c57f112d2..3900f072bb 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -130,11 +130,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0); extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); @@ -648,6 +646,30 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, float *const amax_ptr = reinterpret_cast(output->amax.dptr); const float *noop_ptr = reinterpret_cast(noop->data.dptr); + // Clear padding before either the generic or specialized kernel writes + // directly into the GEMM-swizzled scale layout. + if (with_gemm_swizzled_scales && (cols % 128 != 0 || rows % 128 != 0)) { + constexpr size_t zero_threads = 256; + if (use_rowwise_scaling) { + const size_t size_bytes = output->scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + if (use_colwise_scaling) { + const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->columnwise_scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + } + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( @@ -669,17 +691,26 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, bidimensional_traits::blockDIM::M) <= max_grid_dim_y; const bool is_full_rowwise_chunk = (cols % 128 == 0); + const bool has_full_bidimensional_chunks = + (rows % bidimensional_traits::colChunkElems == 0) && + (cols % bidimensional_traits::rowChunkElems == 0); + // Both rowwise and bidimensional cast-only kernels select their + // scale layout from WITH_GEMM_SWIZZLED_SCALES. const bool scaling_type_has_specialized_support = (scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk && rowwise_specialized_grid_fits) || - (scaling_type == ScalingType::BIDIMENSIONAL && + (scaling_type == ScalingType::BIDIMENSIONAL && has_full_bidimensional_chunks && bidimensional_specialized_grid_fits); - if (specialized::hasSpec() && - !WITH_GEMM_SWIZZLED_SCALES && scaling_type_has_specialized_support) { + // Specialized cast-only kernels do not consume the device noop flag. + // Preserve cached outputs by keeping noop-aware calls on the generic path. + if (noop_ptr == nullptr && + specialized::hasSpec() && + scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { - using traits = specialized::CastTraits; + using traits = specialized::CastTraits; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -698,7 +729,11 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, break; } case ScalingType::BIDIMENSIONAL: { - using traits = specialized::CastTraits; + using traits = + specialized::CastTraitsSwizzle; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -788,38 +823,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - // Zero out swizzled scales if padding is needed - /// TODO (tmoon) Handle this within the cast kernel - if (with_gemm_swizzled_scales) { - constexpr size_t TILE_DIM_X = 128; // Tile dim in data buffer - constexpr size_t TILE_DIM_Y = 128; - if (cols % TILE_DIM_X != 0 || rows % TILE_DIM_Y != 0) { - // Use a noop-aware zero kernel so that the clear is skipped - // when quantization is a noop (e.g. FP8 weight caching). - constexpr size_t zero_threads = 256; - if (use_rowwise_scaling) { - const size_t size_bytes = output->scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->scale_inv.dptr), size_bytes, - noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - if (use_colwise_scaling) { - const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->columnwise_scale_inv.dptr), - size_bytes, noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - } - } - switch (scaling_type) { case ScalingType::ROWWISE: { auto kernel = quantize_mxfp8_kernel #include "../../../util/ptx.cuh" +#include "../swizzle.cuh" // gemm_swizzled_scale_idx (parent dir, GEMM scale swizzle) #include "state_counter.cuh" -#include "swizzle.cuh" +#include "swizzle.cuh" // specialized/swizzle.cuh (TMA input bank-conflict swizzle) namespace transformer_engine { namespace dispatch { @@ -24,6 +25,10 @@ namespace quantize_kernel { namespace specialized { namespace ptx = transformer_engine::ptx; + +// Bring in the GEMM-swizzled scale index helper (from ../swizzle.cuh). +// Used only when the kernel is instantiated with CastTraits::_with_swizzled_scales=true. +using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; namespace { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -122,19 +127,20 @@ struct Layout { static constexpr int32_t num = M * N; }; -template +template struct CastTraits; // 1x32 -template -struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { +template +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false, _kSwizzled> { static constexpr bool isRowwise = true; static constexpr bool isColwise = false; using IType = _IType; using OType = _OType; static constexpr int32_t chunkElems = 32; - using threadLayout = Layout<1, 32>; + using threadLayout = Layout<1, THREADS_PER_WARP>; static constexpr int32_t numThreadsPerChunk = 1; static constexpr int32_t warpDimM = threadLayout::M; static constexpr int32_t warpDimN = threadLayout::N * chunkElems; @@ -151,14 +157,22 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { using iterLayout = Layout<1, 1>; static constexpr int32_t blockDimM = iterLayout::M * blockIterDimM; static constexpr int32_t blockDimN = iterLayout::N * blockIterDimN; + static constexpr int32_t rowwiseScaleStride = blockDimN / chunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 1; static constexpr int32_t numPrefetch = numStages - 1; static constexpr bool _use_cvt_4x = true; static constexpr bool _cache_rowwise_scale_in_smem = true; + static constexpr bool _with_swizzled_scales = _kSwizzled; - static constexpr int32_t numThreads = warpLayout::num * 32; + static constexpr int32_t numThreads = warpLayout::num * THREADS_PER_WARP; static constexpr size_t smem_rowwise_scale = _cache_rowwise_scale_in_smem ? (blockDimM * (blockDimN / chunkElems) * sizeof(e8m0_t)) : 0ul; @@ -498,13 +512,8 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re block_coords.y = blockIdx.y * CastTraits::blockDimM; block_coords.x = blockIdx.x * CastTraits::blockDimN; - constexpr int32_t stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDimM, rows); @@ -514,7 +523,40 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::chunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + if constexpr (CastTraits::_with_swizzled_scales) { + // Four adjacent rowwise scale columns are contiguous in the GEMM scale + // layout, so write them as one uint32_t whenever possible. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::chunkElems; + const size_t num_tiles_x = DIVUP(cols, static_cast(128)); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / groups_per_row; + const int32_t col = (i % groups_per_row) * cols_per_group; + const uint32_t value = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + *reinterpret_cast(&scales_rowwise[idx]) = value; + } + + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / remaining_per_row; + const int32_t col = remaining_start + (i % remaining_per_row); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + scales_rowwise[idx] = sRowwiseScale[row * stride_in_smem + col]; + } + } + } else if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { using DataType = int32_t; constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; @@ -527,8 +569,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -546,8 +589,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -577,11 +621,12 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr int32_t rowChunkElems = 32; static constexpr int32_t colChunkElems = 32; - using rowThreadLayout = Layout<32, 1>; // 32x1 + using rowThreadLayout = Layout; // 32x1 using colThreadLayout = Layout; // 1x32 static_assert(rowThreadLayout::num == colThreadLayout::num, "rowThreadLayout::num must be equal to colThreadLayout::num"); - static_assert(rowThreadLayout::num == 32, "rowThreadLayout::num must be 32"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); using rowWarpDim = Layout; using colWarpDim = Layout; @@ -603,6 +648,13 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { using iterLayout = Layout<1, 4>; using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 2; @@ -636,7 +688,7 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { "It requires aligned smem pointer"); static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; - static constexpr int32_t numThreads = numWarps * 32; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); @@ -660,7 +712,9 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; // && _colwise_reduce_max != ColwiseReduceMax::Redux; static constexpr size_t smem_colwise_reduce = - _need_smem_for_colwise_reduce ? 32 * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul; + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; static constexpr size_t smem = _reuse_input_out_smem @@ -670,9 +724,138 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { smem_alignment + smem_rowwise_scale + smem_colwise_reduce); }; -__device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) { - return (x + align - 1) & ~((align)-1); -} +// Standalone trait for the non-warp-specialized rowwise+colwise cast_only kernel. +// Exposes numStages, iterM, iterN, and the two colwise-scale features as +// caller-controllable template axes. Both colwise flags default to true, +// giving callers the swizzled + colwise-scale-cached fast path by default. +// +// This trait duck-types the same interface CastTraits<_, _, true, true> exposes, +// so it drops into quantize_mxfp8_kernel_cast_only without any +// changes to the kernel signature. Kernel #1 (rowwise-only) and Kernel #2 +// (warp-specialized row+col) don't accept it: kernel #1 requires isColwise=false, +// kernel #2 requires _use_warp_specialization=true - both are wrong here. +template +struct CastTraitsSwizzle { + static constexpr bool isRowwise = true; + static constexpr bool isColwise = true; + using IType = _IType; + using OType = _OType; + + static constexpr int32_t rowChunkElems = 32; + static constexpr int32_t colChunkElems = 32; + + using rowThreadLayout = Layout; // 32x1 + using colThreadLayout = Layout; // 1x32 + static_assert(rowThreadLayout::num == colThreadLayout::num, + "rowThreadLayout::num must be equal to colThreadLayout::num"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); + + using rowWarpDim = Layout; + using colWarpDim = Layout; + using warpDim = + Layout; + + static constexpr bool _tma_swizzle = true; + using warpLayout = Layout<1, 2>; + static_assert(_tma_swizzle ? (warpLayout::N == 2) : true); + static constexpr CUtensorMapSwizzle input_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + static constexpr CUtensorMapSwizzle output_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + using blockIterDim = Layout; + + using iterLayout = Layout<_IterM, _IterN>; + using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; + + static constexpr int32_t numStages = _NumStages; + + using inputUnitType = uint4; + static constexpr int32_t rowNumElemsPerUnit = sizeof(inputUnitType) / sizeof(IType); + static constexpr int32_t rowNumUnitsPerChunk = rowChunkElems / rowNumElemsPerUnit; + using inputElemSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>; + using inputUnitSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>; + + using colIndexSwz = swz::Swizzle<5, 0, 5>; + + using rowOutputUnitType = uint4; + static constexpr int32_t rowNumOutUnitsPerChunk = + rowChunkElems * sizeof(OType) / sizeof(rowOutputUnitType); + static constexpr int32_t rowOutNumElemsPerUnit = sizeof(rowOutputUnitType) / sizeof(OType); + + using rowOutputChunkSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>; + using colOutputSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>; + + static constexpr bool _use_cvt_4x = true; + static constexpr bool _use_warp_specialization = false; + static constexpr bool _need_wait_group = iterLayout::num > numStages; + static constexpr bool _reuse_input_out_smem = false; + static_assert(_reuse_input_out_smem == false, "Just don't use it"); + static constexpr bool _cache_rowwise_scale_in_smem = true; + + static constexpr bool _colwise_source_coming_from_rowwise = true; + static constexpr ColwiseReduceMax _colwise_reduce_max = ColwiseReduceMax::Redux; + static_assert(_colwise_reduce_max != ColwiseReduceMax::RedAsync, + "It requires aligned smem pointer"); + + // The two colwise-scale features exposed as caller-controllable trait axes. + // Both default to true so callers get the vectorized swizzled path by default. + static constexpr bool _cache_colwise_scale_in_smem = _kCacheColwise; + static constexpr bool _with_swizzled_scales = _kSwizzled; + + static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; + static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); + + static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); + static constexpr size_t smemInputPerBlock = smemInputPerWarp * warpLayout::num; + + static constexpr size_t smemRowwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemRowwiseOutputPerBlock = smemRowwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemColwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemColwiseOutputPerBlock = smemColwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemInput = smemInputPerBlock * numStages; + static constexpr size_t smemRowwiseOutput = smemRowwiseOutputPerBlock * numStages; + static constexpr size_t smemColwiseOutput = smemColwiseOutputPerBlock * numStages; + + static constexpr size_t smem_rowwise_scale = + _cache_rowwise_scale_in_smem ? (blockDIM::M * (blockDIM::N / rowChunkElems) * sizeof(e8m0_t)) + : 0ul; + + // Extra shmem for cached colwise scales - only when the flag is on. + static constexpr size_t smem_colwise_scale = + _cache_colwise_scale_in_smem ? (blockDIM::M / colChunkElems) * blockDIM::N * sizeof(e8m0_t) + : 0ul; + + using ColwiseReduceDataType = float; + static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; + static constexpr size_t smem_colwise_reduce = + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; + + static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; + static constexpr size_t smem = + _reuse_input_out_smem + ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce) + : (smemInput + smemRowwiseOutput + smemColwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce); +}; // 32x32 template ( - align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + char *smemAligned = align_up(smem, CastTraits::smem_alignment); IType *sInput = reinterpret_cast(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(sInput); @@ -728,12 +910,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } // TODO: maybe we can assign a different barrier for each warp @@ -744,8 +926,8 @@ __global__ void quantize_mxfp8_kernel_cast_only( #pragma unroll for (int32_t i = 0; i < CastTraits::numStages; i++) { ptx::mbarrier_init(&ldg_producer[i], 1); - ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * 32); - ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * 32); + ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); + ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); ptx::mbarrier_init(&stg_consumer[i], 1); } ptx::fence_mbarrier_init_release_cluster(); @@ -1085,15 +1267,10 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { - ptx::numbered_barrier_sync(CastTraits::warpLayout::num * 32, 0u); + ptx::numbered_barrier_sync(CastTraits::warpLayout::num * THREADS_PER_WARP, 0u); - constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); @@ -1117,8 +1294,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1137,8 +1315,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1179,8 +1358,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( block_coords.x = blockIdx.x * CastTraits::blockDIM::N; extern __shared__ char smem[]; - char *smemAligned = reinterpret_cast( - align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + char *smemAligned = align_up(smem, CastTraits::smem_alignment); IType *sInput = reinterpret_cast(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(sInput); @@ -1191,15 +1369,20 @@ __global__ void quantize_mxfp8_kernel_cast_only( // colwise output will reuse input buffer OType *sColOutput; e8m0_t *sRowwiseScale = nullptr; + e8m0_t *sColwiseScale = nullptr; ColwiseReduceDataType *sColwiseReduce = nullptr; if constexpr (CastTraits::_reuse_input_out_smem) { sColOutput = reinterpret_cast(sInput); if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sRowOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1211,9 +1394,13 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1223,7 +1410,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( rowOutputUnitType *sColOutputUnit = reinterpret_cast(sColOutput); if constexpr (CastTraits::_need_smem_for_colwise_reduce) { - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } __shared__ uint64_t producer[CastTraits::numStages]; @@ -1243,7 +1430,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_colwise_source_coming_from_rowwise && CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { - ptx::mbarrier_init(colwise_reduce_barrier, 32); + ptx::mbarrier_init(colwise_reduce_barrier, THREADS_PER_WARP); } ptx::fence_mbarrier_init_release_cluster(); @@ -1263,18 +1450,35 @@ __global__ void quantize_mxfp8_kernel_cast_only( (threadIdx.x % CastTraits::rowThreadLayout::N) * (CastTraits::rowChunkElems / CastTraits::rowNumElemsPerUnit); - size_t rowwise_scale_base_offset = - (block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N)) * - static_cast(scale_stride_rowwise) + + // Scale coordinates in absolute (compact) scale-tensor space. Shared by both + // compact-layout offsets and by the CastTraits::_with_swizzled_scales branches below. + const int32_t row_scale_row_base = + block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N); + const int32_t row_scale_col_base = (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::rowThreadLayout::N) * CastTraits::rowChunkElems) / - CastTraits::rowChunkElems; + CastTraits::rowChunkElems; + const int32_t col_scale_row_base = + (block_coords.y + warp_coords.y + + (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / + CastTraits::colChunkElems; + const int32_t col_scale_col_base = + block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N); + + size_t rowwise_scale_base_offset = + static_cast(row_scale_row_base) * static_cast(scale_stride_rowwise) + + row_scale_col_base; size_t colwise_scale_base_offset = - ((block_coords.y + warp_coords.y + - (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / - CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N)); + static_cast(col_scale_row_base) * static_cast(scale_stride_colwise) + + col_scale_col_base; + + // Precomputed swizzle constants (each swizzle tile is 128 rows x 4 cols in scale space). + // Rowwise scale tensor has DIVUP(cols, 128) tiles across; + // colwise scale tensor has DIVUP(rows, 128) tiles across (X/Y axes are transposed). + const size_t row_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(cols, static_cast(128)) : 0; + const size_t col_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(rows, static_cast(128)) : 0; constexpr int32_t rowwise_scale_stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; @@ -1401,6 +1605,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( iter_m * CastTraits::blockIterDim::M * rowwise_scale_stride_in_smem + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); sRowwiseScale[rowwise_scale_offset] = row_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = row_scale_row_base + iter_m * CastTraits::blockIterDim::M; + int32_t abs_col = row_scale_col_base + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = row_biased_exponent; } else { size_t rowwise_scale_offset = rowwise_scale_base_offset + @@ -1416,12 +1626,32 @@ __global__ void quantize_mxfp8_kernel_cast_only( e8m0_t col_biased_exponent = to_e8m0(col_amax); float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); sColwiseReduce[threadIdx.x] = col_scale_inverse; - size_t colwise_scale_offset = - colwise_scale_base_offset + - iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - iter_n * CastTraits::blockIterDim::N; - scales_colwise[colwise_scale_offset] = col_biased_exponent; + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + // Cache in shmem; end-of-kernel flush handles gmem indexing. + int32_t smem_row = (warp_coords.y + (threadIdx.x / CastTraits::colThreadLayout::N) * + CastTraits::colChunkElems) / + CastTraits::colChunkElems + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t smem_col = warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N) + + iter_n * CastTraits::blockIterDim::N; + sColwiseScale[smem_row * CastTraits::blockDIM::N + smem_col] = col_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = col_scale_row_base + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t abs_col = col_scale_col_base + iter_n * CastTraits::blockIterDim::N; + // Colwise scale tensor's X/Y axes are transposed vs rowwise + // (see col_swz_num_tiles_X = DIVUP(rows, 128)), so pass + // (abs_col, abs_row) - abs_col is the swizzle "row" dim. + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + scales_colwise[idx] = col_biased_exponent; + } else { + size_t colwise_scale_offset = + colwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + iter_n * CastTraits::blockIterDim::N; + scales_colwise[colwise_scale_offset] = col_biased_exponent; + } __syncwarp(); } } @@ -1546,58 +1776,210 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); - end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, - scale_stride_rowwise); + if constexpr (CastTraits::_with_swizzled_scales) { + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + DIVUP(cols, static_cast(CastTraits::rowChunkElems))); + } else { + // The compact layout's padded entries are consumed by a later swizzle. + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); + } int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { - using DataType = int32_t; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + if constexpr (CastTraits::_with_swizzled_scales) { + // Swizzled flush: within a 128x4 swizzle tile, 4 consecutive column entries + // for the same row live at 4 consecutive gmem bytes. Group by 4 so each + // thread writes a uint32_t when col%4 == 0, then scalar tail for remainder. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::rowChunkElems; + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + + uint32_t val4 = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); + + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + *reinterpret_cast(&scales_rowwise[idx]) = val4; + } - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + // Tail (valid_coords.x % 4 != 0) + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_per_row; + int32_t col = remaining_start + (i % remaining_per_row); + e8m0_t val = sRowwiseScale[row * stride_in_smem + col]; + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = val; + } } } else { - using DataType = PreferredDataType; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + using PreferredDataType = typename CastTraits::PreferredDataType; - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + using DataType = int32_t; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } else { + using DataType = PreferredDataType; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } + } + } + + // Cached colwise scale flush (swizzled path only). Same barrier semantics as + // the rowwise flush above: the last-iter __syncthreads already ordered every + // in-loop sColwiseScale byte store before this block reads them. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + // In GEMM swizzle, contiguous bytes run over four scale-row indices for a + // fixed logical column. A 64-row CTA owns two such rows; pack them as a + // uint16 store only when the pair stays inside the same 4-row swizzle group. + const int32_t row_pairs = valid_rows / 2; + const int32_t total_pairs = row_pairs * valid_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_pairs; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = (i / valid_cols) * 2; + int32_t col = i % valid_cols; + const int32_t abs_col = block_coords.x + col; + const int32_t abs_row = scale_row_base + row; + e8m0_t val0 = sColwiseScale[row * CastTraits::blockDIM::N + col]; + e8m0_t val1 = sColwiseScale[(row + 1) * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + if (((abs_row & 3) != 3) && + ((reinterpret_cast(&scales_colwise[idx]) & (alignof(uint16_t) - 1)) == 0)) { + uint16_t val2 = static_cast(val0) | (static_cast(val1) << 8); + *reinterpret_cast(&scales_colwise[idx]) = val2; + } else { + scales_colwise[idx] = val0; + size_t idx1 = gemm_swizzled_scale_idx(abs_col, abs_row + 1, col_swz_num_tiles_X); + scales_colwise[idx1] = val1; + } + } + // Odd-row tail. + if ((valid_rows & 1) != 0) { + const int32_t row = valid_rows - 1; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < valid_cols; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t col = i; + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(block_coords.x + col, scale_row_base + row, + col_swz_num_tiles_X); + scales_colwise[idx] = val; + } + } + } + + // Cached colwise scale flush (non-swizzled linear layout). + // Rows in scales_colwise are indexed by (block_coords.y / colChunkElems), + // columns are logical input columns; 4 adjacent columns for the same scale + // row live at 4 consecutive gmem bytes, so pack as uint32 stores. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && !CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_cols / cols_per_group; + const int32_t total_groups = valid_rows * groups_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + uint32_t val4 = + *reinterpret_cast(&sColwiseScale[row * CastTraits::blockDIM::N + col]); + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + *reinterpret_cast(&scales_colwise[idx]) = val4; + } + // Column tail (valid_cols not multiple of 4). + const int32_t remaining_start = groups_per_row * cols_per_group; + if (remaining_start < valid_cols) { + const int32_t remaining_cols = valid_cols - remaining_start; + const int32_t total_remaining = valid_rows * remaining_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_cols; + int32_t col = remaining_start + (i % remaining_cols); + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + scales_colwise[idx] = val; } } } @@ -1605,7 +1987,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( ptx::cp_async_bulk_wait_group_read<0>(); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} +} // NOLINT(readability/fn_size) } // namespace specialized } // namespace quantize_kernel diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index 91c6af26b5..0b980c5086 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -262,11 +262,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a38a620ebe..8fc4cc7c46 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -410,11 +410,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); @@ -952,11 +950,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..8fd7d7480c 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -1153,6 +1153,26 @@ inline bool is_aligned_ptr(const void *ptr, size_t alignment) { return reinterpret_cast(ptr) % alignment == 0; } +/*! \brief Align a shared-memory base pointer up to `align` bytes. + * + * The result is derived from `p` by pointer arithmetic on purpose, without losing its + * identity as a pointer in between, so the address is never rounded through an integer + * -- in which case the compiler would lose the link back to the `extern __shared__` + * object, and ptxas could no longer prove the address lives in the shared window and + * would fall back to generic address-space accesses (`LD.E`/`ST.E`) instead of + * `LDS`/`STS`. + * + * `align` must be a power of two. + */ +__device__ __forceinline__ char *align_up(char *p, uintptr_t align) { + const uintptr_t misalign = reinterpret_cast(p) & (align - 1); + // If p is not aligned, (align - misalign) & (align - 1) is the number of bytes to fill the gap between p and + // the next aligned address. + // If p is aligned, misalign is 0 and (align - misalign) is align itself, so we use & (align - 1) + // to make it 0 and return p itself. + return p + ((align - misalign) & (align - 1)); +} + inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) { return is_aligned_ptr(static_cast(t.data.dptr), alignment); } diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index f29b9990f2..0981289ffe 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -15,8 +15,6 @@ #include #include -#include -#include #include #include "../util/logging.h" @@ -59,17 +57,6 @@ inline void* handle_mem_ptr(NVTETensor mem) { NVTE_CHECK(p != nullptr, "handle_mem tensor data must not be null"); return p; } - -bool trace_handle_mem_ptrs() { - const char* value = std::getenv("NVTE_EP_TRACE_HANDLE_MEM_PTRS"); - return value != nullptr && std::strcmp(value, "0") != 0; -} - -void trace_handle_mem_ptr(const char* operation, NVTETensor handle_mem) { - if (!trace_handle_mem_ptrs()) return; - std::printf("%s handle_mem_ptr: %p\n", operation, handle_mem_ptr(handle_mem)); - std::fflush(stdout); -} } // namespace void nvte_ep_initialize(void* ep_comm, const NVTEEpGroupConfig* group_config) { @@ -88,7 +75,6 @@ size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg) { void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv_tokens_per_expert, NVTETensor total_recv_tokens_per_rank, const NVTEEpLayerConfig* layer_cfg, cudaStream_t stream) { - trace_handle_mem_ptr("nvte_ep_prepare", handle_mem); NVTEEpLayerConfig cfg = normalize_ep_config(layer_cfg, kLayerConfigMinSize, "layer_cfg"); EPBackend::get().prepare(handle_mem_ptr(handle_mem), topk_idx, recv_tokens_per_expert, total_recv_tokens_per_rank, cfg, stream); @@ -99,7 +85,6 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok NVTECommWindow topk_weights_win, NVTETensor recv_tokens, NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, NVTECommWindow recv_topk_weights_win, cudaStream_t stream) { - trace_handle_mem_ptr("nvte_ep_dispatch", handle_mem); EPBackend::get().dispatch(handle_mem_ptr(handle_mem), topk_idx, tokens, tokens_win, topk_weights, topk_weights_win, recv_tokens, recv_tokens_win, recv_topk_weights, recv_topk_weights_win, stream); @@ -107,7 +92,6 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok void nvte_ep_combine(NVTETensor handle_mem, NVTETensor expert_out, NVTECommWindow expert_out_win, NVTETensor result, cudaStream_t stream) { - trace_handle_mem_ptr("nvte_ep_combine", handle_mem); EPBackend::get().combine(handle_mem_ptr(handle_mem), expert_out, expert_out_win, result, stream); } @@ -115,7 +99,6 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow NVTETensor g_recv_topk_weights, NVTECommWindow g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream) { - trace_handle_mem_ptr("nvte_ep_dispatch_bwd", handle_mem); EPBackend::get().dispatch_bwd(handle_mem_ptr(handle_mem), grad, grad_win, g_recv_topk_weights, g_recv_topk_weights_win, grad_tokens, grad_topk_weights, stream); } @@ -123,7 +106,6 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, cudaStream_t stream) { - trace_handle_mem_ptr("nvte_ep_combine_bwd", handle_mem); EPBackend::get().combine_bwd(handle_mem_ptr(handle_mem), grad, grad_win, grad_expert_out, grad_expert_out_win, stream); } diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 402d3c5b05..2725008257 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -18,7 +18,6 @@ #include #include #include -#include #include #include "../common.h" @@ -30,63 +29,6 @@ namespace ep { namespace { -std::atomic ep_trace_sequence{0}; -std::atomic ep_prepare_update_sequence{0}; - -bool trace_ep_handles() { - static const bool enabled = [] { - const char* value = std::getenv("NVTE_EP_TRACE_HANDLES"); - return value != nullptr && std::strcmp(value, "0") != 0; - }(); - return enabled; -} - -uint64_t fingerprint_routing_descriptor(const ncclEpTensor_t& routing) { - // This intentionally fingerprints only host-visible descriptor metadata. - // Reading routing values would require extra device work or host completion - // and could perturb CUDA graph capture or EP stream ordering. - constexpr uint64_t kFnvOffset = 1469598103934665603ULL; - constexpr uint64_t kFnvPrime = 1099511628211ULL; - uint64_t hash = kFnvOffset; - const auto mix = [&hash](uint64_t value) { - hash ^= value; - hash *= kFnvPrime; - }; - mix(static_cast(reinterpret_cast(routing.data))); - mix(static_cast(routing.ndim)); - mix(static_cast(routing.datatype)); - for (int i = 0; i < routing.ndim; ++i) { - mix(static_cast(routing.sizes[i])); - } - return hash; -} - -void trace_ep_handle(const char* op, void* handle_mem_ptr, ncclEpHandle_t handle, - cudaStream_t stream, const char* cache_state = "", - const ncclEpTensor_t* routing = nullptr, uint64_t update_id = 0) { - if (!trace_ep_handles()) return; - const uint64_t sequence = ep_trace_sequence.fetch_add(1, std::memory_order_relaxed); - const void* routing_ptr = routing == nullptr ? nullptr : routing->data; - const int routing_ndim = routing == nullptr ? 0 : routing->ndim; - const int64_t routing_dim0 = - routing == nullptr || routing->ndim < 1 ? 0 : routing->sizes[0]; - const int64_t routing_dim1 = - routing == nullptr || routing->ndim < 2 ? 0 : routing->sizes[1]; - const uint64_t routing_fingerprint = - routing == nullptr ? 0 : fingerprint_routing_descriptor(*routing); - std::fprintf(stderr, - "NVTE_EP_HANDLE_TRACE pid=%d seq=%lu op=%s handle_mem_ptr=%p " - "nccl_ep_handle=%p cache=%s stream=%p update_id=%lu " - "routing_ptr=%p routing_shape=[%ld,%ld] routing_ndim=%d " - "routing_descriptor_fingerprint=0x%016lx\n", - static_cast(getpid()), static_cast(sequence), op, handle_mem_ptr, - static_cast(handle), cache_state, static_cast(stream), - static_cast(update_id), routing_ptr, - static_cast(routing_dim0), static_cast(routing_dim1), routing_ndim, - static_cast(routing_fingerprint)); - std::fflush(stderr); -} - ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) { switch (dtype) { case kNVTEFloat32: @@ -317,8 +259,7 @@ size_t EPBackend::cache_cap_locked() { return handle_cache_cap_; } -ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg, - cudaStream_t stream) { +ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg) { // Update the program-wide fallback cfg so dispatch/combine/_bwd can // reconstruct the handle on a pointer-cache miss (WAR for XLA buffer reloc // between runs; one cfg per process). Remove this once XLA preserves the @@ -338,8 +279,6 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon auto it = index_.find(handle_mem); if (it != index_.end()) { lru_.splice(lru_.begin(), lru_, it->second); - trace_ep_handle("prepare_handle", handle_mem, it->second->handle, stream, "hit", nullptr, - it->second->last_update_id); return it->second->handle; } ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; @@ -349,8 +288,7 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon layer_cfg.top_k)); ncclEpHandle_t h = open_handle(handle_mem, hm_size, layer_cfg.top_k, layer_cfg.dispatch_output_per_expert_alignment); - trace_ep_handle("prepare_handle", handle_mem, h, stream, "miss"); - lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size, 0}); + lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size}); index_.emplace(handle_mem, lru_.begin()); while (lru_.size() > cache_cap_locked()) { HandleEntry& victim = lru_.back(); @@ -361,14 +299,10 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon return h; } -ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem, cudaStream_t stream, - uint64_t* last_update_id) { +ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem) { auto it = index_.find(handle_mem); if (it != index_.end()) { lru_.splice(lru_.begin(), lru_, it->second); - *last_update_id = it->second->last_update_id; - trace_ep_handle("lookup_handle", handle_mem, it->second->handle, stream, "hit", nullptr, - *last_update_id); return it->second->handle; } // Miss: reconstruct from the process-wide cached cfg. XLA may relocate @@ -378,9 +312,7 @@ ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem, cudaStream_t st const uintptr_t hm_addr = reinterpret_cast(handle_mem); NVTE_CHECK(fallback_layer_cfg_.has_value(), "ep op on handle_mem=0x", hm_addr, " with no cached entry and no prior nvte_ep_prepare; call prepare first."); - ncclEpHandle_t handle = prepare_handle_locked(handle_mem, *fallback_layer_cfg_, stream); - *last_update_id = 0; - return handle; + return prepare_handle_locked(handle_mem, *fallback_layer_cfg_); } // --------------------------------------------------------------------------- @@ -430,13 +362,7 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg, stream); - const uint64_t update_id = - ep_prepare_update_sequence.fetch_add(1, std::memory_order_relaxed) + 1; - auto entry = index_.find(handle_mem); - NVTE_CHECK(entry != index_.end(), "EP handle cache entry disappeared during prepare"); - entry->second->last_update_id = update_id; - trace_ep_handle("prepare_update", handle_mem, h, stream, "", &nccl_topk_idx, update_id); + ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } @@ -499,10 +425,7 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - uint64_t update_id = 0; - ncclEpHandle_t h = lookup_handle_locked(handle_mem, stream, &update_id); - trace_ep_handle(is_forward ? "dispatch_fwd" : "combine_bwd", handle_mem, h, stream, "", nullptr, - update_id); + ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, /*layout_info=*/nullptr, &dispatch_cfg, stream)); } @@ -526,9 +449,7 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - uint64_t update_id = 0; - ncclEpHandle_t h = lookup_handle_locked(handle_mem, stream, &update_id); - trace_ep_handle("combine_fwd", handle_mem, h, stream, "", nullptr, update_id); + ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } @@ -566,9 +487,7 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); - uint64_t update_id = 0; - ncclEpHandle_t h = lookup_handle_locked(handle_mem, stream, &update_id); - trace_ep_handle("dispatch_bwd", handle_mem, h, stream, "", nullptr, update_id); + ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); } diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index b6eb4101b2..80c9b9cea3 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -93,7 +93,6 @@ class EPBackend { ncclEpHandle_t handle; NVTEEpLayerConfig layer_cfg; size_t handle_mem_size; - uint64_t last_update_id; }; ncclEpGroup_t ep_group_{nullptr}; @@ -107,10 +106,8 @@ class EPBackend { std::optional fallback_layer_cfg_; // Caller must hold mutex_. - ncclEpHandle_t prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg, - cudaStream_t stream); - ncclEpHandle_t lookup_handle_locked(void* handle_mem, cudaStream_t stream, - uint64_t* last_update_id); + ncclEpHandle_t prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg); + ncclEpHandle_t lookup_handle_locked(void* handle_mem); size_t cache_cap_locked(); }; diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 54ec369b90..ca70ea145c 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -900,12 +900,7 @@ def partition( result_infos, ): del is_outer, result_infos - # The combine cotangent must have the same token sharding as the - # forward combine output. Transpose propagation can otherwise infer a - # replicated grad and pass the global token count to a handle prepared - # for only the rank-local tokens. - grad_sharding = NamedSharding(mesh, _ep_output_spec()) - arg_shardings = (arg_infos[0].sharding, grad_sharding) + arg_shardings = tuple(a.sharding for a in arg_infos) # EP-output leading (trailing dims auto-pad to None). out_sharding = NamedSharding(mesh, _ep_output_spec()) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 176a98e85b..fe3515ba03 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -6,7 +6,6 @@ import math import operator import os -import sys from collections.abc import Iterable from dataclasses import dataclass from functools import partial, reduce, cache @@ -52,16 +51,7 @@ ) from .misc import get_padded_spec, is_all_reduce_in_float32, get_min_device_compute_capability from ..sharding import ( - common_spec_axis, - filter_spec_axes, global_mesh_resource, - local_2d_sizes_from_spec, - local_shape_from_spec, - merge_axis_specs, - spec_axes, - spec_contains_axis, - strip_axis_from_spec, - supported_grouped_partition_axes, tpsp_axis_size, dp_or_fsdp_axis_size, ) @@ -84,16 +74,6 @@ num_cublas_streams = get_num_compute_streams() -_debug_python_patch = os.getenv("NVTE_DEBUG_PYTHON_PATCH", "0") == "1" -_debug_grouped_gemm_partition_count = 0 -if _debug_python_patch: - print( - "[TE patch debug] imported transformer_engine.jax.cpp_extensions.gemm " - f"from {__file__} (rank={os.getenv('SLURM_PROCID', 'unknown')})", - file=sys.stderr, - flush=True, - ) - # Cache whether the CUDA-graphable grouped GEMM implementation is available at import time. # Calling get_grouped_gemm_setup_workspace_size raises a RuntimeError mentioning "cublas" when # compiled against cuBLAS < 13.2, in which case the cuda-graphable path is unavailable. @@ -233,20 +213,6 @@ def _get_nvfp4_tensor_scale_inv(amax): return amax / (DATA_DTYPE_MAX * SCALE_DTYPE_MAX) -def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): - ignored_axes = tuple( - axis for axis in spec_axes(original_spec) if axis not in spec_axes(partition_spec) - ) - if ignored_axes: - warnings.warn( - "Grouped GEMM custom partitioning will ignore/replicate sharding " - f"axes {ignored_axes} from {arg_name}; only DP/FSDP/EP grouped " - "partitioning axes are preserved.", - RuntimeWarning, - stacklevel=3, - ) - - def collective_gemm_bootstrap( num_total_devices, num_devices_per_process, @@ -1817,374 +1783,6 @@ def impl( ) return (out,) - @staticmethod - def _parse_partition_specs( - mesh, - arg_infos, - result_infos, - out_shape=None, - lhs_is_trans=None, - lhs_axis_boundary=None, - ): - gsr = global_mesh_resource(validate=False) - fsdp_axis = gsr.fsdp_resource - allowed_axes = supported_grouped_partition_axes(mesh) - - original_arg_specs = tuple(get_padded_spec(arg_info) for arg_info in arg_infos) - lhs_data_spec = filter_spec_axes(original_arg_specs[0], allowed_axes) - lhs_scale_spec = filter_spec_axes(original_arg_specs[1], allowed_axes) - rhs_data_spec = filter_spec_axes(original_arg_specs[2], allowed_axes) - rhs_scale_spec = filter_spec_axes(original_arg_specs[3], allowed_axes) - bias_spec = filter_spec_axes(original_arg_specs[4], allowed_axes) - - lhs_first_dims_spec = filter_spec_axes(original_arg_specs[5], allowed_axes) - lhs_last_dims_spec = filter_spec_axes(original_arg_specs[6], allowed_axes) - rhs_first_dims_spec = filter_spec_axes(original_arg_specs[7], allowed_axes) - rhs_last_dims_spec = filter_spec_axes(original_arg_specs[8], allowed_axes) - out_first_dims_spec = filter_spec_axes(original_arg_specs[9], allowed_axes) - out_last_dims_spec = filter_spec_axes(original_arg_specs[10], allowed_axes) - additional_arg_0_spec = filter_spec_axes(original_arg_specs[11], allowed_axes) - additional_arg_1_spec = filter_spec_axes(original_arg_specs[12], allowed_axes) - - grouped_dim_specs = ( - lhs_first_dims_spec, - lhs_last_dims_spec, - rhs_first_dims_spec, - rhs_last_dims_spec, - out_first_dims_spec, - out_last_dims_spec, - ) - grouped_dim_infos = arg_infos[5:11] - active_group_spec = next( - (spec for spec, info in zip(grouped_dim_specs, grouped_dim_infos) if info.size > 0), - (None,), - ) - if arg_infos[11].size > 1: - additional_arg_0_spec = active_group_spec - if arg_infos[12].size > 1: - additional_arg_1_spec = active_group_spec - - rhs_is_ragged = arg_infos[7].size > 0 or arg_infos[8].size > 0 - ep_axis = gsr.ep_resource - if ( - ep_axis is not None - and not rhs_is_ragged - and spec_contains_axis(active_group_spec, ep_axis) - ): - if len(rhs_data_spec) > 0 and not spec_contains_axis(rhs_data_spec, ep_axis): - rhs_data_spec = ( - merge_axis_specs(rhs_data_spec[0], ep_axis), - *rhs_data_spec[1:], - ) - if len(rhs_scale_spec) > 0 and not spec_contains_axis(rhs_scale_spec, ep_axis): - rhs_scale_spec = ( - merge_axis_specs(rhs_scale_spec[0], ep_axis), - *rhs_scale_spec[1:], - ) - if len(bias_spec) > 0 and not spec_contains_axis(bias_spec, ep_axis): - bias_spec = (merge_axis_specs(bias_spec[0], ep_axis), *bias_spec[1:]) - - # A compound leading group dimension may use FSDP as the outer - # data-parallel axis (for example MoE groups ordered - # (fsdp, ep, local_expert)). Normally FSDP then describes distinct - # groups. MoE model weights are the exception: they retain one - # global expert set while the active group array has one copy per - # FSDP rank. Gather that shared, already-quantized RHS only when the - # group-count ratio exactly matches the FSDP mesh size. The exact - # ratio prevents unrelated unequal group shapes from selecting this - # specialized mapping. - fsdp_is_group_axis = spec_contains_axis(active_group_spec, fsdp_axis) - active_group_count = next( - (info.shape[0] for info in grouped_dim_infos if info.size > 0), - None, - ) - rhs_group_count = arg_infos[2].shape[0] if len(arg_infos[2].shape) > 0 else None - rhs_is_fsdp_shared_group_set = ( - fsdp_axis is not None - and active_group_count is not None - and rhs_group_count is not None - and rhs_group_count * mesh.shape[fsdp_axis] == active_group_count - ) - gather_rhs_fsdp = ( - fsdp_axis is not None - and not rhs_is_ragged - and (not fsdp_is_group_axis or rhs_is_fsdp_shared_group_set) - and ( - spec_contains_axis(rhs_data_spec, fsdp_axis) - or spec_contains_axis(rhs_scale_spec, fsdp_axis) - or spec_contains_axis(bias_spec, fsdp_axis) - ) - ) - - global _debug_grouped_gemm_partition_count - if _debug_python_patch and _debug_grouped_gemm_partition_count < 20: - _debug_grouped_gemm_partition_count += 1 - print( - "[TE patch debug] GroupedGemmPrimitive._parse_partition_specs " - f"call={_debug_grouped_gemm_partition_count} " - f"active_group_spec={active_group_spec} " - f"active_group_count={active_group_count} " - f"rhs_group_count={rhs_group_count} " - f"fsdp_axis={fsdp_axis!r} " - f"fsdp_mesh_size={mesh.shape.get(fsdp_axis) if fsdp_axis is not None else None} " - f"rhs_data_spec_before={rhs_data_spec} " - f"rhs_scale_spec_before={rhs_scale_spec} " - f"rhs_is_fsdp_shared_group_set={rhs_is_fsdp_shared_group_set} " - f"gather_rhs_fsdp={gather_rhs_fsdp}", - file=sys.stderr, - flush=True, - ) - - if gather_rhs_fsdp: - rhs_data_spec = strip_axis_from_spec(rhs_data_spec, fsdp_axis) - rhs_scale_spec = strip_axis_from_spec(rhs_scale_spec, fsdp_axis) - bias_spec = strip_axis_from_spec(bias_spec, fsdp_axis) - - reducible_axes = tuple( - axis for axis in (gsr.dp_resource, gsr.fsdp_resource) if axis is not None - ) - reduce_axis = common_spec_axis(lhs_data_spec, rhs_data_spec, reducible_axes) - # A common DP/FSDP group axis represents independent groups, not a - # partitioned contraction. Keep reductions for genuinely sharded - # contracting dimensions (including grouped wgrad), but not here. - if reduce_axis is not None and spec_contains_axis(active_group_spec, reduce_axis): - reduce_axis = None - if reduce_axis is not None and gather_rhs_fsdp: - reduce_axis = None - - if result_infos: - original_out_spec = get_padded_spec(result_infos[0]) - out_spec = filter_spec_axes(original_out_spec, allowed_axes) - else: - original_out_spec = None - out_spec = (None,) * (len(out_shape) if out_shape is not None else 1) - - if rhs_is_ragged and lhs_is_trans is not None and lhs_axis_boundary is not None: - lhs_non_contracting_dims = ( - range(lhs_axis_boundary, len(lhs_data_spec)) - if lhs_is_trans - else range(0, lhs_axis_boundary) - ) - lhs_data_spec = list(lhs_data_spec) - for out_idx, lhs_dim in enumerate(lhs_non_contracting_dims, start=1): - if out_idx < len(out_spec): - lhs_data_spec[lhs_dim] = merge_axis_specs( - lhs_data_spec[lhs_dim], out_spec[out_idx] - ) - lhs_data_spec = tuple(lhs_data_spec) - - final_arg_specs = ( - lhs_data_spec, - lhs_scale_spec, - rhs_data_spec, - rhs_scale_spec, - bias_spec, - lhs_first_dims_spec, - lhs_last_dims_spec, - rhs_first_dims_spec, - rhs_last_dims_spec, - out_first_dims_spec, - out_last_dims_spec, - additional_arg_0_spec, - additional_arg_1_spec, - ) - for arg_name, original_spec, partition_spec in zip( - ( - "lhs_data", - "lhs_scale_inv", - "rhs_data", - "rhs_scale_inv", - "bias", - "lhs_first_dims", - "lhs_last_dims", - "rhs_first_dims", - "rhs_last_dims", - "out_first_dims", - "out_last_dims", - "additional_arg_0", - "additional_arg_1", - ), - original_arg_specs, - final_arg_specs, - ): - _warn_if_axes_ignored(arg_name, original_spec, partition_spec) - if original_out_spec is not None: - _warn_if_axes_ignored("output", original_out_spec, out_spec) - - return ( - final_arg_specs, - out_spec, - reduce_axis, - ) - - @staticmethod - def partition( - lhs_is_trans, - rhs_is_trans, - scaling_mode, - out_dtype, - has_bias, - use_async_d2h_group_sizes, - use_v2_ffi, - lhs_axis_boundary, - rhs_axis_boundary, - out_shape, - lhs_left_size, - lhs_right_size, - rhs_left_size, - rhs_right_size, - mesh, - arg_infos, - result_infos, - ): - arg_specs, out_spec, reduce_axis = GroupedGemmPrimitive._parse_partition_specs( - mesh, - arg_infos, - result_infos, - out_shape, - lhs_is_trans=lhs_is_trans, - lhs_axis_boundary=lhs_axis_boundary, - ) - arg_shardings = tuple(NamedSharding(mesh, PartitionSpec(*spec)) for spec in arg_specs) - out_sharding = (NamedSharding(mesh, PartitionSpec(*out_spec)),) - local_out_shape = local_shape_from_spec(out_shape, out_spec, mesh) - local_lhs_left_size, local_lhs_right_size = local_2d_sizes_from_spec( - arg_infos[0].shape, - arg_specs[0], - lhs_axis_boundary, - lhs_left_size, - lhs_right_size, - mesh, - ) - local_rhs_left_size, local_rhs_right_size = local_2d_sizes_from_spec( - arg_infos[2].shape, - arg_specs[2], - rhs_axis_boundary, - rhs_left_size, - rhs_right_size, - mesh, - ) - - def sharded_impl( - lhs_data, - lhs_scale_inv, - rhs_data, - rhs_scale_inv, - bias, - lhs_first_dims, - lhs_last_dims, - rhs_first_dims, - rhs_last_dims, - out_first_dims, - out_last_dims, - additional_arg_0, - additional_arg_1, - ): - # Grouped quantization may expose rank-3 scale carriers so Shardy - # can gather expert and FSDP axes independently. At this point the - # requested local shardings/collectives have already been applied; - # the FFI consumes the same contiguous pre-swizzled bytes as 1D. - if lhs_scale_inv.ndim > 2: - lhs_scale_inv = lhs_scale_inv.reshape(-1) - if rhs_scale_inv.ndim > 2: - rhs_scale_inv = rhs_scale_inv.reshape(-1) - (out,) = GroupedGemmPrimitive.impl( - lhs_data, - lhs_scale_inv, - rhs_data, - rhs_scale_inv, - bias, - lhs_first_dims, - lhs_last_dims, - rhs_first_dims, - rhs_last_dims, - out_first_dims, - out_last_dims, - additional_arg_0, - additional_arg_1, - lhs_is_trans=lhs_is_trans, - rhs_is_trans=rhs_is_trans, - scaling_mode=scaling_mode, - out_dtype=out_dtype, - has_bias=has_bias, - use_async_d2h_group_sizes=use_async_d2h_group_sizes, - use_v2_ffi=use_v2_ffi, - lhs_axis_boundary=lhs_axis_boundary, - rhs_axis_boundary=rhs_axis_boundary, - out_shape=local_out_shape, - lhs_left_size=local_lhs_left_size, - lhs_right_size=local_lhs_right_size, - rhs_left_size=local_rhs_left_size, - rhs_right_size=local_rhs_right_size, - ) - - if reduce_axis is not None: - if is_all_reduce_in_float32(): - out = jax.lax.psum(out.astype(jnp.float32), reduce_axis).astype(out_dtype) - else: - out = jax.lax.psum(out, reduce_axis) - return (out,) - - return mesh, sharded_impl, out_sharding, arg_shardings - - @staticmethod - def shardy_sharding_rule( - lhs_is_trans, - rhs_is_trans, - scaling_mode, - out_dtype, - has_bias, - use_async_d2h_group_sizes, - use_v2_ffi, - lhs_axis_boundary, - rhs_axis_boundary, - out_shape, - lhs_left_size, - lhs_right_size, - rhs_left_size, - rhs_right_size, - mesh, - operand_types, - result_types, - ): - del ( - lhs_is_trans, - rhs_is_trans, - scaling_mode, - out_dtype, - has_bias, - use_async_d2h_group_sizes, - use_v2_ffi, - lhs_axis_boundary, - rhs_axis_boundary, - out_shape, - lhs_left_size, - lhs_right_size, - rhs_left_size, - rhs_right_size, - mesh, - ) - - prefix = "GroupedGemm" - - def spec_for(name, rank): - if rank == 0: - return () - return tuple(f"{prefix}_{name}_{i}" for i in range(rank)) - - operand_mappings = tuple( - spec_for(f"arg{i}", len(operand_type.shape)) - for i, operand_type in enumerate(operand_types) - ) - result_mappings = tuple( - spec_for(f"out{i}", len(result_type.shape)) - for i, result_type in enumerate(result_types) - ) - return SdyShardingRule( - operand_mappings=operand_mappings, - result_mappings=result_mappings, - ) - register_primitive(GroupedGemmPrimitive) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index a888c43ded..7138cfcf40 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -6,13 +6,12 @@ from functools import reduce from typing import Tuple, Optional, Union import math -import warnings import jax import jax.numpy as jnp from jax import dtypes, ffi -from jax.experimental.custom_partitioning import SdyShardingRule, BATCHING, CompoundFactor +from jax.experimental.custom_partitioning import SdyShardingRule, BATCHING from jax.sharding import PartitionSpec import transformer_engine_jax @@ -32,15 +31,7 @@ from ..sharding import ( all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp, - axis_spec_size, - filter_spec_axes, get_num_devices_in_mesh, - global_mesh_resource, - lax_paral_op, - local_shape_from_spec, - merge_axis_specs, - spec_axes, - supported_grouped_partition_axes, ) from ..quantize import ( ScaledTensor2x, @@ -61,90 +52,6 @@ __all__ = ["quantize", "quantize_dbias", "grouped_quantize", "grouped_dbias"] -def _flat_data_spec(input_spec): - return (merge_axis_specs(*input_spec),) - - -def _grouped_data_spec(input_spec, flatten_axis): - return _contiguous_flat_input_spec(input_spec, flatten_axis) - - -def _uniform_mxfp8_scale_carrier_shapes(x_shape, n_groups, flatten_axis): - """Return sharding-friendly scale carriers for uniform 3D MXFP8 kernels.""" - flatten_axis = _normalize_flatten_axis(flatten_axis, len(x_shape)) - if len(x_shape) != 3 or flatten_axis != 2 or x_shape[0] != n_groups: - return None - - group_rows = x_shape[1] - columns = x_shape[2] - if group_rows % 128 != 0 or columns % 128 != 0: - return None - - # The fused MXFP8 swizzle groups scales in 128-row/column tiles. These - # carriers expose group and column ownership as separate dimensions while - # preserving the exact contiguous byte order consumed by grouped GEMM. - rowwise_shape = (n_groups, group_rows // 128, columns * 4) - colwise_shape = (n_groups, columns // 128, group_rows * 4) - return rowwise_shape, colwise_shape - - -def _uniform_mxfp8_scale_carrier_specs(x_spec, flatten_axis): - """Map a uniform 3D kernel's data sharding onto its scale carriers.""" - flatten_axis = _normalize_flatten_axis(flatten_axis, len(x_spec)) - assert len(x_spec) == 3 and flatten_axis == 2 - group_spec = x_spec[0] - row_spec = merge_axis_specs(*x_spec[1:flatten_axis]) - column_spec = merge_axis_specs(*x_spec[flatten_axis:]) - return ( - (group_spec, row_spec, column_spec), - (group_spec, column_spec, row_spec), - ) - - -def _normalize_flatten_axis(flatten_axis, ndim): - return flatten_axis + ndim if flatten_axis < 0 else flatten_axis - - -def _contiguous_flat_input_spec(input_spec, flatten_axis): - flatten_axis = _normalize_flatten_axis(flatten_axis, len(input_spec)) - if flatten_axis <= 0 or len(input_spec) == 0: - return (None,) * len(input_spec) - return (*input_spec[:flatten_axis], *((None,) * (len(input_spec) - flatten_axis))) - - -def _warn_if_axes_ignored(arg_name, original_spec, partition_spec): - ignored_axes = tuple( - axis for axis in spec_axes(original_spec) if axis not in spec_axes(partition_spec) - ) - if ignored_axes: - warnings.warn( - "Grouped quantize custom partitioning will ignore/replicate sharding " - f"axes {ignored_axes} from {arg_name}; only supported packed grouped " - "data axes are preserved.", - RuntimeWarning, - stacklevel=3, - ) - - -def _pad_or_slice_to_shape(x, target_shape): - if target_shape is None or x.shape == target_shape: - return x - target_size = math.prod(target_shape) - current_size = math.prod(x.shape) - x = x.reshape(-1) - if current_size > target_size: - return x[:target_size].reshape(target_shape) - return jnp.pad(x, (0, target_size - current_size)).reshape(target_shape) - - -def _all_reduce_grouped_amax_along_dp_fsdp(amax, mesh): - gsr = global_mesh_resource() - for axis in (gsr.dp_resource, gsr.fsdp_resource): - if axis is not None and axis in mesh.axis_names: - amax = lax_paral_op(amax, jax.lax.pmax, axis, mesh) - return amax - - class BaseDBiasQuantizePrimitive(BasePrimitive): """ Cast Primitive wrapping nvte_quantize and nvte_quantize_dbias @@ -1096,8 +1003,7 @@ class GroupedQuantizePrimitive(BasePrimitive): 5, 6, 7, - 8, - ) # out_dtype, scaling_mode, q_layout, flatten_axis, scale_dtype, uniform_groups + ) # out_dtype, scaling_mode, q_layout, flatten_axis, scale_dtype inner_primitive = None outer_primitive = None @@ -1160,16 +1066,13 @@ def abstract( q_layout, flatten_axis, scale_dtype, - uniform_groups, ): """ te_dbias_quantize_p abstract """ dtype = dtypes.canonicalize_dtype(x_aval.dtype) assert dtype in [jnp.float32, jnp.float16, jnp.bfloat16] - # Preserve logical rank for custom partitioning. The FFI still consumes - # the same contiguous buffer, while avoiding oversized flat Shardy dims. - out_shape = x_aval.shape + out_shape = math.prod(x_aval.shape) # TODO(Phuong): can scale_aval be None? assert scale_aval is None or scale_aval.dtype == jnp.float32 @@ -1178,22 +1081,14 @@ def abstract( f" be one of {ScalingMode(scaling_mode).get_compatible_q_dtypes()}" ) - scale_carrier_shapes = None - if uniform_groups and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING: - scale_carrier_shapes = _uniform_mxfp8_scale_carrier_shapes( - x_aval.shape, group_sizes_aval.size, flatten_axis - ) - if scale_carrier_shapes is None: - rowwise_scale_inv_shape, colwise_scale_inv_shape = ScalingMode( - scaling_mode - ).get_grouped_scale_shape_2x( - x_aval.shape, - group_sizes_aval.size, - is_padded=True, - flatten_axis=flatten_axis, - ) - else: - rowwise_scale_inv_shape, colwise_scale_inv_shape = scale_carrier_shapes + rowwise_scale_inv_shape, colwise_scale_inv_shape = ScalingMode( + scaling_mode + ).get_grouped_scale_shape_2x( + x_aval.shape, + group_sizes_aval.size, + is_padded=True, + flatten_axis=flatten_axis, + ) if q_layout.has_rowwise: rowwise_out_shape = out_shape @@ -1273,12 +1168,11 @@ def lowering( q_layout, flatten_axis, scale_dtype, - uniform_groups, ): """ te_dbias_quantize_p lowering rules """ - del out_dtype, scale_dtype, uniform_groups + del out_dtype, scale_dtype x_aval, scale_aval, group_sizes_aval = ctx.avals_in assert x_aval.dtype in [jnp.float32, jnp.float16, jnp.bfloat16] assert scale_aval.dtype == jnp.float32 @@ -1318,7 +1212,6 @@ def impl( q_layout, flatten_axis, scale_dtype, - uniform_groups, ): """ te_dbias_quantize_p implementation @@ -1340,296 +1233,9 @@ def impl( q_layout=q_layout, flatten_axis=flatten_axis, scale_dtype=scale_dtype, - uniform_groups=uniform_groups, ) return rowwise_out, colwise_out, rowwise_scale_inv, colwise_scale_inv, updated_amax - @staticmethod - def _parse_partition_specs( - scaling_mode, q_layout, flatten_axis, uniform_groups, mesh, arg_infos - ): - allowed_axes = supported_grouped_partition_axes(mesh) - original_x_spec = get_padded_spec(arg_infos[0]) - x_spec = filter_spec_axes(original_x_spec, allowed_axes) - use_scale_carrier = ( - uniform_groups - and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING - and _uniform_mxfp8_scale_carrier_shapes( - arg_infos[0].shape, arg_infos[2].size, flatten_axis - ) - is not None - ) - if use_scale_carrier: - x_spec = list(x_spec) - for dim in (1, 2): - local_dim = arg_infos[0].shape[dim] // axis_spec_size(x_spec[dim], mesh) - if local_dim % 128 != 0: - x_spec[dim] = None - x_spec = tuple(x_spec) - else: - x_spec = _contiguous_flat_input_spec(x_spec, flatten_axis) - _warn_if_axes_ignored("x", original_x_spec, x_spec) - - original_group_spec = get_padded_spec(arg_infos[2]) - group_spec = filter_spec_axes(original_group_spec, allowed_axes) - if group_spec == (None,) and len(x_spec) > 0: - group_spec = (x_spec[0],) - _warn_if_axes_ignored("group_sizes", original_group_spec, group_spec) - flat_spec = _flat_data_spec(x_spec) - data_spec = tuple(x_spec) if use_scale_carrier else _grouped_data_spec(x_spec, flatten_axis) - replicated_spec = (None,) - - rowwise_out_spec = data_spec if q_layout.has_rowwise else replicated_spec - colwise_out_spec = data_spec if q_layout.has_colwise else replicated_spec - - rowwise_scale_inv_spec = replicated_spec - colwise_scale_inv_spec = replicated_spec - if ScalingMode(scaling_mode).is_block_scaling: - if use_scale_carrier: - rowwise_carrier_spec, colwise_carrier_spec = _uniform_mxfp8_scale_carrier_specs( - x_spec, flatten_axis - ) - rowwise_scale_inv_spec = ( - rowwise_carrier_spec if q_layout.has_rowwise else replicated_spec - ) - colwise_scale_inv_spec = ( - colwise_carrier_spec if q_layout.has_colwise else replicated_spec - ) - else: - rowwise_scale_inv_spec = flat_spec if q_layout.has_rowwise else replicated_spec - colwise_scale_inv_spec = flat_spec if q_layout.has_colwise else replicated_spec - elif ScalingMode(scaling_mode).is_tensor_scaling(): - rowwise_scale_inv_spec = group_spec if q_layout.has_rowwise else replicated_spec - colwise_scale_inv_spec = group_spec if q_layout.has_colwise else replicated_spec - - updated_amax_spec = group_spec - return ( - x_spec, - group_spec, - ( - rowwise_out_spec, - colwise_out_spec, - rowwise_scale_inv_spec, - colwise_scale_inv_spec, - updated_amax_spec, - ), - ) - - @staticmethod - def partition( - out_dtype, - scaling_mode, - q_layout, - flatten_axis, - scale_dtype, - uniform_groups, - mesh, - arg_infos, - result_infos, - ): - x_spec, group_spec, out_specs = GroupedQuantizePrimitive._parse_partition_specs( - scaling_mode, q_layout, flatten_axis, uniform_groups, mesh, arg_infos - ) - use_scale_carrier = ( - uniform_groups - and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING - and _uniform_mxfp8_scale_carrier_shapes( - arg_infos[0].shape, arg_infos[2].size, flatten_axis - ) - is not None - ) - local_out_shapes = ( - tuple( - local_shape_from_spec(info.shape, spec, mesh) - for info, spec in zip(result_infos, out_specs) - ) - if result_infos - else (None,) * len(out_specs) - ) - - arg_shardings = ( - NamedSharding(mesh, PartitionSpec(*x_spec)), - NamedSharding(mesh, PartitionSpec(*group_spec)), - NamedSharding(mesh, PartitionSpec(*group_spec)), - ) - out_shardings = tuple(NamedSharding(mesh, PartitionSpec(*spec)) for spec in out_specs) - - def sharded_impl(x, scale, group_sizes): - ( - rowwise_out, - colwise_out, - rowwise_scale_inv, - colwise_scale_inv, - updated_amax, - ) = GroupedQuantizePrimitive.impl( - x, - scale, - group_sizes, - out_dtype=out_dtype, - scaling_mode=scaling_mode, - q_layout=q_layout, - flatten_axis=flatten_axis, - scale_dtype=scale_dtype, - # The rank-3 carrier is a partitioning-only representation. Keep - # the inner FFI contract flat and reshape its local scale buffers - # after quantization. - uniform_groups=False if use_scale_carrier else uniform_groups, - ) - if ScalingMode(scaling_mode).is_block_scaling: - rowwise_scale_inv = _pad_or_slice_to_shape(rowwise_scale_inv, local_out_shapes[2]) - colwise_scale_inv = _pad_or_slice_to_shape(colwise_scale_inv, local_out_shapes[3]) - if ScalingMode(scaling_mode).is_tensor_scaling(): - updated_amax = _all_reduce_grouped_amax_along_dp_fsdp(updated_amax, mesh) - return ( - rowwise_out, - colwise_out, - rowwise_scale_inv, - colwise_scale_inv, - updated_amax, - ) - - return mesh, sharded_impl, out_shardings, arg_shardings - - @staticmethod - def shardy_sharding_rule( - out_dtype, - scaling_mode, - q_layout, - flatten_axis, - scale_dtype, - uniform_groups, - mesh, - value_types, - result_types, - ): - del out_dtype, scale_dtype, mesh - - prefix = "GroupedQuantize" - input_shape = value_types[0].shape - input_spec = tuple(f"{prefix}_x_{i}" for i in range(len(input_shape))) - normalized_flatten_axis = _normalize_flatten_axis(flatten_axis, len(input_spec)) - use_scale_carrier = ( - uniform_groups - and ScalingMode(scaling_mode) == ScalingMode.MXFP8_1D_SCALING - and _uniform_mxfp8_scale_carrier_shapes( - input_shape, value_types[2].shape[0], flatten_axis - ) - is not None - ) - if use_scale_carrier: - group_factor = f"{prefix}_kernel_group" - row_factor = f"{prefix}_kernel_row_tiles" - column_factor = f"{prefix}_kernel_column_tiles" - row_block_factor = f"{prefix}_kernel_row_block" - column_block_factor = f"{prefix}_kernel_column_block" - row_pack_factor = f"{prefix}_kernel_row_pack" - column_pack_factor = f"{prefix}_kernel_column_pack" - row_unit_factor = f"{prefix}_kernel_row_unit" - column_unit_factor = f"{prefix}_kernel_column_unit" - - row_tiles = input_shape[1] // 128 - column_tiles = input_shape[2] // 128 - factor_sizes = {} - if row_tiles == 1: - row_input_factor = f"{prefix}_kernel_row" - row_scale_factor = row_unit_factor - row_packed_factors = (row_input_factor,) - else: - row_input_factor = CompoundFactor(row_factor, row_block_factor) - row_scale_factor = row_factor - row_packed_factors = (row_factor, row_block_factor) - factor_sizes[row_block_factor] = 128 - if column_tiles == 1: - column_input_factor = f"{prefix}_kernel_column" - column_scale_factor = column_unit_factor - column_packed_factors = (column_input_factor,) - else: - column_input_factor = CompoundFactor(column_factor, column_block_factor) - column_scale_factor = column_factor - column_packed_factors = (column_factor, column_block_factor) - factor_sizes[column_block_factor] = 128 - - input_spec = ( - group_factor, - row_input_factor, - column_input_factor, - ) - data_spec = input_spec - rowwise_scale_spec = ( - ( - group_factor, - row_scale_factor, - CompoundFactor(*column_packed_factors, row_pack_factor), - ) - if q_layout.has_rowwise - else (BATCHING + f"{prefix}_scalar",) - ) - colwise_scale_spec = ( - ( - group_factor, - column_scale_factor, - CompoundFactor(*row_packed_factors, column_pack_factor), - ) - if q_layout.has_colwise - else (BATCHING + f"{prefix}_scalar",) - ) - if q_layout.has_rowwise: - factor_sizes[row_pack_factor] = 4 - elif row_tiles > 1: - factor_sizes[row_factor] = input_shape[1] // 128 - if q_layout.has_colwise: - factor_sizes[column_pack_factor] = 4 - elif column_tiles > 1: - factor_sizes[column_factor] = input_shape[2] // 128 - - scalar_spec = (BATCHING + f"{prefix}_scalar",) - rowwise_out_spec = data_spec if q_layout.has_rowwise else scalar_spec - colwise_out_spec = data_spec if q_layout.has_colwise else scalar_spec - group_spec = (BATCHING + f"{prefix}_group",) - return SdyShardingRule( - operand_mappings=(input_spec, group_spec, group_spec), - result_mappings=( - rowwise_out_spec, - colwise_out_spec, - rowwise_scale_spec, - colwise_scale_spec, - group_spec, - ), - **factor_sizes, - ) - - data_spec = tuple( - input_spec[i] if i < normalized_flatten_axis else f"{prefix}_data_{i}" - for i in range(len(input_spec)) - ) - flat_spec = (f"{prefix}_flat",) - group_spec = (BATCHING + f"{prefix}_group",) - scalar_spec = (BATCHING + f"{prefix}_scalar",) - - rowwise_out_spec = data_spec if q_layout.has_rowwise else scalar_spec - colwise_out_spec = data_spec if q_layout.has_colwise else scalar_spec - - if ScalingMode(scaling_mode).is_block_scaling: - rowwise_scale_spec = flat_spec if q_layout.has_rowwise else scalar_spec - colwise_scale_spec = flat_spec if q_layout.has_colwise else scalar_spec - elif ScalingMode(scaling_mode).is_tensor_scaling(): - rowwise_scale_spec = group_spec if q_layout.has_rowwise else scalar_spec - colwise_scale_spec = group_spec if q_layout.has_colwise else scalar_spec - else: - rowwise_scale_spec = scalar_spec - colwise_scale_spec = scalar_spec - - return SdyShardingRule( - operand_mappings=(input_spec, group_spec, group_spec), - result_mappings=( - rowwise_out_spec, - colwise_out_spec, - rowwise_scale_spec, - colwise_scale_spec, - group_spec, - ), - ) - register_primitive(GroupedQuantizePrimitive) @@ -1639,7 +1245,6 @@ def grouped_quantize( quantizer: GroupedQuantizer, group_sizes: jnp.ndarray = None, flatten_axis: int = -1, - ragged_scale_sharding: NamedSharding | None = None, ) -> Union[GroupedScaledTensor1x, GroupedNoScaleTensor]: """Quantize a tensor in grouped manner. @@ -1652,8 +1257,6 @@ def grouped_quantize( quantizer: The quantizer to use for quantization group_sizes: Array of ints containing the size of each group (default: None) flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1) - ragged_scale_sharding: Data sharding to preserve for flat ragged scale buffers. - Its first partition-spec entry is applied to each scale buffer. Returns: A GroupedScaledTensor1x containing the quantized data @@ -1683,7 +1286,6 @@ def grouped_quantize( ), f"Only flatten_axis = -1 is supported for now, got {flatten_axis}" ragged_first_dims = group_sizes # None if no explicit group_sizes (kernel case) - uniform_groups = group_sizes is None if group_sizes is None: group_sizes = jnp.ones(x.shape[0], dtype=jnp.int32) @@ -1731,7 +1333,6 @@ def grouped_quantize( q_layout=q_layout, flatten_axis=flatten_axis, scale_dtype=quantizer.get_scale_dtype(), - uniform_groups=uniform_groups, ) # For DelayedScaling2x and CurrentScaling2x, the scale buffer @@ -1739,23 +1340,6 @@ def grouped_quantize( if is_tensor_scaling and quantizer.q_layout.is_rowwise_colwise or apply_colwise_war: colwise_scale_inv = rowwise_scale_inv - if ragged_scale_sharding is not None: - if ragged_first_dims is None: - raise ValueError("ragged_scale_sharding requires explicit group_sizes") - if not quantizer.scaling_mode.is_block_scaling: - raise ValueError("ragged_scale_sharding requires block scaling") - data_spec = ragged_scale_sharding.spec - if len(data_spec) == 0: - raise ValueError("ragged_scale_sharding must have a token/group dimension") - scale_sharding = NamedSharding( - ragged_scale_sharding.mesh, - PartitionSpec(data_spec[0]), - ) - if q_layout.has_rowwise: - rowwise_scale_inv = jax.lax.with_sharding_constraint(rowwise_scale_inv, scale_sharding) - if q_layout.has_colwise: - colwise_scale_inv = jax.lax.with_sharding_constraint(colwise_scale_inv, scale_sharding) - # TODO(Phuong): store the whole updated_amax in the grouped_quantize instead? if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: for i, quantizer_i in enumerate(quantizer.quantizers): diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index b60810f5eb..f8c30ffccb 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -26,6 +26,34 @@ ) +def _all_gather_kernel(kernel, mesh_axis, axis_idx): + assert mesh_axis is not None + assert 0 < axis_idx < len(kernel.shape) + + # TODO(Ming Hunag): Add a condition branch for with/without shmap. + kernel_shape = kernel.shape + kernel_whole_shape = (*kernel_shape[:axis_idx], -1, *kernel_shape[axis_idx + 1 :]) + global_kernel = jax.lax.all_gather(kernel, mesh_axis, axis=axis_idx) + global_kernel = global_kernel.reshape(*kernel_whole_shape) + return global_kernel + + +def _psum_scatter_kernel(kernel, scattered_kernel_shape, mesh_axis, axis_idx): + assert mesh_axis is not None + assert 0 < axis_idx < len(scattered_kernel_shape) + + # TODO(Ming Hunag): Add a condition branch for with/without shmap. + kernel = kernel.reshape( + *scattered_kernel_shape[:axis_idx], + -1, + scattered_kernel_shape[axis_idx], + *scattered_kernel_shape[axis_idx + 1 :], + ) + kernel = jax.lax.psum_scatter(kernel, mesh_axis, scatter_dimension=axis_idx) + kernel = kernel.reshape(scattered_kernel_shape) + return kernel + + def dense( x: jnp.ndarray, kernel: jnp.ndarray, @@ -297,6 +325,7 @@ def grouped_dense( preferred_element_type: jnp.dtype = None, group_offset: jnp.array = None, quantizer_set: QuantizerSet = noop_quantizer_set, + kernel_fsdp_info: Tuple[str, int] = (None, -1), ): """ Perform grouped dense (linear) layer transformation with optional quantization. @@ -312,15 +341,14 @@ def grouped_dense( preferred_element_type: Preferred data type for the output tensor group_offset: 1D array containing offsets for each group (not yet implemented) quantizer_set: Set of quantizers for FP8 quantization of the input and output + kernel_fsdp_info: A tuple containing FSDP-related information for a weight matrix + represented in the format (str, int). The first element is the + FSDP mesh axis, and the second element is the dimension along + which the weight is sharded. Returns: A jnp.ndarray containing the result of the grouped linear operation """ - x_contracting_dims, kernel_contracting_dims = contracting_dims - x_contracting_dims = tex.sanitize_dims(x.ndim, x_contracting_dims) - kernel_contracting_dims = tex.sanitize_dims(kernel.ndim, kernel_contracting_dims) - contracting_dims = (x_contracting_dims, kernel_contracting_dims) - output = _grouped_dense( x, kernel, @@ -331,11 +359,12 @@ def grouped_dense( preferred_element_type, group_offset, quantizer_set, + kernel_fsdp_info, ) return output -@partial(jax.custom_vjp, nondiff_argnums=(3, 5, 6, 7)) +@partial(jax.custom_vjp, nondiff_argnums=(3, 5, 6, 7, 9)) def _grouped_dense( x, kernel, @@ -346,6 +375,7 @@ def _grouped_dense( preferred_element_type, group_offset, quantizer_set, + kernel_fsdp_info, ): output, _ = _grouped_dense_fwd_rule( x, @@ -357,6 +387,7 @@ def _grouped_dense( preferred_element_type, group_offset, quantizer_set, + kernel_fsdp_info, ) return output @@ -371,11 +402,16 @@ def _grouped_dense_fwd_rule( preferred_element_type, group_offset, quantizer_set, + kernel_fsdp_info, ): use_bias = bias is not None - x_contracting_dims, k_contracting_dims = contracting_dims + kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info + kernel_fsdp_enabled = kernel_fsdp_mesh_axis is not None + assert not kernel_fsdp_enabled, "FSDP sharding for grouped_dense is not supported yet." + del kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx, kernel_fsdp_info, kernel_fsdp_enabled + x_contracting_dims, k_contracting_dims = contracting_dims flatten_axis_x = -len(x_contracting_dims) flatten_axis_k = len(k_contracting_dims) - len(kernel.shape) + 1 # +1 for G axis @@ -440,8 +476,12 @@ def _grouped_dense_fwd_rule( def _grouped_dense_bwd_rule( - contracting_dims, precision, preferred_element_type, group_offset, ctx, grad + contracting_dims, precision, preferred_element_type, group_offset, kernel_fsdp_info, ctx, grad ): + kernel_fsdp_mesh_axis, _ = kernel_fsdp_info + kernel_fsdp_enabled = kernel_fsdp_mesh_axis is not None + assert not kernel_fsdp_enabled, "FSDP sharding for grouped_dense is not supported yet." + fwd_x_contracting_dims, fwd_k_contracting_dims = contracting_dims ( diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 14783ecbe2..17c9a242f0 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -1471,7 +1471,7 @@ def te_grouped_dot_general(generate_quantizer_set, x, kernel, group_sizes, **kwa x, kernel, group_sizes=group_sizes, - contracting_dims=((-1,), (1,)), + contracting_dims=((1,), (1,)), quantizer_set=quantizer_set, ) return out diff --git a/transformer_engine/jax/quantize/dequantizer.py b/transformer_engine/jax/quantize/dequantizer.py index 269c824089..ca44c2e4af 100644 --- a/transformer_engine/jax/quantize/dequantizer.py +++ b/transformer_engine/jax/quantize/dequantizer.py @@ -303,13 +303,8 @@ def _grouped_dequantize(grouped_scaled_tensor): Returns: List of dequantized tensors for each group """ - # Group offsets are scalar-element offsets even though grouped quantization - # preserves the logical N-D carrier shape to avoid oversized 1-D dimensions. - data = grouped_scaled_tensor.data.reshape(-1) - # Uniform grouped kernels may use a multidimensional carrier so expert and - # FSDP ownership survive custom partitioning. Group offsets still address - # the same contiguous pre-swizzled byte stream. - scale_inv = grouped_scaled_tensor.scale_inv.reshape(-1) + data = grouped_scaled_tensor.data + scale_inv = grouped_scaled_tensor.scale_inv group_sizes = ( grouped_scaled_tensor.first_dims if grouped_scaled_tensor.first_dims is not None diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index 95bb8ae51c..c5ad0451fd 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -8,10 +8,9 @@ both single-scale (1x) and double-scale (2x) quantization schemes. It supports rowwise and colwise quantization modes with proper scaling and dequantization. """ -from abc import ABC, abstractmethod from dataclasses import dataclass -import math 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 @@ -428,14 +427,8 @@ def group_sizes(self) -> jnp.ndarray: return jnp.ones((self.original_shape[0],), dtype=jnp.int32) def __post_init__(self): - assert self.scale_inv.ndim in (1, 3), ( - "Grouped scale_inv must be flat or use the uniform-kernel 3D carrier, " - f"got shape {self.scale_inv.shape}" - ) - assert self.data.size == math.prod(self.original_shape), ( - f"Quantized data has {self.data.size} elements, expected " - f"{math.prod(self.original_shape)} for original shape {self.original_shape}" - ) + assert self.scale_inv.ndim == 1, "Only support flattened scale_inv" + assert self.data.ndim == 1, "Only support flattened data" assert self.flatten_axis > 0 data_ndim = len(self.original_shape) @@ -453,23 +446,13 @@ def __post_init__(self): else: num_groups = self.original_shape[0] - if self.scale_inv.ndim == 3: - assert self.scaling_mode == ScalingMode.MXFP8_1D_SCALING - assert len(self.original_shape) == 3 and self.flatten_axis == 2 - groups, rows, columns = self.original_shape - expected_scale_shape = ( - (groups, columns // 128, rows * 4) - if self.is_colwise - else (groups, rows // 128, columns * 4) - ) - else: - expected_scale_shape = self.scaling_mode.get_grouped_scale_shape( - self.original_shape, - num_groups, - self.is_colwise, - is_padded=True, - flatten_axis=self.flatten_axis, - ) + expected_scale_shape = self.scaling_mode.get_grouped_scale_shape( + self.original_shape, + num_groups, + self.is_colwise, + is_padded=True, + flatten_axis=self.flatten_axis, + ) assert self.scale_inv.shape == expected_scale_shape, ( f"Unexpected scale_inv shape! \nExpect {expected_scale_shape} for padded" diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index b2319073b3..2e8e611fa3 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -12,7 +12,6 @@ from contextlib import contextmanager from dataclasses import dataclass from typing import Callable, Optional -import math import warnings import jax @@ -134,8 +133,8 @@ def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec): """ A wrapper function to jax.lax.with_sharding_constraint 1. Does nothing if mesh is empty. - 2. Keeps only auto axes in pspec. - 3. Returns x unchanged if no auto axes remain. + 2. If all mesh axes are manual axes, replaces pspec with all Nones. + 3. Otherwise, strips only the manual axes. """ if pspec is None: return x @@ -144,21 +143,22 @@ def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec): if mesh.empty: return x - # with_sharding_constraint can only refer to auto axes. Explicit axes are - # already fixed by the active mesh, and manual axes are managed by shard_map. - abstract_mesh = get_abstract_mesh() - auto_axis_names = set(abstract_mesh.auto_axes) + # We want to exclude the axes that already used by shard_map and shard_map + # only sets those in the abstract_mesh, not the physical one + manual_axis_names = get_abstract_mesh().manual_axes # Multiple mesh axes can be mapped to a single shape axis, so we need to unpack and process tuples here too - def filter_non_auto_axes(name_or_tuple): + def filter_manual_axes(name_or_tuple): if isinstance(name_or_tuple, tuple): - out = tuple(n for n in name_or_tuple if n in auto_axis_names) + out = tuple(n for n in name_or_tuple if n not in manual_axis_names) if len(out) == 0: return None return out - return name_or_tuple if name_or_tuple in auto_axis_names else None + if name_or_tuple in manual_axis_names: + return None + return name_or_tuple - cleaned_axis_names = tuple(filter_non_auto_axes(name_or_tuple) for name_or_tuple in pspec) + cleaned_axis_names = tuple(filter_manual_axes(name_or_tuple) for name_or_tuple in pspec) if cleaned_axis_names == (None,) * len(cleaned_axis_names): return x @@ -233,155 +233,6 @@ def get_padded_spec(spec, ndim): return spec + (None,) * (ndim - len(spec)) -def spec_axes(spec): - """Return unique non-None mesh axes used by a PartitionSpec-like tuple.""" - axes = [] - for axis_spec in spec: - if axis_spec is None: - continue - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - for axis in axis_tuple: - if axis is not None and axis not in axes: - axes.append(axis) - return axes - - -def axis_spec_contains(axis_spec, axis): - """Return whether one dimension's axis spec contains a mesh axis.""" - if axis is None or axis_spec is None: - return False - if isinstance(axis_spec, tuple): - return axis in axis_spec - return axis_spec == axis - - -def spec_contains_axis(spec, axis): - """Return whether a PartitionSpec-like tuple contains a mesh axis.""" - return any(axis_spec_contains(axis_spec, axis) for axis_spec in spec) - - -def filter_axis_spec(axis_spec, allowed_axes): - """Keep only allowed axes in one dimension's axis spec.""" - if axis_spec is None: - return None - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - axes = tuple(axis for axis in axis_tuple if axis in allowed_axes) - if len(axes) == 0: - return None - return axes[0] if len(axes) == 1 else axes - - -def filter_spec_axes(spec, allowed_axes): - """Keep only allowed axes in a PartitionSpec-like tuple.""" - return tuple(filter_axis_spec(axis_spec, allowed_axes) for axis_spec in spec) - - -def supported_grouped_partition_axes(mesh): - """Return mesh axes supported by grouped quantize/GEMM custom partitioning.""" - gsr = global_mesh_resource(validate=False) - return { - axis - for axis in (gsr.ep_resource, gsr.dp_resource, gsr.fsdp_resource) - if axis is not None and axis in mesh.axis_names - } - - -def strip_axis_from_axis_spec(axis_spec, axis): - """Remove one mesh axis from one dimension's axis spec.""" - if axis is None or axis_spec is None: - return axis_spec - if isinstance(axis_spec, tuple): - stripped = tuple(a for a in axis_spec if a != axis) - if len(stripped) == 0: - return None - return stripped[0] if len(stripped) == 1 else stripped - return None if axis_spec == axis else axis_spec - - -def strip_axis_from_spec(spec, axis): - """Remove one mesh axis from a PartitionSpec-like tuple.""" - return tuple(strip_axis_from_axis_spec(axis_spec, axis) for axis_spec in spec) - - -def merge_axis_specs(*axis_specs): - """Merge dimension axis specs while preserving first-seen axis order.""" - axes = [] - for axis_spec in axis_specs: - if axis_spec is None: - continue - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - for axis in axis_tuple: - if axis is not None and axis not in axes: - axes.append(axis) - if len(axes) == 0: - return None - return axes[0] if len(axes) == 1 else tuple(axes) - - -def common_spec_axis(spec_a, spec_b, allowed_axes=None): - """Return the first mesh axis that appears in both specs.""" - axes = [] - for spec in (spec_a, spec_b): - for axis in spec_axes(spec): - if axis not in axes: - axes.append(axis) - for axis in axes: - if allowed_axes is not None and axis not in allowed_axes: - continue - if spec_contains_axis(spec_a, axis) and spec_contains_axis(spec_b, axis): - return axis - return None - - -def axis_spec_size(axis_spec, mesh): - """Return the device count represented by one dimension's axis spec.""" - axis_tuple = axis_spec if isinstance(axis_spec, tuple) else (axis_spec,) - axis_size = 1 - for axis in axis_tuple: - if axis is not None: - axis_size *= mesh.shape[axis] - return axis_size - - -def spec_size(spec, mesh): - """Return the total device count represented by a PartitionSpec-like tuple.""" - axis_size = 1 - for axis_spec in spec: - axis_size *= axis_spec_size(axis_spec, mesh) - return axis_size - - -def local_shape_from_spec(global_shape, spec, mesh): - """Derive a local shape from a global shape and PartitionSpec-like tuple.""" - local_shape = [] - for dim, axis_spec in zip(global_shape, spec): - local_shape.append(dim // axis_spec_size(axis_spec, mesh)) - return tuple(local_shape) - - -def local_2d_sizes_from_spec(shape, spec, axis_boundary, left_size, right_size, mesh): - """Derive local collapsed 2D dimensions from a global shape and sharding spec.""" - if len(shape) == len(spec) and len(shape) > 1: - local_shape = local_shape_from_spec(shape, spec, mesh) - return ( - math.prod(local_shape[:axis_boundary]), - math.prod(local_shape[axis_boundary:]), - ) - - size = spec_size(spec, mesh) - if size == 1: - return left_size, right_size - if left_size % size == 0: - return left_size // size, right_size - if right_size % size == 0: - return left_size, right_size // size - raise ValueError( - "Cannot derive local 2D sizes from sharding spec. " - f"shape={shape}, spec={spec}, axis_boundary={axis_boundary}, " - f"left_size={left_size}, right_size={right_size}, spec_size={size}" - ) - - def lax_paral_op( x: jnp.array, ops: Callable, mesh_resource: str, mesh: jax.sharding.Mesh, **kwargs ): @@ -478,7 +329,6 @@ class MeshResource: tp_resource: Axis name for tensor parallelism (hidden dimension sharding), default is None tpsp_resource: Axis name for tensor sequence parallelism (hidden and sequence sharding), default is None fsdp_resource: Axis name for full-sharded data parallelism, default is None - ep_resource: Axis name for expert parallelism (expert sharding), default is None pp_resource: Axis name for pipeline parallelism (layer sharding), default is None cp_resource: Axis name for context parallelism (sequence sharding), default is None ep_resource: Axis name for expert parallelism. Dispatch input tokens @@ -493,7 +343,6 @@ class MeshResource: tp_resource: str = None tpsp_resource: str = None fsdp_resource: str = None - ep_resource: str = None pp_resource: str = None cp_resource: str = None ep_resource: str = None @@ -521,7 +370,7 @@ def global_shard_guard(resource: MeshResource): _GLOBAL_MESH_RESOURCE = old_resources -def global_mesh_resource(validate: bool = True) -> MeshResource: +def global_mesh_resource() -> MeshResource: """Get the current global mesh resource configuration. Returns: @@ -532,8 +381,7 @@ def global_mesh_resource(validate: bool = True) -> MeshResource: " context. If you are not using multiple GPUs, you can use an empty MeshResource by" " wrapping your program in 'with global_shard_guard(MeshResource()):'" ) - if validate: - _validate_mesh_resource_configuration(_GLOBAL_MESH_RESOURCE) + _validate_mesh_resource_configuration(_GLOBAL_MESH_RESOURCE) return _GLOBAL_MESH_RESOURCE diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 76d51673f0..909e5a8a9b 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -101,9 +101,13 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() -def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool: - """Check cuDNN FE min version for single-group runtime offsets.""" - return _cudnn_frontend_version_at_least("1.27.0") +def _cudnn_frontend_supports_single_group_runtime_offsets( + activation_type: type[FusibleOperation], +) -> bool: + """Check cuDNN FE support for single-group runtime offsets.""" + return not issubclass(activation_type, ScaledSReLU) and _cudnn_frontend_version_at_least( + "1.27.0" + ) def _wrap_single_quantized_as_grouped( @@ -1073,7 +1077,7 @@ def fuser_forward( activation_kernel = self.grouped_gemm_activation_kernel() supports_single_group_runtime_offsets = ( - _cudnn_frontend_supports_single_group_runtime_offsets() + _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) ) # Shared experts have one dense group and all optimized kernels derive M @@ -2032,7 +2036,7 @@ def fuser_backward( "use_dynamic_sched": True, } dactivation_kernel = self.grouped_gemm_dactivation_kernel() - if _cudnn_frontend_supports_single_group_runtime_offsets(): + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor @@ -2377,7 +2381,7 @@ def fuser_backward( "use_dynamic_sched": True, } fc1_dgrad_kernel = self.grouped_gemm_quant_kernel() - if _cudnn_frontend_supports_single_group_runtime_offsets(): + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if fc1_op.single_grouped_weight: From abba769f30a7fbae0e88dbf3a11b9ccad198e120 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 12 Aug 2026 08:44:43 -0700 Subject: [PATCH 34/44] Fix JAX MoE against upstream grouped quantize API --- transformer_engine/jax/moe.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 3a70589c1a..014743e26c 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -743,7 +743,6 @@ def _broadcast_bias(value): fc1_quantizer_set.x, group_sizes, flatten_axis=-1, - ragged_scale_sharding=flat_token_sharding, ) casted_wi = tex.grouped_quantize(wi, fc1_quantizer_set.kernel, flatten_axis=-1) if _use_reference_fwd: @@ -784,7 +783,6 @@ def _broadcast_bias(value): fc2_quantizer_set.x, group_sizes, flatten_axis=-1, - ragged_scale_sharding=flat_token_sharding, ) casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) if _use_reference_fwd: @@ -914,7 +912,6 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): fc2_quantizer_set.dgrad, group_sizes, flatten_axis=-1, - ragged_scale_sharding=flat_token_sharding, ) _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) @@ -978,7 +975,6 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): fc1_quantizer_set.dgrad, group_sizes, flatten_axis=-1, - ragged_scale_sharding=flat_token_sharding, ) if _use_reference_dgrad: d_sorted_x = _reference_ragged_dot( From 169b8a86be6aec01fe77b304183cea13fa713f16 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 12 Aug 2026 08:54:15 -0700 Subject: [PATCH 35/44] Restore shard-mapped JAX MoE VJP --- transformer_engine/jax/moe.py | 461 ++++++++++++---------------------- 1 file changed, 163 insertions(+), 298 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 014743e26c..f63807a5ea 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -622,14 +622,14 @@ class _Ctx: # ============================================================================= -# Global-view FFN body +# Per-shard FFN body # ============================================================================= -def _ffn_fwd_global( - recv_tokens: jnp.ndarray, - recv_topk_weights: jnp.ndarray, - token_counts: jnp.ndarray, +def _ffn_fwd_per_shard( + recv_tokens_local: jnp.ndarray, + recv_topk_weights_local: jnp.ndarray, + token_counts_local: jnp.ndarray, wi: jnp.ndarray, wo: jnp.ndarray, wi_0_bias: Optional[jnp.ndarray], @@ -637,99 +637,22 @@ def _ffn_fwd_global( wo_bias: Optional[jnp.ndarray], quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, - dp_size: int, - num_ep: int, num_local_experts: int, activation_type: str, apply_topk_weights_early: bool, - flat_token_sharding: NamedSharding, - flat_group_sharding: NamedSharding, - grouped_weight_sharding: NamedSharding, - grouped_bias_sharding: NamedSharding, ): - """Run the FFN on global EP-dispatch buffers. - - Grouped-operation custom partitioning lowers the global operands to - the per-device problem. ``token_counts`` from ``tex.ep_prepare`` is - passed through as the dynamic grouped-GEMM group sizes, so cuBLAS - skips both 0-token experts and dispatch-buffer over-allocation. - """ - global _debug_ffn_fwd_global_count - if _debug_python_patch and _debug_ffn_fwd_global_count < 10: - _debug_ffn_fwd_global_count += 1 - print( - "[TE patch debug] _ffn_fwd_global " - f"call={_debug_ffn_fwd_global_count} " - f"recv_tokens.shape={recv_tokens.shape} " - f"token_counts.shape={token_counts.shape} " - f"wi.shape={wi.shape} wo.shape={wo.shape} " - f"dp_size={dp_size} num_ep={num_ep} " - f"num_local_experts={num_local_experts} " - f"flat_group_sharding={flat_group_sharding}", - file=sys.stderr, - flush=True, - ) - hidden = recv_tokens.shape[-1] - sorted_x = recv_tokens.reshape(-1, hidden) - recv_w_flat = recv_topk_weights.reshape(-1) - group_sizes = token_counts.reshape(-1).astype(jnp.int32) - sorted_x = jax.lax.with_sharding_constraint(sorted_x, flat_token_sharding) - recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) - group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) + """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) + group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) wi = wi.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) - # Dispatch groups flatten in (dp, ep, local_expert) order. Keep the - # model weights in their original global [expert, ...] layout: grouped - # quantize must see the FSDP shard before grouped_gemm's custom - # partitioning gathers it. This avoids materializing a global dp*expert - # weight tensor and keeps the FSDP collective inside the grouped-GEMM - # custom partitioning boundary. - # - # Bias is the one exception. grouped_gemm's public bias contract is one - # row per global GEMM, so retain its inexpensive logical expansion here. - num_groups = dp_size * num_ep * num_local_experts - - def _weights_in_dispatch_group_order(weight): - expanded = jnp.broadcast_to( - weight[None, ...], - (dp_size, *weight.shape), - ).reshape(num_groups, *weight.shape[1:]) - return jax.lax.with_sharding_constraint(expanded, grouped_weight_sharding) - - def _reference_ragged_dot(lhs, rhs): - def _per_shard(local_lhs, local_rhs, local_group_sizes): - return jax.lax.ragged_dot(local_lhs, local_rhs, local_group_sizes) - - return jax.shard_map( - _per_shard, - mesh=flat_token_sharding.mesh, - in_specs=( - flat_token_sharding.spec, - grouped_weight_sharding.spec, - flat_group_sharding.spec, - ), - out_specs=flat_token_sharding.spec, - )(lhs, rhs, group_sizes) - if _use_reference_fwd: if wi_0_bias is not None: raise ValueError("NVTE_MOE_REFERENCE_FWD does not support expert biases.") - wi_for_fwd = _weights_in_dispatch_group_order(wi) - wo_for_fwd = _weights_in_dispatch_group_order(wo) - - def _broadcast_bias(value): - value = jnp.broadcast_to( - value.reshape(1, num_ep, num_local_experts, *value.shape[1:]), - (dp_size, num_ep, num_local_experts, *value.shape[1:]), - ).reshape(num_groups, *value.shape[1:]) - return jax.lax.with_sharding_constraint(value, grouped_bias_sharding) - - if wi_0_bias is not None: - wi_0_bias = _broadcast_bias(wi_0_bias) - wi_1_bias = _broadcast_bias(wi_1_bias) - wo_bias = _broadcast_bias(wo_bias) # ``wi`` is stored in its gated-SwiGLU layout [expert, hidden, 2*mlp]. # Keeping it contiguous lets grouped quantize/GEMM consume it directly. @@ -746,7 +669,7 @@ def _broadcast_bias(value): ) casted_wi = tex.grouped_quantize(wi, fc1_quantizer_set.kernel, flatten_axis=-1) if _use_reference_fwd: - combined_out = _reference_ragged_dot(sorted_x, wi_for_fwd) + combined_out = jax.lax.ragged_dot(sorted_x, wi, group_sizes) else: combined_out = tex.grouped_gemm( casted_sorted_x.get_tensor(usage=TensorUsage.LHS), @@ -754,7 +677,6 @@ def _broadcast_bias(value): contracting_dims=((1,), (1,)), bias=wi_combined_bias, ) - combined_out = jax.lax.with_sharding_constraint(combined_out, flat_token_sharding) 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).checkpoint( fc1_quantizer_set.x @@ -770,8 +692,6 @@ def _broadcast_bias(value): # transitions to the target precision. act_fn = _convert_to_activation_function(activation_type) intermediate = act_fn(gate_proj_out) * up_proj_out - intermediate = jax.lax.with_sharding_constraint(intermediate, flat_token_sharding) - 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 @@ -786,7 +706,7 @@ def _broadcast_bias(value): ) casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) if _use_reference_fwd: - expert_outputs = _reference_ragged_dot(intermediate, wo_for_fwd) + expert_outputs = jax.lax.ragged_dot(intermediate, wo, group_sizes) else: expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), @@ -794,7 +714,6 @@ def _broadcast_bias(value): contracting_dims=((1,), (1,)), bias=wo_bias, ) - expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, flat_token_sharding) casted_intermediate_lhs_trans = casted_intermediate.get_tensor( usage=TensorUsage.LHS_TRANS ).checkpoint(fc2_quantizer_set.x) @@ -802,8 +721,8 @@ def _broadcast_bias(value): fc2_quantizer_set.kernel ) - expert_outputs_3d = expert_outputs.reshape(*recv_tokens.shape[:-1], expert_outputs.shape[-1]) - group_sizes_nd = group_sizes.reshape(token_counts.shape) + expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) + group_sizes_2d = group_sizes.reshape(1, num_local_experts) residuals = ( casted_sorted_x_lhs_trans, casted_wi_rhs_trans, @@ -811,100 +730,35 @@ def _broadcast_bias(value): up_proj_out, casted_intermediate_lhs_trans, casted_wo_rhs_trans, - group_sizes_nd, + group_sizes_2d, ) return expert_outputs_3d, residuals -def _ffn_bwd_global( - d_expert_outputs: jnp.ndarray, +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, wi: jnp.ndarray, wo: jnp.ndarray, - local_group_sizes: jnp.ndarray, - recv_topk_weights: jnp.ndarray, quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, activation_type: str, apply_topk_weights_early: bool, has_bias: bool, - flat_token_sharding: NamedSharding, - flat_group_sharding: NamedSharding, - grouped_weight_sharding: NamedSharding, - grouped_bias_sharding: NamedSharding, ): - """Run the FFN backward on global residuals. - - Mirrors :func:`_ffn_fwd_global`. Returns - ``(d_sorted_x [num_procs, recv_pr, H], d_recv_w [num_procs, recv_pr], - d_wi, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. - """ - # Each leading process row has its own packed expert prefix and trailing - # over-allocation. Preserve that boundary for diagnostics before flattening - # the group sizes for the shard-local ragged operations. - active_rows_2d = ( - jnp.arange(recv_topk_weights.shape[-1])[None, :] - < jnp.sum(local_group_sizes, axis=-1, dtype=jnp.int32)[:, None] - ) + """Backward mirror of :func:`_ffn_fwd_per_shard`.""" group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) - d_eo_2d = d_expert_outputs.reshape(-1, d_expert_outputs.shape[-1]) - recv_w_flat = recv_topk_weights.reshape(-1) - d_eo_2d = jax.lax.with_sharding_constraint(d_eo_2d, flat_token_sharding) - recv_w_flat = jax.lax.with_sharding_constraint(recv_w_flat, flat_group_sharding) - group_sizes = jax.lax.with_sharding_constraint(group_sizes, flat_group_sharding) + d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) + recv_w_flat = recv_topk_weights_local.reshape(-1) fc1_quantizer_set, fc2_quantizer_set = quantizer_sets - - def _weights_in_dispatch_group_order(weight): - """Repeat one global expert set for each outer FSDP token replica.""" - if group_sizes.shape[0] % weight.shape[0] != 0: - raise ValueError( - "Reference MoE dgrad requires dispatch groups to contain an " - "integer number of global expert sets, but got " - f"{group_sizes.shape[0]} groups and {weight.shape[0]} experts." - ) - num_replicas = group_sizes.shape[0] // weight.shape[0] - expanded = jnp.broadcast_to( - weight[None, ...], - (num_replicas, *weight.shape), - ).reshape(group_sizes.shape[0], *weight.shape[1:]) - return jax.lax.with_sharding_constraint(expanded, grouped_weight_sharding) - - def _reference_ragged_dot(lhs, rhs): - """Run the JAX reference on matching shard-local rows and expert groups.""" - def _per_shard(local_lhs, local_rhs, local_group_sizes): - return jax.lax.ragged_dot(local_lhs, local_rhs, local_group_sizes) - - return jax.shard_map( - _per_shard, - mesh=flat_token_sharding.mesh, - in_specs=( - flat_token_sharding.spec, - grouped_weight_sharding.spec, - flat_group_sharding.spec, - ), - out_specs=flat_token_sharding.spec, - )(lhs, rhs, group_sizes) - - if _use_reference_dgrad: - global _debug_reference_dgrad_count - if _debug_python_patch and _debug_reference_dgrad_count < 10: - _debug_reference_dgrad_count += 1 - print( - "[TE reference dgrad] tracing " - f"call={_debug_reference_dgrad_count} " - f"group_sizes={group_sizes.shape} wi={wi.shape} wo={wo.shape} " - f"group_order=(fsdp,ep,local_expert) " - f"output_constraint={flat_token_sharding}", - file=sys.stderr, - flush=True, - ) - wi_for_dgrad = _weights_in_dispatch_group_order(wi) - wo_for_dgrad = _weights_in_dispatch_group_order(wo) + wgrad_group_active = (group_sizes > 0)[:, None, None] # wo bwd casted_d_eo = tex.grouped_quantize( @@ -916,26 +770,20 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) if _use_reference_dgrad: - d_intermediate = _reference_ragged_dot( - d_eo_2d, - jnp.swapaxes(wo_for_dgrad, -1, -2), - ) + d_intermediate = jax.lax.ragged_dot(d_eo_2d, jnp.swapaxes(wo, -1, -2), group_sizes) else: d_intermediate = tex.grouped_gemm( _casted_d_eo_lhs, casted_wo_rhs_trans, contracting_dims=((1,), (2,)), ) - d_intermediate = jax.lax.with_sharding_constraint(d_intermediate, flat_token_sharding) d_wo = tex.grouped_gemm( casted_intermediate_lhs_trans, _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) - d_wo = jax.lax.with_sharding_constraint(d_wo, grouped_weight_sharding) + d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None - if has_bias: - d_wo_bias = jax.lax.with_sharding_constraint(d_wo_bias, grouped_bias_sharding) act_fn = _convert_to_activation_function(activation_type) if apply_topk_weights_early: @@ -969,7 +817,6 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): # against the fused casted_wi_rhs_trans residual, then split the # wgrad result remains in the contiguous gated-SwiGLU ``wi`` layout. d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) - d_combined = jax.lax.with_sharding_constraint(d_combined, flat_token_sharding) casted_d_combined = tex.grouped_quantize( d_combined, fc1_quantizer_set.dgrad, @@ -977,75 +824,30 @@ def _per_shard(local_lhs, local_rhs, local_group_sizes): flatten_axis=-1, ) if _use_reference_dgrad: - d_sorted_x = _reference_ragged_dot( - d_combined, - jnp.swapaxes(wi_for_dgrad, -1, -2), - ) + d_sorted_x = jax.lax.ragged_dot(d_combined, jnp.swapaxes(wi, -1, -2), group_sizes) else: d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), casted_wi_rhs_trans, contracting_dims=((1,), (2,)), ) - d_sorted_x = jax.lax.with_sharding_constraint(d_sorted_x, flat_token_sharding) d_wi_combined = tex.grouped_gemm( casted_sorted_x_lhs_trans, casted_d_combined.get_tensor(usage=TensorUsage.RHS), contracting_dims=((0,), (0,)), ) - d_wi_combined = jax.lax.with_sharding_constraint(d_wi_combined, grouped_weight_sharding) + d_wi_combined = jnp.where( + wgrad_group_active, d_wi_combined, jnp.zeros_like(d_wi_combined) + ) if has_bias: d_wi_combined_bias = tex.grouped_dbias(d_combined, group_sizes) - d_wi_combined_bias = jax.lax.with_sharding_constraint( - d_wi_combined_bias, grouped_bias_sharding - ) 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 - if _debug_moe_numerics: - # ragged_dot/grouped_gemm consume one packed prefix per process row. - # recv_w padding is not guaranteed initialized, so recv_w != 0 is not - # a valid activity mask. - active_rows = active_rows_2d.reshape(-1) - deo_stats = _debug_stable_stats(d_eo_2d, active_rows) - dint_stats = _debug_stable_stats(d_intermediate, active_rows) - dcombined_stats = _debug_stable_stats(d_combined, active_rows) - dx_stats = _debug_stable_stats(d_sorted_x, active_rows) - drw_stats = _debug_stable_stats(d_recv_w_from_intermediate, active_rows) - jax.debug.print( - "[TE bwd stats] " - "dout(mean={deo_mean:.3e},absmean={deo_absmean:.3e},std={deo_std:.3e}," - "absmax={deo_absmax:.3e},finite={deo_finite:.6f}) " - "dfc2(absmean={dint_absmean:.3e},std={dint_std:.3e},absmax={dint_absmax:.3e}) " - "dact(absmean={dc_absmean:.3e},std={dc_std:.3e},absmax={dc_absmax:.3e}) " - "din(absmean={dx_absmean:.3e},std={dx_std:.3e},absmax={dx_absmax:.3e}," - "finite={dx_finite:.6f}) " - "dweight(absmean={drw_absmean:.3e},std={drw_std:.3e},absmax={drw_absmax:.3e})", - deo_mean=deo_stats[0], - deo_absmean=deo_stats[1], - deo_std=deo_stats[2], - deo_absmax=deo_stats[3], - deo_finite=deo_stats[4], - dint_absmean=dint_stats[1], - dint_std=dint_stats[2], - dint_absmax=dint_stats[3], - dc_absmean=dcombined_stats[1], - dc_std=dcombined_stats[2], - dc_absmax=dcombined_stats[3], - dx_absmean=dx_stats[1], - dx_std=dx_stats[2], - dx_absmax=dx_stats[3], - dx_finite=dx_stats[4], - drw_absmean=drw_stats[1], - drw_std=drw_stats[2], - drw_absmax=drw_stats[3], - ordered=False, - ) - - d_sorted_x_3d = d_sorted_x.reshape(*d_expert_outputs.shape[:-1], d_sorted_x.shape[-1]) - d_recv_w_3d = d_recv_w_from_intermediate.reshape(recv_topk_weights.shape) + 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, @@ -1097,6 +899,7 @@ def _moe_fwd_rule( ``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) @@ -1198,10 +1001,6 @@ def _moe_fwd_rule( batch_pspec_axis = (*data_parallelism_axes, ep_axis) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) - flat_token_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) - flat_group_sharding = NamedSharding(mesh, P(batch_pspec_axis)) - grouped_weight_sharding = NamedSharding(mesh, P(batch_pspec_axis, None, None)) - grouped_bias_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, ep3_spec)) # ---------------- Gate (global view) ---------------- @@ -1312,32 +1111,54 @@ def _moe_fwd_rule( recv_topk_weights, NamedSharding(mesh, ep2_spec) ) - # ---------------- FFN (global view, custom-partitioned primitives) ---------------- + # ---------------- FFN (per-shard via shard_map) ---------------- has_bias = wi_0_bias is not None - # The NCCL EP receive buffer may contain uninitialized padded slots. - # Dynamic group sizes keep grouped GEMMs from reading those rows; the - # backward masks skipped wgrad groups, while EP combine/dispatch bwd - # consume only positions described by handle_mem. - expert_outputs, ffn_residuals = _ffn_fwd_global( - recv_tokens, - recv_topk_weights, - token_counts, - wi, - wo, - wi_0_bias if has_bias else None, - wi_1_bias if has_bias else None, - wo_bias if has_bias else None, - quantizer_sets, - dp_size=dp_size, - num_ep=num_ep, - num_local_experts=num_local_experts, - activation_type=activation_type, - apply_topk_weights_early=apply_topk_weights_early, - flat_token_sharding=flat_token_sharding, - flat_group_sharding=flat_group_sharding, - grouped_weight_sharding=grouped_weight_sharding, - grouped_bias_sharding=grouped_bias_sharding, + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi, wo] + if has_bias: + ffn_in_specs += (bias_spec, bias_spec, bias_spec) + ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) + + residuals_spec = ( + P(), + P(ep_axis, None, None), + P(), + P(), + P(), + P(ep_axis, None, None), + ep2_spec, ) + + def _ffn_fwd_body(*args): + if has_bias: + r_tok, r_w, tc, local_wi, local_wo, w0b, w1b, wob = args + else: + r_tok, r_w, tc, local_wi, local_wo = args + w0b = w1b = wob = None + return _ffn_fwd_per_shard( + r_tok, + r_w, + tc, + local_wi, + local_wo, + w0b, + w1b, + wob, + quantizer_sets, + num_local_experts=num_local_experts, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + ) + + expert_outputs, ffn_residuals = shard_map( + _ffn_fwd_body, + mesh=mesh, + in_specs=ffn_in_specs, + out_specs=(ep3_spec, residuals_spec), + check_rep=False, + )(*ffn_in_args) expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) # ---------------- TE EP combine (global view) ---------------- @@ -1562,6 +1383,7 @@ def _moe_bwd_rule( ): """Backward mirror of :func:`_moe_fwd_rule`.""" 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. d_output, d_aux_loss, _d_total_recv_tokens = cotangents @@ -1626,10 +1448,6 @@ def _moe_bwd_rule( raise ValueError("moe(...) requires an active jax.sharding.Mesh.") num_ep = mesh.shape[ep_axis] num_local_experts = num_experts // num_ep - 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: @@ -1638,10 +1456,6 @@ def _moe_bwd_rule( batch_pspec_axis = (*data_parallelism_axes, ep_axis) ep3_spec = P(batch_pspec_axis, None, None) ep2_spec = P(batch_pspec_axis, None) - flat_token_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) - flat_group_sharding = NamedSharding(mesh, P(batch_pspec_axis)) - grouped_weight_sharding = NamedSharding(mesh, P(batch_pspec_axis, None, None)) - grouped_bias_sharding = NamedSharding(mesh, P(batch_pspec_axis, None)) out_partition_spec = (batch_pspec_axis, None, None) # A scanned layer can reuse the same physical handle_mem buffer for @@ -1745,16 +1559,20 @@ def _moe_bwd_rule( ordered=False, ) - # ---------------- FFN bwd (global view, custom-partitioned primitives) ---------------- - ( - d_sorted_x, - d_recv_w_from_intermediate, - d_wi, - d_wo, - d_wi_0_bias, - d_wi_1_bias, - d_wo_bias, - ) = _ffn_bwd_global( + # ---------------- FFN bwd (per-shard via shard_map) ---------------- + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) + residuals_specs = ( + P(), + P(ep_axis, None, None), + P(), + P(), + P(), + P(ep_axis, None, None), + ep2_spec, + ) + bwd_in_specs = (ep3_spec, *residuals_specs, ep2_spec) + bwd_in_args = [ d_expert_outputs, ctx.casted_sorted_x_lhs_trans, ctx.casted_wi_rhs_trans, @@ -1762,36 +1580,83 @@ def _moe_bwd_rule( ctx.up_proj_out, ctx.casted_intermediate_lhs_trans, ctx.casted_wo_rhs_trans, - ctx.wi, - ctx.wo, ctx.local_group_sizes, ctx.recv_topk_weights, - ctx.quantizer_sets, - activation_type=activation_type, - apply_topk_weights_early=apply_topk_weights_early, - has_bias=has_bias, - flat_token_sharding=flat_token_sharding, - flat_group_sharding=flat_group_sharding, - grouped_weight_sharding=grouped_weight_sharding, - grouped_bias_sharding=grouped_bias_sharding, - ) + ] + if _use_reference_dgrad: + bwd_in_specs += (kernel_spec, kernel_spec) + bwd_in_args.extend([ctx.wi, ctx.wo]) - # Wgrad has one expert-gradient group per outer data replica, ordered - # (dp, ep, local_expert). Sum that dimension to recover the public - # parameter shapes [num_experts, ...]. - def _fold_dp_groups(grad): + def _ffn_bwd_body(*args): + if _use_reference_dgrad: + *common_args, local_wi, local_wo = args + else: + common_args = args + local_wi = local_wo = None + grads = _ffn_bwd_per_shard( + *common_args, + local_wi, + local_wo, + ctx.quantizer_sets, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + has_bias=has_bias, + ) + ( + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, + ) = grads + if data_parallelism_axes: + dp_axes = tuple(data_parallelism_axes) + d_wi_local = jax.lax.psum(d_wi_local, axis_name=dp_axes) + d_wo_local = jax.lax.psum(d_wo_local, axis_name=dp_axes) + if has_bias: + d_wi_0_bias_local = jax.lax.psum(d_wi_0_bias_local, axis_name=dp_axes) + d_wi_1_bias_local = jax.lax.psum(d_wi_1_bias_local, axis_name=dp_axes) + d_wo_bias_local = jax.lax.psum(d_wo_bias_local, axis_name=dp_axes) return ( - grad.reshape(dp_size, num_ep, num_local_experts, *grad.shape[1:]) - .sum(axis=0) - .reshape(num_experts, *grad.shape[1:]) + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, ) - d_wi = _fold_dp_groups(d_wi) - d_wo = _fold_dp_groups(d_wo) if has_bias: - d_wi_0_bias = _fold_dp_groups(d_wi_0_bias) - d_wi_1_bias = _fold_dp_groups(d_wi_1_bias) - d_wo_bias = _fold_dp_groups(d_wo_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, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = shard_map( + _ffn_bwd_body, + mesh=mesh, + in_specs=bwd_in_specs, + out_specs=bwd_out_specs, + check_rep=False, + )(*bwd_in_args) d_recv_w_total = d_recv_w_from_combine + d_recv_w_from_intermediate From 1c98135dd9dfe699f05d0b039d87ce6d43ae8c2f Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 12 Aug 2026 09:10:05 -0700 Subject: [PATCH 36/44] Support MXFP8 quantization in shard-mapped MoE VJP --- tests/jax/run_te_ep_moe.sh | 2 +- tests/jax/test_te_ep_moe_mxfp8.py | 228 ++++++++++++++++++++++++++++++ transformer_engine/jax/moe.py | 110 +++++++++----- 3 files changed, 302 insertions(+), 38 deletions(-) create mode 100644 tests/jax/test_te_ep_moe_mxfp8.py 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_mxfp8.py b/tests/jax/test_te_ep_moe_mxfp8.py new file mode 100644 index 0000000000..aa1432b3ee --- /dev/null +++ b/tests/jax/test_te_ep_moe_mxfp8.py @@ -0,0 +1,228 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Multiprocess MXFP8BlockScaling VJP coverage for the TE EP MoE path.""" + +import os +import sys +from contextlib import ExitStack + +os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") +os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") +os.environ.setdefault("NVTE_JAX_ENFORCE_V2_GROUPED_GEMM", "1") + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from flax.linen import partitioning as nn_partitioning +from jax.experimental import mesh_utils +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P + + +def _read_mp_options(): + num_processes = 0 + process_id = 0 + for index, arg in enumerate(sys.argv): + if arg.startswith("--num-process="): + num_processes = int(arg.split("=", 1)[1]) + elif arg == "--num-process" and index + 1 < len(sys.argv): + num_processes = int(sys.argv[index + 1]) + elif arg.startswith("--process-id="): + process_id = int(arg.split("=", 1)[1]) + elif arg == "--process-id" and index + 1 < len(sys.argv): + process_id = int(sys.argv[index + 1]) + return num_processes, process_id + + +_NUM_PROCESSES, _PROCESS_ID = _read_mp_options() +if _NUM_PROCESSES <= 1: + pytest.skip("requires tests/jax/run_te_ep_moe.sh", allow_module_level=True) + +jax.distributed.initialize( + coordinator_address=os.environ.get("TE_EP_MOE_COORDINATOR_ADDRESS", "127.0.0.1:13457"), + num_processes=_NUM_PROCESSES, + process_id=_PROCESS_ID, + local_device_ids=_PROCESS_ID, +) + +from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.jax.ep import ep_bootstrap +from transformer_engine.jax.moe import ( + get_moe_recv_capacity_per_rank, + moe, + record_ep_bootstrap_signature_for_moe, +) +from transformer_engine.jax.quantize import ( + QuantizeMeta, + QuantizeMetaSet, + QuantizerFactory, + QuantizerSet, + TensorSource, + get_quantize_config_with_recipe, +) +from transformer_engine.jax.sharding import MeshResource, global_shard_guard +from transformer_engine_jax import get_device_compute_capability + + +if get_device_compute_capability(0) < 100: + pytest.skip("MXFP8 grouped GEMM requires Blackwell", allow_module_level=True) + +EP_AXIS = "ep" +FSDP_AXIS = "fsdp" +EP_SIZE = 2 +FSDP_SIZE = jax.device_count() // EP_SIZE +NUM_EXPERTS = 8 +TOPK = 2 +BATCH = jax.device_count() * 2 +SEQ = 32 +HIDDEN = 128 +INTERMEDIATE = 128 +DTYPE = jnp.bfloat16 + +LOGICAL_AXIS_RULES = ( + ("batch", (FSDP_AXIS, EP_AXIS)), + ("exp", EP_AXIS), + ("embed", FSDP_AXIS), + ("mlp", None), +) + + +def _make_mxfp8_quantizer_sets(): + """Construct MaxText-shaped quantizers from MXFP8BlockScaling.""" + recipe = MXFP8BlockScaling() + config = get_quantize_config_with_recipe(recipe) + meta_set = QuantizeMetaSet(x=QuantizeMeta(), kernel=QuantizeMeta(), grad=QuantizeMeta()) + + def _set(n_token_groups, n_expert_groups): + token_set = QuantizerFactory.create_set( + fp8_recipe=recipe, + n_groups=n_token_groups, + quantize_meta_set=meta_set, + ) + expert_set = QuantizerFactory.create_set( + fp8_recipe=recipe, + n_groups=n_expert_groups, + quantize_meta_set=meta_set, + ) + for source in TensorSource: + assert config.get_scaling_mode(source).is_mxfp8_scaling + return QuantizerSet(x=token_set.x, kernel=expert_set.kernel, dgrad=token_set.dgrad) + + # Match MaxText: global dispatch groups include FSDP replicas, while + # expert weights have one group per global expert. + return tuple(_set(FSDP_SIZE * NUM_EXPERTS, NUM_EXPERTS) for _ in range(2)) + + +@pytest.fixture(scope="module") +def mesh(): + devices = mesh_utils.create_device_mesh((FSDP_SIZE, EP_SIZE)) + mesh_obj = Mesh(devices, axis_names=(FSDP_AXIS, EP_AXIS)) + max_tokens_per_rank = (BATCH // jax.process_count()) * SEQ + recv_capacity = get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=max_tokens_per_rank, + ep_size=EP_SIZE, + ) + with mesh_obj, global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ): + ep_bootstrap( + world_size=jax.process_count(), + rank=jax.process_index(), + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_capacity, + 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, + hidden_dim=HIDDEN, + ep_size=EP_SIZE, + ) + return mesh_obj + + +def _context(mesh): + stack = ExitStack() + stack.enter_context(mesh) + stack.enter_context( + global_shard_guard(MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS)) + ) + stack.enter_context(nn_partitioning.axis_rules(LOGICAL_AXIS_RULES)) + return stack + + +def _global_array(value, mesh): + with mesh: + value = jax.jit( + lambda x: jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P())) + )(value) + value.block_until_ready() + return np.asarray(jax.device_get(value.addressable_data(0))) + + +def test_mxfp8_block_scaling_forward_and_vjp(mesh): + keys = jax.random.split(jax.random.PRNGKey(123), 5) + x = jax.random.normal(keys[0], (BATCH, SEQ, HIDDEN), DTYPE) + gate = jax.random.normal(keys[1], (HIDDEN, NUM_EXPERTS), DTYPE) / jnp.sqrt(HIDDEN) + wi = jax.random.normal(keys[2], (NUM_EXPERTS, HIDDEN, 2 * INTERMEDIATE), DTYPE) / jnp.sqrt( + HIDDEN + ) + wo = jax.random.normal(keys[3], (NUM_EXPERTS, INTERMEDIATE, HIDDEN), DTYPE) / jnp.sqrt( + INTERMEDIATE + ) + cotangent = jax.random.normal(keys[4], x.shape, DTYPE) + quantizer_sets = _make_mxfp8_quantizer_sets() + + def forward(x_arg, gate_arg, wi_arg, wo_arg): + output, _, total_recv_tokens = moe( + x_arg, + gate_arg, + wi_arg, + wo_arg, + None, + None, + None, + None, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + apply_topk_weights_early=True, + quantizer_sets=quantizer_sets, + ep_axis=EP_AXIS, + data_parallelism_axes=(FSDP_AXIS,), + input_axes=("batch", None, None), + gate_kernel_axes=("embed", "exp"), + wi_kernel_axes=("exp", "embed", "mlp"), + wo_kernel_axes=("exp", "mlp", "embed"), + dtype=DTYPE, + ) + return output, total_recv_tokens + + def value_and_vjp(x_arg, gate_arg, wi_arg, wo_arg, cotangent_arg): + output, pullback = jax.vjp(lambda a, b, c, d: forward(a, b, c, d)[0], x_arg, gate_arg, wi_arg, wo_arg) + return output, pullback(cotangent_arg) + + with _context(mesh): + x = jax.lax.with_sharding_constraint( + x, NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None, None)) + ) + output, grads = jax.jit(value_and_vjp)(x, gate, wi, wo, cotangent) + jax.block_until_ready((output, grads)) + + assert output.shape == x.shape + assert output.dtype == DTYPE + output_np = _global_array(output, mesh) + assert np.all(np.isfinite(output_np)) + assert np.any(output_np != 0) + + expected_shapes = (x.shape, gate.shape, wi.shape, wo.shape) + for name, grad, shape in zip(("x", "gate", "wi", "wo"), grads, expected_shapes): + assert grad.shape == shape, f"{name} gradient shape mismatch" + grad_np = _global_array(grad, mesh) + assert np.all(np.isfinite(grad_np)), f"{name} gradient has NaN/Inf" + assert np.any(grad_np != 0), f"{name} gradient is identically zero" diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index f63807a5ea..ae84a24ea5 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -34,6 +34,7 @@ import os import sys import warnings +from dataclasses import replace from functools import partial from typing import Any, Optional, Tuple, Union @@ -44,6 +45,7 @@ from . import cpp_extensions as tex from .quantize import ( + GroupedQuantizer, QuantizerSet, TensorUsage, noop_quantizer_set, @@ -626,6 +628,39 @@ class _Ctx: # ============================================================================= +def _localize_grouped_quantizer_set( + quantizer_set: QuantizerSet, num_local_groups: int +) -> QuantizerSet: + """Resize stateless grouped quantizers for one shard-local FFN. + + The public MoE API receives quantizers sized for the global dispatch and + expert layouts. Under ``shard_map``, however, grouped quantize/GEMM sees + only the local experts. MXFP8 block quantizers are stateless and identical + per group, so selecting the corresponding number of entries preserves the + recipe while matching the local grouped operation. + """ + + def _localize(quantizer): + if quantizer is None or not isinstance(quantizer, GroupedQuantizer): + return quantizer + if len(quantizer.quantizers) < num_local_groups: + raise ValueError( + "MoE grouped quantizer has fewer entries than the shard-local " + f"FFN requires: {len(quantizer.quantizers)} < {num_local_groups}." + ) + return replace( + quantizer, + n_groups=num_local_groups, + quantizers=quantizer.quantizers[:num_local_groups], + ) + + return QuantizerSet( + x=_localize(quantizer_set.x), + kernel=_localize(quantizer_set.kernel), + dgrad=_localize(quantizer_set.dgrad), + ) + + def _ffn_fwd_per_shard( recv_tokens_local: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, @@ -660,7 +695,10 @@ def _ffn_fwd_per_shard( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) - fc1_quantizer_set, fc2_quantizer_set = quantizer_sets + fc1_quantizer_set, fc2_quantizer_set = ( + _localize_grouped_quantizer_set(qset, num_local_experts) + for qset in quantizer_sets + ) casted_sorted_x = tex.grouped_quantize( sorted_x, fc1_quantizer_set.x, @@ -678,12 +716,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).checkpoint( - fc1_quantizer_set.x - ) - casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint( - fc1_quantizer_set.kernel - ) # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM # output dtype; the activation output (`intermediate`) stays in the @@ -714,22 +746,15 @@ def _ffn_fwd_per_shard( contracting_dims=((1,), (1,)), bias=wo_bias, ) - casted_intermediate_lhs_trans = casted_intermediate.get_tensor( - usage=TensorUsage.LHS_TRANS - ).checkpoint(fc2_quantizer_set.x) - casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint( - fc2_quantizer_set.kernel - ) - expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) group_sizes_2d = group_sizes.reshape(1, num_local_experts) residuals = ( - casted_sorted_x_lhs_trans, - casted_wi_rhs_trans, + sorted_x, + wi, gate_proj_out, up_proj_out, - casted_intermediate_lhs_trans, - casted_wo_rhs_trans, + intermediate, + wo, group_sizes_2d, ) return expert_outputs_3d, residuals @@ -737,29 +762,50 @@ def _ffn_fwd_per_shard( def _ffn_bwd_per_shard( d_expert_outputs_local: jnp.ndarray, - casted_sorted_x_lhs_trans, - casted_wi_rhs_trans, + sorted_x: jnp.ndarray, + wi: jnp.ndarray, gate_proj_out: jnp.ndarray, up_proj_out: jnp.ndarray, - casted_intermediate_lhs_trans, - casted_wo_rhs_trans, + intermediate: jnp.ndarray, + wo: jnp.ndarray, local_group_sizes: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, - wi: jnp.ndarray, - wo: jnp.ndarray, quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, activation_type: str, apply_topk_weights_early: bool, has_bias: bool, + num_local_experts: int, ): """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) - fc1_quantizer_set, fc2_quantizer_set = quantizer_sets + fc1_quantizer_set, fc2_quantizer_set = ( + _localize_grouped_quantizer_set(qset, num_local_experts) + for qset in quantizer_sets + ) wgrad_group_active = (group_sizes > 0)[:, None, None] + # Recompute stateless grouped quantization inside backward. Besides being + # natural rematerialization for MXFP8BlockScaling, this keeps shard_map + # residuals in their logical BF16 shapes instead of exposing flattened + # quantized data and scale layouts to partition specs. + 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) + casted_intermediate = tex.grouped_quantize( + intermediate, fc2_quantizer_set.x, group_sizes, flatten_axis=-1 + ) + casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) + 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) + casted_intermediate_lhs_trans = casted_intermediate.get_tensor( + usage=TensorUsage.LHS_TRANS + ) + casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) + # wo bwd casted_d_eo = tex.grouped_quantize( d_eo_2d, @@ -1583,24 +1629,14 @@ def _moe_bwd_rule( ctx.local_group_sizes, ctx.recv_topk_weights, ] - if _use_reference_dgrad: - bwd_in_specs += (kernel_spec, kernel_spec) - bwd_in_args.extend([ctx.wi, ctx.wo]) - def _ffn_bwd_body(*args): - if _use_reference_dgrad: - *common_args, local_wi, local_wo = args - else: - common_args = args - local_wi = local_wo = None grads = _ffn_bwd_per_shard( - *common_args, - local_wi, - local_wo, + *args, ctx.quantizer_sets, activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, has_bias=has_bias, + num_local_experts=num_local_experts, ) ( d_sorted_x_local, From 4a97eba4a93979e6337b477764abf126abdfb51a Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 12 Aug 2026 09:37:12 -0700 Subject: [PATCH 37/44] Remove JAX MoE debug switches --- transformer_engine/jax/moe.py | 867 ++-------------------------------- 1 file changed, 27 insertions(+), 840 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index ae84a24ea5..79a6b504c6 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -31,8 +31,6 @@ """ import math -import os -import sys import warnings from dataclasses import replace from functools import partial @@ -128,348 +126,6 @@ def get_moe_recv_capacity_per_rank( return min(requested, worst_case) -_debug_python_patch = os.getenv("NVTE_DEBUG_PYTHON_PATCH", "0") == "1" -_debug_moe_numerics = os.getenv("NVTE_DEBUG_MOE_NUMERICS", "0") == "1" -_debug_moe_input_grad = os.getenv("NVTE_DEBUG_MOE_INPUT_GRAD", "0") == "1" -_use_reference_fwd = os.getenv("NVTE_MOE_REFERENCE_FWD", "0") == "1" -_use_reference_dgrad = os.getenv("NVTE_MOE_REFERENCE_DGRAD", "0") == "1" -_zero_dispatch_weight_grad = os.getenv("NVTE_MOE_ZERO_DISPATCH_WEIGHT_GRAD", "0") == "1" -_refresh_ep_handle_in_bwd = os.getenv("NVTE_MOE_REFRESH_EP_HANDLE_IN_BWD", "0") == "1" -_refresh_ep_handle_before_combine = ( - os.getenv("NVTE_MOE_REFRESH_EP_HANDLE_BEFORE_COMBINE", "0") == "1" -) -_validate_ep_routing = os.getenv("NVTE_MOE_VALIDATE_EP_ROUTING", "0") == "1" -_debug_moe_fwd_recompute = os.getenv("NVTE_MOE_DEBUG_FWD_RECOMPUTE", "0") == "1" -_validate_ep_token_roundtrip = ( - os.getenv("NVTE_MOE_VALIDATE_EP_TOKEN_ROUNDTRIP", "0") == "1" -) -_zero_moe_input_grad = os.getenv("NVTE_MOE_ZERO_INPUT_GRAD", "0") == "1" -_skip_moe_backward = os.getenv("NVTE_MOE_SKIP_BACKWARD", "0") == "1" -_debug_handle_mem = os.getenv("NVTE_MOE_DEBUG_HANDLE_MEM", "0") == "1" -_validate_ep_forward_roundtrip = ( - os.getenv("NVTE_MOE_VALIDATE_EP_FORWARD_ROUNDTRIP", "0") == "1" -) -_zero_moe_output = os.getenv("NVTE_MOE_ZERO_OUTPUT", "0") == "1" -_skip_moe_forward = os.getenv("NVTE_MOE_SKIP_FORWARD", "0") == "1" -_EP_ROUTING_PROBE_WIDTH = 16 -_debug_ffn_fwd_global_count = 0 -_debug_reference_dgrad_count = 0 -if _debug_python_patch: - print( - "[TE patch debug] imported transformer_engine.jax.moe " - f"from {__file__} (rank={os.getenv('SLURM_PROCID', 'unknown')}, " - f"reference_fwd={_use_reference_fwd}, reference_dgrad={_use_reference_dgrad})", - file=sys.stderr, - flush=True, - ) -if _use_reference_dgrad: - print( - "[TE reference dgrad] enabled: activation dgrad uses jax.lax.ragged_dot; " - "forward and weight gradients remain on the TE grouped-GEMM path", - file=sys.stderr, - flush=True, - ) -if _use_reference_fwd: - print( - "[TE reference fwd] enabled: FFN forward uses shard-local " - "jax.lax.ragged_dot; weight gradients remain on the TE grouped-GEMM path", - file=sys.stderr, - flush=True, - ) -if _zero_dispatch_weight_grad: - print( - "[TE dispatch diagnostic] routing-weight cotangent into ep_dispatch_bwd " - "is forced to zero; token dgrad remains enabled", - file=sys.stderr, - flush=True, - ) -if _refresh_ep_handle_in_bwd: - print( - "[TE EP handle diagnostic] backward refreshes the NCCL EP routing " - "handle from the saved routing map before combine/dispatch", - file=sys.stderr, - flush=True, - ) -if _refresh_ep_handle_before_combine: - print( - "[TE EP handle diagnostic] forward refreshes the NCCL EP routing " - "handle immediately before combine", - file=sys.stderr, - flush=True, - ) -if _validate_ep_routing: - print( - "[TE EP routing probe] enabled: validates effective combine-fwd and " - "dispatch-bwd handle mappings with deterministic expert codes", - file=sys.stderr, - flush=True, - ) -if _debug_moe_fwd_recompute: - print( - "[TE MoE fwd recompute debug] enabled: logs order-sensitive input/output " - "statistics keyed by the routing signature", - file=sys.stderr, - flush=True, - ) -if _validate_ep_token_roundtrip: - print( - "[TE EP token roundtrip] enabled: validates intra-expert token ordering " - "across dispatch-fwd and dispatch-bwd", - file=sys.stderr, - flush=True, - ) -if _zero_moe_input_grad: - print( - "[TE MoE input-grad diagnostic] the complete MoE input cotangent is " - "forced to zero after routing and gate gradients are combined", - file=sys.stderr, - flush=True, - ) -if _skip_moe_backward: - print( - "[TE MoE backward diagnostic] bypassing the complete TE MoE backward " - "body and returning zero activation/parameter cotangents", - file=sys.stderr, - flush=True, - ) -if _debug_handle_mem: - print( - "[TE EP handle-value diagnostic] logging each handle_mem's complete-byte " - "signatures and exact head/tail byte samples after ep_prepare", - file=sys.stderr, - flush=True, - ) -if _validate_ep_forward_roundtrip: - print( - "[TE EP forward roundtrip] validating real dispatched token values and " - "ordering through combine_fwd", - file=sys.stderr, - flush=True, - ) -if _zero_moe_output: - print( - "[TE MoE output diagnostic] forcing the routed-MoE forward output to " - "zero after all TE forward operations and validation probes", - file=sys.stderr, - flush=True, - ) -if _skip_moe_forward: - print( - "[TE MoE forward diagnostic] bypassing the complete TE routed-MoE " - "forward body and returning a zero output", - file=sys.stderr, - flush=True, - ) - - -def _debug_stable_stats(value, row_active=None): - """Return bounded sampled stats without materializing a full float32 copy.""" - value = jnp.asarray(value) - # A dispatch tensor is commonly [num_procs, recv_capacity, hidden] while - # its activity mask is [num_procs, recv_capacity]. Treat every mask entry - # as one logical row; using only value.shape[0] here would accidentally - # apply the first few token-mask entries to whole process-sized slabs. - if row_active is None: - matrix = value.reshape(value.shape[0], -1) - else: - row_active = jnp.asarray(row_active, jnp.bool_).reshape(-1) - if value.size % row_active.size != 0: - raise ValueError( - "Debug row mask must divide the sampled tensor size, but got " - f"value.shape={value.shape} and row_active.shape={row_active.shape}." - ) - matrix = value.reshape(row_active.size, -1) - max_samples = 65536 - stride = max((matrix.size + max_samples - 1) // max_samples, 1) - sample = matrix.reshape(-1)[::stride][:max_samples].astype(jnp.float32) - # Avoid forming flattened indices that can exceed int32 for production - # dispatch buffers (>180B logical elements). - stride_rows, stride_remainder = divmod(stride, matrix.shape[1]) - sample_indices = jnp.arange(sample.size, dtype=jnp.int32) - sampled_rows = ( - sample_indices * stride_rows - + (sample_indices * stride_remainder) // matrix.shape[1] - ) - if row_active is None: - active = jnp.ones(sample.shape, dtype=jnp.bool_) - else: - active = row_active[sampled_rows] - finite = jnp.isfinite(sample) & active - finite_value = jnp.where(finite, sample, 0.0) - absmax = jnp.max(jnp.abs(finite_value)) - safe_scale = jnp.where(absmax > 0, absmax, 1.0) - scaled = finite_value / safe_scale - element_count = jnp.maximum(jnp.sum(active, dtype=jnp.float32), 1.0) - scaled_mean = jnp.sum(scaled) / element_count - scaled_square_mean = jnp.sum(jnp.square(scaled)) / element_count - mean = absmax * scaled_mean - abs_mean = absmax * (jnp.sum(jnp.abs(scaled)) / element_count) - stddev = absmax * jnp.sqrt(jnp.maximum(scaled_square_mean - jnp.square(scaled_mean), 0.0)) - finite_fraction = jnp.sum(finite, dtype=jnp.float32) / element_count - return mean, abs_mean, stddev, absmax, finite_fraction - - -def _ep_routing_probe_code_table(num_experts, dtype): - """Return deterministic ±1 codes that identify experts across 16 channels.""" - expert = jnp.arange(num_experts, dtype=jnp.uint32)[:, None] - channel = jnp.arange(_EP_ROUTING_PROBE_WIDTH, dtype=jnp.uint32)[None, :] - value = (expert + jnp.uint32(1)) * jnp.uint32(0x9E3779B1) - value ^= (channel + jnp.uint32(1)) * jnp.uint32(0x85EBCA77) - value ^= value >> jnp.uint32(16) - value *= jnp.uint32(0xC2B2AE3D) - value ^= value >> jnp.uint32(13) - return jnp.where((value & jnp.uint32(1)) != 0, 1, -1).astype(dtype) - - -def _ep_routing_probe_packed_codes( - token_counts, - recv_capacity_per_rank, - num_ep, - num_local_experts, - dtype, -): - """Build expert-major probe rows matching the native EP receive layout.""" - leading_size = token_counts.shape[0] - ep_rank = jnp.arange(leading_size, dtype=jnp.int32) % num_ep - local_expert = jnp.arange(num_local_experts, dtype=jnp.int32) - global_expert = ep_rank[:, None] * num_local_experts + local_expert[None, :] - code_table = _ep_routing_probe_code_table(num_ep * num_local_experts, dtype) - codes_by_group = code_table[global_expert] - - def _repeat_one_leading_group(group_codes, group_counts): - return jnp.repeat( - group_codes, - group_counts.astype(jnp.int32), - axis=0, - total_repeat_length=recv_capacity_per_rank, - ) - - packed = jax.vmap(_repeat_one_leading_group)(codes_by_group, token_counts) - active_rows = ( - jnp.arange(recv_capacity_per_rank, dtype=jnp.int32)[None, :] - < jnp.sum(token_counts, axis=-1, dtype=jnp.int32)[:, None] - ) - return jnp.where(active_rows[..., None], packed, jnp.zeros_like(packed)) - - -def _ep_routing_probe_signature(topk_idx): - """Return two order-sensitive uint32 signatures for the intended map.""" - flat = topk_idx.reshape(-1).astype(jnp.uint32) - position = jnp.arange(flat.size, dtype=jnp.uint32) - signature_0 = jnp.sum( - (flat + jnp.uint32(1)) - * (position * jnp.uint32(0x9E3779B1) + jnp.uint32(0x85EBCA77)), - dtype=jnp.uint32, - ) - signature_1 = jnp.sum( - (flat + jnp.uint32(17)) - * (position * jnp.uint32(0xC2B2AE3D) + jnp.uint32(0x27D4EB2F)), - dtype=jnp.uint32, - ) - return signature_0, signature_1 - - -def _print_handle_mem_value(label, handle_mem, topk_idx): - """Log exact handle bytes plus full-buffer fingerprints after ep_prepare. - - ``handle_mem`` is an opaque uint8 tensor with one row per global EP/DP - rank. Printing every byte for every scanned layer would make the - multi-host log impractically large, so the exact first/last 64 bytes are - printed and two order-sensitive signatures cover every byte in each row. - The routing-map signature in the same record makes repeated-forward - comparisons unambiguous. - """ - if not _debug_handle_mem: - return - rows = handle_mem.reshape(-1, handle_mem.shape[-1]).astype(jnp.uint32) - position = jnp.arange(rows.shape[-1], dtype=jnp.uint32) - signature_0 = jnp.sum( - (rows + jnp.uint32(1)) - * (position * jnp.uint32(0x9E3779B1) + jnp.uint32(0x85EBCA77)), - axis=-1, - dtype=jnp.uint32, - ) - signature_1 = jnp.sum( - (rows + jnp.uint32(17)) - * (position * jnp.uint32(0xC2B2AE3D) + jnp.uint32(0x27D4EB2F)), - axis=-1, - dtype=jnp.uint32, - ) - byte_sum = jnp.sum(rows, axis=-1, dtype=jnp.uint32) - route_signature_0, route_signature_1 = _ep_routing_probe_signature(topk_idx) - sample_width = min(64, handle_mem.shape[-1]) - jax.debug.print( - f"[TE EP handle value] label={label} " - f"shape={handle_mem.shape} " - "route_sig=({route_sig0},{route_sig1}) " - "byte_sum={byte_sum} handle_sig0={handle_sig0} " - "handle_sig1={handle_sig1} head={head} tail={tail}", - route_sig0=route_signature_0, - route_sig1=route_signature_1, - byte_sum=byte_sum, - handle_sig0=signature_0, - handle_sig1=signature_1, - head=handle_mem[..., :sample_width], - tail=handle_mem[..., -sample_width:], - ordered=False, - ) - - -def _debug_ordered_tensor_stats(value): - """Return stable scalar stats plus two order-sensitive sampled projections.""" - value = jnp.asarray(value) - max_samples = 65536 - stride = max((value.size + max_samples - 1) // max_samples, 1) - sample = value.reshape(-1)[::stride][:max_samples].astype(jnp.float32) - finite = jnp.isfinite(sample) - finite_value = jnp.where(finite, sample, 0.0) - absmax = jnp.max(jnp.abs(finite_value)) - safe_scale = jnp.where(absmax > 0, absmax, 1.0) - scaled = finite_value / safe_scale - count = jnp.maximum(jnp.sum(finite, dtype=jnp.float32), 1.0) - mean = absmax * jnp.sum(scaled) / count - square_mean = jnp.sum(jnp.square(scaled)) / count - stddev = absmax * jnp.sqrt( - jnp.maximum(square_mean - jnp.square(jnp.sum(scaled) / count), 0.0) - ) - position = jnp.arange(sample.size, dtype=jnp.uint32) - sign_0 = jnp.where( - ((position * jnp.uint32(0x9E3779B1) + jnp.uint32(0x85EBCA77)) >> 31) != 0, - 1.0, - -1.0, - ) - sign_1 = jnp.where( - ((position * jnp.uint32(0xC2B2AE3D) + jnp.uint32(0x27D4EB2F)) >> 31) != 0, - 1.0, - -1.0, - ) - projection_0 = jnp.sum(scaled * sign_0) / count - projection_1 = jnp.sum(scaled * sign_1) / count - return mean, stddev, absmax, projection_0, projection_1, jnp.mean(finite) - - -def _print_ep_routing_probe_result(label, topk_idx, actual, expected, tolerance): - """Print an elementwise comparison for an expert-code routing probe.""" - difference = actual.astype(jnp.float32) - expected.astype(jnp.float32) - abs_difference = jnp.abs(difference) - mismatch = abs_difference > tolerance - signature_0, signature_1 = _ep_routing_probe_signature(topk_idx) - jax.debug.print( - "[TE EP routing probe] {label} route_sig=({signature_0},{signature_1}) " - "match={match} " - "mismatch_fraction={mismatch_fraction:.6e} " - "absmean={absmean:.6e} absmax={absmax:.6e}", - label=label, - signature_0=signature_0, - signature_1=signature_1, - match=jnp.all(~mismatch), - mismatch_fraction=jnp.mean(mismatch.astype(jnp.float32)), - absmean=jnp.mean(abs_difference), - absmax=jnp.max(abs_difference), - ordered=False, - ) - - def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: """Sharding constraint that keeps bwd cotangents in the primal dtype. @@ -603,18 +259,12 @@ class _Ctx: handle_mem: jnp.ndarray token_counts: jnp.ndarray recv_topk_weights: jnp.ndarray - recv_token_probe: Any 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 - wi: Any - wo: Any - wi_0_bias: Any - wi_1_bias: Any - wo_bias: Any expert_outputs: jnp.ndarray local_group_sizes: jnp.ndarray quantizer_sets: Any @@ -685,10 +335,6 @@ def _ffn_fwd_per_shard( wi = wi.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) - if _use_reference_fwd: - if wi_0_bias is not None: - raise ValueError("NVTE_MOE_REFERENCE_FWD does not support expert biases.") - # ``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 = ( @@ -706,15 +352,12 @@ def _ffn_fwd_per_shard( flatten_axis=-1, ) casted_wi = tex.grouped_quantize(wi, fc1_quantizer_set.kernel, flatten_axis=-1) - if _use_reference_fwd: - combined_out = jax.lax.ragged_dot(sorted_x, wi, group_sizes) - else: - combined_out = tex.grouped_gemm( - casted_sorted_x.get_tensor(usage=TensorUsage.LHS), - casted_wi.get_tensor(usage=TensorUsage.RHS), - contracting_dims=((1,), (1,)), - bias=wi_combined_bias, - ) + combined_out = tex.grouped_gemm( + casted_sorted_x.get_tensor(usage=TensorUsage.LHS), + casted_wi.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((1,), (1,)), + bias=wi_combined_bias, + ) gate_proj_out, up_proj_out = jnp.split(combined_out, 2, axis=-1) # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM @@ -737,15 +380,12 @@ def _ffn_fwd_per_shard( flatten_axis=-1, ) casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) - if _use_reference_fwd: - expert_outputs = jax.lax.ragged_dot(intermediate, wo, group_sizes) - else: - 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, - ) + 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, + ) expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) group_sizes_2d = group_sizes.reshape(1, num_local_experts) residuals = ( @@ -815,14 +455,11 @@ def _ffn_bwd_per_shard( ) _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) - if _use_reference_dgrad: - d_intermediate = jax.lax.ragged_dot(d_eo_2d, jnp.swapaxes(wo, -1, -2), group_sizes) - else: - d_intermediate = tex.grouped_gemm( - _casted_d_eo_lhs, - casted_wo_rhs_trans, - contracting_dims=((1,), (2,)), - ) + d_intermediate = tex.grouped_gemm( + _casted_d_eo_lhs, + casted_wo_rhs_trans, + contracting_dims=((1,), (2,)), + ) d_wo = tex.grouped_gemm( casted_intermediate_lhs_trans, _casted_d_eo_rhs, @@ -869,14 +506,11 @@ def _ffn_bwd_per_shard( group_sizes, flatten_axis=-1, ) - if _use_reference_dgrad: - d_sorted_x = jax.lax.ragged_dot(d_combined, jnp.swapaxes(wi, -1, -2), group_sizes) - else: - d_sorted_x = tex.grouped_gemm( - casted_d_combined.get_tensor(usage=TensorUsage.LHS), - casted_wi_rhs_trans, - contracting_dims=((1,), (2,)), - ) + d_sorted_x = tex.grouped_gemm( + casted_d_combined.get_tensor(usage=TensorUsage.LHS), + casted_wi_rhs_trans, + contracting_dims=((1,), (2,)), + ) d_wi_combined = tex.grouped_gemm( casted_sorted_x_lhs_trans, casted_d_combined.get_tensor(usage=TensorUsage.RHS), @@ -949,50 +583,6 @@ def _moe_fwd_rule( x = with_sharding_constraint_by_logical_axes(x, input_axes) - if _skip_moe_forward: - if not _skip_moe_backward: - raise RuntimeError( - "NVTE_MOE_SKIP_FORWARD requires NVTE_MOE_SKIP_BACKWARD=1." - ) - has_bias = wi_0_bias is not None - ctx = _Ctx( - x=x, - gate_kernel=gate_kernel, - expert_bias=expert_bias, - logits_2d=None, - saved_scores=None, - routing_map=None, - cfg=None, - handle_mem=None, - token_counts=None, - recv_topk_weights=None, - recv_token_probe=None, - casted_sorted_x_lhs_trans=None, - casted_wi_rhs_trans=None, - gate_proj_out=None, - up_proj_out=None, - casted_intermediate_lhs_trans=None, - casted_wo_rhs_trans=None, - wi=wi, - wo=wo, - wi_0_bias=wi_0_bias if has_bias else None, - wi_1_bias=wi_1_bias if has_bias else None, - wo_bias=wo_bias if has_bias else None, - expert_outputs=None, - local_group_sizes=None, - quantizer_sets=quantizer_sets, - ) - static = { - "has_bias": has_bias, - "x_shape": x.shape, - "recv_pr": 0, - } - return ( - jnp.zeros_like(x), - jnp.zeros((), dtype=x.dtype), - jnp.zeros((1,), dtype=jnp.int32), - ), (ctx, static) - mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") @@ -1147,7 +737,6 @@ 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) - _print_handle_mem_value("forward_prepare", handle_mem, 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 @@ -1209,31 +798,11 @@ def _ffn_fwd_body(*args): # ---------------- TE EP combine (global view) ---------------- out_partition_spec = (batch_pspec_axis, None, None) - combine_handle_mem = handle_mem - if _refresh_ep_handle_before_combine: - # A scanned forward can reuse one physical handle_mem allocation across - # loop iterations. Re-prepare the current routing map at the combine - # consumption point so combine cannot rely on routing state left by a - # different iteration. Keep the original handle in the VJP residual: - # this diagnostic intentionally isolates forward combine. - refreshed_token_counts, combine_handle_mem = tex.ep_prepare(cfg, topk_idx_3d) - _print_handle_mem_value( - "forward_pre_combine_refresh", combine_handle_mem, topk_idx_3d - ) - refreshed_token_counts = jax.lax.with_sharding_constraint( - refreshed_token_counts, NamedSharding(mesh, ep2_spec) - ) - if _debug_moe_numerics: - jax.debug.print( - "[TE EP pre-combine refresh] token_counts_match={match}", - match=jnp.all(refreshed_token_counts == token_counts), - ordered=False, - ) if apply_topk_weights_early: # expert_outputs is already weighted upstream. output = tex.ep_combine_fwd( cfg, - combine_handle_mem, + handle_mem, expert_outputs, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, @@ -1245,7 +814,7 @@ def _ffn_fwd_body(*args): weighted = expert_outputs * w output = tex.ep_combine_fwd( cfg, - combine_handle_mem, + handle_mem, weighted, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, @@ -1253,102 +822,6 @@ def _ffn_fwd_body(*args): # output of MLP should be sharded the same way as the activation input output = with_sharding_constraint_by_logical_axes(output, input_axes) - if _validate_ep_forward_roundtrip: - # This validates more than the synthetic expert-code probe below: - # dispatch the actual input token values, combine them immediately, - # and compare against the weighted identity analytically. It detects - # a mutually consistent but wrong dispatch/combine token permutation. - probe_width = min(_EP_ROUTING_PROBE_WIDTH, H) - roundtrip_weighted = ( - recv_tokens[..., :probe_width].astype(jnp.float32) - * recv_topk_weights[..., None] - ).astype(x.dtype) - roundtrip_output = tex.ep_combine_fwd( - cfg, - combine_handle_mem, - roundtrip_weighted, - num_local_tokens=(B, S), - out_partition_spec=out_partition_spec, - ) - expected_roundtrip = ( - x[..., :probe_width].astype(jnp.float32) - * jnp.sum(topk_w_3d, axis=-1, keepdims=True) - ).astype(x.dtype) - _print_ep_routing_probe_result( - "forward_token_roundtrip", - topk_idx_3d, - roundtrip_output, - expected_roundtrip, - tolerance=6.25e-2, - ) - - if _validate_ep_routing: - probe_code_table = _ep_routing_probe_code_table(num_experts, x.dtype) - probe_packed = _ep_routing_probe_packed_codes( - token_counts, - recv_pr, - num_ep, - num_local_experts, - x.dtype, - ) - probe_packed = jax.lax.with_sharding_constraint( - probe_packed, NamedSharding(mesh, ep3_spec) - ) - probe_weighted = ( - probe_packed.astype(jnp.float32) * recv_topk_weights[..., None] - ).astype(x.dtype) - probe_combined = tex.ep_combine_fwd( - cfg, - combine_handle_mem, - probe_weighted, - num_local_tokens=(B, S), - out_partition_spec=out_partition_spec, - ) - expected_probe_terms = ( - probe_code_table[topk_idx_3d].astype(jnp.float32) - * topk_w_3d[..., None] - ).astype(x.dtype) - expected_probe_combined = jnp.sum( - expected_probe_terms.astype(jnp.float32), axis=-2 - ).astype(x.dtype) - _print_ep_routing_probe_result( - "combine_fwd", - topk_idx_3d, - probe_combined, - expected_probe_combined, - tolerance=6.25e-2, - ) - - if _zero_moe_output: - output = jnp.zeros_like(output) - - if _debug_moe_fwd_recompute: - signature_0, signature_1 = _ep_routing_probe_signature(topk_idx_3d) - x_stats = _debug_ordered_tensor_stats(x) - output_stats = _debug_ordered_tensor_stats(output) - jax.debug.print( - "[TE MoE fwd recompute] route_sig=({signature_0},{signature_1}) " - "input(mean={x_mean:.6e},std={x_std:.6e},absmax={x_absmax:.6e}," - "proj=({x_proj0:.9e},{x_proj1:.9e}),finite={x_finite:.6f}) " - "output(mean={out_mean:.6e},std={out_std:.6e},absmax={out_absmax:.6e}," - "proj=({out_proj0:.9e},{out_proj1:.9e}),finite={out_finite:.6f})", - signature_0=signature_0, - signature_1=signature_1, - x_mean=x_stats[0], - x_std=x_stats[1], - x_absmax=x_stats[2], - x_proj0=x_stats[3], - x_proj1=x_stats[4], - x_finite=x_stats[5], - out_mean=output_stats[0], - out_std=output_stats[1], - out_absmax=output_stats[2], - out_proj0=output_stats[3], - out_proj1=output_stats[4], - out_finite=output_stats[5], - ordered=False, - ) - ( casted_sorted_x_lhs_trans, casted_wi_rhs_trans, @@ -1370,25 +843,12 @@ def _ffn_fwd_body(*args): handle_mem=handle_mem, token_counts=token_counts, recv_topk_weights=recv_topk_weights, - recv_token_probe=( - jax.lax.with_sharding_constraint( - recv_tokens[..., :_EP_ROUTING_PROBE_WIDTH], - NamedSharding(mesh, ep3_spec), - ) - if _validate_ep_token_roundtrip - else None - ), 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, - wi=wi if (_use_reference_dgrad or _skip_moe_backward) else None, - wo=wo if (_use_reference_dgrad or _skip_moe_backward) else None, - wi_0_bias=wi_0_bias if (has_bias and _skip_moe_backward) else None, - wi_1_bias=wi_1_bias if (has_bias and _skip_moe_backward) else None, - wo_bias=wo_bias if (has_bias and _skip_moe_backward) else None, expert_outputs=expert_outputs, local_group_sizes=local_group_sizes, quantizer_sets=quantizer_sets, @@ -1439,56 +899,6 @@ def _moe_bwd_rule( x_shape = static["x_shape"] recv_pr = static["recv_pr"] - if _skip_moe_backward: - # Strong isolation diagnostic: do not execute combine_bwd, grouped - # GEMM backward, dispatch_bwd, or router backward. This differs from - # NVTE_MOE_ZERO_INPUT_GRAD, which discards d_x only after all of those - # operations have already run and therefore cannot rule out an - # asynchronous side effect from a backward custom call. - if ctx.wi is None or ctx.wo is None: - raise RuntimeError( - "NVTE_MOE_SKIP_BACKWARD requires wi/wo in the VJP residual." - ) - d_x = with_sharding_constraint_by_logical_axes( - jnp.zeros_like(ctx.x), input_axes - ) - d_gate_kernel = with_sharding_constraint_by_logical_axes( - jnp.zeros_like(ctx.gate_kernel), gate_kernel_axes - ) - d_wi = with_sharding_constraint_by_logical_axes( - jnp.zeros_like(ctx.wi), wi_kernel_axes - ) - d_wo = with_sharding_constraint_by_logical_axes( - jnp.zeros_like(ctx.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( - jnp.zeros_like(ctx.wi_0_bias), wi_bias_axes - ) - d_wi_1_bias = with_sharding_constraint_by_logical_axes( - jnp.zeros_like(ctx.wi_1_bias), wi_bias_axes - ) - d_wo_bias = with_sharding_constraint_by_logical_axes( - jnp.zeros_like(ctx.wo_bias), wo_bias_axes - ) - else: - d_wi_0_bias = None - d_wi_1_bias = None - d_wo_bias = None - return ( - d_x, - d_gate_kernel, - d_wi, - d_wo, - d_wi_0_bias, - d_wi_1_bias, - d_wo_bias, - jnp.zeros_like(ctx.expert_bias), - ctx.quantizer_sets, - ) - mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") @@ -1504,36 +914,9 @@ def _moe_bwd_rule( ep2_spec = P(batch_pspec_axis, None) out_partition_spec = (batch_pspec_axis, None, None) - # A scanned layer can reuse the same physical handle_mem buffer for - # different loop iterations. The native NCCL EP cache is keyed by that - # buffer pointer and retains routing state established by ep_prepare. - # Refreshing here updates the cached handle to the routing plan for the - # current reverse-scan iteration before either backward EP operation. - bwd_handle_mem = ctx.handle_mem - if _refresh_ep_handle_in_bwd: - bwd_selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] - bwd_topk_idx = bwd_selected_experts.reshape(B, S, K).astype(jnp.int32) - bwd_topk_idx = jax.lax.with_sharding_constraint( - bwd_topk_idx, NamedSharding(mesh, ep3_spec) - ) - refreshed_token_counts, bwd_handle_mem = tex.ep_prepare(ctx.cfg, bwd_topk_idx) - _print_handle_mem_value( - "backward_refresh", bwd_handle_mem, bwd_topk_idx - ) - refreshed_token_counts = jax.lax.with_sharding_constraint( - refreshed_token_counts, NamedSharding(mesh, ep2_spec) - ) - if _debug_moe_numerics: - token_count_match = jnp.all(refreshed_token_counts == ctx.token_counts) - jax.debug.print( - "[TE EP handle refresh] token_counts_match={match}", - match=token_count_match, - ordered=False, - ) - # ---------------- 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, bwd_handle_mem, d_output, recv_pr) + 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) ) @@ -1565,46 +948,6 @@ def _moe_bwd_rule( 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) - if _debug_moe_numerics: - d_output_stats = _debug_stable_stats(d_output) - grad_pre_combine_stats = _debug_stable_stats( - grad_pre_combine, active_recv_rows - ) - d_expert_output_stats = _debug_stable_stats( - d_expert_outputs, active_recv_rows - ) - d_recv_w_stats = _debug_stable_stats( - d_recv_w_from_combine, active_recv_rows - ) - jax.debug.print( - "[TE combine-bwd stats] " - "upstream(absmean={up_absmean:.3e},std={up_std:.3e}," - "absmax={up_absmax:.3e},finite={up_finite:.6f}) " - "combine(absmean={combine_absmean:.3e},std={combine_std:.3e}," - "absmax={combine_absmax:.3e},finite={combine_finite:.6f}) " - "weighted(absmean={weighted_absmean:.3e},std={weighted_std:.3e}," - "absmax={weighted_absmax:.3e},finite={weighted_finite:.6f}) " - "dweight(absmean={dw_absmean:.3e},std={dw_std:.3e}," - "absmax={dw_absmax:.3e},finite={dw_finite:.6f})", - up_absmean=d_output_stats[1], - up_std=d_output_stats[2], - up_absmax=d_output_stats[3], - up_finite=d_output_stats[4], - combine_absmean=grad_pre_combine_stats[1], - combine_std=grad_pre_combine_stats[2], - combine_absmax=grad_pre_combine_stats[3], - combine_finite=grad_pre_combine_stats[4], - weighted_absmean=d_expert_output_stats[1], - weighted_std=d_expert_output_stats[2], - weighted_absmax=d_expert_output_stats[3], - weighted_finite=d_expert_output_stats[4], - dw_absmean=d_recv_w_stats[1], - dw_std=d_recv_w_stats[2], - dw_absmax=d_recv_w_stats[3], - dw_finite=d_recv_w_stats[4], - ordered=False, - ) - # ---------------- FFN bwd (per-shard via shard_map) ---------------- kernel_spec = P(ep_axis, None, None) bias_spec = P(ep_axis, None) @@ -1699,125 +1042,15 @@ def _ffn_bwd_body(*args): # ---------------- 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)) - dispatch_weight_cotangent = ( - jnp.zeros_like(d_recv_w_total) if _zero_dispatch_weight_grad else d_recv_w_total - ) d_x_from_dispatch, d_topk_w = tex.ep_dispatch_bwd( ctx.cfg, - bwd_handle_mem, + ctx.handle_mem, d_sorted_x, - dispatch_weight_cotangent, + d_recv_w_total, num_local_tokens=(B, S), out_partition_spec=out_partition_spec, ) - if _validate_ep_token_roundtrip: - roundtrip_topk_idx = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] - roundtrip_topk_idx_3d = roundtrip_topk_idx.reshape(B, S, K).astype(jnp.int32) - roundtrip_topk_idx_3d = jax.lax.with_sharding_constraint( - roundtrip_topk_idx_3d, NamedSharding(mesh, ep3_spec) - ) - roundtrip_tokens, _ = tex.ep_dispatch_bwd( - ctx.cfg, - bwd_handle_mem, - ctx.recv_token_probe, - jnp.zeros_like(ctx.recv_topk_weights), - num_local_tokens=(B, S), - out_partition_spec=out_partition_spec, - ) - expected_roundtrip = ( - ctx.x[..., :_EP_ROUTING_PROBE_WIDTH].astype(jnp.float32) * float(K) - ).astype(roundtrip_tokens.dtype) - _print_ep_routing_probe_result( - "dispatch_token_roundtrip", - roundtrip_topk_idx_3d, - roundtrip_tokens, - expected_roundtrip, - tolerance=6.25e-2, - ) - - if _validate_ep_routing: - probe_topk_idx = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] - probe_topk_idx_3d = probe_topk_idx.reshape(B, S, K).astype(jnp.int32) - probe_topk_idx_3d = jax.lax.with_sharding_constraint( - probe_topk_idx_3d, NamedSharding(mesh, ep3_spec) - ) - probe_code_table = _ep_routing_probe_code_table(num_experts, d_sorted_x.dtype) - probe_packed = _ep_routing_probe_packed_codes( - ctx.token_counts, - recv_pr, - num_ep, - num_local_experts, - d_sorted_x.dtype, - ) - probe_packed = jax.lax.with_sharding_constraint( - probe_packed, NamedSharding(mesh, ep3_spec) - ) - probe_dispatch, _ = tex.ep_dispatch_bwd( - ctx.cfg, - bwd_handle_mem, - probe_packed, - jnp.zeros_like(ctx.recv_topk_weights), - num_local_tokens=(B, S), - out_partition_spec=out_partition_spec, - ) - expected_probe_dispatch = jnp.sum( - probe_code_table[probe_topk_idx_3d].astype(jnp.float32), - axis=-2, - ).astype(d_sorted_x.dtype) - _print_ep_routing_probe_result( - "dispatch_bwd", - probe_topk_idx_3d, - probe_dispatch, - expected_probe_dispatch, - tolerance=1.0e-3, - ) - - if _debug_moe_numerics or _debug_moe_input_grad: - d_sorted_x_stats = _debug_stable_stats(d_sorted_x, active_recv_rows) - d_recv_combine_stats = _debug_stable_stats( - d_recv_w_from_combine, active_recv_rows - ) - d_recv_intermediate_stats = _debug_stable_stats( - d_recv_w_from_intermediate, active_recv_rows - ) - d_recv_input_stats = _debug_stable_stats( - dispatch_weight_cotangent, active_recv_rows - ) - d_dispatch_stats = _debug_stable_stats(d_x_from_dispatch) - d_topk_stats = _debug_stable_stats(d_topk_w) - jax.debug.print( - "[TE dispatch-bwd stats] " - "token_in(absmean={token_in_absmean:.3e},std={token_in_std:.3e}," - "absmax={token_in_absmax:.3e},finite={token_in_finite:.6f}) " - "weight_combine(absmean={wc_absmean:.3e},absmax={wc_absmax:.3e}) " - "weight_ffn(absmean={wf_absmean:.3e},absmax={wf_absmax:.3e}) " - "weight_input(absmean={wi_absmean:.3e},absmax={wi_absmax:.3e}) " - "token_out(absmean={token_out_absmean:.3e},std={token_out_std:.3e}," - "absmax={token_out_absmax:.3e},finite={token_out_finite:.6f}) " - "weight_out(absmean={weight_out_absmean:.3e},std={weight_out_std:.3e}," - "absmax={weight_out_absmax:.3e},finite={weight_out_finite:.6f})", - token_in_absmean=d_sorted_x_stats[1], - token_in_std=d_sorted_x_stats[2], - token_in_absmax=d_sorted_x_stats[3], - token_in_finite=d_sorted_x_stats[4], - wc_absmean=d_recv_combine_stats[1], - wc_absmax=d_recv_combine_stats[3], - wf_absmean=d_recv_intermediate_stats[1], - wf_absmax=d_recv_intermediate_stats[3], - wi_absmean=d_recv_input_stats[1], - wi_absmax=d_recv_input_stats[3], - token_out_absmean=d_dispatch_stats[1], - token_out_std=d_dispatch_stats[2], - token_out_absmax=d_dispatch_stats[3], - token_out_finite=d_dispatch_stats[4], - weight_out_absmean=d_topk_stats[1], - weight_out_std=d_topk_stats[2], - weight_out_absmax=d_topk_stats[3], - weight_out_finite=d_topk_stats[4], - ordered=False, - ) - # ---------------- Routing bwd (global view) ---------------- # The cotangent on routing_weights is a sparse scatter into sparse_probs # at the selected_experts indices. @@ -1839,26 +1072,6 @@ def _ffn_bwd_body(*args): compute_aux_scores=False, ) - if _debug_moe_numerics or _debug_moe_input_grad: - sparse_prob_stats = _debug_stable_stats(d_sparse_probs) - logits_stats = _debug_stable_stats(d_logits_2d) - jax.debug.print( - "[TE router-bwd stats] " - "sparse_in(absmean={sparse_absmean:.3e},std={sparse_std:.3e}," - "absmax={sparse_absmax:.3e},finite={sparse_finite:.6f}) " - "logits_out(absmean={logits_absmean:.3e},std={logits_std:.3e}," - "absmax={logits_absmax:.3e},finite={logits_finite:.6f})", - sparse_absmean=sparse_prob_stats[1], - sparse_std=sparse_prob_stats[2], - sparse_absmax=sparse_prob_stats[3], - sparse_finite=sparse_prob_stats[4], - logits_absmean=logits_stats[1], - logits_std=logits_stats[2], - logits_absmax=logits_stats[3], - logits_finite=logits_stats[4], - ordered=False, - ) - # ---------------- 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 @@ -1894,32 +1107,6 @@ def _ffn_bwd_body(*args): 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 - if _zero_moe_input_grad: - d_x = jnp.zeros_like(d_x) - - if _debug_moe_numerics or _debug_moe_input_grad: - dispatch_stats = _debug_stable_stats(d_x_from_dispatch) - gate_stats = _debug_stable_stats(d_x_from_gate) - total_stats = _debug_stable_stats(d_x) - jax.debug.print( - "[TE input-grad stats] " - "dispatch(absmean={dispatch_absmean:.3e},std={dispatch_std:.3e}," - "absmax={dispatch_absmax:.3e}) " - "gate(absmean={gate_absmean:.3e},std={gate_std:.3e},absmax={gate_absmax:.3e}) " - "total(absmean={total_absmean:.3e},std={total_std:.3e}," - "absmax={total_absmax:.3e},finite={total_finite:.6f})", - dispatch_absmean=dispatch_stats[1], - dispatch_std=dispatch_stats[2], - dispatch_absmax=dispatch_stats[3], - gate_absmean=gate_stats[1], - gate_std=gate_stats[2], - gate_absmax=gate_stats[3], - total_absmean=total_stats[1], - total_std=total_stats[2], - total_absmax=total_stats[3], - total_finite=total_stats[4], - ordered=False, - ) # Pin output grads to the declared logical axes so downstream # optimizers see consistent shardings. From 421e2dd6dd52eb274009f9a9429bbdd78ba167a1 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 12 Aug 2026 16:44:45 -0700 Subject: [PATCH 38/44] Checkpoint quantized JAX MoE VJP residuals --- transformer_engine/jax/moe.py | 78 +++++++++++------------ transformer_engine/jax/quantize/tensor.py | 12 +++- 2 files changed, 50 insertions(+), 40 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 79a6b504c6..5179a7ce98 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -389,12 +389,20 @@ def _ffn_fwd_per_shard( expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) group_sizes_2d = group_sizes.reshape(1, num_local_experts) residuals = ( - sorted_x, - wi, + 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, - intermediate, - wo, + 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 @@ -402,12 +410,12 @@ def _ffn_fwd_per_shard( def _ffn_bwd_per_shard( d_expert_outputs_local: jnp.ndarray, - sorted_x: jnp.ndarray, - wi: jnp.ndarray, + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, gate_proj_out: jnp.ndarray, up_proj_out: jnp.ndarray, - intermediate: jnp.ndarray, - wo: jnp.ndarray, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, local_group_sizes: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, quantizer_sets: Tuple[QuantizerSet, QuantizerSet], @@ -427,25 +435,6 @@ def _ffn_bwd_per_shard( ) wgrad_group_active = (group_sizes > 0)[:, None, None] - # Recompute stateless grouped quantization inside backward. Besides being - # natural rematerialization for MXFP8BlockScaling, this keeps shard_map - # residuals in their logical BF16 shapes instead of exposing flattened - # quantized data and scale layouts to partition specs. - 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) - casted_intermediate = tex.grouped_quantize( - intermediate, fc2_quantizer_set.x, group_sizes, flatten_axis=-1 - ) - casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) - 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) - casted_intermediate_lhs_trans = casted_intermediate.get_tensor( - usage=TensorUsage.LHS_TRANS - ) - casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) - # wo bwd casted_d_eo = tex.grouped_quantize( d_eo_2d, @@ -756,13 +745,21 @@ def _moe_fwd_rule( ffn_in_specs += (bias_spec, bias_spec, bias_spec) ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) + # Quantized grouped tensors store their data, scales, and group metadata + # as physical buffers rather than in the source tensor's logical shape. + # A PartitionSpec used as a pytree prefix applies the same ownership to + # every array leaf of the grouped tensor: dispatched-token buffers belong + # to the compound batch shard, while expert-weight buffers belong to EP. + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) residuals_spec = ( - P(), - P(ep_axis, None, None), - P(), - P(), - P(), - P(ep_axis, None, None), + token_buffer_spec, + expert_buffer_spec, + token_matrix_spec, + token_matrix_spec, + token_buffer_spec, + expert_buffer_spec, ep2_spec, ) @@ -951,13 +948,16 @@ def _moe_bwd_rule( # ---------------- FFN bwd (per-shard via shard_map) ---------------- kernel_spec = P(ep_axis, None, None) bias_spec = P(ep_axis, None) + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) residuals_specs = ( - P(), - P(ep_axis, None, None), - P(), - P(), - P(), - P(ep_axis, None, None), + 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) 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 From 53dc6b2f1b4f329b1c6a72ca83a4b7a14b2f8285 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 12 Aug 2026 17:22:15 -0700 Subject: [PATCH 39/44] Test MXFP8 MoE VJP against JAX reference --- tests/jax/test_te_ep_moe.py | 128 +++++++++++----- tests/jax/test_te_ep_moe_mxfp8.py | 228 ----------------------------- transformer_engine/jax/flax/moe.py | 47 +++++- 3 files changed, 132 insertions(+), 271 deletions(-) delete mode 100644 tests/jax/test_te_ep_moe_mxfp8.py diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 0e9cb6598d..f178ffc6e8 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. @@ -45,6 +46,7 @@ os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") +os.environ.setdefault("NVTE_JAX_ENFORCE_V2_GROUPED_GEMM", "1") import sys from functools import partial @@ -118,8 +120,13 @@ def _read_mp_options(): ) 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 +155,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 @@ -194,10 +201,16 @@ def _compute_worst_case_recv_pr(): 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 + 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 + num_local_experts + * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) + * _ALIGN_SIZE ) return min(per_expert_bound, aligned_total_bound) @@ -222,7 +235,9 @@ 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(), @@ -275,8 +290,7 @@ def mesh(): def _pure_jax_moe_reference( x, gate_kernel, - wi_0, - wi_1, + wi, wo, expert_bias=None, *, @@ -312,11 +326,14 @@ 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 # 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 -- @@ -324,7 +341,9 @@ def _pure_jax_moe_reference( # 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_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: @@ -339,7 +358,9 @@ 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] @@ -365,6 +386,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 +399,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 +459,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 +533,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 +583,11 @@ def _make_inputs(key): ), ] +_QUANTIZATION_CASES = [ + pytest.param("bf16", id="bf16"), + 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 +606,11 @@ 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 +629,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 +640,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 +649,11 @@ 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 +670,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,11 +686,15 @@ 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}]" - assert np.any(g_te != 0.0), f"{name} grad identically zero [config={config}]" + 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" @@ -655,7 +705,10 @@ 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} " + f"[config={config}, quantization={quantization}]" + ), ) # d_x: the gradient propagated back to the previous layer. Checks @@ -677,7 +730,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 +763,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, @@ -729,7 +781,9 @@ 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" @@ -741,7 +795,9 @@ 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"): - g_local = np.asarray(jax.device_get(_unwrap(grads["params"][name]).addressable_data(0))) + 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/tests/jax/test_te_ep_moe_mxfp8.py b/tests/jax/test_te_ep_moe_mxfp8.py deleted file mode 100644 index aa1432b3ee..0000000000 --- a/tests/jax/test_te_ep_moe_mxfp8.py +++ /dev/null @@ -1,228 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. -"""Multiprocess MXFP8BlockScaling VJP coverage for the TE EP MoE path.""" - -import os -import sys -from contextlib import ExitStack - -os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") -os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") -os.environ.setdefault("NVTE_JAX_ENFORCE_V2_GROUPED_GEMM", "1") - -import jax -import jax.numpy as jnp -import numpy as np -import pytest -from flax.linen import partitioning as nn_partitioning -from jax.experimental import mesh_utils -from jax.sharding import Mesh, NamedSharding, PartitionSpec as P - - -def _read_mp_options(): - num_processes = 0 - process_id = 0 - for index, arg in enumerate(sys.argv): - if arg.startswith("--num-process="): - num_processes = int(arg.split("=", 1)[1]) - elif arg == "--num-process" and index + 1 < len(sys.argv): - num_processes = int(sys.argv[index + 1]) - elif arg.startswith("--process-id="): - process_id = int(arg.split("=", 1)[1]) - elif arg == "--process-id" and index + 1 < len(sys.argv): - process_id = int(sys.argv[index + 1]) - return num_processes, process_id - - -_NUM_PROCESSES, _PROCESS_ID = _read_mp_options() -if _NUM_PROCESSES <= 1: - pytest.skip("requires tests/jax/run_te_ep_moe.sh", allow_module_level=True) - -jax.distributed.initialize( - coordinator_address=os.environ.get("TE_EP_MOE_COORDINATOR_ADDRESS", "127.0.0.1:13457"), - num_processes=_NUM_PROCESSES, - process_id=_PROCESS_ID, - local_device_ids=_PROCESS_ID, -) - -from transformer_engine.common.recipe import MXFP8BlockScaling -from transformer_engine.jax.ep import ep_bootstrap -from transformer_engine.jax.moe import ( - get_moe_recv_capacity_per_rank, - moe, - record_ep_bootstrap_signature_for_moe, -) -from transformer_engine.jax.quantize import ( - QuantizeMeta, - QuantizeMetaSet, - QuantizerFactory, - QuantizerSet, - TensorSource, - get_quantize_config_with_recipe, -) -from transformer_engine.jax.sharding import MeshResource, global_shard_guard -from transformer_engine_jax import get_device_compute_capability - - -if get_device_compute_capability(0) < 100: - pytest.skip("MXFP8 grouped GEMM requires Blackwell", allow_module_level=True) - -EP_AXIS = "ep" -FSDP_AXIS = "fsdp" -EP_SIZE = 2 -FSDP_SIZE = jax.device_count() // EP_SIZE -NUM_EXPERTS = 8 -TOPK = 2 -BATCH = jax.device_count() * 2 -SEQ = 32 -HIDDEN = 128 -INTERMEDIATE = 128 -DTYPE = jnp.bfloat16 - -LOGICAL_AXIS_RULES = ( - ("batch", (FSDP_AXIS, EP_AXIS)), - ("exp", EP_AXIS), - ("embed", FSDP_AXIS), - ("mlp", None), -) - - -def _make_mxfp8_quantizer_sets(): - """Construct MaxText-shaped quantizers from MXFP8BlockScaling.""" - recipe = MXFP8BlockScaling() - config = get_quantize_config_with_recipe(recipe) - meta_set = QuantizeMetaSet(x=QuantizeMeta(), kernel=QuantizeMeta(), grad=QuantizeMeta()) - - def _set(n_token_groups, n_expert_groups): - token_set = QuantizerFactory.create_set( - fp8_recipe=recipe, - n_groups=n_token_groups, - quantize_meta_set=meta_set, - ) - expert_set = QuantizerFactory.create_set( - fp8_recipe=recipe, - n_groups=n_expert_groups, - quantize_meta_set=meta_set, - ) - for source in TensorSource: - assert config.get_scaling_mode(source).is_mxfp8_scaling - return QuantizerSet(x=token_set.x, kernel=expert_set.kernel, dgrad=token_set.dgrad) - - # Match MaxText: global dispatch groups include FSDP replicas, while - # expert weights have one group per global expert. - return tuple(_set(FSDP_SIZE * NUM_EXPERTS, NUM_EXPERTS) for _ in range(2)) - - -@pytest.fixture(scope="module") -def mesh(): - devices = mesh_utils.create_device_mesh((FSDP_SIZE, EP_SIZE)) - mesh_obj = Mesh(devices, axis_names=(FSDP_AXIS, EP_AXIS)) - max_tokens_per_rank = (BATCH // jax.process_count()) * SEQ - recv_capacity = get_moe_recv_capacity_per_rank( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - max_tokens_per_rank=max_tokens_per_rank, - ep_size=EP_SIZE, - ) - with mesh_obj, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ): - ep_bootstrap( - world_size=jax.process_count(), - rank=jax.process_index(), - num_experts=NUM_EXPERTS, - max_tokens_per_rank=max_tokens_per_rank, - recv_capacity_per_rank=recv_capacity, - 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, - hidden_dim=HIDDEN, - ep_size=EP_SIZE, - ) - return mesh_obj - - -def _context(mesh): - stack = ExitStack() - stack.enter_context(mesh) - stack.enter_context( - global_shard_guard(MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS)) - ) - stack.enter_context(nn_partitioning.axis_rules(LOGICAL_AXIS_RULES)) - return stack - - -def _global_array(value, mesh): - with mesh: - value = jax.jit( - lambda x: jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P())) - )(value) - value.block_until_ready() - return np.asarray(jax.device_get(value.addressable_data(0))) - - -def test_mxfp8_block_scaling_forward_and_vjp(mesh): - keys = jax.random.split(jax.random.PRNGKey(123), 5) - x = jax.random.normal(keys[0], (BATCH, SEQ, HIDDEN), DTYPE) - gate = jax.random.normal(keys[1], (HIDDEN, NUM_EXPERTS), DTYPE) / jnp.sqrt(HIDDEN) - wi = jax.random.normal(keys[2], (NUM_EXPERTS, HIDDEN, 2 * INTERMEDIATE), DTYPE) / jnp.sqrt( - HIDDEN - ) - wo = jax.random.normal(keys[3], (NUM_EXPERTS, INTERMEDIATE, HIDDEN), DTYPE) / jnp.sqrt( - INTERMEDIATE - ) - cotangent = jax.random.normal(keys[4], x.shape, DTYPE) - quantizer_sets = _make_mxfp8_quantizer_sets() - - def forward(x_arg, gate_arg, wi_arg, wo_arg): - output, _, total_recv_tokens = moe( - x_arg, - gate_arg, - wi_arg, - wo_arg, - None, - None, - None, - None, - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - apply_topk_weights_early=True, - quantizer_sets=quantizer_sets, - ep_axis=EP_AXIS, - data_parallelism_axes=(FSDP_AXIS,), - input_axes=("batch", None, None), - gate_kernel_axes=("embed", "exp"), - wi_kernel_axes=("exp", "embed", "mlp"), - wo_kernel_axes=("exp", "mlp", "embed"), - dtype=DTYPE, - ) - return output, total_recv_tokens - - def value_and_vjp(x_arg, gate_arg, wi_arg, wo_arg, cotangent_arg): - output, pullback = jax.vjp(lambda a, b, c, d: forward(a, b, c, d)[0], x_arg, gate_arg, wi_arg, wo_arg) - return output, pullback(cotangent_arg) - - with _context(mesh): - x = jax.lax.with_sharding_constraint( - x, NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None, None)) - ) - output, grads = jax.jit(value_and_vjp)(x, gate, wi, wo, cotangent) - jax.block_until_ready((output, grads)) - - assert output.shape == x.shape - assert output.dtype == DTYPE - output_np = _global_array(output, mesh) - assert np.all(np.isfinite(output_np)) - assert np.any(output_np != 0) - - expected_shapes = (x.shape, gate.shape, wi.shape, wo.shape) - for name, grad, shape in zip(("x", "gate", "wi", "wo"), grads, expected_shapes): - assert grad.shape == shape, f"{name} gradient shape mismatch" - grad_np = _global_array(grad, mesh) - assert np.all(np.isfinite(grad_np)), f"{name} gradient has NaN/Inf" - assert np.any(grad_np != 0), f"{name} gradient is identically zero" diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 3941a00b04..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 @@ -120,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 @@ -160,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: @@ -180,7 +185,6 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: ---------- inputs : jnp.ndarray ``[batch, sequence, hidden]``. - Returns ------- output : jnp.ndarray @@ -255,6 +259,34 @@ 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, @@ -275,6 +307,7 @@ 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, From 0dbf8a541b19dd8c67d6245c096018ccadd11fcf Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 13 Aug 2026 08:31:30 -0700 Subject: [PATCH 40/44] Keep MoE quantizer sets in global view --- .../jax/cpp_extensions/quantization.py | 17 ++- transformer_engine/jax/moe.py | 110 +++++++++++------- 2 files changed, 84 insertions(+), 43 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 7138cfcf40..858b662eff 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -1293,9 +1293,20 @@ 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/moe.py b/transformer_engine/jax/moe.py index 5179a7ce98..74e979a901 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -32,7 +32,6 @@ import math import warnings -from dataclasses import replace from functools import partial from typing import Any, Optional, Tuple, Union @@ -278,37 +277,69 @@ class _Ctx: # ============================================================================= -def _localize_grouped_quantizer_set( - quantizer_set: QuantizerSet, num_local_groups: int -) -> QuantizerSet: - """Resize stateless grouped quantizers for one shard-local FFN. +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. - The public MoE API receives quantizers sized for the global dispatch and - expert layouts. Under ``shard_map``, however, grouped quantize/GEMM sees - only the local experts. MXFP8 block quantizers are stateless and identical - per group, so selecting the corresponding number of entries preserves the - recipe while matching the local grouped operation. + 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.") - def _localize(quantizer): - if quantizer is None or not isinstance(quantizer, GroupedQuantizer): - return quantizer - if len(quantizer.quantizers) < num_local_groups: - raise ValueError( - "MoE grouped quantizer has fewer entries than the shard-local " - f"FFN requires: {len(quantizer.quantizers)} < {num_local_groups}." + 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." ) - return replace( - quantizer, - n_groups=num_local_groups, - quantizers=quantizer.quantizers[:num_local_groups], - ) - return QuantizerSet( - x=_localize(quantizer_set.x), - kernel=_localize(quantizer_set.kernel), - dgrad=_localize(quantizer_set.dgrad), - ) + 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( @@ -341,10 +372,7 @@ def _ffn_fwd_per_shard( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) - fc1_quantizer_set, fc2_quantizer_set = ( - _localize_grouped_quantizer_set(qset, num_local_experts) - for qset in quantizer_sets - ) + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets casted_sorted_x = tex.grouped_quantize( sorted_x, fc1_quantizer_set.x, @@ -423,16 +451,12 @@ def _ffn_bwd_per_shard( activation_type: str, apply_topk_weights_early: bool, has_bias: bool, - num_local_experts: int, ): """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) - fc1_quantizer_set, fc2_quantizer_set = ( - _localize_grouped_quantizer_set(qset, num_local_experts) - for qset in quantizer_sets - ) + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets wgrad_group_active = (group_sizes > 0)[:, None, None] # wo bwd @@ -586,6 +610,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 @@ -899,8 +928,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] - num_local_experts = num_experts // num_ep B, S, _ = x_shape K = num_experts_per_tok if not data_parallelism_axes: @@ -979,7 +1006,6 @@ def _ffn_bwd_body(*args): activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, has_bias=has_bias, - num_local_experts=num_local_experts, ) ( d_sorted_x_local, @@ -1265,7 +1291,11 @@ def moe( ``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. They are differentiable + 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`` From 0645751bec33779a050d844ce47fde783c9de35c Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 13 Aug 2026 08:48:54 -0700 Subject: [PATCH 41/44] Remove redundant MoE padding masks --- transformer_engine/jax/moe.py | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 74e979a901..76d7f72f8b 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -256,7 +256,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 @@ -457,7 +456,6 @@ def _ffn_bwd_per_shard( d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) recv_w_flat = recv_topk_weights_local.reshape(-1) fc1_quantizer_set, fc2_quantizer_set = quantizer_sets - wgrad_group_active = (group_sizes > 0)[:, None, None] # wo bwd casted_d_eo = tex.grouped_quantize( @@ -478,7 +476,6 @@ 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, group_sizes) if has_bias else None act_fn = _convert_to_activation_function(activation_type) @@ -529,9 +526,6 @@ 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) - ) if has_bias: 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) @@ -867,7 +861,6 @@ def _ffn_fwd_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, @@ -944,21 +937,6 @@ def _moe_bwd_rule( grad_pre_combine = jax.lax.with_sharding_constraint( grad_pre_combine, NamedSharding(mesh, ep3_spec) ) - # The EP kernel writes only the per-process packed expert prefix. Its - # over-allocation tail is intentionally left uninitialized, which is safe - # only while every downstream consumer is perfectly handle/group aware. - # Materialize the contract here so padding cannot leak through elementwise - # weighting, compiler fusion, or a later dispatch backward. - active_recv_rows = ( - jnp.arange(ctx.recv_topk_weights.shape[-1])[None, :] - < jnp.sum(ctx.token_counts, axis=-1, dtype=jnp.int32)[:, None] - ) - grad_pre_combine = jnp.where( - active_recv_rows[..., None], - grad_pre_combine, - jnp.zeros_like(grad_pre_combine), - ) - 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. From bc28b24b20613d3a4fafdec9cac972710e30206e Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 13 Aug 2026 16:47:11 -0700 Subject: [PATCH 42/44] Remove V2 grouped GEMM enforcement env var from MoE EP tests --- tests/jax/test_te_ep_moe.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index f178ffc6e8..18cdec1b62 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -46,7 +46,6 @@ os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") -os.environ.setdefault("NVTE_JAX_ENFORCE_V2_GROUPED_GEMM", "1") import sys from functools import partial @@ -110,12 +109,10 @@ 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 ec09e5a10b480c0957434a122d7f9501431b7cfa Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Fri, 14 Aug 2026 07:13:03 -0700 Subject: [PATCH 43/44] Gate mxfp8 test on blackwell --- tests/jax/test_te_ep_moe.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 18cdec1b62..1d9ad1034b 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -582,9 +582,11 @@ def _quantization_recipe(quantization): _QUANTIZATION_CASES = [ pytest.param("bf16", id="bf16"), - pytest.param("mxfp8", id="mxfp8"), ] +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.""" From 79d7cda6df4ba3b7cdd08b8eef9c82a0d213e0a7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:14:24 +0000 Subject: [PATCH 44/44] [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 | 55 +++++-------------- .../jax/cpp_extensions/quantization.py | 10 ++-- transformer_engine/jax/moe.py | 44 ++++++--------- 3 files changed, 35 insertions(+), 74 deletions(-) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index 1d9ad1034b..f92519fe6b 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -198,16 +198,10 @@ def _compute_worst_case_recv_pr(): 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 + 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 + num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE ) return min(per_expert_bound, aligned_total_bound) @@ -232,9 +226,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 +315,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 @@ -338,9 +328,7 @@ def _pure_jax_moe_reference( # 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_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: @@ -355,9 +343,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] @@ -607,9 +593,7 @@ class TestTeEpMoeForward: @pytest.mark.parametrize("config", _CONFIGS) @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) def test_forward(self, mesh, config, quantization): - block = _make_block( - **config, quantization_recipe=_quantization_recipe(quantization) - ) + 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)) @@ -650,9 +634,7 @@ class TestTeEpMoeBackward: @pytest.mark.parametrize("config", _CONFIGS) @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) def test_backward(self, mesh, config, quantization): - block = _make_block( - **config, quantization_recipe=_quantization_recipe(quantization) - ) + 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) @@ -688,12 +670,8 @@ def loss_fn(params, x): 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}]" - assert np.any( - g_te != 0.0 - ), f"{name} grad identically zero [config={config}]" + 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" @@ -705,8 +683,7 @@ def loss_fn(params, x): atol=atol, rtol=rtol, err_msg=( - f"grad parity breach on {name} " - f"[config={config}, quantization={quantization}]" + f"grad parity breach on {name} [config={config}, quantization={quantization}]" ), ) @@ -780,9 +757,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" @@ -795,8 +770,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", "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" diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 858b662eff..2d56eb4769 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -1300,13 +1300,11 @@ def grouped_quantize( # 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}" - ) + 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}" - ) + 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/moe.py b/transformer_engine/jax/moe.py index 76d7f72f8b..eae56872e8 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -80,7 +80,9 @@ def get_moe_recv_capacity_per_rank( 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") + 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: @@ -95,18 +97,12 @@ def get_moe_recv_capacity_per_rank( 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_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 + aligned_total_bound = ((padded_total_bound + alignment - 1) // alignment) * alignment per_expert_bound = ( - num_local_experts - * ((tokens_per_ep_group + alignment - 1) // alignment) - * alignment + 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: @@ -116,9 +112,7 @@ def get_moe_recv_capacity_per_rank( 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 + num_local_experts * ((balanced_per_expert + alignment - 1) // alignment) * alignment ) requested = math.ceil(balanced_aligned * recv_capacity_factor) requested = ((requested + alignment - 1) // alignment) * alignment @@ -416,20 +410,12 @@ def _ffn_fwd_per_shard( expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) group_sizes_2d = group_sizes.reshape(1, num_local_experts) residuals = ( - casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint( - fc1_quantizer_set.x - ), - casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint( - fc1_quantizer_set.kernel - ), + 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.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint( - fc2_quantizer_set.x - ), - casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint( - fc2_quantizer_set.kernel - ), + 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 @@ -629,7 +615,8 @@ def _moe_fwd_rule( 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 {recv_pr}" + f"recv_capacity_per_rank must be a positive multiple of {_ALIGN_SIZE}, got" + f" {recv_pr}" ) _te_ep_assert_compatible_bootstrap( @@ -977,6 +964,7 @@ def _moe_bwd_rule( ctx.local_group_sizes, ctx.recv_topk_weights, ] + def _ffn_bwd_body(*args): grads = _ffn_bwd_per_shard( *args, @@ -1039,7 +1027,9 @@ def _ffn_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