checkpoint engine manager 'colocate' mode support npu

src/twinkle/checkpoint_engine/ipc_checkpoint_engine.py

@@ -31,7 +31,8 @@
 import zmq
 from typing import Any, AsyncGenerator, Generator
 
-from twinkle import get_logger
+from twinkle import Platform, get_logger
+from twinkle.utils.framework import Torch
 from .base import CheckpointEngine
 
 logger = get_logger()
@@ -49,7 +50,7 @@
 
 
 class IPCCheckpointEngine(CheckpointEngine):
-    """Hand weights to a sampler on the same GPU by mapping memory instead of copying it."""
+    """Hand weights to a sampler on the same device by mapping memory instead of copying it."""
 
     def __init__(self, bucket_size: int = 512 << 20, **kwargs) -> None:
         # Smaller default than the NCCL engine's 3 GB: a bigger bucket buys nothing when the transfer
@@ -65,21 +66,22 @@ def __init__(self, bucket_size: int = 512 << 20, **kwargs) -> None:
         self.socket = None
         self._context = None
         self._handle = None
+        self._shm = None
         # Receiver side: the mapping of the sender's buffer, kept across buckets. Re-mapping per
         # bucket is what makes device memory appear to grow during a sync.
         self._mapped: torch.Tensor | None = None
+        self._mapped_shms = []
         self._mapped_signature = None
 
     # ── rendezvous ───────────────────────────────────────────────────────
 
     @staticmethod
     def endpoint() -> str:
-        """The socket both peers derive independently from the GPU they share.
+        """The socket both peers derive independently from the device they share.
 
-        The device's UUID rather than its index: under Ray each role sees its own GPU as index 0, so
-        indices collide across ranks while UUIDs do not.
+        The platform helper obtains the physical device UUID for the current local device.
         """
-        uuid = str(torch.cuda.get_device_properties(torch.cuda.current_device()).uuid)
+        uuid = str(Platform.get_vllm_device_uuid(Torch.get_current_device()))
         return f'ipc:///tmp/twinkle-colocate-{uuid}.sock'
 
     def prepare(self) -> dict[str, Any]:
@@ -165,9 +167,17 @@ def finalize(self):
             path = self.endpoint().removeprefix('ipc://')
             if os.path.exists(path):
                 os.unlink(path)
+        if self._shm is not None:
+            self.send_buf = None
+            self._shm.close()
+            self._shm.unlink()
+            self._shm = None
         self.send_buf = None
         self._handle = None
         self._mapped = None
+        for shm in self._mapped_shms:
+            shm.close()
+        self._mapped_shms.clear()
         self._mapped_signature = None
         self.rank = None
 
@@ -178,9 +188,29 @@ def _ensure_buffer(self, min_size: int) -> None:
         if self.send_buf is not None and self.send_buf.numel() >= min_size:
             return
         size = max(self.bucket_size, min_size)
-        self.send_buf = torch.empty(size, dtype=torch.uint8, device=torch.cuda.current_device())
+        platform = Platform.get_platform()
+        if platform.device_prefix() == 'npu' and not platform.is_ipc_supported():
+            from multiprocessing import shared_memory
+
+            if self._shm is not None:
+                self.send_buf = None
+                self._shm.close()
+                self._shm.unlink()
+                self._shm = None
+            self._shm = shared_memory.SharedMemory(create=True, size=size)
+            self.send_buf = torch.frombuffer(self._shm.buf, dtype=torch.uint8, count=size)
+            self._handle = {'name': self._shm.name, 'size': size}
+            return
+
+        self.send_buf = torch.empty(
+            size,
+            dtype=torch.uint8,
+            device=f'{platform.device_prefix()}:{Torch.get_current_device()}',
+        )
         # One handle per buffer, reused for every bucket: the buffer is refilled, not reallocated, so
         # the mapping stays valid and the receiver can keep it.
+        if platform.device_prefix() == 'npu':
+            import torch_npu  # noqa: F401
         from torch.multiprocessing.reductions import reduce_tensor
         self._handle = reduce_tensor(self.send_buf)
 
@@ -223,7 +253,7 @@ async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None,
     def _flush(self, bucket_meta: list[dict], is_last: bool) -> None:
         """Publish the filled part of the buffer and wait until the receiver is done with it."""
         # The copies above are non_blocking; without this the receiver could map bytes not yet written.
-        torch.cuda.synchronize()
+        Torch.synchronize()
         self.socket.send(pickle.dumps({'handle': self._handle, 'bucket_meta': bucket_meta, 'is_last': is_last}))
         # The receiver copies out of this buffer, so it must say so before we overwrite it.
         self.socket.recv()
@@ -245,7 +275,7 @@ async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor], None
                 yield meta['name'], buffer[start:start + nbytes].view(meta['dtype']).view(meta['shape'])
             # Consumers copy with non_blocking=True, so the acknowledgement has to wait for the copies
             # and not merely for the loop above.
-            torch.cuda.synchronize()
+            Torch.synchronize()
             self.socket.send(b'ack')
             if message['is_last']:
                 break
@@ -259,13 +289,25 @@ def _map(self, handle) -> torch.Tensor:
         signature = self._handle_signature(handle)
         if self._mapped is not None and signature == self._mapped_signature:
             return self._mapped
-        from torch.multiprocessing.reductions import rebuild_cuda_tensor
+        if isinstance(handle, dict):
+            from multiprocessing import shared_memory
+
+            mapped_shm = shared_memory.SharedMemory(name=handle['name'])
+            self._mapped_shms.append(mapped_shm)
+            self._mapped = torch.frombuffer(
+                mapped_shm.buf,
+                dtype=torch.uint8,
+                count=handle['size'],
+            )
+            self._mapped_signature = signature
+            return self._mapped
+
         func, args = handle
         args = list(args)
-        # Both peers see the shared GPU as their own device 0, but be explicit rather than trust the
-        # index the sender happened to record.
-        args[6] = torch.cuda.current_device()
-        self._mapped = func(*args) if callable(func) else rebuild_cuda_tensor(*args)
+        if Platform.device_prefix() == 'npu':
+            import torch_npu  # noqa: F401
+        args[6] = Torch.get_current_device()
+        self._mapped = func(*args)
         self._mapped_signature = signature
         return self._mapped
 
@@ -276,6 +318,8 @@ def _handle_signature(handle) -> tuple:
         Locally implemented rather than shared with the sampler's worker extension, which has the same
         helper: the sampler imports this package, so importing it back would be circular.
         """
+        if isinstance(handle, dict):
+            return tuple(handle.items())
         _, args = handle
         return tuple((type(v).__name__, bytes(v) if isinstance(v, (bytes, bytearray)) else v) for v in args
                      if isinstance(v, (bytes, bytearray, int, float, bool, str)) or v is None)

src/twinkle/checkpoint_engine/manager.py

@@ -127,8 +127,6 @@ def decide_backend_engine(
 
         platform_name = Platform.get_platform(platform).__name__
         if mode == 'colocate':
-            if platform_name != 'GPU':
-                raise NotImplementedError("mode='colocate' currently requires the GPU platform.")
             from twinkle.checkpoint_engine import IPCCheckpointEngine
             return IPCCheckpointEngine
         if mode != 'standalone':
@@ -166,8 +164,8 @@ def sync_weights(self, merge_and_sync=True):
             self._sync_weights_naive(merge_and_sync)
             return
 
-        is_master = [True] + [False] * (self.model.device_mesh.world_size - 1)
-        model_metadata = self.model.prepare_checkpoint_engine(is_master)
+        model_metadata = self.model.prepare_checkpoint_engine([True]
+                                                              + [False] * (self.model.device_mesh.world_size - 1))
         self.sampler.prepare_checkpoint_engine(False)
         model_kwargs, sampler_kwargs = self.backend_cls.build_topology(
             self.model.device_mesh.world_size,

src/twinkle/sampler/vllm_sampler/vllm_engine.py

@@ -51,7 +51,7 @@ class VLLMEngine(BaseSamplerEngine):
     This engine uses vLLM v1's AsyncLLM and supports:
     - Tinker-compatible sample() API with logprobs
     - Multi-tenant LoRA adapters for client-server mode
-    - Weight synchronization via load_weights (colocated) or CUDA IPC
+    - Weight synchronization via load_weights (colocated) or device IPC
     - Sleep/wake_up for GPU memory management in colocated training
 
     Deployment scenarios:
@@ -109,10 +109,10 @@ def __init__(
         # ``list_loras()`` per request.
         self._synced_lora_request: Optional[Any] = None
 
-        # Long-lived CUDA IPC bucket reused across all update_weights()
+        # Long-lived device IPC bucket reused across all update_weights()
         # calls. Allocating a new IPC buffer (and hence a new IPC handle)
-        # per sync forces every worker to create a new CUDA IPC mapping via
-        # ``rebuild_cuda_tensor`` because PyTorch's ``shared_cache`` cannot
+        # per sync forces every worker to create a new device IPC mapping via
+        # the reducer callable because PyTorch's ``shared_cache`` cannot
         # hit on unseen storage handles. The driver reclaims those mappings
         # lazily, which is the root cause of the slow GPU memory drift we
         # observed under frequent LoRA syncs. By pinning a single buffer
@@ -578,14 +578,14 @@ async def update_weights(
         bucket_size_mb: int = 2048,
         **kwargs,
     ) -> None:
-        """Update model weights via ZMQ + CUDA IPC to worker extension.
+        """Update model weights via ZMQ + device IPC to worker extension.
 
         Accepts **either** a ``dict[str, Tensor]`` (legacy) **or** an async
         generator / sync generator of ``(name, tensor)`` pairs (streaming).
 
         The streaming path avoids accumulating a full model copy on GPU:
         tensors are consumed one-by-one from the generator, copied into a
-        GPU IPC bucket, and flushed to the vLLM worker subprocess when the
+        device IPC bucket, and flushed to the vLLM worker subprocess when the
         bucket is full.
 
         Args:
@@ -621,15 +621,19 @@ async def _sync_iter():
 
             weight_aiter = _sync_iter()
 
-        # Peek first tensor to detect device (GPU → IPC, CPU → SHM).
+        # Peek first tensor to detect device (supported accelerator → IPC, CPU → SHM).
         try:
             first_name, first_tensor = await weight_aiter.__anext__()
         except StopAsyncIteration:
             logger.warning('update_weights called with empty weights')
             return
 
-        use_gpu_ipc = first_tensor.is_cuda
-        use_shm = not use_gpu_ipc
+        use_device_ipc = first_tensor.is_cuda
+        if first_tensor.device.type == 'npu':
+            from twinkle.utils.platforms import NPU
+
+            use_device_ipc = NPU.is_ipc_supported()
+        use_shm = not use_device_ipc
 
         # Use a per-sync unique IPC endpoint to avoid cross-actor collisions
         # when multiple sampler actors share the same device UUID.
@@ -650,13 +654,16 @@ async def _sync_iter():
         buffer = None
         shm = None
 
-        if use_gpu_ipc:
+        if use_device_ipc:
+            if first_tensor.device.type == 'npu':
+                # torch_npu registers the NPU reducer used by reduce_tensor.
+                import torch_npu  # noqa: F401
             from torch.multiprocessing.reductions import reduce_tensor
 
             # Reuse a long-lived IPC bucket whenever the requested size
             # fits. The handle is produced once and shipped to every
             # subsequent sync so each worker's ``shared_cache`` stays warm
-            # and no new CUDA IPC mapping is created per sync.
+            # and no new device IPC mapping is created per sync.
             need_realloc = (
                 self._ipc_buffer is None or self._ipc_buffer_size < bucket_size
                 or self._ipc_buffer.device != first_tensor.device)
@@ -714,7 +721,7 @@ def _zmq_send_recv(payload, where: str):
                 ))
 
             # Send IPC/SHM handle, wait for worker ready (non-blocking)
-            handle_payload = ipc_handle if use_gpu_ipc else {'name': shm_name, 'size': bucket_size}
+            handle_payload = ipc_handle if use_device_ipc else {'name': shm_name, 'size': bucket_size}
             await loop.run_in_executor(None, _zmq_send_recv, handle_payload, 'handle handshake')
 
             # Stream weights into buckets and send to worker
@@ -821,7 +828,7 @@ async def _flush_bucket(is_last: bool) -> None:
         elapsed = time.time() - start_time
         mode = 'LoRA' if base_sync_done and peft_config else 'base'
         logger.info(f'Updated {n_weights} {mode} weights via '
-                    f"{'IPC' if use_gpu_ipc else 'SHM'} in {elapsed:.2f}s")
+                    f"{'IPC' if use_device_ipc else 'SHM'} in {elapsed:.2f}s")
 
     async def shutdown(self) -> None:
         """Shutdown the vLLM engine and release all resources.

src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py

@@ -45,26 +45,20 @@ def set_death_signal():
 
 
 def _rebuild_ipc(handle, device_id: Optional[int] = None) -> torch.Tensor:
-    """Rebuild CUDA tensor from IPC handle."""
-    from torch.multiprocessing.reductions import rebuild_cuda_tensor
-
+    """Rebuild an accelerator tensor from an IPC reducer handle."""
     func, args = handle
     list_args = list(args)
     if device_id is not None:
         list_args[6] = device_id
-
-    if callable(func):
-        return func(*list_args)
-    else:
-        return rebuild_cuda_tensor(*list_args)
+    return func(*list_args)
 
 
 def _ipc_handle_signature(handle) -> Optional[tuple]:
-    """Derive a stable signature for a CUDA IPC handle.
+    """Derive a stable signature for an accelerator IPC handle.
 
     ``reduce_tensor`` returns ``(func, args)`` where ``args`` contains the
-    CUDA IPC storage handle bytes, storage size, ref-counter handle, etc.
-    Two handles are equivalent (i.e. map the same CUDA memory region) when
+    IPC storage handle bytes, storage size, ref-counter handle, etc.
+    Two handles are equivalent (i.e. map the same device memory region) when
     these inner fields match. We hash only the parts that are picklable and
     comparable to avoid accidental mismatches due to local objects.
     """
@@ -127,12 +121,12 @@ def update_weights_from_ipc(
         use_shm: bool = False,
         zmq_handle: Optional[str] = None,
     ) -> None:
-        """Receive and load weights via ZMQ + CUDA IPC/SHM.
+        """Receive and load weights via ZMQ + device IPC/SHM.
 
         Called via ``collective_rpc("update_weights_from_ipc", ...)`` from
         :meth:`VLLMEngine.update_weights`.  The VLLMEngine sends weights
-        in buckets over a ZMQ REQ/REP channel backed by CUDA IPC (GPU
-        tensors) or shared memory (CPU tensors).
+        in buckets over a ZMQ REQ/REP channel backed by device IPC
+        (accelerator tensors) or shared memory (CPU tensors).
 
         For TP > 1, only TP rank 0 communicates with the VLLMEngine over
         ZMQ.  It broadcasts the IPC handle and bucket metadata to other
@@ -142,7 +136,7 @@ def update_weights_from_ipc(
         Args:
             peft_config: If provided with base_sync_done, loads as LoRA.
             base_sync_done: If True and peft_config, replaces existing LoRA.
-            use_shm: If True, use shared memory instead of CUDA IPC.
+            use_shm: If True, use shared memory instead of device IPC.
             zmq_handle: Optional ZMQ IPC endpoint. If None, uses _get_zmq_handle().
         """
         import torch.distributed as dist
@@ -196,6 +190,10 @@ def _broadcast_obj(obj):
         # ── Step 2: Receive and broadcast IPC/SHM handle ──
         buffer, shm = None, None
 
+        if not use_shm and self.device.type == 'npu':
+            # Register the NPU reducer before recv_pyobj() unpickles its callable.
+            import torch_npu  # noqa: F401
+
         if is_driver:
             try:
                 comm_metadata = socket.recv_pyobj()
@@ -210,9 +208,9 @@ def _broadcast_obj(obj):
         if not use_shm:
             handle = comm_metadata
             # All TP ranks rebuild the IPC buffer from the same handle.
-            # CUDA IPC allows any process on the same node to map the memory.
+            # Device IPC allows any process on the same node to map the memory.
             # Reuse a cached buffer across syncs when the sender reuses the
-            # same IPC handle: this avoids creating a fresh CUDA IPC mapping
+            # same IPC handle: this avoids creating a fresh device IPC mapping
             # per sync, which the driver releases lazily and is the root
             # cause of the apparent GPU memory growth under frequent syncs.
             handle_signature = _ipc_handle_signature(handle)