MORI + Permute/Unpermute zero token fix

[sudhu2k]

Summary

Under imbalanced routing, an expert-parallel (EP) rank can legitimately receive zero
tokens — a "cold" rank. With moe_permute_fusion enabled, TransformerEngine's fused
permute/unpermute kernels short-circuit on empty [0, H] tensors, which corrupted the MORI
combine buffer on the forward pass and raised in the permute backward, deadlocking the job.

This PR makes cold ranks a first-class case on both the dispatch and combine seams, in both
the forward and backward passes. The fix is entirely on the Megatron side and requires no
TransformerEngine changes.

Motivation

Cold ranks arise in normal training whenever the router is skewed. The prior code masked the problem by clamping num_out_tokens to a minimum of 1, which fabricated a dummy token row on cold ranks. That workaround avoided the immediate crash but left the forward combine buffer semantically wrong and did not address the backward
pass at all, so training under real skew was unsafe.

Root cause

The empty-tensor handling in TE's fused kernels fails only in the unpermute/combine direction —
which shows up on two distinct autograd functions:

1. Combine seam — fused_unpermute (forward).
   On a cold rank the expert output is [0, H]. TE's fused unpermute short-circuits
   (if not inp.numel(): return inp) and returns [0, H] instead of zeros(restore_shape).
   The collapsed shape then flows into op.combine (x=[0, H], total_recv=0), corrupting the
   combine forward and zeroing the combine backward.

2. Dispatch seam — fused_permute (backward).
   On a cold rank the permute input is non-empty (local tokens exist) but the output is
   empty, so the incoming gradient is [0, H]. TE's permute backward short-circuits on that empty
   gradient and dereferences ctx.probs, which is never set on the non-empty-input path, then
   returns a gradient of the wrong shape. The resulting exception fires only on cold ranks; they
   abort backward and proceed to teardown while hot ranks block in MORI finalize
   (cuda.synchronize) waiting on collective work the cold ranks never posted — a hang that
   surfaces in process-group finalize/destroy but originates in the cold-rank backward.

Changes

megatron/core/transformer/moe/token_dispatcher.py

Remove the num_out_tokens = max(num_out_tokens, 1) clamp in _MoriManager. With the guards
below, an honest empty permute ([0, H]) is handled correctly end-to-end, so the dummy row is no
longer needed.

megatron/core/transformer/moe/moe_utils.py

Add two symmetric cold-rank guards that fall back to the native (non-fused) paths, which have
correct native autograd for empty selections:

# permute() — dispatch seam
if fused and num_out_tokens == 0:
    fused = False

# unpermute() — combine seam
if fused and permuted_tokens.numel() == 0:
    fused = False

• permute: the native index_select path yields [0, H] on the forward (a no-op over zero
  rows) and, on the backward, scatters the incoming [0, H] gradient into zeros_like(tokens),
  producing the correct [R, H] dispatch-buffer gradient (a cold rank contributes no gradient).
• unpermute: the native scatter path fills zeros(restore_shape) on the forward and, on the
  backward, computes grad.gather(empty) = [0, H], keeping the empty input on the autograd graph
  so the combine backward stays symmetric across ranks. This replaces the previous workaround that
  synthesized zeros(restore_shape) and manually reattached it via + permuted_tokens.sum() * 0.

Both guards are no-ops on non-cold ranks and carry no performance cost (the fallback executes over
zero rows only).

Tests

• tests/unit_tests/transformer/moe/test_moe_unpermute.py (new): unit test asserting that fused
  unpermute on [0, H] input restores zeros(restore_shape) — verifying shape, dtype, device,
  and all-zero contents.
• tests/unit_tests/transformer/moe/test_token_dispatcher.py: adds TestMoriColdRank, a
  distributed (8-way EP) test that forces all tokens onto the first topk experts so higher EP
  ranks receive zero tokens, runs dispatch → SequentialMLP → combine, and asserts that
  .backward() completes with the correct [R, H] dispatch-buffer gradient. MORI shmem is
  finalized once at teardown, since MORI cannot finalize and reinitialize shmem within the same
  process.