| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Weight cache daemon for fast engine restarts. |
| 4 | |
| 5 | One daemon process per GPU holds the post-quantized, TP-sharded weights of its |
| 6 | rank in GPU memory and serves CUDA IPC handles to vLLM engines over a Unix |
| 7 | domain socket. Restarting engines map the weights via zero-copy IPC instead of |
| 8 | reloading from disk. |
| 9 | |
| 10 | Launch one daemon per TP rank with a single command: |
| 11 | |
| 12 | vllm preload \\ |
| 13 | --model /path/to/model --tensor-parallel-size 4 |
| 14 | |
| 15 | Engines then load from the daemons with: |
| 16 | |
| 17 | vllm serve /path/to/model --tensor-parallel-size 4 \\ |
| 18 | --load-format ipc_cache |
| 19 | |
| 20 | Tensor, expert, data and pipeline parallelism are supported. |
| 21 | |
| 22 | For multi-node tensor parallelism, run one launcher per node with a shared |
| 23 | rendezvous so the global TP group forms across nodes (CUDA IPC handles are |
| 24 | node-local, so each node serves only its local GPUs' shards). Reuse the same |
| 25 | ``--nnodes``/``--node-rank``/``--master-addr`` flags you pass the engine, plus a |
| 26 | ``--weight-cache-master-port`` distinct from the engine's ``--master-port``: |
| 27 | |
| 28 | # node 0 (8 local GPUs) |
| 29 | vllm preload \\ |
| 30 | --model /path/to/model --tensor-parallel-size 16 \\ |
| 31 | --nnodes 2 --node-rank 0 --master-addr 10.0.0.1 \\ |
| 32 | --weight-cache-master-port 29600 |
| 33 | # node 1 (8 local GPUs) |
| 34 | vllm preload \\ |
| 35 | --model /path/to/model --tensor-parallel-size 16 \\ |
| 36 | --nnodes 2 --node-rank 1 --master-addr 10.0.0.1 \\ |
| 37 | --weight-cache-master-port 29600 |
| 38 | |
| 39 | The global TP rank of local GPU ``i`` on node ``r`` is |
| 40 | ``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank |
| 41 | assignment, so each engine worker maps its shard from the daemon on its own |
| 42 | node. With pipeline parallelism each node additionally runs one daemon per PP |
| 43 | stage; the global rank is ``pp_rank * tp_size + tp_rank``. |
| 44 | |
| 45 | For data parallelism (e.g. a TP1 x DP16 x EP decode fleet) run one launcher per |
| 46 | node with the engine's DP placement flags. Local GPU ``i`` serves DP rank |
| 47 | ``start_rank + i // tp_size`` and TP rank ``i % tp_size``, and all |
| 48 | ``dp_size * tp_size`` daemons form one world group on |
| 49 | ``--data-parallel-address``/``--weight-cache-master-port`` so the expert |
| 50 | shards are laid out exactly as in the engine: |
| 51 | |
| 52 | # node r (4 local GPUs) |
| 53 | vllm preload \\ |
| 54 | --model /path/to/model --tensor-parallel-size 1 --enable-expert-parallel \\ |
| 55 | --data-parallel-size 16 --data-parallel-size-local 4 \\ |
| 56 | --data-parallel-start-rank 4r --data-parallel-address 10.0.0.1 \\ |
| 57 | --weight-cache-master-port 29600 |
| 58 | |
| 59 | Data parallelism also combines with multi-node tensor parallelism: pass both |
| 60 | flag sets (``--nnodes``/``--node-rank``/``--master-addr`` and the DP flags with |
| 61 | a reachable ``--data-parallel-address``). Each node then serves a contiguous |
| 62 | block of the ``dp_size * tp_size`` global ranks, e.g. TP8 x DP2 on 4 nodes: |
| 63 | |
| 64 | # node r (4 local GPUs) |
| 65 | vllm preload \\ |
| 66 | --model /path/to/model --tensor-parallel-size 8 --enable-expert-parallel \\ |
| 67 | --nnodes 4 --node-rank r --master-addr 10.0.0.1 \\ |
| 68 | --data-parallel-size 2 --data-parallel-address 10.0.0.1 \\ |
| 69 | --weight-cache-master-port 29600 |
| 70 | |
| 71 | With MTP, EAGLE or EAGLE3 speculative decoding the launcher additionally |
| 72 | starts a draft daemon group that caches the draft model. It uses its own cache |
| 73 | key, Unix sockets (``*_draft.sock``) and rendezvous port |
| 74 | (``--weight-cache-draft-master-port``, default ``--weight-cache-master-port + |
| 75 | 1``), so each process serves exactly one model role. Other draft types are not |
| 76 | cached and keep loading from disk in the engine. |
| 77 | """ |
| 78 | |
| 79 | import contextlib |
| 80 | import fcntl |
| 81 | import gc |
| 82 | import multiprocessing |
| 83 | import os |
| 84 | import socket |
| 85 | from collections.abc import Callable |
| 86 | |
| 87 | import torch |
| 88 | |
| 89 | from vllm.config import ( |
| 90 | ParallelConfig, |
| 91 | VllmConfig, |
| 92 | replace, |
| 93 | set_current_vllm_config, |
| 94 | ) |
| 95 | from vllm.distributed import ( |
| 96 | ensure_model_parallel_initialized, |
| 97 | init_distributed_environment, |
| 98 | ) |
| 99 | from vllm.logger import init_logger |
| 100 | from vllm.model_executor.model_loader import get_model_loader |
| 101 | from vllm.model_executor.model_loader.utils import process_weights_after_loading |
| 102 | from vllm.model_executor.model_loader.weight_cache.protocol import ( |
| 103 | TensorEntry, |
| 104 | WeightCacheKey, |
| 105 | WeightCacheUnavailableError, |
| 106 | check_ipc_quant_support, |
| 107 | ensure_private_socket_dir, |
| 108 | get_current_device_uuid, |
| 109 | get_socket_path, |
| 110 | recv_msg, |
| 111 | send_msg, |
| 112 | verify_peer_is_owner, |
| 113 | ) |
| 114 | from vllm.model_executor.model_loader.weight_cache.utils import ( |
| 115 | export_model_attrs, |
| 116 | format_daemon_role, |
| 117 | is_draft_model_cacheable, |
| 118 | ) |
| 119 | from vllm.platforms import current_platform |
| 120 | from vllm.utils.mem_utils import format_gib |
| 121 | from vllm.utils.torch_utils import set_default_torch_dtype |
| 122 | from vllm.v1.worker.workspace import init_workspace_manager |
| 123 | |
| 124 | logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon") |
| 125 | |
| 126 | |
| 127 | def export_entries( |
| 128 | model: torch.nn.Module, |
| 129 | ) -> tuple[dict[str, TensorEntry], dict[str, str]]: |
| 130 | """Export a model's tensors, preserving tied-parameter aliases. |
| 131 | |
| 132 | ``named_parameters``/``named_buffers`` are iterated with |
| 133 | ``remove_duplicate=False`` so tied weights (e.g. ``lm_head.weight`` sharing |
| 134 | storage with ``embed_tokens.weight``) are not silently dropped. Each unique |
| 135 | tensor is exported once per call; every additional name that refers to the same |
| 136 | tensor object is recorded in the returned alias map so the client can |
| 137 | re-establish the shared identity instead of allocating uninitialized |
| 138 | memory for it. |
| 139 | |
| 140 | CUDA reduction arguments must be exported separately for each consumer so |
| 141 | that PyTorch registers a reference for each IPC mapping's lifetime. |
| 142 | |
| 143 | Returns: |
| 144 | A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to |
| 145 | its ``TensorEntry`` and ``aliases`` maps each duplicate name to its |
| 146 | canonical name. |
| 147 | |
| 148 | """ |
| 149 | entries: dict[str, TensorEntry] = {} |
| 150 | aliases: dict[str, str] = {} |
| 151 | canonical_by_id: dict[int, str] = {} |
| 152 | |
| 153 | def _add(name: str, tensor: torch.Tensor, kind: str) -> None: |
| 154 | canonical = canonical_by_id.get(id(tensor)) |
| 155 | if canonical is not None: |
| 156 | aliases[name] = canonical |
| 157 | return |
| 158 | canonical_by_id[id(tensor)] = name |
| 159 | entries[name] = TensorEntry.from_tensor(tensor, kind) |
| 160 | |
| 161 | for name, param in model.named_parameters(remove_duplicate=False): |
| 162 | _add(name, param, "param") |
| 163 | # named_buffers includes non-persistent buffers (e.g. rotary embedding |
| 164 | # caches) that state_dict would miss. |
| 165 | for name, buffer in model.named_buffers(remove_duplicate=False): |
| 166 | if name in entries or name in aliases: |
| 167 | continue |
| 168 | _add(name, buffer, "buffer") |
| 169 | return entries, aliases |
| 170 | |
| 171 | |
| 172 | class WeightCacheDaemon: |
| 173 | """Per-GPU process that loads one TP/PP shard and serves CUDA IPC handles.""" |
| 174 | |
| 175 | def __init__( |
| 176 | self, |
| 177 | vllm_config: VllmConfig, |
| 178 | global_rank: int, |
| 179 | local_rank: int, |
| 180 | distributed_init_method: str, |
| 181 | socket_dir: str | None = None, |
| 182 | is_draft: bool = False, |
| 183 | dp_rank: int = 0, |
| 184 | pp_rank: int = 0, |
| 185 | ): |
| 186 | parallel_config = vllm_config.parallel_config |
| 187 | self.tp_size = parallel_config.tensor_parallel_size |
| 188 | self.pp_size = parallel_config.pipeline_parallel_size |
| 189 | self.dp_size = parallel_config.data_parallel_size |
| 190 | if self.dp_size > 1: |
| 191 | # The engine sets these on each DP rank's workers; the MoE/EP |
| 192 | # layers read them to place their expert shards. |
| 193 | if parallel_config.nnodes > 1: |
| 194 | world_size = self.dp_size * self.tp_size |
| 195 | local_world = world_size // parallel_config.nnodes |
| 196 | start_dp_rank = parallel_config.node_rank * local_world // self.tp_size |
| 197 | else: |
| 198 | start_dp_rank = parallel_config.data_parallel_rank |
| 199 | parallel_config = replace( |
| 200 | parallel_config, |
| 201 | data_parallel_rank=dp_rank, |
| 202 | data_parallel_rank_local=dp_rank - start_dp_rank, |
| 203 | ) |
| 204 | vllm_config = replace(vllm_config, parallel_config=parallel_config) |
| 205 | self.vllm_config = vllm_config |
| 206 | # A draft daemon builds the draft with the target's VllmConfig, like |
| 207 | # the engine does; only the ModelConfig differs. |
| 208 | if is_draft: |
| 209 | spec = vllm_config.speculative_config |
| 210 | assert spec is not None and spec.draft_model_config is not None |
| 211 | self.model_config = spec.draft_model_config |
| 212 | else: |
| 213 | self.model_config = vllm_config.model_config |
| 214 | self.global_rank = global_rank |
| 215 | rank_in_dp = global_rank - dp_rank * self.pp_size * self.tp_size |
| 216 | self.tp_rank = rank_in_dp % self.tp_size |
| 217 | self.pp_rank = rank_in_dp // self.tp_size |
| 218 | self.dp_rank = dp_rank |
| 219 | self.world_size = self.dp_size * self.pp_size * self.tp_size |
| 220 | self.local_rank = local_rank |
| 221 | self.distributed_init_method = distributed_init_method |
| 222 | self.socket_dir = socket_dir |
| 223 | self.is_draft = is_draft |
| 224 | self.role = format_daemon_role(is_draft) |
| 225 | self.model: torch.nn.Module | None = None |
| 226 | # Fingerprint before loading: process_weights_after_loading may |
| 227 | # mutate hf_config.quantization_config. |
| 228 | self.cache_config = WeightCacheKey.from_model_config( |
| 229 | self.model_config, |
| 230 | tp_size=self.tp_size, |
| 231 | tp_rank=self.tp_rank, |
| 232 | pp_size=self.pp_size, |
| 233 | pp_rank=self.pp_rank, |
| 234 | dp_size=self.dp_size, |
| 235 | dp_rank=dp_rank, |
| 236 | is_draft=is_draft, |
| 237 | ) |
| 238 | |
| 239 | def load_model(self) -> None: |
| 240 | torch.accelerator.set_device_index(self.local_rank) |
| 241 | # Outside the config context so the daemon keeps its own rendezvous |
| 242 | # instead of the engine's; model parallel groups (TP/DP/EP) are then |
| 243 | # carved out of this world group like in the engine. |
| 244 | init_distributed_environment( |
| 245 | world_size=self.world_size, |
| 246 | rank=self.global_rank, |
| 247 | distributed_init_method=self.distributed_init_method, |
| 248 | local_rank=self.local_rank, |
| 249 | backend=current_platform.dist_backend, |
| 250 | ) |
| 251 | with set_current_vllm_config(self.vllm_config): |
| 252 | ensure_model_parallel_initialized(self.tp_size, self.pp_size) |
| 253 | self.model = self.get_model() |
| 254 | # Loading and post-processing leave freed transients in the caching |
| 255 | # allocator; return them so engines sharing the GPU can use them. |
| 256 | gc.collect() |
| 257 | torch.accelerator.empty_cache() |
| 258 | logger.info( |
| 259 | "Weight cache %s daemon rank %d loaded model (%s GiB allocated, " |
| 260 | "%s GiB reserved)", |
| 261 | self.role, |
| 262 | self.global_rank, |
| 263 | format_gib(torch.accelerator.memory_allocated()), |
| 264 | format_gib(torch.accelerator.memory_reserved()), |
| 265 | ) |
| 266 | |
| 267 | def get_model(self) -> torch.nn.Module: |
| 268 | """Load the daemon's model, composed from the configured loader. |
| 269 | |
| 270 | Runs the quantization check after model creation but before the slow |
| 271 | weight load, so an unsupported method fails fast. Online quantization |
| 272 | always fails the check, so load_model's finalize step for it is |
| 273 | unnecessary here. |
| 274 | """ |
| 275 | vllm_config = self.vllm_config |
| 276 | model_config = self.model_config |
| 277 | load_config = vllm_config.load_config |
| 278 | loader = get_model_loader(load_config) |
| 279 | device_config = vllm_config.device_config |
| 280 | target_device = torch.device( |
| 281 | device_config.device if load_config.device is None else load_config.device |
| 282 | ) |
| 283 | # Attention backends may take workspace at construction, like in the |
| 284 | # engine's worker init. |
| 285 | init_workspace_manager(target_device) |
| 286 | with set_default_torch_dtype(model_config.dtype): |
| 287 | with target_device: |
| 288 | model = loader.create_model(vllm_config, model_config) |
| 289 | check_ipc_quant_support(model) |
| 290 | loader.load_weights(model, model_config) |
| 291 | process_weights_after_loading(model, model_config, target_device) |
| 292 | return model.eval() |
| 293 | |
| 294 | def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None: |
| 295 | """Serve requests until terminated. |
| 296 | |
| 297 | The socket is only bound once the model is fully cached, so clients |
| 298 | get a connection error (and fall back to disk) until the daemon is |
| 299 | ready. |
| 300 | |
| 301 | Args: |
| 302 | ready_callback: Invoked once the socket is bound and listening, |
| 303 | so the launcher can report overall readiness. |
| 304 | |
| 305 | """ |
| 306 | socket_path = self._socket_path |
| 307 | ensure_private_socket_dir( |
| 308 | os.path.dirname(socket_path), strict_perms=self.socket_dir is None |
| 309 | ) |
| 310 | # Hold an exclusive per-GPU lock for the daemon's lifetime so a second |
| 311 | # daemon cannot remove this daemon's live socket and hijack the path. |
| 312 | lock_fd = self._acquire_gpu_lock(socket_path) |
| 313 | if os.path.exists(socket_path): |
| 314 | os.unlink(socket_path) |
| 315 | server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) |
| 316 | server.bind(socket_path) |
| 317 | os.chmod(socket_path, 0o600) |
| 318 | server.listen() |
| 319 | logger.info( |
| 320 | "Weight cache %s daemon rank %d serving on %s", |
| 321 | self.role, |
| 322 | self.global_rank, |
| 323 | socket_path, |
| 324 | ) |
| 325 | if ready_callback is not None: |
| 326 | ready_callback() |
| 327 | try: |
| 328 | while True: |
| 329 | conn, _ = server.accept() |
| 330 | with conn: |
| 331 | try: |
| 332 | verify_peer_is_owner(conn) |
| 333 | self._handle_connection(conn) |
| 334 | except (ConnectionError, EOFError): |
| 335 | logger.warning("Client disconnected mid-request") |
| 336 | except Exception as e: |
| 337 | # Report the error back instead of just closing the |
| 338 | # socket, but don't let it take the daemon down. |
| 339 | logger.exception( |
| 340 | "Error handling weight cache client; continuing" |
| 341 | ) |
| 342 | with contextlib.suppress(OSError): |
| 343 | send_msg(conn, {"status": "error", "message": str(e)}) |
| 344 | finally: |
| 345 | server.close() |
| 346 | if os.path.exists(socket_path): |
| 347 | os.unlink(socket_path) |
| 348 | os.close(lock_fd) |
| 349 | |
| 350 | def _acquire_gpu_lock(self, socket_path: str) -> int: |
| 351 | """Take an exclusive lock guarding this GPU's socket path. |
| 352 | |
| 353 | The lock is advisory and released automatically when the daemon exits |
| 354 | (or crashes), so a stale socket is only ever removed by whoever owns |
| 355 | the lock. A running daemon holding it makes a second daemon fail fast |
| 356 | instead of clobbering the live socket. |
| 357 | """ |
| 358 | lock_fd = os.open(f"{socket_path}.lock", os.O_CREAT | os.O_RDWR, 0o600) |
| 359 | try: |
| 360 | fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) |
| 361 | except OSError as e: |
| 362 | os.close(lock_fd) |
| 363 | raise WeightCacheUnavailableError( |
| 364 | f"Another weight cache daemon already owns {socket_path}" |
| 365 | ) from e |
| 366 | return lock_fd |
| 367 | |
| 368 | @property |
| 369 | def _socket_path(self) -> str: |
| 370 | return get_socket_path( |
| 371 | get_current_device_uuid(), |
| 372 | self.socket_dir, |
| 373 | is_draft=self.is_draft, |
| 374 | ) |
| 375 | |
| 376 | def _handle_connection(self, conn: socket.socket) -> None: |
| 377 | request = recv_msg(conn) |
| 378 | cmd = request.get("cmd") |
| 379 | if cmd == "get_state": |
| 380 | self._handle_get_state(conn, request) |
| 381 | elif cmd == "get_memory": |
| 382 | send_msg( |
| 383 | conn, |
| 384 | { |
| 385 | "status": "ok", |
| 386 | "memory_bytes": torch.accelerator.memory_allocated(), |
| 387 | }, |
| 388 | ) |
| 389 | elif cmd == "release": |
| 390 | self._handle_release(conn) |
| 391 | else: |
| 392 | send_msg(conn, {"status": "error", "message": f"Unknown command {cmd!r}"}) |
| 393 | |
| 394 | def _handle_get_state(self, conn: socket.socket, request: dict) -> None: |
| 395 | client_config = request.get("cache_config") |
| 396 | if not isinstance(client_config, WeightCacheKey): |
| 397 | send_msg(conn, {"status": "error", "message": "Missing cache_config"}) |
| 398 | return |
| 399 | mismatched = self.cache_config.mismatched_fields(client_config) |
| 400 | if mismatched: |
| 401 | logger.warning("WeightCacheKey mismatch on fields: %s", mismatched) |
| 402 | send_msg(conn, {"status": "mismatch", "fields": mismatched}) |
| 403 | return |
| 404 | if self.model is None: |
| 405 | send_msg(conn, {"status": "error", "message": "Weights were released"}) |
| 406 | return |
| 407 | gpu_uuid = get_current_device_uuid() |
| 408 | entries, aliases = export_entries(self.model) |
| 409 | send_msg( |
| 410 | conn, |
| 411 | { |
| 412 | "status": "ok", |
| 413 | "entries": entries, |
| 414 | "aliases": aliases, |
| 415 | "attrs": export_model_attrs(self.model), |
| 416 | "gpu_uuid": gpu_uuid, |
| 417 | }, |
| 418 | ) |
| 419 | logger.info_once( |
| 420 | "Weight cache %s daemon rank %d sent %d tensors (+%d aliases) to engine", |
| 421 | self.role, |
| 422 | self.global_rank, |
| 423 | len(entries), |
| 424 | len(aliases), |
| 425 | ) |
| 426 | |
| 427 | def _handle_release(self, conn: socket.socket) -> None: |
| 428 | self.model = None |
| 429 | torch.accelerator.empty_cache() |
| 430 | logger.info( |
| 431 | "Weight cache %s daemon rank %d released cached weights", |
| 432 | self.role, |
| 433 | self.global_rank, |
| 434 | ) |
| 435 | send_msg(conn, {"status": "ok"}) |
| 436 | |
| 437 | |
| 438 | def _run_daemon( |
| 439 | global_rank: int, |
| 440 | local_rank: int, |
| 441 | vllm_config: VllmConfig, |
| 442 | distributed_init_method: str, |
| 443 | socket_dir: str | None, |
| 444 | ready_queue: "multiprocessing.Queue[tuple[str, int]]", |
| 445 | is_draft: bool = False, |
| 446 | dp_rank: int = 0, |
| 447 | pp_rank: int = 0, |
| 448 | ) -> None: |
| 449 | daemon = WeightCacheDaemon( |
| 450 | vllm_config, |
| 451 | global_rank, |
| 452 | local_rank, |
| 453 | distributed_init_method, |
| 454 | socket_dir, |
| 455 | is_draft, |
| 456 | dp_rank, |
| 457 | pp_rank, |
| 458 | ) |
| 459 | daemon.load_model() |
| 460 | daemon.serve_forever( |
| 461 | ready_callback=lambda: ready_queue.put((daemon.role, daemon.global_rank)) |
| 462 | ) |
| 463 | |
| 464 | |
| 465 | def plan_local_ranks( |
| 466 | parallel_config: ParallelConfig, |
| 467 | ) -> list[tuple[int, int, int, int]]: |
| 468 | """``(local_rank, dp_rank, pp_rank, tp_rank)`` for every local daemon. |
| 469 | |
| 470 | Global ranks enumerate DP, PP, then TP: ``global = dp_rank * pp_size * |
| 471 | tp_size + pp_rank * tp_size + tp_rank``. |
| 472 | With ``--nnodes`` the engine hands each node a contiguous block of global |
| 473 | ranks, so node ``r`` serves ``node_rank * local + i``; without it a |
| 474 | launcher's block starts at ``--data-parallel-start-rank * tp_size``. |
| 475 | """ |
| 476 | tp_size = parallel_config.tensor_parallel_size |
| 477 | pp_size = parallel_config.pipeline_parallel_size |
| 478 | world_size = parallel_config.data_parallel_size * pp_size * tp_size |
| 479 | if parallel_config.nnodes > 1: |
| 480 | local_world_size = world_size // parallel_config.nnodes |
| 481 | base = parallel_config.node_rank * local_world_size |
| 482 | else: |
| 483 | local_world_size = parallel_config.data_parallel_size_local * pp_size * tp_size |
| 484 | base = parallel_config.data_parallel_rank * pp_size * tp_size |
| 485 | placements = [] |
| 486 | for local_rank in range(local_world_size): |
| 487 | global_rank = base + local_rank |
| 488 | rank_in_dp = global_rank % (pp_size * tp_size) |
| 489 | dp_rank = global_rank // (pp_size * tp_size) |
| 490 | pp_rank, tp_rank = divmod(rank_in_dp, tp_size) |
| 491 | placements.append((local_rank, dp_rank, pp_rank, tp_rank)) |
| 492 | return placements |
| 493 | |
| 494 | |
| 495 | def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None: |
| 496 | """VllmConfig for the draft daemon group""" |
| 497 | if not is_draft_model_cacheable(vllm_config.speculative_config): |
| 498 | return None |
| 499 | assert vllm_config.speculative_config is not None |
| 500 | return vllm_config.speculative_config.apply_draft_overrides(vllm_config) |
| 501 | |
| 502 | |
| 503 | def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None: |
| 504 | """Reject placements the daemon cannot map.""" |
| 505 | dp_size = parallel_config.data_parallel_size |
| 506 | tp_size = parallel_config.tensor_parallel_size |
| 507 | pp_size = parallel_config.pipeline_parallel_size |
| 508 | if parallel_config.nnodes > 1: |
| 509 | world_size = dp_size * pp_size * tp_size |
| 510 | if world_size % parallel_config.nnodes != 0: |
| 511 | raise ValueError( |
| 512 | f"--nnodes ({parallel_config.nnodes}) must evenly divide the " |
| 513 | f"daemon world size ({world_size} = data-parallel-size " |
| 514 | f"{dp_size} x pipeline-parallel-size {pp_size} x " |
| 515 | f"tensor-parallel-size {tp_size})" |
| 516 | ) |
| 517 | return |
| 518 | if dp_size == 1: |
| 519 | return |
| 520 | end_rank = ( |
| 521 | parallel_config.data_parallel_rank + parallel_config.data_parallel_size_local |
| 522 | ) |
| 523 | if end_rank > dp_size: |
| 524 | raise ValueError( |
| 525 | "--data-parallel-start-rank + --data-parallel-size-local " |
| 526 | f"({end_rank}) exceeds --data-parallel-size ({dp_size})" |
| 527 | ) |