fix(moe): preserve DeepEP autograd for empty expert ranks

[Slyne]

Summary

• keep the zero-local-token GroupedExpertsDeepEP path connected to the tensors returned by DeepEP dispatch
• attach zero-valued dependencies for all local expert parameters so empty EP ranks still enter reverse dispatch and materialize optimizer gradients
• stabilize the empty routing-probability dtype across activation-checkpoint recomputation
• add CPU coverage and a scheduled two-GPU DeepEP regression with deliberately imbalanced routing

Motivation

When one EP rank receives no local expert tokens, the previous fallback built its dummy computation from the original input. That bypassed FusedDispatch autograd on that rank while peers entered the reverse-dispatch collective, causing either a no-grad error or a distributed timeout during backward.

Verification

• ruff format --check and ruff check on the changed Python files
• 3 focused GroupedExpertsDeepEP CPU regression cases passed
• shell syntax check for the new L2 launcher
• end-to-end Hero5c CP=1/EP=8 head-only MTP training completed five 12,288-token optimizer steps on eight H100s (Slurm 6351147)

The new standalone two-GPU L2 guard is included for CI execution. A CP=2/EP=4 end-to-end smoke remains a draft-PR gate.

Related

• NeMo Speech LK/MTP integration: feat(speechlm2): distill MTP heads with LK loss Speech#16140

[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.

[HuiyingLi]

Hi @Slyne, what failure did you see? Hang?