[JAX] Grouped quant+GEMM custom partitioning rules

[jberchtold-nvidia]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

[greptile-apps]

Greptile Summary

This PR adds JAX custom partitioning rules (partition and shardy_sharding_rule) to GroupedGemmPrimitive and GroupedQuantizePrimitive, enabling Expert Parallelism (EP) + FSDP sharding for grouped quant+GEMM operations outside of shard_map. It also removes the "FSDP not supported" assertion from _grouped_dense_bwd_rule and wires up the actual allgather/reduce-scatter logic for FSDP in the backward pass.

• gemm.py / quantization.py: Partition spec helpers propagate EP/FSDP axes and compute local sizes; GroupedGemmPrimitive.partition handles FSDP all-gather on the RHS weight and optional psum on the output.
• dense.py: Lifts the FSDP assertion, adds a _is_manual_mesh_axis guard, and adds a with_sharding_constraint on wgrad to propagate the correct EP+FSDP spec.
• module.py / sharding.py: Contracting dims switched to (-1,) for 3-D EP-leading input support; ep_resource added to MeshResource.

Confidence Score: 3/5

The core partitioning logic is new and complex; the backward FSDP path in dense.py has a silent EP-axis overwrite that could produce an incorrect sharding constraint in an edge case, and the 1-D size heuristic in gemm.py assumes a specific sharding orientation without enforcing it.

The EP-axis overwrite in _grouped_dense_bwd_rule silently produces a wrong sharding constraint when kernel_fsdp_axis_idx == 0, and the _local_2d_sizes_from_spec heuristic could compute incorrect local sizes if sharding falls on the right dimension. Both issues are in newly-written paths lacking targeted test coverage.

transformer_engine/jax/dense.py (backward FSDP wgrad sharding constraint) and transformer_engine/jax/cpp_extensions/gemm.py (_local_2d_sizes_from_spec fallback logic) deserve a close second look.

Important Files Changed

Filename	Overview
transformer_engine/jax/cpp_extensions/gemm.py	Adds spec-manipulation helpers and GroupedGemmPrimitive.partition / shardy_sharding_rule; helper functions are duplicated with quantization.py and the _local_2d_sizes_from_spec fallback may divide the wrong 2D size.
transformer_engine/jax/cpp_extensions/quantization.py	Adds GroupedQuantizePrimitive.partition and shardy_sharding_rule; logic looks correct for the supported MXFP8/tensor-scaling paths.
transformer_engine/jax/dense.py	Removes the FSDP assertion and implements backward FSDP allgather/reduce-scatter; has a silent EP-axis overwrite when kernel_fsdp_axis_idx == 0.
transformer_engine/jax/flax/module.py	Contracting dims switched from ((1,), (1,)) to ((-1,), (1,)) to support the new 3-D EP reshape path.
transformer_engine/jax/sharding.py	Adds ep_resource field to MeshResource dataclass; trivial, safe change.
tests/jax/test_grouped_gemm_partitioning.py	New unit tests for partition/shardy spec correctness and an end-to-end MXFP8 EP+FSDP test on a 1x1 mesh.
tests/jax/test_multi_process_distributed_grouped_gemm.py	Switches to MXFP8_1D_SCALING, fixes the pre-existing jnp.allclose no-assert bug, adds a new helper that references a module-level mesh global.

Sequence Diagram

sequenceDiagram
    participant JAX as JAX Compiler
    participant QP as GroupedQuantizePrimitive.partition
    participant GP as GroupedGemmPrimitive.partition
    participant SI as sharded_impl (per device)

    JAX->>QP: (x, scale, group_sizes) with EP+FSDP specs
    QP->>QP: _parse_partition_specs() derive flat/group output specs
    QP-->>JAX: arg_shardings, out_shardings, sharded_impl
    JAX->>SI: local x shard
    SI->>SI: GroupedQuantizePrimitive.impl()
    SI->>SI: _pad_or_slice_to_shape(scale_inv, local_shape)
    SI->>SI: pmax amax over dp/fsdp
    SI-->>JAX: rowwise_out, colwise_out, scale_invs, amax

    JAX->>GP: lhs, rhs_weight with EP+FSDP specs
    GP->>GP: inject EP into rhs dim-0 if missing
    GP->>GP: strip FSDP from rhs gather_rhs_fsdp
    GP-->>JAX: arg_shardings, out_sharding, sharded_impl
    JAX->>SI: local lhs/rhs shards
    SI->>SI: GroupedGemmPrimitive.impl() local GEMM
    SI->>SI: psum(out, reduce_axis) if needed
    SI-->>JAX: out

    Note over JAX,SI: Backward pass inside shard_map only
    JAX->>SI: _grouped_dense_bwd_rule
    SI->>SI: _is_manual_mesh_axis check
    alt FSDP axis is manual
        SI->>SI: allgather or psum dgrad
        SI->>SI: psum_scatter or psum wgrad
    end
    SI->>SI: with_sharding_constraint wgrad EP+FSDP spec

Loading

Comments Outside Diff (2)

1. tests/jax/test_multi_process_distributed_grouped_gemm.py, line 427-429 (link)

   Function implicitly depends on module-level mesh variable

   run_grouped_dense_mxfp8_ep_fsdp_outside_shard_map references mesh as a free variable. It is only defined in the if __name__ == "__main__": block (line 243), which assigns it to the module global scope. If this function is ever renamed to start with test_ (or called from any other context), it will raise NameError: name 'mesh' is not defined. Consider accepting mesh as a parameter to make the dependency explicit.

2. transformer_engine/jax/cpp_extensions/gemm.py, line 601-621 (link)

   Ambiguous heuristic when both 2D sizes are divisible by the shard count

   When the tensor is 1-D, _local_2d_sizes_from_spec picks left_size // spec_size whenever left_size % spec_size == 0, even if the actual sharding is along the right/hidden dimension. A mismatch would silently produce incorrect local sizes passed to the GEMM primitive. A comment documenting the assumption (EP always shards the left/token dimension) would at least make the contract explicit.

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

[jberchtold-nvidia]

/te-ci L1 jax