Skip to content

Fix double application of per_expert_scale in Gemma 4 MoE. #4880

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Aug 14, 2026
+8 −3
Merged

Fix double application of per_expert_scale in Gemma 4 MoE. #4880

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions src/maxtext/layers/moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3119,9 +3119,14 @@ def __call__(
w0_kernel = jnp.asarray(self.wi_0[...], self.dtype)
w1_kernel = jnp.asarray(self.wi_1[...], self.dtype)

# Only apply per expert scales if we have not fused with the out-projections at init time.
if self.per_expert_scale is not None and cfg.model_call_mode != "inference" and not cfg.fuse_expert_scales:
wo_kernel = wo_kernel * jnp.asarray(self.per_expert_scale[...], self.dtype)[:, None, None]
# For fused MoE path (inference only), if we have not fused expert
# scales at init, we must apply them to wo_kernel here because
# fused_moe_func doesn't support them. Other paths (dense/sparse
# matmul) apply them to top_k_weights in get_topk.
is_fused_moe_path = cfg.attention in ("vllm_rpa", "vllm_batched_rpa") and not self.is_hash_routing
if is_fused_moe_path:
if self.per_expert_scale is not None and not (cfg.model_call_mode == "inference" and cfg.fuse_expert_scales):
wo_kernel = wo_kernel * jnp.asarray(self.per_expert_scale[...], self.dtype)[:, None, None]

if self.wi_0_sparsity_module is not None:
_, w0_kernel = self.wi_0_sparsity_module(jnp.zeros_like(w0_kernel), w0_kernel)
Expand Down
Loading