Add HybridModel hash MoE routing

[FDecaYed]

Add HybridModel hash MoE routing

Summary

Add complete Hash MoE behavior for HybridModel, from token-ID transport through router selection, MoE execution, metrics, recompute, and gradient finalization.

Scope

• Add persistent token-to-expert mappings and the supported hash combination-score modes.
• Carry token IDs through HybridModel, Hybrid layers, MTP, recompute, and partial CUDA-graph paths.
• Route only the intended leading MoE layers and preserve normal routing elsewhere.
• Keep router/MoE metric slots and final gradient handling consistent with the Hybrid execution path.
• Add focused router, HybridModel, recompute, and gradient-finalization tests.

Non-goals

• mHC implementation; P2 owns that prerequisite.
• Variable-token HybridEP correctness, packed-data scheduling, or static THD CUDA graphs.
• GPTModel integration, DSv4 attention, or recipes.

Review boundary and dependencies

Requires #6401 (P2). The local branch temporarily carries P2 and its #5929/#4531 prerequisites, but this PR’s own review commit is:

c9657bcd64c91a375e242a013aab3cc3a6731c71

Its exact own diff is 16 files, +946/-45. Please review that commit only. This will target main, not a pull-request/<P2> base; after P2 merges, the own commit will be rebased onto current main.

Provenance

Reconstructed from the Hash MoE behavior tracked by #5795, including the focused recuts #5938 and #5945. The original feature and forced-routing correction came from #4481 and #5130.

Validation

Completed locally:

• clean worktree;
• signed-commit verification and required Signed-off-by trailer;
• git diff --check;
• Python bytecode compilation;
• Ruff on all 16 files;
• AST behavior comparison.

Black 26.3.1 passes 14 of 16 files. The two whole-file failures in megatron/training/argument_utils.py and megatron/training/arguments.py are identical on the P2 parent; P4 adds only one non-formatting line to each. Isort passes 15 of 16 files; the arguments.py failure is also present on the P2 parent, and P4 changes no imports there.

Runtime pytest is unavailable in the bare macOS environment because PyTorch/CUDA is not installed. Focused CUDA tests and full functional CI remain required after publication. No dependency or lock-file change is part of this commit.

Local artifact

P4-hybrid-hash-moe-routing.patch, generated from the own commit only. Verify it against SHA256SUMS.

Before publication

After P2 merges, replay this own commit onto then-current main; recheck current router, partial CUDA-graph, MTP, and gradient-finalization contracts; rerun focused/static validation; then verify signing/trailers and launch full functional CI.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.