diff --git a/README.md b/README.md index a656c61..df64369 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,28 @@ Use the flexible P2P implementation, notice this will install `mooncake-transfer pip install 'checkpoint-engine[p2p]' ``` +### Intel XPU (build from source) + +Intel XPU is supported for the **broadcast** update path. The cross-process weight handoff uses a native SYCL `ipc_memory` extension (the XPU counterpart of CUDA IPC) that is JIT-compiled at runtime. Install from source since XPU support is not yet in the released package. P2P is not supported on XPU (Mooncake has no Level Zero backend for XPU device memory). + +Requirements: +- Intel XPU build of PyTorch — `torch.xpu.is_available()` returns `True`, `torch>=2.9` (for the device `.uuid` property; the SYCL extension build also needs `torch>=2.7`). See the [PyTorch XPU install guide](https://pytorch.org/docs/stable/notes/get_start_xpu.html); this is not the default PyPI `torch`. +- Intel oneAPI 2026.0+ providing the `icpx` compiler with SYCL IPC memory support. Needed at runtime (first weight update), not at `pip install` time. + +```Bash +git clone https://github.com/MoonshotAI/checkpoint-engine.git +cd checkpoint-engine +pip install -e . # no [p2p] extra on XPU +``` + +Make `icpx` discoverable in the runtime environment, either by sourcing oneAPI (`source /opt/intel/oneapi/setvars.sh`) or by setting `CMPLR_ROOT`. If neither is set, `icpx` is auto-detected under `/opt/intel/oneapi/compiler/*/bin` and then on `PATH`. The extension then builds automatically on first use; `ParameterServer` also prebuilds it at startup so the one-time compile stays out of the update window. + +Verify the build and the IPC path on the target machine: + +```Bash +pytest tests/test_xpu_ipc.py # hardware-gated; skipped without an Intel GPU + buildable extension +``` + ## Getting Started Prepare an H800 or H20 machine with 8 GPUs with vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights. vLLM version `v0.10.2` is fully tested and recommended. diff --git a/checkpoint_engine/device_utils.py b/checkpoint_engine/device_utils.py index be5b81f..72d76d4 100644 --- a/checkpoint_engine/device_utils.py +++ b/checkpoint_engine/device_utils.py @@ -1,4 +1,5 @@ import ctypes +import gc import os import re import socket @@ -209,9 +210,20 @@ def _is_torch_npu_available(self) -> bool: except ImportError: return False + def _is_torch_xpu_available(self) -> bool: + try: + if hasattr(torch, "xpu") and callable(getattr(torch.xpu, "is_available", None)): + return torch.xpu.is_available() + else: + return False + except Exception: # noqa: BLE001 + return False + def _detect_device_type(self) -> str: if self._is_torch_npu_available(): return "npu" + elif self._is_torch_xpu_available(): + return "xpu" elif torch.cuda.is_available(): return "cuda" else: @@ -222,6 +234,8 @@ def _setup_device_module(self): import torch_npu self.device_module = torch_npu.npu + elif self.device_type == "xpu": + self.device_module = torch.xpu elif self.device_type == "cuda": self.device_module = torch.cuda else: @@ -231,6 +245,8 @@ def _setup_device_module(self): def backend(self) -> str: if self.device_type == "npu": return "hccl" + elif self.device_type == "xpu": + return "xccl" elif self.device_type == "cuda": return "nccl" else: @@ -240,7 +256,7 @@ def backend(self) -> str: def transfer_engine_protocol(self) -> str: if self.device_type == "npu": return "ascend_direct" - elif self.device_type == "cuda": + elif self.device_type in ("cuda", "xpu"): if has_efa_pci(): return "efa" else: @@ -255,3 +271,42 @@ def rdma_device(self, rank: int) -> str: return _get_my_rdma_device(rank, self.device_module.device_count(), _get_rdma_devices()) else: raise TypeError("The current transfer engine protocol is not supported") + + def ipc_collect(self) -> None: + """Reclaim memory held by stale IPC handles.""" + if self.device_type in ("cuda", "npu"): + self.device_module.ipc_collect() + elif self.device_type == "xpu": + # SYCL ipc_memory frees on close_handle; there is no cache to collect. + pass + else: + raise TypeError("The current device type is not supported") + + def supports_inplace_pin(self) -> bool: + """Whether in-place host-memory pinning (cudaHostRegister) is available -- CUDA only.""" + return self.device_type == "cuda" + + def supports_device_ipc(self) -> bool: + """Whether cross-process IPC of *device* tensors works for this backend. + + CUDA/NPU use ``torch.multiprocessing.reductions``; XPU uses the native SYCL + ``ipc_memory`` extension, available when it can be built (oneAPI >= 2026.0). + """ + if self.device_type in ("cuda", "npu"): + return True + if self.device_type == "xpu": + from checkpoint_engine import xpu_ipc + + return xpu_ipc.is_available() + return False + + def supports_device_p2p(self) -> bool: + """Whether P2P (Mooncake) transfer of *device* memory works for this backend (CUDA/NPU only).""" + return self.device_type in ("cuda", "npu") + + def host_empty_cache(self) -> None: + """Release cached pinned host memory (``_host_emptyCache`` on CUDA; else ``gc.collect``).""" + if self.device_type == "cuda": + torch._C._host_emptyCache() + else: + gc.collect() diff --git a/checkpoint_engine/distributed/base.py b/checkpoint_engine/distributed/base.py index 4299394..ea28af7 100644 --- a/checkpoint_engine/distributed/base.py +++ b/checkpoint_engine/distributed/base.py @@ -229,7 +229,12 @@ def use_backend(backend: str | None): "vllm_hccl": ".vllm_hccl.DistributedHccl", } if backend not in mapping: - raise ValueError(f"Unsupported custom backend: {backend}") + # XPU has no custom backend; leave custom_dist unset to use the default xccl TorchBackend. + raise ValueError( + f"Unsupported custom backend: {backend}. " + f"Supported custom backends: {sorted(mapping)}. " + "XPU is not supported here; leave custom_dist unset to use the default xccl backend." + ) module_path, class_name = mapping[backend].rsplit(".", 1) module = importlib.import_module(module_path, "checkpoint_engine.distributed") diff --git a/checkpoint_engine/ipc_handler.py b/checkpoint_engine/ipc_handler.py new file mode 100644 index 0000000..4806997 --- /dev/null +++ b/checkpoint_engine/ipc_handler.py @@ -0,0 +1,134 @@ +"""Pluggable IPC-handle exchange between the ParameterServer and the worker. + +The broadcast path shares a device buffer with the colocated worker. Nothing here +copies or moves the buffer: it only exchanges the IPC handle that lets the worker +map the same device memory. CUDA/NPU use :class:`TorchIPCHandler` +(``torch.multiprocessing`` CUDA IPC, wire-format unchanged); XPU uses +:class:`XpuIPCHandler` (native SYCL ``ipc_memory``). The handle is always a +picklable, self-contained value, so the producer's ``export`` -> ZMQ +``send_pyobj`` -> consumer ``attach`` flow is identical for both; each side calls +``detach`` on cleanup. +""" + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +import torch +from loguru import logger +from torch.multiprocessing.reductions import reduce_tensor + + +if TYPE_CHECKING: + from collections.abc import Callable + + from typing_extensions import Self + + from checkpoint_engine.device_utils import DeviceManager + + +def _rebuild_ipc(handle: tuple["Callable", tuple], device_id: int | None = None) -> torch.Tensor: + func, args = handle + list_args = list(args) + if device_id is not None: + # the key is to change device id to the current device id + # in case two processes have different CUDA_VISIBLE_DEVICES + list_args[6] = device_id + return func(*list_args) + + +class IPCHandler(ABC): + """Hands an IPC handle for a device buffer from the producer (ps) to the consumer (worker).""" + + @abstractmethod + def export(self, buffer: torch.Tensor) -> Any: + """Producer: return the picklable handle to send over ZMQ.""" + + @abstractmethod + def attach(self, handle: Any, device_id: int) -> torch.Tensor: + """Consumer: reconstruct the shared device buffer from ``handle``.""" + + def detach(self) -> None: + """Release IPC resources on either side. No-op by default.""" + + # Used as a context manager so the handle is always released, without the + # caller needing its own try/finally. + def __enter__(self) -> "Self": + return self + + def __exit__(self, *exc_info: object) -> None: + self.detach() + + +class TorchIPCHandler(IPCHandler): + """CUDA/NPU zero-copy handoff via torch.multiprocessing CUDA IPC (unchanged).""" + + def export(self, buffer: torch.Tensor) -> Any: + return reduce_tensor(buffer) + + def attach(self, handle: Any, device_id: int) -> torch.Tensor: + assert isinstance(handle, tuple), f"expected reduce_tensor tuple, got {type(handle)}" + buffer = _rebuild_ipc(handle, device_id) + assert buffer.dtype == torch.uint8 + return buffer + + +class XpuIPCHandler(IPCHandler): + """Intel XPU zero-copy handoff via native SYCL ``ipc_memory`` (portable byte blob).""" + + kind = "xpu_sycl" # wire tag: identifies this handle format to the consumer + + def __init__(self) -> None: + self._opened_ptr: int | None = None # consumer: the mapping to unmap + self._exported_ptr: int | None = None # producer: the retained exporter handle + + def export(self, buffer: torch.Tensor) -> Any: + from checkpoint_engine import xpu_ipc + + ptr = buffer.data_ptr() + handle_bytes = xpu_ipc.get_handle(ptr) + # Release only in detach(): freeing before the consumer opens can drop the + # fd under the level-zero-v2 UR adapter. + self._exported_ptr = ptr + return { + "kind": self.kind, + "handle_bytes": handle_bytes, + "nbytes": buffer.nbytes, + } + + def attach(self, handle: Any, device_id: int) -> torch.Tensor: + from checkpoint_engine import xpu_ipc + + assert isinstance(handle, dict) and handle.get("kind") == self.kind, ( + f"expected {self.kind} handle dict, got {type(handle)}" + ) + ptr = xpu_ipc.open_handle(handle["handle_bytes"], device_id) + self._opened_ptr = ptr + buffer = xpu_ipc.wrap_tensor(ptr, handle["nbytes"], device_id) + assert buffer.dtype == torch.uint8 + return buffer + + def detach(self) -> None: + # Consumer unmaps its opened pointer; producer releases its exported handle. + # At most one of the two is set on any given instance. + from checkpoint_engine import xpu_ipc + + if self._opened_ptr is not None: + try: + torch.xpu.synchronize() # no in-flight reads before unmapping + xpu_ipc.close_handle(self._opened_ptr) + except Exception as e: # noqa: BLE001 + logger.debug(f"xpu ipc close_handle failed during detach: {e}") + self._opened_ptr = None + if self._exported_ptr is not None: + try: + xpu_ipc.release_handle(self._exported_ptr) + except Exception as e: # noqa: BLE001 + logger.debug(f"xpu ipc release_handle failed during detach: {e}") + self._exported_ptr = None + + +def build_ipc_handler(device_manager: "DeviceManager") -> IPCHandler: + """Select the IPC handler for the current device backend.""" + if device_manager.device_type == "xpu": + return XpuIPCHandler() + return TorchIPCHandler() diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py index 1d8c5cf..d9cfe25 100644 --- a/checkpoint_engine/ps.py +++ b/checkpoint_engine/ps.py @@ -11,7 +11,6 @@ import torch.distributed import zmq from loguru import logger -from torch.multiprocessing.reductions import reduce_tensor import checkpoint_engine.distributed as dist from checkpoint_engine.data_types import ( @@ -24,6 +23,7 @@ ParameterMeta, ) from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid +from checkpoint_engine.ipc_handler import build_ipc_handler from checkpoint_engine.p2p_store import P2PStore from checkpoint_engine.pin_memory import _ALIGN_SIZE, _register_checkpoint @@ -53,7 +53,14 @@ def _get_physical_gpu_id(device_manager: DeviceManager, device_index: int | None if device_manager.device_type == "npu": return f"NPU-{npu_generate_uuid()}" else: - return f"GPU-{device_manager.device_module.get_device_properties(device_index).uuid!s}" + # CUDA and XPU both expose get_device_properties(idx).uuid. + props = device_manager.device_module.get_device_properties(device_index) + if not hasattr(props, "uuid"): + raise ValueError( + f"{device_manager.device_type} device properties do not expose a 'uuid' " + f"attribute; a newer PyTorch is required (xpu .uuid needs torch>=2.9)" + ) + return f"GPU-{props.uuid!s}" except AssertionError as e: raise ValueError(f"fail to get physical gpu id {device_index}") from e @@ -223,15 +230,39 @@ def __init__( # NPU transfer engine initialization requires prior set_device. device_index = self._local_rank self.device_manager.device_module.set_device(device_index) - try: - self._p2p_store = P2PStore(self.device_manager) - except ImportError as e: - logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}") + # P2P (Mooncake) transfer is only supported on CUDA/NPU; XPU has no Level Zero + # backend for device memory. Skip the store entirely on unsupported backends + # rather than eagerly initializing something that can never be used (and whose + # engine.initialize() may fail with more than just ImportError). + if self.device_manager.supports_device_p2p(): + try: + self._p2p_store = P2PStore(self.device_manager) + except ImportError as e: + logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}") + self._p2p_store = None + else: + logger.info( + f"[rank{self._rank}] p2p store disabled: not supported on device type " + f"'{self.device_manager.device_type}'" + ) self._p2p_store = None self._device_uuid = _get_physical_gpu_id(self.device_manager, device_index) self._rdma_device = None if self._p2p_store is None else self._p2p_store.device + # Build the JIT SYCL IPC extension now, so its multi-second compile is outside + # the first weight-update window. + if self.device_manager.device_type == "xpu": + from checkpoint_engine import xpu_ipc + + if xpu_ipc.prewarm(): + logger.info(f"[rank{self._rank}] XPU SYCL ipc_memory extension prebuilt") + else: + logger.warning( + f"[rank{self._rank}] XPU SYCL ipc_memory extension unavailable at init; " + "weight updates will fail until it can be built" + ) + master_addr = master_addr or os.getenv("MASTER_ADDR") assert master_addr, "master_addr is required" self._store = torch.distributed.TCPStore( @@ -297,7 +328,7 @@ def register_checkpoint( use_inplace_pin_memory: If True (default), allows inplace pin memory for /dev/shm/ safetensors files. This option is ignored when ``use_shared_memory_pool`` is True. """ - if self.device_manager.device_type != "cuda" and use_inplace_pin_memory: + if not self.device_manager.supports_inplace_pin() and use_inplace_pin_memory: logger.warning( f"[rank{self._rank}] Only cuda devices support in-place pin memory, set use_inplace_pin_memory to False" ) @@ -425,14 +456,8 @@ def _unpin(t: torch.Tensor): # we won't delete the memory pool if unpinning fails. del self._memory_pool[checkpoint_name] # see https://github.com/pytorch/pytorch/blob/31d5c675394705f8a6bc767f80ae14bf4f01246b/torch/csrc/cuda/Module.cpp#L2018 - # this works by using torch>=2.5.0 - if self.device_manager.device_type == "cuda": - torch._C._host_emptyCache() - else: - # torch._C._host_emptyCache() is not supported on NPU, so we call gc.collect() to empty host cache. - import gc - - gc.collect() + # this works by using torch>=2.5.0 on cuda; NPU/XPU fall back to gc.collect(). + self.device_manager.host_empty_cache() def gather_metas(self, checkpoint_name: str): """ @@ -730,12 +755,33 @@ def _update_per_bucket( assert len(self._current_global_parameter_metas) != 0, "parameter metas is empty" assert dist.is_initialized(), "process group is not initialized" + # The broadcast shares a device buffer with the colocated worker via cross-process + # device-tensor IPC (CUDA/NPU: torch.multiprocessing; XPU: native SYCL ipc_memory). + # Fail loudly here rather than with an opaque "_share_fd_: only available on CPU" + # deeper in the update. + if not self.device_manager.supports_device_ipc(): + raise RuntimeError( + f"[rank{self._rank}] weight update requires cross-process device-tensor IPC, which " + f"is not available for device type '{self.device_manager.device_type}' in this " + f"environment. On XPU this needs the native SYCL ipc_memory extension " + f"(an oneAPI 'icpx' compiler with SYCL ipc_memory support, i.e. oneAPI >= 2026.0)." + ) + p2p_update = False # if both ranks is None or [], it will use fully broadcast to update to all ranks if not ranks: logger.info(f"[rank{self._rank}] update checkpoint {checkpoint_name}") # if ranks is set, it will use p2p to update to the ranks else: + # The P2P path RDMA-transfers device memory via Mooncake, which has no Level Zero + # backend for XPU. Reject it clearly rather than failing inside p2p_store registration. + if not self.device_manager.supports_device_p2p(): + raise RuntimeError( + f"[rank{self._rank}] P2P weight update (ranks={ranks}) is not supported on " + f"device type '{self.device_manager.device_type}': the Mooncake transfer engine " + f"cannot register {self.device_manager.device_type} device memory. Use the " + f"broadcast update (leave ranks unset) instead." + ) assert self._p2p_store is not None, "p2p store is not initialized" assert ranks, "ranks should be set" @@ -780,112 +826,118 @@ def _update_per_bucket( self._p2p_store.register_named_tensors( {p2p_ipc_buffer_name: buffer if disable_h2d_buffer else h2d_buffer} ) - handle = reduce_tensor(buffer) - - buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list) - max_len = 0 - for receiver_rank, _, bucket in buckets: - buckets_by_receiver_rank[receiver_rank].append(bucket) - if len(buckets_by_receiver_rank[receiver_rank]) > max_len: - max_len = len(buckets_by_receiver_rank[receiver_rank]) - - socket, socket_paths = self._bind_zmq_socket() - req_thread = threading.Thread( - target=req_func, - args=(socket_paths,), - ) - req_thread.start() - socket.send_pyobj(handle) + # `with` releases the exported handle even if the bind or first send below + # fails; the barrier stays inside so an early failure cannot deadlock peers. + with build_ipc_handler(self.device_manager) as ipc_handler: + handle = ipc_handler.export(buffer) + + buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list) + max_len = 0 + for receiver_rank, _, bucket in buckets: + buckets_by_receiver_rank[receiver_rank].append(bucket) + if len(buckets_by_receiver_rank[receiver_rank]) > max_len: + max_len = len(buckets_by_receiver_rank[receiver_rank]) + + socket, socket_paths = self._bind_zmq_socket() + req_thread = threading.Thread( + target=req_func, + args=(socket_paths,), + ) + req_thread.start() + # The handle is self-contained for every handler, so one ZMQ send completes the handoff. + socket.send_pyobj(handle) - gidx = 0 - ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64) - try: - for i in range(max_len): - if i < len(receiver_rank_buckets) and not disable_h2d_buffer: - self._copy_to_buffer( - checkpoint_name, - receiver_rank_buckets[i][1], - h2d_buffer, - receiver_rank_buckets[i][0] if ranks else None, - ) - for receiver_rank, _buckets in buckets_by_receiver_rank.items(): - if i >= len(_buckets): - continue - bucket = _buckets[i] - alloc, reserved = ( - self.device_manager.device_module.memory_allocated() / 1024 / 1024, - self.device_manager.device_module.memory_reserved() / 1024 / 1024, - ) - self._logger_rank0( - f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. " - f"Current device allocated {alloc:.2f} MB, " - f"reserved {reserved:.2f} MB." - ) - start = gidx % 2 * bucket_size - buffer_b: torch.Tensor = buffer[start : start + bucket.size] - if receiver_rank == self._rank: - if disable_h2d_buffer: - if p2p_update: - assert bucket == receiver_rank_buckets[i][1] - self._copy_to_buffer( - checkpoint_name, - bucket, - buffer_b, - receiver_rank_buckets[i][0] if p2p_update else None, + gidx = 0 + ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64) + try: + for i in range(max_len): + if i < len(receiver_rank_buckets) and not disable_h2d_buffer: + self._copy_to_buffer( + checkpoint_name, + receiver_rank_buckets[i][1], + h2d_buffer, + receiver_rank_buckets[i][0] if ranks else None, + ) + for receiver_rank, _buckets in buckets_by_receiver_rank.items(): + if i >= len(_buckets): + continue + bucket = _buckets[i] + alloc, reserved = ( + self.device_manager.device_module.memory_allocated() / 1024 / 1024, + self.device_manager.device_module.memory_reserved() / 1024 / 1024, + ) + self._logger_rank0( + f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. " + f"Current device allocated {alloc:.2f} MB, " + f"reserved {reserved:.2f} MB." + ) + start = gidx % 2 * bucket_size + buffer_b: torch.Tensor = buffer[start : start + bucket.size] + if receiver_rank == self._rank: + if disable_h2d_buffer: + if p2p_update: + assert bucket == receiver_rank_buckets[i][1] + self._copy_to_buffer( + checkpoint_name, + bucket, + buffer_b, + receiver_rank_buckets[i][0] if p2p_update else None, + ) + else: + buffer_b.data.copy_(h2d_buffer[: bucket.size]) + dist.broadcast(buffer_b, src=receiver_rank, group=ranks_group) + resp = socket.recv() + if resp != b"": + msg = resp.decode("utf-8") + logger.error( + f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}" ) - else: - buffer_b.data.copy_(h2d_buffer[: bucket.size]) - dist.broadcast(buffer_b, src=receiver_rank, group=ranks_group) - resp = socket.recv() - if resp != b"": - msg = resp.decode("utf-8") - logger.error( - f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}" + ret_code.fill_(1) + dist.all_reduce( + ret_code, op=torch.distributed.ReduceOp.SUM, group=ranks_group ) - ret_code.fill_(1) - dist.all_reduce(ret_code, op=torch.distributed.ReduceOp.SUM, group=ranks_group) - self.device_manager.device_module.synchronize() - if ret_code.item() != 0: - # quit early if any rank failed - socket.send_pyobj(RuntimeError("Some workers failed to update weights")) - raise RuntimeError("Failed to update weights due to remote errors") - socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size)) - gidx += 1 - - socket.recv() - device_mem = self.device_manager.device_module.mem_get_info() - logger.info( - f"[rank{self._rank}] weights broadcast done, device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, allocated memory: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, reserved memory: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" - ) - # Notify worker to release handle - socket.send_pyobj(None) - socket.recv() - # Set to None in correct order (views first, then base tensors) - del buffer_b, h2d_buffer, buffer, handle - self.device_manager.device_module.synchronize() - gc.collect() - self.device_manager.device_module.ipc_collect() - self.device_manager.device_module.empty_cache() - self.device_manager.device_module.synchronize() - - # Log actual memory usage - device_mem = self.device_manager.device_module.mem_get_info() - logger.info( - f"[rank{self._rank}] post-release: device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, " - f"allocated: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, " - f"reserved: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" - ) - # Notify worker to call post_hook - socket.send_pyobj(None) - socket.recv() - finally: - req_thread.join() - dist.barrier(group=ranks_group) - socket.close() - if p2p_update: - self._p2p_store.unregister_named_tensors([p2p_ipc_buffer_name]) - - self.device_manager.device_module.empty_cache() + self.device_manager.device_module.synchronize() + if ret_code.item() != 0: + # quit early if any rank failed + socket.send_pyobj(RuntimeError("Some workers failed to update weights")) + raise RuntimeError("Failed to update weights due to remote errors") + socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size)) + gidx += 1 + + socket.recv() + device_mem = self.device_manager.device_module.mem_get_info() + logger.info( + f"[rank{self._rank}] weights broadcast done, device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, allocated memory: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, reserved memory: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" + ) + # Notify worker to release handle + socket.send_pyobj(None) + socket.recv() + # Set to None in correct order (views first, then base tensors) + del buffer_b, h2d_buffer, buffer, handle + self.device_manager.device_module.synchronize() + gc.collect() + self.device_manager.ipc_collect() + self.device_manager.device_module.empty_cache() + self.device_manager.device_module.synchronize() + + # Log actual memory usage + device_mem = self.device_manager.device_module.mem_get_info() + logger.info( + f"[rank{self._rank}] post-release: device mem usage: {(device_mem[1] - device_mem[0]) / 1024 / 1024:.2f} MB, " + f"allocated: {self.device_manager.device_module.memory_allocated() / 1024 / 1024:.2f} MB, " + f"reserved: {self.device_manager.device_module.memory_reserved() / 1024 / 1024:.2f} MB" + ) + # Notify worker to call post_hook + socket.send_pyobj(None) + socket.recv() + finally: + req_thread.join() + dist.barrier(group=ranks_group) + socket.close() + if p2p_update: + self._p2p_store.unregister_named_tensors([p2p_ipc_buffer_name]) + + self.device_manager.device_module.empty_cache() # we need this CLI entry point for compatibility with former versions diff --git a/checkpoint_engine/worker.py b/checkpoint_engine/worker.py index ea170ca..180cc09 100644 --- a/checkpoint_engine/worker.py +++ b/checkpoint_engine/worker.py @@ -8,20 +8,24 @@ import zmq from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid +from checkpoint_engine.ipc_handler import ( + IPCHandler, + TorchIPCHandler, + XpuIPCHandler, +) _WEIGHTS_TYPE = list[tuple[str, torch.Tensor]] -def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor: - func, args = handle - list_args = list(args) - if device_id is not None: - # the key is to change device id to the current device id - # in case two processes have different CUDA_VISIBLE_DEVICES - list_args[6] = device_id - buffer = func(*list_args) - return buffer +def _ipc_handler_for_handle(handle: object) -> IPCHandler: + """Pick the consumer-side IPC handler based on the handle wire format. + + CUDA/NPU send a ``reduce_tensor`` tuple; XPU sends a dict tagged with its kind. + """ + if isinstance(handle, dict) and handle.get("kind") == XpuIPCHandler.kind: + return XpuIPCHandler() + return TorchIPCHandler() class FlattenedTensorMetadata(TypedDict): @@ -59,10 +63,11 @@ def update_weights_from_ipc( socket.connect(zmq_handle) buffer: torch.Tensor | None = None device_manager = DeviceManager() + ipc_handler: IPCHandler | None = None try: - ipc_handle: tuple[Callable, tuple] = socket.recv_pyobj() - assert isinstance(ipc_handle, tuple) - buffer = _rebuild_ipc(ipc_handle, device_id) + ipc_handle = socket.recv_pyobj() + ipc_handler = _ipc_handler_for_handle(ipc_handle) + buffer = ipc_handler.attach(ipc_handle, device_id) assert buffer.dtype == torch.uint8 socket.send(b"") except Exception as e: @@ -91,10 +96,11 @@ def update_weights_from_ipc( device_manager.device_module.synchronize() released = True buffer = None - del ipc_handle + if ipc_handler is not None: + ipc_handler.detach() gc.collect() - device_manager.device_module.ipc_collect() + device_manager.ipc_collect() device_manager.device_module.empty_cache() device_manager.device_module.synchronize() socket.send(b"") @@ -119,6 +125,8 @@ def update_weights_from_ipc( finally: socket.close() del buffer + if ipc_handler is not None: + ipc_handler.detach() gc.collect() device_manager.device_module.empty_cache() @@ -147,6 +155,9 @@ def _device_uuid(self) -> str: return current_platform.get_device_uuid(self.device.index) elif current_platform.device_type == "npu": return f"NPU-{npu_generate_uuid()}" + elif current_platform.device_type == "xpu": + # Must match ps.py::_get_physical_gpu_id ("GPU-") for the ZMQ key to resolve. + return f"GPU-{torch.xpu.get_device_properties(self.device.index).uuid!s}" else: raise ValueError(f"Unsupported device type: {current_platform.device_type}") @@ -170,9 +181,10 @@ def update_weights_from_ipc(self, zmq_handles: dict[str, str]): The device UUID is platform-specific: - For CUDA: UUID from `current_platform.get_device_uuid()` - For NPU: Format "NPU-{generated_uuid}" + - For XPU: Format "GPU-{torch.xpu device uuid}" Raises: - ValueError: If the device type is not supported (not CUDA or NPU). + ValueError: If the device type is not supported (not CUDA, NPU, or XPU). AssertionError: If the device is not properly initialized. Note: @@ -185,6 +197,8 @@ def update_weights_from_ipc(self, zmq_handles: dict[str, str]): # vllm-ascend not init device if current_platform.device_type == "npu" and self.device is None: self.device = torch.device(f"npu:{self.local_rank}") + elif current_platform.device_type == "xpu" and self.device is None: + self.device = torch.device(f"xpu:{self.local_rank}") assert self.device is not None def _load_weights(weights: _WEIGHTS_TYPE): diff --git a/checkpoint_engine/xpu_ipc/__init__.py b/checkpoint_engine/xpu_ipc/__init__.py new file mode 100644 index 0000000..34c2a2a --- /dev/null +++ b/checkpoint_engine/xpu_ipc/__init__.py @@ -0,0 +1,145 @@ +"""Cross-process device-buffer IPC for Intel XPU via SYCL IPC memory. + +``sycl_ipc.cpp`` wraps the SYCL IPC memory API (``get``/``open``/``close``), +exported by torch's own libsycl (oneAPI >= 2026.0); this module JIT-compiles it +with ``with_sycl``. The handle is a self-contained portable byte blob (no dma-buf +fd, no offset to carry), so it rides the existing ZMQ channel like CUDA's +``reduce_tensor`` tuple -- see ``XpuIPCHandler``. +""" + +import functools +import glob +import os +import shutil +from pathlib import Path +from typing import TYPE_CHECKING + +from loguru import logger + + +if TYPE_CHECKING: + from types import ModuleType + + import torch + + +def _has_ipc_memory(icpx: str) -> bool: + """Whether this icpx ships the SYCL IPC memory header (oneAPI >= 2026.0).""" + root = os.path.dirname(os.path.dirname(icpx)) + header = os.path.join( + root, "include", "sycl", "ext", "oneapi", "experimental", "ipc_memory.hpp" + ) + return os.path.exists(header) + + +def _find_icpx() -> str | None: + """Locate an icpx (SYCL) compiler new enough for the SYCL IPC memory build. + + Compilers without the header are skipped: they build a device image that + torch's newer libsycl cannot load, aborting the process on dlopen rather + than raising something we could catch. + """ + candidates: list[str] = [] + root = os.getenv("CMPLR_ROOT") + if root: + candidates.append(os.path.join(root, "bin", "icpx")) + candidates += sorted(glob.glob("/opt/intel/oneapi/compiler/*/bin/icpx"), reverse=True) + # Fallback to PATH: covers oneAPI layouts outside /opt and a sourced setvars.sh + # that puts icpx on PATH without exporting CMPLR_ROOT. + which = shutil.which("icpx") + if which: + candidates.append(which) + return next( + (c for c in candidates if os.path.exists(c) and _has_ipc_memory(c)), + None, + ) + + +@functools.lru_cache(maxsize=1) +def load_ext() -> "ModuleType": + """JIT-compile (``with_sycl``, linking torch's libsycl) and cache the SYCL IPC extension. + + Raises on any failure; callers treat an exception as "XPU IPC unavailable". + """ + icpx = _find_icpx() + if icpx is None: + raise RuntimeError( + "no icpx with SYCL ipc_memory support found (needs oneAPI >= 2026.0); " + "cannot build XPU IPC extension" + ) + + from torch.utils.cpp_extension import load + + src = Path(__file__).with_name("sycl_ipc.cpp") + + # with_sycl=True supplies the SYCL include paths and device link, but torch invokes + # a bare "icpx", so it must be on PATH; keep -O2 or the host object is built -O0. + prev_path = os.environ.get("PATH", "") + os.environ["PATH"] = os.path.dirname(icpx) + os.pathsep + prev_path + try: + module = load( + name="checkpoint_engine_sycl_ipc", + sources=[str(src)], + extra_cflags=["-O2"], + with_sycl=True, + verbose=False, + ) + finally: + os.environ["PATH"] = prev_path + return module + + +# Cache only a *successful* probe so a transient first failure can be retried +# (the build itself is memoised by load_ext()'s lru_cache). +_AVAILABLE: bool = False + + +def is_available() -> bool: + """Whether native XPU SYCL IPC can be built and used here (successes cached, failures retried).""" + global _AVAILABLE + if _AVAILABLE: + return True + try: + import torch + + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + return False + load_ext() + except Exception as e: # noqa: BLE001 + logger.warning(f"xpu sycl ipc unavailable: {e}") + return False + _AVAILABLE = True + return True + + +def prewarm() -> bool: + """Build the extension ahead of time (outside any weight-update timeout); safe on non-XPU hosts.""" + return is_available() + + +def get_handle(ptr: int) -> bytes: + """Portable IPC handle bytes for a device pointer (interior pointers ok; offset is in the blob).""" + return bytes(load_ext().ipc_get_handle(ptr)) + + +def open_handle(handle_bytes: bytes, device: int) -> int: + """Open another process's handle -> device pointer (offset included); free via :func:`close_handle`.""" + return load_ext().ipc_open_handle(list(handle_bytes), device) + + +def release_handle(ptr: int) -> None: + """Release the exporter handle from :func:`get_handle`; no-op if ``ptr`` was never exported. + + Deferred until the consumer has opened: releasing earlier can free the fd under + the level-zero-v2 UR adapter. + """ + load_ext().ipc_release_handle(ptr) + + +def close_handle(ptr: int) -> None: + load_ext().ipc_close_handle(ptr) + + +def wrap_tensor(ptr: int, nbytes: int, device: int) -> "torch.Tensor": + """Wrap an IPC-mapped device pointer as a non-owning torch XPU uint8 tensor.""" + return load_ext().ipc_wrap_tensor(ptr, nbytes, device) diff --git a/checkpoint_engine/xpu_ipc/sycl_ipc.cpp b/checkpoint_engine/xpu_ipc/sycl_ipc.cpp new file mode 100644 index 0000000..ef9b933 --- /dev/null +++ b/checkpoint_engine/xpu_ipc/sycl_ipc.cpp @@ -0,0 +1,111 @@ +// Cross-process IPC for Intel XPU tensors via SYCL ipc_memory. The handle is a +// self-contained, portable byte blob: the consumer opens it from the bytes alone +// (no dma-buf fd, no offset to carry -- get() takes torch's interior data_ptr and +// open() restores it). Context and device come from torch (c10::xpu) so the +// mapping lands on torch's own SYCL context. +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +// Upstream split this API (functions -> ipc::memory, types -> parent ipc) and +// deprecated flat ipc_memory; no oneAPI release ships it yet, so probe for it. +#if __has_include() +namespace ipc = sycl::ext::oneapi::experimental::ipc::memory; +namespace ipc_types = sycl::ext::oneapi::experimental::ipc; +#else +namespace ipc = sycl::ext::oneapi::experimental::ipc_memory; +namespace ipc_types = sycl::ext::oneapi::experimental::ipc_memory; +#endif + +namespace { + +std::vector to_bytes(const std::vector& in) { + return {reinterpret_cast(in.data()), + reinterpret_cast(in.data()) + in.size()}; +} + +// Exporter-side handles kept alive until ipc_release_handle(). We must not +// ipc::put() before the consumer opens: under the UR level-zero-v2 adapter +// put_ipc_handle frees the exporter fd and can race the consumer's open. +// The handle is a copyable, non-owning value (freed only via ipc::put), so +// storing it by value needs no manual new/delete. +std::mutex g_handles_mu; +std::unordered_map g_handles; + +} // namespace + +// Portable IPC handle bytes for the allocation backing `ptr` (interior pointers +// are fine -- the offset is in the blob). Handle retained until ipc_release_handle(). +std::vector ipc_get_handle(uintptr_t ptr) { + sycl::context ctx = c10::xpu::get_device_context(); + ipc_types::handle h = ipc::get(reinterpret_cast(ptr), ctx); + ipc_types::handle_data_t data = h.data(); // owning copy of the blob, independent of `h` + { + std::lock_guard lk(g_handles_mu); + auto it = g_handles.find(ptr); + if (it != g_handles.end()) { + ipc::put(it->second, ctx); // release stale handle for a reused address + it->second = h; + } else { + g_handles.emplace(ptr, h); + } + } + return {reinterpret_cast(data.data()), + reinterpret_cast(data.data()) + data.size()}; +} + +// Release the exporter handle from ipc_get_handle(ptr); no-op if unregistered. +// Call only after all consumers have opened their mappings (see level-zero-v2 note). +void ipc_release_handle(uintptr_t ptr) { + std::optional h; + { + std::lock_guard lk(g_handles_mu); + auto it = g_handles.find(ptr); + if (it == g_handles.end()) { + return; + } + h = it->second; + g_handles.erase(it); + } + ipc::put(*h, c10::xpu::get_device_context()); +} + +// Open a handle from another process -> mapped device pointer (offset included); +// pass it back to ipc_close_handle to release the mapping. +uintptr_t ipc_open_handle(const std::vector& blob, int64_t device) { + sycl::context ctx = c10::xpu::get_device_context(); + sycl::device dev = c10::xpu::get_raw_device(static_cast(device)); + std::vector data = to_bytes(blob); + void* p = ipc::open(data, ctx, dev); + return reinterpret_cast(p); +} + +// Close an IPC mapping; `ptr` must be from ipc_open_handle. +void ipc_close_handle(uintptr_t ptr) { + ipc::close(reinterpret_cast(ptr), c10::xpu::get_device_context()); +} + +// Wrap an external device pointer as a non-owning torch XPU uint8 tensor (the XPU +// analogue of rebuild_cuda_tensor) so the worker can read weight slices. +torch::Tensor ipc_wrap_tensor(uintptr_t dptr, int64_t nbytes, int64_t device) { + auto opts = torch::TensorOptions().dtype(torch::kUInt8).device( + torch::kXPU, static_cast(device)); + return torch::from_blob(reinterpret_cast(dptr), {nbytes}, [](void*) {}, opts); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ipc_get_handle", &ipc_get_handle, "Get portable IPC handle bytes for a device ptr"); + m.def("ipc_release_handle", &ipc_release_handle, "Release the retained exporter handle for a ptr"); + m.def("ipc_open_handle", &ipc_open_handle, "Open an IPC handle blob -> device ptr"); + m.def("ipc_close_handle", &ipc_close_handle, "Close an IPC mapping"); + m.def("ipc_wrap_tensor", &ipc_wrap_tensor, "Wrap an external XPU ptr as a torch tensor"); +} diff --git a/pyproject.toml b/pyproject.toml index c200382..b5a9f28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,11 @@ requires = ["setuptools", "setuptools-scm>=8"] build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] -include = ["checkpoint_engine"] +include = ["checkpoint_engine", "checkpoint_engine.*"] + +[tool.setuptools.package-data] +# Ship the SYCL ipc_memory source; it is JIT-compiled at runtime on XPU hosts. +"checkpoint_engine.xpu_ipc" = ["*.cpp"] [tool.setuptools_scm] version_file = "checkpoint_engine/_version.py" diff --git a/tests/test_device_manager.py b/tests/test_device_manager.py new file mode 100644 index 0000000..f7cdb03 --- /dev/null +++ b/tests/test_device_manager.py @@ -0,0 +1,183 @@ +"""Unit tests for DeviceManager multi-accelerator dispatch (cuda / npu / xpu). + +These tests mock the device backends so they run on CPU-only CI (``-m "not gpu"``). A separate +hardware-gated test exercises the real ``torch.xpu`` path when an Intel GPU is present. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from checkpoint_engine.device_utils import DeviceManager + + +def _make_manager(device_type: str, device_module: object) -> DeviceManager: + """Build a DeviceManager without touching real hardware by stubbing detection/setup.""" + dm = DeviceManager.__new__(DeviceManager) + dm.device_type = device_type + dm.device_module = device_module + return dm + + +@pytest.mark.parametrize( + "device_type,expected_backend", + [("cuda", "nccl"), ("npu", "hccl"), ("xpu", "xccl")], +) +def test_backend_mapping(device_type: str, expected_backend: str): + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.backend == expected_backend + + +def test_backend_unsupported(): + dm = _make_manager("tpu", SimpleNamespace()) + with pytest.raises(TypeError): + _ = dm.backend + + +@pytest.mark.parametrize("device_type", ["cuda", "xpu"]) +def test_transfer_engine_protocol_rdma(device_type: str): + dm = _make_manager(device_type, SimpleNamespace()) + with patch("checkpoint_engine.device_utils.has_efa_pci", return_value=False): + assert dm.transfer_engine_protocol == "rdma" + with patch("checkpoint_engine.device_utils.has_efa_pci", return_value=True): + assert dm.transfer_engine_protocol == "efa" + + +def test_transfer_engine_protocol_npu(): + dm = _make_manager("npu", SimpleNamespace()) + assert dm.transfer_engine_protocol == "ascend_direct" + + +def test_ipc_collect_present_is_called(): + called = [] + module = SimpleNamespace(ipc_collect=lambda: called.append(True)) + dm = _make_manager("cuda", module) + dm.ipc_collect() + assert called == [True] + + +def test_ipc_collect_xpu_is_noop(): + # SYCL frees on close_handle, so XPU has no handle cache to collect (and + # torch.xpu has no ipc_collect); it must be a no-op, not an error. + dm = _make_manager("xpu", SimpleNamespace()) + dm.ipc_collect() + + +def test_ipc_collect_rejects_unsupported_device(): + # An unsupported backend must fail loudly rather than silently skipping the + # collect, matching backend/transfer_engine_protocol/_setup_device_module. + dm = _make_manager("tpu", SimpleNamespace()) + with pytest.raises(TypeError, match="not supported"): + dm.ipc_collect() + + +@pytest.mark.parametrize( + "device_type,expected", + [("cuda", True), ("npu", False), ("xpu", False)], +) +def test_supports_inplace_pin(device_type: str, expected: bool): + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.supports_inplace_pin() is expected + + +@pytest.mark.parametrize("device_type", ["cuda", "npu"]) +def test_supports_device_ipc_true_for_cuda_npu(device_type: str): + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.supports_device_ipc() is True + + +def test_supports_device_ipc_xpu_uses_sycl_extension(): + # On XPU we do not rely on torch reductions (PyTorch has no XPU tensor IPC); + # instead we detect our native SYCL ipc_memory extension via xpu_ipc.is_available(). + dm = _make_manager("xpu", SimpleNamespace()) + with patch("checkpoint_engine.xpu_ipc.is_available", return_value=True): + assert dm.supports_device_ipc() is True + with patch("checkpoint_engine.xpu_ipc.is_available", return_value=False): + assert dm.supports_device_ipc() is False + + +def test_supports_device_ipc_unknown_device(): + dm = _make_manager("tpu", SimpleNamespace()) + assert dm.supports_device_ipc() is False + + +@pytest.mark.parametrize( + "device_type,expected", + [("cuda", True), ("npu", True), ("xpu", False), ("tpu", False)], +) +def test_supports_device_p2p(device_type: str, expected: bool): + # Mooncake has no Level Zero backend, so XPU device-memory P2P is unsupported. + dm = _make_manager(device_type, SimpleNamespace()) + assert dm.supports_device_p2p() is expected + + +def test_host_empty_cache_noncuda_uses_gc(): + dm = _make_manager("xpu", SimpleNamespace()) + with patch("checkpoint_engine.device_utils.gc.collect") as gc_collect: + dm.host_empty_cache() + gc_collect.assert_called_once() + + +def test_host_empty_cache_cuda_uses_torch(): + dm = _make_manager("cuda", SimpleNamespace()) + with patch.object(torch._C, "_host_emptyCache", create=True) as host_empty: + dm.host_empty_cache() + host_empty.assert_called_once() + + +# -------------------------------------------------------------------------------------------- +# Hardware-gated: real torch.xpu behavior on an Intel GPU host. +# -------------------------------------------------------------------------------------------- + +_HAS_XPU = hasattr(torch, "xpu") and torch.xpu.is_available() + + +@pytest.mark.gpu +@pytest.mark.skipif(not _HAS_XPU, reason="requires an Intel XPU device") +def test_real_xpu_device_manager(): + dm = DeviceManager() + assert dm.device_type == "xpu" + assert dm.backend == "xccl" + assert dm.device_module is torch.xpu + # ipc_collect must be a harmless no-op (torch.xpu has no ipc_collect). + dm.ipc_collect() + # XPU has no torch-native device-tensor IPC, but checkpoint-engine ships its own + # native SYCL IPC memory handler; supports_device_ipc() must agree with whether + # that extension can actually be built/loaded in this environment. + from checkpoint_engine import xpu_ipc + + assert dm.supports_device_ipc() is xpu_ipc.is_available() + # Mooncake has no Level Zero backend, so device-memory P2P stays unsupported on XPU. + assert dm.supports_device_p2p() is False + # cudaHostRegister-style in-place pinning is CUDA-only. + assert dm.supports_inplace_pin() is False + + +@pytest.mark.gpu +@pytest.mark.skipif(not _HAS_XPU, reason="requires an Intel XPU device") +def test_real_xpu_device_ipc_available_when_extension_builds(): + """When the SYCL ipc_memory extension builds on real XPU hardware, the + broadcast path must be reported as supported. This guards the torch>=2.14 + c++20 build regression that silently disabled XPU broadcast.""" + from checkpoint_engine import xpu_ipc + + if not xpu_ipc.is_available(): + pytest.skip("SYCL ipc_memory extension could not be built in this environment") + dm = DeviceManager() + assert dm.supports_device_ipc() is True + + +@pytest.mark.gpu +@pytest.mark.skipif(not _HAS_XPU, reason="requires an Intel XPU device") +def test_real_xpu_physical_uuid_matches_worker_format(): + from checkpoint_engine.ps import _get_physical_gpu_id + + dm = DeviceManager() + uuid0 = _get_physical_gpu_id(dm, 0) + # Same format the vLLM worker derives from torch.xpu.get_device_properties(idx).uuid. + expected = f"GPU-{torch.xpu.get_device_properties(0).uuid!s}" + assert uuid0 == expected + if dm.device_module.device_count() > 1: + assert _get_physical_gpu_id(dm, 1) != uuid0, "per-device uuids must be distinct" diff --git a/tests/test_ipc_handler.py b/tests/test_ipc_handler.py new file mode 100644 index 0000000..7acbdfd --- /dev/null +++ b/tests/test_ipc_handler.py @@ -0,0 +1,110 @@ +"""Unit tests for the IPC-handler seam (CPU-only, no accelerator required). + +The zero-copy device handoff itself is exercised by a hardware-gated end-to-end +test on real XPU/CUDA. Here we cover the dispatch logic and wire formats. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from checkpoint_engine.ipc_handler import ( + TorchIPCHandler, + XpuIPCHandler, + build_ipc_handler, +) +from checkpoint_engine.worker import _ipc_handler_for_handle + + +def _dm(device_type: str) -> object: + return SimpleNamespace(device_type=device_type) + + +@pytest.mark.parametrize( + "device_type,expected", + [("cuda", TorchIPCHandler), ("npu", TorchIPCHandler), ("xpu", XpuIPCHandler)], +) +def test_build_ipc_handler_dispatch(device_type: str, expected: type): + assert isinstance(build_ipc_handler(_dm(device_type)), expected) + + +def test_consumer_dispatch_by_handle_shape(): + # CUDA/NPU send a reduce_tensor tuple; XPU sends a tagged dict. + tuple_handle = (lambda *a: None, (1, 2, 3)) + assert isinstance(_ipc_handler_for_handle(tuple_handle), TorchIPCHandler) + + xpu_handle = {"kind": XpuIPCHandler.kind, "handle_bytes": b"", "nbytes": 0} + assert isinstance(_ipc_handler_for_handle(xpu_handle), XpuIPCHandler) + + # An unrelated dict must not be mistaken for the XPU handle. + assert isinstance(_ipc_handler_for_handle({"foo": "bar"}), TorchIPCHandler) + + +def test_torch_handler_export_uses_reduce_tensor(): + sentinel = ("REDUCED",) + with patch("checkpoint_engine.ipc_handler.reduce_tensor", return_value=sentinel) as m: + t = TorchIPCHandler() + out = t.export(SimpleNamespace()) + assert out is sentinel + m.assert_called_once() + + +def test_xpu_export_returns_self_contained_handle(): + # The SYCL handle bytes travel over ZMQ as a picklable dict; no fd, no offset, + # no companion socket. Mock the native extension so this runs on CPU CI. + buffer = SimpleNamespace(data_ptr=lambda: 0xDEAD, nbytes=256) + with patch("checkpoint_engine.xpu_ipc.get_handle", return_value=b"HANDLE") as m: + handle = XpuIPCHandler().export(buffer) + m.assert_called_once_with(0xDEAD) + assert handle == {"kind": "xpu_sycl", "handle_bytes": b"HANDLE", "nbytes": 256} + + +def test_xpu_export_defers_release_until_detach(): + # The exporter handle must be released only in detach() (not export()), against + # the exact pointer exported -- releasing early can free the fd under UR v2. + buffer = SimpleNamespace(data_ptr=lambda: 0xBEEF, nbytes=128) + with ( + patch("checkpoint_engine.xpu_ipc.get_handle", return_value=b"H"), + patch("checkpoint_engine.xpu_ipc.release_handle") as release, + ): + t = XpuIPCHandler() + t.export(buffer) + release.assert_not_called() # not released during export + t.detach() + release.assert_called_once_with(0xBEEF) + # Idempotent: a second detach must not double-release. + t.detach() + release.assert_called_once_with(0xBEEF) + + +def test_xpu_handler_detach_is_safe_when_unused(): + # detach() before any export/attach must not raise (and must not touch the ext). + with ( + patch("checkpoint_engine.xpu_ipc.release_handle") as release, + patch("checkpoint_engine.xpu_ipc.close_handle") as close, + ): + XpuIPCHandler().detach() + release.assert_not_called() + close.assert_not_called() + + +def test_xpu_consumer_detach_closes_opened_mapping(): + # The consumer (attach) side unmaps its opened pointer on detach, and must not + # try to release an exporter handle it never took. + handle = {"kind": "xpu_sycl", "handle_bytes": b"H", "nbytes": 64} + with ( + patch("checkpoint_engine.xpu_ipc.open_handle", return_value=0x7000), + patch("checkpoint_engine.xpu_ipc.wrap_tensor") as wrap, + patch("checkpoint_engine.ipc_handler.torch.xpu.synchronize"), + patch("checkpoint_engine.xpu_ipc.close_handle") as close, + patch("checkpoint_engine.xpu_ipc.release_handle") as release, + ): + import torch as _torch + + wrap.return_value = SimpleNamespace(dtype=_torch.uint8) + t = XpuIPCHandler() + t.attach(handle, device_id=0) + t.detach() + close.assert_called_once_with(0x7000) + release.assert_not_called() diff --git a/tests/test_p2p_guard.py b/tests/test_p2p_guard.py new file mode 100644 index 0000000..0b2b2d8 --- /dev/null +++ b/tests/test_p2p_guard.py @@ -0,0 +1,45 @@ +"""The P2P update path must reject devices whose memory Mooncake cannot register. + +XPU device memory has no Level Zero backend in Mooncake, so a P2P update +(``ranks`` set) must raise a clear error before touching the transfer engine. +CPU-only: we stub the ParameterServer internals up to the guard. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import checkpoint_engine.distributed as dist +from checkpoint_engine.ps import ParameterServer + + +def _ps_with_device(device_type: str, *, supports_ipc: bool, supports_p2p: bool) -> ParameterServer: + ps = ParameterServer.__new__(ParameterServer) + ps._rank = 0 + ps.device_manager = SimpleNamespace( + device_type=device_type, + supports_device_ipc=lambda: supports_ipc, + supports_device_p2p=lambda: supports_p2p, + ) + # Non-empty metas so the leading assert passes; content is irrelevant (guard fires first). + ps._current_global_parameter_metas = {0: object()} + return ps + + +def test_p2p_update_rejected_on_xpu(): + ps = _ps_with_device("xpu", supports_ipc=True, supports_p2p=False) + with ( + patch.object(dist, "is_initialized", return_value=True), + pytest.raises(RuntimeError, match=r"P2P weight update .* is not supported"), + ): + ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=[0]) + + +def test_ipc_unavailable_rejected(): + ps = _ps_with_device("xpu", supports_ipc=False, supports_p2p=False) + with ( + patch.object(dist, "is_initialized", return_value=True), + pytest.raises(RuntimeError, match="cross-process device-tensor IPC"), + ): + ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=None) diff --git a/tests/test_xpu_ipc.py b/tests/test_xpu_ipc.py new file mode 100644 index 0000000..598c458 --- /dev/null +++ b/tests/test_xpu_ipc.py @@ -0,0 +1,154 @@ +"""Hardware-gated tests for the native SYCL IPC memory handler on Intel XPU. + +These are skipped unless an Intel GPU is present and the SYCL IPC extension can +be built (needs an oneAPI ``icpx`` with SYCL ipc_memory support). They are marked +``gpu`` so CPU-only CI (``-m "not gpu"``) skips them. +""" + +import os + +import pytest +import torch + + +pytestmark = pytest.mark.gpu + + +def _xpu_ipc_available() -> bool: + try: + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + return False + from checkpoint_engine import xpu_ipc + + return xpu_ipc.is_available() + except Exception: # noqa: BLE001 + return False + + +skip_no_xpu_ipc = pytest.mark.skipif( + not _xpu_ipc_available(), reason="Intel XPU with buildable SYCL ipc_memory extension required" +) + +_N_TENSORS = 8 + + +def _gen_tensors() -> dict[str, torch.Tensor]: + gen = torch.Generator().manual_seed(0) + return { + f"w{i}": (torch.randn(128 + i, 64, generator=gen) * 50).to(torch.bfloat16) + for i in range(_N_TENSORS) + } + + +def _worker_proc( + device_uuid: str, + expected: dict[str, torch.Tensor], + inq: object, + outq: object, +) -> None: + # Module-level so the "spawn" start method can pickle it. + import zmq + + from checkpoint_engine.worker import update_weights_from_ipc + + torch.xpu.set_device(0) + exp = {k: v.to("xpu:0") for k, v in expected.items()} + ctx = zmq.Context() + state = {"n": 0, "ok": True} + + def run(weights: list[tuple[str, torch.Tensor]]) -> None: + for name, w in weights: + if name in exp and not torch.equal(w.to(torch.bfloat16), exp[name]): + state["ok"] = False + elif name in exp: + state["n"] += 1 + + while True: + socket_paths = inq.get() + if socket_paths is None: + break + update_weights_from_ipc( + ctx, + dict(socket_paths)[device_uuid], + device_id=0, + run=run, + post_hook=lambda: torch.xpu.synchronize(), + ) + outq.put((state["n"], state["ok"])) + + +@skip_no_xpu_ipc +def test_sycl_ipc_same_process_roundtrip(): + """get_handle -> open_handle (same process) maps back to the original bytes.""" + from checkpoint_engine import xpu_ipc + + torch.xpu.set_device(0) + t = torch.arange(256, device="xpu:0", dtype=torch.uint8) + torch.xpu.synchronize() + + handle_bytes = xpu_ipc.get_handle(t.data_ptr()) + ptr = xpu_ipc.open_handle(handle_bytes, 0) + wrapped = xpu_ipc.wrap_tensor(ptr, t.numel(), 0) + try: + assert torch.equal(wrapped, t) + finally: + xpu_ipc.close_handle(ptr) + + +@skip_no_xpu_ipc +def test_sycl_ipc_interior_pointer_offset_preserved(): + """A SYCL IPC handle for an interior (sub-allocation) pointer must reopen at + the same offset -- the offset is encoded in the portable handle bytes, not + carried separately.""" + from checkpoint_engine import xpu_ipc + + torch.xpu.set_device(0) + big = torch.zeros(4096, device="xpu:0", dtype=torch.uint8) + view = big[1024:1280] + view.fill_(0xAB) + torch.xpu.synchronize() + + handle_bytes = xpu_ipc.get_handle(view.data_ptr()) + ptr = xpu_ipc.open_handle(handle_bytes, 0) + wrapped = xpu_ipc.wrap_tensor(ptr, view.numel(), 0) + try: + assert torch.equal(wrapped, view) + finally: + xpu_ipc.close_handle(ptr) + + +@skip_no_xpu_ipc +def test_sycl_ipc_cross_process_broadcast(): + """Full ParameterServer broadcast -> colocated worker over the XPU SYCL handler.""" + from torch.multiprocessing import get_context + + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29570") + + from checkpoint_engine.ps import ParameterServer, _get_physical_gpu_id + + tensors = _gen_tensors() + + torch.xpu.set_device(0) + ps = ParameterServer(auto_pg=True) + uuid = _get_physical_gpu_id(ps.device_manager, 0) + + mp = get_context("spawn") + inq, outq = mp.Queue(), mp.Queue() + proc = mp.Process(target=_worker_proc, args=(uuid, tensors, inq, outq)) + proc.start() + try: + ps.register_checkpoint("ckpt", named_tensors=tensors) + ps.init_process_group() + ps.gather_metas("ckpt") + ps.update("ckpt", inq.put) + inq.put(None) + n, ok = outq.get(timeout=60) + assert ok, "received weights did not match originals" + assert n == _N_TENSORS, f"expected {_N_TENSORS} tensors checked, got {n}" + finally: + proc.join(timeout=30) + if proc.is_alive(): + proc.terminate() diff --git a/tests/test_xpu_parity.py b/tests/test_xpu_parity.py new file mode 100644 index 0000000..bb1078e --- /dev/null +++ b/tests/test_xpu_parity.py @@ -0,0 +1,274 @@ +"""CPU-only parity tests for the XPU support paths that diverge from CUDA/NPU. + +These cover device-agnostic logic that would otherwise have no unit coverage: +* The portable-handle contract of ``xpu_ipc.get_handle``/``open_handle`` (bytes + in, bytes out -- no fd, no offset). +* The custom-distributed backend rejection for XPU. +* ``register_checkpoint`` forcing in-place pinning off on non-CUDA devices. + +The zero-copy device handoff itself is exercised by the hardware-gated tests in +``test_xpu_ipc.py``; here we isolate the surrounding logic so it runs on CPU-only +CI (``-m "not gpu"``). +""" + +import os +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import checkpoint_engine.distributed as dist +from checkpoint_engine import xpu_ipc +from checkpoint_engine.distributed.base import TorchBackend, use_backend + + +# ------------------------------------------------------------------------------ +# get_handle / open_handle: a SYCL ipc_memory handle is self-contained portable +# bytes -- get_handle returns the raw blob and open_handle passes it straight to +# the extension (offset and any fd are encoded inside the bytes, not carried). +# ------------------------------------------------------------------------------ + + +def test_get_handle_returns_raw_portable_bytes(): + fake_ext = MagicMock() + blob = list(b"\xab" * 120) # SYCL handle blob (opaque, self-contained) + fake_ext.ipc_get_handle.return_value = blob + with patch("checkpoint_engine.xpu_ipc.load_ext", return_value=fake_ext): + handle_bytes = xpu_ipc.get_handle(0xDEAD) + fake_ext.ipc_get_handle.assert_called_once_with(0xDEAD) + assert handle_bytes == bytes(blob) + + +def test_open_handle_passes_bytes_and_device_through(): + fake_ext = MagicMock() + fake_ext.ipc_open_handle.return_value = 0x5000 # mapped ptr (offset already applied) + with patch("checkpoint_engine.xpu_ipc.load_ext", return_value=fake_ext): + ptr = xpu_ipc.open_handle(b"\xab" * 120, device=2) + assert ptr == 0x5000 + (blob_arg, device_arg) = fake_ext.ipc_open_handle.call_args.args + assert device_arg == 2 + assert bytes(blob_arg) == b"\xab" * 120 + + +# ------------------------------------------------------------------------------ +# load_ext: with_sycl=True makes torch shell out to a bare "icpx", so the located +# compiler must be on PATH for the build and PATH restored afterwards. +# ------------------------------------------------------------------------------ + + +def test_load_ext_puts_icpx_on_path_and_restores_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("PATH", "/usr/bin") + + captured: dict[str, object] = {} + + def fake_load(**kwargs: object) -> MagicMock: + # Record what torch.utils.cpp_extension.load would see. + captured["PATH"] = os.environ.get("PATH") + captured["kwargs"] = kwargs + return MagicMock() + + xpu_ipc.load_ext.cache_clear() + try: + with ( + patch("checkpoint_engine.xpu_ipc._find_icpx", return_value="/opt/oneapi/bin/icpx"), + patch("torch.utils.cpp_extension.load", side_effect=fake_load), + ): + xpu_ipc.load_ext() + finally: + xpu_ipc.load_ext.cache_clear() + + # torch runs `icpx --version`, so its directory must lead PATH during the build. + assert captured["PATH"] == "/opt/oneapi/bin:/usr/bin" + # The SYCL toolchain comes from with_sycl; -O2 must stay or the host object is -O0. + kwargs = captured["kwargs"] + assert kwargs["with_sycl"] is True + assert kwargs["extra_cflags"] == ["-O2"] + # ...and the caller's environment must be restored afterwards. + assert os.environ["PATH"] == "/usr/bin" + + +def test_find_icpx_falls_back_to_path(monkeypatch: pytest.MonkeyPatch) -> None: + # A sourced setvars.sh may put icpx on PATH without exporting CMPLR_ROOT, and + # oneAPI need not live under /opt. Without the PATH fallback the compiler is + # reported missing and XPU IPC is wrongly declared unavailable. + monkeypatch.delenv("CMPLR_ROOT", raising=False) + with ( + patch("checkpoint_engine.xpu_ipc.glob.glob", return_value=[]), + patch("checkpoint_engine.xpu_ipc.shutil.which", return_value="/custom/bin/icpx") as which, + patch("checkpoint_engine.xpu_ipc.os.path.exists", return_value=True), + patch("checkpoint_engine.xpu_ipc._has_ipc_memory", return_value=True), + ): + assert xpu_ipc._find_icpx() == "/custom/bin/icpx" + which.assert_called_once_with("icpx") + + +def test_find_icpx_skips_compiler_without_ipc_memory(monkeypatch: pytest.MonkeyPatch) -> None: + # An icpx older than oneAPI 2026.0 builds a device image that torch's libsycl + # cannot load, aborting the process (SIGABRT) on dlopen -- uncatchable from + # Python. Such compilers must be rejected up front, before any build. + monkeypatch.setenv("CMPLR_ROOT", "/opt/intel/oneapi/compiler/2025.3") + with ( + patch("checkpoint_engine.xpu_ipc.glob.glob", return_value=[]), + patch("checkpoint_engine.xpu_ipc.shutil.which", return_value=None), + patch("checkpoint_engine.xpu_ipc.os.path.exists", return_value=True), + patch("checkpoint_engine.xpu_ipc._has_ipc_memory", return_value=False), + ): + assert xpu_ipc._find_icpx() is None + + +def test_has_ipc_memory_probes_the_header(tmp_path: Path) -> None: + # The discriminator is the header's presence in the compiler's own tree + # (/bin/icpx -> /include/sycl/.../ipc_memory.hpp). + icpx = tmp_path / "bin" / "icpx" + icpx.parent.mkdir(parents=True) + icpx.touch() + assert xpu_ipc._has_ipc_memory(str(icpx)) is False + + header = tmp_path / "include" / "sycl" / "ext" / "oneapi" / "experimental" / "ipc_memory.hpp" + header.parent.mkdir(parents=True) + header.touch() + assert xpu_ipc._has_ipc_memory(str(icpx)) is True + + +def test_is_available_caches_only_success_and_retries_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import checkpoint_engine.xpu_ipc as mod + + monkeypatch.setattr(mod, "_AVAILABLE", False) + calls = {"n": 0} + + def flaky_load() -> MagicMock: + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("transient: compiler env not warm yet") + return MagicMock() + + fake_torch = SimpleNamespace(xpu=SimpleNamespace(is_available=lambda: True)) + with ( + patch.dict("sys.modules", {"torch": fake_torch}), + patch.object(mod, "load_ext", side_effect=flaky_load), + ): + assert mod.is_available() is False # transient failure NOT cached + assert mod.is_available() is True # retried, now succeeds + assert mod.is_available() is True # success cached (no third load_ext) + assert calls["n"] == 2 + + +# ------------------------------------------------------------------------------ +# Custom distributed backend: XPU must reject custom_dist and fall back to the +# native "xccl" TorchBackend (there is no vLLM PyXcclCommunicator to subclass). +# ------------------------------------------------------------------------------ + + +def test_use_backend_rejects_xpu_custom_dist(): + with pytest.raises(ValueError, match="XPU is not supported here"): + use_backend("vllm_xccl") + + +def test_use_backend_none_keeps_default_torch_backend(): + # A falsy backend must leave the (default) TorchBackend in place, which is what + # XPU relies on for xccl. + before = dist.is_initialized # attribute presence sanity + use_backend(None) + from checkpoint_engine.distributed.base import _BACKEND_INSTANCE + + assert isinstance(_BACKEND_INSTANCE, TorchBackend) + assert before is dist.is_initialized + + +# ------------------------------------------------------------------------------ +# register_checkpoint: in-place pinning (cudaHostRegister) is CUDA-only; on XPU +# it must be silently disabled rather than attempted. +# ------------------------------------------------------------------------------ + + +# ------------------------------------------------------------------------------ +# _update_per_bucket: the exported IPC handle is retained until detach(). A +# failure after export but before the broadcast loop (e.g. the ZMQ bind) must +# still release it -- otherwise the exporter handle leaks on every failed update. +# ------------------------------------------------------------------------------ + + +def test_update_per_bucket_detaches_handler_on_early_failure(): + from checkpoint_engine.ps import ParameterServer + + ps = ParameterServer.__new__(ParameterServer) + ps._rank = 0 + ps.device_manager = SimpleNamespace( + device_type="cpu", + supports_device_ipc=lambda: True, + supports_device_p2p=lambda: False, + ) + ps._current_global_parameter_metas = {0: object()} + ps._local_rdma_devices = None + ps._remote_rdma_devices = None + + from checkpoint_engine.ipc_handler import IPCHandler + + # A real IPCHandler (not a MagicMock) so detach() must be reached through + # the context manager's __exit__ rather than by mocked attribute access. + class RecordingHandler(IPCHandler): + def __init__(self) -> None: + self.exported = 0 + self.detached = 0 + + def export(self, buffer: object) -> dict: + self.exported += 1 + return {"kind": "fake"} + + def attach(self, handle: object, device_id: int) -> None: + raise AssertionError("attach is the consumer side; not used here") + + def detach(self) -> None: + self.detached += 1 + + handler = RecordingHandler() + + with ( + patch.object(dist, "is_initialized", return_value=True), + patch("checkpoint_engine.ps.build_ipc_handler", return_value=handler), + patch("checkpoint_engine.ps._gen_h2d_buckets", return_value=[]), + patch.object(ps, "_detect_bucket_size", return_value=(16, False)), + patch.object(ps, "_bind_zmq_socket", side_effect=RuntimeError("bind failed")), + pytest.raises(RuntimeError, match="bind failed"), + ): + ps._update_per_bucket("ckpt", req_func=lambda _paths: None, ranks_group=None, ranks=None) + + # export happened, so the exporter handle is live and must be released even + # though the failure struck before the broadcast loop's own cleanup. + assert handler.exported == 1 + assert handler.detached == 1 + + +def test_register_checkpoint_disables_inplace_pin_on_xpu(): + from checkpoint_engine.ps import ParameterServer + + ps = ParameterServer.__new__(ParameterServer) + ps._rank = 0 + ps.device_manager = SimpleNamespace( + device_type="xpu", + supports_inplace_pin=lambda: False, + ) + ps._memory_pool = {} + ps._current_shared_memory_pool_user = "" + ps._p2p_store = None + ps.shared_memory_pool_name = ParameterServer.shared_memory_pool_name + + captured = {} + + def fake_register(*, inplace_pin: bool, **kwargs: object) -> list: + captured["inplace_pin"] = inplace_pin + return [] + + with patch("checkpoint_engine.ps._register_checkpoint", side_effect=fake_register): + ps.register_checkpoint("ckpt", named_tensors={}, use_inplace_pin_memory=True) + # Requested True, but XPU cannot in-place pin -> must be forced False. + assert captured["inplace_pin"] is False + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"]))