checkpoint-engine 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

checkpoint_engine/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.3'
-__version_tuple__ = version_tuple = (0, 3, 3)
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)
 
 __commit_id__ = commit_id = None

checkpoint_engine/distributed/__init__.py

@@ -0,0 +1,28 @@
+from .base import (
+    Distributed,
+    DistributedProcessGroup,
+    all_gather_object,
+    all_reduce,
+    barrier,
+    broadcast,
+    destroy_process_group,
+    init_process_group,
+    is_initialized,
+    new_group,
+    use_backend,
+)
+
+
+__all__ = [
+    "Distributed",
+    "DistributedProcessGroup",
+    "all_gather_object",
+    "all_reduce",
+    "barrier",
+    "broadcast",
+    "destroy_process_group",
+    "init_process_group",
+    "is_initialized",
+    "new_group",
+    "use_backend",
+]

checkpoint_engine/distributed/base.py

@@ -0,0 +1,288 @@
+import importlib
+import io
+import pickle
+from abc import ABC, abstractmethod
+from datetime import timedelta
+from typing import Any, Protocol
+
+import torch
+import torch.distributed as torch_dist
+
+
+class CommunicatorProtocol(Protocol):
+    def all_gather(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
+
+
+class CommGroup:
+    def __init__(self, comm_handle: int, ranks: list[int]):
+        self._comm = comm_handle
+        self._ranks = ranks
+
+    @property
+    def handle(self) -> int:
+        return self._comm
+
+    @property
+    def ranks(self) -> list[int]:
+        return self._ranks
+
+
+DistributedProcessGroup = torch_dist.ProcessGroup | CommGroup
+
+
+class Distributed(ABC):
+    @abstractmethod
+    def init_process_group(
+        self,
+        rank: int,
+        world_size: int,
+        store: torch_dist.TCPStore,
+        **kwargs,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def destroy_process_group(
+        self,
+        group: DistributedProcessGroup | None = None,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def is_initialized(self) -> bool:
+        raise NotImplementedError
+
+    @abstractmethod
+    def all_gather_object(
+        self,
+        object_list: list[Any],
+        obj: Any,
+        group: DistributedProcessGroup | None = None,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op: torch_dist.ReduceOp.RedOpType,
+        group: DistributedProcessGroup | None = None,
+        **kwargs,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def broadcast(
+        self,
+        tensor: torch.Tensor,
+        src: int,
+        group: DistributedProcessGroup | None = None,
+        **kwargs,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def barrier(
+        self,
+        group: DistributedProcessGroup | None = None,
+        **kwargs,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def new_group(
+        self,
+        ranks: list[int],
+        **kwargs,
+    ):
+        raise NotImplementedError
+
+
+class TorchBackend(Distributed):
+    def init_process_group(
+        self,
+        rank: int,
+        world_size: int,
+        store: torch_dist.TCPStore,
+        **kwargs,
+    ):
+        backend = kwargs.get("backend", "nccl")
+        timeout = kwargs.get("timeout", timedelta(minutes=10))
+
+        torch_dist.init_process_group(
+            backend=backend,
+            world_size=world_size,
+            rank=rank,
+            timeout=timeout,
+            store=store,
+        )
+
+    def destroy_process_group(self, group: DistributedProcessGroup | None = None):
+        torch_dist.destroy_process_group(group)
+
+    def is_initialized(self) -> bool:
+        return torch_dist.is_initialized()
+
+    def all_gather_object(
+        self, object_list: list[Any], obj: Any, group: DistributedProcessGroup | None = None
+    ):
+        torch_dist.all_gather_object(object_list, obj, group)
+
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op: torch_dist.ReduceOp.RedOpType = torch_dist.ReduceOp.SUM,
+        group: DistributedProcessGroup | None = None,
+        **kwargs,
+    ):
+        torch_dist.all_reduce(tensor, op, group, **kwargs)
+
+    def broadcast(
+        self,
+        tensor: torch.Tensor,
+        src: int = 0,
+        group: DistributedProcessGroup | None = None,
+        **kwargs,
+    ):
+        torch_dist.broadcast(tensor, src, group, **kwargs)
+
+    def barrier(self, group: DistributedProcessGroup | None = None, **kwargs):
+        torch_dist.barrier(group, **kwargs)
+
+    def new_group(self, ranks: list[int], **kwargs) -> DistributedProcessGroup | None:
+        return torch_dist.new_group(ranks, **kwargs)
+
+
+# specific device instance
+_BACKEND_INSTANCE: Distributed = TorchBackend()
+
+_pickler = pickle.Pickler
+_unpickler = pickle.Unpickler
+
+
+def _object_to_tensor(obj: Any, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+    f = io.BytesIO()
+    _pickler(f).dump(obj)
+    byte_storage = torch.ByteStorage._from_buffer(f.getvalue())
+    byte_tensor = torch.ByteTensor(byte_storage).to(device)
+    local_size = torch.LongTensor([byte_tensor.numel()]).to(device)
+    return byte_tensor, local_size
+
+
+def _tensor_to_object(tensor: torch.Tensor, tensor_size: int) -> Any:
+    tensor = tensor.cpu()
+    buf = tensor.numpy().tobytes()[:tensor_size]
+    return _unpickler(io.BytesIO(buf)).load()
+
+
+def _flatten_for_scatter_gather(
+    tensor_list: list[torch.Tensor], copy: bool = False
+) -> torch.Tensor:
+    if not tensor_list:
+        raise RuntimeError("Received an empty list.")
+    t = tensor_list[0]
+    buffer_shape = [len(tensor_list)] + list(t.shape)
+
+    buffer = torch.empty(tuple(buffer_shape), dtype=t.dtype, device=t.device)
+    if copy:
+        for i, tensor in enumerate(tensor_list):
+            buffer[i].copy_(tensor)
+    return buffer
+
+
+def _common_all_gather_object(
+    comm: CommunicatorProtocol,
+    device: torch.device,
+    world_size: int,
+    object_list: list[Any],
+    object: Any,
+):
+    input_tensor, local_size = _object_to_tensor(object, device)
+    object_sizes_tensor = torch.empty(world_size, dtype=torch.long, device=device)
+    comm.all_gather(object_sizes_tensor, local_size)
+    object_size_list = [object_sizes_tensor[i].unsqueeze(dim=0) for i in range(world_size)]
+    max_object_size = int(max(object_size_list).item())
+    input_tensor.resize_(max_object_size)
+    coalesced_output_tensor = torch.empty(
+        max_object_size * world_size, dtype=torch.uint8, device=device
+    )
+
+    comm.all_gather(coalesced_output_tensor, input_tensor)
+    output_tensors = [
+        coalesced_output_tensor[max_object_size * i : max_object_size * (i + 1)]
+        for i in range(world_size)
+    ]
+    for i, tensor in enumerate(output_tensors):
+        tensor = tensor.type(torch.uint8)
+        tensor_size = object_size_list[i]
+        object_list[i] = _tensor_to_object(tensor, tensor_size)
+
+
+def use_backend(backend: str | None):
+    global _BACKEND_INSTANCE
+
+    if not backend:
+        return
+
+    mapping = {
+        "vllm_nccl": ".vllm_nccl.DistributedNccl",
+        "vllm_hccl": ".vllm_hccl.DistributedHccl",
+    }
+    if backend not in mapping:
+        raise ValueError(f"Unsupported custom backend: {backend}")
+
+    module_path, class_name = mapping[backend].rsplit(".", 1)
+    module = importlib.import_module(module_path, "checkpoint_engine.distributed")
+    backend_class = getattr(module, class_name)
+    _BACKEND_INSTANCE = backend_class()
+
+
+def init_process_group(
+    rank: int,
+    world_size: int,
+    store: torch_dist.TCPStore,
+    **kwargs,
+):
+    _BACKEND_INSTANCE.init_process_group(rank, world_size, store, **kwargs)
+
+
+def destroy_process_group(group: DistributedProcessGroup | None = None):
+    _BACKEND_INSTANCE.destroy_process_group(group)
+
+
+def is_initialized() -> bool:
+    return _BACKEND_INSTANCE.is_initialized()
+
+
+def all_gather_object(
+    object_list: list[Any],
+    obj: Any,
+    group: DistributedProcessGroup | None = None,
+):
+    _BACKEND_INSTANCE.all_gather_object(object_list, obj, group)
+
+
+def all_reduce(
+    tensor: torch.Tensor,
+    op: torch_dist.ReduceOp.RedOpType = torch_dist.ReduceOp.SUM,
+    group: DistributedProcessGroup | None = None,
+    **kwargs,
+):
+    _BACKEND_INSTANCE.all_reduce(tensor, op, group, **kwargs)
+
+
+def broadcast(
+    tensor: torch.Tensor,
+    src: int = 0,
+    group: DistributedProcessGroup | None = None,
+    **kwargs,
+):
+    _BACKEND_INSTANCE.broadcast(tensor, src, group, **kwargs)
+
+
+def barrier(group: DistributedProcessGroup | None = None, **kwargs):
+    _BACKEND_INSTANCE.barrier(group, **kwargs)
+
+
+def new_group(ranks: list[int], **kwargs) -> DistributedProcessGroup | None:
+    return _BACKEND_INSTANCE.new_group(ranks, **kwargs)

checkpoint_engine/distributed/vllm_hccl.py

@@ -0,0 +1,323 @@
+import ctypes
+from contextlib import contextmanager
+from typing import Any, ClassVar
+
+import torch
+from torch.distributed import ReduceOp
+from vllm.distributed.utils import StatelessProcessGroup
+from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator
+from vllm_ascend.distributed.device_communicators.pyhccl_wrapper import (
+    Function,
+    HCCLLibrary,
+    aclrtStream_t,
+    buffer_type,
+    hcclComm_t,
+    hcclDataType_t,
+    hcclDataTypeEnum,
+    hcclResult_t,
+)
+from vllm_ascend.utils import current_stream
+
+from checkpoint_engine.distributed.base import CommGroup, Distributed, _common_all_gather_object
+
+
+class HcclCommConfig(ctypes.Structure):
+    _fields_: ClassVar[list[tuple[str, Any]]] = [
+        ("size", ctypes.c_size_t),
+        ("magic_word", ctypes.c_uint32),
+        ("version", ctypes.c_uint32),
+        ("reserved", ctypes.c_uint64),
+        ("hccl_buffer_size", ctypes.c_uint32),
+        ("hccl_deterministic", ctypes.c_uint32),
+        ("hccl_comm_name", ctypes.c_char * 128),
+        ("hccl_udi", ctypes.c_char * 128),
+        ("hccl_op_expansion_mode", ctypes.c_uint32),
+        ("hccl_rdma_traffic_class", ctypes.c_uint32),
+        ("hccl_rdma_service_level", ctypes.c_uint32),
+        ("hcll_world_rank_id", ctypes.c_uint32),
+        ("hccl_job_id", ctypes.c_uint64),
+        ("comm_engine", ctypes.c_int32),
+        ("thread_num", ctypes.c_uint32),
+        ("notify_num_per_thread", ctypes.c_uint32),
+        ("acl_graph_zero_copy_enable", ctypes.c_uint8),
+    ]
+
+
+orig_exported_functions = HCCLLibrary.exported_functions
+extended_functions = [
+    # HcclResult HcclAllGather(
+    #   void *sendBuf, void *recvBuf, uint64_t sendCount, HcclDataType dataType,
+    #   HcclComm comm, alcrtStream stream
+    # )
+    Function(
+        "HcclAllGather",
+        hcclResult_t,
+        [
+            buffer_type,
+            buffer_type,
+            ctypes.c_uint64,
+            hcclDataType_t,
+            hcclComm_t,
+            aclrtStream_t,
+        ],
+    ),
+    # HcclResult HcclCreateSubCommConfig(
+    #   HcclComm *comm, uin32_t rankNum, uint32_t *rankIds, uint64_t subCommId,
+    #   uint32_t subCommRankId, HcclCommConfig *config, HcclComm *subComm
+    # )
+    Function(
+        "HcclCreateSubCommConfig",
+        hcclResult_t,
+        [
+            ctypes.POINTER(hcclComm_t),
+            ctypes.c_uint32,
+            ctypes.POINTER(ctypes.c_uint32),
+            ctypes.c_uint64,
+            ctypes.c_uint32,
+            ctypes.POINTER(HcclCommConfig),
+            ctypes.POINTER(hcclComm_t),
+        ],
+    ),
+]
+
+
+def hccl_all_gather(
+    self,  # noqa: ANN001
+    send_buf: buffer_type,
+    recv_buf: buffer_type,
+    count: ctypes.c_uint64,
+    data_type: hcclDataType_t,
+    comm: hcclComm_t,
+    stream: aclrtStream_t,
+):
+    self.HCCL_CHECK(
+        self._funcs["HcclAllGather"](send_buf, recv_buf, count, data_type, comm, stream)
+    )
+
+
+def hccl_create_subcomm_config(
+    self,  # noqa: ANN001
+    comm: hcclComm_t,
+    ranks_size: ctypes.c_uint32,
+    c_rank_ids: ctypes.POINTER(ctypes.c_uint32),
+    subcomm_id: ctypes.c_uint64,
+    subcomm_rank: ctypes.c_uint64,
+    comm_config: HcclCommConfig,
+) -> hcclComm_t:
+    subcomm = hcclComm_t()
+    self.HCCL_CHECK(
+        self._funcs["HcclCreateSubCommConfig"](
+            ctypes.byref(comm),
+            ranks_size,
+            c_rank_ids,
+            subcomm_id,
+            subcomm_rank,
+            ctypes.byref(comm_config),
+            ctypes.byref(subcomm),
+        )
+    )
+    return subcomm
+
+
+# extend HCCLLibrary
+HCCLLibrary.exported_functions = orig_exported_functions + extended_functions
+HCCLLibrary.hcclAllGather = hccl_all_gather
+HCCLLibrary.hcclCreateSubCommConfig = hccl_create_subcomm_config
+
+
+class PyHcclCommunicatorEx(PyHcclCommunicator):
+    def __init__(self, group: StatelessProcessGroup, device: torch.device):
+        super().__init__(group, device)
+        self.subcomm_id = 1
+
+    def destroy_comm(self, comm: hcclComm_t = None):
+        if comm:
+            self.hccl.hcclCommDestroy(comm)
+        else:
+            self.hccl.hcclCommDestroy(self.comm)
+
+    def all_gather(
+        self, out_tensor: torch.Tensor, in_tensor: torch.Tensor, stream: torch.npu.Stream = None
+    ) -> torch.Tensor:
+        if self.disabled:
+            return
+        assert in_tensor.device == self.device, (
+            f"this hccl communicator is created to work on {self.device}, "
+            f"but the input tensor in on {in_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.hccl.hcclAllGather(
+            buffer_type(in_tensor.data_ptr()),
+            buffer_type(out_tensor.data_ptr()),
+            in_tensor.numel(),
+            hcclDataTypeEnum.from_torch(in_tensor.dtype),
+            self.comm,  # todo
+            aclrtStream_t(stream.npu_stream),
+        )
+        return out_tensor
+
+    def create_subcomm(self, ranks: list[int]) -> hcclComm_t:
+        comm_config = HcclCommConfig(
+            size=312,
+            magic_word=0xF0F0F0F0,
+            version=6,
+            reserved=0,
+            hccl_buffer_size=0xFFFFFFFF,
+            hccl_deterministic=0xFFFFFFFF,
+            hccl_comm_name=b"\0",
+            hccl_udi=b"\0",
+            hccl_op_expansize_mode=0,
+            hccl_rdma_traffic_class=0xFFFFFFFF,
+            hccl_rdma_service_level=0xFFFFFFFF,
+            hccl_world_rank_id=0,
+            hccl_job_id=0,
+            comm_engine=-1,
+            thread_num=0xFFFFFFFF,
+            notify_num_per_thread=0xFFFFFFFF,
+            acl_graph_zero_copy_enable=0,
+        )
+        uint32_array = ctypes.c_uint32 * len(ranks)
+        c_rank_ids = uint32_array(*ranks)
+        subcomm_rank = ranks.index(self.rank)
+        ranks_size = len(ranks)
+        subcomm_id = self.subcomm_id
+
+        subcomm = self.hccl.hcclCreateSubCommConfig(
+            self.comm, ranks_size, c_rank_ids, subcomm_id, subcomm_rank, comm_config
+        )
+        self.subcomm_id += 1
+        return subcomm
+
+
+class DistributedHccl(Distributed):
+    def __init__(self):
+        self.pg: StatelessProcessGroup = None
+        self.pyhccl: PyHcclCommunicatorEx = None
+        self.sub_groups: dict[int, CommGroup] = {}
+        self.comm: hcclComm_t = None
+
+        self.host: str = None
+        self.port: int = None
+        self.rank: int = None
+        self.world_size: int = None
+        self.device: torch.device = None
+
+        self.initialized: bool = False
+
+    @contextmanager
+    def _use_group(self, group: CommGroup | None, src: int | None = None):
+        active_src = src
+        if group:
+            assert group.handle in self.sub_groups, "invalid sub_group"
+            newcomm = ctypes.c_void_p(group.handle)
+            self.pyhccl.comm = newcomm
+
+            if src is not None:
+                assert src in group.ranks, "src rank not in group"
+                # convert src rank id in default world to newcomm
+                active_src = group.ranks.index(src)
+                self.pyhccl.rank = group.ranks.index(self.rank)
+
+        try:
+            yield active_src
+        finally:
+            if group:
+                self.pyhccl.comm = self.comm
+                if src is not None:
+                    self.pyhccl.rank = self.rank
+
+    def init_process_group(
+        self,
+        rank: int,
+        world_size: int,
+        store: torch.distributed.TCPStore,
+        **kwargs,
+    ):
+        assert not self.initialized, "already initialized"
+
+        self.rank = rank
+        self.world_size = world_size
+        self.device = torch.device("npu", torch.npu.current_device())
+
+        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
+        self.pyhccl = PyHcclCommunicatorEx(group=self.pg, device=self.device)
+        self.comm = self.pyhccl.comm
+        self.initialized = True
+
+    def destroy_process_group(
+        self,
+        group: CommGroup | None = None,
+    ):
+        assert self.initialized, "not initialized"
+
+        if group and group.handle in self.sub_groups:
+            subcomm = ctypes.c_void_p(group.handle)
+            self.pyhccl.destroy_comm(subcomm)
+            del self.sub_groups[group.handle]
+            return
+
+        self.pyhccl.destroy_comm()
+        self.pyhccl = None
+        self.pg = None
+        self.initialized = False
+
+    def is_initialized(self) -> bool:
+        return self.initialized
+
+    def all_gather_object(self, object_list: list[Any], obj: Any, group: CommGroup | None = None):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group):
+            _common_all_gather_object(self.pyhccl, self.device, self.world_size, object_list, obj)
+            current_stream().synchronize()
+
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op: ReduceOp.RedOpType = ReduceOp.SUM,
+        group: CommGroup | None = None,
+        **kwargs,
+    ):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group):
+            out_tensor = self.pyhccl.all_reduce(tensor, op)
+            current_stream().synchronize()
+            tensor.copy_(out_tensor)
+
+    def broadcast(
+        self, tensor: torch.Tensor, src: int | None = None, group: CommGroup | None = None, **kwargs
+    ):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group, src) as local_rank:
+            self.pyhccl.broadcast(tensor, local_rank)
+            current_stream().synchronize()
+
+    def barrier(self, group: CommGroup | None = None, **kwargs):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group):
+            data = torch.zeros(1, device=self.device)
+            self.pyhccl.all_reduce(data)
+            current_stream().synchronize()
+
+    def new_group(self, ranks: list[int], **kwargs) -> CommGroup | None:
+        assert self.initialized, "not initialized"
+
+        # ranks is None or []
+        if not ranks:
+            ranks = list(range(self.world_size))
+        else:
+            ranks.sort()
+
+        group: CommGroup = None
+        if self.rank not in ranks:
+            return group
+
+        subcomm = self.pyhccl.create_subcomm(ranks)
+        if subcomm:
+            group = CommGroup(subcomm.value, ranks)
+            self.sub_groups[subcomm.value] = group
+        return group

checkpoint_engine/distributed/vllm_nccl.py

@@ -0,0 +1,223 @@
+import ctypes
+from contextlib import contextmanager
+from typing import Any, ClassVar
+
+import torch
+from torch.distributed import ReduceOp
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+from vllm.distributed.device_communicators.pynccl_wrapper import (
+    Function,
+    NCCLLibrary,
+    ncclComm_t,
+    ncclResult_t,
+)
+from vllm.distributed.utils import StatelessProcessGroup
+from vllm.utils import current_stream
+
+from checkpoint_engine.distributed.base import CommGroup, Distributed, _common_all_gather_object
+
+
+class NcclConfigT(ctypes.Structure):
+    _fields_: ClassVar[list[tuple[str, Any]]] = [
+        ("size", ctypes.c_size_t),
+        ("magic", ctypes.c_uint),
+        ("version", ctypes.c_uint),
+        ("blocking", ctypes.c_int),
+        ("cgaClusterSize", ctypes.c_int),
+        ("minCTAs", ctypes.c_int),
+        ("maxCTAs", ctypes.c_int),
+        ("netName", ctypes.c_char_p),
+        ("splitShare", ctypes.c_int),
+        ("trafficClass", ctypes.c_int),
+        ("commName", ctypes.c_char_p),
+        ("collnetEnable", ctypes.c_int),
+        ("CTAPolicy", ctypes.c_int),
+        ("shrinkShare", ctypes.c_int),
+        ("nvlsCTAs", ctypes.c_int),
+        ("nChannelsPerNetPeer", ctypes.c_int),
+        ("nvlinkCentricSched", ctypes.c_int),
+        ("graphUsageMode", ctypes.c_int),
+        ("numRmaCtx", ctypes.c_int),
+    ]
+
+
+nccl_orig_exported_functions = NCCLLibrary.exported_functions
+nccl_extended_functions = [
+    # ncclResult_t ncclCommSplit(
+    # ncclComm_t comm, int color, int key, ncclComm_t *newcomm, NcclConfigT *config
+    # )
+    Function(
+        "ncclCommSplit",
+        ncclResult_t,
+        [
+            ncclComm_t,
+            ctypes.c_int,
+            ctypes.c_int,
+            ctypes.POINTER(ncclComm_t),
+            ctypes.POINTER(NcclConfigT),
+        ],
+    ),
+]
+
+
+def nccl_comm_split(
+    self,  # noqa: ANN001
+    comm: ncclComm_t,
+    color: int,
+    key: int,
+) -> ncclComm_t:
+    newcomm = ncclComm_t()
+
+    self.NCCL_CHECK(self._funcs["ncclCommSplit"](comm, color, key, ctypes.byref(newcomm), None))
+    return newcomm
+
+
+# extend NCCLLibrary
+NCCLLibrary.exported_functions = nccl_orig_exported_functions + nccl_extended_functions
+NCCLLibrary.ncclCommSplit = nccl_comm_split
+
+
+class PyNcclCommunicatorEx(PyNcclCommunicator):
+    def destroy_comm(self, comm: ncclComm_t = None):
+        if comm:
+            self.nccl.ncclCommDestroy(comm)
+        else:
+            self.nccl.ncclCommDestroy(self.comm)
+
+    def create_newcomm(self, ranks: list[int]) -> ncclComm_t:
+        if self.rank in ranks:
+            color = 0
+        else:
+            color = -1  # NCCL_SPLIT_NOCOLOR
+        newcomm = self.nccl.ncclCommSplit(self.comm, color, self.rank)
+        return newcomm
+
+
+class DistributedNccl(Distributed):
+    def __init__(self):
+        self.pg: StatelessProcessGroup = None
+        self.pynccl: PyNcclCommunicatorEx = None
+        self.sub_groups: dict[int, list[int]] = {}
+        self.comm: ncclComm_t = None
+
+        self.host: str = None
+        self.port: int = None
+        self.rank: int = None
+        self.world_size: int = None
+        self.device: torch.device = None
+
+        self.initialized: bool = False
+
+    @contextmanager
+    def _use_group(self, group: CommGroup | None, src: int | None = None):
+        active_src = src
+        if group:
+            assert group.handle in self.sub_groups, "invalid sub_group"
+            newcomm = ctypes.c_void_p(group.handle)
+            self.pynccl.comm = newcomm
+
+            if src is not None:
+                assert src in group.ranks, "src rank not in group"
+                # convert src rank id in default world to newcomm
+                active_src = group.ranks.index(src)
+                self.pynccl.rank = group.ranks.index(self.rank)
+
+        try:
+            yield active_src
+        finally:
+            if group:
+                self.pynccl.comm = self.comm
+                if src is not None:
+                    self.pynccl.rank = self.rank
+
+    def init_process_group(
+        self,
+        rank: int,
+        world_size: int,
+        store: torch.distributed.TCPStore,
+        **kwargs,
+    ):
+        assert not self.initialized, "already initialized"
+
+        self.rank = rank
+        self.world_size = world_size
+        self.device = torch.device("cuda", torch.cuda.current_device())
+
+        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
+        self.pynccl = PyNcclCommunicatorEx(group=self.pg, device=self.device)
+        self.comm = self.pynccl.comm
+        self.initialized = True
+
+    def destroy_process_group(
+        self,
+        group: CommGroup | None = None,
+    ):
+        assert self.initialized, "not initialized"
+
+        if group and group.handle in self.sub_groups:
+            newcomm = ctypes.c_void_p(group.handle)
+            self.pynccl.destroy_comm(newcomm)
+            del self.sub_groups[group.handle]
+            return
+
+        self.pynccl.destroy_comm()
+        self.pynccl = None
+        self.pg = None
+        self.initialized = False
+
+    def is_initialized(self) -> bool:
+        return self.initialized
+
+    def all_gather_object(self, object_list: list[Any], obj: Any, group: CommGroup | None = None):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group):
+            _common_all_gather_object(self.pynccl, self.device, self.world_size, object_list, obj)
+            current_stream().synchronize()
+
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op: ReduceOp.RedOpType = ReduceOp.SUM,
+        group: CommGroup | None = None,
+        **kwargs,
+    ):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group):
+            out_tensor = self.pynccl.all_reduce(in_tensor=tensor, op=op)
+            current_stream().synchronize()
+            tensor.copy_(out_tensor)
+
+    def broadcast(
+        self, tensor: torch.Tensor, src: int | None = None, group: CommGroup | None = None, **kwargs
+    ):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group, src) as local_src:
+            self.pynccl.broadcast(tensor, local_src)
+            current_stream().synchronize()
+
+    def barrier(self, group: CommGroup | None = None, **kwargs):
+        assert self.initialized, "not initialized"
+
+        with self._use_group(group):
+            data = torch.zeros(1, device=self.device)
+            self.pynccl.all_reduce(data)
+            current_stream().synchronize()
+
+    def new_group(self, ranks: list[int], **kwargs) -> CommGroup | None:
+        assert self.initialized, "not initialized"
+
+        # ranks is None or []
+        if not ranks:
+            ranks = list(range(self.world_size))
+        else:
+            ranks.sort()
+
+        group: CommGroup = None
+        newcomm = self.pynccl.create_newcomm(ranks)
+        if newcomm:
+            group = CommGroup(newcomm.value, ranks)
+            self.sub_groups[newcomm.value] = group
+        return group

checkpoint_engine/ps.py

@@ -1,4 +1,5 @@
 import ctypes
+import gc
 import os
 import threading
 from collections import defaultdict
@@ -7,11 +8,12 @@ from datetime import timedelta
 from typing import TYPE_CHECKING
 
 import torch
-import torch.distributed as dist
+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 (
     BucketRange,
     DataToGather,
@@ -175,6 +177,8 @@ class ParameterServer:
         auto_pg: bool = True,
         gpu_count: int | None = None,
         mem_fraction: float | None = None,
+        master_addr: str | None = None,
+        master_port: int | None = None,
     ):
         """
         Initialize the parameter server. env RANK, WORLD_SIZE and MASTER_ADDR must be set.
@@ -228,6 +232,17 @@ class ParameterServer:
         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
 
+        master_addr = master_addr or os.getenv("MASTER_ADDR")
+        assert master_addr, "master_addr is required"
+        self._store = torch.distributed.TCPStore(
+            master_addr,
+            _get_master_port(master_port),
+            self._world_size,
+            timeout=timedelta(minutes=10),
+            is_master=self._rank == 0,
+        )
+        self._store_counter = 0
+
     def _get_memory_pool(self, checkpoint_name: str) -> list[MemoryBuffer]:
         if checkpoint_name == self._current_shared_memory_pool_user:
             assert self._memory_pool[self.shared_memory_pool_name], (
@@ -487,8 +502,6 @@ class ParameterServer:
     def init_process_group(
         self,
         *,
-        master_addr: str | None = None,
-        master_port: int | None = None,
         timeout: timedelta = timedelta(minutes=10),
     ):
         """
@@ -498,27 +511,18 @@ class ParameterServer:
             master_port: The specified port of the master node. If not set, will use _get_master_port to get the port.
             timeout: The timeout of the process group.
         """
-        master_addr = master_addr or os.getenv("MASTER_ADDR")
-        assert master_addr, "master_addr is required"
-        store = dist.TCPStore(
-            master_addr,
-            _get_master_port(master_port),
-            self._world_size,
-            timeout=timeout,
-            is_master=self._rank == 0,
-        )
+        self._store_counter += 1
+        sub_store = torch.distributed.PrefixStore(f"prefix-{self._store_counter}", self._store)
         dist.init_process_group(
             backend=self.device_manager.backend,
             world_size=self._world_size,
             rank=self._rank,
             timeout=timeout,
-            store=store,
+            store=sub_store,
         )
         logger.info(f"[rank{self._rank}] init process group successfully.")
 
-    def store_based_barrier(
-        self, store: dist.TCPStore, timeout: timedelta = timedelta(minutes=5)
-    ) -> None:
+    def store_based_barrier(self, timeout: timedelta = timedelta(minutes=5)) -> None:
         """
         Perform a store-based barrier synchronization across all ranks.
 
@@ -529,9 +533,9 @@ class ParameterServer:
         Args:
             store: The TCPStore instance to use for synchronization.
         """
-        dist.distributed_c10d._store_based_barrier(
+        torch.distributed.distributed_c10d._store_based_barrier(
             rank=self._rank,
-            store=store,
+            store=self._store,
             group_name="parameter_server_barrier",
             rendezvous_count=self._world_size,
             timeout=timeout,
@@ -544,8 +548,6 @@ class ParameterServer:
         *,
         timeout: timedelta = timedelta(minutes=10),
         ranks: list[int] | None = None,
-        master_addr: str | None = None,
-        master_port: int | None = None,
     ) -> None:
         """
         Update the checkpoint to inference engine. This function should be called after gather_metas.
@@ -566,28 +568,12 @@ class ParameterServer:
         assert req_func is not None, "req_func is required"
         ranks_group = None
         try:
-            master_addr = os.getenv("MASTER_ADDR") or master_addr
-            assert master_addr, "master_addr is required"
-            if self._auto_pg:
-                if not dist.is_initialized():
-                    self.init_process_group(
-                        timeout=timeout, master_addr=master_addr, master_port=master_port
-                    )
-                manager_store = dist.distributed_c10d._get_default_store()
-            else:
-                # HACK: MASTER_PORT+2 for barrier store if master_port is not provided, _get_master_port() returns MASTER_PORT+1
-                # If master_port is provided, use master_port+1 for barrier store
-                manager_store = dist.TCPStore(
-                    master_addr,
-                    _get_master_port(master_port) + 1,
-                    self._world_size,
-                    timeout=timeout,
-                    is_master=self._rank == 0,
-                )
+            if self._auto_pg and not dist.is_initialized():
+                self.init_process_group(timeout=timeout)
             # if ranks is None or [], it will use fully broadcast to update to all ranks
             ranks_group = dist.new_group(ranks) if ranks else None
             self._update_per_bucket(checkpoint_name, req_func, ranks_group, ranks)
-            self.store_based_barrier(manager_store)
+            self.store_based_barrier()
         except Exception as e:
             logger.exception(
                 f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} error {e}"
@@ -616,7 +602,10 @@ class ParameterServer:
         return socket, socket_paths
 
     def _detect_bucket_size(
-        self, ranks_group: dist.ProcessGroup | None, *, disable_h2d_buffer: bool = False
+        self,
+        ranks_group: dist.DistributedProcessGroup | None,
+        *,
+        disable_h2d_buffer: bool = False,
     ) -> tuple[int, bool]:
         GiB = 1 << 30  # noqa: N806
         # auto detect bucket size
@@ -633,7 +622,7 @@ class ParameterServer:
             dtype=torch.int64,
             device=self.device_manager.device_type,
         )
-        dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=ranks_group)
+        dist.all_reduce(tensor, op=torch.distributed.ReduceOp.MIN, group=ranks_group)
         tensor = tensor.cpu()
         free_bytes, self._zmq_addr_counter = tensor[0].item(), -tensor[1].item()
         max_tensor_bytes = 0
@@ -735,7 +724,7 @@ class ParameterServer:
         self,
         checkpoint_name: str,
         req_func: Callable[[list[tuple[str, str]]], None],
-        ranks_group: dist.ProcessGroup | None,
+        ranks_group: dist.DistributedProcessGroup | None,
         ranks: list[int] | None = None,
     ):
         assert len(self._current_global_parameter_metas) != 0, "parameter metas is empty"
@@ -854,7 +843,7 @@ class ParameterServer:
                             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=dist.ReduceOp.SUM, group=ranks_group)
+                    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
@@ -864,6 +853,29 @@ class ParameterServer:
                     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:

checkpoint_engine/worker.py

@@ -10,6 +10,9 @@ import zmq
 from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid
 
 
+_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)
@@ -29,11 +32,9 @@ class FlattenedTensorMetadata(TypedDict):
     offset: int
 
 
-def _extract_weights(
-    payload: list[FlattenedTensorMetadata], buffer: torch.Tensor
-) -> list[tuple[str, torch.Tensor]]:
+def _extract_weights(payload: list[FlattenedTensorMetadata], buffer: torch.Tensor) -> _WEIGHTS_TYPE:
     assert buffer is not None
-    weights: list[tuple[str, torch.Tensor]] = []
+    weights: _WEIGHTS_TYPE = []
     for item in payload:
         shape = item["shape"]
         if isinstance(shape, list | tuple):
@@ -69,15 +70,35 @@ def update_weights_from_ipc(
         socket.send_string(msg)
         socket.recv()  # wait for ack
         raise
+    # State machine:
+    # + receive tensor_metadata -> update_weights
+    # + receive Exception -> raise and stop
+    # + receive None first time -> release resources
+    # + receive None second time -> call post_hook and stop
     try:
+        released = False
         while True:
             payload: list[FlattenedTensorMetadata] | Exception | None = socket.recv_pyobj()
-            if payload is None:  # done signal
+            if released:
+                assert payload is None, "Should not receive any payload after released"
                 if post_hook is not None:
                     post_hook()
                 device_manager.device_module.synchronize()
                 socket.send(b"")
                 break
+            if payload is None:  # done signal
+                # TODO: wrap all messages to an object instead of None and Exception
+                device_manager.device_module.synchronize()
+                released = True
+                buffer = None
+                del ipc_handle
+
+                gc.collect()
+                device_manager.device_module.ipc_collect()
+                device_manager.device_module.empty_cache()
+                device_manager.device_module.synchronize()
+                socket.send(b"")
+                continue
             if isinstance(payload, list):  # still updating weights
                 try:
                     run(_extract_weights(payload, buffer))
@@ -166,12 +187,31 @@ class VllmColocateWorkerExtension:
             self.device = torch.device(f"npu:{self.local_rank}")
         assert self.device is not None
 
+        def _load_weights(weights: _WEIGHTS_TYPE):
+            # Load main model weights
+            self.model_runner.model.load_weights(weights)
+            # Load drafter model weights if MTP/speculative decoding is enabled
+            if (
+                getattr(self.model_runner, "drafter", None) is not None
+                and getattr(self.model_runner.drafter, "model", None) is not None
+            ):
+                self.model_runner.drafter.model.load_weights(weights=weights)
+
+        def _post_hook():
+            process_weights_after_loading(self.model_runner.model, self.model_config, self.device)
+            # Also trigger drafter model's post processing if MTP is enabled
+            if (
+                getattr(self.model_runner, "drafter", None) is not None
+                and getattr(self.model_runner.drafter, "model", None) is not None
+            ):
+                process_weights_after_loading(
+                    self.model_runner.drafter.model, self.model_config, self.device
+                )
+
         update_weights_from_ipc(
             self._zmq_ctx,
             zmq_handles[self._device_uuid],
             device_id=self.device.index,
-            run=self.model_runner.model.load_weights,
-            post_hook=lambda: process_weights_after_loading(
-                self.model_runner.model, self.model_config, self.device
-            ),
+            run=_load_weights,
+            post_hook=_post_hook,
         )

{checkpoint_engine-0.3.3.dist-info → checkpoint_engine-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.3.3
+Version: 0.4.0
 Summary: checkpoint-engine is a lightweight, decoupling and efficient weight update middleware
 Project-URL: Homepage, https://github.com/MoonshotAI/checkpoint-engine
 Project-URL: Repository, https://github.com/MoonshotAI/checkpoint-engine

checkpoint_engine-0.4.0.dist-info/RECORD

@@ -0,0 +1,19 @@
+checkpoint_engine/__init__.py,sha256=OeWxe9mxl2sZ6cW-blSTg6JbFlOMpGbBghLZtxGOqXk,942
+checkpoint_engine/__main__.py,sha256=yzQlApuYo6eIOqtqM018RosyxNzXzB5a-stxUvsh-dg,709
+checkpoint_engine/_version.py,sha256=2_0GUP7yBCXRus-qiJKxQD62z172WSs1sQ6DVpPsbmM,704
+checkpoint_engine/api.py,sha256=JDiQ4i3Gb6GoaBhlp8lNuUPaVURoFFdeGJY9ZDDGvPc,3518
+checkpoint_engine/data_types.py,sha256=O9uAXjwB20iwrOHfEEQd8Y9CmaFspNJ9ks9noHqwQKk,2716
+checkpoint_engine/device_utils.py,sha256=iKrof60j3CY3fStRTq3DRTt_kE1vYoEWHhAeyh0lByA,3020
+checkpoint_engine/p2p_store.py,sha256=abiCDVmRISPt9QFfavHB9Jo7ZpBbSjUS1NevGuB-AVA,8721
+checkpoint_engine/pin_memory.py,sha256=b7nABKJV2bSIsOfX2YTHzUk1OkOze6AQjCaOIFaQnbA,16708
+checkpoint_engine/ps.py,sha256=DQ9-hvZJW0eA9d6bU1glIbSwYl4cZCmVRjFPPRF41YY,41957
+checkpoint_engine/worker.py,sha256=fTWiF6Gggehzjx4mnIFTDZFR-GwkEBUdTAC_ZLmsgZE,8649
+checkpoint_engine/distributed/__init__.py,sha256=fC0EEX1nfWkg8OolzAj5vd2P0x6s4hScOlwV8q8Uiik,492
+checkpoint_engine/distributed/base.py,sha256=dpdjcGXNdCdAUDPnX-vdJmCGXbGS4A69yNsd60t-UgA,7800
+checkpoint_engine/distributed/vllm_hccl.py,sha256=bLE-GrnOxu1GTw_2GIqu2o67_Sw7vgjzJnlMvvQz_8c,10313
+checkpoint_engine/distributed/vllm_nccl.py,sha256=nHnlY1jk--xNEjKDDnywx36FgrnjEGc9lrBBC3o-YzE,7015
+checkpoint_engine-0.4.0.dist-info/licenses/LICENCE,sha256=D3gPmHKpGtF1yxYNhqjtBtZY_brZjDotJTzpnmClzlY,1067
+checkpoint_engine-0.4.0.dist-info/METADATA,sha256=qtS4bAI6SC3nBatKqAi5EVkty2zkvEZOWfshpaswF6k,11559
+checkpoint_engine-0.4.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+checkpoint_engine-0.4.0.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
+checkpoint_engine-0.4.0.dist-info/RECORD,,

{checkpoint_engine-0.3.3.dist-info → checkpoint_engine-0.4.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

checkpoint_engine-0.3.3.dist-info/RECORD

@@ -1,15 +0,0 @@
-checkpoint_engine/__init__.py,sha256=OeWxe9mxl2sZ6cW-blSTg6JbFlOMpGbBghLZtxGOqXk,942
-checkpoint_engine/__main__.py,sha256=yzQlApuYo6eIOqtqM018RosyxNzXzB5a-stxUvsh-dg,709
-checkpoint_engine/_version.py,sha256=lemL_4Kl75FgrO6lVuFrrtw6-Dcf9wtXBalKkXuzkO4,704
-checkpoint_engine/api.py,sha256=JDiQ4i3Gb6GoaBhlp8lNuUPaVURoFFdeGJY9ZDDGvPc,3518
-checkpoint_engine/data_types.py,sha256=O9uAXjwB20iwrOHfEEQd8Y9CmaFspNJ9ks9noHqwQKk,2716
-checkpoint_engine/device_utils.py,sha256=iKrof60j3CY3fStRTq3DRTt_kE1vYoEWHhAeyh0lByA,3020
-checkpoint_engine/p2p_store.py,sha256=abiCDVmRISPt9QFfavHB9Jo7ZpBbSjUS1NevGuB-AVA,8721
-checkpoint_engine/pin_memory.py,sha256=b7nABKJV2bSIsOfX2YTHzUk1OkOze6AQjCaOIFaQnbA,16708
-checkpoint_engine/ps.py,sha256=wBsHu2qWy5oRBrvLc7aEOroG_j58UJoWT6lFH4ylMRk,41092
-checkpoint_engine/worker.py,sha256=ghj9d2u8hY_U2uiOZWIN2CqRNZH6PrzujT22fHUFBWI,6879
-checkpoint_engine-0.3.3.dist-info/licenses/LICENCE,sha256=D3gPmHKpGtF1yxYNhqjtBtZY_brZjDotJTzpnmClzlY,1067
-checkpoint_engine-0.3.3.dist-info/METADATA,sha256=WyyGLw1qrteQgRGOWhAm15NN2nzklTqw4iiQ9U2nYpQ,11559
-checkpoint_engine-0.3.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-checkpoint_engine-0.3.3.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
-checkpoint_engine-0.3.3.dist-info/RECORD,,