406 lines
16 KiB
Python
406 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""NPU IPC-based weight transfer engine using Ascend IPC for communication."""
|
|
|
|
import os
|
|
import pickle
|
|
import socket
|
|
from collections.abc import Callable, Iterator
|
|
from dataclasses import asdict, dataclass
|
|
from functools import lru_cache
|
|
from typing import Any
|
|
|
|
import pybase64 as base64
|
|
import requests
|
|
import torch
|
|
from torch.multiprocessing.reductions import reduce_tensor
|
|
from vllm import envs
|
|
from vllm.config.parallel import ParallelConfig
|
|
from vllm.config.weight_transfer import WeightTransferConfig
|
|
from vllm.distributed.weight_transfer.base import (
|
|
WeightTransferEngine,
|
|
WeightTransferInitInfo,
|
|
)
|
|
from vllm.distributed.weight_transfer.ipc_engine import (
|
|
IPCTrainerSendWeightsArgs,
|
|
IPCWeightTransferUpdateInfo,
|
|
)
|
|
|
|
from vllm_ascend.distributed.weight_transfer.packed_tensor import (
|
|
packed_npu_ipc_consumer,
|
|
packed_npu_ipc_producer,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class NPUIPCTrainerSendWeightsArgs(IPCTrainerSendWeightsArgs):
|
|
"""NPU IPC variant — inherits all fields and validation from the CUDA IPC
|
|
base class. Only the ``send_mode`` callable type is widened to accept the
|
|
NPU update-info type."""
|
|
|
|
send_mode: str | Callable[["NPUIPCWeightTransferUpdateInfo"], None]
|
|
|
|
|
|
@dataclass
|
|
class NPUIPCWeightTransferInitInfo(WeightTransferInitInfo):
|
|
"""Initialization info for NPU IPC weight transfer backend.
|
|
|
|
No initialization needed for NPU IPC.
|
|
"""
|
|
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class NPUIPCWeightTransferUpdateInfo(IPCWeightTransferUpdateInfo):
|
|
"""NPU IPC variant — inherits all fields and validation from the CUDA IPC
|
|
base class. No overrides needed; the field types and ``__post_init__`` are
|
|
identical."""
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_ip() -> str:
|
|
try:
|
|
# try to get ip from network interface
|
|
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
|
|
s.connect(("8.8.8.8", 80))
|
|
return s.getsockname()[0]
|
|
except Exception: # noqa: BLE001
|
|
# fallback to get ip from hostname
|
|
return socket.gethostbyname(socket.gethostname())
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def npu_generate_uuid() -> str:
|
|
"""Generate a unique identifier for the current process's physical NPU chip.
|
|
|
|
Returns ``{host_ip}-{physical_chip_id}`` where ``host_ip`` is the local
|
|
machine's IP address and ``physical_chip_id`` is derived from the current
|
|
logical device index mapped through ``ASCEND_RT_VISIBLE_DEVICES``.
|
|
|
|
On Ascend NPU, ``torch.accelerator.current_device_index()`` returns the
|
|
*logical* device index. When ``ASCEND_RT_VISIBLE_DEVICES`` is set, it
|
|
maps logical indices to physical chip IDs (e.g., ``ASCEND_RT_VISIBLE_DEVICES=2,3``
|
|
means logical device 0 → physical chip 2, logical device 1 → physical chip 3).
|
|
If the env var is not set, the logical index is used directly as the
|
|
physical chip ID (identity mapping).
|
|
|
|
The result is cached because it is constant for the lifetime of the
|
|
process. Both the trainer and inference worker processes co-located
|
|
on the same physical NPU chip will produce the same UUID, which is
|
|
required for NPU IPC handle matching.
|
|
"""
|
|
logical_device = torch.accelerator.current_device_index()
|
|
visible_devices = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", None)
|
|
if visible_devices:
|
|
physical_device = int(visible_devices.split(",")[logical_device].strip())
|
|
else:
|
|
physical_device = logical_device
|
|
return f"{get_ip()}-{physical_device}"
|
|
|
|
|
|
class NPUIPCWeightTransferEngine(WeightTransferEngine[NPUIPCWeightTransferInitInfo, NPUIPCWeightTransferUpdateInfo]):
|
|
"""
|
|
Weight transfer engine using NPU IPC for communication between
|
|
trainer and workers.
|
|
|
|
This implementation uses Ascend NPU IPC to transfer weights from the
|
|
trainer (rank 0) to all inference workers. IPC handles are used to
|
|
share memory between processes on the same node.
|
|
|
|
Requires ``torch_npu`` to be imported (which patches
|
|
``torch.multiprocessing.reductions.reduce_tensor`` to support
|
|
NPU tensors via ``_share_npu_()`` / ``rebuild_npu_tensor``).
|
|
"""
|
|
|
|
init_info_cls = NPUIPCWeightTransferInitInfo
|
|
update_info_cls = NPUIPCWeightTransferUpdateInfo
|
|
|
|
def __init__(
|
|
self,
|
|
config: WeightTransferConfig,
|
|
parallel_config: ParallelConfig,
|
|
model: torch.nn.Module | None = None,
|
|
) -> None:
|
|
super().__init__(config, parallel_config, model)
|
|
|
|
def parse_update_info(self, update_dict: dict[str, Any]) -> NPUIPCWeightTransferUpdateInfo:
|
|
"""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``.
|
|
"""
|
|
if "ipc_handles_pickled" in update_dict:
|
|
if "ipc_handles" in update_dict:
|
|
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"
|
|
)
|
|
|
|
pickled = update_dict.pop("ipc_handles_pickled")
|
|
update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
|
|
|
|
return super().parse_update_info(update_dict)
|
|
|
|
def init_transfer_engine(self, init_info: NPUIPCWeightTransferInitInfo) -> None:
|
|
"""No initialization needed for NPU IPC backend."""
|
|
pass
|
|
|
|
def receive_weights(
|
|
self,
|
|
update_info: NPUIPCWeightTransferUpdateInfo,
|
|
load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
|
|
) -> None:
|
|
"""Receive weights from the trainer via NPU IPC handles.
|
|
|
|
Args:
|
|
update_info: NPU IPC update info containing parameter names,
|
|
dtypes, shapes, and IPC handles.
|
|
load_weights: Callable that loads weights into the model.
|
|
"""
|
|
device_index = torch.accelerator.current_device_index()
|
|
physical_npu_id = npu_generate_uuid()
|
|
|
|
if update_info.packed:
|
|
assert update_info.tensor_sizes is not None
|
|
assert isinstance(update_info.ipc_handles, dict)
|
|
weights = packed_npu_ipc_consumer(
|
|
ipc_handle=update_info.ipc_handles,
|
|
physical_npu_id=physical_npu_id,
|
|
names=update_info.names,
|
|
shapes=update_info.shapes,
|
|
dtype_names=update_info.dtype_names,
|
|
tensor_sizes=update_info.tensor_sizes,
|
|
device_index=device_index,
|
|
)
|
|
load_weights(weights)
|
|
else:
|
|
# Lazy import: ``rebuild_npu_tensor`` lives in ``torch_npu`` and
|
|
# must not be imported at module load time on non-NPU hosts.
|
|
from torch_npu.multiprocessing.reductions import rebuild_npu_tensor
|
|
|
|
assert isinstance(update_info.ipc_handles, list)
|
|
weights = []
|
|
for name, ipc_handle in zip(
|
|
update_info.names,
|
|
update_info.ipc_handles,
|
|
):
|
|
if physical_npu_id not in ipc_handle:
|
|
raise ValueError(
|
|
f"IPC handle not found for NPU UUID {physical_npu_id}. "
|
|
f"Available UUIDs: {list(ipc_handle.keys())}. "
|
|
f"This may indicate that the trainer and worker are "
|
|
f"not co-located on the same physical NPU (node)."
|
|
)
|
|
|
|
args = ipc_handle[physical_npu_id]
|
|
list_args = list(args)
|
|
# Index 6 is the device_index parameter in torch's
|
|
# IPC handle tuple (rebuild_npu_tensor). Update it
|
|
# to the current device since the logical index can
|
|
# differ between sender and receiver.
|
|
list_args[6] = device_index
|
|
weight = rebuild_npu_tensor(*list_args)
|
|
weights.append((name, weight))
|
|
|
|
load_weights(weights)
|
|
|
|
def shutdown(self) -> None:
|
|
pass
|
|
|
|
@staticmethod
|
|
def trainer_send_weights(
|
|
iterator: Iterator[tuple[str, torch.Tensor]],
|
|
trainer_args: dict[str, Any] | NPUIPCTrainerSendWeightsArgs,
|
|
) -> None:
|
|
"""Send weights from trainer to inference workers via NPU 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 NPU memory.
|
|
|
|
For multi-NPU 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 NPU UUID. Only rank 0
|
|
sends the payload to vLLM.
|
|
|
|
.. note::
|
|
This method calls ``update_weights`` internally. The caller must
|
|
handle ``pause`` / ``start_weight_update`` / ``finish_weight_update``
|
|
/ ``resume`` before and after this method.
|
|
|
|
Args:
|
|
iterator: Iterator of (name, tensor) pairs.
|
|
trainer_args: NPUIPCTrainerSendWeightsArgs or equivalent dict.
|
|
"""
|
|
args = NPUIPCTrainerSendWeightsArgs(**trainer_args) if isinstance(trainer_args, dict) else trainer_args
|
|
npu_uuid = npu_generate_uuid()
|
|
if args.packed:
|
|
NPUIPCWeightTransferEngine._send_packed(iterator, args, npu_uuid)
|
|
else:
|
|
NPUIPCWeightTransferEngine._send_unpacked(iterator, args, npu_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.
|
|
|
|
Each rank contributes a list of ``{npu_uuid: ipc_args}`` dicts.
|
|
Rank 0 collects and merges per-index; other ranks receive a list
|
|
of empty dicts. No-op 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.npu.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 + synchronize after a send; no-op if single-NPU."""
|
|
if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:
|
|
torch.distributed.barrier()
|
|
torch.npu.synchronize()
|
|
|
|
@staticmethod
|
|
def _send_unpacked(
|
|
iterator: Iterator[tuple[str, torch.Tensor]],
|
|
args: NPUIPCTrainerSendWeightsArgs,
|
|
npu_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.
|
|
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)
|
|
# Store only the rebuild args (drop the func); the consumer rebuilds
|
|
# with the well-known ``rebuild_npu_tensor``, mirroring upstream's
|
|
# CUDA IPC engine.
|
|
_, ipc_args = reduce_tensor(weight)
|
|
ipc_handles.append({npu_uuid: ipc_args})
|
|
|
|
ipc_handles = NPUIPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles)
|
|
|
|
if NPUIPCWeightTransferEngine._is_rank_zero():
|
|
NPUIPCWeightTransferEngine._do_send(
|
|
args=args,
|
|
names=names,
|
|
dtype_names=dtype_names,
|
|
shapes=shapes,
|
|
ipc_handles=ipc_handles,
|
|
)
|
|
|
|
NPUIPCWeightTransferEngine._post_send_sync()
|
|
|
|
@staticmethod
|
|
def _send_packed(
|
|
iterator: Iterator[tuple[str, torch.Tensor]],
|
|
args: NPUIPCTrainerSendWeightsArgs,
|
|
npu_uuid: str,
|
|
) -> None:
|
|
"""Send weights in bounded-memory chunks (packed mode)."""
|
|
post_iter_func: Callable = lambda item: item[1]
|
|
|
|
for chunk in packed_npu_ipc_producer(
|
|
iterator=iterator,
|
|
npu_uuid=npu_uuid,
|
|
post_iter_func=post_iter_func,
|
|
buffer_size_bytes=args.packed_buffer_size_bytes,
|
|
):
|
|
ipc_handle = NPUIPCWeightTransferEngine._all_gather_and_merge_handles([chunk["ipc_handle"]])[0]
|
|
|
|
if NPUIPCWeightTransferEngine._is_rank_zero():
|
|
NPUIPCWeightTransferEngine._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,
|
|
)
|
|
|
|
NPUIPCWeightTransferEngine._post_send_sync()
|
|
|
|
@staticmethod
|
|
def _do_send(
|
|
args: NPUIPCTrainerSendWeightsArgs,
|
|
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 = NPUIPCWeightTransferUpdateInfo(**update_fields)
|
|
|
|
if callable(args.send_mode):
|
|
args.send_mode(update_info)
|
|
elif args.send_mode == "ray":
|
|
import 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()
|