From a3ac52540b715401f6b4839d54e3bdc7800bc0fa Mon Sep 17 00:00:00 2001 From: Yuchen Yao Date: Mon, 19 Oct 2026 21:13:34 +0000 Subject: [PATCH] Compute DeepSeek-V4 router logits in fp32 The DeepSeek-V4 reference Gate always projects router logits in fp32 (linear(x.float(), weight.float())), but deepseek4-284b.yml inherits float32_gate_logits: false, so bf16 runs route on bf16 logits and flip experts. - Set float32_gate_logits: true in src/maxtext/configs/models/deepseek4-284b.yml (config-only; same setting as the qwen3.5 fp8 model configs). - Validated on v5p-8 bf16 against the fp32 reference on real Flash weights: layers 3-6 h_out rel_l2 8.21e-2 -> 2.44e-2, router.3 top-k slot overlap 0.886 -> 0.996, dh_in 0.213 -> 0.064; layers 41-42 + head pass all checks. FIXES: b/566474598 --- src/maxtext/configs/models/deepseek4-284b.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/maxtext/configs/models/deepseek4-284b.yml b/src/maxtext/configs/models/deepseek4-284b.yml index 5623878b6c..5bef652d42 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/src/maxtext/configs/models/deepseek4-284b.yml @@ -56,6 +56,8 @@ routed_bias_update_rate: 0.001 load_balance_loss_weight: 0.0001 log_moe_bias_norms: false routed_scaling_factor: 1.5 +# Reference Gate computes router logits in fp32 (linear(x.float(), weight.float())) regardless of model dtype. +float32_gate_logits: true # --- Attention configuration ---