vllm.distributed.weight_transfer.ipc_engine ¶

@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}"
            )

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.config import set_current_vllm_config
        from vllm.model_executor.model_loader.reload import (
            initialize_layerwise_reload,
        )

        # set_current_vllm_config is required because the layerwise reload path
        # may instantiate kernels (e.g. FlashInfer CUTLASS MoE) whose __init__
        # reads get_current_vllm_config().
        with set_current_vllm_config(self.vllm_config), torch.device(self.device):
            initialize_layerwise_reload(self.model)

    def finish_weight_update(self) -> None:
        """Finalize layerwise reloading after all weights have been received."""
        from vllm.config import set_current_vllm_config
        from vllm.model_executor.model_loader.reload import (
            finalize_layerwise_reload,
        )

        with set_current_vllm_config(self.vllm_config), torch.device(self.device):
            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))

        from vllm.config import set_current_vllm_config

        # set_current_vllm_config is required because model.load_weights may run
        # per-layer postprocessing (process_weights_after_loading) that
        # instantiates kernels reading get_current_vllm_config() — e.g.
        # FlashInfer CUTLASS MoE for unquantized MoE models.
        with set_current_vllm_config(self.vllm_config), torch.device(self.device):
            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()

@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")