[JAX] Grouped quant+GEMM custom partitioning rules by jberchtold-nvidia · Pull Request #3058 · NVIDIA/TransformerEngine

jberchtold-nvidia · 2026-05-28T20:59:21Z

Description

Adds custom partitioning rules to the grouped quantization and grouped GEMM primitives to support DP/FSDP and EP shardings

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Add custom partitioning to grouped quantize and grouped GEMM primitives
Add tests to validate custom partitioning functions directly, then a real test on a JIT'd program to verify the shardings are correct

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

…mm-custom-partition-rules

greptile-apps · 2026-05-28T21:07:32Z

Greptile Summary

This PR adds JAX custom partitioning rules to GroupedQuantizePrimitive and GroupedGemmPrimitive to support DP/FSDP/EP shardings without requiring shard_map. The MoE FFN is refactored from a per-shard shard_map body to a global-view execution path backed by these new rules, and the gated-SwiGLU wi_0/wi_1 pair is merged into a single contiguous wi buffer.

GroupedQuantizePrimitive.partition and GroupedGemmPrimitive.partition are added with spec-filtering helpers in sharding.py to propagate only EP/DP/FSDP axes.
moe.py drops shard_map in favour of with_sharding_constraint pinning; quantizer_sets are threaded through custom_vjp residuals for FP8/MXFP8 recipe state.
flax/module.py fixes a contracting-dim bug and dense.py removes the superseded kernel_fsdp_info parameter.

Confidence Score: 3/5

Two correctness issues need fixing before merging: a duplicate dataclass field that reorders the MeshResource constructor, and a removed wgrad zero-fill that can produce NaN gradients when any expert receives zero tokens.

The MeshResource dataclass has ep_resource declared twice, silently changing positional argument ordering. More critically, the explicit wgrad zero-fill guard for 0-token expert groups was removed from both d_wo and d_wi_combined in _ffn_bwd_global, yet the forward-rule comment still says 'the backward masks skipped wgrad groups' — indicating an accidental omission. Under normal EP configs where some experts receive zero tokens, cuBLAS leaves those wgrad slices uninitialised, which NaN-propagates through _fold_dp_groups and corrupts the optimizer state.

transformer_engine/jax/sharding.py for the duplicate ep_resource field, and transformer_engine/jax/moe.py (_ffn_bwd_global) for the missing wgrad masking.

Important Files Changed

Filename	Overview
transformer_engine/jax/sharding.py	Adds spec-manipulation utilities; also adds a duplicate `ep_resource` field to `MeshResource` that silently reorders the dataclass constructor's positional arguments.
transformer_engine/jax/cpp_extensions/quantization.py	Adds `partition` and `shardy_sharding_rule` to `GroupedQuantizePrimitive`; introduces `uniform_groups` and 3D MXFP8 scale carrier shapes.
transformer_engine/jax/cpp_extensions/gemm.py	Adds `partition` and `shardy_sharding_rule` to `GroupedGemmPrimitive` with EP/FSDP/DP spec inference.
transformer_engine/jax/moe.py	Replaces per-shard shard_map FFN with global-view custom-partitioned primitives; removal of wgrad zero-fill for 0-token expert groups is a regression.
transformer_engine/jax/flax/moe.py	Merges wi_0/wi_1 into one wi parameter.
transformer_engine/jax/flax/module.py	Fixes contracting_dims to ((-1,), (1,)) for correct H-axis contraction on 3-D inputs.
transformer_engine/jax/dense.py	Removes obsolete kernel_fsdp_info parameter.
transformer_engine/jax/quantize/tensor.py	Relaxes scale_inv ndim assertion to allow 3-D carriers for the uniform MXFP8 path.
tests/jax/test_distributed_grouped_gemm.py	New partition tests for EP/FSDP/DP specs; 0-token-group edge case not covered.

_{Reviews (5): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

jberchtold-nvidia · 2026-06-02T22:02:05Z

/te-ci L1 jax

phu0ngng · 2026-07-06T10:14:16Z

+def _flat_data_spec(input_spec):
+    return (merge_axis_specs(*input_spec),)


Do we really need a helper for this simple task?

phu0ngng · 2026-07-06T10:15:15Z

+def merge_axis_specs(*axis_specs):
+    """Merge dimension axis specs while preserving first-seen axis order."""


For (tupleA, ) + (tupleB,) ?

phu0ngng · 2026-07-06T10:19:34Z

+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)))


Not sure if this is needed.

In every partition(), we already have something like this:

x_spec = get_padded_spec(arg_infos[0])

You can always do x_spec[:flatten_axis].

phu0ngng · 2026-07-06T10:27:24Z

+                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])


Is it correct that this is needed for the weight only?

phu0ngng · 2026-07-06T10:30:18Z

        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):


This helper is called only once, so I don't think we need it as a separate helper.

The reason GemmPrimitive has it is that in the past, we needed to generate the partition spec twice (i.e., in infer_sharding_for_operands besides partition).

phu0ngng · 2026-07-06T12:01:09Z

+        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)
+        _warn_if_axes_ignored("x", original_x_spec, x_spec)


I would merge supported_grouped_partition_axes, filter_spec_axes, and _warn_if_axes_ignored into a single helper and make it more generic, i.e.

def filter_unsupported_partition_axes(input_specs, supported_axes)

Then we can define a global supported_axes for MoE inputs or local specs case by case. We can still reuse this generic function for other cases. We should print the warning when we overwrite an axis, not after, to avoid TOCTOU.

The other question here is whether we really should filter the unsupported spec out and replicate that axis, or should we validate and error out. I'm leaning toward erroring out so that users will really know what they are doing.

phu0ngng · 2026-07-06T12:12:51Z

+                for info, spec in zip(result_infos, out_specs)
+            )
+            if result_infos
+            else (None,) * len(out_specs)


When will result_infos be None?

phu0ngng · 2026-07-06T12:35:35Z

+            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:],
+                )


I'm not sure if we should insert/force the rhs to be partitioned by ep. I'm in favor of validating if the inputs are partitioned correctly, then continue.

phu0ngng · 2026-07-06T12:37:14Z

+                    *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:])


Same, bias should be expected to be partitioned as in the RHS non-contracting specs.

phu0ngng · 2026-07-06T12:42:42Z

+                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)


I think the scale and the bias should be validated to ensure that it has the same/the right partition wrt to rhs_data, then here we should check for one instead of OR.

phu0ngng · 2026-07-06T12:54:06Z

+        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)


Here, we should check if there is a common axis in the contracting dims, not in all dims.

phu0ngng · 2026-07-06T12:55:28Z

+        if reduce_axis is not None and gather_rhs_fsdp:
+            reduce_axis = None


I think we can add a comment here that when father_rhs_fsdp = True, it's not the dW GGEMM, thus, we should not expect an AllReduce or ReduceScatter. But in this case, we should error out if reduce_axis is not None.

phu0ngng · 2026-07-06T13:00:39Z

+        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)


Not sure about what we are doing here for the dW GEMM.
At least, should we have a validation on what is the input specs we are expecting here?

phu0ngng · 2026-07-06T13:02:48Z

+        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(


I could not follow why we need 2D sizes here.

This was to attempt to WAR a Shardy bug that limited the num elements on a single data axis in the tensor's global shape to INT32_MAX. There is now a PR to fix this here so I will go back to 1D
openxla/shardy#1468

tdophung · 2026-07-08T18:56:14Z

+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,
+        )


This looks exactly like the one in gemm.py, maybe we define it somewhere common and reuse for both

tdophung · 2026-07-08T22:28:54Z

+        assert not _spec_contains_axis(spec, "tp")
+
+
+def test_grouped_gemm_reduce_axis_skips_ep_and_uses_dp():


there is no full test (execute a JIT + VJP run) that uses "dp" axis. The only full test test_grouped_dense_mxfp8_ep_fsdp_outside_shard_map_single_process I see does not use the dp axis. And this test on this line only check if reduce axis is "dp"

tdophung · 2026-07-08T22:49:02Z

+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, 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", None, "fsdp")),
+                _arg_info(mesh, (8,), ("expert",)),
+                _arg_info(mesh, (8,), ("expert",)),
+            ),
+            (),
+        )
+
+    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",)
+    assert _normalize_spec(specs[2]) == ("expert",)
+    assert _normalize_spec(specs[4]) == ("expert",)
+
+
+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, arg_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",)),
+            ),
+            (),
+        )
+
+    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",)
+    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_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(
+        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,
+                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", "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", "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")
+
+
+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_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"


This was really verbose to read. I wonder if you can refactor it into:

a wrapper over the partition() call that takes in (mesh, input spec, layout, scaling) and return plain spec tuples that can apply for all the different cases here

parametrize with pytest.mark.parametrize

a common helper function to build the big arg_infos tuple for each case

rename test cases

tdophung · 2026-07-09T03:35:15Z

+        )
+
+    assert reduce_axis == "dp"
+


Also, I notice that the tests here focuses on testing the subtlety of the axis stripping/merging logic, which diverges from the usual pattern of other tests/jax/test_distributed_*.py, where we mostly have a pure jax implementation results as ground truth to compare with the defined primitive's result. Is this inconsistency intentional as you implement way more axis handling logic in the _parse_partition_specs in partition than in other primitives?

tdophung · 2026-07-09T03:47:16Z

            additional_args: Either
                * group_offsets: 1D array containing offsets for each group (not yet implemented)
                OR
                * alpha: 1D array of shape (G,) containing alpha values for each group
                * beta: 1D array of shape (G,) containing beta values for each group


according to this, the addittional_arg_0 and additional_arg_1 used in the partition func definition below are alpha and beta here, should we name it so, so it is consistent and less confusing? it also make it hard to infer why these seem to be per-group array with sharding spec inherited from active_group_spec

…mm-custom-partition-rules # Conflicts: # qa/L1_jax_distributed_unittest/test.sh

for more information, see https://pre-commit.ci

greptile-apps · 2026-07-21T17:18:33Z

    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_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1)
+    d_wi_combined = jax.lax.with_sharding_constraint(d_wi_combined, grouped_weight_sharding)


Wgrad masking removed for 0-token expert groups

The old shard_map path explicitly zeroed wgrad output for 0-token groups. That masking is gone here (and identically for d_wi_combined at line ~467). cuBLAS grouped GEMM still leaves uninitialized memory for size-0 groups, so when _fold_dp_groups sums those slices across DP replicas, a single NaN-filled shard propagates NaN to all replicas and ultimately to the optimizer. The forward rule comment at line 700–703 still says "the backward masks skipped wgrad groups", confirming this is an unintended regression. A jnp.where(group_sizes.reshape(-1) > 0, ...) guard on d_wo and d_wi_combined before the with_sharding_constraint call is needed.

jberchtold-nvidia added 2 commits May 26, 2026 16:12

Add grouped GMM custom partitioning rules

b14df81

Add outside shard_map grouped GMM backward test

60a0b50

jberchtold-nvidia marked this pull request as draft May 28, 2026 20:59

Merge remote-tracking branch 'github-upstream/main' into jberchtold/g…

c3b00f5

…mm-custom-partition-rules

greptile-apps Bot reviewed May 28, 2026

View reviewed changes

Comment thread transformer_engine/jax/dense.py Outdated

Comment thread transformer_engine/jax/cpp_extensions/gemm.py Outdated

jberchtold-nvidia added 2 commits May 28, 2026 16:07

Use 4 GPU mesh for grouped GEMM partitioning test

786fa1d

progress

ff0407d

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

jberchtold-nvidia force-pushed the jberchtold/gmm-custom-partition-rules branch from 027b3e6 to ff0407d Compare June 1, 2026 15:46

Add warnings

1bd6b54

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

jberchtold-nvidia force-pushed the jberchtold/gmm-custom-partition-rules branch from 70893b7 to 1bd6b54 Compare June 2, 2026 20:11

Remove kernel_fsdp_info

3c30c9b

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

jberchtold-nvidia force-pushed the jberchtold/gmm-custom-partition-rules branch from 59ff8e0 to 3c30c9b Compare June 2, 2026 20:25

pre-commit-ci Bot and others added 5 commits June 2, 2026 20:26

[pre-commit.ci] auto fixes from pre-commit.com hooks

273e066

for more information, see https://pre-commit.ci

Remove unnecessary reshape

fe906fd

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

Remove multi-process grouped GEMM tests

e76d20a

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

Refactor helpers into sharding.py

1aa1e82

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

Rename distributed grouped GEMM tests

6fd04b5

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

jberchtold-nvidia requested review from phu0ngng and tdophung June 3, 2026 23:15

jberchtold-nvidia marked this pull request as ready for review June 3, 2026 23:15

phu0ngng reviewed Jul 6, 2026

View reviewed changes

tdophung reviewed Jul 9, 2026

View reviewed changes

jberchtold-nvidia added 5 commits July 16, 2026 14:45

Merge remote-tracking branch 'github-upstream/main' into jberchtold/g…

5e8531b

…mm-custom-partition-rules # Conflicts: # qa/L1_jax_distributed_unittest/test.sh

Run JAX TE-EP MoE FFN without shard_map

81d45af

Support MXFP8 quantization in global-view JAX MoE

3672316

Quantize JAX MoE weights before FSDP gather

330b2cc

Simplify TE MoE early weighting padding handling

d88cbc1

Remove redundant masks from TE MoE ragged paths

1b783b9

jberchtold-nvidia force-pushed the jberchtold/gmm-custom-partition-rules branch from 7e6baeb to 1b783b9 Compare July 20, 2026 22:56

Keep ragged MoE scales on local shards

35cef80

jberchtold-nvidia force-pushed the jberchtold/gmm-custom-partition-rules branch from c4d4555 to 35cef80 Compare July 21, 2026 16:52

Use contiguous wi in JAX MoE VJP

f4beba5

jberchtold-nvidia force-pushed the jberchtold/gmm-custom-partition-rules branch from 024bb80 to f4beba5 Compare July 21, 2026 17:09

[pre-commit.ci] auto fixes from pre-commit.com hooks

2b9079b

for more information, see https://pre-commit.ci

jberchtold-nvidia marked this pull request as draft July 21, 2026 17:12

greptile-apps Bot reviewed Jul 21, 2026

View reviewed changes

		def _flat_data_spec(input_spec):
		return (merge_axis_specs(*input_spec),)

		def merge_axis_specs(*axis_specs):
		"""Merge dimension axis specs while preserving first-seen axis order."""

		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 reduce_axis is not None and gather_rhs_fsdp:
		reduce_axis = None

		assert not _spec_contains_axis(spec, "tp")


		def test_grouped_gemm_reduce_axis_skips_ep_and_uses_dp():

Uh oh!

Conversation

jberchtold-nvidia commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Uh oh!

Uh oh!

Uh oh!

jberchtold-nvidia commented Jun 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jberchtold-nvidia commented May 28, 2026 •

edited

Loading

greptile-apps Bot commented May 28, 2026 •

edited

Loading