| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """IPC model loader: maps post-quantized weights from a local weight cache |
| 4 | daemon via CUDA IPC instead of loading from disk.""" |
| 5 | |
| 6 | import dataclasses |
| 7 | import socket |
| 8 | import time |
| 9 | from collections.abc import Callable |
| 10 | from copy import copy |
| 11 | from typing import TypeVar |
| 12 | |
| 13 | import torch |
| 14 | import torch.nn as nn |
| 15 | |
| 16 | from vllm.config import ModelConfig, VllmConfig |
| 17 | from vllm.config.load import LoadConfig |
| 18 | from vllm.distributed import ( |
| 19 | get_dp_group, |
| 20 | get_pp_group, |
| 21 | get_tensor_model_parallel_rank, |
| 22 | get_tensor_model_parallel_world_size, |
| 23 | ) |
| 24 | from vllm.logger import init_logger |
| 25 | from vllm.model_executor.model_loader.base_loader import BaseModelLoader |
| 26 | from vllm.model_executor.model_loader.default_loader import DefaultModelLoader |
| 27 | from vllm.model_executor.model_loader.utils import ( |
| 28 | get_draft_load_config, |
| 29 | initialize_model, |
| 30 | process_weights_after_loading, |
| 31 | ) |
| 32 | from vllm.model_executor.model_loader.weight_cache.protocol import ( |
| 33 | CacheConfigMismatchError, |
| 34 | UnsupportedQuantForIPCError, |
| 35 | WeightCacheKey, |
| 36 | WeightCacheState, |
| 37 | WeightCacheUnavailableError, |
| 38 | check_ipc_platform_support, |
| 39 | check_ipc_quant_support, |
| 40 | get_current_device_uuid, |
| 41 | get_socket_path, |
| 42 | recv_msg, |
| 43 | send_msg, |
| 44 | verify_socket_owner, |
| 45 | ) |
| 46 | from vllm.model_executor.model_loader.weight_cache.utils import ( |
| 47 | is_draft_model_cacheable, |
| 48 | ) |
| 49 | from vllm.model_executor.utils import weights_already_processed |
| 50 | from vllm.tracing import instrument |
| 51 | from vllm.utils.torch_utils import set_default_torch_dtype |
| 52 | |
| 53 | logger = init_logger(__name__) |
| 54 | |
| 55 | _CONNECT_TIMEOUT_S = 5.0 |
| 56 | _STATE_TIMEOUT_S = 300.0 |
| 57 | _STARTUP_RETRY_INTERVAL_S = 0.5 |
| 58 | |
| 59 | _T = TypeVar("_T") |
| 60 | |
| 61 | |
| 62 | class IpcModelLoader(BaseModelLoader): |
| 63 | """Loads a model by mapping the weight cache daemon's tensors via CUDA IPC. |
| 64 | |
| 65 | The model is initialized on the meta device and every parameter/buffer is |
| 66 | replaced by the daemon's post-quantized tensor, so |
| 67 | process_weights_after_loading is skipped entirely. In "zero_copy" mode the |
| 68 | engine shares the daemon's GPU memory; in "copy" mode the tensors are |
| 69 | cloned into engine-owned memory and the daemon is asked to release its |
| 70 | cache afterwards. |
| 71 | |
| 72 | Extra config keys (via --model-loader-extra-config): |
| 73 | |
| 74 | - socket_path: explicit daemon socket path. Defaults to a per-GPU path |
| 75 | derived from the physical GPU uuid and the cache role (target/draft). |
| 76 | - socket_dir: directory containing the daemon sockets. |
| 77 | - mode: "zero_copy" (default) or "copy". |
| 78 | - fallback: fall back to disk loading when the daemon is unavailable or |
| 79 | the fingerprints mismatch (default: True). |
| 80 | - connect_timeout_s: socket connect timeout (default: 5.0). |
| 81 | - state_timeout_s: timeout for the weight-transfer request (default: 300.0). |
| 82 | |
| 83 | Note: in zero-copy mode the weights live in the daemon's CUDA IPC |
| 84 | allocations, so sleep mode (CuMemAllocator weight offloading) must not be |
| 85 | used with this loader. |
| 86 | """ |
| 87 | |
| 88 | def __init__(self, load_config: LoadConfig): |
| 89 | super().__init__(load_config) |
| 90 | extra_config = copy(load_config.model_loader_extra_config or {}) |
| 91 | self.socket_path: str | None = extra_config.pop("socket_path", None) |
| 92 | self.socket_dir: str | None = extra_config.pop("socket_dir", None) |
| 93 | # Internal: set by the engine when routing a speculative draft to the |
| 94 | # daemon's draft group. |
| 95 | self.is_draft = bool(extra_config.pop("is_draft", False)) |
| 96 | if self.is_draft and self.socket_path is not None: |
| 97 | raise ValueError( |
| 98 | "socket_path cannot be combined with the draft weight cache role; " |
| 99 | "use socket_dir so the target and draft sockets are derived " |
| 100 | "independently" |
| 101 | ) |
| 102 | self.mode: str = extra_config.pop("mode", "zero_copy") |
| 103 | self.fallback: bool = extra_config.pop("fallback", True) |
| 104 | self.connect_timeout_s: float = float( |
| 105 | extra_config.pop("connect_timeout_s", _CONNECT_TIMEOUT_S) |
| 106 | ) |
| 107 | self.state_timeout_s: float = float( |
| 108 | extra_config.pop("state_timeout_s", _STATE_TIMEOUT_S) |
| 109 | ) |
| 110 | if self.mode not in ("zero_copy", "copy"): |
| 111 | raise ValueError( |
| 112 | f"Invalid weight cache mode {self.mode!r}, " |
| 113 | "expected 'zero_copy' or 'copy'" |
| 114 | ) |
| 115 | if extra_config: |
| 116 | raise ValueError( |
| 117 | f"Unexpected extra config keys for load format " |
| 118 | f"{load_config.load_format}: {sorted(extra_config)}" |
| 119 | ) |
| 120 | |
| 121 | def get_external_weight_memory(self, vllm_config: VllmConfig) -> int: |
| 122 | # Copy mode clones the weights into this process; nothing external. |
| 123 | if self.mode != "zero_copy": |
| 124 | return 0 |
| 125 | total = self._daemon_memory() |
| 126 | if is_draft_model_cacheable(vllm_config.speculative_config): |
| 127 | # The draft group is queried with the same config the engine |
| 128 | # would load the draft with. |
| 129 | draft_config = get_draft_load_config(vllm_config) |
| 130 | if draft_config.load_format == "ipc_cache": |
| 131 | draft_loader = IpcModelLoader(draft_config) |
| 132 | # Only a zero-copy draft-group daemon holds external weights; |
| 133 | # an explicit draft_load_config without is_draft would resolve |
| 134 | # back to the target socket and double-count it. |
| 135 | if draft_loader.is_draft and draft_loader.mode == "zero_copy": |
| 136 | total += draft_loader._daemon_memory() |
| 137 | return total |
| 138 | |
| 139 | def _daemon_memory(self) -> int: |
| 140 | if not self.fallback: |
| 141 | # The loader waits for the daemon at load time, so the weights |
| 142 | # will be zero-copy mapped for sure; mirror that wait here since |
| 143 | # returning 0 would over-grant the memory budget. |
| 144 | return self._with_startup_wait(self._query_daemon_memory) |
| 145 | # An unreachable daemon means a disk load, i.e. nothing external. |
| 146 | try: |
| 147 | return self._query_daemon_memory() |
| 148 | except (WeightCacheUnavailableError, ConnectionError, OSError) as e: |
| 149 | logger.warning( |
| 150 | "Cannot query weight cache daemon memory (%s); " |
| 151 | "assuming the weights are not externally held", |
| 152 | e, |
| 153 | ) |
| 154 | return 0 |
| 155 | |
| 156 | def _query_daemon_memory(self) -> int: |
| 157 | with self._connect(self.connect_timeout_s) as conn: |
| 158 | send_msg(conn, {"cmd": "get_memory"}) |
| 159 | response = recv_msg(conn) |
| 160 | if response.get("status") != "ok": |
| 161 | raise WeightCacheUnavailableError( |
| 162 | "Weight cache daemon rejected the memory query: " |
| 163 | f"{response.get('message')}" |
| 164 | ) |
| 165 | return int(response.get("memory_bytes", 0)) |
| 166 | |
| 167 | def download_model(self, model_config: ModelConfig) -> None: |
| 168 | DefaultModelLoader(self._fallback_load_config()).download_model(model_config) |
| 169 | |
| 170 | def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: |
| 171 | """Best-effort in-place reload for an already-initialized model. |
| 172 | |
| 173 | Copies daemon tensors into matching parameters/buffers. The model is |
| 174 | expected to already be in the post-quantized layout (e.g. previously |
| 175 | loaded through this loader). |
| 176 | """ |
| 177 | device_index = torch.accelerator.current_device_index() |
| 178 | entries = self._fetch_entries(model_config).entries |
| 179 | params = dict(model.named_parameters()) |
| 180 | buffers = dict(model.named_buffers()) |
| 181 | for name, entry in entries.items(): |
| 182 | target = params.get(name, buffers.get(name)) |
| 183 | source = entry.rebuild(device_index) |
| 184 | if target is None or target.shape != source.shape: |
| 185 | logger.warning("Skipping mismatched cached tensor %s", name) |
| 186 | continue |
| 187 | target.data.copy_(source) |
| 188 | |
| 189 | @instrument(span_name="Load model") |
| 190 | def load_model( |
| 191 | self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" |
| 192 | ) -> nn.Module: |
| 193 | # An unsupported platform is a permanent misconfiguration rather than |
| 194 | # a transient daemon outage, so it is raised even when fallback is on. |
| 195 | check_ipc_platform_support() |
| 196 | state_fetched = False |
| 197 | try: |
| 198 | # Cross-check the routing flag against the identity of the model |
| 199 | # being loaded: a draft load that lost its flag (or a target load |
| 200 | # that got one) would hit the wrong daemon group and |
| 201 | # fingerprint-mismatch. |
| 202 | spec = vllm_config.speculative_config |
| 203 | inferred = spec is not None and model_config is spec.draft_model_config |
| 204 | if inferred != self.is_draft: |
| 205 | raise CacheConfigMismatchError( |
| 206 | f"Weight cache role mismatch: loading " |
| 207 | f"{'draft' if inferred else 'target'} model but the loader " |
| 208 | f"was configured for the " |
| 209 | f"{'draft' if self.is_draft else 'target'} group" |
| 210 | ) |
| 211 | state = self._fetch_entries(model_config) |
| 212 | state_fetched = True |
| 213 | return self._build_model(vllm_config, model_config, prefix, state) |
| 214 | except (WeightCacheUnavailableError, CacheConfigMismatchError) as e: |
| 215 | if not self.fallback: |
| 216 | raise |
| 217 | logger.warning( |
| 218 | "Weight cache unusable (%s); falling back to disk loading", e |
| 219 | ) |
| 220 | except UnsupportedQuantForIPCError: |
| 221 | # Unsupported quantization is a permanent misconfiguration rather |
| 222 | # than a transient daemon outage, so it is raised even when |
| 223 | # fallback is on. |
| 224 | raise |
| 225 | except Exception: |
| 226 | if not self.fallback: |
| 227 | raise |
| 228 | logger.exception( |
| 229 | "Weight cache IPC loading failed; falling back to disk loading" |
| 230 | ) |
| 231 | # _build_model failed after fetching state without reaching its |
| 232 | # copy-mode release, so the daemon still holds the full cache; |
| 233 | # release it so the disk fallback does not OOM against it. |
| 234 | if state_fetched and self.mode == "copy": |
| 235 | self._send_release() |
| 236 | torch.accelerator.empty_cache() |
| 237 | return self._fallback_load(vllm_config, model_config, prefix) |
| 238 | |
| 239 | def _build_model( |
| 240 | self, |
| 241 | vllm_config: VllmConfig, |
| 242 | model_config: ModelConfig, |
| 243 | prefix: str, |
| 244 | state: WeightCacheState, |
| 245 | ) -> nn.Module: |
| 246 | device_config = vllm_config.device_config |
| 247 | load_device = ( |
| 248 | device_config.device |
| 249 | if self.load_config.device is None |
| 250 | else self.load_config.device |
| 251 | ) |
| 252 | target_device = torch.device(load_device) |
| 253 | device_index = ( |
| 254 | target_device.index |
| 255 | if target_device.index is not None |
| 256 | else torch.accelerator.current_device_index() |
| 257 | ) |
| 258 | with set_default_torch_dtype(model_config.dtype): |
| 259 | with torch.device("meta"): |
| 260 | model = initialize_model( |
| 261 | vllm_config=vllm_config, |
| 262 | model_config=model_config, |
| 263 | prefix=prefix, |
| 264 | ) |
| 265 | check_ipc_quant_support(model) |
| 266 | self._apply_entries(model, state, device_index) |
| 267 | # Flags that load_weights would have set (e.g. EAGLE ownership of |
| 268 | # embed_tokens / lm_head); the daemon ran it, this process did not. |
| 269 | for name, value in state.attrs.items(): |
| 270 | setattr(model, name, value) |
| 271 | # The daemon exports tensors that already went through |
| 272 | # process_weights_after_loading; re-run it in pre-processed mode |
| 273 | # so quant methods only rebuild Python-side state (e.g. the MoE |
| 274 | # kernel). Leftovers are materialized afterwards so that |
| 275 | # placeholders the daemon-side post-processing consumed are |
| 276 | # dropped rather than filled with uninitialized memory. |
| 277 | with weights_already_processed(): |
| 278 | process_weights_after_loading(model, model_config, target_device) |
| 279 | _materialize_remaining_meta_tensors( |
| 280 | model, torch.device(target_device.type, device_index) |
| 281 | ) |
| 282 | if self.mode == "copy": |
| 283 | self._send_release() |
| 284 | logger.info( |
| 285 | "Mapped %d tensors from the weight cache daemon (%s mode)", |
| 286 | len(state.entries), |
| 287 | self.mode, |
| 288 | ) |
| 289 | return model.eval() |
| 290 | |
| 291 | def _apply_entries( |
| 292 | self, |
| 293 | model: nn.Module, |
| 294 | state: WeightCacheState, |
| 295 | device_index: int, |
| 296 | ) -> None: |
| 297 | # remove_duplicate=False keeps tied module aliases reachable by name: |
| 298 | # a tied lm_head *is* the embedding module, so the deduplicated view |
| 299 | # would not contain "lm_head" at all. |
| 300 | modules = dict(model.named_modules(remove_duplicate=False)) |
| 301 | registered: dict[str, torch.Tensor] = {} |
| 302 | |
| 303 | def _register(name: str, tensor: torch.Tensor, is_param: bool) -> None: |
| 304 | module_name, _, leaf = name.rpartition(".") |
| 305 | module = modules.get(module_name) |
| 306 | if module is None: |
| 307 | raise RuntimeError(f"Cached tensor {name} has no matching module") |
| 308 | # Replace via registration rather than param.data assignment, |
| 309 | # which fails for meta tensors. Entries may also introduce |
| 310 | # post-quantization tensors absent from the meta model. |
| 311 | module._parameters.pop(leaf, None) |
| 312 | module._buffers.pop(leaf, None) |
| 313 | if is_param: |
| 314 | obj: torch.Tensor = ( |
| 315 | tensor |
| 316 | if isinstance(tensor, nn.Parameter) |
| 317 | else nn.Parameter(tensor, requires_grad=False) |
| 318 | ) |
| 319 | module.register_parameter(leaf, obj) |
| 320 | else: |
| 321 | obj = tensor |
| 322 | module.register_buffer(leaf, obj) |
| 323 | registered[name] = obj |
| 324 | |
| 325 | for name, entry in state.entries.items(): |
| 326 | tensor = entry.rebuild(device_index) |
| 327 | if self.mode == "copy": |
| 328 | tensor = tensor.clone() |
| 329 | _register(name, tensor, entry.kind == "param") |
| 330 | |
| 331 | # Re-establish tied-weight aliases by registering the *same* object the |
| 332 | # canonical name resolved to, so parameter identity (and the tie) is |
| 333 | # preserved instead of allocating uninitialized memory. |
| 334 | for alias_name, canonical_name in state.aliases.items(): |
| 335 | obj = registered.get(canonical_name) |
| 336 | if obj is None: |
| 337 | logger.warning( |
| 338 | "Cached alias %s references missing canonical tensor %s", |
| 339 | alias_name, |
| 340 | canonical_name, |
| 341 | ) |
| 342 | continue |
| 343 | _register(alias_name, obj, isinstance(obj, nn.Parameter)) |
| 344 | |
| 345 | def _fetch_entries(self, model_config: ModelConfig) -> WeightCacheState: |
| 346 | dp_group = get_dp_group() |
| 347 | pp_group = get_pp_group() |
| 348 | cache_config = WeightCacheKey.from_model_config( |
| 349 | model_config, |
| 350 | tp_size=get_tensor_model_parallel_world_size(), |
| 351 | tp_rank=get_tensor_model_parallel_rank(), |
| 352 | pp_size=pp_group.world_size, |
| 353 | pp_rank=pp_group.rank_in_group, |
| 354 | dp_size=dp_group.world_size, |
| 355 | dp_rank=dp_group.rank_in_group, |
| 356 | is_draft=self.is_draft, |
| 357 | ) |
| 358 | if not self.fallback: |
| 359 | return self._request_state_with_startup_wait(cache_config) |
| 360 | return self._request_state(cache_config) |
| 361 | |
| 362 | def _request_state_with_startup_wait( |
| 363 | self, cache_config: WeightCacheKey |
| 364 | ) -> WeightCacheState: |
| 365 | return self._with_startup_wait(lambda: self._request_state(cache_config)) |
| 366 | |
| 367 | def _with_startup_wait(self, op: Callable[[], _T]) -> _T: |
| 368 | """Retry op until the daemon answers or the state timeout elapses; |
| 369 | the daemon may still be loading the model when the engine starts.""" |
| 370 | deadline = time.monotonic() + self.state_timeout_s |
| 371 | while True: |
| 372 | try: |
| 373 | return op() |
| 374 | except (WeightCacheUnavailableError, ConnectionError, OSError) as e: |
| 375 | if time.monotonic() >= deadline: |
| 376 | raise WeightCacheUnavailableError( |
| 377 | "Weight cache daemon did not become ready within " |
| 378 | f"{self.state_timeout_s:.1f}s: {e}" |
| 379 | ) from e |
| 380 | logger.info_once( |
| 381 | "Waiting up to %.1fs for the weight cache daemon to start", |
| 382 | self.state_timeout_s, |
| 383 | ) |
| 384 | time.sleep( |
| 385 | max( |
| 386 | 0.0, |
| 387 | min(_STARTUP_RETRY_INTERVAL_S, deadline - time.monotonic()), |
| 388 | ) |
| 389 | ) |
| 390 | |
| 391 | def _request_state(self, cache_config: WeightCacheKey) -> WeightCacheState: |
| 392 | with self._connect(self.state_timeout_s) as conn: |
| 393 | send_msg(conn, {"cmd": "get_state", "cache_config": cache_config}) |
| 394 | response = recv_msg(conn) |
| 395 | status = response.get("status") |
| 396 | if status == "mismatch": |
| 397 | raise CacheConfigMismatchError( |
| 398 | f"WeightCacheKey mismatch on fields: {response.get('fields')}" |
| 399 | ) |
| 400 | if status != "ok": |
| 401 | raise WeightCacheUnavailableError( |
| 402 | f"Weight cache daemon error: {response.get('message')}" |
| 403 | ) |
| 404 | self._check_gpu_uuid(response.get("gpu_uuid")) |
| 405 | return WeightCacheState( |
| 406 | entries=response["entries"], |
| 407 | aliases=response.get("aliases", {}), |
| 408 | attrs=response.get("attrs", {}), |
| 409 | ) |
| 410 | |
| 411 | def _connect(self, timeout: float) -> socket.socket: |
| 412 | socket_path = self._resolve_socket_path() |
| 413 | # The auto-derived per-user directory is locked to 0700 and checked |
| 414 | # strictly. When the operator explicitly configures a path they own the |
| 415 | # trust decision, so only ownership/symlink safety is enforced. |
| 416 | strict_perms = self.socket_path is None and self.socket_dir is None |
| 417 | try: |
| 418 | verify_socket_owner(socket_path, strict_perms=strict_perms) |
| 419 | except OSError as e: |
| 420 | raise WeightCacheUnavailableError( |
| 421 | f"Weight cache socket {socket_path} is unavailable: {e}" |
| 422 | ) from e |
| 423 | sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) |
| 424 | sock.settimeout(timeout) |
| 425 | try: |
| 426 | sock.connect(socket_path) |
| 427 | except OSError as e: |
| 428 | sock.close() |
| 429 | raise WeightCacheUnavailableError( |
| 430 | f"Cannot connect to weight cache daemon at {socket_path}: {e}" |
| 431 | ) from e |
| 432 | return sock |
| 433 | |
| 434 | def _resolve_socket_path(self) -> str: |
| 435 | if self.socket_path is not None: |
| 436 | return self.socket_path |
| 437 | return get_socket_path( |
| 438 | get_current_device_uuid(), |
| 439 | self.socket_dir, |
| 440 | is_draft=self.is_draft, |
| 441 | ) |
| 442 | |
| 443 | def _check_gpu_uuid(self, daemon_uuid: str | None) -> None: |
| 444 | if daemon_uuid is None: |
| 445 | return |
| 446 | local_uuid = get_current_device_uuid() |
| 447 | if daemon_uuid != local_uuid: |
| 448 | raise CacheConfigMismatchError( |
| 449 | f"Daemon GPU {daemon_uuid} != engine GPU {local_uuid}; " |
| 450 | "check the socket path / GPU mapping" |
| 451 | ) |
| 452 | |
| 453 | def _send_release(self) -> None: |
| 454 | try: |
| 455 | with self._connect(self.connect_timeout_s) as conn: |
| 456 | send_msg(conn, {"cmd": "release"}) |
| 457 | recv_msg(conn) |
| 458 | except (WeightCacheUnavailableError, ConnectionError, OSError): |
| 459 | logger.warning("Failed to ask the weight cache daemon to release") |
| 460 | |
| 461 | def _fallback_load_config(self) -> LoadConfig: |
| 462 | # DefaultModelLoader must not see load_format="ipc_cache" or the ipc |
| 463 | # extra config keys. |
| 464 | return dataclasses.replace( |
| 465 | self.load_config, |
| 466 | load_format="auto", |
| 467 | model_loader_extra_config={}, |
| 468 | ) |
| 469 | |
| 470 | def _fallback_load( |
| 471 | self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str |
| 472 | ) -> nn.Module: |
| 473 | loader = DefaultModelLoader(self._fallback_load_config()) |
| 474 | return loader.load_model( |
| 475 | vllm_config=vllm_config, model_config=model_config, prefix=prefix |
| 476 | ) |
| 477 | |
| 478 | |
| 479 | def _materialize_remaining_meta_tensors(model: nn.Module, device: torch.device) -> None: |
| 480 | """Allocate any tensors the daemon did not provide. |
| 481 | |
| 482 | These are typically parameters removed on the daemon side by |
| 483 | process_weights_after_loading; they are not expected to be read at |
| 484 | runtime, so they are left uninitialized. |
| 485 | """ |
| 486 | for module_name, module in model.named_modules(): |
| 487 | for leaf, param in list(module._parameters.items()): |
| 488 | if param is not None and param.device.type == "meta": |
| 489 | logger.warning( |
| 490 | "Materializing empty parameter %s.%s missing from the weight cache", |
| 491 | module_name, |
| 492 | leaf, |
| 493 | ) |
| 494 | module._parameters[leaf] = nn.Parameter( |
| 495 | torch.empty_like(param, device=device), requires_grad=False |
| 496 | ) |
| 497 | for leaf, buffer in list(module._buffers.items()): |
| 498 | if buffer is not None and buffer.device.type == "meta": |
| 499 | module._buffers[leaf] = torch.empty_like(buffer, device=device) |