From f5ed3d9c03ec634fbfe670baaf234569bbdecf96 Mon Sep 17 00:00:00 2001 From: siyu Date: Wed, 30 Sep 2026 01:58:35 +0800 Subject: [PATCH] [Fast Start] Charge daemon-held weights against `gpu_memory_utilization` (#57298) Signed-off-by: liusy58 Signed-off-by: Isotr0py Signed-off-by: Xun Sun Co-authored-by: Isotr0py Co-authored-by: Xun Sun --- .../model_loader/test_weight_cache.py | 15 ++++ tests/v1/worker/test_utils.py | 38 +++++++++- .../model_loader/base_loader.py | 5 ++ .../model_loader/weight_cache/daemon.py | 19 ++++- .../model_loader/weight_cache/ipc_loader.py | 76 ++++++++++++++++--- vllm/v1/worker/gpu_worker.py | 9 ++- vllm/v1/worker/utils.py | 39 +++++++++- 7 files changed, 187 insertions(+), 14 deletions(-) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 7706522f6e..106b9690ac 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -354,3 +354,18 @@ def test_weight_cache_key_distinguishes_dp_ranks(): assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"] assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"] assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"] + + +def test_ipc_loader_copy_mode_reports_no_external_weight_memory(): + """Copy mode clones the weights into the engine, so nothing is external + (vllm_config is never touched).""" + from vllm.config import LoadConfig + from vllm.model_executor.model_loader.weight_cache.ipc_loader import ( + IpcModelLoader, + ) + + copy_mode = LoadConfig( + load_format="ipc_cache", model_loader_extra_config={"mode": "copy"} + ) + loader = IpcModelLoader(copy_mode) + assert loader.get_external_weight_memory(None) == 0 # type: ignore[arg-type] diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 9a345e926c..40ac8b56ee 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -13,7 +13,7 @@ import torch import vllm.v1.attention.backends.mla.index_group as index_group_module import vllm.v1.hisparse.runtime as hisparse_runtime_module -from vllm.config import CUDAGraphMode +from vllm.config import CacheConfig, CUDAGraphMode from vllm.config.mamba import MambaBackendEnum, MambaConfig from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import ( worker as hisparse_worker_module, @@ -25,12 +25,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( _SlotMappingStaging, ) from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 +from vllm.utils.mem_constants import GiB_bytes +from vllm.utils.mem_utils import MemorySnapshot from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror from vllm.v1.worker.utils import ( bind_kv_cache, bind_kv_cache_to_layers, copy_kv_cache_blocks_inplace, + request_memory, ) @@ -1390,3 +1393,36 @@ def test_bind_kv_cache_draft_model(default_vllm_config): assert runner_kv_caches[1] is kv_cache["draft_model.layers.0.attn"] assert runner_kv_caches[2] is kv_cache["model.layers.1.attn"] assert runner_kv_caches[3] is kv_cache["draft_model.layers.1.attn"] + + +def _memory_snapshot(total_gib: int, free_gib: int) -> MemorySnapshot: + return MemorySnapshot( + free_memory=free_gib * GiB_bytes, + total_memory=total_gib * GiB_bytes, + device="cpu", + auto_measure=False, + ) + + +def test_request_memory_charges_external_weights(): + """Externally held weights are charged against the utilization budget; + the engine is granted only the remainder.""" + cache_config = CacheConfig(gpu_memory_utilization=0.9) + # 70 GiB of a 100 GiB device is held externally before the worker starts. + snapshot = _memory_snapshot(total_gib=100, free_gib=30) + + with pytest.raises(ValueError, match="less than desired"): + request_memory(snapshot, cache_config) + + # 90 GiB budget - 70 GiB external = 20 GiB for the engine. + assert request_memory(snapshot, cache_config, 70 * GiB_bytes) == 20 * GiB_bytes + + # The external weights alone exceed the utilization budget. + with pytest.raises(ValueError, match="exceed the desired"): + request_memory(snapshot, cache_config, 95 * GiB_bytes) + + # Other tenants squeeze free memory below the engine's remainder. + with pytest.raises(ValueError, match="less than the engine's budget"): + request_memory( + _memory_snapshot(total_gib=100, free_gib=10), cache_config, 70 * GiB_bytes + ) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 6f35646458..7712efbfd1 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -28,6 +28,11 @@ class BaseModelLoader(ABC): def __init__(self, load_config: LoadConfig): self.load_config = load_config + def get_external_weight_memory(self, vllm_config: VllmConfig) -> int: + """Get weights memory from external process; + 0 when the weights are not external.""" + return 0 + @abstractmethod def download_model(self, model_config: ModelConfig) -> None: """Download a model so that it can be immediately loaded.""" diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index b56cbbea8f..2c6944f1eb 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -78,6 +78,7 @@ cached and keep loading from disk in the engine. import contextlib import fcntl +import gc import multiprocessing import os import socket @@ -116,6 +117,7 @@ from vllm.model_executor.model_loader.weight_cache.utils import ( is_draft_model_cacheable, ) from vllm.platforms import current_platform +from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.worker.workspace import init_workspace_manager @@ -249,10 +251,17 @@ class WeightCacheDaemon: with set_current_vllm_config(self.vllm_config): ensure_model_parallel_initialized(self.tp_size, self.pp_size) self.model = self.get_model() + # Loading and post-processing leave freed transients in the caching + # allocator; return them so engines sharing the GPU can use them. + gc.collect() + torch.accelerator.empty_cache() logger.info( - "Weight cache %s daemon rank %d loaded model", + "Weight cache %s daemon rank %d loaded model (%s GiB allocated, " + "%s GiB reserved)", self.role, self.global_rank, + format_gib(torch.accelerator.memory_allocated()), + format_gib(torch.accelerator.memory_reserved()), ) def get_model(self) -> torch.nn.Module: @@ -369,6 +378,14 @@ class WeightCacheDaemon: cmd = request.get("cmd") if cmd == "get_state": self._handle_get_state(conn, request) + elif cmd == "get_memory": + send_msg( + conn, + { + "status": "ok", + "memory_bytes": torch.accelerator.memory_allocated(), + }, + ) elif cmd == "release": self._handle_release(conn) else: diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index 185c11b83a..5fd755bf3f 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -6,7 +6,9 @@ daemon via CUDA IPC instead of loading from disk.""" import dataclasses import socket import time +from collections.abc import Callable from copy import copy +from typing import TypeVar import torch import torch.nn as nn @@ -23,6 +25,7 @@ from vllm.logger import init_logger from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.utils import ( + get_draft_load_config, initialize_model, process_weights_after_loading, ) @@ -40,6 +43,9 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( send_msg, verify_socket_owner, ) +from vllm.model_executor.model_loader.weight_cache.utils import ( + is_draft_model_cacheable, +) from vllm.model_executor.utils import weights_already_processed from vllm.tracing import instrument from vllm.utils.torch_utils import set_default_torch_dtype @@ -50,6 +56,8 @@ _CONNECT_TIMEOUT_S = 5.0 _STATE_TIMEOUT_S = 300.0 _STARTUP_RETRY_INTERVAL_S = 0.5 +_T = TypeVar("_T") + class IpcModelLoader(BaseModelLoader): """Loads a model by mapping the weight cache daemon's tensors via CUDA IPC. @@ -110,6 +118,52 @@ class IpcModelLoader(BaseModelLoader): f"{load_config.load_format}: {sorted(extra_config)}" ) + def get_external_weight_memory(self, vllm_config: VllmConfig) -> int: + # Copy mode clones the weights into this process; nothing external. + if self.mode != "zero_copy": + return 0 + total = self._daemon_memory() + if is_draft_model_cacheable(vllm_config.speculative_config): + # The draft group is queried with the same config the engine + # would load the draft with. + draft_config = get_draft_load_config(vllm_config) + if draft_config.load_format == "ipc_cache": + draft_loader = IpcModelLoader(draft_config) + # Only a zero-copy draft-group daemon holds external weights; + # an explicit draft_load_config without is_draft would resolve + # back to the target socket and double-count it. + if draft_loader.is_draft and draft_loader.mode == "zero_copy": + total += draft_loader._daemon_memory() + return total + + def _daemon_memory(self) -> int: + if not self.fallback: + # The loader waits for the daemon at load time, so the weights + # will be zero-copy mapped for sure; mirror that wait here since + # returning 0 would over-grant the memory budget. + return self._with_startup_wait(self._query_daemon_memory) + # An unreachable daemon means a disk load, i.e. nothing external. + try: + return self._query_daemon_memory() + except (WeightCacheUnavailableError, ConnectionError, OSError) as e: + logger.warning( + "Cannot query weight cache daemon memory (%s); " + "assuming the weights are not externally held", + e, + ) + return 0 + + def _query_daemon_memory(self) -> int: + with self._connect(self.connect_timeout_s) as conn: + send_msg(conn, {"cmd": "get_memory"}) + response = recv_msg(conn) + if response.get("status") != "ok": + raise WeightCacheUnavailableError( + "Weight cache daemon rejected the memory query: " + f"{response.get('message')}" + ) + return int(response.get("memory_bytes", 0)) + def download_model(self, model_config: ModelConfig) -> None: DefaultModelLoader(self._fallback_load_config()).download_model(model_config) @@ -308,23 +362,25 @@ class IpcModelLoader(BaseModelLoader): def _request_state_with_startup_wait( self, cache_config: WeightCacheKey ) -> WeightCacheState: + return self._with_startup_wait(lambda: self._request_state(cache_config)) + + def _with_startup_wait(self, op: Callable[[], _T]) -> _T: + """Retry op until the daemon answers or the state timeout elapses; + the daemon may still be loading the model when the engine starts.""" deadline = time.monotonic() + self.state_timeout_s - waiting_logged = False while True: try: - return self._request_state(cache_config) - except WeightCacheUnavailableError as e: + return op() + except (WeightCacheUnavailableError, ConnectionError, OSError) as e: if time.monotonic() >= deadline: raise WeightCacheUnavailableError( "Weight cache daemon did not become ready within " - f"{self.state_timeout_s:.1f}s" + f"{self.state_timeout_s:.1f}s: {e}" ) from e - if not waiting_logged: - logger.info( - "Waiting up to %.1fs for the weight cache daemon to start", - self.state_timeout_s, - ) - waiting_logged = True + logger.info_once( + "Waiting up to %.1fs for the weight cache daemon to start", + self.state_timeout_s, + ) time.sleep( max( 0.0, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 778d71bdbc..c6da90b077 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -57,6 +57,7 @@ from vllm.distributed.weight_transfer import ( ) from vllm.logger import init_logger from vllm.lora.request import LoRARequest +from vllm.model_executor.model_loader import get_model_loader from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory from vllm.platforms import current_platform @@ -487,7 +488,13 @@ class Worker(WorkerBase): # take current memory snapshot self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) - self.requested_memory = request_memory(init_snapshot, self.cache_config) + # Weights from external model loader process + external_weight_memory = get_model_loader( + self.load_config + ).get_external_weight_memory(self.vllm_config) + self.requested_memory = request_memory( + init_snapshot, self.cache_config, external_weight_memory + ) logger.debug("worker init memory snapshot: %r", self.init_snapshot) logger.debug( "worker requested memory: %sGiB", format_gib(self.requested_memory) diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index a6ef84bbd3..422ba8885f 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -536,7 +536,11 @@ def sanity_check_mm_encoder_outputs( ) -def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> int: +def request_memory( + init_snapshot: MemorySnapshot, + cache_config: CacheConfig, + external_weight_memory: int = 0, +) -> int: """Calculate the amount of memory required by vLLM, then validate that the current amount of free memory is sufficient for that. """ @@ -544,6 +548,39 @@ def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> init_snapshot.total_memory * cache_config.gpu_memory_utilization ) + if external_weight_memory > 0: + engine_memory = requested_memory - external_weight_memory + if engine_memory <= 0: + raise ValueError( + f"Externally held weights on device {init_snapshot.device_} " + f"({format_gib(external_weight_memory)}/" + f"{format_gib(init_snapshot.total_memory)} GiB) exceed the " + "desired GPU memory utilization " + f"({cache_config.gpu_memory_utilization}, " + f"{format_gib(requested_memory)} GiB). Increase GPU memory " + "utilization or reduce GPU memory used by other processes." + ) + if init_snapshot.free_memory < engine_memory: + raise ValueError( + f"Free memory on device {init_snapshot.device_} " + f"({format_gib(init_snapshot.free_memory)}/" + f"{format_gib(init_snapshot.total_memory)} GiB) on startup " + "is less than the engine's budget after excluding " + "external process's weights " + f"({cache_config.gpu_memory_utilization}, " + f"{format_gib(engine_memory)} GiB). Decrease GPU memory " + "utilization or reduce GPU memory used by other processes." + ) + logger.info_once( + "Weights are held outside this process: of the %s GiB " + "utilization budget, %s GiB is externally held and " + "%s GiB remains for the engine's own allocations.", + format_gib(requested_memory), + format_gib(external_weight_memory), + format_gib(engine_memory), + ) + return engine_memory + if init_snapshot.free_memory < requested_memory: raise ValueError( f"Free memory on device {init_snapshot.device_} "