diff --git a/tests/lora/test_lora_checkpoints.py b/tests/lora/test_lora_checkpoints.py index 8db529223b..9dd0296a88 100644 --- a/tests/lora/test_lora_checkpoints.py +++ b/tests/lora/test_lora_checkpoints.py @@ -147,10 +147,10 @@ def test_gemma4_lora_weights_mapping(): def test_gemma4_moe_lora_weights_mapping(): mapper = Gemma4ForCausalLM.hf_to_vllm_mapper name = ( - "base_model.model.model.language_model.layers.9.moe.experts." + "base_model.model.model.language_model.layers.9.experts." "gate_up_proj.lora_B.weight" ) assert parse_fine_tuned_lora_name(name, mapper) == ( - "model.layers.9.moe.gate_up_proj", + "model.layers.9.experts.gate_up_proj", False, ) diff --git a/tests/models/language/generation/test_gemma.py b/tests/models/language/generation/test_gemma.py index c2c8e88e42..5cffe6fb19 100644 --- a/tests/models/language/generation/test_gemma.py +++ b/tests/models/language/generation/test_gemma.py @@ -14,7 +14,10 @@ from vllm.model_executor.models.gemma3n import ( Gemma3nTextModel, _kv_sharing_weights_mapper, ) -from vllm.model_executor.models.gemma4 import Gemma4Model +from vllm.model_executor.models.gemma4 import ( + Gemma4ForCausalLM, + _gemma4_layer_weights_mapper, +) MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"] @@ -57,36 +60,70 @@ def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None: @pytest.mark.cpu_test -def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None: - """KV-shared layers have q_proj instead of a packed qkv_proj; their - redundant K/V tensors in original checkpoints have no parameter and are - skipped rather than failing the load.""" - model = torch.nn.Module() - model.config = SimpleNamespace(num_experts=0) - model.start_layer, model.end_layer = 0, 2 - model.layers = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()]) - for layer in model.layers: - layer.self_attn = torch.nn.Module() - model.layers[0].self_attn.qkv_proj = torch.nn.Linear(2, 6, bias=False) - model.layers[1].self_attn.q_proj = torch.nn.Linear(2, 2, bias=False) - shards: list[str] = [] - model.layers[0].self_attn.qkv_proj.weight.weight_loader = ( - lambda param, weight, shard_id: shards.append(shard_id) +def test_gemma4_attention_mapper() -> None: + """Layers with a qkv_proj pack q/k/v; `attention_k_eq_v` full-attention + layers also load K as the V shard, and leave any v_proj such a checkpoint + ships unmapped so it fails the load rather than silently overwriting V; + KV-shared layers keep q_proj and drop the K/V tensors original checkpoints + still ship for them.""" + config = SimpleNamespace( + num_hidden_layers=3, + num_kv_shared_layers=1, + attention_k_eq_v=True, + layer_types=["sliding_attention", "full_attention", "sliding_attention"], ) - - q_weight = torch.full((2, 2), 2.0) weights = [ - (f"layers.{i}.self_attn.{tensor}.weight", q_weight) - for i in (0, 1) - for tensor in ("q_proj", "k_proj", "v_proj") - ] + [("layers.1.self_attn.k_norm.weight", torch.ones(2))] - loaded = Gemma4Model.load_weights(cast(Gemma4Model, model), weights) + (f"model.layers.{i}.self_attn.{tensor}.weight", torch.full((2, 2), i + 1.0)) + for i in range(3) + for tensor in ("q_proj", "k_proj", "k_norm") + ] + [ + ("model.layers.0.mlp.up_proj.weight", torch.empty(0)), + ("model.layers.1.self_attn.v_proj.weight", torch.empty(0)), + ] - assert shards == ["q", "k", "v"] - assert "layers.1.self_attn.q_proj.weight" in loaded - assert not any(name.startswith("layers.1.self_attn.k") for name in loaded) - assert not any(name.startswith("layers.1.self_attn.v") for name in loaded) - assert torch.equal(model.layers[1].self_attn.q_proj.weight, q_weight) + mapper = _gemma4_layer_weights_mapper(config) + mapped = list(mapper.apply(weights)) + + assert [(name, getattr(w, "shard_id", None)) for name, w in mapped] == [ + ("model.layers.0.self_attn.qkv_proj.weight", "q"), + ("model.layers.0.self_attn.qkv_proj.weight", "k"), + ("model.layers.0.self_attn.k_norm.weight", None), + ("model.layers.1.self_attn.qkv_proj.weight", "q"), + ("model.layers.1.self_attn.qkv_proj.weight", "k"), + ("model.layers.1.self_attn.qkv_proj.weight", "v"), + ("model.layers.1.self_attn.k_norm.weight", None), + ("model.layers.2.self_attn.q_proj.weight", None), + ("model.layers.0.mlp.gate_up_proj.weight", 1), + ("model.layers.1.self_attn.v_proj.weight", None), + ] + k_weight, v_weight = weights[4][1], mapped[5][1] + assert torch.equal(v_weight, k_weight) and v_weight is not k_weight + + +@pytest.mark.cpu_test +def test_gemma4_expert_names_strip_language_model_prefix() -> None: + """The text-only path reuses the conditional wrapper's checkpoint naming, + so fused and per-expert tensors reach the experts under `model.*`.""" + prefix = "model.language_model.layers.0." + weights = [ + (prefix + name, torch.empty(0)) + for name in ( + "experts.gate_up_proj", + "experts.3.down_proj.weight_packed", + "router.per_expert_scale", + ) + ] + + mapped = [ + (name, getattr(w, "shard_id", None)) + for name, w in Gemma4ForCausalLM.hf_to_vllm_mapper.apply(weights) + ] + + assert mapped == [ + ("model.layers.0.experts.gate_up_proj", None), + ("model.layers.0.experts.3.down_proj.weight_packed", None), + ("model.layers.0.router.per_expert_scale", None), + ] @pytest.mark.cpu_test diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index 02903504c6..0dac2dbf2a 100644 --- a/tests/models/test_utils.py +++ b/tests/models/test_utils.py @@ -217,3 +217,25 @@ def test_get_rename_mapper_keeps_only_renames(): for name in ("drop_regex.w", "drop_substr.w", "drop_prefix.w", "w.drop_suffix"): assert mapper._map_name(name) is None assert renames._map_name(name) == name + + +def test_weights_mapper_stacks_one_weight_into_several_shards(): + """One checkpoint tensor can feed more than one shard of a stacked + parameter, and a dropped name stays dropped rather than reappearing as a + shard of the stacked one.""" + mapper = WeightsMapper( + orig_to_new_substr={"layers.1.k_proj.": None}, + orig_to_new_stacked={".k_proj.": (".qkv_proj.", ["k", "v"])}, + ) + weight = torch.ones(2) + weights = [(f"layers.{i}.k_proj.weight", weight) for i in (0, 1)] + + mapped = list(mapper.apply(weights)) + + assert [(name, w.shard_id) for name, w in mapped] == [ + ("layers.0.qkv_proj.weight", "k"), + ("layers.0.qkv_proj.weight", "v"), + ] + # Each shard needs its own tensor object, but they alias one allocation. + assert mapped[0][1] is not mapped[1][1] + assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() == weight.data_ptr() diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 60b207c184..4ca1af53a2 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -341,7 +341,7 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers(): "quant_algo": "NVFP4", "group_size": 16, }, - "model.language_model.layers.1.moe.experts.gate_up_proj": { + "model.language_model.layers.1.experts.gate_up_proj": { "quant_algo": "NVFP4", "group_size": 16, }, @@ -353,10 +353,10 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers(): ) config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_rename_mapper()) - expected_prefix = "language_model.model.layers.0.moe.experts" + expected_prefix = "language_model.model.layers.0.experts" assert set(config.quantized_layers) == { expected_prefix, - "language_model.model.layers.1.moe.gate_up_proj", + "language_model.model.layers.1.experts.gate_up_proj", } assert config._resolve_quant_algo(expected_prefix) == "NVFP4" diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index af0cb35168..558bc82bcb 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -46,7 +46,11 @@ from vllm.model_executor.models.gemma4_mm import ( ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.models.transformers.utils import recursive_replace_linear -from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.utils.torch_utils import async_tensor_h2d @@ -160,14 +164,12 @@ class DiffusionGemmaForConditionalGeneration( hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ + "model.decoder.self_conditioning.": "self_conditioning.", "model.decoder.": "model.", "model.encoder.language_model.": "model.", "model.encoder.vision_tower.": "vision_tower.", "model.encoder.embed_vision.": "embed_vision.", }, - orig_to_new_substr={ - ".experts.": ".moe.experts.", - }, ) packed_modules_mapping = { @@ -336,113 +338,12 @@ class DiffusionGemmaForConditionalGeneration( logits = _softcap_logits(logits, self.final_logit_softcapping) return logits - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - """Load weights from checkpoint. - - Checkpoint layout (HF DiffusionGemma): - model.encoder.vision_tower.* → vision tower - model.encoder.embed_vision.* → vision embedder - model.encoder.language_model.layers.* → backbone - model.decoder.layers.* → backbone (tied) - model.decoder.embed_tokens.* → embeddings - model.decoder.self_conditioning.* → self-conditioning MLP - lm_head.* → LM head (tied) - - We load encoder weights into our single ``Gemma4Model`` backbone, - skip duplicate decoder backbone weights, handle vision tower and - self-conditioning separately. - """ - sc_params = dict( - (n, p) - for n, p in self.named_parameters() - if n.startswith("self_conditioning.") + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Some checkpoints carry a vestigial Gemma3n-style embedding table. + loader = AutoWeightsLoader( + self, ignore_unexpected_prefixes=["embed_vision.embedding."] ) - - # Collect vision tower + embedder parameters AND buffers for manual - # loading. The HF vision tower registers std_bias / std_scale as - # buffers (not parameters) when config.standardize is True, so we - # must include named_buffers() to avoid "not found in model" warnings. - vision_params: dict[str, torch.Tensor] = {} - for n, p in self.named_parameters(): - if n.startswith(("vision_tower.", "embed_vision.")): - vision_params[n] = p - for n, b in self.named_buffers(): - if n.startswith(("vision_tower.", "embed_vision.")): - vision_params[n] = b - - def _remap_weights(): - # Use full weight names (including suffixes like .weight_scale, - # .weight_packed) for dedup instead of just the base layer name. Critical - # for quantized checkpoints where each weight has multiple tensors; - # tracking only base names skips scales as duplicates. - seen_weights: set[str] = set() - for name, weight in weights: - # Self-conditioning lives under model.decoder.self_conditioning.* - # in the checkpoint but at self_conditioning.* in our model. - if "self_conditioning" in name: - sc_name = name.split("self_conditioning.", 1)[1] - sc_name = "self_conditioning." + sc_name - if sc_name in sc_params: - sc_params[sc_name].data.copy_(weight) - continue - - # Vision tower: model.encoder.vision_tower.* → vision_tower.* - # In HF, the vision tower is a sibling of language_model - # under the encoder module. - if name.startswith("model.encoder.vision_tower."): - vt_name = name[len("model.encoder.") :] - if vt_name in vision_params: - vision_params[vt_name].data.copy_(weight) - else: - logger.warning( - "Vision tower weight %s (mapped to %s) not found in model", - name, - vt_name, - ) - continue - - # Vision embedder: model.encoder.embed_vision.* → embed_vision.* - if name.startswith("model.encoder.embed_vision."): - ev_name = name[len("model.encoder.") :] - if ev_name in vision_params: - vision_params[ev_name].data.copy_(weight) - else: - logger.warning( - "Embed vision weight %s (mapped to %s) not found in model", - name, - ev_name, - ) - continue - - # Skip vestigial embed_vision.embedding weights. - if "embed_vision.embedding." in name: - continue - - # Encoder backbone → model.* - if name.startswith("model.encoder.language_model."): - name = name.replace("model.encoder.language_model.", "model.") - # Decoder backbone → model.* (skip exact duplicates) - elif name.startswith("model.decoder."): - name = name.replace("model.decoder.", "model.") - - # Skip only if we've seen the exact same weight name (including scales) - if name in seen_weights: - continue - seen_weights.add(name) - yield name, weight - - # Delegate to Gemma4ForCausalLM.load_weights for the backbone, - # which handles stacked params, MoE, k_eq_v, etc. - # Temporarily set self.config to text_config since Gemma4's - # load_weights expects it (e.g. tie_word_embeddings, layer_types). - from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM - - saved_config = self.config - self.config = self.model.config - try: - Gemma4ForCausalLM.load_weights(self, _remap_weights()) - finally: - self.config = saved_config + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 04e22523c2..e3dfe65c2f 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -22,12 +22,12 @@ from collections.abc import Iterable from dataclasses import replace from itertools import islice -import regex as re import torch from torch import nn from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig +from vllm.config.utils import getattr_iter from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -40,7 +40,7 @@ from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, GateLinear, - fused_moe_make_expert_params_mapping, + MoERunner, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -57,10 +57,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config @@ -73,27 +69,63 @@ from .interfaces import ( SupportsEagle3, SupportsLoRA, SupportsPP, + SupportsQuant, ) from .utils import ( AutoWeightsLoader, + PPMissingLayer, + ShardIds, WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_layers, maybe_prefix, ) logger = init_logger(__name__) -_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper( - orig_to_new_regex={ - re.compile(r"(? WeightsMapper: + """Stack each decoder layer's projections the way its modules are built. + + gate/up pack into `mlp.gate_up_proj`. Layers with a qkv_proj pack q/k/v + into it; `attention_k_eq_v` full-attention layers ship k_proj only, so K + is loaded into both the K and V shards and no v_proj is expected. KV-shared + layers only have q_proj; the K/V tensors original checkpoints still ship + for them are dropped. + """ + num_layers = config.num_hidden_layers + first_kv_shared = num_layers - getattr(config, "num_kv_shared_layers", 0) + k_eq_v_layers = ( + { + i + for i, layer_type in enumerate(config.layer_types[:first_kv_shared]) + if layer_type == "full_attention" + } + if getattr(config, "attention_k_eq_v", False) + else set() + ) + qkv_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", "k"), ("v", "v")) + k_eq_v_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", ["k", "v"])) + stacked: dict[str, tuple[str, ShardIds]] = { + ".mlp.gate_proj.": (".mlp.gate_up_proj.", 0), + ".mlp.up_proj.": (".mlp.gate_up_proj.", 1), } -) - - -def _remap_gemma4_expert_weight_name(name: str) -> str: - return re.sub(r"(? tuple[torch.Tensor, torch.Tensor]: + """Softmax over ALL experts → top-k → renormalize. + + MoERunner's built-in fused_topk scopes softmax differently, so Gemma4 + needs its own routing for numerical correctness. + + NOTE: self.per_expert_scale is read at call time (not captured into a + local) so that torch.func.functional_call parameter substitution + reaches the routing function correctly. + """ + if current_platform.is_cuda_alike() or current_platform.is_xpu(): + return gemma4_fused_routing_kernel_triton( + gating_output, topk, self.per_expert_scale + ) + + return gemma4_routing_function_torch(gating_output, topk, self.per_expert_scale) + def forward(self, x: torch.Tensor) -> torch.Tensor: """Returns raw router logits [T, E].""" x = self.norm(x) @@ -305,77 +363,6 @@ class Gemma4Router(nn.Module): return router_logits -class Gemma4MoE(nn.Module): - """Mixture of Experts for Gemma4 using vLLM's MoERunner. - - Wraps MoERunner with custom routing. The router projection is - external (Gemma4Router) — this class only handles expert dispatch. - - Gemma4 routing: softmax over ALL experts → top-k → renormalize. - per_expert_scale is folded into routing weights for mathematical - correctness with MoERunner's fused kernel. - """ - - def __init__( - self, - config, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - self.num_experts = config.num_experts - - # Per-expert output scale folded into routing weights so that - # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e) - self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts)) - - # Gemma4 routing: softmax over ALL experts → top-k → renormalize. - # MoERunner's built-in fused_topk scopes softmax differently, so - # a custom routing function is needed for numerical correctness. - # NOTE: self.per_expert_scale is read at call time (not captured into - # a local) so that torch.func.functional_call parameter substitution - # reaches the routing function correctly. - def routing_function( - hidden_states: torch.Tensor, - gating_output: torch.Tensor, - topk: int, - renormalize: bool, - ) -> tuple[torch.Tensor, torch.Tensor]: - if current_platform.is_cuda_alike() or current_platform.is_xpu(): - return gemma4_fused_routing_kernel_triton( - gating_output, topk, self.per_expert_scale - ) - - return gemma4_routing_function_torch( - gating_output, topk, self.per_expert_scale - ) - - # MoERunner experts with custom Gemma4 routing - intermediate_size = getattr( - config, - "moe_intermediate_size", - getattr(config, "expert_intermediate_size", None), - ) - if intermediate_size is None: - raise ValueError("Gemma4 MoE requires an expert intermediate size") - - self.experts = FusedMoEFactory( - num_experts=config.num_experts, - top_k=config.top_k_experts, - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - renormalize=True, - quant_config=quant_config, - prefix=f"{prefix}.experts", - custom_routing_function=routing_function, - activation="gelu_tanh", - ) - - def forward(self, x: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor: - return self.experts(x, router_logits) - - class Gemma4Attention(nn.Module): def __init__( self, @@ -433,7 +420,7 @@ class Gemma4Attention(nn.Module): else: # QKVParallelLinear handles GQA correctly for all layer types. # k_eq_v layers load K weights into both K and V slots via - # _weight_iterator remapping — no structural difference needed. + # hf_to_vllm_mapper — no structural difference needed. self.qkv_proj = QKVParallelLinear( hidden_size, self.head_dim, @@ -645,7 +632,7 @@ class Gemma4DecoderLayer(nn.Module): config, "use_second_mlp_block", False ) self.router: Gemma4Router | None - self.moe: Gemma4MoE | None + self.experts: MoERunner | None self.post_feedforward_layernorm_1: RMSNorm | None self.post_feedforward_layernorm_2: RMSNorm | None self.pre_feedforward_layernorm_2: RMSNorm | None @@ -655,10 +642,20 @@ class Gemma4DecoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.router", ) - self.moe = Gemma4MoE( - config, + names = ("moe_intermediate_size", "expert_intermediate_size") + intermediate_size = getattr_iter(config, names) + if intermediate_size is None: + raise ValueError("Gemma4 MoE requires an expert intermediate size") + self.experts = FusedMoEFactory( + num_experts=config.num_experts, + top_k=config.top_k_experts, + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + renormalize=True, quant_config=quant_config, - prefix=f"{prefix}.moe", + prefix=f"{prefix}.experts", + custom_routing_function=self.router.routing_function, + activation="gelu_tanh", ) self.post_feedforward_layernorm_1 = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -671,7 +668,7 @@ class Gemma4DecoderLayer(nn.Module): ) else: self.router = None - self.moe = None + self.experts = None self.post_feedforward_layernorm_1 = None self.post_feedforward_layernorm_2 = None self.pre_feedforward_layernorm_2 = None @@ -747,13 +744,13 @@ class Gemma4DecoderLayer(nn.Module): assert self.post_feedforward_layernorm_1 is not None assert self.pre_feedforward_layernorm_2 is not None assert self.router is not None - assert self.moe is not None + assert self.experts is not None assert self.post_feedforward_layernorm_2 is not None hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states) hidden_states_2 = self.pre_feedforward_layernorm_2(residual) router_logits = self.router(residual) - hidden_states_2 = self.moe(hidden_states_2, router_logits) + hidden_states_2 = self.experts(hidden_states_2, router_logits) hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2) # Combine MLP and MoE outputs @@ -983,7 +980,7 @@ class Gemma4CrossDecoderLayers(nn.Module): @support_torch_compile( enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill ) -class Gemma4Model(nn.Module, EagleModelMixin): +class Gemma4Model(nn.Module, EagleModelMixin, SupportsQuant): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = _get_text_config(vllm_config.model_config.hf_config) @@ -991,6 +988,7 @@ class Gemma4Model(nn.Module, EagleModelMixin): quant_config = vllm_config.quant_config self.config = config self.quant_config = quant_config + self.hf_to_vllm_mapper = _gemma4_layer_weights_mapper(config) # PLE config values (default to 0 if not present — disables PLE) self.hidden_size_per_layer_input = getattr( @@ -1391,152 +1389,21 @@ class Gemma4Model(nn.Module, EagleModelMixin): return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - - # MoE expert weight mapping: checkpoint can have either: - # 1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D) - # 2. Already per-expert 2D weights (if quantized) - # Map to MoERunner parameters: - # moe.experts.{id}.gate_proj → MoERunner w1 (shard of w13) - # moe.experts.{id}.up_proj → MoERunner w3 (shard of w13) - # moe.experts.{id}.down_proj → MoERunner w2 - num_experts = getattr(self.config, "num_experts", None) or 0 - # Strategy A: dot-separated suffix - # (standard AWQ/GPTQ e.g. .qweight, .scales, .weight) - dot_suffix_expert_params_mapping = fused_moe_make_expert_params_mapping( - self, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=num_experts, - ) - # Strategy B: underscore-separated suffix - # (CompressedTensors-format AWQ/W4A16 _packed, _scale) - underscore_suffix_expert_params_mapping = [ - ( - f"{param_name}weight_", - f"{weight_name.rstrip('.')}_", - expert_id, - shard_id, - ) - for ( - param_name, - weight_name, - expert_id, - shard_id, - ) in dot_suffix_expert_params_mapping - ] - expert_params_mapping = ( - dot_suffix_expert_params_mapping + underscore_suffix_expert_params_mapping - ) - params_dict = dict(self.named_parameters()) - # Include buffers (e.g. layer_scalar) so they can be loaded too - params_dict.update(dict(self.named_buffers())) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): - remapped_name = maybe_remap_kv_scale_name(name, params_dict) - if remapped_name is not None and remapped_name in params_dict: - param = params_dict[remapped_name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(remapped_name) - continue - - for param_name, shard_name, shard_id in stacked_params_mapping: - if shard_name not in name: - continue - stacked_name = name.replace(shard_name, param_name) - # KV-shared layers have a plain q_proj instead of a packed - # qkv_proj: fall through to the direct load, which also - # skips the redundant K/V tensors original checkpoints ship. - if stacked_name not in params_dict: - continue - if is_pp_missing_parameter(stacked_name, self): - continue - param = params_dict[stacked_name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(stacked_name) - break - else: - for ( - param_name, - weight_name, - expert_id, - shard_id, - ) in expert_params_mapping: - # Match both: - # - Bare weights: "experts.0.down_proj" (from 3D explosion) - # - With suffix: "experts.0.down_proj.weight_scale" (2D quantized) - # weight_name has trailing dot, so check with and without it - weight_name_base = weight_name.rstrip(".") - if weight_name in name: - # Has suffix (e.g., .weight_scale) - moe_name = name.replace(weight_name, param_name) - elif name.endswith(weight_name_base): - # Bare weight (no suffix) - moe_name = name.replace( - weight_name_base, param_name.rstrip("_") + "_weight" - ) - else: - continue - if moe_name not in params_dict: - continue - if is_pp_missing_parameter(moe_name, self): - continue - param = params_dict[moe_name] - # Expert weights are already in the correct - # orientation for MoERunner after _weight_iterator: - # gate/up: [I, H] → w1/w3 expects [I, H] - # down: [H, I] → w2 expects [H, I] - # Scales and other quantization params may be 1D or scalar. - weight_loader = param.weight_loader - weight_loader( - param, - loaded_weight, - moe_name, # Pass mapped name (handles both weights and scales) - shard_id=shard_id, - expert_id=expert_id, - ) - loaded_params.add(moe_name) - break - else: - if name.endswith(".bias") and name not in params_dict: - continue - remapped_name = maybe_remap_kv_scale_name(name, params_dict) - if remapped_name is None: - continue - name = remapped_name - if is_pp_missing_parameter(name, self): - continue - # Skip if name doesn't exist in params_dict (e.g., individual - # expert weights that should have been handled above) - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - return loaded_params + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class Gemma4ForCausalLM( nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3 ): - hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper( + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + # Multimodal weights are handled by the multimodal wrapper. + "audio_tower.": None, + "vision_tower.": None, + "embed_audio.": None, + "embed_vision.": None, + }, orig_to_new_prefix={ # Gemma4ForConditionalGeneration already loads the text stack # from `model.language_model.*`. We reuse that same checkpoint @@ -1544,13 +1411,6 @@ class Gemma4ForCausalLM( # so LoRA keys from the conditional wrapper map onto `model.*`. "model.language_model.": "model.", }, - orig_to_new_substr={ - # Gemma4ForConditionalGeneration names MoE adapter targets under - # `...moe.experts.*`, while the text-only model exposes them - # under `...moe.*`. - ".moe.experts.gate_up_proj": ".moe.gate_up_proj", - ".moe.experts.down_proj": ".moe.down_proj", - }, ) # KV-shared layers only have q_proj, so qkv_proj packing applies to # the other layers. @@ -1596,26 +1456,18 @@ class Gemma4ForCausalLM( ) # --- MixtureOfExperts protocol --- - self.moe_layers: list[nn.Module] = [] - example_moe: Gemma4MoE | None = None - - for layer in self.model.layers: - if hasattr(layer, "moe") and isinstance(layer.moe, Gemma4MoE): - example_moe = layer.moe - self.moe_layers.append(layer.moe.experts) - + self.moe_layers: list[nn.Module] = [ + layer.experts + for layer in self.model.layers + if not isinstance(layer, PPMissingLayer) and layer.experts is not None + ] self.num_moe_layers = len(self.moe_layers) - if example_moe is not None: - self.num_logical_experts = example_moe.num_experts - self.num_physical_experts = example_moe.num_experts - self.num_local_physical_experts = example_moe.num_experts - self.num_routed_experts = example_moe.num_experts - else: - self.num_logical_experts = 0 - self.num_physical_experts = 0 - self.num_local_physical_experts = 0 - self.num_routed_experts = 0 + num_experts = config.num_experts if self.num_moe_layers else 0 + self.num_logical_experts = num_experts + self.num_physical_experts = num_experts + self.num_local_physical_experts = num_experts + self.num_routed_experts = num_experts self.num_expert_groups = 1 self.num_shared_experts = 0 @@ -1644,106 +1496,5 @@ class Gemma4ForCausalLM( return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Checkpoint weight names use "language_model." prefix (from the - # Gemma4ForConditionalGeneration wrapper). Strip it to map to our - # model tree which is just "model.*". - def _weight_iterator(): - use_k_eq_v = getattr(self.config, "attention_k_eq_v", False) - # Build set of k_eq_v layer indices (full_attention layers - # when attention_k_eq_v is enabled). These layers have k_proj - # but no v_proj in checkpoint — we duplicate k_proj as v_proj. - k_eq_v_layer_indices: set[int] = set() - if use_k_eq_v: - for idx, lt in enumerate(self.config.layer_types): - if lt == "full_attention": - k_eq_v_layer_indices.add(idx) - - for name, weight in weights: - # Remap "language_model." → "" to match our model tree. - # Checkpoint: model.language_model.layers.X.* - # Our model: model.layers.X.* - name = name.replace("language_model.", "") - - # Remap new HF checkpoint naming to internal vLLM - # naming: HF moved per_expert_scale to router and - # renamed moe → experts in the MoE block. - name = name.replace( - ".router.per_expert_scale", - ".moe.per_expert_scale", - ) - if ".experts.gate_up_proj" in name: - name = name.replace( - ".experts.gate_up_proj", - ".moe.gate_up_proj", - ) - elif ".experts.down_proj" in name: - name = name.replace( - ".experts.down_proj", - ".moe.down_proj", - ) - - # Remap individual 2D expert weights: - # .experts.{id}.{proj} → .moe.experts.{id}.{proj} - # (This handles per-expert 2D quantized weights) - name = _remap_gemma4_expert_weight_name(name) - - # MoE expert weights: checkpoint stores as 3D packed - # tensors. Explode into per-expert 2D weights for - # MoERunner weight_loader. - # - # Checkpoint format: - # moe.gate_up_proj: [E, 2*I, H] (fused gate + up) - # moe.down_proj: [E, H, I] - # - # MoERunner expects per-expert: - # w1 (gate): [I, H] — first half of gate_up - # w3 (up): [I, H] — second half of gate_up - # w2 (down): [H, I] — as-is from checkpoint - # - # No transpose needed: checkpoint orientation already - # matches MoERunner's expected layout. - if "moe.gate_up_proj" in name and weight.dim() == 3: - num_experts = weight.size(0) - intermediate_size = weight.size(1) // 2 - for expert_id in range(num_experts): - gate_weight = weight[expert_id, :intermediate_size, :] - up_weight = weight[expert_id, intermediate_size:, :] - base = name.replace("moe.", f"moe.experts.{expert_id}.") - yield base.replace("gate_up_proj", "gate_proj"), gate_weight - yield base.replace("gate_up_proj", "up_proj"), up_weight - continue - - if "moe.down_proj" in name and weight.dim() == 3: - num_experts = weight.size(0) - for expert_id in range(num_experts): - expert_name = name.replace("moe.", f"moe.experts.{expert_id}.") - yield expert_name, weight[expert_id] - continue - - # k_eq_v layers: checkpoint has k_proj but no v_proj. - # QKVParallelLinear expects both, so duplicate k_proj - # as v_proj so V gets identical weights to K. - # ONLY for full_attention layers — sliding layers have - # their own real v_proj weights. - if "self_attn.k_proj" in name and k_eq_v_layer_indices: - m = re.search(r"layers\.(\d+)\.", name) - if m and int(m.group(1)) in k_eq_v_layer_indices: - yield name, weight - yield name.replace("k_proj", "v_proj"), weight.clone() - continue - - yield name, weight - - # Drop multimodal weights, which are handled by the multimodal wrapper. - # `_weight_iterator` already applies this model's renames by hand, so - # `hf_to_vllm_mapper` is deliberately not passed here. - mapper = WeightsMapper( - orig_to_new_substr={ - "audio_tower.": None, - "vision_tower.": None, - "embed_audio.": None, - "embed_vision.": None, - } - ) loader = AutoWeightsLoader(self) - return loader.load_weights(_weight_iterator(), mapper=mapper) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 6ad98cb650..f8cb9752c8 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -46,10 +46,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.gemma3n_mm import batch_audio_features -from vllm.model_executor.models.gemma4 import ( - _GEMMA4_EXPERT_PARENT_MAPPER, - Gemma4ForCausalLM, -) +from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.models.transformers.utils import recursive_replace_linear from vllm.multimodal import MULTIMODAL_REGISTRY @@ -1060,7 +1057,7 @@ class Gemma4ForConditionalGeneration( } # Maps checkpoint prefixes to vLLM module paths. - hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper( + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ # vision tower "model.vision_tower": "vision_tower", diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index eaa64793de..947a6f6276 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -41,6 +41,9 @@ if TYPE_CHECKING: logger = init_logger(__name__) ShardId: TypeAlias = str | int | tuple[int, ...] +"""One shard of a stacked parameter. A tuple is a single contiguous span.""" +ShardIds: TypeAlias = ShardId | list[ShardId] +"""One shard, or a list of shards the same weight is loaded into in turn.""" @dataclass @@ -52,7 +55,9 @@ class WeightsMapper: orig_to_new_renaming: list["WeightRenaming"] = field(default_factory=list) orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict) - orig_to_new_stacked: Mapping[str, tuple[str, ShardId]] = field(default_factory=dict) + orig_to_new_stacked: Mapping[str, tuple[str, ShardIds]] = field( + default_factory=dict + ) orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict) orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict) @@ -78,11 +83,12 @@ class WeightsMapper: result = self._map_name_with_shard(key) return result[0] if result is not None else None - def _map_name_with_shard(self, key: str) -> tuple[str, ShardId | None] | None: + def _map_name_with_shard(self, key: str) -> tuple[str, ShardIds | None] | None: """Map a weight name and extract any shard_id metadata. Returns: - (mapped_name, shard_id) if the name should be kept. + (mapped_name, shard_id) if the name should be kept. A list of shard + ids means the weight is loaded into each of those shards in turn. None if the name should be dropped. """ @@ -113,7 +119,7 @@ class WeightsMapper: key = key.replace(substr, new_key, 1) - shard_id: ShardId | None = None + shard_id: ShardIds | None = None for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items(): if substr in key: key = key.replace(substr, new_key, 1) @@ -143,9 +149,16 @@ class WeightsMapper: if result is None: continue out_name, shard_id = result - if shard_id is not None: - data.shard_id = shard_id - yield out_name, data + if shard_id is None: + yield out_name, data + continue + shard_ids = shard_id if isinstance(shard_id, list) else [shard_id] + for i, one_shard_id in enumerate(shard_ids): + # Each shard carries its own id, so extra shards need their own + # tensor object; `detach` aliases the storage instead of copying. + shard_data = data if i == 0 else data.detach() + shard_data.shard_id = one_shard_id + yield out_name, shard_data def apply_list(self, values: list[str]) -> list[str]: return [