| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the Transformers backend's hw-agnostic layer resolution. |
| 4 | |
| 5 | `layers._resolve` imports a layer symbol from |
| 6 | `vllm.model_executor.hw_agnostic.layers.<module>` when `VLLM_USE_HW_AGNOSTIC` |
| 7 | is set and the symbol exists, and otherwise falls back to |
| 8 | `vllm.model_executor.layers.<module>`. These tests pin that contract and the |
| 9 | logging that reports which source was used. |
| 10 | """ |
| 11 | |
| 12 | import logging |
| 13 | import sys |
| 14 | import types |
| 15 | |
| 16 | import pytest |
| 17 | import torch |
| 18 | |
| 19 | from vllm.model_executor.models.transformers import layers |
| 20 | |
| 21 | HW_MODULE = "vllm.model_executor.hw_agnostic.layers.layernorm" |
| 22 | |
| 23 | |
| 24 | @pytest.fixture |
| 25 | def fake_hw_layernorm(monkeypatch): |
| 26 | """Inject a hw-agnostic `layernorm` module exposing a sentinel `RMSNorm`. |
| 27 | |
| 28 | A `SimpleNamespace` stands in for the module: `importlib.import_module` |
| 29 | returns it from `sys.modules` and `getattr` resolves `RMSNorm`, while its |
| 30 | attributes are set at construction (no `ModuleType` attribute-set that mypy |
| 31 | rejects, no constant `setattr` that ruff rejects).""" |
| 32 | module = types.SimpleNamespace(RMSNorm=type("HwRMSNorm", (), {})) |
| 33 | monkeypatch.setitem(sys.modules, HW_MODULE, module) |
| 34 | return module |
| 35 | |
| 36 | |
| 37 | def test_falls_back_to_vllm_when_disabled(monkeypatch, fake_hw_layernorm): |
| 38 | """Disabled: the vLLM class is used even if a hw-agnostic one exists.""" |
| 39 | monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0") |
| 40 | from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm |
| 41 | |
| 42 | assert layers._resolve("layernorm", "RMSNorm") is VllmRMSNorm |
| 43 | |
| 44 | |
| 45 | def test_uses_hw_agnostic_when_enabled(monkeypatch, fake_hw_layernorm, caplog): |
| 46 | """Enabled and available: the hw-agnostic class is used and logged.""" |
| 47 | monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") |
| 48 | with caplog.at_level(logging.INFO): |
| 49 | resolved = layers._resolve("layernorm", "RMSNorm") |
| 50 | assert resolved is fake_hw_layernorm.RMSNorm |
| 51 | assert "Using hw-agnostic layer: RMSNorm" in caplog.text |
| 52 | |
| 53 | |
| 54 | def test_falls_back_when_symbol_missing(monkeypatch, caplog): |
| 55 | """Enabled but the symbol is not ported: fall back to vLLM and warn.""" |
| 56 | monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") |
| 57 | # A hw-agnostic module without the requested attribute triggers fallback. |
| 58 | empty = types.ModuleType(HW_MODULE) |
| 59 | monkeypatch.setitem(sys.modules, HW_MODULE, empty) |
| 60 | from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm |
| 61 | |
| 62 | with caplog.at_level(logging.WARNING): |
| 63 | resolved = layers._resolve("layernorm", "RMSNorm") |
| 64 | assert resolved is VllmRMSNorm |
| 65 | assert "falling back to default" in caplog.text |
| 66 | |
| 67 | |
| 68 | def test_act_and_mul_falls_back_for_unknown_activation( |
| 69 | monkeypatch, default_vllm_config |
| 70 | ): |
| 71 | """An activation with no hw-agnostic equivalent falls back to vLLM's. |
| 72 | |
| 73 | `default_vllm_config` supplies the config context the CustomOp needs. |
| 74 | """ |
| 75 | monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") |
| 76 | from vllm.model_executor.layers.activation import GeluAndMul |
| 77 | |
| 78 | assert isinstance(layers.get_act_and_mul_fn("gelu"), GeluAndMul) |
| 79 | |
| 80 | |
| 81 | @pytest.fixture(scope="module") |
| 82 | def tiny_llama_path(tmp_path_factory): |
| 83 | """A randomly-initialized microscopic Llama saved to disk (with an ungated |
| 84 | tokenizer) so vLLM can load it like any local checkpoint.""" |
| 85 | from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM |
| 86 | |
| 87 | tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") |
| 88 | config = LlamaConfig( |
| 89 | vocab_size=tokenizer.vocab_size, |
| 90 | hidden_size=128, |
| 91 | intermediate_size=256, |
| 92 | num_hidden_layers=2, |
| 93 | num_attention_heads=4, |
| 94 | rms_norm_eps=1e-6, |
| 95 | hidden_act="silu", |
| 96 | ) |
| 97 | torch.manual_seed(0) |
| 98 | model = LlamaForCausalLM(config) |
| 99 | |
| 100 | path = tmp_path_factory.mktemp("tiny_llama") |
| 101 | model.save_pretrained(path) |
| 102 | tokenizer.save_pretrained(path) |
| 103 | return str(path) |
| 104 | |
| 105 | |
| 106 | # Registered names of the layers the backend can |
| 107 | # currently route to hw-agnostic implementations. |
| 108 | _COVERED_LAYERS = ("rms_norm", "silu_and_mul") |
| 109 | |
| 110 | |
| 111 | def _layer_providers(model) -> dict[str, str]: |
| 112 | """Map each covered layer type present in the model to the provider its |
| 113 | implementation came from (``hw_agnostic`` or ``vllm``). |
| 114 | """ |
| 115 | |
| 116 | def provider_of(module) -> str | None: |
| 117 | for cls in type(module).__mro__: |
| 118 | if "hw_agnostic.layers" in cls.__module__: |
| 119 | return "hw_agnostic" |
| 120 | if ".model_executor.layers." in cls.__module__: |
| 121 | return "vllm" |
| 122 | return None |
| 123 | |
| 124 | providers: dict[str, str] = {} |
| 125 | for module in model.modules(): |
| 126 | name = getattr(module, "name", None) |
| 127 | if ( |
| 128 | name in _COVERED_LAYERS |
| 129 | and name not in providers |
| 130 | and (prov := provider_of(module)) is not None |
| 131 | ): |
| 132 | providers[name] = prov |
| 133 | return providers |
| 134 | |
| 135 | |
| 136 | def _serve(vllm_runner, model_path, prompts): |
| 137 | """Serve the model through the backend; return (layer_providers, logprobs).""" |
| 138 | with vllm_runner( |
| 139 | model_path, |
| 140 | model_impl="transformers", |
| 141 | max_model_len=64, |
| 142 | enforce_eager=True, |
| 143 | gpu_memory_utilization=0.3, |
| 144 | ) as runner: |
| 145 | assert runner.llm.llm_engine.model_config.using_transformers_backend() |
| 146 | providers = runner.apply_model(_layer_providers)[0] |
| 147 | outputs = runner.generate_greedy_logprobs( |
| 148 | prompts, max_tokens=32, num_logprobs=5 |
| 149 | ) |
| 150 | return providers, outputs |
| 151 | |
| 152 | |
| 153 | def test_hw_agnostic_matches_vllm_end_to_end(monkeypatch, vllm_runner, tiny_llama_path): |
| 154 | """Serving the tiny model with hw-agnostic layers matches the vLLM baseline.""" |
| 155 | # spawn: worker re-imports layers with the env set (see docstring). |
| 156 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 157 | # apply_model pickles the introspection function. |
| 158 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 159 | from ..utils import check_logprobs_close |
| 160 | |
| 161 | prompts = ["The capital of France is", "vLLM is"] |
| 162 | |
| 163 | monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0") |
| 164 | vllm_providers, vllm_outputs = _serve(vllm_runner, tiny_llama_path, prompts) |
| 165 | # Both replaceable layers present in a Llama block must be vLLM's here. |
| 166 | assert vllm_providers == {"rms_norm": "vllm", "silu_and_mul": "vllm"} |
| 167 | |
| 168 | monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") |
| 169 | hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_path, prompts) |
| 170 | assert hw_providers == {"rms_norm": "hw_agnostic", "silu_and_mul": "hw_agnostic"} |
| 171 | |
| 172 | check_logprobs_close( |
| 173 | outputs_0_lst=vllm_outputs, |
| 174 | outputs_1_lst=hw_outputs, |
| 175 | name_0="vllm", |
| 176 | name_1="hw_agnostic", |
| 177 | ) |