[Model][Gemma4] Load Weights with AutoWeightsLoader (#55911)
Signed-off-by: Josephasafg <ajgard7@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Asaf Gardin committed
Sep 18, 2026 at 15:07 UTC
44dd18fe0bb0f13157f97a5aa029b6604468fca6
8 files changed
+238
-517
tests/lora/test_lora_checkpoints.py
+2
-2
@@ -147,10 +147,10 @@ def test_gemma4_lora_weights_mapping():
147
def test_gemma4_moe_lora_weights_mapping():
148
mapper = Gemma4ForCausalLM.hf_to_vllm_mapper
149
name = (
150
- "base_model.model.model.language_model.layers.9.moe.experts."
150
+ "base_model.model.model.language_model.layers.9.experts."
151
"gate_up_proj.lora_B.weight"
152
)
153
assert parse_fine_tuned_lora_name(name, mapper) == (
154
- "model.layers.9.moe.gate_up_proj",
154
+ "model.layers.9.experts.gate_up_proj",
155
False,
156
)
tests/models/language/generation/test_gemma.py
+65
-28
@@ -14,7 +14,10 @@ from vllm.model_executor.models.gemma3n import (
14
Gemma3nTextModel,
15
_kv_sharing_weights_mapper,
16
)
17
-from vllm.model_executor.models.gemma4 import Gemma4Model
17
+from vllm.model_executor.models.gemma4 import (
18
+ Gemma4ForCausalLM,
19
+ _gemma4_layer_weights_mapper,
20
+)
21
22
MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"]
23
@@ -57,36 +60,70 @@ def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None:
60
61
62
@pytest.mark.cpu_test
60
-def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None:
61
- """KV-shared layers have q_proj instead of a packed qkv_proj; their
62
- redundant K/V tensors in original checkpoints have no parameter and are
63
- skipped rather than failing the load."""
64
- model = torch.nn.Module()
65
- model.config = SimpleNamespace(num_experts=0)
66
- model.start_layer, model.end_layer = 0, 2
67
- model.layers = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()])
68
- for layer in model.layers:
69
- layer.self_attn = torch.nn.Module()
70
- model.layers[0].self_attn.qkv_proj = torch.nn.Linear(2, 6, bias=False)
71
- model.layers[1].self_attn.q_proj = torch.nn.Linear(2, 2, bias=False)
72
- shards: list[str] = []
73
- model.layers[0].self_attn.qkv_proj.weight.weight_loader = (
74
- lambda param, weight, shard_id: shards.append(shard_id)
63
+def test_gemma4_attention_mapper() -> None:
64
+ """Layers with a qkv_proj pack q/k/v; `attention_k_eq_v` full-attention
65
+ layers also load K as the V shard, and leave any v_proj such a checkpoint
66
+ ships unmapped so it fails the load rather than silently overwriting V;
67
+ KV-shared layers keep q_proj and drop the K/V tensors original checkpoints
68
+ still ship for them."""
69
+ config = SimpleNamespace(
70
+ num_hidden_layers=3,
71
+ num_kv_shared_layers=1,
72
+ attention_k_eq_v=True,
73
+ layer_types=["sliding_attention", "full_attention", "sliding_attention"],
74
)
75
+ weights = [
76
+ (f"model.layers.{i}.self_attn.{tensor}.weight", torch.full((2, 2), i + 1.0))
77
+ for i in range(3)
78
+ for tensor in ("q_proj", "k_proj", "k_norm")
79
+ ] + [
80
+ ("model.layers.0.mlp.up_proj.weight", torch.empty(0)),
81
+ ("model.layers.1.self_attn.v_proj.weight", torch.empty(0)),
82
+ ]
83
77
- q_weight = torch.full((2, 2), 2.0)
84
+ mapper = _gemma4_layer_weights_mapper(config)
85
+ mapped = list(mapper.apply(weights))
86
+
87
+ assert [(name, getattr(w, "shard_id", None)) for name, w in mapped] == [
88
+ ("model.layers.0.self_attn.qkv_proj.weight", "q"),
89
+ ("model.layers.0.self_attn.qkv_proj.weight", "k"),
90
+ ("model.layers.0.self_attn.k_norm.weight", None),
91
+ ("model.layers.1.self_attn.qkv_proj.weight", "q"),
92
+ ("model.layers.1.self_attn.qkv_proj.weight", "k"),
93
+ ("model.layers.1.self_attn.qkv_proj.weight", "v"),
94
+ ("model.layers.1.self_attn.k_norm.weight", None),
95
+ ("model.layers.2.self_attn.q_proj.weight", None),
96
+ ("model.layers.0.mlp.gate_up_proj.weight", 1),
97
+ ("model.layers.1.self_attn.v_proj.weight", None),
98
+ ]
99
+ k_weight, v_weight = weights[4][1], mapped[5][1]
100
+ assert torch.equal(v_weight, k_weight) and v_weight is not k_weight
101
+
102
+
103
+@pytest.mark.cpu_test
104
+def test_gemma4_expert_names_strip_language_model_prefix() -> None:
105
+ """The text-only path reuses the conditional wrapper's checkpoint naming,
106
+ so fused and per-expert tensors reach the experts under `model.*`."""
107
+ prefix = "model.language_model.layers.0."
108
weights = [
79
- (f"layers.{i}.self_attn.{tensor}.weight", q_weight)
80
- for i in (0, 1)
81
- for tensor in ("q_proj", "k_proj", "v_proj")
82
- ] + [("layers.1.self_attn.k_norm.weight", torch.ones(2))]
83
- loaded = Gemma4Model.load_weights(cast(Gemma4Model, model), weights)
84
-
85
- assert shards == ["q", "k", "v"]
86
- assert "layers.1.self_attn.q_proj.weight" in loaded
87
- assert not any(name.startswith("layers.1.self_attn.k") for name in loaded)
88
- assert not any(name.startswith("layers.1.self_attn.v") for name in loaded)
89
- assert torch.equal(model.layers[1].self_attn.q_proj.weight, q_weight)
109
+ (prefix + name, torch.empty(0))
110
+ for name in (
111
+ "experts.gate_up_proj",
112
+ "experts.3.down_proj.weight_packed",
113
+ "router.per_expert_scale",
114
+ )
115
+ ]
116
+
117
+ mapped = [
118
+ (name, getattr(w, "shard_id", None))
119
+ for name, w in Gemma4ForCausalLM.hf_to_vllm_mapper.apply(weights)
120
+ ]
121
+
122
+ assert mapped == [
123
+ ("model.layers.0.experts.gate_up_proj", None),
124
+ ("model.layers.0.experts.3.down_proj.weight_packed", None),
125
+ ("model.layers.0.router.per_expert_scale", None),
126
+ ]
127
128
129
@pytest.mark.cpu_test
tests/models/test_utils.py
+22
@@ -217,3 +217,25 @@ def test_get_rename_mapper_keeps_only_renames():
217
for name in ("drop_regex.w", "drop_substr.w", "drop_prefix.w", "w.drop_suffix"):
218
assert mapper._map_name(name) is None
219
assert renames._map_name(name) == name
220
+
221
+
222
+def test_weights_mapper_stacks_one_weight_into_several_shards():
223
+ """One checkpoint tensor can feed more than one shard of a stacked
224
+ parameter, and a dropped name stays dropped rather than reappearing as a
225
+ shard of the stacked one."""
226
+ mapper = WeightsMapper(
227
+ orig_to_new_substr={"layers.1.k_proj.": None},
228
+ orig_to_new_stacked={".k_proj.": (".qkv_proj.", ["k", "v"])},
229
+ )
230
+ weight = torch.ones(2)
231
+ weights = [(f"layers.{i}.k_proj.weight", weight) for i in (0, 1)]
232
+
233
+ mapped = list(mapper.apply(weights))
234
+
235
+ assert [(name, w.shard_id) for name, w in mapped] == [
236
+ ("layers.0.qkv_proj.weight", "k"),
237
+ ("layers.0.qkv_proj.weight", "v"),
238
+ ]
239
+ # Each shard needs its own tensor object, but they alias one allocation.
240
+ assert mapped[0][1] is not mapped[1][1]
241
+ assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() == weight.data_ptr()
tests/quantization/test_modelopt.py
+3
-3
@@ -341,7 +341,7 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers():
341
"quant_algo": "NVFP4",
342
"group_size": 16,
343
},
344
- "model.language_model.layers.1.moe.experts.gate_up_proj": {
344
+ "model.language_model.layers.1.experts.gate_up_proj": {
345
"quant_algo": "NVFP4",
346
"group_size": 16,
347
},
@@ -353,10 +353,10 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers():
353
)
354
config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_rename_mapper())
355
356
- expected_prefix = "language_model.model.layers.0.moe.experts"
356
+ expected_prefix = "language_model.model.layers.0.experts"
357
assert set(config.quantized_layers) == {
358
expected_prefix,
359
- "language_model.model.layers.1.moe.gate_up_proj",
359
+ "language_model.model.layers.1.experts.gate_up_proj",
360
}
361
assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
362
vllm/model_executor/models/diffusion_gemma.py
+11
-110
@@ -46,7 +46,11 @@ from vllm.model_executor.models.gemma4_mm import (
46
)
47
from vllm.model_executor.models.module_mapping import MultiModelKeys
48
from vllm.model_executor.models.transformers.utils import recursive_replace_linear
49
-from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
49
+from vllm.model_executor.models.utils import (
50
+ AutoWeightsLoader,
51
+ WeightsMapper,
52
+ maybe_prefix,
53
+)
54
from vllm.multimodal import MULTIMODAL_REGISTRY
55
from vllm.platforms import current_platform
56
from vllm.utils.torch_utils import async_tensor_h2d
@@ -160,14 +164,12 @@ class DiffusionGemmaForConditionalGeneration(
164
165
hf_to_vllm_mapper = WeightsMapper(
166
orig_to_new_prefix={
167
+ "model.decoder.self_conditioning.": "self_conditioning.",
168
"model.decoder.": "model.",
169
"model.encoder.language_model.": "model.",
170
"model.encoder.vision_tower.": "vision_tower.",
171
"model.encoder.embed_vision.": "embed_vision.",
172
},
168
- orig_to_new_substr={
169
- ".experts.": ".moe.experts.",
170
- },
173
)
174
175
packed_modules_mapping = {
@@ -336,113 +338,12 @@ class DiffusionGemmaForConditionalGeneration(
338
logits = _softcap_logits(logits, self.final_logit_softcapping)
339
return logits
340
339
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
340
- """Load weights from checkpoint.
341
-
342
- Checkpoint layout (HF DiffusionGemma):
343
- model.encoder.vision_tower.* → vision tower
344
- model.encoder.embed_vision.* → vision embedder
345
- model.encoder.language_model.layers.* → backbone
346
- model.decoder.layers.* → backbone (tied)
347
- model.decoder.embed_tokens.* → embeddings
348
- model.decoder.self_conditioning.* → self-conditioning MLP
349
- lm_head.* → LM head (tied)
350
-
351
- We load encoder weights into our single ``Gemma4Model`` backbone,
352
- skip duplicate decoder backbone weights, handle vision tower and
353
- self-conditioning separately.
354
- """
355
- sc_params = dict(
356
- (n, p)
357
- for n, p in self.named_parameters()
358
- if n.startswith("self_conditioning.")
341
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
342
+ # Some checkpoints carry a vestigial Gemma3n-style embedding table.
343
+ loader = AutoWeightsLoader(
344
+ self, ignore_unexpected_prefixes=["embed_vision.embedding."]
345
)
360
-
361
- # Collect vision tower + embedder parameters AND buffers for manual
362
- # loading. The HF vision tower registers std_bias / std_scale as
363
- # buffers (not parameters) when config.standardize is True, so we
364
- # must include named_buffers() to avoid "not found in model" warnings.
365
- vision_params: dict[str, torch.Tensor] = {}
366
- for n, p in self.named_parameters():
367
- if n.startswith(("vision_tower.", "embed_vision.")):
368
- vision_params[n] = p
369
- for n, b in self.named_buffers():
370
- if n.startswith(("vision_tower.", "embed_vision.")):
371
- vision_params[n] = b
372
-
373
- def _remap_weights():
374
- # Use full weight names (including suffixes like .weight_scale,
375
- # .weight_packed) for dedup instead of just the base layer name. Critical
376
- # for quantized checkpoints where each weight has multiple tensors;
377
- # tracking only base names skips scales as duplicates.
378
- seen_weights: set[str] = set()
379
- for name, weight in weights:
380
- # Self-conditioning lives under model.decoder.self_conditioning.*
381
- # in the checkpoint but at self_conditioning.* in our model.
382
- if "self_conditioning" in name:
383
- sc_name = name.split("self_conditioning.", 1)[1]
384
- sc_name = "self_conditioning." + sc_name
385
- if sc_name in sc_params:
386
- sc_params[sc_name].data.copy_(weight)
387
- continue
388
-
389
- # Vision tower: model.encoder.vision_tower.* → vision_tower.*
390
- # In HF, the vision tower is a sibling of language_model
391
- # under the encoder module.
392
- if name.startswith("model.encoder.vision_tower."):
393
- vt_name = name[len("model.encoder.") :]
394
- if vt_name in vision_params:
395
- vision_params[vt_name].data.copy_(weight)
396
- else:
397
- logger.warning(
398
- "Vision tower weight %s (mapped to %s) not found in model",
399
- name,
400
- vt_name,
401
- )
402
- continue
403
-
404
- # Vision embedder: model.encoder.embed_vision.* → embed_vision.*
405
- if name.startswith("model.encoder.embed_vision."):
406
- ev_name = name[len("model.encoder.") :]
407
- if ev_name in vision_params:
408
- vision_params[ev_name].data.copy_(weight)
409
- else:
410
- logger.warning(
411
- "Embed vision weight %s (mapped to %s) not found in model",
412
- name,
413
- ev_name,
414
- )
415
- continue
416
-
417
- # Skip vestigial embed_vision.embedding weights.
418
- if "embed_vision.embedding." in name:
419
- continue
420
-
421
- # Encoder backbone → model.*
422
- if name.startswith("model.encoder.language_model."):
423
- name = name.replace("model.encoder.language_model.", "model.")
424
- # Decoder backbone → model.* (skip exact duplicates)
425
- elif name.startswith("model.decoder."):
426
- name = name.replace("model.decoder.", "model.")
427
-
428
- # Skip only if we've seen the exact same weight name (including scales)
429
- if name in seen_weights:
430
- continue
431
- seen_weights.add(name)
432
- yield name, weight
433
-
434
- # Delegate to Gemma4ForCausalLM.load_weights for the backbone,
435
- # which handles stacked params, MoE, k_eq_v, etc.
436
- # Temporarily set self.config to text_config since Gemma4's
437
- # load_weights expects it (e.g. tie_word_embeddings, layer_types).
438
- from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
439
-
440
- saved_config = self.config
441
- self.config = self.model.config
442
- try:
443
- Gemma4ForCausalLM.load_weights(self, _remap_weights())
444
- finally:
445
- self.config = saved_config
346
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
347
348
@classmethod
349
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
vllm/model_executor/models/gemma4.py
+113
-362
@@ -22,12 +22,12 @@ from collections.abc import Iterable
22
from dataclasses import replace
23
from itertools import islice
24
25
-import regex as re
25
import torch
26
from torch import nn
27
28
from vllm.compilation.decorators import support_torch_compile
29
from vllm.config import CacheConfig, VllmConfig
30
+from vllm.config.utils import getattr_iter
31
from vllm.distributed import (
32
get_pp_group,
33
get_tensor_model_parallel_rank,
@@ -40,7 +40,7 @@ from vllm.model_executor.layers.attention import Attention
40
from vllm.model_executor.layers.fused_moe import (
41
FusedMoEFactory,
42
GateLinear,
43
- fused_moe_make_expert_params_mapping,
43
+ MoERunner,
44
)
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
@@ -57,10 +57,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
57
ParallelLMHead,
58
VocabParallelEmbedding,
59
)
60
-from vllm.model_executor.model_loader.weight_utils import (
61
- default_weight_loader,
62
- maybe_remap_kv_scale_name,
63
-)
60
from vllm.platforms import current_platform
61
from vllm.sequence import IntermediateTensors
62
from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
@@ -73,27 +69,63 @@ from .interfaces import (
69
SupportsEagle3,
70
SupportsLoRA,
71
SupportsPP,
72
+ SupportsQuant,
73
)
74
from .utils import (
75
AutoWeightsLoader,
76
+ PPMissingLayer,
77
+ ShardIds,
78
WeightsMapper,
79
extract_layer_index,
81
- is_pp_missing_parameter,
80
make_layers,
81
maybe_prefix,
82
)
83
84
logger = init_logger(__name__)
85
88
-_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
89
- orig_to_new_regex={
90
- re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
91
- }
92
-)
86
87
+def _gemma4_layer_weights_mapper(config) -> WeightsMapper:
88
+ """Stack each decoder layer's projections the way its modules are built.
89
95
-def _remap_gemma4_expert_weight_name(name: str) -> str:
96
- return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
90
+ gate/up pack into `mlp.gate_up_proj`. Layers with a qkv_proj pack q/k/v
91
+ into it; `attention_k_eq_v` full-attention layers ship k_proj only, so K
92
+ is loaded into both the K and V shards and no v_proj is expected. KV-shared
93
+ layers only have q_proj; the K/V tensors original checkpoints still ship
94
+ for them are dropped.
95
+ """
96
+ num_layers = config.num_hidden_layers
97
+ first_kv_shared = num_layers - getattr(config, "num_kv_shared_layers", 0)
98
+ k_eq_v_layers = (
99
+ {
100
+ i
101
+ for i, layer_type in enumerate(config.layer_types[:first_kv_shared])
102
+ if layer_type == "full_attention"
103
+ }
104
+ if getattr(config, "attention_k_eq_v", False)
105
+ else set()
106
+ )
107
+ qkv_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", "k"), ("v", "v"))
108
+ k_eq_v_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", ["k", "v"]))
109
+ stacked: dict[str, tuple[str, ShardIds]] = {
110
+ ".mlp.gate_proj.": (".mlp.gate_up_proj.", 0),
111
+ ".mlp.up_proj.": (".mlp.gate_up_proj.", 1),
112
+ }
113
+ stacked |= {
114
+ f"layers.{i}.self_attn.{shard}_proj.": (
115
+ f"layers.{i}.self_attn.qkv_proj.",
116
+ shard_id,
117
+ )
118
+ for i in range(first_kv_shared)
119
+ for shard, shard_id in (k_eq_v_shards if i in k_eq_v_layers else qkv_shards)
120
+ }
121
+ return WeightsMapper(
122
+ orig_to_new_substr={
123
+ f"layers.{i}.self_attn.{name}.": None
124
+ for i in range(first_kv_shared, num_layers)
125
+ for name in ("k_proj", "v_proj", "k_norm")
126
+ },
127
+ orig_to_new_stacked=stacked,
128
+ )
129
130
131
@triton.jit
@@ -279,6 +311,9 @@ class Gemma4Router(nn.Module):
311
self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps, has_weight=False)
312
# Per-dimension learned scale, applied after norm + root_size
313
self.scale = nn.Parameter(torch.ones(self.hidden_size))
314
+ # Per-expert output scale folded into routing weights so that
315
+ # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
316
+ self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts))
317
# Constant 1/sqrt(hidden_size) scaling factor
318
self.register_buffer(
319
"root_size",
@@ -296,84 +331,36 @@ class Gemma4Router(nn.Module):
331
prefix=f"{prefix}.proj",
332
)
333
299
- def forward(self, x: torch.Tensor) -> torch.Tensor:
300
- """Returns raw router logits [T, E]."""
301
- x = self.norm(x)
302
- x = x * self.root_size.to(x.dtype)
303
- x = x * self.scale.to(x.dtype)
304
- router_logits, _ = self.proj(x)
305
- return router_logits
306
-
307
-
308
-class Gemma4MoE(nn.Module):
309
- """Mixture of Experts for Gemma4 using vLLM's MoERunner.
310
-
311
- Wraps MoERunner with custom routing. The router projection is
312
- external (Gemma4Router) — this class only handles expert dispatch.
313
-
314
- Gemma4 routing: softmax over ALL experts → top-k → renormalize.
315
- per_expert_scale is folded into routing weights for mathematical
316
- correctness with MoERunner's fused kernel.
317
- """
318
-
319
- def __init__(
334
+ def routing_function(
335
self,
321
- config,
322
- quant_config: QuantizationConfig | None = None,
323
- prefix: str = "",
324
- ) -> None:
325
- super().__init__()
326
- self.hidden_size = config.hidden_size
327
- self.num_experts = config.num_experts
328
-
329
- # Per-expert output scale folded into routing weights so that
330
- # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
331
- self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts))
336
+ hidden_states: torch.Tensor,
337
+ gating_output: torch.Tensor,
338
+ topk: int,
339
+ renormalize: bool,
340
+ ) -> tuple[torch.Tensor, torch.Tensor]:
341
+ """Softmax over ALL experts → top-k → renormalize.
342
333
- # Gemma4 routing: softmax over ALL experts → top-k → renormalize.
334
- # MoERunner's built-in fused_topk scopes softmax differently, so
335
- # a custom routing function is needed for numerical correctness.
336
- # NOTE: self.per_expert_scale is read at call time (not captured into
337
- # a local) so that torch.func.functional_call parameter substitution
338
- # reaches the routing function correctly.
339
- def routing_function(
340
- hidden_states: torch.Tensor,
341
- gating_output: torch.Tensor,
342
- topk: int,
343
- renormalize: bool,
344
- ) -> tuple[torch.Tensor, torch.Tensor]:
345
- if current_platform.is_cuda_alike() or current_platform.is_xpu():
346
- return gemma4_fused_routing_kernel_triton(
347
- gating_output, topk, self.per_expert_scale
348
- )
343
+ MoERunner's built-in fused_topk scopes softmax differently, so Gemma4
344
+ needs its own routing for numerical correctness.
345
350
- return gemma4_routing_function_torch(
346
+ NOTE: self.per_expert_scale is read at call time (not captured into a
347
+ local) so that torch.func.functional_call parameter substitution
348
+ reaches the routing function correctly.
349
+ """
350
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
351
+ return gemma4_fused_routing_kernel_triton(
352
gating_output, topk, self.per_expert_scale
353
)
354
354
- # MoERunner experts with custom Gemma4 routing
355
- intermediate_size = getattr(
356
- config,
357
- "moe_intermediate_size",
358
- getattr(config, "expert_intermediate_size", None),
359
- )
360
- if intermediate_size is None:
361
- raise ValueError("Gemma4 MoE requires an expert intermediate size")
362
-
363
- self.experts = FusedMoEFactory(
364
- num_experts=config.num_experts,
365
- top_k=config.top_k_experts,
366
- hidden_size=config.hidden_size,
367
- intermediate_size=intermediate_size,
368
- renormalize=True,
369
- quant_config=quant_config,
370
- prefix=f"{prefix}.experts",
371
- custom_routing_function=routing_function,
372
- activation="gelu_tanh",
373
- )
355
+ return gemma4_routing_function_torch(gating_output, topk, self.per_expert_scale)
356
375
- def forward(self, x: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor:
376
- return self.experts(x, router_logits)
357
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
358
+ """Returns raw router logits [T, E]."""
359
+ x = self.norm(x)
360
+ x = x * self.root_size.to(x.dtype)
361
+ x = x * self.scale.to(x.dtype)
362
+ router_logits, _ = self.proj(x)
363
+ return router_logits
364
365
366
class Gemma4Attention(nn.Module):
@@ -433,7 +420,7 @@ class Gemma4Attention(nn.Module):
420
else:
421
# QKVParallelLinear handles GQA correctly for all layer types.
422
# k_eq_v layers load K weights into both K and V slots via
436
- # _weight_iterator remapping — no structural difference needed.
423
+ # hf_to_vllm_mapper — no structural difference needed.
424
self.qkv_proj = QKVParallelLinear(
425
hidden_size,
426
self.head_dim,
@@ -645,7 +632,7 @@ class Gemma4DecoderLayer(nn.Module):
632
config, "use_second_mlp_block", False
633
)
634
self.router: Gemma4Router | None
648
- self.moe: Gemma4MoE | None
635
+ self.experts: MoERunner | None
636
self.post_feedforward_layernorm_1: RMSNorm | None
637
self.post_feedforward_layernorm_2: RMSNorm | None
638
self.pre_feedforward_layernorm_2: RMSNorm | None
@@ -655,10 +642,20 @@ class Gemma4DecoderLayer(nn.Module):
642
quant_config=quant_config,
643
prefix=f"{prefix}.router",
644
)
658
- self.moe = Gemma4MoE(
659
- config,
645
+ names = ("moe_intermediate_size", "expert_intermediate_size")
646
+ intermediate_size = getattr_iter(config, names)
647
+ if intermediate_size is None:
648
+ raise ValueError("Gemma4 MoE requires an expert intermediate size")
649
+ self.experts = FusedMoEFactory(
650
+ num_experts=config.num_experts,
651
+ top_k=config.top_k_experts,
652
+ hidden_size=config.hidden_size,
653
+ intermediate_size=intermediate_size,
654
+ renormalize=True,
655
quant_config=quant_config,
661
- prefix=f"{prefix}.moe",
656
+ prefix=f"{prefix}.experts",
657
+ custom_routing_function=self.router.routing_function,
658
+ activation="gelu_tanh",
659
)
660
self.post_feedforward_layernorm_1 = RMSNorm(
661
config.hidden_size, eps=config.rms_norm_eps
@@ -671,7 +668,7 @@ class Gemma4DecoderLayer(nn.Module):
668
)
669
else:
670
self.router = None
674
- self.moe = None
671
+ self.experts = None
672
self.post_feedforward_layernorm_1 = None
673
self.post_feedforward_layernorm_2 = None
674
self.pre_feedforward_layernorm_2 = None
@@ -747,13 +744,13 @@ class Gemma4DecoderLayer(nn.Module):
744
assert self.post_feedforward_layernorm_1 is not None
745
assert self.pre_feedforward_layernorm_2 is not None
746
assert self.router is not None
750
- assert self.moe is not None
747
+ assert self.experts is not None
748
assert self.post_feedforward_layernorm_2 is not None
749
hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states)
750
751
hidden_states_2 = self.pre_feedforward_layernorm_2(residual)
752
router_logits = self.router(residual)
756
- hidden_states_2 = self.moe(hidden_states_2, router_logits)
753
+ hidden_states_2 = self.experts(hidden_states_2, router_logits)
754
hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2)
755
756
# Combine MLP and MoE outputs
@@ -983,7 +980,7 @@ class Gemma4CrossDecoderLayers(nn.Module):
980
@support_torch_compile(
981
enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill
982
)
986
-class Gemma4Model(nn.Module, EagleModelMixin):
983
+class Gemma4Model(nn.Module, EagleModelMixin, SupportsQuant):
984
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
985
super().__init__()
986
config = _get_text_config(vllm_config.model_config.hf_config)
@@ -991,6 +988,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
988
quant_config = vllm_config.quant_config
989
self.config = config
990
self.quant_config = quant_config
991
+ self.hf_to_vllm_mapper = _gemma4_layer_weights_mapper(config)
992
993
# PLE config values (default to 0 if not present — disables PLE)
994
self.hidden_size_per_layer_input = getattr(
@@ -1391,152 +1389,21 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1389
return hidden_states
1390
1391
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1394
- stacked_params_mapping = [
1395
- # (param_name, shard_name, shard_id)
1396
- ("qkv_proj", "q_proj", "q"),
1397
- ("qkv_proj", "k_proj", "k"),
1398
- ("qkv_proj", "v_proj", "v"),
1399
- ("gate_up_proj", "gate_proj", 0),
1400
- ("gate_up_proj", "up_proj", 1),
1401
- ]
1402
-
1403
- # MoE expert weight mapping: checkpoint can have either:
1404
- # 1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D)
1405
- # 2. Already per-expert 2D weights (if quantized)
1406
- # Map to MoERunner parameters:
1407
- # moe.experts.{id}.gate_proj → MoERunner w1 (shard of w13)
1408
- # moe.experts.{id}.up_proj → MoERunner w3 (shard of w13)
1409
- # moe.experts.{id}.down_proj → MoERunner w2
1410
- num_experts = getattr(self.config, "num_experts", None) or 0
1411
- # Strategy A: dot-separated suffix
1412
- # (standard AWQ/GPTQ e.g. .qweight, .scales, .weight)
1413
- dot_suffix_expert_params_mapping = fused_moe_make_expert_params_mapping(
1414
- self,
1415
- ckpt_gate_proj_name="gate_proj",
1416
- ckpt_down_proj_name="down_proj",
1417
- ckpt_up_proj_name="up_proj",
1418
- num_experts=num_experts,
1419
- )
1420
- # Strategy B: underscore-separated suffix
1421
- # (CompressedTensors-format AWQ/W4A16 _packed, _scale)
1422
- underscore_suffix_expert_params_mapping = [
1423
- (
1424
- f"{param_name}weight_",
1425
- f"{weight_name.rstrip('.')}_",
1426
- expert_id,
1427
- shard_id,
1428
- )
1429
- for (
1430
- param_name,
1431
- weight_name,
1432
- expert_id,
1433
- shard_id,
1434
- ) in dot_suffix_expert_params_mapping
1435
- ]
1436
- expert_params_mapping = (
1437
- dot_suffix_expert_params_mapping + underscore_suffix_expert_params_mapping
1438
- )
1439
- params_dict = dict(self.named_parameters())
1440
- # Include buffers (e.g. layer_scalar) so they can be loaded too
1441
- params_dict.update(dict(self.named_buffers()))
1442
- loaded_params: set[str] = set()
1443
- for name, loaded_weight in weights:
1444
- if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
1445
- remapped_name = maybe_remap_kv_scale_name(name, params_dict)
1446
- if remapped_name is not None and remapped_name in params_dict:
1447
- param = params_dict[remapped_name]
1448
- weight_loader = getattr(
1449
- param, "weight_loader", default_weight_loader
1450
- )
1451
- weight_loader(param, loaded_weight)
1452
- loaded_params.add(remapped_name)
1453
- continue
1454
-
1455
- for param_name, shard_name, shard_id in stacked_params_mapping:
1456
- if shard_name not in name:
1457
- continue
1458
- stacked_name = name.replace(shard_name, param_name)
1459
- # KV-shared layers have a plain q_proj instead of a packed
1460
- # qkv_proj: fall through to the direct load, which also
1461
- # skips the redundant K/V tensors original checkpoints ship.
1462
- if stacked_name not in params_dict:
1463
- continue
1464
- if is_pp_missing_parameter(stacked_name, self):
1465
- continue
1466
- param = params_dict[stacked_name]
1467
- weight_loader = param.weight_loader
1468
- weight_loader(param, loaded_weight, shard_id)
1469
- loaded_params.add(stacked_name)
1470
- break
1471
- else:
1472
- for (
1473
- param_name,
1474
- weight_name,
1475
- expert_id,
1476
- shard_id,
1477
- ) in expert_params_mapping:
1478
- # Match both:
1479
- # - Bare weights: "experts.0.down_proj" (from 3D explosion)
1480
- # - With suffix: "experts.0.down_proj.weight_scale" (2D quantized)
1481
- # weight_name has trailing dot, so check with and without it
1482
- weight_name_base = weight_name.rstrip(".")
1483
- if weight_name in name:
1484
- # Has suffix (e.g., .weight_scale)
1485
- moe_name = name.replace(weight_name, param_name)
1486
- elif name.endswith(weight_name_base):
1487
- # Bare weight (no suffix)
1488
- moe_name = name.replace(
1489
- weight_name_base, param_name.rstrip("_") + "_weight"
1490
- )
1491
- else:
1492
- continue
1493
- if moe_name not in params_dict:
1494
- continue
1495
- if is_pp_missing_parameter(moe_name, self):
1496
- continue
1497
- param = params_dict[moe_name]
1498
- # Expert weights are already in the correct
1499
- # orientation for MoERunner after _weight_iterator:
1500
- # gate/up: [I, H] → w1/w3 expects [I, H]
1501
- # down: [H, I] → w2 expects [H, I]
1502
- # Scales and other quantization params may be 1D or scalar.
1503
- weight_loader = param.weight_loader
1504
- weight_loader(
1505
- param,
1506
- loaded_weight,
1507
- moe_name, # Pass mapped name (handles both weights and scales)
1508
- shard_id=shard_id,
1509
- expert_id=expert_id,
1510
- )
1511
- loaded_params.add(moe_name)
1512
- break
1513
- else:
1514
- if name.endswith(".bias") and name not in params_dict:
1515
- continue
1516
- remapped_name = maybe_remap_kv_scale_name(name, params_dict)
1517
- if remapped_name is None:
1518
- continue
1519
- name = remapped_name
1520
- if is_pp_missing_parameter(name, self):
1521
- continue
1522
- # Skip if name doesn't exist in params_dict (e.g., individual
1523
- # expert weights that should have been handled above)
1524
- if name not in params_dict:
1525
- continue
1526
- param = params_dict[name]
1527
- weight_loader = getattr(
1528
- param, "weight_loader", default_weight_loader
1529
- )
1530
- weight_loader(param, loaded_weight)
1531
- loaded_params.add(name)
1532
-
1533
- return loaded_params
1392
+ loader = AutoWeightsLoader(self)
1393
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1394
1395
1396
class Gemma4ForCausalLM(
1397
nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
1398
):
1539
- hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
1399
+ hf_to_vllm_mapper = WeightsMapper(
1400
+ orig_to_new_substr={
1401
+ # Multimodal weights are handled by the multimodal wrapper.
1402
+ "audio_tower.": None,
1403
+ "vision_tower.": None,
1404
+ "embed_audio.": None,
1405
+ "embed_vision.": None,
1406
+ },
1407
orig_to_new_prefix={
1408
# Gemma4ForConditionalGeneration already loads the text stack
1409
# from `model.language_model.*`. We reuse that same checkpoint
@@ -1544,13 +1411,6 @@ class Gemma4ForCausalLM(
1411
# so LoRA keys from the conditional wrapper map onto `model.*`.
1412
"model.language_model.": "model.",
1413
},
1547
- orig_to_new_substr={
1548
- # Gemma4ForConditionalGeneration names MoE adapter targets under
1549
- # `...moe.experts.*`, while the text-only model exposes them
1550
- # under `...moe.*`.
1551
- ".moe.experts.gate_up_proj": ".moe.gate_up_proj",
1552
- ".moe.experts.down_proj": ".moe.down_proj",
1553
- },
1414
)
1415
# KV-shared layers only have q_proj, so qkv_proj packing applies to
1416
# the other layers.
@@ -1596,26 +1456,18 @@ class Gemma4ForCausalLM(
1456
)
1457
1458
# --- MixtureOfExperts protocol ---
1599
- self.moe_layers: list[nn.Module] = []
1600
- example_moe: Gemma4MoE | None = None
1601
-
1602
- for layer in self.model.layers:
1603
- if hasattr(layer, "moe") and isinstance(layer.moe, Gemma4MoE):
1604
- example_moe = layer.moe
1605
- self.moe_layers.append(layer.moe.experts)
1606
-
1459
+ self.moe_layers: list[nn.Module] = [
1460
+ layer.experts
1461
+ for layer in self.model.layers
1462
+ if not isinstance(layer, PPMissingLayer) and layer.experts is not None
1463
+ ]
1464
self.num_moe_layers = len(self.moe_layers)
1465
1609
- if example_moe is not None:
1610
- self.num_logical_experts = example_moe.num_experts
1611
- self.num_physical_experts = example_moe.num_experts
1612
- self.num_local_physical_experts = example_moe.num_experts
1613
- self.num_routed_experts = example_moe.num_experts
1614
- else:
1615
- self.num_logical_experts = 0
1616
- self.num_physical_experts = 0
1617
- self.num_local_physical_experts = 0
1618
- self.num_routed_experts = 0
1466
+ num_experts = config.num_experts if self.num_moe_layers else 0
1467
+ self.num_logical_experts = num_experts
1468
+ self.num_physical_experts = num_experts
1469
+ self.num_local_physical_experts = num_experts
1470
+ self.num_routed_experts = num_experts
1471
1472
self.num_expert_groups = 1
1473
self.num_shared_experts = 0
@@ -1644,106 +1496,5 @@ class Gemma4ForCausalLM(
1496
return self.logits_processor(self.lm_head, hidden_states)
1497
1498
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1647
- # Checkpoint weight names use "language_model." prefix (from the
1648
- # Gemma4ForConditionalGeneration wrapper). Strip it to map to our
1649
- # model tree which is just "model.*".
1650
- def _weight_iterator():
1651
- use_k_eq_v = getattr(self.config, "attention_k_eq_v", False)
1652
- # Build set of k_eq_v layer indices (full_attention layers
1653
- # when attention_k_eq_v is enabled). These layers have k_proj
1654
- # but no v_proj in checkpoint — we duplicate k_proj as v_proj.
1655
- k_eq_v_layer_indices: set[int] = set()
1656
- if use_k_eq_v:
1657
- for idx, lt in enumerate(self.config.layer_types):
1658
- if lt == "full_attention":
1659
- k_eq_v_layer_indices.add(idx)
1660
-
1661
- for name, weight in weights:
1662
- # Remap "language_model." → "" to match our model tree.
1663
- # Checkpoint: model.language_model.layers.X.*
1664
- # Our model: model.layers.X.*
1665
- name = name.replace("language_model.", "")
1666
-
1667
- # Remap new HF checkpoint naming to internal vLLM
1668
- # naming: HF moved per_expert_scale to router and
1669
- # renamed moe → experts in the MoE block.
1670
- name = name.replace(
1671
- ".router.per_expert_scale",
1672
- ".moe.per_expert_scale",
1673
- )
1674
- if ".experts.gate_up_proj" in name:
1675
- name = name.replace(
1676
- ".experts.gate_up_proj",
1677
- ".moe.gate_up_proj",
1678
- )
1679
- elif ".experts.down_proj" in name:
1680
- name = name.replace(
1681
- ".experts.down_proj",
1682
- ".moe.down_proj",
1683
- )
1684
-
1685
- # Remap individual 2D expert weights:
1686
- # .experts.{id}.{proj} → .moe.experts.{id}.{proj}
1687
- # (This handles per-expert 2D quantized weights)
1688
- name = _remap_gemma4_expert_weight_name(name)
1689
-
1690
- # MoE expert weights: checkpoint stores as 3D packed
1691
- # tensors. Explode into per-expert 2D weights for
1692
- # MoERunner weight_loader.
1693
- #
1694
- # Checkpoint format:
1695
- # moe.gate_up_proj: [E, 2*I, H] (fused gate + up)
1696
- # moe.down_proj: [E, H, I]
1697
- #
1698
- # MoERunner expects per-expert:
1699
- # w1 (gate): [I, H] — first half of gate_up
1700
- # w3 (up): [I, H] — second half of gate_up
1701
- # w2 (down): [H, I] — as-is from checkpoint
1702
- #
1703
- # No transpose needed: checkpoint orientation already
1704
- # matches MoERunner's expected layout.
1705
- if "moe.gate_up_proj" in name and weight.dim() == 3:
1706
- num_experts = weight.size(0)
1707
- intermediate_size = weight.size(1) // 2
1708
- for expert_id in range(num_experts):
1709
- gate_weight = weight[expert_id, :intermediate_size, :]
1710
- up_weight = weight[expert_id, intermediate_size:, :]
1711
- base = name.replace("moe.", f"moe.experts.{expert_id}.")
1712
- yield base.replace("gate_up_proj", "gate_proj"), gate_weight
1713
- yield base.replace("gate_up_proj", "up_proj"), up_weight
1714
- continue
1715
-
1716
- if "moe.down_proj" in name and weight.dim() == 3:
1717
- num_experts = weight.size(0)
1718
- for expert_id in range(num_experts):
1719
- expert_name = name.replace("moe.", f"moe.experts.{expert_id}.")
1720
- yield expert_name, weight[expert_id]
1721
- continue
1722
-
1723
- # k_eq_v layers: checkpoint has k_proj but no v_proj.
1724
- # QKVParallelLinear expects both, so duplicate k_proj
1725
- # as v_proj so V gets identical weights to K.
1726
- # ONLY for full_attention layers — sliding layers have
1727
- # their own real v_proj weights.
1728
- if "self_attn.k_proj" in name and k_eq_v_layer_indices:
1729
- m = re.search(r"layers\.(\d+)\.", name)
1730
- if m and int(m.group(1)) in k_eq_v_layer_indices:
1731
- yield name, weight
1732
- yield name.replace("k_proj", "v_proj"), weight.clone()
1733
- continue
1734
-
1735
- yield name, weight
1736
-
1737
- # Drop multimodal weights, which are handled by the multimodal wrapper.
1738
- # `_weight_iterator` already applies this model's renames by hand, so
1739
- # `hf_to_vllm_mapper` is deliberately not passed here.
1740
- mapper = WeightsMapper(
1741
- orig_to_new_substr={
1742
- "audio_tower.": None,
1743
- "vision_tower.": None,
1744
- "embed_audio.": None,
1745
- "embed_vision.": None,
1746
- }
1747
- )
1499
loader = AutoWeightsLoader(self)
1749
- return loader.load_weights(_weight_iterator(), mapper=mapper)
1500
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/gemma4_mm.py
+2
-5
@@ -46,10 +46,7 @@ from vllm.logger import init_logger
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import ReplicatedLinear
48
from vllm.model_executor.models.gemma3n_mm import batch_audio_features
49
-from vllm.model_executor.models.gemma4 import (
50
- _GEMMA4_EXPERT_PARENT_MAPPER,
51
- Gemma4ForCausalLM,
52
-)
49
+from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
50
from vllm.model_executor.models.module_mapping import MultiModelKeys
51
from vllm.model_executor.models.transformers.utils import recursive_replace_linear
52
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -1060,7 +1057,7 @@ class Gemma4ForConditionalGeneration(
1057
}
1058
1059
# Maps checkpoint prefixes to vLLM module paths.
1063
- hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
1060
+ hf_to_vllm_mapper = WeightsMapper(
1061
orig_to_new_prefix={
1062
# vision tower
1063
"model.vision_tower": "vision_tower",
vllm/model_executor/models/utils.py
+20
-7
@@ -41,6 +41,9 @@ if TYPE_CHECKING:
41
logger = init_logger(__name__)
42
43
ShardId: TypeAlias = str | int | tuple[int, ...]
44
+"""One shard of a stacked parameter. A tuple is a single contiguous span."""
45
+ShardIds: TypeAlias = ShardId | list[ShardId]
46
+"""One shard, or a list of shards the same weight is loaded into in turn."""
47
48
49
@dataclass
@@ -52,7 +55,9 @@ class WeightsMapper:
55
orig_to_new_renaming: list["WeightRenaming"] = field(default_factory=list)
56
orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict)
57
orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict)
55
- orig_to_new_stacked: Mapping[str, tuple[str, ShardId]] = field(default_factory=dict)
58
+ orig_to_new_stacked: Mapping[str, tuple[str, ShardIds]] = field(
59
+ default_factory=dict
60
+ )
61
orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict)
62
orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict)
63
@@ -78,11 +83,12 @@ class WeightsMapper:
83
result = self._map_name_with_shard(key)
84
return result[0] if result is not None else None
85
81
- def _map_name_with_shard(self, key: str) -> tuple[str, ShardId | None] | None:
86
+ def _map_name_with_shard(self, key: str) -> tuple[str, ShardIds | None] | None:
87
"""Map a weight name and extract any shard_id metadata.
88
89
Returns:
85
- (mapped_name, shard_id) if the name should be kept.
90
+ (mapped_name, shard_id) if the name should be kept. A list of shard
91
+ ids means the weight is loaded into each of those shards in turn.
92
None if the name should be dropped.
93
94
"""
@@ -113,7 +119,7 @@ class WeightsMapper:
119
120
key = key.replace(substr, new_key, 1)
121
116
- shard_id: ShardId | None = None
122
+ shard_id: ShardIds | None = None
123
for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items():
124
if substr in key:
125
key = key.replace(substr, new_key, 1)
@@ -143,9 +149,16 @@ class WeightsMapper:
149
if result is None:
150
continue
151
out_name, shard_id = result
146
- if shard_id is not None:
147
- data.shard_id = shard_id
148
- yield out_name, data
152
+ if shard_id is None:
153
+ yield out_name, data
154
+ continue
155
+ shard_ids = shard_id if isinstance(shard_id, list) else [shard_id]
156
+ for i, one_shard_id in enumerate(shard_ids):
157
+ # Each shard carries its own id, so extra shards need their own
158
+ # tensor object; `detach` aliases the storage instead of copying.
159
+ shard_data = data if i == 0 else data.detach()
160
+ shard_data.shard_id = one_shard_id
161
+ yield out_name, shard_data
162
163
def apply_list(self, values: list[str]) -> list[str]:
164
return [