Use the sequence-wise balance loss for DeepSeek-V4 routers

[AusarYao28]

Description

uses_megatron_seq_aux_loss accepted only routed_score_func == "sigmoid", so DeepSeek-V4 (routed_score_func: "sqrtsoftplus") fell back to the Switch-style load-balancing loss (softmax probabilities and bias-adjusted top-$k$ averaged across layers) even though its configs set load_balance_loss_weight: 1e-4 for the DeepSeek-V3 sequence-wise auxiliary load-balancing loss (DeepSeek-V3 Eqs. 17–20).

• Accept "sqrtsoftplus" in uses_megatron_seq_aux_loss (src/maxtext/layers/moe.py). Both callers already pass the pre-bias $\sqrt{\text{Softplus}}$ affinities (MoeGate applies score_func before recording pre_bias_logits) and skip hash-routed layers (is_hash_routing), so no other path changes.
• Set moe_use_megatron_seq_aux_loss: true in deepseek4-284b.yml, deepseek4-1600b.yml, and deepseek4-tiny.yml (matching the existing sigmoid convention in base.yml and types.py where per-layer sequence-wise auxiliary losses are summed over routed MoE layers).
• Add tests/unit/deepseek_v4_aux_loss_spec_test.py verifying deepseek4-tiny RoutedMoE auxiliary loss value and router gradient against an independent float64 reference of DeepSeek-V3 Eqs. 17–20 with $\sqrt{\text{Softplus}}$ scores, and update the predicate test in tests/unit/moe_megatron_aux_loss_test.py.

FIXES: b/521990776

Tests

• Unit (CPU, 4 devices): tests/unit/deepseek_v4_aux_loss_spec_test.py and tests/unit/moe_megatron_aux_loss_test.py (14 passed), tests/unit/configs_test.py (81 passed).
• Revert control: reverting src/maxtext/layers/moe.py and src/maxtext/configs/models/deepseek4-tiny.yml causes test_config_matches_spec, test_loss_value_matches_spec, and test_router_gradient_matches_spec in tests/unit/deepseek_v4_aux_loss_spec_test.py to fail (3 failed, 1 passed).
• pyink, pylint (10.00/10), yamllint, and codespell clean.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[github-actions]

🤖 Hi @AusarYao28, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[gemini-code-assist]

Code Review

This pull request extends the Megatron-LM sequence-wise auxiliary load-balancing loss to support Sqrt(Softplus) routers, which are used in DeepSeek-V4 models. It updates the configuration files for deepseek4-1600b, deepseek4-284b, and deepseek4-tiny to enable this loss, adjusts the routing logic in moe.py to recognize sqrtsoftplus as a valid routing score function, and updates relevant documentation. Additionally, a new unit test spec is introduced to validate the DeepSeek-V4 balance loss against a reference implementation, and existing tests are updated. There are no review comments, so no further feedback is provided.

[github-actions]

📋 Review Summary

Verdict: Spec met; the implementation correctly extends the DeepSeek-V3 seq aux loss to sqrtsoftplus routers, and the tests accurately validate it against a spec reference.

Top issues:

1. [Alignment] src/maxtext/experimental/lineage/lineage_adapter.py has a mirrored copy of uses_megatron_seq_aux_loss that was not updated. See inline comment on moe.py. → Update the Lineage adapter's copy for consistency, even if Lineage currently asserts sigmoid for other reasons.

Inferred spec & validation

• Sources: PR description, code before/after, and newly added tests.
• Problem: DeepSeek-V4 uses sqrtsoftplus affinities, which fell back to the Switch-style loss because uses_megatron_seq_aux_loss only checked for sigmoid, leading to incorrect load balancing.

Spec item	Implementation	Validating test
R1 (stated): Allow sqrtsoftplus to trigger uses_megatron_seq_aux_loss.	✅	test_config_matches_spec ✅
R2 (stated): Apply V3-style seq aux loss on Sqrt(Softplus) outputs correctly.	✅	test_loss_value_matches_spec, test_router_gradient_matches_spec, test_score_gradient_matches_spec ✅
R3 (stated): Update configs to enable moe_use_megatron_seq_aux_loss.	✅	test_config_matches_spec ✅
C1 (inferred): Keep uses_megatron_seq_aux_loss synced across the codebase.	❌ partial divergence	none

[github-actions]

🟠 Please also update uses_megatron_seq_aux_loss in src/maxtext/experimental/lineage/lineage_adapter.py (which mirrors this function) to include "sqrtsoftplus". Even though Lineage doesn't yet support V4 (it explicitly asserts sigmoid), keeping the mirrored function in sync prevents silent divergence when V4 support is eventually added to Lineage.

Suggested change

	and getattr(config, "routed_score_func", "") in ("sigmoid", "sqrtsoftplus")
	and getattr(config, "routed_score_func", "") in ("sigmoid", "sqrtsoftplus")

[AusarYao28]

Done — updated uses_megatron_seq_aux_loss in src/maxtext/experimental/lineage/lineage_adapter.py to include "sqrtsoftplus". (Commit e35cb28cd)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!