[ckpt, rollout, vllm] feat: support delta weight sync on Ascend NPU (HCCL transport)

src/twinkle/checkpoint_engine/__init__.py

@@ -16,6 +16,7 @@
 """
 
 from .base import CheckpointEngine, TensorMeta
+from .delta import DeltaShardedCheckpointEngine, DeltaSyncConfig
 from .hccl_checkpoint_engine import HCCLCheckpointEngine
 from .ipc_checkpoint_engine import IPCCheckpointEngine
 from .manager import CheckpointEngineManager, CheckpointEngineMode
@@ -33,6 +34,8 @@
     'HCCLCheckpointEngine',
     'IPCCheckpointEngine',
     'TensorMeta',
+    'DeltaSyncConfig',
+    'DeltaShardedCheckpointEngine',
     'generate_auth_key',
     'send_auth_pyobj',
     'recv_auth_pyobj',

src/twinkle/checkpoint_engine/delta/__init__.py

@@ -0,0 +1,30 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Sparse delta weight synchronization.
+
+Layout:
+    config.py      DeltaSyncConfig and the TypedDicts exchanged between trainer and rollout
+    platform.py    resolve_delta_engine_cls: GPU -> NCCL engine, NPU -> HCCL engine (lazy import)
+    protocol/      wire format shared by both sides
+        wire.py        DeltaFlush / DeltaParam structures and checksum
+        decode.py      device-agnostic decode of one received flush (no vLLM dependency)
+    trainer/       trainer side: what changed and how it is exported
+        mixin.py       DeltaTrainerMixin: trainer-side Ray lifecycle
+        spec.py        ShardSpec / BlockPlacement: where a local shard sits in the full tensor
+        export.py      snapshot, diff and HF-coordinate delta export helpers
+        backends/      per-training-backend export hooks (base.py contract, fsdp.py implementation)
+    transport/     sender pipeline and its transports
+        pipeline.py    DeltaSyncMixin: transport-agnostic bucket/flush pipeline and seed/steady state machine
+        gather.py      sparse gather of per-rank deltas to rank 0
+        nccl_engine.py DeltaShardedCheckpointEngine: NCCL transport (hooks + rollout receive)
+        hccl_engine.py DeltaHCCLCheckpointEngine: HCCL (Ascend NPU) sender transport
+    rollout/       rollout side
+        mixin.py       DeltaSamplerMixin: rollout-side lifecycle RPCs driven by vLLMSampler
+        cuda_receiver.py  CudaDeltaReceiver: Ray actor per vLLM worker, NCCL receive + CUDA IPC relay (GPU)
+        hccl_receiver.py  HCCL receive pump that runs inside each vLLM worker process (NPU)
+    ../zmq_pubsub.py   zmq PUB/SUB manifest-channel helpers shared by the NCCL and HCCL engines
+"""
+from .config import DeltaSyncConfig
+from .transport.nccl_engine import DeltaShardedCheckpointEngine
+from .platform import resolve_delta_engine_cls
+
+__all__ = ['DeltaSyncConfig', 'DeltaShardedCheckpointEngine', 'resolve_delta_engine_cls']

src/twinkle/checkpoint_engine/delta/config.py

@@ -0,0 +1,71 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Configuration for the delta checkpoint engine."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypedDict, Union
+from typing_extensions import NotRequired
+
+if TYPE_CHECKING:
+    from ..hccl_checkpoint_engine import MasterMetadata as HCCLMasterMetadata
+    from ..nccl_checkpoint_engine import MasterMetadata
+
+    DeltaMasterMetadata = Union[MasterMetadata, HCCLMasterMetadata]
+else:
+    DeltaMasterMetadata = object
+
+SPARSE_ELEMENT_BYTES = 6
+
+
+@dataclass(frozen=True)
+class DeltaSyncConfig:
+    bucket_bytes: int = 64 << 20
+    timeout_seconds: float = 300
+    gather_batch_size: int = 32
+    delta_pin_snapshots: bool = True
+    free_cache_engine: bool = True
+
+    def __post_init__(self):
+        if not isinstance(self.gather_batch_size, int) or self.gather_batch_size < 1:
+            raise ValueError('delta gather batch size must be a positive integer')
+        if not isinstance(self.bucket_bytes, int) or self.bucket_bytes < SPARSE_ELEMENT_BYTES:
+            raise ValueError('delta bucket must hold at least one int32 index and BF16 value')
+        if not 0 < self.timeout_seconds < float('inf'):
+            raise ValueError('delta timeout must be finite and positive')
+
+
+class DeltaTransportEndpoint(TypedDict):
+    world_size: int
+    master_metadata: DeltaMasterMetadata
+    # CUDA only: NPU workers join the trainer's HCCL group themselves, so there is no worker list.
+    rollout_workers: NotRequired[list[DeltaWorkerMetadata]]
+    # NPU only: trainer rank 0's HCCL socket env; workers apply it before joining the group.
+    hccl_env: NotRequired[dict[str, str]]
+
+
+class DeltaTrainerMetadata(TypedDict):
+    device: str
+    master_metadata: DeltaMasterMetadata | None
+    hccl_env: NotRequired[dict[str, str]]
+
+
+class DeltaWorkerCapabilities(TypedDict):
+    tp: int
+    device: str
+    # NPU only: data-parallel index of the reporting sampler replica.
+    replica_rank: NotRequired[int]
+
+
+class DeltaWorkerMetadata(TypedDict):
+    node_ip: str
+    device_uuid: str
+    cuda_visible_device: NotRequired[str]
+    zmq_handle: str
+
+
+class DeltaReceiverPlacement(TypedDict):
+    rank: int
+    node_id: str
+    device_uuid: str
+    cuda_visible_device: NotRequired[str]
+    zmq_handle: str

src/twinkle/checkpoint_engine/delta/platform.py

@@ -0,0 +1,28 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Platform -> delta transport resolution."""
+from __future__ import annotations
+
+_HCCL_MODULE = 'twinkle.checkpoint_engine.delta.transport.hccl_engine'
+
+# Trainer-side torch device type per supported platform.
+DELTA_DEVICES = {'GPU': 'cuda', 'NPU': 'npu'}
+
+
+def resolve_delta_engine_cls(platform_name: str) -> type:
+    """Return the delta transport class for ``platform_name`` ('GPU' or 'NPU').
+
+    Imports are lazy so the GPU path never touches NPU-specific modules.
+    """
+    if platform_name == 'GPU':
+        from .transport.nccl_engine import DeltaShardedCheckpointEngine
+        return DeltaShardedCheckpointEngine
+    if platform_name == 'NPU':
+        try:
+            from .transport.hccl_engine import DeltaHCCLCheckpointEngine
+        except ModuleNotFoundError as error:
+            if error.name != _HCCL_MODULE:
+                raise
+            raise NotImplementedError(
+                'delta_sharded on NPU requires the HCCL delta transport, which is not implemented yet') from error
+        return DeltaHCCLCheckpointEngine
+    raise ValueError(f'delta_sharded does not support platform {platform_name!r}; supported platforms: GPU, NPU')

src/twinkle/checkpoint_engine/delta/protocol/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Wire format shared by both sides: flush schema, checksum and device-agnostic decode."""

src/twinkle/checkpoint_engine/delta/protocol/decode.py

@@ -0,0 +1,80 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Device-agnostic decoding of one received ``DeltaFlush`` (no vLLM dependency).
+
+Shared by the vLLM weight-transfer adapter (which converts :class:`DeltaPatch` into vLLM's
+``CheckpointWeightPatch``) and by tests.
+"""
+from __future__ import annotations
+
+import json
+import torch
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+from .wire import checksum
+
+
+@dataclass
+class DeltaPatch:
+    """One parameter's patch. ``positions`` are flat int32 positions, ``None`` for dense."""
+
+    name: str
+    shape: tuple
+    dtype: torch.dtype
+    values: torch.Tensor
+    positions: torch.Tensor | None = None
+
+
+def decode_delta_wire(
+    named_tensors: Iterable[tuple[str, torch.Tensor]], ) -> tuple[str, list[DeltaPatch]]:
+    """Validate one DeltaFlush and slice it into per-parameter patches."""
+    tensors = dict(named_tensors)
+    try:
+        spec_tensor = tensors['__delta_spec__']
+        values = tensors['__values__']
+        spec = json.loads(bytes(spec_tensor.detach().cpu().numpy().tobytes()).decode())
+    except KeyError as exc:
+        raise ValueError(f'DeltaFlush is missing {exc.args[0]}') from exc
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise ValueError('DeltaFlush contains an invalid JSON spec') from exc
+
+    encoding = spec['encoding']
+    if encoding not in {'dense', 'indices'}:
+        raise ValueError(f'Unsupported DeltaFlush encoding: {encoding!r}')
+    positions = tensors.get('__positions__')
+    if positions is None:
+        positions = torch.empty(0, dtype=torch.uint8, device=values.device)
+    if encoding == 'dense' and positions.numel():
+        raise ValueError('Dense DeltaFlush must not contain positions')
+    if encoding == 'indices' and '__positions__' not in tensors:
+        raise ValueError('Indices DeltaFlush requires __positions__')
+
+    got_checksum = checksum(positions, values)
+    expected_checksum = int(spec['checksum'])
+    if got_checksum != expected_checksum:
+        raise RuntimeError(
+            f'Delta checksum mismatch in vLLM consumer: got {got_checksum}, expected {expected_checksum}')
+
+    patches: list[DeltaPatch] = []
+    for param_spec in spec['params']:
+        name = param_spec['name']
+        dtype = getattr(torch, param_spec['dtype'])
+        patch_values = values[param_spec['val_start']:param_spec['val_end']]
+        if encoding == 'dense':
+            patch_positions = None
+        else:
+            if param_spec['pos_width'] != 4:
+                raise ValueError(f'{name}: indices encoding requires pos_width=4')
+            pos_bytes = positions[param_spec['pos_start']:param_spec['pos_end']]
+            if pos_bytes.numel() != patch_values.numel() * 4:
+                raise ValueError(f'{name}: position and value slices have inconsistent lengths')
+            patch_positions = pos_bytes.view(torch.int32)
+        patches.append(
+            DeltaPatch(
+                name=name,
+                shape=tuple(param_spec['shape']),
+                dtype=dtype,
+                values=patch_values,
+                positions=patch_positions,
+            ))
+    return encoding, patches

src/twinkle/checkpoint_engine/delta/protocol/wire.py

@@ -0,0 +1,90 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Delta payload schema and checksum based on verl."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+
+DeltaEncodingName = Literal["indices"]
+
+
+# ---------- diff ----------------------------------------------------------
+
+
+@dataclass
+class DeltaParam:
+    """Per-parameter manifest entry for a single chunk / bucket.
+
+    Offsets are byte offsets into the surrounding ``__positions__`` blob and
+    element offsets into the surrounding ``__values__`` tensor.
+    """
+
+    name: str
+    dtype: str
+    shape: list[int]
+    pos_start: int
+    pos_end: int
+    pos_width: int  # 2 or 4
+    val_start: int
+    val_end: int
+
+
+def _xor_tree(t: torch.Tensor) -> int:
+    """Device-side XOR fold of a tensor's raw bytes (log2(n) rounds, one sync).
+
+    Works on any device (including NPU), unlike ``torch.hash_tensor`` which only
+    has CPU/CUDA kernels. The fold is out-of-place, so the argument's storage is
+    never mutated, and the result is a pure function of the bytes, independent
+    of dtype, so sender and receiver agree anywhere.
+    """
+    if t.numel() == 0:
+        return 0
+    flat = t.contiguous().reshape(-1)
+    v = flat.view(torch.uint8) if flat.dtype != torch.uint8 else flat
+    # Halving XOR fold; an odd middle element carries to the next round.
+    while v.numel() > 1:
+        n = v.numel()
+        half = n // 2
+        if n % 2 == 0:
+            v = v[:half] ^ v[half:]
+        else:
+            v = torch.cat((v[:half] ^ v[half + 1 :], v[half : half + 1]))
+    return int(v.item())
+
+
+def checksum(positions: torch.Tensor, values: torch.Tensor) -> int:
+    """Wire-corruption check; sender computes pre-flush, receiver post-recv.
+
+    XOR tree over the bytes of each argument (device-side, one ``.item()`` sync
+    per argument). The same function runs on both ends, so the value is stable
+    across dtypes and devices.
+    """
+    p = _xor_tree(positions)
+    v = _xor_tree(values)
+    return (p ^ (v << 1)) & ((1 << 64) - 1)
+
+
+# ---------- encode --------------------------------------------------------
+
+
+@dataclass
+class DeltaFlush:
+    """One ready-to-dispatch flush.
+
+    * ``positions_cpu`` is a uint8 positions blob. Despite the name it lives on
+      the GPU in the sharded engine (the wire broadcasts from the GPU, so a
+      host round-trip would be pure overhead).
+    * ``values_gpu`` stays on the GPU until the checkpoint engine broadcasts it
+      over NCCL.
+    * ``params`` carries the per-parameter manifest the receiver needs to
+      decode the blob (sent alongside the data over the zmq side-channel).
+    """
+
+    encoding: DeltaEncodingName
+    params: list[DeltaParam]
+    positions_cpu: torch.Tensor
+    values_gpu: torch.Tensor
+    checksum: int

src/twinkle/checkpoint_engine/delta/rollout/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Rollout side: the sampler Ray mixin and the HCCL receiver that runs inside each vLLM worker."""