| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import logging |
| 5 | from typing import Any |
| 6 | |
| 7 | import regex as re |
| 8 | import torch |
| 9 | from pydantic import TypeAdapter |
| 10 | |
| 11 | from vllm.config import ( |
| 12 | CompilationConfig, |
| 13 | CompilationMode, |
| 14 | ModelConfig, |
| 15 | VllmConfig, |
| 16 | set_current_vllm_config, |
| 17 | ) |
| 18 | from vllm.config.quantization import QuantizationConfigArgs, QuantSpec |
| 19 | from vllm.model_executor.layers.quantization import get_quantization_config |
| 20 | from vllm.model_executor.model_loader.default_loader import DefaultModelLoader |
| 21 | from vllm.platforms import current_platform |
| 22 | |
| 23 | |
| 24 | def quant_spec(**kwargs: Any) -> QuantSpec: |
| 25 | """Build a `QuantSpec` from user-facing names such as `"mxfp8"`.""" |
| 26 | return TypeAdapter(QuantSpec).validate_python(kwargs) |
| 27 | |
| 28 | |
| 29 | def quant_config_args(**kwargs: Any) -> QuantizationConfigArgs: |
| 30 | """Build `QuantizationConfigArgs` from user-facing shorthand values.""" |
| 31 | return TypeAdapter(QuantizationConfigArgs).validate_python(kwargs) |
| 32 | |
| 33 | |
| 34 | def _limit_num_hidden_layers( |
| 35 | model: torch.nn.Module, num_hidden_layers: int | None |
| 36 | ) -> None: |
| 37 | if num_hidden_layers is None: |
| 38 | return |
| 39 | |
| 40 | original_load_weights = model.load_weights |
| 41 | |
| 42 | def should_load_weight(name: str) -> bool: |
| 43 | for prefix in ("model.layers.", "layers."): |
| 44 | if name.startswith(prefix): |
| 45 | layer_idx = int(name.removeprefix(prefix).split(".", 1)[0]) |
| 46 | return layer_idx < num_hidden_layers |
| 47 | return True |
| 48 | |
| 49 | def load_weights(weights): |
| 50 | weights = ( |
| 51 | (name, weight) for name, weight in weights if should_load_weight(name) |
| 52 | ) |
| 53 | return original_load_weights(weights) |
| 54 | |
| 55 | model.load_weights = load_weights |
| 56 | |
| 57 | |
| 58 | def load_model_without_vllm_runner( |
| 59 | model_path: str, |
| 60 | *, |
| 61 | dtype: str | torch.dtype = "bfloat16", |
| 62 | quantization: str | None = None, |
| 63 | model_config_kwargs: dict[str, Any] | None = None, |
| 64 | vllm_config_kwargs: dict[str, Any] | None = None, |
| 65 | model_loader_cls: type = DefaultModelLoader, |
| 66 | ) -> tuple[torch.nn.Module, VllmConfig]: |
| 67 | """Instantiate a model, load weights, and process them for inference.""" |
| 68 | model_config = ModelConfig( |
| 69 | model=model_path, |
| 70 | dtype=dtype, |
| 71 | quantization=quantization, |
| 72 | **(model_config_kwargs or {}), |
| 73 | ) |
| 74 | vllm_config_args = dict(vllm_config_kwargs or {}) |
| 75 | vllm_config_args.setdefault( |
| 76 | "compilation_config", CompilationConfig(mode=CompilationMode.NONE) |
| 77 | ) |
| 78 | vllm_config = VllmConfig(model_config=model_config, **vllm_config_args) |
| 79 | hf_overrides = (model_config_kwargs or {}).get("hf_overrides") or {} |
| 80 | num_hidden_layers = hf_overrides.get("num_hidden_layers") |
| 81 | |
| 82 | with set_current_vllm_config(vllm_config): |
| 83 | model_loader = model_loader_cls(vllm_config.load_config) |
| 84 | if num_hidden_layers is not None: |
| 85 | original_load_weights = model_loader.load_weights |
| 86 | |
| 87 | def load_weights(model, model_config): |
| 88 | _limit_num_hidden_layers(model, num_hidden_layers) |
| 89 | original_load_weights(model, model_config) |
| 90 | |
| 91 | model_loader.load_weights = load_weights |
| 92 | model = model_loader.load_model(vllm_config, model_config) |
| 93 | |
| 94 | return model, vllm_config |
| 95 | |
| 96 | |
| 97 | def is_quant_method_supported(quant_method: str) -> bool: |
| 98 | # Currently, quantization tests only run GPUs |
| 99 | if current_platform.is_cpu(): |
| 100 | return False |
| 101 | try: |
| 102 | current_platform.verify_quantization(quant_method) |
| 103 | except ValueError: |
| 104 | return False |
| 105 | if current_platform.is_xpu(): |
| 106 | return True |
| 107 | capability = current_platform.get_device_capability() |
| 108 | assert capability is not None |
| 109 | |
| 110 | min_capability = get_quantization_config(quant_method).get_min_capability() |
| 111 | |
| 112 | return capability.to_int() >= min_capability |
| 113 | |
| 114 | |
| 115 | def _test_online_quant_peak_mem_impl( |
| 116 | quantization_arg_value, |
| 117 | vllm_runner, |
| 118 | caplog_mp_spawn, |
| 119 | monkeypatch, |
| 120 | ) -> None: |
| 121 | # Note: `allenai/OLMoE-1B-7B-0125-Instruct` was selected because: |
| 122 | # 1. it covers both Linear and MoE paths |
| 123 | # 2. it is already used by other tests in CI, so adding it here |
| 124 | # does not increase disk space for CI runners |
| 125 | # I really wanted to use `ibm-granite/granite-3.0-1b-a400m-base` |
| 126 | # which I think is the smallest MoE model in vLLM (2.5 GiB bf16, |
| 127 | # 1.3 GiB fp8), but could not as adding one more model makes CI |
| 128 | # run out of disk space. |
| 129 | model_name = "allenai/OLMoE-1B-7B-0125-Instruct" |
| 130 | |
| 131 | # Force spawn to ensure caplog_mp_spawn works consistently |
| 132 | # (it relies on VLLM_LOGGING_CONFIG_PATH which spawn reads but fork ignores) |
| 133 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 134 | |
| 135 | with ( |
| 136 | caplog_mp_spawn(logging.DEBUG) as log_holder, |
| 137 | vllm_runner( |
| 138 | model_name, |
| 139 | quantization=quantization_arg_value, |
| 140 | enforce_eager=True, |
| 141 | ) as llm, |
| 142 | ): |
| 143 | outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4) |
| 144 | print(outputs[0][1]) |
| 145 | |
| 146 | log_text = log_holder.text |
| 147 | |
| 148 | # Parse memory usage from captured logs |
| 149 | model_memory_gib = None |
| 150 | peak_memory_gib = None |
| 151 | for line in log_text.splitlines(): |
| 152 | if model_memory_gib is None: |
| 153 | match = re.search(r"Model loading took ([\d.]+) GiB memory", line) |
| 154 | if match: |
| 155 | model_memory_gib = float(match.group(1)) |
| 156 | if peak_memory_gib is None: |
| 157 | match = re.search( |
| 158 | r"Peak GPU memory after loading weights: ([\d.]+) GiB", line |
| 159 | ) |
| 160 | if match: |
| 161 | peak_memory_gib = float(match.group(1)) |
| 162 | |
| 163 | assert model_memory_gib is not None, "Could not find model loading memory log" |
| 164 | assert peak_memory_gib is not None, "Could not find peak memory log" |
| 165 | print(f"GPU memory used after loading weights: {model_memory_gib} GiB") |
| 166 | print(f"Peak GPU memory usage while loading weights: {peak_memory_gib} GiB") |
| 167 | |
| 168 | expected_model_memory_gib = 6.7 |
| 169 | |
| 170 | # for allenai/OLMoE-1B-7B-0125-Instruct the number we see today is 9.06 |
| 171 | # GiB on CUDA, which is 1.36x above model_memory_gib. A slightly higher |
| 172 | # number is expected as when we load and quantize weights in a streaming |
| 173 | # fashion we need to have individual weights in bf16 + fp8 alive at the |
| 174 | # same time. |
| 175 | expected_peak_memory_gib = expected_model_memory_gib * 1.4 |
| 176 | |
| 177 | assert model_memory_gib < expected_model_memory_gib, ( |
| 178 | f"{model_memory_gib=} higher than {expected_model_memory_gib}" |
| 179 | ) |
| 180 | assert peak_memory_gib < expected_peak_memory_gib, ( |
| 181 | f"{peak_memory_gib=} higher than {expected_peak_memory_gib}" |
| 182 | ) |