253 lines
9.7 KiB
Python
253 lines
9.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""The `preload` subcommand for the vLLM CLI.
|
|
|
|
Launches one weight cache daemon per GPU; each daemon holds the
|
|
post-quantized, TP-sharded weights of its rank in GPU memory and serves
|
|
CUDA IPC handles to vLLM engines over a Unix domain socket, so restarting
|
|
engines can map the weights via zero-copy IPC instead of reloading from disk:
|
|
|
|
vllm preload --model /path/to/model --tensor-parallel-size 4
|
|
|
|
Engines then load from the daemons with:
|
|
|
|
vllm serve /path/to/model --tensor-parallel-size 4 --load-format ipc_cache
|
|
|
|
Tensor, expert and data parallelism are supported, including across nodes;
|
|
see vllm/model_executor/model_loader/weight_cache/daemon.py for the rank
|
|
layout and rendezvous details.
|
|
"""
|
|
|
|
import argparse
|
|
import multiprocessing
|
|
import queue
|
|
import signal
|
|
import sys
|
|
import typing
|
|
from itertools import product
|
|
|
|
from vllm.engine.arg_utils import EngineArgs
|
|
from vllm.entrypoints.cli.types import CLISubcommand
|
|
from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.model_loader.weight_cache.protocol import (
|
|
check_ipc_platform_support,
|
|
)
|
|
from vllm.utils.network_utils import get_distributed_init_method, get_open_port
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from vllm.config import VllmConfig
|
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
|
else:
|
|
FlexibleArgumentParser = argparse.ArgumentParser
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class PreloadSubcommand(CLISubcommand):
|
|
"""The `preload` subcommand for the vLLM CLI."""
|
|
|
|
name = "preload"
|
|
|
|
@staticmethod
|
|
def add_cli_args(parser: FlexibleArgumentParser) -> None:
|
|
EngineArgs.add_cli_args(parser)
|
|
parser.add_argument(
|
|
"--weight-cache-socket-dir",
|
|
type=str,
|
|
default=None,
|
|
help="Directory for the daemon Unix sockets (default: tempdir).",
|
|
)
|
|
parser.add_argument(
|
|
"--weight-cache-master-port",
|
|
type=int,
|
|
default=None,
|
|
help="Rendezvous port for the daemon's own TP group. Must differ "
|
|
"from the engine's --master-port (the daemon holds its group open "
|
|
"while serving) and match across nodes. Required when --nnodes > "
|
|
"1; defaults to a free port for single-node.",
|
|
)
|
|
parser.add_argument(
|
|
"--weight-cache-draft-master-port",
|
|
type=int,
|
|
default=None,
|
|
help="Rendezvous port for the MTP draft daemon group. Defaults to "
|
|
"--weight-cache-master-port + 1 for multi-node, or a free port "
|
|
"for single-node.",
|
|
)
|
|
|
|
@staticmethod
|
|
def cmd(args: argparse.Namespace) -> None:
|
|
# Imported here so `vllm --help` and friends stay free of torch.
|
|
from vllm.model_executor.model_loader.weight_cache.daemon import (
|
|
_reject_unsupported_parallelism,
|
|
_run_daemon,
|
|
get_draft_daemon_config,
|
|
plan_local_ranks,
|
|
)
|
|
from vllm.model_executor.model_loader.weight_cache.utils import (
|
|
format_daemon_role,
|
|
)
|
|
|
|
engine_args = EngineArgs.from_cli_args(args)
|
|
vllm_config = engine_args.create_engine_config()
|
|
if vllm_config.load_config.load_format == "ipc_cache":
|
|
raise ValueError(
|
|
"The weight cache daemon itself must load from disk; use the "
|
|
"default --load-format"
|
|
)
|
|
# Config-only so it can fail before any model loading; the quant
|
|
# method check needs the created model and runs in
|
|
# WeightCacheDaemon.get_model.
|
|
check_ipc_platform_support()
|
|
parallel_config = vllm_config.parallel_config
|
|
_reject_unsupported_parallelism(parallel_config)
|
|
tp_size = parallel_config.tensor_parallel_size
|
|
dp_size = parallel_config.data_parallel_size
|
|
|
|
placements = plan_local_ranks(parallel_config)
|
|
local_world_size = len(placements)
|
|
if dp_size == 1:
|
|
nnodes = parallel_config.nnodes
|
|
node_rank = parallel_config.node_rank
|
|
master_addr = parallel_config.master_addr
|
|
else:
|
|
master_addr = parallel_config.data_parallel_master_ip
|
|
if parallel_config.nnodes > 1:
|
|
nnodes = parallel_config.nnodes
|
|
node_rank = parallel_config.node_rank
|
|
else:
|
|
nnodes = dp_size // parallel_config.data_parallel_size_local
|
|
node_rank = parallel_config.data_parallel_rank // (
|
|
parallel_config.data_parallel_size_local
|
|
)
|
|
if nnodes > 1 and master_addr in ("127.0.0.1", "localhost"):
|
|
raise ValueError(
|
|
"Data parallelism across nodes requires a reachable "
|
|
"--data-parallel-address for the daemon rendezvous"
|
|
)
|
|
|
|
# The daemon forms its own world group and holds it open while
|
|
# serving, so it needs a rendezvous port distinct from the engine's.
|
|
# All nodes must agree on it; single-node can auto-pick a free port.
|
|
# The master address is the engine's --master-addr (TP across nodes)
|
|
# or --data-parallel-address.
|
|
if nnodes > 1 and args.weight_cache_master_port is None:
|
|
raise ValueError(
|
|
"--weight-cache-master-port is required when the daemons span nodes"
|
|
)
|
|
master_port = args.weight_cache_master_port or get_open_port()
|
|
distributed_init_method = get_distributed_init_method(master_addr, master_port)
|
|
|
|
# (is_draft, vllm_config, rendezvous) per daemon group.
|
|
groups: list[tuple[bool, VllmConfig, str]] = [
|
|
(False, vllm_config, distributed_init_method)
|
|
]
|
|
draft_vllm_config = get_draft_daemon_config(vllm_config)
|
|
if draft_vllm_config is not None:
|
|
draft_master_port = args.weight_cache_draft_master_port or (
|
|
master_port + 1 if nnodes > 1 else get_open_port()
|
|
)
|
|
if draft_master_port == master_port:
|
|
raise ValueError(
|
|
"--weight-cache-draft-master-port must differ from "
|
|
"--weight-cache-master-port"
|
|
)
|
|
groups.append(
|
|
(
|
|
True,
|
|
draft_vllm_config,
|
|
get_distributed_init_method(master_addr, draft_master_port),
|
|
)
|
|
)
|
|
|
|
ctx = multiprocessing.get_context("spawn")
|
|
ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue()
|
|
# Local index == device index; global rank enumerates DP then TP.
|
|
expected_ready = {
|
|
(format_daemon_role(is_draft), dp_rank * tp_size + tp_rank)
|
|
for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements)
|
|
}
|
|
procs = [
|
|
ctx.Process(
|
|
target=_run_daemon,
|
|
args=(
|
|
tp_rank,
|
|
local_rank,
|
|
config,
|
|
init_method,
|
|
args.weight_cache_socket_dir,
|
|
ready_queue,
|
|
is_draft,
|
|
dp_rank,
|
|
),
|
|
name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-"
|
|
f"{dp_rank * tp_size + tp_rank}",
|
|
)
|
|
for (is_draft, config, init_method), (
|
|
local_rank,
|
|
dp_rank,
|
|
tp_rank,
|
|
) in product(groups, placements)
|
|
]
|
|
for proc in procs:
|
|
proc.start()
|
|
|
|
def _shutdown(signum, frame):
|
|
for proc in procs:
|
|
proc.terminate()
|
|
|
|
signal.signal(signal.SIGINT, _shutdown)
|
|
signal.signal(signal.SIGTERM, _shutdown)
|
|
|
|
ready: set[tuple[str, int]] = set()
|
|
while len(ready) < len(expected_ready):
|
|
try:
|
|
ready.add(ready_queue.get(timeout=1.0))
|
|
except queue.Empty:
|
|
dead = [p for p in procs if p.exitcode is not None]
|
|
if dead:
|
|
logger.error(
|
|
"Weight cache daemon rank(s) exited during startup "
|
|
"(exitcodes=%s); shutting down.",
|
|
[p.exitcode for p in dead],
|
|
)
|
|
for proc in procs:
|
|
proc.terminate()
|
|
for proc in procs:
|
|
proc.join()
|
|
sys.exit(max((p.exitcode or 0) for p in procs))
|
|
socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
|
|
logger.info_once(
|
|
"===== Weight cache daemon READY: node %d/%d serving %d local "
|
|
"rank(s) x %d role(s) in %s =====",
|
|
node_rank,
|
|
nnodes,
|
|
local_world_size,
|
|
len(groups),
|
|
socket_dir_msg,
|
|
)
|
|
|
|
for proc in procs:
|
|
proc.join()
|
|
sys.exit(max(proc.exitcode or 0 for proc in procs))
|
|
|
|
def subparser_init(
|
|
self, subparsers: argparse._SubParsersAction
|
|
) -> FlexibleArgumentParser:
|
|
preload_parser = subparsers.add_parser(
|
|
self.name,
|
|
help="Launch weight cache daemons (one per GPU) for fast engine restarts.",
|
|
description="Launch weight cache daemons (one per GPU) that "
|
|
"keep post-quantized weights in GPU memory and serve them to "
|
|
"engines over CUDA IPC.",
|
|
usage="vllm preload --model <model> [options]",
|
|
)
|
|
self.add_cli_args(preload_parser)
|
|
preload_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)
|
|
return preload_parser
|
|
|
|
|
|
def cmd_init() -> list[CLISubcommand]:
|
|
return [PreloadSubcommand()]
|