fix: handle zero-size tensors in MoE token dispatchers

[callum-ward-inflection]

What does this PR do ?

Fixes MoE token dispatcher crash when an Expert Parallelism rank receives zero tokens from the router.

When an EP rank receives zero tokens from the router, the fused permute/unpermute autograd pair breaks and .view() crashes on a zero-size tensor. Switch to unfused path symmetrically, reconnect gradient graph for backward collectives, and guard .view() in all three dispatcher classes.

Fixes: #1877

Changes

Three changes to megatron/core/transformer/moe/token_dispatcher.py:

1. Symmetric unfused permute/unpermute for empty EP ranks — TE's fused_permute saves state that fused_unpermute reads during backward. With zero tokens this state is invalid. Both dispatch_postprocess and combine_preprocess detect zero-token ranks and fall back to the unfused PyTorch path together (they must match — mixing fused permute with unfused unpermute crashes due to incompatible index formats).

2. Gradient connectivity for backward collectives — Unfused unpermute with zero tokens returns a tensor disconnected from the autograd graph. During backward, distributed collectives (AllGather/ReduceScatter) need every rank to participate. A detached tensor means one rank never triggers its collective and all others hang (NCCL timeout). Reconnected with unpermuted_local_hidden + hidden_states.sum() * 0.

3. Safety guard on .view() for all three dispatcher classes — MoEAllGatherTokenDispatcher, MoEAlltoAllTokenDispatcher, and MoEFlexTokenDispatcher all have .view(self.hidden_shape) that crashes on zero-size input. Guarded with a zero tensor of the correct shape.

Reproduction

Reproducible with openai/gpt-oss-20b (32 experts, hidden_size=2880) SFT with TP=4, EP=2 on 16 GPUs. The pretrained router's weight distribution combined with reduced tokens-per-rank from TP=4 sequence parallelism causes certain EP ranks to consistently receive zero tokens.

Testing

Config	Result
TP=4, EP=2 (previously crashing)	PASS — 20 steps
TP=2, EP=2	PASS
TP=1, EP=2	PASS
TP=4, EP=4	PASS

Contribution process

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Victarry]

Hi @callum-ward-inflection , I think the ideal way to resolve this issue is to fix the fused permute kernel. Adding more conditions in the token dispatcher will make it more complex.

Could you post the detailed error you facing with zero-tokens? IIRC, the TE fused permute kernel should support the zero-token case. cc @hxbai

[hxbai]

I agree with @Victarry , it is better to create a fix to TE's permute function rather than to MCore.