checkpoint-engine 0.3.0rc0__py3-none-any.whl → 0.3.1__py3-none-any.whl

checkpoint_engine/ps.py

@@ -1,142 +1,33 @@
-import argparse
-import concurrent.futures
 import ctypes
-import json
 import os
-import pickle
-import random
 import threading
-import time
 from collections import defaultdict
 from collections.abc import Callable
 from datetime import timedelta
-from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, NamedTuple
+from typing import TYPE_CHECKING
 
-import httpx
-import numpy as np
 import torch
 import torch.distributed as dist
 import zmq
 from loguru import logger
-from pydantic import BaseModel, PlainSerializer, PlainValidator, WithJsonSchema
-from safetensors.torch import _getdtype, safe_open
 from torch.multiprocessing.reductions import reduce_tensor
 
+from checkpoint_engine.data_types import (
+    BucketRange,
+    DataToGather,
+    H2DBucket,
+    MemoryBuffer,
+    MemoryBufferMetaList,
+    MemoryBufferMetas,
+    ParameterMeta,
+)
 from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid
+from checkpoint_engine.p2p_store import P2PStore
+from checkpoint_engine.pin_memory import _ALIGN_SIZE, _register_checkpoint
 
 
 if TYPE_CHECKING:
-    from typing import TypeVar
-
-    from typing_extensions import TypedDict
-
-    class FileMeta(TypedDict):
-        key: str  # parameter name
-        dtype: torch.dtype
-        shape: torch.Size
-        type: type
-        tp_concat_dim: int
-
-    T = TypeVar("T")
-
-
-def _dt_validate(value: Any) -> torch.dtype:
-    if isinstance(value, str):
-        if not value.startswith("torch."):
-            raise ValueError(f"dtype {value} should start with torch.")
-        try:
-            value = getattr(torch, value.split(".")[1])
-        except AttributeError as e:
-            raise ValueError(f"unknown dtype: {value}") from e
-    if not isinstance(value, torch.dtype):
-        raise TypeError(f"dtype {value} should be torch.dtype, got {type(value)}")
-    return value
-
-
-_TorchDtype = Annotated[
-    torch.dtype,
-    PlainValidator(_dt_validate),
-    PlainSerializer(lambda x: str(x), return_type=str),
-    WithJsonSchema({"type": "string"}, mode="serialization"),
-]
-
-
-def _size_validate(value: Any) -> torch.Size:
-    if isinstance(value, list | tuple):
-        return torch.Size(value)
-    if not isinstance(value, torch.Size):
-        raise TypeError(f"size {value} should be torch.Size, got {type(value)}")
-    return value
-
-
-_TorchSize = Annotated[
-    torch.Size,
-    PlainValidator(_size_validate),
-    PlainSerializer(lambda x: tuple(x), return_type=tuple),
-    WithJsonSchema({"type": "array", "items": {"type": "integer"}}, mode="serialization"),
-]
-
-
-def _tensor_validate(value: Any) -> torch.Tensor:
-    if isinstance(value, torch.Tensor):
-        return value
-    raise TypeError(f"tensor {value} should be torch.Tensor, got {type(value)}")
-
-
-_TorchTensor = Annotated[
-    torch.Tensor,
-    PlainValidator(_tensor_validate),
-]
-
-
-class ParameterMeta(BaseModel):
-    name: str
-    dtype: _TorchDtype
-    shape: _TorchSize
-    aligned_size: int
-
-
-class BucketRange(NamedTuple):
-    idx: int  # bucket_idx of MemoryBucket in memory_pool
-    offset: int
-    size: int
-
-
-class H2DBucket(BaseModel):
-    size: int
-    ranges: list[BucketRange]
-    items: list[ParameterMeta]
-
-
-class MemoryBufferMetas(BaseModel):
-    metas: list[ParameterMeta]
-    ptr: int
-    size: int
-
-
-class MemoryBuffer(BaseModel):
-    buffer: _TorchTensor
-    size: int
-    metas: list[ParameterMeta]
-
-
-class MemoryBufferMetaList(BaseModel):
-    p2p_store_addr: str | None
-    memory_buffer_metas_list: list[MemoryBufferMetas]
-    rdma_device: str
-
-
-class DataToGather(MemoryBufferMetaList):
-    host_ip: str
-    device_uuid: str
-
-
-# 256 bytes alignment when flatten torch tensors to uint8 buffer
-_ALIGN_SIZE = 256
-
-
-def _align_size(dtype: torch.dtype, shape: torch.Size) -> int:
-    return (dtype.itemsize * shape.numel() + _ALIGN_SIZE - 1) // _ALIGN_SIZE * _ALIGN_SIZE
+    from checkpoint_engine.data_types import T
 
 
 def _to_named_tensor(metas: list[ParameterMeta], offset: int = 0) -> list[dict]:
@@ -155,107 +46,6 @@ def _to_named_tensor(metas: list[ParameterMeta], offset: int = 0) -> list[dict]:
     return ret
 
 
-def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple["FileMeta", torch.Tensor]]]:
-    def _safetensors_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]:
-        ret = {}
-        with safe_open(fn, framework="pt") as f:
-            for name in f.keys():  # noqa: SIM118
-                weight = f.get_tensor(name)
-                meta = {
-                    "key": name,
-                    "dtype": weight.dtype,
-                    "shape": weight.shape,
-                    "type": type(weight),
-                    "tp_concat_dim": -1,  # safetensors does not support tp_concat_dim
-                }
-                ret[name] = (meta, weight)
-        return ret
-
-    # deprecated, will be removed in the future
-    def _fast_np_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]:
-        """load *.np file and return memmap and related tensor meta"""
-
-        def parse_npy_header(fin: BinaryIO) -> dict[str, Any]:
-            start = fin.tell()
-            major, minor = np.lib.format.read_magic(fin)
-            if major == 1 and minor == 0:
-                read_header_fn = np.lib.format.read_array_header_1_0
-            elif major == 2 and minor == 0:
-                read_header_fn = np.lib.format.read_array_header_2_0
-            else:
-                raise ValueError(
-                    f"unknown version {major}.{minor} when parsing npy header from {fn}"
-                )
-            shape, is_fortran, dtype = read_header_fn(fin)
-            return {
-                "shape": shape,
-                "is_fortran": is_fortran,
-                "dtype": dtype,
-                "header_length": fin.tell() - start,
-            }
-
-        meta_fn = fn + ".meta"
-        with open(meta_fn, "rb") as fin:
-            meta_lst = pickle.load(fin)
-
-        tensors = []
-        offset = 0
-        with open(fn, "rb") as fin:
-            fin.seek(0, os.SEEK_END)
-            filesize = fin.tell()
-            fin.seek(0)
-            while fin.tell() < filesize:
-                tensor_meta = parse_npy_header(fin)
-                tensor = np.memmap(
-                    fn,
-                    dtype=tensor_meta["dtype"],
-                    mode="c",
-                    offset=offset + tensor_meta["header_length"],
-                    shape=tensor_meta["shape"],
-                )
-                offset += tensor_meta["header_length"] + tensor.nbytes
-                fin.seek(offset)
-                tensors.append(tensor)
-
-        assert len(meta_lst) == len(tensors)
-        ret = {}
-        for meta, tensor in zip(meta_lst, tensors):
-            if meta["type"] == torch.Tensor:
-                tensor = torch.from_numpy(tensor)
-            tensor = tensor.view(dtype=meta["dtype"]).view(*meta["shape"])
-            ret[meta["key"]] = (meta, tensor)
-        return ret
-
-    tp_rank = 0
-    if file_path.endswith(".npy"):
-        logger.warning("numpy model file is deprecated, will be removed in the future")
-        filename_split = os.path.basename(file_path).split(".")
-        # if using numpy and want to specify tp rank
-        # file should be in model.{layer}.{tp}[.{ep}].npy format
-        tp_rank = int(filename_split[2]) if len(filename_split) > 3 else 0
-        ret = _fast_np_load(file_path)
-    elif file_path.endswith(".safetensors"):
-        ret = _safetensors_load(file_path)
-    else:
-        raise ValueError(f"unsupported file format: {file_path}")
-    return tp_rank, ret
-
-
-def _concat_tp_weights(
-    tp_weights: list[torch.Tensor], tp_concat_dim: int, tp_size: int
-) -> torch.Tensor:
-    """Concat tp weights with meta info.
-    If meta.concat_dim is -1, meas this is shared tp weights, just use the first weights.
-    Else we will cat weights in concat_dim.
-    """
-    if tp_concat_dim == -1:
-        return tp_weights[0]
-    assert tp_size == len(tp_weights)
-    if len(tp_weights) == 1:
-        return tp_weights[0]
-    return torch.cat([w for w in tp_weights], dim=tp_concat_dim)
-
-
 def _get_physical_gpu_id(device_manager: DeviceManager, device_index: int | None = None) -> str:
     try:
         if device_manager.device_type == "npu":
@@ -266,420 +56,6 @@ def _get_physical_gpu_id(device_manager: DeviceManager, device_index: int | None
         raise ValueError(f"fail to get physical gpu id {device_index}") from e
 
 
-def _ibv_get_device_list() -> list[str]:
-    lib = ctypes.CDLL("libibverbs.so.1")
-    lib.ibv_get_device_list.argtypes = [ctypes.POINTER(ctypes.c_int)]  # int *num_devices
-    lib.ibv_get_device_list.restype = ctypes.POINTER(ctypes.c_void_p)  # struct ibv_device **
-
-    lib.ibv_free_device_list.argtypes = [ctypes.POINTER(ctypes.c_void_p)]
-    lib.ibv_get_device_name.argtypes = [ctypes.c_void_p]  # struct ibv_device *
-    lib.ibv_get_device_name.restype = ctypes.c_char_p  # const char *
-
-    num = ctypes.c_int()
-    dev_array = lib.ibv_get_device_list(ctypes.byref(num))
-    if not dev_array or num.value <= 0:
-        return []
-
-    devices = []
-    for i in range(num.value):
-        dev_ptr = dev_array[i]  # struct ibv_device *
-        name = lib.ibv_get_device_name(dev_ptr)  # const char *
-        devices.append(name.decode())
-    lib.ibv_free_device_list(dev_array)
-    return devices
-
-
-def _get_rdma_devices() -> list[str]:
-    """
-    use _ibv_get_device_list to get RDMA devices, if NCCL_IB_HCA has multiple values, just return
-    """
-    devices_str = os.getenv("PS_P2P_STORE_RDMA_DEVICES")
-    if devices_str:
-        return devices_str.split(",")
-    # if PS_P2P_STORE_RDMA_DEVICES is not set, try to use NCCL_IB_HCA to get RDMA devices
-    hca = os.getenv("NCCL_IB_HCA", None)
-    return _parse_NCCL_IB_HCA(hca or "", _ibv_get_device_list()) or _ibv_get_device_list()
-
-
-def _get_my_rdma_device(local_rank: int, gpu_count: int, devices: list[str]) -> str:
-    """
-    implement network card device allocation, if network card is "mlx5_0,mlx5_1", then 0-3 will share mlx5_0, 4-7 will share mlx5_1, etc.
-    """
-    if not devices:
-        raise RuntimeError("no rdma devices found")
-    try:
-        assert len(devices) <= gpu_count, (
-            f"rdma devices count {len(devices)} should be less than or equal to gpu count {gpu_count}"
-        )
-        assert gpu_count % len(devices) == 0, (
-            f"gpu count {gpu_count} should be divisible by rdma devices count {len(devices)}"
-        )
-        return devices[local_rank // (gpu_count // len(devices))]
-    except AssertionError:
-        logger.error(
-            "Please set 'NCCL_IB_HCA' or 'PS_P2P_STORE_RDMA_DEVICES' environment variable to choose proper number of RDMA devices."
-            "The number of RDMA devices should be less than or equal to GPU count, and GPU count should be divisible by the number of RDMA devices."
-            "The acceptable value by NCCL_IB_HCA is documented in 'https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8'."
-        )
-        raise
-
-
-def _parse_NCCL_IB_HCA(value: str, available_devices: list[str]) -> list[str]:
-    """
-    The acceptable value by NCCL_IB_HCA is documented in https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8.
-    The Python version parser is referred to the CPP parser in NCCL: https://github.com/NVIDIA/nccl/blob/v2.28.3-1/src/transport/net_ib.cc#L658-L662.
-
-    The list is comma-separated; port numbers are NOT supported yet.
-    An optional prefix '^' indicates the list is an exclude list.
-    A second optional prefix '=' indicates that the tokens are exact names, otherwise by default NCCL would treat each token as a prefix.
-    Please note that when '^' and '=' appear together, only '^=' is allowed, '=^' is not supported.
-
-    Examples:
-    - `NCCL_IB_HCA="mlx5"`: Use all cards starting with `mlx5`.
-    - `NCCL_IB_HCA="=mlx5_0,mlx5_1"`: Use specific cards `mlx5_0` and `mlx5_1`.
-    - `NCCL_IB_HCA="^mlx5"`: Use all cards except those starting with `mlx5`.
-    - `NCCL_IB_HCA="^=mlx5_0,mlx5_1"`: Use all cards except `mlx5_0` and `mlx5_1`.
-    """
-    max_hcas = 32
-    if not value or value.strip() == "":
-        return available_devices[:max_hcas]
-
-    value = value.strip()
-    result = []
-    is_exclude = value.startswith("^")
-    if is_exclude:
-        value = value.removeprefix("^")
-    is_exact_match = value.startswith("=")
-    if is_exact_match:
-        value = value.removeprefix("=")
-
-    device_specs = [spec.strip() for spec in value.split(",") if spec.strip()]
-
-    result = _resolve_device_specs(device_specs, is_exact_match, available_devices)
-    if is_exclude:
-        result = [dev for dev in available_devices if dev not in result]
-    if len(result) > max_hcas:
-        result = result[:max_hcas]
-
-    logger.info(f"RDMA Devices from 'NCCL_IB_HCA': {result}")
-
-    return result
-
-
-def _resolve_device_specs(
-    device_specs: list[str], is_exact_match: bool, available_devices: list[str]
-) -> list[str]:
-    devices = set()
-    for spec in device_specs:
-        parts = spec.split(":", 1)
-        device_name = parts[0].strip()
-        # HACK: mooncake transfer engine does not support port specification yet, so we ignore it
-        # port = parts[1].strip() if len(parts) > 1 else None
-        base_devices = (
-            [device_name]
-            if device_name in available_devices
-            else []
-            if is_exact_match
-            else [dev for dev in available_devices if dev.startswith(device_name)]
-        )
-
-        if not base_devices:
-            logger.warning(f"No RDMA device match {device_name=} where {is_exact_match=}.")
-            continue
-
-        for base_dev in base_devices:
-            devices.add(base_dev)
-
-    return sorted(devices)
-
-
-def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
-    class TPMeta(BaseModel):
-        concat_dim: int
-        size: int
-
-    parameters: dict[str, torch.Tensor] = {}
-    parameter_metas: dict[str, ParameterMeta] = {}
-    tp_metas: dict[str, TPMeta] = {}
-    parameters_with_tp: dict[str, dict[int, torch.Tensor]] = {}
-    for file in files:
-        tp_rank, ret = _load_checkpoint_file(file)
-        for parameter_name, (meta, weight) in ret.items():
-            if parameter_name not in parameters_with_tp:
-                parameters_with_tp[parameter_name] = {}
-            parameters_with_tp[parameter_name][tp_rank] = weight
-            if parameter_name not in tp_metas:
-                tp_metas[parameter_name] = TPMeta(
-                    concat_dim=meta["tp_concat_dim"],
-                    size=1,
-                )
-            if parameter_name not in parameter_metas:
-                assert isinstance(meta["dtype"], torch.dtype), (
-                    f"meta {meta} dtype should be torch.dtype"
-                )
-                assert isinstance(meta["shape"], torch.Size), (
-                    f"meta {meta} shape should be torch.Size"
-                )
-                parameter_metas[parameter_name] = ParameterMeta(
-                    name=parameter_name,
-                    shape=meta["shape"],
-                    dtype=meta["dtype"],
-                    aligned_size=_align_size(meta["dtype"], meta["shape"]),
-                )
-            tp_meta = tp_metas[parameter_name]
-            if tp_meta.concat_dim != -1:
-                tp_meta.size = max(tp_meta.size, tp_rank + 1)
-    for name, tp_meta in tp_metas.items():
-        if tp_meta.concat_dim != -1:
-            shape = list(parameter_metas[name].shape)
-            shape[tp_meta.concat_dim] = shape[tp_meta.concat_dim] * tp_meta.size
-            parameter_metas[name] = ParameterMeta(
-                name=name,
-                shape=torch.Size(shape),
-                dtype=parameter_metas[name].dtype,
-                aligned_size=_align_size(parameter_metas[name].dtype, torch.Size(shape)),
-            )
-        weights_in_cpu = [parameters_with_tp[name][key] for key in sorted(parameters_with_tp[name])]
-        # TODO: here concat is serial, which may be slow
-        # but since tp storage is not used in the future
-        # we ignore this performance issue for now
-        parameters[name] = _concat_tp_weights(weights_in_cpu, tp_meta.concat_dim, tp_meta.size)
-    for name, parameter in parameters.items():
-        assert name in parameter_metas, f"parameter {name} not found in parameter_metas"
-        assert parameter_metas[name].shape == parameter.shape, (
-            f"parameter {name} shape mismatch, {parameter_metas[name].shape} != {parameter.shape}"
-        )
-        assert parameter_metas[name].dtype == parameter.dtype, (
-            f"parameter {name} dtype mismatch, {parameter_metas[name].dtype} != {parameter.dtype}"
-        )
-    return parameters
-
-
-def _inplace_pin_memory(files: list[str], rank: int | None = None) -> list[MemoryBuffer]:
-    def _parse_and_pin_from_safetensors(file_path: str) -> MemoryBuffer:
-        """
-        safetensors format see https://huggingface.co/docs/safetensors/en/index#format.
-        We load the safetensors file as bytes, then parse the header manually to get parameter metas.
-        The actual tensor data is in the remaining bytes and is naturally aligned.
-        We pin the remaining bytes as the buffer, making pinning faster.
-        """
-
-        def _pin(t: torch.Tensor):
-            """
-            Pin the memory of tensor in-place.
-            See: https://github.com/pytorch/pytorch/issues/32167
-            """
-            cudart = torch.cuda.cudart()
-            r = cudart.cudaHostRegister(t.data_ptr(), t.numel() * t.element_size(), 0)
-            assert r == 0, f"pin memory error, error code: {r}"
-
-        # TODO: should only support /dev/shm? but we found files in disk also work?
-        size = os.stat(file_path).st_size
-        flag_size = 8
-        t = torch.from_file(file_path, True, size, dtype=torch.uint8)
-        assert t.nbytes > flag_size, (
-            f"tensor nbytes {t.nbytes} should be greater than flag_size {flag_size}"
-        )
-        start_pos = (
-            int.from_bytes(t[0:flag_size].numpy().tobytes(), byteorder="little", signed=False)
-            + flag_size
-        )
-        header_tensor = t[flag_size:start_pos]
-        header = json.loads(header_tensor.numpy().tobytes())
-        if "__metadata__" in header:
-            header.pop("__metadata__")
-
-        metas: list[ParameterMeta] = []
-        offset = 0
-        try:
-            for name, meta in sorted(header.items(), key=lambda x: x[1]["data_offsets"]):
-                start, end = meta["data_offsets"]
-                # safetensors format ensures offsets are aligned
-                assert offset == start, f"offset {offset} should be equal to start {start}"
-                metas.append(
-                    ParameterMeta(
-                        name=name,
-                        dtype=_getdtype(meta["dtype"]),
-                        shape=torch.Size(meta["shape"]),
-                        aligned_size=end - start,
-                    )
-                )
-                offset = end
-        except Exception as e:
-            logger.error(f"fail to parse safetensors header from {file_path}: {e}")
-            raise
-
-        buffer = t[start_pos:]
-        assert offset == buffer.nbytes, (
-            f"offset {offset} should be equal to buffer.nbytes {buffer.nbytes}"
-        )
-        # Remove the file after successfully loading. This will avoid doubling the memory usage.
-        # We assume files in /dev/shm/ are temporary files. So it's safe to remove them after loading.
-        os.remove(file_path)
-        _pin(buffer)
-        logger.info(
-            f"[rank{rank}] inplace pin memory for file {file_path} finished, size {buffer.nbytes / 1024 / 1024:.2f}MiB"
-        )
-        return MemoryBuffer(buffer=buffer, size=buffer.nbytes, metas=metas)
-
-    memory_buffers: list[MemoryBuffer] = []
-    with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
-        memory_buffers = list(executor.map(_parse_and_pin_from_safetensors, files))
-    return memory_buffers
-
-
-def _normal_pin_memory(
-    files: list[str],
-    named_tensors: dict[str, torch.Tensor],
-    rank: int | None = None,
-    shared_pin_memory: list[MemoryBuffer] | None = None,
-) -> list[MemoryBuffer]:
-    parameters = _load_checkpoint(files)
-    if named_tensors:
-        parameters.update(named_tensors)
-    bucket_size = max(4 << 30, max(_align_size(x.dtype, x.shape) for x in parameters.values()))
-
-    class MemoryBucket(BaseModel):
-        size: int
-        metas: list[ParameterMeta]
-
-    buckets: list[MemoryBucket] = []
-    buckets.append(MemoryBucket(size=0, metas=[]))
-    for name, tensor in sorted(parameters.items()):
-        size = _align_size(tensor.dtype, tensor.shape)
-        if buckets[-1].size + size > bucket_size:
-            assert buckets[-1], f"buckets[{len(buckets) - 1}] should not be empty"
-            buckets.append(MemoryBucket(size=0, metas=[]))
-        buckets[-1].metas.append(
-            ParameterMeta(name=name, shape=tensor.shape, dtype=tensor.dtype, aligned_size=size)
-        )
-        buckets[-1].size += size
-
-    memory_buffers = [
-        MemoryBuffer(buffer=torch.empty(0), size=bucket.size, metas=bucket.metas)
-        for bucket in buckets
-    ]
-
-    def register_pin_memory(
-        idx: int, size: int, shared_pin_memory: list[MemoryBuffer] | None = None
-    ) -> tuple[int, torch.Tensor]:
-        if shared_pin_memory:
-            # If shared_pin_memory is provided, reuse the pin memory buffer, do not allocate new one
-            # Reusing pin memory only support fixed shape of checkpoints, which is registered the first time
-            assert idx < len(shared_pin_memory), (
-                f"idx {idx} should be less than shared_pin_memory length {len(shared_pin_memory)}"
-            )
-            assert shared_pin_memory[idx].size == size, (
-                f"shared_pin_memory[{idx}].size {shared_pin_memory[idx].size} should be equal to {size}"
-            )
-            return idx, shared_pin_memory[idx].buffer
-        else:
-            buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
-            return idx, buffer
-
-    def register_tensor(buffer: torch.Tensor, offset: int, tensor: torch.Tensor):
-        buffer[offset : offset + tensor.nbytes] = tensor.view(-1).view(dtype=torch.uint8)
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
-        futures = [
-            executor.submit(
-                register_pin_memory,
-                idx,
-                bucket.size,
-                shared_pin_memory,
-            )
-            for idx, bucket in enumerate(buckets)
-        ]
-        new_futures = []
-        for future in concurrent.futures.as_completed(futures):
-            idx, buffer = future.result()
-            assert buffer.numel() == buckets[idx].size, (
-                f"buffer numel {buffer.numel()} should be equal to bucket size {buckets[idx].size}"
-            )
-            memory_buffers[idx].buffer = buffer
-            logger.info(
-                f"[rank{rank}] register pin_memory for bucket {idx + 1}/{len(buckets)} finished, "
-                f"size {buffer.numel() / 1024 / 1024:.2f}MiB, start to copy tensors to buffer"
-            )
-            offset = 0
-            for meta in buckets[idx].metas:
-                name = meta.name
-                tensor = parameters[name]
-                size = _align_size(tensor.dtype, tensor.shape)
-                assert size == _align_size(meta.dtype, meta.shape), (
-                    f"tensor {name} size {size} should be equal to meta size {_align_size(meta.dtype, meta.shape)}"
-                )
-                new_futures.append(executor.submit(register_tensor, buffer, offset, tensor))
-                offset += size
-        for future in concurrent.futures.as_completed(new_futures):
-            future.result()
-        return memory_buffers
-
-
-def _register_checkpoint(
-    *,
-    files: list[str],
-    named_tensors: dict[str, torch.Tensor],
-    rank: int | None = None,
-    shared_pin_memory: list[MemoryBuffer] | None = None,
-) -> list[MemoryBuffer]:
-    logger.info(
-        f"[rank{rank}] start to register checkpoint with {len(files)} files and {len(named_tensors)} named_tensors"
-    )
-    if not files and not named_tensors:
-        return []
-    memory_buffers: list[MemoryBuffer] = []
-    files_to_inplace_pin = [
-        file
-        for file in files
-        if file.startswith("/dev/shm/") and file.endswith(".safetensors")  # noqa: S108
-    ]
-    files_to_normal_pin = [file for file in files if file not in files_to_inplace_pin]
-    if files_to_normal_pin or named_tensors:
-        memory_buffers.extend(
-            _normal_pin_memory(
-                files=files_to_normal_pin,
-                named_tensors=named_tensors,
-                rank=rank,
-                shared_pin_memory=shared_pin_memory,
-            )
-        )
-    if files_to_inplace_pin:
-        memory_buffers.extend(_inplace_pin_memory(files_to_inplace_pin, rank=rank))
-    return memory_buffers
-
-
-def request_inference_to_update(
-    url: str,
-    socket_paths: dict[str, str],
-    timeout: float = 300.0,
-    uds: str | None = None,
-):
-    """Send an inference update request to inference server via HTTP or Unix socket.
-
-    Args:
-        url (str): The HTTP URL or request path (e.g., "http://localhost:19730/inference") to send the request to.
-        socket_paths (dict[str, str]): A dictionary containing device uuid and IPC socket paths for updating weights.
-        timeout (float, optional): Request timeout in seconds. Defaults to 300.0.
-        uds (str, optional): Path to a Unix domain socket. If provided, the request
-            will be sent via the Unix socket instead of HTTP. Defaults to None.
-
-    Raises:
-        httpx.HTTPStatusError: If the response contains an HTTP error status.
-        httpx.RequestError: If there was an issue while making the request.
-    """
-    resp = httpx.Client(transport=httpx.HTTPTransport(uds=uds)).post(
-        url,
-        json={
-            "method": "update_weights_from_ipc",
-            "args": [socket_paths],
-            "timeout": timeout,
-        },
-        timeout=timeout,
-    )
-    resp.raise_for_status()
-
-
 def _gen_h2d_buckets(
     global_metas: dict[int, MemoryBufferMetaList],
     bucket_size: int,
@@ -782,84 +158,12 @@ def _get_master_port(master_port: int | None = None) -> int:
     if master_port is None:
         # HACK: use MASTER_PORT + 1 as master_port, avoid conflict with torchrun's rendezvous port
         # TODO: check whether master_port is available or use a more elegant way
-        master_port = int(os.getenv("MASTER_PORT")) + 1
+        master_port_str = os.getenv("MASTER_PORT")
+        assert master_port_str, "MASTER_PORT is required if no master_port is provided."
+        master_port = int(master_port_str) + 1
     return master_port
 
 
-class P2PStore:
-    def __init__(self, device_manager: DeviceManager):
-        from mooncake.engine import TransferEngine
-
-        self.rank = int(os.getenv("RANK"))
-        gpu_count = device_manager.device_module.device_count()
-        local_rank = self.rank % gpu_count
-        device_type = device_manager.device_type
-        if device_type == "npu" and os.getenv("PS_P2P_STORE_RDMA_DEVICES") is None:
-            self.device = ""
-        else:
-            self.device = _get_my_rdma_device(local_rank, gpu_count, _get_rdma_devices())
-        self.ip = get_ip()
-
-        # we will start at most 8 ps processes, so we use 8 retries to avoid port conflicts in extreme cases
-        retry_count = 8
-        for i in range(retry_count):
-            self.engine = TransferEngine()
-            ret = self.engine.initialize(
-                self.ip,
-                "P2PHANDSHAKE",
-                "ascend_direct" if device_type == "npu" else "rdma",
-                self.device,
-            )
-            if ret == 0:
-                break
-            # sleep 0.5 ~ 2.0s, to avoid port conflicts when two processes retry at the same time
-            sleep_ms = random.randint(500, 2000)
-            logger.warning(
-                f"[rank{self.rank}] fail to initialize transfer engine, ret {ret}, retry {i + 1}/{retry_count} in {sleep_ms}ms"
-            )
-            time.sleep(sleep_ms / 1000)
-        else:
-            raise RuntimeError(f"[rank{self.rank}] fail to initialize transfer engine")
-        self.port = self.engine.get_rpc_port()
-        self.named_tensors: dict[str, torch.Tensor] = {}
-        logger.info(
-            f"[rank{self.rank}] p2p store initialized, addr is {self.addr}, rdma device is {self.device}"
-        )
-
-    @property
-    def addr(self) -> str:
-        return f"{self.ip}:{self.port}"
-
-    def register_named_tensors(self, named_tensors: dict[str, torch.Tensor]):
-        buffer_addresses = [tensor.data_ptr() for tensor in named_tensors.values()]
-        capacities = [tensor.nbytes for tensor in named_tensors.values()]
-        self.named_tensors.update(named_tensors)
-        for i, name in enumerate(named_tensors.keys()):
-            logger.info(
-                f"[rank{self.rank}] p2p store register tensor {name} with addr {hex(buffer_addresses[i])} and capacity {capacities[i]}"
-            )
-        assert self.engine.batch_register_memory(buffer_addresses, capacities) == 0
-
-    def unregister_named_tensors(self, names: list[str]) -> int:
-        buffer_addresses = [self.named_tensors[name].data_ptr() for name in names]
-        assert self.engine.batch_unregister_memory(buffer_addresses) == 0
-        num_unregistered = 0
-        for i, name in enumerate(names):
-            del self.named_tensors[name]
-            logger.info(
-                f"[rank{self.rank}] p2p store unregister tensor {name} with addr {hex(buffer_addresses[i])}"
-            )
-            num_unregistered += 1
-        return num_unregistered
-
-    def batch_transfer_sync_read(
-        self, target_hostname: str, buf_ptrs: list[int], remote_ptrs: list[int], lens: list[int]
-    ):
-        assert (
-            self.engine.batch_transfer_sync_read(target_hostname, buf_ptrs, remote_ptrs, lens) == 0
-        )
-
-
 class ParameterServer:
     shared_memory_pool_name = "__shared_memory_pool__"
 
@@ -868,7 +172,7 @@ class ParameterServer:
         *,
         rank: int | None = None,
         world_size: int | None = None,
-        auto_pg: bool = False,
+        auto_pg: bool = True,
         gpu_count: int | None = None,
         mem_fraction: float | None = None,
     ):
@@ -877,11 +181,11 @@ class ParameterServer:
 
         Args:
             auto_pg: Whether to automatically initialize the process group.
-                Notice that if auto_pg is True, will destroy the process group after update.
+                Notice that if auto_pg is True, will destroy the process group after update. It is recommended to set auto_pg to True!
             mem_fraction: The proportion (as a fraction) of the current free device memory for allocation.
         """
-        self._rank = rank or int(os.environ.get("RANK", None))
-        self._world_size = world_size or int(os.environ.get("WORLD_SIZE", None))
+        self._rank = rank or int(os.environ["RANK"])
+        self._world_size = world_size or int(os.environ["WORLD_SIZE"])
         self.device_manager = DeviceManager()
         self._gpu_count = gpu_count or self.device_manager.device_module.device_count()
         self._local_rank = self._rank % self._gpu_count
@@ -890,7 +194,7 @@ class ParameterServer:
         self._global_device_uuids: list[str] = []
         self._local_rdma_devices: dict[str, set[int]] = defaultdict(set)
         self._remote_rdma_devices: dict[str, set[int]] = defaultdict(set)
-        self._mem_fraction = mem_fraction or 0.9
+        self._mem_fraction = mem_fraction or float(os.getenv("PS_MEM_FRACTION", "0.9"))
 
         assert self._rank is not None and self._rank >= 0, self._rank
         assert self._world_size and self._world_size > 0, self._world_size
@@ -959,11 +263,12 @@ class ParameterServer:
         files: list[str] | None = None,
         named_tensors: dict[str, torch.Tensor] | None = None,
         use_shared_memory_pool: bool = False,
+        use_inplace_pin_memory: bool = True,
     ) -> None:
         """
         Register a checkpoint to the parameter server. Both files and named_tensors will be registered together.
-        Warning: .safetensors files in /dev/shm/ will be pinned in-place, and the files will be REMOVED after pinning.
-        Please make sure to copy the files to disks if you need to keep them.
+        Warning: if `use_inplace_pin_memory` is True, .safetensors files in /dev/shm/ will be pinned in-place, and the files will be REMOVED after pinning.
+        Please make sure to copy the files to disks if you need to keep them. NPU does not support inplace pin memory.
 
         Args:
             checkpoint_name: The name of the checkpoint.
@@ -974,7 +279,14 @@ class ParameterServer:
                 cannot accommodate checkpoints with different memory requirements.
                 To free the actual memory of the shared pool or to modify its shape,
                 please unregister the current user of the shared memory pool using `unregister_checkpoint` with `force=True`.
+            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:
+            logger.warning(
+                f"[rank{self._rank}] Only cuda devices support in-place pin memory, set use_inplace_pin_memory to False"
+            )
+            use_inplace_pin_memory = False
         try:
             if use_shared_memory_pool:
                 logger.info(
@@ -993,6 +305,7 @@ class ParameterServer:
                     named_tensors=named_tensors or {},
                     rank=self._rank,
                     shared_pin_memory=self._memory_pool[self.shared_memory_pool_name],
+                    inplace_pin=False,  # inplace pin memory is not compatible with shared memory pool
                 )
                 self._current_shared_memory_pool_user = checkpoint_name
                 if self._p2p_store is not None and _is_first_time:
@@ -1002,7 +315,10 @@ class ParameterServer:
                     f"checkpoint {checkpoint_name} already registered"
                 )
                 self._memory_pool[checkpoint_name] = _register_checkpoint(
-                    files=files or [], named_tensors=named_tensors or {}, rank=self._rank
+                    files=files or [],
+                    named_tensors=named_tensors or {},
+                    rank=self._rank,
+                    inplace_pin=use_inplace_pin_memory,
                 )
                 if self._p2p_store is not None:
                     self._register_parameters_to_p2p_store(checkpoint_name)
@@ -1048,6 +364,46 @@ class ParameterServer:
             del self._memory_pool[self.shared_memory_pool_name]
             self._memory_pool[self.shared_memory_pool_name] = []
         else:
+
+            def _unpin(t: torch.Tensor):
+                """
+                Un-pin the pinned memory.
+                """
+                p_flags = ctypes.c_uint()
+                try:
+                    libc = ctypes.CDLL(None)  # get all symbols from the current process
+                    cuda_host_get_flags = libc.cudaHostGetFlags
+                    cuda_host_get_flags.argtypes = [ctypes.POINTER(ctypes.c_uint), ctypes.c_void_p]
+                    cuda_host_get_flags.restype = ctypes.c_int
+                except AttributeError:
+                    logger.error("cudaHostGetFlags not found in libc, cannot unpin memory manually")
+                    raise
+                r = cuda_host_get_flags(ctypes.byref(p_flags), ctypes.c_void_p(t.data_ptr()))
+                assert r == 0, f"get pin flags error, error code: {r}"
+                # p_flags value meaning from cuda/include/driver_types.h
+                # cudaHostRegisterDefault             0x00  /**< Default host memory registration flag */
+                # cudaHostRegisterPortable            0x01  /**< Pinned memory accessible by all CUDA contexts */
+                # cudaHostRegisterMapped              0x02  /**< Map registered memory into device space */
+                # cudaHostRegisterIoMemory            0x04  /**< Memory-mapped I/O space */
+                # cudaHostRegisterReadOnly            0x08  /**< Memory-mapped read-only */
+                assert p_flags.value == 0x02, (
+                    f"pin memory flag error, expected: 0x02 (cudaHostRegisterMapped), got flag: {p_flags.value}"
+                )
+                cudart = torch.cuda.cudart()
+                r = cudart.cudaHostUnregister(t.data_ptr())
+                assert r == 0, f"unpin memory error, error code: {r}"
+
+            # if the checkpoint is pinned by cudaHostRegister manually, we need to unpin it manually
+            try:
+                for memory_buffer in self._memory_pool.get(checkpoint_name, []):
+                    if memory_buffer.manually_pinned:
+                        _unpin(memory_buffer.buffer)
+            except Exception as e:
+                logger.error(
+                    f"[rank{self._rank}] fail to unpin memory for checkpoint {checkpoint_name}: {e}"
+                )
+                raise
+            # 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
@@ -1183,6 +539,8 @@ class ParameterServer:
     ) -> None:
         """
         Update the checkpoint to inference engine. This function should be called after gather_metas.
+        Warning: if _auto_pg is False when initializing ParameterServer, please make sure ALL ranks in the WORLD_SIZE call `update` function,
+        otherwise, it will hang.
 
         Args:
             checkpoint_name: The name of the checkpoint.
@@ -1217,7 +575,7 @@ class ParameterServer:
                     is_master=self._rank == 0,
                 )
             # 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)
+            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)
         except Exception as e:
@@ -1248,7 +606,7 @@ class ParameterServer:
         return socket, socket_paths
 
     def _detect_bucket_size(
-        self, ranks_group: dist.ProcessGroup, *, disable_h2d_buffer: bool = False
+        self, ranks_group: dist.ProcessGroup | None, *, disable_h2d_buffer: bool = False
     ) -> tuple[int, bool]:
         GiB = 1 << 30  # noqa: N806
         # auto detect bucket size
@@ -1291,7 +649,7 @@ class ParameterServer:
                 f"max_tensor_bytes {max_tensor_bytes} should be less than free_bytes {free_bytes}"
             )
             disable_h2d_buffer = True
-        max_bytes = int(os.getenv("PS_MAX_BUCKET_SIZE_GB", 8)) * GiB
+        max_bytes = int(float(os.getenv("PS_MAX_BUCKET_SIZE_GB", "8")) * GiB)
         bucket_size = min(max(max_bytes, max_tensor_bytes), free_bytes)
         logger.info(f"[rank{self._rank}] auto detect bucket size {bucket_size / GiB:.2f} GiB")
         return bucket_size, disable_h2d_buffer
@@ -1367,7 +725,7 @@ class ParameterServer:
         self,
         checkpoint_name: str,
         req_func: Callable[[list[tuple[str, str]]], None],
-        ranks_group: dist.ProcessGroup,
+        ranks_group: dist.ProcessGroup | None,
         ranks: list[int] | None = None,
     ):
         assert len(self._current_global_parameter_metas) != 0, "parameter metas is empty"
@@ -1498,79 +856,8 @@ class ParameterServer:
             self.device_manager.device_module.empty_cache()
 
 
-def _init_api(ps: ParameterServer) -> Any:
-    import fastapi
-    from fastapi import Request
-    from fastapi.responses import JSONResponse, Response
-
-    app = fastapi.FastAPI()
-
-    class RegisterRequest(BaseModel):
-        files: list[str]
-
-    class UpdateRequest(BaseModel):
-        ranks: list[int] = []
-        update_url: str | None = None
-        inference_group_ranks: list[int] = []
-        timeout: float = 300.0
-        uds: str | None = None
-
-    def wrap_exception(func: Callable[[], None]) -> Response:
-        try:
-            func()
-        except Exception as e:  # noqa: BLE001
-            logger.exception(f"wrap exception {func} failed")
-            return JSONResponse(content=str(e), status_code=500)
-        return Response(status_code=200)
-
-    @app.post("/v1/checkpoints/{checkpoint_name}/files")
-    async def register_files(checkpoint_name: str, req: RegisterRequest, raw: Request) -> Response:
-        return wrap_exception(lambda: ps.register_checkpoint(checkpoint_name, files=req.files))
-
-    @app.delete("/v1/checkpoints/{checkpoint_name}")
-    async def unregister_checkpoint(checkpoint_name: str) -> Response:
-        return wrap_exception(lambda: ps.unregister_checkpoint(checkpoint_name))
-
-    @app.get("/v1/healthz")
-    async def healthz() -> Response:
-        return Response(status_code=200)
-
-    @app.post("/v1/checkpoints/{checkpoint_name}/gather-metas")
-    async def gather_metas(checkpoint_name: str) -> Response:
-        return wrap_exception(lambda: ps.gather_metas(checkpoint_name))
-
-    @app.post("/v1/checkpoints/{checkpoint_name}/update")
-    async def update(checkpoint_name: str, req: UpdateRequest) -> Response:
-        def update_func(socket_paths: list[tuple[str, str]]):
-            if req.update_url is None:
-                return
-            if req.inference_group_ranks:
-                socket_paths = [socket_paths[i] for i in req.inference_group_ranks]
-            request_inference_to_update(
-                req.update_url, dict(socket_paths), timeout=req.timeout, uds=req.uds
-            )
-
-        return wrap_exception(lambda: ps.update(checkpoint_name, update_func, ranks=req.ranks))
-
-    return app
-
-
-@logger.catch(reraise=True)
-def run_from_cli():
-    import uvicorn
-
-    parser = argparse.ArgumentParser(description="Parameter Server")
-    parser.add_argument("--uds", type=str)
-
-    args = parser.parse_args()
-    logger.info(
-        f"Parameter Server {args=}, master addr: {os.getenv('MASTER_ADDR')}, master port {os.getenv('MASTER_PORT')}"
-    )
-
-    assert args.uds and len(args.uds) > 0, args.uds
-    ps = ParameterServer(auto_pg=True)
-    uvicorn.run(_init_api(ps), uds=args.uds, timeout_keep_alive=60)
-
-
+# we need this CLI entry point for compatibility with former versions
 if __name__ == "__main__":
+    from .__main__ import run_from_cli
+
     run_from_cli()