Compute DeepSeek-V4 router logits in fp32

[AusarYao28]

Description

The official DeepSeek-V4 reference Gate always projects MoE router logits in float32 (linear(x.float(), weight.float())). While deepseek4-1600b.yml sets float32_gate_logits: true, src/maxtext/configs/models/deepseek4-284b.yml inherited float32_gate_logits: false from base.yml, causing bfloat16 runs on DeepSeek-V4-Flash (deepseek4-284b) to project router logits in bfloat16 and flip top-$k$ expert selections near decision boundaries.

• Set float32_gate_logits: true in src/maxtext/configs/models/deepseek4-284b.yml, matching deepseek4-1600b.yml and the official reference implementation.

FIXES: b/566474598

Tests

• Unit (CPU, 4 devices): tests/unit/configs_test.py (81 passed), tests/unit/deepseek_v4_routed_bias_test.py (4 passed). yamllint and codespell clean.
• TPU forward-logit check on v5p-128 (deepseek4-284b, float32_gate_logits: true, official bf16 golden): PASS, KL[0] = 0.0429, max KL divergence = 0.116 (Log: https://cloudlogging.app.goo.gl/XtFn6fS4w15PdhQk9 ).

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.

[gemini-code-assist]

Code Review

This pull request updates the configuration for the DeepSeek4-284B model by setting float32_gate_logits to true, ensuring that the router logits are computed in float32 precision. There are no review comments, and I have no feedback to provide.

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

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

[github-actions]

📋 Review Summary

Verdict: Spec met; core config setting lacks automated regression test coverage.

Top issues:

1. [Tests] Core requirement is not covered by automated regression tests in CI (see inline comment on src/maxtext/configs/models/deepseek4-284b.yml). → Add a test verifying cfg.float32_gate_logits is True for deepseek4-284b (mirroring DeepSeekV4ProConfigSmokeTest in tests/unit/flop_calculation_test.py) or register deepseek4-284b.yml in DEEPSEEK_CONFIGS in tests/unit/configs_test.py.
2. [Spec] src/maxtext/configs/models/deepseek4-tiny.yml still inherits float32_gate_logits: false from base.yml, leaving CPU testing on deepseek4-tiny divergent from V4-Pro and V4-Flash router behavior. → Confirm whether deepseek4-tiny.yml should also enable float32_gate_logits: true for consistency across all DeepSeek-V4 configs.

Inferred spec & validation

• Sources: PR description; linked Buganizer ticket b/566474598 (not readable here); src/maxtext/configs/models/deepseek4-284b.yml; src/maxtext/configs/models/deepseek4-1600b.yml; src/maxtext/configs/base.yml.
• Problem: deepseek4-284b.yml omitted float32_gate_logits, inheriting false from base.yml and causing router logits to project in bfloat16 instead of float32, which diverged from the official reference and flipped top-$k$ expert selection near boundaries.

Spec item	Implementation	Validating test
R1 (stated): deepseek4-284b.yml sets float32_gate_logits: true so MoE router projections run in fp32	✅	❌ none (configs_test.py omits deepseek4-284b.yml; deepseek_v4_routed_bias_test.py overrides the flag via CLI)
R2 (inferred): DeepSeek-V4 test configs mirror production router precision	⚠️ partial (deepseek4-tiny.yml still inherits false)	❌ none
C1 (inferred): YAML syntax and schema validation remain valid	✅ preserved	❌ none (configs_test.py does not include deepseek4-284b.yml)
C2 (inferred): Other model families without explicit overrides retain default behavior	✅ preserved	tests/unit/configs_test.py ✅

[github-actions]

🟡 The new config setting is unvalidated by automated regression tests in CI. While the PR description cites tests/unit/configs_test.py and tests/unit/deepseek_v4_routed_bias_test.py, deepseek4-284b.yml is omitted from DEEPSEEK_CONFIGS in configs_test.py, and deepseek_v4_routed_bias_test.py sets float32_gate_logits explicitly via CLI args rather than validating model config defaults.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!