| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test the functionality of the Transformers modeling backend.""" |
| 4 | |
| 5 | import ast |
| 6 | import contextlib |
| 7 | import os |
| 8 | import tempfile |
| 9 | from functools import partial |
| 10 | from types import SimpleNamespace |
| 11 | from typing import Any |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | import torch.nn as nn |
| 16 | from transformers import AutoConfig, AutoModel, PreTrainedConfig |
| 17 | |
| 18 | from vllm.config import ModelConfig, VllmConfig |
| 19 | from vllm.model_executor.models.interfaces import SupportsMultiModal |
| 20 | from vllm.model_executor.models.transformers.base import Base |
| 21 | from vllm.model_executor.models.transformers.fusers import AttentionFuser |
| 22 | from vllm.model_executor.models.transformers.fusers.attention import VLLM_ATTN_IMPL |
| 23 | from vllm.model_executor.models.transformers.multimodal import MultiModalMixin |
| 24 | from vllm.model_executor.models.utils import StageMissingLayer |
| 25 | |
| 26 | from ...conftest import HfRunner, VllmRunner |
| 27 | from ...utils import multi_gpu_test, prep_prompts |
| 28 | from ..registry import HF_EXAMPLE_MODELS |
| 29 | from ..utils import check_embeddings_close, check_logprobs_close |
| 30 | |
| 31 | |
| 32 | @pytest.fixture(scope="function", autouse=True) |
| 33 | def enable_pickle(monkeypatch): |
| 34 | """`LLM.apply_model` requires pickling a function.""" |
| 35 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 36 | |
| 37 | |
| 38 | def get_model(arch: str) -> str: |
| 39 | model_info = HF_EXAMPLE_MODELS.get_hf_info(arch) |
| 40 | model_info.check_transformers_version(on_fail="skip") |
| 41 | return model_info.default |
| 42 | |
| 43 | |
| 44 | def count_modules(model, cls: type) -> int: |
| 45 | return sum(isinstance(m, cls) for m in model.modules()) |
| 46 | |
| 47 | |
| 48 | def get_num_fused(model) -> tuple[int, int]: |
| 49 | from vllm.model_executor.layers.linear import ( |
| 50 | MergedColumnParallelLinear, |
| 51 | QKVParallelLinear, |
| 52 | ) |
| 53 | |
| 54 | glu = count_modules(model, MergedColumnParallelLinear) |
| 55 | qkv = count_modules(model, QKVParallelLinear) |
| 56 | return glu, qkv |
| 57 | |
| 58 | |
| 59 | def check_implementation( |
| 60 | runner_ref: type[HfRunner | VllmRunner], |
| 61 | runner_test: type[VllmRunner], |
| 62 | example_prompts: list[str], |
| 63 | model: str, |
| 64 | kwargs_ref: dict[str, Any] | None = None, |
| 65 | kwargs_test: dict[str, Any] | None = None, |
| 66 | num_fused: tuple[int, int] = (1, 1), |
| 67 | **kwargs, |
| 68 | ): |
| 69 | if kwargs_ref is None: |
| 70 | kwargs_ref = {} |
| 71 | if kwargs_test is None: |
| 72 | kwargs_test = {} |
| 73 | |
| 74 | max_tokens = 32 |
| 75 | num_logprobs = 5 |
| 76 | |
| 77 | args = (example_prompts, max_tokens, num_logprobs) |
| 78 | |
| 79 | with runner_test(model, **kwargs_test, **kwargs) as model_test: |
| 80 | model_config = model_test.llm.llm_engine.model_config |
| 81 | assert model_config.using_transformers_backend() |
| 82 | |
| 83 | num_layers = model_config.hf_config.get_text_config().num_hidden_layers |
| 84 | expected_glu, expected_qkv = num_fused |
| 85 | for num_glu, num_qkv in model_test.apply_model(get_num_fused): |
| 86 | assert num_glu == expected_glu * num_layers |
| 87 | assert num_qkv == expected_qkv * num_layers |
| 88 | |
| 89 | tp_size = kwargs_test.get("tensor_parallel_size", 1) |
| 90 | from vllm.model_executor.layers.attention import Attention |
| 91 | |
| 92 | counts = model_test.apply_model(partial(count_modules, cls=Attention)) |
| 93 | assert counts == [num_layers] * tp_size |
| 94 | |
| 95 | outputs_test = model_test.generate_greedy_logprobs(*args) |
| 96 | |
| 97 | with runner_ref(model, **kwargs_ref) as model_ref: |
| 98 | if isinstance(model_ref, VllmRunner): |
| 99 | outputs_ref = model_ref.generate_greedy_logprobs(*args) |
| 100 | else: |
| 101 | outputs_ref = model_ref.generate_greedy_logprobs_limit(*args) |
| 102 | |
| 103 | check_logprobs_close( |
| 104 | outputs_0_lst=outputs_ref, |
| 105 | outputs_1_lst=outputs_test, |
| 106 | name_0="ref", |
| 107 | name_1="test", |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | @pytest.mark.parametrize( |
| 112 | "model,model_impl,num_fused", |
| 113 | [ |
| 114 | ("meta-llama/Llama-3.2-1B-Instruct", "transformers", (1, 1)), |
| 115 | ("hmellor/Ilama-3.2-1B", "auto", (1, 1)), # CUSTOM CODE |
| 116 | ("allenai/OLMoE-1B-7B-0924", "transformers", (0, 1)), # MoE |
| 117 | ], |
| 118 | ) # trust_remote_code=True by default |
| 119 | def test_models( |
| 120 | hf_runner: type[HfRunner], |
| 121 | vllm_runner: type[VllmRunner], |
| 122 | example_prompts: list[str], |
| 123 | model: str, |
| 124 | model_impl: str, |
| 125 | num_fused: tuple[int, int], |
| 126 | ) -> None: |
| 127 | check_implementation( |
| 128 | hf_runner, |
| 129 | vllm_runner, |
| 130 | example_prompts, |
| 131 | model, |
| 132 | num_fused=num_fused, |
| 133 | model_impl=model_impl, |
| 134 | ) |
| 135 | |
| 136 | |
| 137 | def test_hybrid_attention(vllm_runner: type[VllmRunner]) -> None: |
| 138 | prompts, _, _ = prep_prompts(4, (800, 801)) |
| 139 | kwargs_ref = {"max_model_len": 8192, "enforce_eager": True} |
| 140 | kwargs_test = {"model_impl": "transformers", **kwargs_ref} |
| 141 | check_implementation( |
| 142 | vllm_runner, |
| 143 | vllm_runner, |
| 144 | prompts, |
| 145 | model="hmellor/tiny-random-Gemma2ForCausalLM", |
| 146 | kwargs_ref=kwargs_ref, |
| 147 | kwargs_test=kwargs_test, |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | def get_sinks(model) -> dict[int, torch.Tensor]: |
| 152 | """The sink tensor each attention layer was handed, keyed by layer index.""" |
| 153 | sinks = {} |
| 154 | for i, (prefix, fuser) in model.attention_fusers.items(): |
| 155 | if (sink := fuser.sinks(model.get_submodule(prefix))) is not None: |
| 156 | sinks[i] = sink.float().cpu() |
| 157 | return sinks |
| 158 | |
| 159 | |
| 160 | def test_sinks(hf_runner: type[HfRunner], vllm_runner: type[VllmRunner]) -> None: |
| 161 | """Learnable attention sinks must reach the attention layers. |
| 162 | |
| 163 | Only the attention impl can apply them, so if they are not passed to |
| 164 | `Attention` they are dropped and every softmax is subtly wrong. |
| 165 | """ |
| 166 | from vllm.platforms import current_platform |
| 167 | |
| 168 | if not current_platform.has_device_capability(90): |
| 169 | pytest.skip("Attention sinks need FlashAttention 3 or TRTLLM attention") |
| 170 | |
| 171 | model = "tiny-random/gpt-oss-bf16" |
| 172 | with vllm_runner( |
| 173 | model, model_impl="transformers", max_model_len=1024, enforce_eager=True |
| 174 | ) as model_test: |
| 175 | assert model_test.llm.llm_engine.model_config.using_transformers_backend() |
| 176 | sinks = model_test.apply_model(get_sinks)[0] |
| 177 | |
| 178 | with hf_runner(model, dtype="bfloat16") as model_ref: |
| 179 | layers = model_ref.model.model.get_decoder().layers |
| 180 | expected = {i: layer.self_attn.sinks.float() for i, layer in enumerate(layers)} |
| 181 | |
| 182 | assert sinks.keys() == expected.keys() |
| 183 | for i, sink in sinks.items(): |
| 184 | torch.testing.assert_close(sink, expected[i].cpu()) |
| 185 | |
| 186 | |
| 187 | def test_mla(vllm_runner: type[VllmRunner], example_prompts: list[str]) -> None: |
| 188 | model = "hmellor/tiny-random-DeepseekV2ForCausalLM" |
| 189 | args = (example_prompts, 32, 5) |
| 190 | kwargs: dict[str, Any] = {"max_model_len": 2048, "enforce_eager": True} |
| 191 | |
| 192 | with vllm_runner( |
| 193 | model, model_impl="transformers", trust_remote_code=False, **kwargs |
| 194 | ) as model_test: |
| 195 | model_config = model_test.llm.llm_engine.model_config |
| 196 | assert model_config.using_transformers_backend() |
| 197 | num_layers = model_config.hf_config.get_text_config().num_hidden_layers |
| 198 | from vllm.model_executor.layers.attention import MLAAttention |
| 199 | |
| 200 | counts = model_test.apply_model(partial(count_modules, cls=MLAAttention)) |
| 201 | assert counts == [num_layers] |
| 202 | outputs_test = model_test.generate_greedy_logprobs(*args) |
| 203 | |
| 204 | with vllm_runner(model, model_impl="auto") as model_ref: |
| 205 | outputs_ref = model_ref.generate_greedy_logprobs(*args) |
| 206 | |
| 207 | check_logprobs_close( |
| 208 | outputs_0_lst=outputs_ref, |
| 209 | outputs_1_lst=outputs_test, |
| 210 | name_0="native", |
| 211 | name_1="transformers", |
| 212 | ) |
| 213 | |
| 214 | |
| 215 | @multi_gpu_test(num_gpus=2) |
| 216 | def test_distributed( |
| 217 | hf_runner: type[HfRunner], |
| 218 | vllm_runner: type[VllmRunner], |
| 219 | example_prompts, |
| 220 | ): |
| 221 | kwargs = {"model_impl": "transformers", "tensor_parallel_size": 2} |
| 222 | check_implementation( |
| 223 | hf_runner, |
| 224 | vllm_runner, |
| 225 | example_prompts, |
| 226 | "meta-llama/Llama-3.2-1B-Instruct", |
| 227 | kwargs_test=kwargs, |
| 228 | ) |
| 229 | |
| 230 | |
| 231 | @pytest.mark.parametrize( |
| 232 | "model, quantization_kwargs", |
| 233 | [ |
| 234 | ("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {}), |
| 235 | ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", {}), |
| 236 | ], |
| 237 | ) |
| 238 | @pytest.mark.parametrize("max_tokens", [32]) |
| 239 | @pytest.mark.parametrize("num_logprobs", [5]) |
| 240 | def test_quantization( |
| 241 | vllm_runner: type[VllmRunner], |
| 242 | example_prompts: list[str], |
| 243 | model: str, |
| 244 | quantization_kwargs: dict[str, str], |
| 245 | max_tokens: int, |
| 246 | num_logprobs: int, |
| 247 | ) -> None: |
| 248 | with vllm_runner( |
| 249 | model, |
| 250 | model_impl="auto", |
| 251 | enforce_eager=True, |
| 252 | **quantization_kwargs, |
| 253 | ) as vllm_model: |
| 254 | vllm_outputs = vllm_model.generate_greedy_logprobs( |
| 255 | example_prompts, max_tokens=max_tokens, num_logprobs=num_logprobs |
| 256 | ) |
| 257 | |
| 258 | with vllm_runner( |
| 259 | model, |
| 260 | model_impl="transformers", |
| 261 | enforce_eager=True, |
| 262 | **quantization_kwargs, |
| 263 | ) as vllm_model: |
| 264 | model_config = vllm_model.llm.llm_engine.model_config |
| 265 | assert model_config.using_transformers_backend() |
| 266 | |
| 267 | transformers_outputs = vllm_model.generate_greedy_logprobs( |
| 268 | example_prompts, max_tokens=max_tokens, num_logprobs=num_logprobs |
| 269 | ) |
| 270 | |
| 271 | check_logprobs_close( |
| 272 | outputs_0_lst=transformers_outputs, |
| 273 | outputs_1_lst=vllm_outputs, |
| 274 | name_0="transformers", |
| 275 | name_1="vllm", |
| 276 | ) |
| 277 | |
| 278 | |
| 279 | @pytest.mark.parametrize( |
| 280 | "model", |
| 281 | [ |
| 282 | # Layers live in `layers` |
| 283 | "Qwen/Qwen3-Embedding-0.6B", |
| 284 | # Layers live in `model.layers` |
| 285 | "meta-llama/Llama-3.2-1B-Instruct", |
| 286 | ], |
| 287 | ) |
| 288 | def test_embed_loading(vllm_runner, model): |
| 289 | with vllm_runner( |
| 290 | model, |
| 291 | max_model_len=1024, |
| 292 | enforce_eager=True, |
| 293 | runner="pooling", |
| 294 | model_impl="transformers", |
| 295 | ) as model_test: |
| 296 | model_config = model_test.llm.llm_engine.model_config |
| 297 | assert model_config.using_transformers_backend() |
| 298 | |
| 299 | |
| 300 | @pytest.mark.parametrize( |
| 301 | "arch", ["TransformersEmbeddingModel", "TransformersForSequenceClassification"] |
| 302 | ) |
| 303 | @pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) |
| 304 | def test_pooling( |
| 305 | hf_runner, |
| 306 | vllm_runner, |
| 307 | example_prompts, |
| 308 | arch, |
| 309 | monkeypatch, |
| 310 | use_v2_model_runner, |
| 311 | ): |
| 312 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) |
| 313 | model = get_model(arch) |
| 314 | |
| 315 | vllm_kwargs = dict(max_model_len=None, model_impl="transformers") |
| 316 | |
| 317 | hf_kwargs = dict() |
| 318 | if arch == "TransformersEmbeddingModel": |
| 319 | hf_kwargs["is_sentence_transformer"] = True |
| 320 | elif arch == "TransformersForSequenceClassification": |
| 321 | from transformers import AutoModelForSequenceClassification |
| 322 | |
| 323 | hf_kwargs["auto_cls"] = AutoModelForSequenceClassification |
| 324 | |
| 325 | # The example_prompts has ending "\n", for example: |
| 326 | # "Write a short story about a robot that dreams for the first time.\n" |
| 327 | # sentence_transformers will strip the input texts, see: |
| 328 | # https://github.com/UKPLab/sentence-transformers/blob/v3.1.1/sentence_transformers/models/Transformer.py#L159 |
| 329 | # This makes the input_ids different between hf_model and vllm_model. |
| 330 | # So we need to strip the input texts to avoid test failing. |
| 331 | example_prompts = [str(s).strip() for s in example_prompts] |
| 332 | |
| 333 | with ( |
| 334 | vllm_runner(model, **vllm_kwargs) as vllm_model, |
| 335 | hf_runner(model, **hf_kwargs) as hf_model, |
| 336 | ): |
| 337 | model_config = vllm_model.llm.llm_engine.model_config |
| 338 | assert model_config.using_transformers_backend() |
| 339 | |
| 340 | if arch == "TransformersEmbeddingModel": |
| 341 | vllm_outputs = vllm_model.embed(example_prompts) |
| 342 | hf_outputs = hf_model.encode(example_prompts) |
| 343 | elif arch == "TransformersForSequenceClassification": |
| 344 | vllm_outputs = vllm_model.classify(example_prompts) |
| 345 | hf_outputs = hf_model.classify(example_prompts) |
| 346 | |
| 347 | check_embeddings_close( |
| 348 | embeddings_0_lst=hf_outputs, |
| 349 | embeddings_1_lst=vllm_outputs, |
| 350 | name_0="hf", |
| 351 | name_1="vllm", |
| 352 | ) |
| 353 | |
| 354 | |
| 355 | VOCAB_SIZE = 64 |
| 356 | PER_LAYER_VOCAB_SIZE = 32 |
| 357 | NUM_POSITIONS = 16 |
| 358 | HIDDEN_SIZE = 8 |
| 359 | EMBED_SCALE = 3.0 |
| 360 | |
| 361 | |
| 362 | class ScaledWordEmbedding(nn.Embedding): |
| 363 | """Mirrors Transformers' `*ScaledWordEmbedding` classes.""" |
| 364 | |
| 365 | def __init__( |
| 366 | self, num_embeddings, embedding_dim, padding_idx=None, embed_scale=1.0 |
| 367 | ): |
| 368 | super().__init__(num_embeddings, embedding_dim, padding_idx) |
| 369 | self.scalar_embed_scale = embed_scale |
| 370 | self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False) |
| 371 | |
| 372 | def forward(self, input_ids): |
| 373 | return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) |
| 374 | |
| 375 | |
| 376 | class ComposedWordEmbedding(nn.Module): |
| 377 | """Scales embeddings, but wraps `nn.Embedding` instead of inheriting from it.""" |
| 378 | |
| 379 | def __init__(self, num_embeddings, embedding_dim, embed_scale=1.0): |
| 380 | super().__init__() |
| 381 | self.embed = nn.Embedding(num_embeddings, embedding_dim) |
| 382 | self.embed_scale = embed_scale |
| 383 | |
| 384 | def forward(self, input_ids): |
| 385 | return self.embed(input_ids) * self.embed_scale |
| 386 | |
| 387 | |
| 388 | @pytest.fixture |
| 389 | def tp_init(): |
| 390 | """Single rank tensor parallel state, so vLLM layers can be constructed.""" |
| 391 | from vllm.distributed import ( |
| 392 | cleanup_dist_env_and_memory, |
| 393 | init_distributed_environment, |
| 394 | initialize_model_parallel, |
| 395 | ) |
| 396 | from vllm.platforms import current_platform |
| 397 | |
| 398 | from ...utils import ensure_current_vllm_config |
| 399 | |
| 400 | fd, temp_file = tempfile.mkstemp() |
| 401 | os.close(fd) |
| 402 | try: |
| 403 | with ensure_current_vllm_config(): |
| 404 | init_distributed_environment( |
| 405 | world_size=1, |
| 406 | rank=0, |
| 407 | distributed_init_method=f"file://{temp_file}", |
| 408 | local_rank=0, |
| 409 | backend=current_platform.dist_backend, |
| 410 | ) |
| 411 | initialize_model_parallel(1, 1) |
| 412 | yield |
| 413 | cleanup_dist_env_and_memory() |
| 414 | finally: |
| 415 | with contextlib.suppress(OSError): |
| 416 | os.unlink(temp_file) |
| 417 | |
| 418 | |
| 419 | @pytest.fixture |
| 420 | def vpe(tp_init): |
| 421 | """`VocabParallelEmbedding`, imported late so collection does not import vLLM.""" |
| 422 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 423 | VocabParallelEmbedding, |
| 424 | ) |
| 425 | |
| 426 | return VocabParallelEmbedding |
| 427 | |
| 428 | |
| 429 | def replace(embedding): |
| 430 | """Replace `embedding` and fill the new weights with recognisable values.""" |
| 431 | from vllm.model_executor.models.transformers.utils import replace_embedding_class |
| 432 | |
| 433 | new_embedding = replace_embedding_class(embedding) |
| 434 | for _, param in new_embedding.named_parameters(): |
| 435 | param.data = torch.arange(param.numel(), dtype=param.dtype).view(param.shape) |
| 436 | return new_embedding |
| 437 | |
| 438 | |
| 439 | def assert_scaled(vpe, module): |
| 440 | """`module`'s output is its own unscaled output times `EMBED_SCALE`.""" |
| 441 | input_ids = torch.arange(VOCAB_SIZE) |
| 442 | unscaled = vpe.forward(module, input_ids) |
| 443 | torch.testing.assert_close(module(input_ids), unscaled * EMBED_SCALE) |
| 444 | |
| 445 | |
| 446 | def test_replace_plain_embedding(vpe): |
| 447 | """A plain `nn.Embedding` is replaced outright, leaving no subclass behind.""" |
| 448 | assert type(replace(nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE))) is vpe |
| 449 | |
| 450 | |
| 451 | def test_replace_infers_shape_and_dtype(tp_init): |
| 452 | """Shape and dtype come from the replaced module, not from the config.""" |
| 453 | embedding = nn.Embedding(VOCAB_SIZE * 2, HIDDEN_SIZE + 1, dtype=torch.float16) |
| 454 | new_embedding = replace(embedding) |
| 455 | |
| 456 | assert new_embedding.num_embeddings == VOCAB_SIZE * 2 |
| 457 | assert new_embedding.org_vocab_size == VOCAB_SIZE * 2 |
| 458 | assert new_embedding.embedding_dim == HIDDEN_SIZE + 1 |
| 459 | assert new_embedding.weight.dtype == torch.float16 |
| 460 | |
| 461 | |
| 462 | def test_replace_inherited_embedding(vpe): |
| 463 | """Subclasses keep their extra state and their scaled `forward`.""" |
| 464 | new_embedding = replace( |
| 465 | ScaledWordEmbedding( |
| 466 | VOCAB_SIZE, HIDDEN_SIZE, padding_idx=0, embed_scale=EMBED_SCALE |
| 467 | ) |
| 468 | ) |
| 469 | |
| 470 | assert isinstance(new_embedding, vpe) |
| 471 | assert new_embedding.scalar_embed_scale == EMBED_SCALE |
| 472 | assert "embed_scale" in new_embedding._non_persistent_buffers_set |
| 473 | assert_scaled(vpe, new_embedding) |
| 474 | |
| 475 | |
| 476 | def test_replace_is_idempotent(tp_init): |
| 477 | """A tied table is reached twice by the recursion, and must not be rebased twice.""" |
| 478 | from vllm.model_executor.models.transformers.utils import replace_embedding_class |
| 479 | |
| 480 | new_embedding = replace(ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=2)) |
| 481 | again = replace_embedding_class(new_embedding) |
| 482 | |
| 483 | assert again is new_embedding |
| 484 | assert type(again) is type(new_embedding) |
| 485 | |
| 486 | |
| 487 | class FakeVocabModel(nn.Module): |
| 488 | """The embedding tables a model can hold, only one of which is a vocab table. |
| 489 | |
| 490 | `embed_tokens_per_layer` mirrors Gemma 3n, whose per-layer table is sized by |
| 491 | `vocab_size_per_layer_input` rather than `vocab_size`. |
| 492 | """ |
| 493 | |
| 494 | def __init__(self, per_layer: bool = False, composed: bool = False): |
| 495 | super().__init__() |
| 496 | embed_cls = ComposedWordEmbedding if composed else ScaledWordEmbedding |
| 497 | self.embed_tokens = embed_cls(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE) |
| 498 | if per_layer: |
| 499 | self.embed_tokens_per_layer = ScaledWordEmbedding( |
| 500 | PER_LAYER_VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE |
| 501 | ) |
| 502 | # Not a vocab table: indexed by position, and read as a whole by models |
| 503 | # which interpolate it (e.g. `SiglipVisionEmbeddings`) |
| 504 | self.position_embedding = nn.Embedding(NUM_POSITIONS, HIDDEN_SIZE) |
| 505 | |
| 506 | def get_input_embeddings(self): |
| 507 | return self.embed_tokens |
| 508 | |
| 509 | |
| 510 | def replace_vocab_embeddings(model, **config_kwargs): |
| 511 | """Replace `model`'s vocab tables as `recursive_replace` would, and return them.""" |
| 512 | from vllm.model_executor.models.transformers.utils import attrsetter |
| 513 | |
| 514 | stub = nn.Module() |
| 515 | stub.model = model |
| 516 | stub.config = PreTrainedConfig( |
| 517 | vocab_size=VOCAB_SIZE, num_positions=NUM_POSITIONS, **config_kwargs |
| 518 | ) |
| 519 | embeddings = Base._vocab_embeddings(stub) |
| 520 | |
| 521 | replaced = [] |
| 522 | for name, module in list(model.named_modules()): |
| 523 | if module in embeddings: |
| 524 | replaced.append(replace(module)) |
| 525 | attrsetter(name)(model, replaced[-1]) |
| 526 | return replaced |
| 527 | |
| 528 | |
| 529 | def test_only_vocab_tables_are_replaced(vpe): |
| 530 | """Position embeddings share the `nn.Embedding` type, but not the treatment.""" |
| 531 | model = FakeVocabModel() |
| 532 | |
| 533 | assert replace_vocab_embeddings(model) == [model.embed_tokens] |
| 534 | assert isinstance(model.embed_tokens, vpe) |
| 535 | assert type(model.position_embedding) is nn.Embedding |
| 536 | |
| 537 | |
| 538 | def test_per_layer_vocab_table_is_replaced(vpe): |
| 539 | """A second vocab table is found by its size, without a bespoke accessor.""" |
| 540 | model = FakeVocabModel(per_layer=True) |
| 541 | replaced = replace_vocab_embeddings( |
| 542 | model, vocab_size_per_layer_input=PER_LAYER_VOCAB_SIZE |
| 543 | ) |
| 544 | |
| 545 | assert replaced == [model.embed_tokens, model.embed_tokens_per_layer] |
| 546 | assert all(isinstance(module, vpe) for module in replaced) |
| 547 | |
| 548 | |
| 549 | def test_composed_input_embeddings_are_replaced(vpe): |
| 550 | """A wrapper is left alone; only the `nn.Embedding` it composes is replaced. |
| 551 | |
| 552 | `CausalMixin` ties `lm_head` to whichever `VocabParallelEmbedding` it finds under |
| 553 | `get_input_embeddings()`, reading a `.weight` the wrapper does not have. |
| 554 | """ |
| 555 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 556 | |
| 557 | model = FakeVocabModel(composed=True) |
| 558 | wrapper = model.embed_tokens |
| 559 | replaced = replace_vocab_embeddings(model) |
| 560 | |
| 561 | assert model.embed_tokens is wrapper |
| 562 | assert replaced == [m for m in wrapper.modules() if isinstance(m, vpe)] |
| 563 | |
| 564 | lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE) |
| 565 | assert lm_head.tie_weights(replaced[0]).weight is replaced[0].weight |
| 566 | |
| 567 | |
| 568 | def test_missing_input_embeddings_are_skipped(tp_init): |
| 569 | """Pipeline ranks without the embeddings have a `PPMissingLayer` in their place.""" |
| 570 | from vllm.model_executor.models.utils import PPMissingLayer |
| 571 | |
| 572 | model = FakeVocabModel() |
| 573 | model.embed_tokens = PPMissingLayer() |
| 574 | |
| 575 | assert replace_vocab_embeddings(model) == [] |
| 576 | |
| 577 | |
| 578 | MULTIMODAL_MODEL = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" |
| 579 | |
| 580 | |
| 581 | class MarkingStub(SupportsMultiModal, nn.Module): |
| 582 | """Just enough of the backend to exercise component marking.""" |
| 583 | |
| 584 | _mark_model_components = MultiModalMixin._mark_model_components |
| 585 | _find_encoder_classes = MultiModalMixin._find_encoder_classes |
| 586 | _from_config_kwargs = Base._from_config_kwargs |
| 587 | _pre_trained_model_classes = Base._pre_trained_model_classes |
| 588 | |
| 589 | |
| 590 | def build_marked_model(image_limit: int, skip_tokenizer_init: bool = False): |
| 591 | """Build the HF model inside the marking context and return it.""" |
| 592 | model_config = ModelConfig( |
| 593 | model=MULTIMODAL_MODEL, |
| 594 | model_impl="transformers", |
| 595 | limit_mm_per_prompt={"image": image_limit}, |
| 596 | ) |
| 597 | # Set after construction: building the config itself needs the tokenizer |
| 598 | model_config.skip_tokenizer_init = skip_tokenizer_init |
| 599 | stub = MarkingStub() |
| 600 | stub.config = AutoConfig.from_pretrained(MULTIMODAL_MODEL) |
| 601 | stub.model_config = model_config |
| 602 | |
| 603 | vllm_config = VllmConfig(model_config=model_config) |
| 604 | with stub._mark_model_components(vllm_config), torch.device("meta"): |
| 605 | stub.model = AutoModel.from_config(**stub._from_config_kwargs) |
| 606 | return stub.model |
| 607 | |
| 608 | |
| 609 | @pytest.mark.parametrize(("image_limit", "skipped"), [(0, True), (4, False)]) |
| 610 | def test_tower_weights_skipped_when_modality_disabled(image_limit, skipped): |
| 611 | """`--limit-mm-per-prompt image=0` should drop the vision tower's weights.""" |
| 612 | vision_tower = build_marked_model(image_limit).vision_tower |
| 613 | assert isinstance(vision_tower, StageMissingLayer) is skipped |
| 614 | |
| 615 | |
| 616 | def test_marking_skipped_without_tokenizer(): |
| 617 | """Marking needs the HF processor, which needs a tokenizer, so it is skipped. |
| 618 | |
| 619 | The tower is built as normal rather than the model failing to load. |
| 620 | """ |
| 621 | vision_tower = build_marked_model(0, skip_tokenizer_init=True).vision_tower |
| 622 | assert not isinstance(vision_tower, StageMissingLayer) |
| 623 | |
| 624 | |
| 625 | NUM_LAYERS = 4 |
| 626 | |
| 627 | |
| 628 | def build_model(model_type: str, **overrides) -> nn.Module: |
| 629 | """A tiny HF model of `NUM_LAYERS` layers, built on the meta device.""" |
| 630 | try: |
| 631 | config = AutoConfig.for_model( |
| 632 | model_type, |
| 633 | num_hidden_layers=NUM_LAYERS, |
| 634 | hidden_size=64, |
| 635 | intermediate_size=128, |
| 636 | num_attention_heads=4, |
| 637 | num_key_value_heads=2, |
| 638 | head_dim=16, |
| 639 | vocab_size=VOCAB_SIZE, |
| 640 | **overrides, |
| 641 | ) |
| 642 | except ValueError: |
| 643 | pytest.skip(f"The installed transformers has no {model_type!r} model") |
| 644 | with torch.device("meta"): |
| 645 | return AutoModel.from_config(config) |
| 646 | |
| 647 | |
| 648 | ATTENTION_MODEL_TYPES = [ |
| 649 | # The default `head_size**-0.5` happens to be right here |
| 650 | "llama", |
| 651 | # Scales the query by a learnable per-dim weight, so it declares 1.0 |
| 652 | "gemma4_text", |
| 653 | # MLA, and declares yarn's `mscale` |
| 654 | "deepseek_v3", |
| 655 | ] |
| 656 | |
| 657 | |
| 658 | @pytest.mark.parametrize("model_type", ATTENTION_MODEL_TYPES) |
| 659 | def test_attention_dispatch_is_matched(model_type: str): |
| 660 | """Exactly the decoder layers' attention modules match an `AttentionFuser`. |
| 661 | |
| 662 | That match is what `recursive_replace` records, and so what |
| 663 | `_create_attention_instances` attaches vLLM's attention layer to. |
| 664 | """ |
| 665 | model = build_model(model_type) |
| 666 | matched = { |
| 667 | name |
| 668 | for name, module in model.named_modules() |
| 669 | if AttentionFuser.match(None, module) is not None |
| 670 | } |
| 671 | assert matched == {f"layers.{i}.self_attn" for i in range(NUM_LAYERS)} |
| 672 | |
| 673 | |
| 674 | def test_attention_layer_index_is_the_modules_own(): |
| 675 | """The layer served comes from the module, not from its position in the stack. |
| 676 | |
| 677 | LongCat Flash gives each decoder layer two attention sublayers numbered |
| 678 | `2i` and `2i + 1`, so `num_hidden_layers` is twice the length of the stack |
| 679 | and the enclosing layer's position is not the index the KV cache is keyed by. |
| 680 | A vision tower is excluded separately, by `validate`: only the text config is |
| 681 | patched to dispatch to vLLM, so a tower's attention is left to Transformers. |
| 682 | """ |
| 683 | model = build_model("longcat_flash", num_layers=2) |
| 684 | text_config = model.config.get_text_config() |
| 685 | assert len(model.layers) * 2 == text_config.num_hidden_layers |
| 686 | |
| 687 | attentions = [ |
| 688 | model.get_submodule(f"layers.{i}.self_attn.{j}") for i in (0, 1) for j in (0, 1) |
| 689 | ] |
| 690 | fusers = [AttentionFuser.match(None, attn) for attn in attentions] |
| 691 | assert all(fuser is not None for fuser in fusers) |
| 692 | assert [f.layer_index(a) for f, a in zip(fusers, attentions)] == [0, 1, 2, 3] |
| 693 | |
| 694 | # Not dispatching to vLLM yet, as a vision tower never would be. The vLLM |
| 695 | # config is not consulted, only the one the module was built with. |
| 696 | assert not any(f.validate(a, None) for f, a in zip(fusers, attentions)) |
| 697 | text_config._attn_implementation = VLLM_ATTN_IMPL |
| 698 | assert all(f.validate(a, None) for f, a in zip(fusers, attentions)) |
| 699 | |
| 700 | |
| 701 | def test_attention_dispatch_is_required(monkeypatch: pytest.MonkeyPatch): |
| 702 | """Models selected by the permissive registry fail clearly without dispatch.""" |
| 703 | monkeypatch.setattr( |
| 704 | "vllm.model_executor.models.transformers.base.get_pp_indices", |
| 705 | lambda *_: (0, 1), |
| 706 | ) |
| 707 | model = SimpleNamespace( |
| 708 | text_config=SimpleNamespace(num_hidden_layers=1), |
| 709 | pp_group=SimpleNamespace(rank_in_group=0, world_size=1), |
| 710 | attention_fusers={}, |
| 711 | _get_attn_cls=lambda: None, |
| 712 | ) |
| 713 | |
| 714 | with pytest.raises( |
| 715 | ValueError, |
| 716 | match="Layer 0 does not dispatch through the Transformers attention interface", |
| 717 | ): |
| 718 | Base._create_attention_instances(model) |
| 719 | |
| 720 | |
| 721 | @pytest.mark.parametrize("model_type", ATTENTION_MODEL_TYPES) |
| 722 | def test_attention_scale_is_the_declared_one(model_type: str): |
| 723 | """The scale comes off the module, not from the `head_size**-0.5` default. |
| 724 | |
| 725 | Two of the three model types declare something the default would silently |
| 726 | replace, so they are what makes this more than a tautology. |
| 727 | """ |
| 728 | attention = build_model(model_type).layers[0].self_attn |
| 729 | fuser = AttentionFuser.match(None, attention) |
| 730 | assert fuser is not None |
| 731 | assert fuser.scale(attention) == attention.scaling |
| 732 | |
| 733 | |
| 734 | def test_attention_scale_is_the_argument_not_the_attribute(): |
| 735 | """OPT applies the scale to the query itself, then declares a literal 1.0. |
| 736 | |
| 737 | Its `self.scaling` is the `head_size**-0.5` it has already applied, so |
| 738 | reading the attribute rather than the argument the module hands the |
| 739 | interface would scale twice. |
| 740 | """ |
| 741 | model = build_model("opt", ffn_dim=128, word_embed_proj_dim=64) |
| 742 | attention = model.get_submodule("decoder.layers.0.self_attn") |
| 743 | fuser = AttentionFuser.match(None, attention) |
| 744 | assert fuser is not None |
| 745 | assert fuser.scale(attention) == 1.0 != attention.scaling |
| 746 | |
| 747 | |
| 748 | def test_attention_scale_rejects_unresolvable_expression(): |
| 749 | """Unresolvable declared scales must not silently use the Llama default.""" |
| 750 | |
| 751 | class Attention(nn.Module): |
| 752 | head_dim = 16 |
| 753 | |
| 754 | fuser = AttentionFuser( |
| 755 | source_cls="Attention", |
| 756 | scale_expr=ast.parse("self.head_dim**-0.5 * self.factor", mode="eval").body, |
| 757 | ) |
| 758 | |
| 759 | with pytest.raises(ValueError, match="Cannot resolve attention scaling expression"): |
| 760 | fuser.scale(Attention()) |
| 761 | |
| 762 | |
| 763 | class _ImageOnlyMRoPEModel: |
| 764 | """`get_rope_index` without `video_grid_thw` or `**kwargs` (e.g. HunYuanVL).""" |
| 765 | |
| 766 | def get_rope_index(self, input_ids, image_grid_thw): |
| 767 | seq_len = input_ids.shape[-1] |
| 768 | positions = torch.arange(seq_len).view(1, 1, -1).expand(4, 1, -1) |
| 769 | return positions, torch.tensor([0]) |
| 770 | |
| 771 | |
| 772 | def test_get_mrope_input_positions_omits_unsupported_grid_kwargs(): |
| 773 | """Optional grids the model can't accept must not be passed when empty.""" |
| 774 | mixin = SimpleNamespace(model=_ImageOnlyMRoPEModel()) |
| 775 | |
| 776 | positions, delta = MultiModalMixin.get_mrope_input_positions(mixin, [1, 2, 3], []) |
| 777 | |
| 778 | assert positions.shape == (4, 3) |
| 779 | assert delta == 0 |