| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | from typing import cast |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import CompilationConfig, VllmConfig |
| 11 | from vllm.config.compilation import CompilationMode |
| 12 | from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding |
| 13 | from vllm.model_executor.models import gemma |
| 14 | from vllm.model_executor.models.gemma3n import ( |
| 15 | Gemma3nTextModel, |
| 16 | _kv_sharing_weights_mapper, |
| 17 | ) |
| 18 | from vllm.model_executor.models.gemma4 import ( |
| 19 | Gemma4ForCausalLM, |
| 20 | _gemma4_layer_weights_mapper, |
| 21 | ) |
| 22 | from vllm.model_executor.models.gemma4_dspark import ( |
| 23 | Gemma4DSparkForCausalLM, |
| 24 | Gemma4DSparkModel, |
| 25 | ) |
| 26 | |
| 27 | MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"] |
| 28 | |
| 29 | |
| 30 | @pytest.mark.usefixtures("dist_init") |
| 31 | @pytest.mark.parametrize( |
| 32 | "enabled,has_weights,with_markov", |
| 33 | [ |
| 34 | pytest.param(True, True, False, id="hidden-only"), |
| 35 | pytest.param(True, True, True, id="with-markov"), |
| 36 | pytest.param(True, False, True, id="missing-weights"), |
| 37 | pytest.param(False, True, True, id="disabled-head"), |
| 38 | ], |
| 39 | ) |
| 40 | def test_gemma4_dspark_loads_confidence_head( |
| 41 | monkeypatch, enabled, has_weights, with_markov |
| 42 | ) -> None: |
| 43 | """Use checkpoint confidence parameters, or disable an unavailable head.""" |
| 44 | config = SimpleNamespace( |
| 45 | vocab_size=64, |
| 46 | hidden_size=8, |
| 47 | target_layer_ids=[0, 1], |
| 48 | num_hidden_layers=0, |
| 49 | rms_norm_eps=1e-6, |
| 50 | markov_rank=4, |
| 51 | enable_confidence_head=enabled, |
| 52 | confidence_head_with_markov=with_markov, |
| 53 | ) |
| 54 | vllm_config = SimpleNamespace( |
| 55 | compilation_config=CompilationConfig(mode=CompilationMode.NONE), |
| 56 | model_config=SimpleNamespace(dtype=torch.bfloat16), |
| 57 | speculative_config=SimpleNamespace( |
| 58 | draft_model_config=SimpleNamespace(hf_config=config), |
| 59 | ), |
| 60 | ) |
| 61 | monkeypatch.setattr(Gemma4DSparkModel, "_build_fused_kv_buffers", lambda _: None) |
| 62 | model = Gemma4DSparkForCausalLM(vllm_config=cast(VllmConfig, vllm_config)) |
| 63 | width = config.hidden_size + (config.markov_rank if with_markov else 0) |
| 64 | weight = torch.arange(width, dtype=torch.bfloat16).reshape(1, -1) / 16 |
| 65 | bias = torch.tensor([-0.5], dtype=torch.bfloat16) |
| 66 | loaded = model.load_weights( |
| 67 | [("confidence_head.proj.weight", weight), ("confidence_head.proj.bias", bias)] |
| 68 | if has_weights |
| 69 | else [] |
| 70 | ) |
| 71 | if not (enabled and has_weights): |
| 72 | assert model.model.confidence_head is None |
| 73 | assert not loaded |
| 74 | return |
| 75 | |
| 76 | assert loaded == { |
| 77 | "model.confidence_head.proj.weight", |
| 78 | "model.confidence_head.proj.bias", |
| 79 | } |
| 80 | assert model.model.confidence_head.proj.weight.dtype == torch.float32 |
| 81 | hidden = torch.full((3, config.hidden_size), 0.25, dtype=torch.bfloat16) |
| 82 | markov = torch.full((3, config.markov_rank), -0.5, dtype=torch.bfloat16) |
| 83 | inputs = torch.cat([hidden, markov], dim=-1) if with_markov else hidden |
| 84 | expected = (inputs.float() @ weight.float().T + bias.float()).sigmoid().squeeze(-1) |
| 85 | torch.testing.assert_close(model.compute_confidence(hidden, markov), expected) |
| 86 | |
| 87 | |
| 88 | @pytest.mark.cpu_test |
| 89 | @pytest.mark.usefixtures("dist_init") |
| 90 | def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None: |
| 91 | """A physical LM head must load after checkpoint-driven untying.""" |
| 92 | |
| 93 | class StubGemmaModel(torch.nn.Module): |
| 94 | def __init__(self, *, vllm_config, prefix): |
| 95 | super().__init__() |
| 96 | self.embed_tokens = VocabParallelEmbedding(4, 2) |
| 97 | self.make_empty_intermediate_tensors = None |
| 98 | |
| 99 | monkeypatch.setattr(gemma, "GemmaModel", StubGemmaModel) |
| 100 | config = SimpleNamespace( |
| 101 | vocab_size=4, |
| 102 | hidden_size=2, |
| 103 | tie_word_embeddings=False, |
| 104 | ) |
| 105 | vllm_config = SimpleNamespace( |
| 106 | model_config=SimpleNamespace(hf_config=config), |
| 107 | quant_config=None, |
| 108 | ) |
| 109 | model = gemma.GemmaForCausalLM(vllm_config=cast(VllmConfig, vllm_config)) |
| 110 | embedding_weight = torch.full((4, 2), 1.0) |
| 111 | lm_head_weight = torch.full((4, 2), 2.0) |
| 112 | |
| 113 | loaded = model.load_weights( |
| 114 | [ |
| 115 | ("model.embed_tokens.weight", embedding_weight), |
| 116 | ("lm_head.weight", lm_head_weight), |
| 117 | ] |
| 118 | ) |
| 119 | |
| 120 | assert loaded == {"model.embed_tokens.weight", "lm_head.weight"} |
| 121 | assert torch.equal(model.model.embed_tokens.weight[:4], embedding_weight) |
| 122 | assert torch.equal(model.lm_head.weight[:4], lm_head_weight) |
| 123 | |
| 124 | |
| 125 | @pytest.mark.cpu_test |
| 126 | def test_gemma4_attention_mapper() -> None: |
| 127 | """Layers with a qkv_proj pack q/k/v; `attention_k_eq_v` full-attention |
| 128 | layers also load K as the V shard, and leave any v_proj such a checkpoint |
| 129 | ships unmapped so it fails the load rather than silently overwriting V; |
| 130 | KV-shared layers keep q_proj and drop the K/V tensors original checkpoints |
| 131 | still ship for them.""" |
| 132 | config = SimpleNamespace( |
| 133 | num_hidden_layers=3, |
| 134 | num_kv_shared_layers=1, |
| 135 | attention_k_eq_v=True, |
| 136 | layer_types=["sliding_attention", "full_attention", "sliding_attention"], |
| 137 | ) |
| 138 | weights = [ |
| 139 | (f"model.layers.{i}.self_attn.{tensor}.weight", torch.full((2, 2), i + 1.0)) |
| 140 | for i in range(3) |
| 141 | for tensor in ("q_proj", "k_proj", "k_norm") |
| 142 | ] + [ |
| 143 | ("model.layers.0.mlp.up_proj.weight", torch.empty(0)), |
| 144 | ("model.layers.1.self_attn.v_proj.weight", torch.empty(0)), |
| 145 | ] |
| 146 | |
| 147 | mapper = _gemma4_layer_weights_mapper(config) |
| 148 | mapped = list(mapper.apply(weights)) |
| 149 | |
| 150 | assert [(name, getattr(w, "shard_id", None)) for name, w in mapped] == [ |
| 151 | ("model.layers.0.self_attn.qkv_proj.weight", "q"), |
| 152 | ("model.layers.0.self_attn.qkv_proj.weight", "k"), |
| 153 | ("model.layers.0.self_attn.k_norm.weight", None), |
| 154 | ("model.layers.1.self_attn.qkv_proj.weight", "q"), |
| 155 | ("model.layers.1.self_attn.qkv_proj.weight", "k"), |
| 156 | ("model.layers.1.self_attn.qkv_proj.weight", "v"), |
| 157 | ("model.layers.1.self_attn.k_norm.weight", None), |
| 158 | ("model.layers.2.self_attn.q_proj.weight", None), |
| 159 | ("model.layers.0.mlp.gate_up_proj.weight", 1), |
| 160 | ("model.layers.1.self_attn.v_proj.weight", None), |
| 161 | ] |
| 162 | k_weight, v_weight = weights[4][1], mapped[5][1] |
| 163 | assert torch.equal(v_weight, k_weight) and v_weight is not k_weight |
| 164 | |
| 165 | |
| 166 | @pytest.mark.cpu_test |
| 167 | def test_gemma4_expert_names_strip_language_model_prefix() -> None: |
| 168 | """The text-only path reuses the conditional wrapper's checkpoint naming, |
| 169 | so fused and per-expert tensors reach the experts under `model.*`.""" |
| 170 | prefix = "model.language_model.layers.0." |
| 171 | weights = [ |
| 172 | (prefix + name, torch.empty(0)) |
| 173 | for name in ( |
| 174 | "experts.gate_up_proj", |
| 175 | "experts.3.down_proj.weight_packed", |
| 176 | "router.per_expert_scale", |
| 177 | ) |
| 178 | ] |
| 179 | |
| 180 | mapped = [ |
| 181 | (name, getattr(w, "shard_id", None)) |
| 182 | for name, w in Gemma4ForCausalLM.hf_to_vllm_mapper.apply(weights) |
| 183 | ] |
| 184 | |
| 185 | assert mapped == [ |
| 186 | ("model.layers.0.experts.gate_up_proj", None), |
| 187 | ("model.layers.0.experts.3.down_proj.weight_packed", None), |
| 188 | ("model.layers.0.router.per_expert_scale", None), |
| 189 | ] |
| 190 | |
| 191 | |
| 192 | @pytest.mark.cpu_test |
| 193 | def test_gemma3n_kv_shared_layer_mapper() -> None: |
| 194 | """Only non-shared layers pack q/k/v into qkv_proj; KV-shared layers keep |
| 195 | q_proj and drop the redundant K/V tensors original checkpoints ship.""" |
| 196 | config = SimpleNamespace(num_hidden_layers=4, num_kv_shared_layers=2) |
| 197 | mapper = Gemma3nTextModel.hf_to_vllm_mapper | _kv_sharing_weights_mapper(config) |
| 198 | weights = [ |
| 199 | (f"layers.{i}.self_attn.{tensor}.weight", torch.empty(0)) |
| 200 | for i in (1, 3) |
| 201 | for tensor in ("q_proj", "k_proj", "v_proj", "k_norm", "o_proj") |
| 202 | ] |
| 203 | |
| 204 | mapped = [ |
| 205 | (name, getattr(weight, "shard_id", None)) |
| 206 | for name, weight in mapper.apply(weights) |
| 207 | ] |
| 208 | |
| 209 | assert mapped == [ |
| 210 | ("layers.1.self_attn.qkv_proj.weight", "q"), |
| 211 | ("layers.1.self_attn.qkv_proj.weight", "k"), |
| 212 | ("layers.1.self_attn.qkv_proj.weight", "v"), |
| 213 | ("layers.1.self_attn.k_norm.weight", None), |
| 214 | ("layers.1.self_attn.o_proj.weight", None), |
| 215 | ("layers.3.self_attn.q_proj.weight", None), |
| 216 | ("layers.3.self_attn.o_proj.weight", None), |
| 217 | ] |
| 218 | |
| 219 | |
| 220 | @pytest.mark.parametrize("model", MODELS) |
| 221 | def test_dummy_loader(vllm_runner, monkeypatch, model: str) -> None: |
| 222 | with monkeypatch.context() as m: |
| 223 | m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 224 | with vllm_runner( |
| 225 | model, |
| 226 | load_format="dummy", |
| 227 | ) as llm: |
| 228 | if model == "google/gemma-3-4b-it": |
| 229 | normalizers = llm.llm.collective_rpc( |
| 230 | lambda self: ( |
| 231 | self.model_runner.model.language_model.model.normalizer.cpu().item() |
| 232 | ) # noqa: E501 |
| 233 | ) |
| 234 | config = llm.llm.llm_engine.model_config.hf_config.text_config |
| 235 | else: |
| 236 | normalizers = llm.llm.collective_rpc( |
| 237 | lambda self: self.model_runner.model.model.normalizer.cpu().item() |
| 238 | ) |
| 239 | config = llm.llm.llm_engine.model_config.hf_config |
| 240 | assert np.allclose(normalizers, config.hidden_size**0.5, rtol=2e-3) |