Fix Gemma4 MoE expert weight remapping (#41206)

Signed-off-by: sunghoon.baek <sunghoon.baek@connectfy.cloud>
Co-authored-by: sunghoon.baek <sunghoon.baek@connectfy.cloud>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Baekpica
2026-04-30 00:12:42 -07:00
committed by GitHub
co-authored by sunghoon.baek OpenAI Codex
parent a04e0cf3b8
commit ca97f7b9bb
+5 -1
View File
@@ -84,6 +84,10 @@ from .utils import (
logger = init_logger(__name__)
def _remap_gemma4_expert_weight_name(name: str) -> str:
return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
@triton.jit
def _gemma4_routing_kernel(
gating_ptr,
@@ -1650,7 +1654,7 @@ class Gemma4ForCausalLM(
# Remap individual 2D expert weights:
# .experts.{id}.{proj} → .moe.experts.{id}.{proj}
# (This handles per-expert 2D quantized weights)
name = re.sub(r"\.experts\.(\d+)\.", r".moe.experts.\1.", name)
name = _remap_gemma4_expert_weight_name(name)
# MoE expert weights: checkpoint stores as 3D packed
# tensors. Explode into per-expert 2D weights for