[port] Skip auxiliary loss for hash MoE layers (upstream #7079)

[HJSang]

Update 2026-09-23 — ready for review. Rebased onto current miles-main (96fe16ecb, which adds #94, #102 and #105); it applied cleanly. The commit content is unchanged.

Stack and merge order: #97 (independent) · #98 → #99 → #100 → #101. Each PR is based on miles-main and contains its parent PR's commits. Review only the top commit of each PR (the Commits tab lists them). After a parent is squash-merged, rebase the child onto miles-main and drop the parent's commit. The order follows upstream's own merge order (NVIDIA#6058 on 2026-09-14 before NVIDIA#7169 on 2026-09-17), so #101 now applies upstream's final NVIDIA#7169 as-is.

Test plan: tracked in the comment on #99.

Test results (2026-09-23) — ✅ passed

Test	Base	This PR
test_routers.py (4 ranks)	43 passed / 2 failed / 14 errors	46 passed / 2 failed / 14 errors
Miles E2E, 256-token, --moe-aux-loss-coeff 1e-3: aux-loss grad_norm, steps 0/1/2	0.00618 / 0.00605 / 0.00612	0.00527 / 0.00552 / 0.00528 (−14%)
Miles E2E: old≠train tokens	0%	0%

The 3 new test_hash_routing_skips_aux_loss[*] cases pass. The 2 failures (test_router_gating_linear*[fp32], an fp32 router GEMM precision check) and the 14 errors (tests that need 8 ranks) are identical on the base. The E2E reproduces the grad-norm drop reported below. No regression.

Environment: 2× new 4× B200 devboxes (b200-verda-k8s), radixark/miles:dev rebuilt after miles#3296 (cuDNN Frontend 1.28.0 with the compact wrapper, TE 2.17, torch 2.13), Miles @ a6b72d31f, base miles-main@96fe16ecb. Unit tests ran under torchrun --nproc_per_node=4, and every failure was compared against the base run the same way. The Miles E2E recipe is DeepSeek-V4-Flash-FP8-4layer, GSM8K, deterministic mode, 3 rollouts, --dsv4-impl megatron --dsa-kernel-backend cudnn, with a per-token old-vs-train log-prob dump (~3.14M response tokens per arm). Rollouts were identical across arms. Full details: #99 (comment)

---

Ports NVIDIA#7079 (2b31f24598, merged to upstream dev 2026-09-07, author Hongxiao Bai) onto miles-main. Cherry-picked with -x; applied cleanly, no conflicts.

Why

Hash-routed MoE layers dispatch tokens by tid2eid, not by the learned router logits. When an auxiliary load-balancing loss is enabled, TopKRouter still recomputes a top-k routing map from the logits and applies the aux loss to it. That loss then optimizes routing decisions that were never executed. DeepSeek V4 uses hash routing on its first moe_n_hash_layers MoE layers (3 of the 4 layers in the Flash 4-layer prune, and the first 3 MoE layers of the full model), so every DSv4 run with a nonzero aux-loss coefficient trains against a spurious objective on those layers.

What changes

router.py: the aux-loss block in TopKRouter.forward gains a not self.is_hash_layer guard. Non-hash layers sharing the same config keep their aux loss. No behavior change when moe_aux_loss_coeff == 0.

test_routers.py: new TestHashRouting::test_hash_routing_skips_aux_loss, parametrized over aux_loss, seq_aux_loss, global_aux_loss, asserting no aux-loss autograd function is attached on a hash layer.

Applicability to Miles

Miles' scripts/run_deepseek_v4.py currently converts and trains with --moe-aux-loss-coeff 0, so existing recipes are unaffected. The fix matters as soon as anyone turns the aux loss on for DSv4, and it keeps miles-main aligned with upstream for the pending dev bump.

Validation

• 4x B200 devbox, radixark/miles:dev image, Megatron miles-main @ 73b54618f + this commit.
• tests/unit_tests/transformer/moe/test_routers.py -k test_hash_routing_skips_aux_loss: 3 passed.
• TestTop2Router::test_aux_loss: passed.
• TestAuxLossFreeTop2Router::test_router_forward_aux_free errors at setup on this devbox both with and without the patch (missing /opt/data/datasets.zip fixtures, torch.distributed not initialized in a single-process run). Pre-existing environment issue, unrelated to this change.

Miles end-to-end A/B (4x B200, radixark/miles:dev, model script scripts/models/deepseek-v4-flash-4layer.py via scripts/run_deepseek_v4.py train)

Same recipe both arms: DeepSeek-V4-Flash-FP8-4layer, GSM8K, --dsv4-impl megatron --dsa-kernel-backend cudnn, TP1/PP1/EP1/CP1, deterministic mode, 3 rollouts, 256-token responses, and --moe-aux-loss-coeff 1e-3 in --extra-args to override the script's 0 so the aux loss is actually live (moe_aux_loss_coeff 0.001, moe_n_hash_layers 3 confirmed in both logs). The pruned model earns zero reward, so train/loss and train/ppo_kl are exactly 0 in both arms and train/grad_norm is entirely aux-loss gradient.

Arm	Megatron	train/grad_norm steps 0 / 1 / 2	Tracebacks
baseline	miles-main 73b5461	0.00609 / 0.00625 / 0.00605	0
this PR	port/7079	0.00524 / 0.00537 / 0.00508	0

The fix arm's aux-loss gradient is ~15% lower on every step (within-arm spread ~0.0002, between-arm gap ~0.0009): the three hash-routed layers no longer contribute an aux-loss gradient. The drop is not 4x because the router gate is frozen in this recipe (--moe-router-freeze-gate), so aux-loss gradient only reaches parameters through the hidden states of earlier layers, and the surviving learned-routing layer is the last one. Direction and consistency are the point; the unit test proves the mechanism exactly. No behavior change when the coefficient is 0, as in current Miles recipes.

Series

First of two ports of DeepSeek V4 commits from upstream dev that miles-main (snapshot of dev at 2026-08-19) lacks. Companion: #98 (upstream NVIDIA#6343, fused RoPE packing). Further DSv4 commits (NVIDIA#7064, NVIDIA#5992) are left for the full dev bump.
🤖 Generated with Claude Code