| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from functools import partial |
| 6 | from pathlib import Path |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | from transformers import Gemma3nTextConfig |
| 11 | |
| 12 | from vllm import LLM, SamplingParams |
| 13 | from vllm.inputs import TokensPrompt |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | from ..models.utils import dummy_hf_overrides |
| 17 | from ..utils import create_new_process_for_each_test |
| 18 | |
| 19 | |
| 20 | @dataclass(frozen=True) |
| 21 | class SleepModelCase: |
| 22 | name: str |
| 23 | architecture: str |
| 24 | model: str | None |
| 25 | tensor_names: tuple[str, ...] |
| 26 | modalities: tuple[str, ...] = () |
| 27 | revision: str | None = None |
| 28 | trust_remote_code: bool = False |
| 29 | tokenizer_mode: str | None = None |
| 30 | |
| 31 | |
| 32 | SLEEP_MODEL_CASES = ( |
| 33 | SleepModelCase( |
| 34 | name="fireredasr2", |
| 35 | architecture="FireRedASR2ForConditionalGeneration", |
| 36 | model="allendou/FireRedASR2-LLM-vllm", |
| 37 | tensor_names=("pe",), |
| 38 | trust_remote_code=True, |
| 39 | ), |
| 40 | SleepModelCase( |
| 41 | name="gemma3n", |
| 42 | architecture="Gemma3nForCausalLM", |
| 43 | model=None, |
| 44 | tensor_names=( |
| 45 | "router_input_scale", |
| 46 | "embed_scale", |
| 47 | "embed_scale_per_layer", |
| 48 | "per_layer_input_scale", |
| 49 | "per_layer_projection_scale", |
| 50 | ), |
| 51 | ), |
| 52 | SleepModelCase( |
| 53 | name="voxtral", |
| 54 | architecture="VoxtralForConditionalGeneration", |
| 55 | model="mistralai/Voxtral-Mini-3B-2507", |
| 56 | tensor_names=("mel_filters",), |
| 57 | revision="3060fe34b35ba5d44202ce9ff3c097642914f8f3", |
| 58 | tokenizer_mode="mistral", |
| 59 | ), |
| 60 | SleepModelCase( |
| 61 | name="ernie45_vl", |
| 62 | architecture="Ernie4_5_VLMoeForConditionalGeneration", |
| 63 | model="baidu/ERNIE-4.5-VL-28B-A3B-PT", |
| 64 | tensor_names=("inv_freq", "_visual_token_ids_tensor_cache"), |
| 65 | revision="refs/pr/17", |
| 66 | trust_remote_code=True, |
| 67 | ), |
| 68 | ) |
| 69 | |
| 70 | |
| 71 | def _create_local_model_config(case: SleepModelCase, model_dir: Path) -> str: |
| 72 | if case.name == "gemma3n": |
| 73 | config = Gemma3nTextConfig( |
| 74 | architectures=[case.architecture], |
| 75 | vocab_size=128, |
| 76 | hidden_size=64, |
| 77 | intermediate_size=[128, 128, 128], |
| 78 | num_hidden_layers=3, |
| 79 | num_attention_heads=2, |
| 80 | num_key_value_heads=2, |
| 81 | head_dim=32, |
| 82 | max_position_embeddings=256, |
| 83 | sliding_window=64, |
| 84 | layer_types=[ |
| 85 | "sliding_attention", |
| 86 | "full_attention", |
| 87 | "sliding_attention", |
| 88 | ], |
| 89 | vocab_size_per_layer_input=128, |
| 90 | hidden_size_per_layer_input=16, |
| 91 | num_kv_shared_layers=1, |
| 92 | laurel_rank=16, |
| 93 | activation_sparsity_pattern=[0.0, 0.0, 0.0], |
| 94 | ) |
| 95 | else: |
| 96 | raise AssertionError(f"No local config builder for {case.name}") |
| 97 | |
| 98 | model_dir.mkdir() |
| 99 | config.save_pretrained(model_dir) |
| 100 | return str(model_dir) |
| 101 | |
| 102 | |
| 103 | def _sleep_test_hf_overrides(hf_config, *, model_arch: str): |
| 104 | hf_config = dummy_hf_overrides(hf_config, model_arch=model_arch) |
| 105 | if model_arch == "FireRedASR2ForConditionalGeneration": |
| 106 | audio_encoder_conf = getattr(hf_config, "audio_encoder_conf", {}) |
| 107 | audio_encoder_conf.update( |
| 108 | { |
| 109 | "idim": 80, |
| 110 | "n_layers_enc": 1, |
| 111 | "n_head": 1, |
| 112 | "d_model": 32, |
| 113 | "kernel_size": 3, |
| 114 | # MM profiling uses the repository's real feature extractor |
| 115 | # and can produce roughly 750 encoder frames. Keep the real |
| 116 | # positional-encoding capacity instead of shrinking it with |
| 117 | # the number of test decoder tokens. |
| 118 | "pe_maxlen": 5000, |
| 119 | } |
| 120 | ) |
| 121 | hf_config.audio_encoder_conf = audio_encoder_conf |
| 122 | if ( |
| 123 | model_arch == "Ernie4_5_VLMoeForConditionalGeneration" |
| 124 | and getattr(hf_config, "im_patch_id", None) is None |
| 125 | ): |
| 126 | # Some revisions expose this id as image_token_id, while the native |
| 127 | # model intentionally caches it under the older im_patch_id name. |
| 128 | hf_config.im_patch_id = getattr(hf_config, "image_token_id", 100295) |
| 129 | return hf_config |
| 130 | |
| 131 | |
| 132 | def _get_sleep_tensor_targets( |
| 133 | model_name: str, |
| 134 | ) -> tuple[tuple[type[torch.nn.Module], tuple[str, ...]], ...]: |
| 135 | if model_name == "fireredasr2": |
| 136 | from vllm.model_executor.models.conformer_encoder import ( |
| 137 | RelPositionalEncoding, |
| 138 | ) |
| 139 | |
| 140 | return ((RelPositionalEncoding, ("pe",)),) |
| 141 | if model_name == "gemma3n": |
| 142 | from vllm.model_executor.models.gemma3n import ( |
| 143 | Gemma3nAltUp, |
| 144 | Gemma3nSelfDecoder, |
| 145 | ) |
| 146 | |
| 147 | return ( |
| 148 | (Gemma3nAltUp, ("router_input_scale",)), |
| 149 | ( |
| 150 | Gemma3nSelfDecoder, |
| 151 | ( |
| 152 | "embed_scale", |
| 153 | "embed_scale_per_layer", |
| 154 | "per_layer_input_scale", |
| 155 | "per_layer_projection_scale", |
| 156 | ), |
| 157 | ), |
| 158 | ) |
| 159 | if model_name == "voxtral": |
| 160 | from vllm.model_executor.models.voxtral import VoxtralEncoderModel |
| 161 | |
| 162 | return ((VoxtralEncoderModel, ("mel_filters",)),) |
| 163 | if model_name == "ernie45_vl": |
| 164 | from vllm.model_executor.models.ernie45_vl import ( |
| 165 | Ernie4_5_VisionRotaryEmbedding, |
| 166 | Ernie4_5_VLMoeForConditionalGeneration, |
| 167 | ) |
| 168 | |
| 169 | return ( |
| 170 | (Ernie4_5_VisionRotaryEmbedding, ("inv_freq",)), |
| 171 | ( |
| 172 | Ernie4_5_VLMoeForConditionalGeneration, |
| 173 | ("_visual_token_ids_tensor_cache",), |
| 174 | ), |
| 175 | ) |
| 176 | raise AssertionError(f"No sleep tensor targets for {model_name}") |
| 177 | |
| 178 | |
| 179 | def _snapshot_sleep_tensors( |
| 180 | model: torch.nn.Module, |
| 181 | *, |
| 182 | model_name: str, |
| 183 | tensor_names: tuple[str, ...], |
| 184 | ) -> dict[str, tuple[int, torch.Tensor]]: |
| 185 | targets = _get_sleep_tensor_targets(model_name) |
| 186 | snapshot: dict[str, tuple[int, torch.Tensor]] = {} |
| 187 | found_tensor_names = set[str]() |
| 188 | |
| 189 | for module_path, module in model.named_modules(): |
| 190 | for module_type, target_names in targets: |
| 191 | if not isinstance(module, module_type): |
| 192 | continue |
| 193 | for tensor_name in target_names: |
| 194 | tensor = getattr(module, tensor_name, None) |
| 195 | if tensor is None: |
| 196 | continue |
| 197 | assert isinstance(tensor, torch.Tensor) |
| 198 | qualified_name = ( |
| 199 | f"{module_path}.{tensor_name}" if module_path else tensor_name |
| 200 | ) |
| 201 | snapshot[qualified_name] = ( |
| 202 | tensor.data_ptr(), |
| 203 | tensor.detach().cpu().clone(), |
| 204 | ) |
| 205 | found_tensor_names.add(tensor_name) |
| 206 | |
| 207 | assert found_tensor_names == set(tensor_names), ( |
| 208 | f"{model_name}: expected tensors {tensor_names}, " |
| 209 | f"found {tuple(sorted(found_tensor_names))}" |
| 210 | ) |
| 211 | return snapshot |
| 212 | |
| 213 | |
| 214 | def _save_dummy_weights(worker) -> None: |
| 215 | model = worker.model_runner.model |
| 216 | worker._sleep_test_dummy_weights = [ |
| 217 | (name, parameter.detach().cpu().clone()) |
| 218 | for name, parameter in model.named_parameters() |
| 219 | ] |
| 220 | |
| 221 | |
| 222 | def _reload_dummy_weights(worker) -> None: |
| 223 | with torch.no_grad(): |
| 224 | worker.reload_weights( |
| 225 | weights_iterator=iter(worker._sleep_test_dummy_weights), |
| 226 | is_checkpoint_format=False, |
| 227 | ) |
| 228 | |
| 229 | |
| 230 | def _generation_signature(output) -> tuple[tuple[int, ...], tuple[float, ...]]: |
| 231 | completion = output[0].outputs[0] |
| 232 | token_ids = tuple(completion.token_ids) |
| 233 | assert completion.logprobs is not None |
| 234 | chosen_logprobs = tuple( |
| 235 | step[token_id].logprob |
| 236 | for token_id, step in zip(token_ids, completion.logprobs, strict=True) |
| 237 | ) |
| 238 | return token_ids, chosen_logprobs |
| 239 | |
| 240 | |
| 241 | def _assert_tensor_snapshots_equal( |
| 242 | before: dict[str, tuple[int, torch.Tensor]], |
| 243 | after: dict[str, tuple[int, torch.Tensor]], |
| 244 | model_name: str, |
| 245 | ) -> None: |
| 246 | assert after.keys() == before.keys(), model_name |
| 247 | for name, (before_ptr, before_value) in before.items(): |
| 248 | after_ptr, after_value = after[name] |
| 249 | assert after_ptr == before_ptr, f"{model_name}: {name} address changed" |
| 250 | torch.testing.assert_close( |
| 251 | after_value, |
| 252 | before_value, |
| 253 | msg=lambda msg, name=name: (f"{model_name}: {name} was corrupted: {msg}"), |
| 254 | ) |
| 255 | |
| 256 | |
| 257 | @create_new_process_for_each_test("spawn") |
| 258 | @pytest.mark.slow_test |
| 259 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA sleep mode test") |
| 260 | @pytest.mark.parametrize("case", SLEEP_MODEL_CASES, ids=lambda case: case.name) |
| 261 | def test_static_model_tensors_survive_level2_restore( |
| 262 | monkeypatch: pytest.MonkeyPatch, |
| 263 | tmp_path: Path, |
| 264 | case: SleepModelCase, |
| 265 | ) -> None: |
| 266 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 267 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 268 | |
| 269 | if case.name == "fireredasr2": |
| 270 | # FireRedASR2 has a regular decoder and compute_logits implementation, |
| 271 | # but advertises only the transcription task. Enable its text path so |
| 272 | # this ownership test can exercise the common LLM.generate lifecycle |
| 273 | # without constructing an unrelated audio fixture. |
| 274 | from vllm.model_executor.models.fireredasr2 import ( |
| 275 | FireRedASR2ForConditionalGeneration, |
| 276 | ) |
| 277 | |
| 278 | monkeypatch.setattr( |
| 279 | FireRedASR2ForConditionalGeneration, |
| 280 | "supports_transcription_only", |
| 281 | False, |
| 282 | ) |
| 283 | |
| 284 | model = case.model or _create_local_model_config(case, tmp_path / case.name) |
| 285 | llm_kwargs = { |
| 286 | "model": model, |
| 287 | "load_format": "dummy", |
| 288 | "enable_sleep_mode": True, |
| 289 | "enforce_eager": True, |
| 290 | "attention_backend": "TRITON_ATTN", |
| 291 | "max_model_len": 128, |
| 292 | "max_num_seqs": 1, |
| 293 | "gpu_memory_utilization": 0.5, |
| 294 | "seed": 0, |
| 295 | "enable_prefix_caching": False, |
| 296 | "disable_log_stats": True, |
| 297 | "trust_remote_code": case.trust_remote_code, |
| 298 | } |
| 299 | if case.revision is not None: |
| 300 | llm_kwargs["revision"] = case.revision |
| 301 | if case.name not in ("fireredasr2", "voxtral", "ernie45_vl"): |
| 302 | llm_kwargs["skip_tokenizer_init"] = True |
| 303 | if case.tokenizer_mode is not None: |
| 304 | llm_kwargs["tokenizer_mode"] = case.tokenizer_mode |
| 305 | if case.model is not None: |
| 306 | llm_kwargs["hf_overrides"] = partial( |
| 307 | _sleep_test_hf_overrides, model_arch=case.architecture |
| 308 | ) |
| 309 | if case.modalities: |
| 310 | llm_kwargs["limit_mm_per_prompt"] = { |
| 311 | modality: 0 for modality in case.modalities |
| 312 | } |
| 313 | |
| 314 | llm = LLM(**llm_kwargs) |
| 315 | prompt = TokensPrompt(prompt_token_ids=[2, 5, 9, 12, 7, 3]) |
| 316 | sampling_params = SamplingParams(temperature=0.0, max_tokens=8, logprobs=1) |
| 317 | |
| 318 | before_generation = _generation_signature(llm.generate(prompt, sampling_params)) |
| 319 | snapshot_fn = partial( |
| 320 | _snapshot_sleep_tensors, |
| 321 | model_name=case.name, |
| 322 | tensor_names=case.tensor_names, |
| 323 | ) |
| 324 | (before_tensors,) = llm.apply_model(snapshot_fn) |
| 325 | llm.collective_rpc(_save_dummy_weights) |
| 326 | |
| 327 | llm.sleep(level=2) |
| 328 | llm.wake_up(tags=["weights"]) |
| 329 | (after_tensors,) = llm.apply_model(snapshot_fn) |
| 330 | _assert_tensor_snapshots_equal(before_tensors, after_tensors, case.name) |
| 331 | |
| 332 | llm.collective_rpc(_reload_dummy_weights) |
| 333 | llm.wake_up(tags=["kv_cache"]) |
| 334 | after_generation = _generation_signature(llm.generate(prompt, sampling_params)) |
| 335 | |
| 336 | assert after_generation[0] == before_generation[0], case.name |
| 337 | assert after_generation[1] == pytest.approx( |
| 338 | before_generation[1], rel=1e-5, abs=1e-5 |
| 339 | ), case.name |