TE MoE Return Expert Counts and Custom Routing for DSv4 Support

[dinodeep]

Description

This PR adds support for DSv4 features in TE MoE block which includes the following

• Providing expert count metadata (this is useful for users to perform routing bias updates, as DSv4 does)
• Adding support for custom routing policies with moe_from_routing (allows for hash routing)

These changes will be useful to make TE MoE block more flexible and configurable for new, upcoming models such as DSv4.

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:

• If desired by the user, allow the caller of moe to pass collect_expert_counts=True which will compute the expert counts and return them to the caller as a part of the metadata.
  • Tested by validating expert counts in MoE block tests
• A new API function moe_from_routing(...) which allows the user to provide their expert assignments and weights. Given this routing information, moe_from_routing will re-use functionality of moe except the gate computation. Backwards gradients will be computed through weight.
  • Tested by implementing a simple hash routing in test and validating the TE MoE block with the pure JAX implementation
  • Had to slightly increase tolerance to adjust for accumulation differences

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]

Confidence Score: 3/5

[Medium risk] Mixture-of-Experts layer gains expert counting and custom routing modes.

The PR does not appear safe to merge while existing callers can fail on the changed return shape and externally routed backward remains unresolved.

Findings

1.  Existing callers cannot unpack results ▶
2.  Externally routed backward fails ▶

Summary

The PR adds per-expert assignment counts and supports caller-supplied routing alongside TE-computed routing.

• Refactors the Flax MoE block to pass router configuration through RouterComputationInfo.
• Adds tests for expert counts, hash routing, and gradients.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[MoE inputs] --> B{Router info}
  B -->|RouterComputationInfo| C[Gate and top-k]
  B -->|RoutingMapInfo| D[Caller-supplied indices and weights]
  C --> E[Expert dispatch]
  D --> E
  E --> F[Expert FFN]
  F --> G[Combine and return output, loss, receive total, counts]

Loading

_{Reviews (11) · Last reviewed commit: "Fixing tolerance assignment" · Reviewed by Greptile}

[greptile-apps]

Existing callers cannot unpack results

moe() now returns four values instead of three, and _MoEBlock passes that result through. Existing callers that unpack the previous three-value result will raise ValueError: too many values to unpack, making this a breaking change to the public API.

[ptrendx]

@dinodeep Please sign every commit in this PR. See https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst#sign-your-work

[greptile-apps]

Externally routed backward fails

When a caller differentiates moe() with RoutingMapInfo, the backward rule returns routing_indices=None even though the forward call supplied an integer array. That changes the JAX pytree structure of the router argument, so the custom VJP cannot match the gradient to its input and the backward pass fails. Return a structurally matching cotangent for the indices.

[jberchtold-nvidia]

/te-ci L1 jax

[jberchtold-nvidia]

/te-ci L1 jax

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[jberchtold-nvidia]

/te-ci L1 jax