wehub-resource-sync
497 lines
19 KiB
Python
497 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""IPC-based weight transfer engine using CUDA IPC for communication."""
|
|
|
|
import pickle
|
|
from collections.abc import Callable, Iterator
|
|
from dataclasses import asdict, dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import pybase64 as base64
|
|
import ray
|
|
import requests
|
|
import torch
|
|
from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
|
|
|
|
from vllm import envs
|
|
from vllm.config.weight_transfer import WeightTransferConfig
|
|
from vllm.distributed.weight_transfer.base import (
|
|
WeightTransferEngine,
|
|
WeightTransferInitInfo,
|
|
WeightTransferUpdateInfo,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config import VllmConfig
|
|
from vllm.distributed.weight_transfer.packed_tensor import (
|
|
DEFAULT_PACKED_BUFFER_SIZE_BYTES,
|
|
packed_ipc_consumer,
|
|
packed_ipc_producer,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class IPCTrainerSendWeightsArgs:
|
|
"""Arguments for IPC trainer_send_weights method."""
|
|
|
|
send_mode: str | Callable[["IPCWeightTransferUpdateInfo"], None]
|
|
"""How to send updates to vLLM. Either a string ('ray' or 'http') for
|
|
built-in transports, or a callable that receives an
|
|
IPCWeightTransferUpdateInfo and performs the send."""
|
|
llm_handle: Any = None
|
|
"""Ray actor handle or list of handles (required for 'ray' send_mode)."""
|
|
url: str | None = None
|
|
"""Base URL for HTTP endpoint (required for 'http' send_mode)."""
|
|
packed: bool = False
|
|
"""Whether to use packed tensor transfer for bounded-memory chunking."""
|
|
packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
|
|
"""Size in bytes for each packed tensor buffer when packed=True."""
|
|
|
|
def __post_init__(self):
|
|
"""Validate that required arguments are provided for the selected mode."""
|
|
if callable(self.send_mode):
|
|
return
|
|
if self.send_mode == "ray" and self.llm_handle is None:
|
|
raise ValueError("llm_handle is required for 'ray' send_mode")
|
|
if self.send_mode == "http" and self.url is None:
|
|
raise ValueError("url is required for 'http' send_mode")
|
|
if self.send_mode not in ("ray", "http"):
|
|
raise ValueError(
|
|
f"send_mode must be 'ray', 'http', or a callable, "
|
|
f"got {self.send_mode!r}"
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class IPCWeightTransferInitInfo(WeightTransferInitInfo):
|
|
"""Initialization info for IPC weight transfer backend. No init needed for IPC."""
|
|
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
|
|
"""Update info for IPC weight transfer backend."""
|
|
|
|
names: list[str]
|
|
dtype_names: list[str]
|
|
shapes: list[list[int]]
|
|
ipc_handles: list[dict[str, tuple]] | dict[str, tuple] | None = None
|
|
"""IPC handles mapping physical GPU UUID to rebuild_cuda_tensor args.
|
|
For non-packed mode: list of per-parameter handle dicts.
|
|
For packed mode: single handle dict for the packed buffer."""
|
|
ipc_handles_pickled: str | None = None
|
|
"""Base64-encoded pickled IPC handles, used for HTTP transport."""
|
|
tensor_sizes: list[int] | None = None
|
|
"""Per-parameter sizes in bytes within the packed buffer.
|
|
Required when packed=True, unused otherwise."""
|
|
packed: bool = False
|
|
"""Whether this update uses packed tensor format."""
|
|
|
|
def __post_init__(self):
|
|
if self.ipc_handles_pickled is not None:
|
|
if self.ipc_handles is not None:
|
|
raise ValueError(
|
|
"Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
|
|
)
|
|
|
|
if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
|
|
raise ValueError(
|
|
"Refusing to deserialize `ipc_handles_pickled` without "
|
|
"VLLM_ALLOW_INSECURE_SERIALIZATION=1"
|
|
)
|
|
|
|
self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled))
|
|
self.ipc_handles_pickled = None
|
|
|
|
if self.ipc_handles is None:
|
|
raise ValueError(
|
|
"Either `ipc_handles` or `ipc_handles_pickled` must be provided"
|
|
)
|
|
num_params = len(self.names)
|
|
if len(self.dtype_names) != num_params:
|
|
raise ValueError(
|
|
f"`dtype_names` should be of the same size as `names`: "
|
|
f"got {len(self.dtype_names)} and {len(self.names)}"
|
|
)
|
|
if len(self.shapes) != num_params:
|
|
raise ValueError(
|
|
f"`shapes` should be of the same size as `names`: "
|
|
f"got {len(self.shapes)} and {len(self.names)}"
|
|
)
|
|
if (
|
|
not self.packed
|
|
and isinstance(self.ipc_handles, list)
|
|
and len(self.ipc_handles) != num_params
|
|
):
|
|
raise ValueError(
|
|
f"`ipc_handles` should be of the same size as `names`: "
|
|
f"got {len(self.ipc_handles)} and {len(self.names)}"
|
|
)
|
|
if self.packed and self.tensor_sizes is None:
|
|
raise ValueError("`tensor_sizes` is required when packed=True")
|
|
|
|
|
|
class IPCWeightTransferEngine(
|
|
WeightTransferEngine[IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo]
|
|
):
|
|
"""
|
|
Weight transfer engine using CUDA IPC for communication between trainer and workers.
|
|
|
|
This implementation uses CUDA IPC to transfer weights from the trainer (rank 0)
|
|
to all inference workers in a process group. IPC handles are used to share
|
|
memory between processes on the same node.
|
|
"""
|
|
|
|
# Define backend-specific dataclass types
|
|
init_info_cls = IPCWeightTransferInitInfo
|
|
update_info_cls = IPCWeightTransferUpdateInfo
|
|
|
|
def __init__(
|
|
self,
|
|
config: WeightTransferConfig,
|
|
vllm_config: "VllmConfig",
|
|
device: torch.device,
|
|
model: torch.nn.Module,
|
|
) -> None:
|
|
"""
|
|
Initialize the IPC weight transfer engine.
|
|
|
|
Args:
|
|
config: The configuration for the weight transfer engine
|
|
vllm_config: The full vLLM config
|
|
device: The device this worker's model lives on
|
|
model: The local model instance which will receive the weights
|
|
"""
|
|
super().__init__(config, vllm_config, device, model)
|
|
|
|
def parse_update_info(
|
|
self, update_dict: dict[str, Any]
|
|
) -> IPCWeightTransferUpdateInfo:
|
|
"""Parse update dict, deserializing pickled IPC handles if present.
|
|
|
|
HTTP transport sends IPC handles as a base64-encoded pickle under the
|
|
key ``ipc_handles_pickled``. This method deserializes them back into
|
|
``ipc_handles`` before constructing the typed dataclass, keeping
|
|
serialization concerns out of the dataclass itself.
|
|
|
|
Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the
|
|
payload is deserialized via ``pickle.loads``.
|
|
"""
|
|
pickled = update_dict.pop("ipc_handles_pickled", None)
|
|
if pickled is not None:
|
|
if update_dict.get("ipc_handles") is not None:
|
|
raise ValueError(
|
|
"Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
|
|
)
|
|
|
|
if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
|
|
raise ValueError(
|
|
"Refusing to deserialize `ipc_handles_pickled` without "
|
|
"VLLM_ALLOW_INSECURE_SERIALIZATION=1"
|
|
)
|
|
|
|
update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
|
|
|
|
return super().parse_update_info(update_dict)
|
|
|
|
def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
|
|
"""
|
|
Initialize the weight transfer mechanism.
|
|
This is called once at the beginning of training.
|
|
No initialization needed for IPC backend.
|
|
|
|
Args:
|
|
init_info: IPC initialization info (empty)
|
|
"""
|
|
pass
|
|
|
|
def start_weight_update(self) -> None:
|
|
"""Initialize layerwise reloading for the incoming checkpoint weights."""
|
|
from vllm.model_executor.model_loader.reload import (
|
|
initialize_layerwise_reload,
|
|
)
|
|
|
|
initialize_layerwise_reload(self.model)
|
|
|
|
def finish_weight_update(self) -> None:
|
|
"""Finalize layerwise reloading after all weights have been received."""
|
|
from vllm.model_executor.model_loader.reload import (
|
|
finalize_layerwise_reload,
|
|
)
|
|
|
|
finalize_layerwise_reload(self.model, self.model_config)
|
|
|
|
def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None:
|
|
"""
|
|
Receive weights from the trainer via CUDA IPC handles and load them.
|
|
|
|
Args:
|
|
update_info: IPC update info containing parameter names, dtypes, shapes,
|
|
and IPC handles. Each IPC handle is a mapping between physical
|
|
GPU UUID and the rebuild_cuda_tensor args tuple.
|
|
"""
|
|
# Use the worker's assigned device rather than the ambient current
|
|
# device: the receive path is no longer wrapped in
|
|
# `with torch.device(self.device)` by the caller, so the current device
|
|
# is not guaranteed to match self.device. The IPC tensors must be
|
|
# rebuilt on the device the model lives on.
|
|
device_index = self.device.index
|
|
|
|
if update_info.packed:
|
|
assert update_info.tensor_sizes is not None
|
|
assert isinstance(update_info.ipc_handles, dict)
|
|
weights = packed_ipc_consumer(
|
|
ipc_handle=update_info.ipc_handles,
|
|
names=update_info.names,
|
|
shapes=update_info.shapes,
|
|
dtype_names=update_info.dtype_names,
|
|
tensor_sizes=update_info.tensor_sizes,
|
|
device_index=device_index,
|
|
)
|
|
else:
|
|
assert isinstance(update_info.ipc_handles, list)
|
|
weights = []
|
|
for name, ipc_handle in zip(
|
|
update_info.names,
|
|
update_info.ipc_handles,
|
|
):
|
|
props = torch.cuda.get_device_properties(device_index)
|
|
physical_gpu_id = str(props.uuid)
|
|
|
|
if physical_gpu_id not in ipc_handle:
|
|
raise ValueError(
|
|
f"IPC handle not found for GPU UUID "
|
|
f"{physical_gpu_id}. "
|
|
f"Available UUIDs: {list(ipc_handle.keys())}"
|
|
)
|
|
|
|
args = ipc_handle[physical_gpu_id]
|
|
list_args = list(args)
|
|
# Index 6 of the args from reduce_tensor is the device_index.
|
|
# We need to overwrite it with the receiver's device index.
|
|
list_args[6] = device_index
|
|
weight = rebuild_cuda_tensor(*list_args)
|
|
weights.append((name, weight))
|
|
|
|
self.model.load_weights(weights)
|
|
|
|
def shutdown(self) -> None:
|
|
pass
|
|
|
|
@staticmethod
|
|
def trainer_send_weights(
|
|
iterator: Iterator[tuple[str, torch.Tensor]],
|
|
trainer_args: dict[str, Any] | IPCTrainerSendWeightsArgs,
|
|
) -> None:
|
|
"""Send weights from trainer to inference workers via CUDA IPC.
|
|
|
|
Supports two transport modes ('ray' and 'http') and two transfer
|
|
strategies:
|
|
- Non-packed (default): all weights in a single API call.
|
|
- Packed (packed=True): chunked transfer with bounded GPU memory.
|
|
|
|
For multi-GPU training, all ranks must call this method in
|
|
parallel. IPC handles are all-gathered across ranks and merged
|
|
so that each vLLM worker can find its own GPU UUID. Only rank 0
|
|
sends the payload to vLLM.
|
|
|
|
.. note::
|
|
This method calls ``update_weights`` internally. The caller must
|
|
call ``start_weight_update`` before and ``finish_weight_update``
|
|
after this method.
|
|
|
|
Args:
|
|
iterator: Iterator of (name, tensor) pairs. For multi-GPU,
|
|
each rank should yield the full tensor on its own GPU
|
|
(e.g. via FSDP full_tensor()).
|
|
trainer_args: IPCTrainerSendWeightsArgs or equivalent dict.
|
|
"""
|
|
args = (
|
|
IPCTrainerSendWeightsArgs(**trainer_args)
|
|
if isinstance(trainer_args, dict)
|
|
else trainer_args
|
|
)
|
|
device_index = torch.accelerator.current_device_index()
|
|
gpu_uuid = str(torch.cuda.get_device_properties(device_index).uuid)
|
|
if args.packed:
|
|
IPCWeightTransferEngine._send_packed(iterator, args, gpu_uuid)
|
|
else:
|
|
IPCWeightTransferEngine._send_unpacked(iterator, args, gpu_uuid)
|
|
|
|
@staticmethod
|
|
def _is_rank_zero() -> bool:
|
|
"""Return True if this is rank 0 or no distributed group exists."""
|
|
if not torch.distributed.is_initialized():
|
|
return True
|
|
return torch.distributed.get_rank() == 0
|
|
|
|
@staticmethod
|
|
def _all_gather_and_merge_handles(
|
|
handles: list[dict[str, tuple]],
|
|
) -> list[dict[str, tuple]]:
|
|
"""All-gather and merge IPC handle dicts across ranks in one call.
|
|
|
|
Each rank contributes a list of {gpu_uuid: ipc_args} dicts (one
|
|
per parameter or one per chunk). A single all_gather_object
|
|
collects every rank's full list, then rank 0 merges per-index so
|
|
each dict maps every GPU UUID to its args.
|
|
|
|
Non-rank-0 returns a list of empty dicts.
|
|
No-op (returns handles unchanged) when no distributed group exists.
|
|
"""
|
|
if (
|
|
not torch.distributed.is_initialized()
|
|
or torch.distributed.get_world_size() == 1
|
|
):
|
|
return handles
|
|
|
|
world_size = torch.distributed.get_world_size()
|
|
gathered: list[list[dict[str, tuple]] | None] = [None] * world_size
|
|
torch.distributed.all_gather_object(gathered, handles)
|
|
torch.distributed.barrier()
|
|
torch.cuda.synchronize()
|
|
|
|
if torch.distributed.get_rank() == 0:
|
|
merged: list[dict[str, tuple]] = []
|
|
for param_idx in range(len(handles)):
|
|
m: dict[str, tuple] = {}
|
|
for rank_handles in gathered:
|
|
if rank_handles is not None:
|
|
m.update(rank_handles[param_idx])
|
|
merged.append(m)
|
|
return merged
|
|
return [{} for _ in handles]
|
|
|
|
@staticmethod
|
|
def _post_send_sync() -> None:
|
|
"""Barrier + ipc_collect after a send; no-op if single-GPU."""
|
|
if (
|
|
torch.distributed.is_initialized()
|
|
and torch.distributed.get_world_size() > 1
|
|
):
|
|
torch.distributed.barrier()
|
|
torch.cuda.ipc_collect()
|
|
|
|
@staticmethod
|
|
def _send_unpacked(
|
|
iterator: Iterator[tuple[str, torch.Tensor]],
|
|
args: IPCTrainerSendWeightsArgs,
|
|
gpu_uuid: str,
|
|
) -> None:
|
|
"""Send all weights in a single API call (non-packed mode)."""
|
|
names: list[str] = []
|
|
dtype_names: list[str] = []
|
|
shapes: list[list[int]] = []
|
|
ipc_handles: list[dict[str, tuple]] = []
|
|
# Hold strong refs to every contiguous copy until the send + post-send
|
|
# sync completes. reduce_tensor's returned args do NOT keep storage
|
|
# alive, and non-contiguous inputs allocate fresh storage in
|
|
# .contiguous() that would otherwise be GC'd before the consumer opens
|
|
# the IPC handle.
|
|
weight_refs: list[torch.Tensor] = []
|
|
|
|
for name, tensor in iterator:
|
|
names.append(name)
|
|
dtype_names.append(str(tensor.dtype).split(".")[-1])
|
|
shapes.append(list(tensor.shape))
|
|
|
|
weight = tensor.detach().contiguous()
|
|
weight_refs.append(weight)
|
|
_, ipc_args = reduce_tensor(weight)
|
|
ipc_handles.append({gpu_uuid: ipc_args})
|
|
|
|
ipc_handles = IPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles)
|
|
|
|
if IPCWeightTransferEngine._is_rank_zero():
|
|
IPCWeightTransferEngine._do_send(
|
|
args=args,
|
|
names=names,
|
|
dtype_names=dtype_names,
|
|
shapes=shapes,
|
|
ipc_handles=ipc_handles,
|
|
)
|
|
|
|
IPCWeightTransferEngine._post_send_sync()
|
|
|
|
@staticmethod
|
|
def _send_packed(
|
|
iterator: Iterator[tuple[str, torch.Tensor]],
|
|
args: IPCTrainerSendWeightsArgs,
|
|
gpu_uuid: str,
|
|
) -> None:
|
|
"""Send weights in bounded-memory chunks (packed mode)."""
|
|
post_iter_func: Callable = lambda item: item[1]
|
|
|
|
for chunk in packed_ipc_producer(
|
|
iterator=iterator,
|
|
gpu_uuid=gpu_uuid,
|
|
post_iter_func=post_iter_func,
|
|
buffer_size_bytes=args.packed_buffer_size_bytes,
|
|
):
|
|
ipc_handle = IPCWeightTransferEngine._all_gather_and_merge_handles(
|
|
[chunk.ipc_handle]
|
|
)[0]
|
|
|
|
if IPCWeightTransferEngine._is_rank_zero():
|
|
IPCWeightTransferEngine._do_send(
|
|
args=args,
|
|
names=chunk.names,
|
|
dtype_names=chunk.dtype_names,
|
|
shapes=chunk.shapes,
|
|
ipc_handles=ipc_handle,
|
|
tensor_sizes=chunk.tensor_sizes,
|
|
packed=True,
|
|
)
|
|
|
|
IPCWeightTransferEngine._post_send_sync()
|
|
|
|
@staticmethod
|
|
def _do_send(
|
|
args: IPCTrainerSendWeightsArgs,
|
|
names: list[str],
|
|
dtype_names: list[str],
|
|
shapes: list[list[int]],
|
|
ipc_handles: list[dict[str, tuple]] | dict[str, tuple],
|
|
tensor_sizes: list[int] | None = None,
|
|
packed: bool = False,
|
|
) -> None:
|
|
"""Send a single update payload via the configured transport."""
|
|
update_fields: dict[str, Any] = {
|
|
"names": names,
|
|
"dtype_names": dtype_names,
|
|
"shapes": shapes,
|
|
"packed": packed,
|
|
}
|
|
if tensor_sizes is not None:
|
|
update_fields["tensor_sizes"] = tensor_sizes
|
|
|
|
update_fields["ipc_handles"] = ipc_handles
|
|
update_info = IPCWeightTransferUpdateInfo(**update_fields)
|
|
|
|
if callable(args.send_mode):
|
|
args.send_mode(update_info)
|
|
elif args.send_mode == "ray":
|
|
handles = (
|
|
args.llm_handle
|
|
if isinstance(args.llm_handle, list)
|
|
else [args.llm_handle]
|
|
)
|
|
ray.get(
|
|
[
|
|
h.update_weights.remote(dict(update_info=asdict(update_info)))
|
|
for h in handles
|
|
]
|
|
)
|
|
elif args.send_mode == "http":
|
|
pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode(
|
|
"utf-8"
|
|
)
|
|
http_fields = {k: v for k, v in update_fields.items() if k != "ipc_handles"}
|
|
http_fields["ipc_handles_pickled"] = pickled_handles
|
|
|
|
url = f"{args.url}/update_weights"
|
|
payload = {"update_info": http_fields}
|
|
response = requests.post(url, json=payload, timeout=300)
|
|
response.raise_for_status()
|