[JAX] Keep MoE aux-loss cotangent scalar

[tdophung]

Description

Fix the distributed JAX MoE auxiliary-loss backward pass by preserving its scalar cotangent instead of reshaping it to (1,).
The Shardy rule correctly declares this operand as rank 0, so the reshape caused distributed lowering to fail with a rank mismatch. The FFI already accepts a scalar buffer.

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

removing reshape of cotangent in MoE aux loss bwd pass to (1,)

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 fixes the distributed JAX MoE auxiliary-loss backward pass by removing an incorrect g.reshape(1) call that was converting the scalar cotangent (g, rank-0) into a shape-(1,) tensor (rank-1). Under Shardy, the FusedMoEAuxLossBwdPrimitive.shardy_sharding_rule declares the grad_aux_loss operand as rank-0, so the reshape caused a rank-mismatch failure during distributed lowering; the FFI C++ backend already accepts a scalar buffer natively.

• Root cause removed: g.reshape(1) in _fused_moe_aux_loss_bwd is deleted; the raw scalar cotangent g is now passed directly to fused_moe_aux_loss_bwd, matching the rank-0 declaration in the Shardy rule.
• No API surface changed: the FusedMoEAuxLossBwdPrimitive.abstract method only reads grad_aux_loss_aval.dtype and is shape-agnostic, so the abstract evaluation is unaffected by the fix.

Confidence Score: 5/5

The change is a one-line deletion that aligns the cotangent shape with the existing Shardy rule and FFI contract; no new behaviour is introduced and the distributed backward test already covers this path.

The fix removes a single reshape that was causing rank mismatches under Shardy distributed lowering. The abstract evaluation in the primitive is shape-agnostic (reads only dtype), the FFI backend already accepts a scalar, and the Shardy rule explicitly documents the rank-0 expectation with an explanatory comment.

No files require special attention; the single changed hunk in transformer_engine/jax/router.py is self-contained and consistent with the primitive's declared sharding contract.

Important Files Changed

Filename	Overview
transformer_engine/jax/router.py	Removes the erroneous g.reshape(1) in _fused_moe_aux_loss_bwd; the scalar cotangent g is now passed directly to fused_moe_aux_loss_bwd, aligning with the rank-0 Shardy declaration and fixing distributed lowering failures.

Sequence Diagram

sequenceDiagram
    participant JAX as JAX Autograd
    participant BWD as _fused_moe_aux_loss_bwd
    participant API as fused_moe_aux_loss_bwd
    participant PRIM as FusedMoEAuxLossBwdPrimitive
    participant FFI as C++ FFI kernel

    JAX->>BWD: g (rank-0 scalar cotangent)
    Note over BWD: BEFORE: grad_aux_loss = g.reshape(1) shape (1,) rank-1, Shardy expects rank-0
    Note over BWD: AFTER: pass g directly shape () rank-0, Shardy rule satisfied
    BWD->>API: fused_moe_aux_loss_bwd(const_buf, tokens_per_expert, g, num_tokens)
    API->>PRIM: outer_primitive.bind(...)
    PRIM->>PRIM: shardy_sharding_rule declares grad_aux_loss as rank-0
    PRIM->>FFI: ffi.ffi_lowering scalar buffer accepted natively
    FFI-->>JAX: grad_probs [num_tokens, num_experts]

Loading

_{Reviews (1): Last reviewed commit: "[JAX] Keep MoE aux-loss cotangent scalar" | Re-trigger Greptile}

[tdophung]

/te-ci jax

[tdophung]

/te-ci L1 jax

[tdophung]

/te-ci L2 jax