iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][Models] Register sleep-managed runtime buffers (#53507)

Browse Source 
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
This commit is contained in:
Ronald
2026-08-29 10:32:24 +08:00
committed by GitHub
co-authored by aoshen02
parent 7fd9cc036e
commit 16d6c376bc
6 changed files with 427 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/basic_correctness/test_mem.py
+38
View File
@@ -18,6 +18,22 @@ from ..utils import create_new_process_for_each_test, requires_fp8
DEVICE_TYPE = current_platform.device_type
def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None:
"""Wake discarded allocations with deterministic nonzero contents."""
original_create_and_map = cumem.create_and_map
def create_and_map_with_poison(handle) -> None:
original_create_and_map(handle)
_, size, ptr, _ = handle
cumem.libcudart.cudaMemset(ptr, byte_value, size)
cumem.create_and_map = create_and_map_with_poison
try:
allocator.wake_up()
finally:
cumem.create_and_map = original_create_and_map
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
def test_python_error():
"""
@@ -110,6 +126,28 @@ def test_discard_tags():
assert torch.allclose(weights, torch.ones_like(weights))
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
def test_level2_discards_ordinary_tensor_with_weights_tag():
"""Reproduce the level-2 variant for an ordinary tensor in weights."""
allocator = get_mem_allocator_instance()
with allocator.use_memory_pool("weights"):
fake_weight = torch.full((4096,), 0x44, dtype=torch.uint8, device=DEVICE_TYPE)
ordinary_tensor = torch.full(
(4096,), 0x55, dtype=torch.uint8, device=DEVICE_TYPE
)
pointers = (fake_weight.data_ptr(), ordinary_tensor.data_ptr())
allocator.sleep(offload_tags=())
_wake_up_with_poisoned_mappings(allocator)
torch.accelerator.synchronize()
assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers
assert torch.all(fake_weight == 0xA5)
assert torch.all(ordinary_tensor == 0xA5)
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
@pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU")
def test_cumem_with_cudagraph():
tests/model_executor/test_sleep_mode_tensor_ownership.py
+339
View File
@@ -0,0 +1,339 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from functools import partial
from pathlib import Path
import pytest
import torch
from transformers import Gemma3nTextConfig
from vllm import LLM, SamplingParams
from vllm.inputs import TokensPrompt
from vllm.platforms import current_platform
from ..models.utils import dummy_hf_overrides
from ..utils import create_new_process_for_each_test
@dataclass(frozen=True)
class SleepModelCase:
name: str
architecture: str
model: str | None
tensor_names: tuple[str, ...]
modalities: tuple[str, ...] = ()
revision: str | None = None
trust_remote_code: bool = False
tokenizer_mode: str | None = None
SLEEP_MODEL_CASES = (
SleepModelCase(
name="fireredasr2",
architecture="FireRedASR2ForConditionalGeneration",
model="allendou/FireRedASR2-LLM-vllm",
tensor_names=("pe",),
trust_remote_code=True,
),
SleepModelCase(
name="gemma3n",
architecture="Gemma3nForCausalLM",
model=None,
tensor_names=(
"router_input_scale",
"embed_scale",
"embed_scale_per_layer",
"per_layer_input_scale",
"per_layer_projection_scale",
),
),
SleepModelCase(
name="voxtral",
architecture="VoxtralForConditionalGeneration",
model="mistralai/Voxtral-Mini-3B-2507",
tensor_names=("mel_filters",),
revision="3060fe34b35ba5d44202ce9ff3c097642914f8f3",
tokenizer_mode="mistral",
),
SleepModelCase(
name="ernie45_vl",
architecture="Ernie4_5_VLMoeForConditionalGeneration",
model="baidu/ERNIE-4.5-VL-28B-A3B-PT",
tensor_names=("inv_freq", "_visual_token_ids_tensor_cache"),
revision="refs/pr/17",
trust_remote_code=True,
),
)
def _create_local_model_config(case: SleepModelCase, model_dir: Path) -> str:
if case.name == "gemma3n":
config = Gemma3nTextConfig(
architectures=[case.architecture],
vocab_size=128,
hidden_size=64,
intermediate_size=[128, 128, 128],
num_hidden_layers=3,
num_attention_heads=2,
num_key_value_heads=2,
head_dim=32,
max_position_embeddings=256,
sliding_window=64,
layer_types=[
"sliding_attention",
"full_attention",
"sliding_attention",
],
vocab_size_per_layer_input=128,
hidden_size_per_layer_input=16,
num_kv_shared_layers=1,
laurel_rank=16,
activation_sparsity_pattern=[0.0, 0.0, 0.0],
)
else:
raise AssertionError(f"No local config builder for {case.name}")
model_dir.mkdir()
config.save_pretrained(model_dir)
return str(model_dir)
def _sleep_test_hf_overrides(hf_config, *, model_arch: str):
hf_config = dummy_hf_overrides(hf_config, model_arch=model_arch)
if model_arch == "FireRedASR2ForConditionalGeneration":
audio_encoder_conf = getattr(hf_config, "audio_encoder_conf", {})
audio_encoder_conf.update(
{
"idim": 80,
"n_layers_enc": 1,
"n_head": 1,
"d_model": 32,
"kernel_size": 3,
# MM profiling uses the repository's real feature extractor
# and can produce roughly 750 encoder frames. Keep the real
# positional-encoding capacity instead of shrinking it with
# the number of test decoder tokens.
"pe_maxlen": 5000,
}
)
hf_config.audio_encoder_conf = audio_encoder_conf
if (
model_arch == "Ernie4_5_VLMoeForConditionalGeneration"
and getattr(hf_config, "im_patch_id", None) is None
):
# Some revisions expose this id as image_token_id, while the native
# model intentionally caches it under the older im_patch_id name.
hf_config.im_patch_id = getattr(hf_config, "image_token_id", 100295)
return hf_config
def _get_sleep_tensor_targets(
model_name: str,
) -> tuple[tuple[type[torch.nn.Module], tuple[str, ...]], ...]:
if model_name == "fireredasr2":
from vllm.model_executor.models.conformer_encoder import (
RelPositionalEncoding,
)
return ((RelPositionalEncoding, ("pe",)),)
if model_name == "gemma3n":
from vllm.model_executor.models.gemma3n import (
Gemma3nAltUp,
Gemma3nSelfDecoder,
)
return (
(Gemma3nAltUp, ("router_input_scale",)),
(
Gemma3nSelfDecoder,
(
"embed_scale",
"embed_scale_per_layer",
"per_layer_input_scale",
"per_layer_projection_scale",
),
),
)
if model_name == "voxtral":
from vllm.model_executor.models.voxtral import VoxtralEncoderModel
return ((VoxtralEncoderModel, ("mel_filters",)),)
if model_name == "ernie45_vl":
from vllm.model_executor.models.ernie45_vl import (
Ernie4_5_VisionRotaryEmbedding,
Ernie4_5_VLMoeForConditionalGeneration,
)
return (
(Ernie4_5_VisionRotaryEmbedding, ("inv_freq",)),
(
Ernie4_5_VLMoeForConditionalGeneration,
("_visual_token_ids_tensor_cache",),
),
)
raise AssertionError(f"No sleep tensor targets for {model_name}")
def _snapshot_sleep_tensors(
model: torch.nn.Module,
*,
model_name: str,
tensor_names: tuple[str, ...],
) -> dict[str, tuple[int, torch.Tensor]]:
targets = _get_sleep_tensor_targets(model_name)
snapshot: dict[str, tuple[int, torch.Tensor]] = {}
found_tensor_names = set[str]()
for module_path, module in model.named_modules():
for module_type, target_names in targets:
if not isinstance(module, module_type):
continue
for tensor_name in target_names:
tensor = getattr(module, tensor_name, None)
if tensor is None:
continue
assert isinstance(tensor, torch.Tensor)
qualified_name = (
f"{module_path}.{tensor_name}" if module_path else tensor_name
)
snapshot[qualified_name] = (
tensor.data_ptr(),
tensor.detach().cpu().clone(),
)
found_tensor_names.add(tensor_name)
assert found_tensor_names == set(tensor_names), (
f"{model_name}: expected tensors {tensor_names}, "
f"found {tuple(sorted(found_tensor_names))}"
)
return snapshot
def _save_dummy_weights(worker) -> None:
model = worker.model_runner.model
worker._sleep_test_dummy_weights = [
(name, parameter.detach().cpu().clone())
for name, parameter in model.named_parameters()
]
def _reload_dummy_weights(worker) -> None:
with torch.no_grad():
worker.reload_weights(
weights_iterator=iter(worker._sleep_test_dummy_weights),
is_checkpoint_format=False,
)
def _generation_signature(output) -> tuple[tuple[int, ...], tuple[float, ...]]:
completion = output[0].outputs[0]
token_ids = tuple(completion.token_ids)
assert completion.logprobs is not None
chosen_logprobs = tuple(
step[token_id].logprob
for token_id, step in zip(token_ids, completion.logprobs, strict=True)
)
return token_ids, chosen_logprobs
def _assert_tensor_snapshots_equal(
before: dict[str, tuple[int, torch.Tensor]],
after: dict[str, tuple[int, torch.Tensor]],
model_name: str,
) -> None:
assert after.keys() == before.keys(), model_name
for name, (before_ptr, before_value) in before.items():
after_ptr, after_value = after[name]
assert after_ptr == before_ptr, f"{model_name}: {name} address changed"
torch.testing.assert_close(
after_value,
before_value,
msg=lambda msg, name=name: (f"{model_name}: {name} was corrupted: {msg}"),
)
@create_new_process_for_each_test("spawn")
@pytest.mark.slow_test
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA sleep mode test")
@pytest.mark.parametrize("case", SLEEP_MODEL_CASES, ids=lambda case: case.name)
def test_static_model_tensors_survive_level2_restore(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
case: SleepModelCase,
) -> None:
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
if case.name == "fireredasr2":
# FireRedASR2 has a regular decoder and compute_logits implementation,
# but advertises only the transcription task. Enable its text path so
# this ownership test can exercise the common LLM.generate lifecycle
# without constructing an unrelated audio fixture.
from vllm.model_executor.models.fireredasr2 import (
FireRedASR2ForConditionalGeneration,
)
monkeypatch.setattr(
FireRedASR2ForConditionalGeneration,
"supports_transcription_only",
False,
)
model = case.model or _create_local_model_config(case, tmp_path / case.name)
llm_kwargs = {
"model": model,
"load_format": "dummy",
"enable_sleep_mode": True,
"enforce_eager": True,
"attention_backend": "TRITON_ATTN",
"max_model_len": 128,
"max_num_seqs": 1,
"gpu_memory_utilization": 0.5,
"seed": 0,
"enable_prefix_caching": False,
"disable_log_stats": True,
"trust_remote_code": case.trust_remote_code,
}
if case.revision is not None:
llm_kwargs["revision"] = case.revision
if case.name not in ("fireredasr2", "voxtral", "ernie45_vl"):
llm_kwargs["skip_tokenizer_init"] = True
if case.tokenizer_mode is not None:
llm_kwargs["tokenizer_mode"] = case.tokenizer_mode
if case.model is not None:
llm_kwargs["hf_overrides"] = partial(
_sleep_test_hf_overrides, model_arch=case.architecture
)
if case.modalities:
llm_kwargs["limit_mm_per_prompt"] = {
modality: 0 for modality in case.modalities
}
llm = LLM(**llm_kwargs)
prompt = TokensPrompt(prompt_token_ids=[2, 5, 9, 12, 7, 3])
sampling_params = SamplingParams(temperature=0.0, max_tokens=8, logprobs=1)
before_generation = _generation_signature(llm.generate(prompt, sampling_params))
snapshot_fn = partial(
_snapshot_sleep_tensors,
model_name=case.name,
tensor_names=case.tensor_names,
)
(before_tensors,) = llm.apply_model(snapshot_fn)
llm.collective_rpc(_save_dummy_weights)
llm.sleep(level=2)
llm.wake_up(tags=["weights"])
(after_tensors,) = llm.apply_model(snapshot_fn)
_assert_tensor_snapshots_equal(before_tensors, after_tensors, case.name)
llm.collective_rpc(_reload_dummy_weights)
llm.wake_up(tags=["kv_cache"])
after_generation = _generation_signature(llm.generate(prompt, sampling_params))
assert after_generation[0] == before_generation[0], case.name
assert after_generation[1] == pytest.approx(
before_generation[1], rel=1e-5, abs=1e-5
), case.name
vllm/model_executor/models/conformer_encoder.py
+3 -1
View File
@@ -68,7 +68,9 @@ class RelPositionalEncoding(nn.Module):
pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
pe_negative = pe_negative[1:].unsqueeze(0)
self.pe = torch.cat([pe_positive, pe_negative], dim=1)
self.register_buffer(
"pe", torch.cat([pe_positive, pe_negative], dim=1), persistent=False
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Tmax = 2 * max_len - 1
vllm/model_executor/models/ernie45_vl.py
+13 -4
View File
@@ -343,8 +343,12 @@ class Ernie4_5_VisionPatchEmbed(nn.Module):
class Ernie4_5_VisionRotaryEmbedding(nn.Module):
def __init__(self, dim: int, theta: float = 10000.0) -> None:
super().__init__()
self.inv_freq = 1.0 / theta ** (
torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim
self.register_buffer(
"inv_freq",
1.0
/ theta
** (torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim),
persistent=False,
)
def forward(self, seqlen: int) -> torch.Tensor:
@@ -1452,11 +1456,16 @@ class Ernie4_5_VLMoeForConditionalGeneration(
]
if token_id is not None
]
self._visual_token_ids_tensor_cache = torch.tensor(
visual_token_ids_tensor_cache = torch.tensor(
visual_token_ids, dtype=torch.long
)
else:
self._visual_token_ids_tensor_cache = None
visual_token_ids_tensor_cache = None
self.register_buffer(
"_visual_token_ids_tensor_cache",
visual_token_ids_tensor_cache,
persistent=False,
)
def compute_logits(
self,
vllm/model_executor/models/gemma3n.py
+29 -13
View File
@@ -114,8 +114,10 @@ class Gemma3nAltUp(nn.Module):
hidden_size=hidden_size,
eps=rms_norm_eps,
)
self.router_input_scale = torch.tensor(
hidden_size**-1.0, dtype=self.modality_router.weight.dtype
self.register_buffer(
"router_input_scale",
torch.tensor(hidden_size**-1.0, dtype=self.modality_router.weight.dtype),
persistent=False,
)
self.correct_output_scale = nn.Parameter(
torch.zeros(hidden_size, dtype=torch.float32)
@@ -597,9 +599,13 @@ class Gemma3nSelfDecoder(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.embed_tokens",
)
self.embed_scale = torch.tensor(
config.hidden_size**0.5,
dtype=self.embed_tokens.weight.dtype,
self.register_buffer(
"embed_scale",
torch.tensor(
config.hidden_size**0.5,
dtype=self.embed_tokens.weight.dtype,
),
persistent=False,
)
# Additional per-layer embeddings (PLE)
self.embed_tokens_per_layer = VocabParallelEmbedding(
@@ -608,9 +614,13 @@ class Gemma3nSelfDecoder(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.per_layer_embed_tokens",
)
self.embed_scale_per_layer = torch.tensor(
config.hidden_size_per_layer_input**0.5,
dtype=self.embed_tokens.weight.dtype,
self.register_buffer(
"embed_scale_per_layer",
torch.tensor(
config.hidden_size_per_layer_input**0.5,
dtype=self.embed_tokens.weight.dtype,
),
persistent=False,
)
self.per_layer_model_projection = ColumnParallelLinear(
config.hidden_size,
@@ -625,12 +635,18 @@ class Gemma3nSelfDecoder(nn.Module):
hidden_size=config.hidden_size_per_layer_input,
eps=config.rms_norm_eps,
)
self.per_layer_input_scale = torch.rsqrt(torch.tensor(2.0)).to(
self.embed_tokens.weight.dtype
self.register_buffer(
"per_layer_input_scale",
torch.rsqrt(torch.tensor(2.0)).to(self.embed_tokens.weight.dtype),
persistent=False,
)
self.per_layer_projection_scale = torch.tensor(
config.hidden_size**0.5,
dtype=self.embed_tokens.weight.dtype,
self.register_buffer(
"per_layer_projection_scale",
torch.tensor(
config.hidden_size**0.5,
dtype=self.embed_tokens.weight.dtype,
),
persistent=False,
)
self.altup_projections = nn.ModuleList(
[
vllm/model_executor/models/voxtral.py
+5 -1
View File
@@ -765,7 +765,11 @@ class VoxtralEncoderModel(nn.Module):
max_frequency=8000.0,
sampling_rate=self.config.sampling_rate,
)
self.mel_filters = torch.tensor(mel_filters, dtype=torch.float32)
self.register_buffer(
"mel_filters",
torch.tensor(mel_filters, dtype=torch.float32),
persistent=False,
)
def compute_whisper_melspec(
self,