diff --git a/src/twinkle/checkpoint_engine/__init__.py b/src/twinkle/checkpoint_engine/__init__.py index 191f205d5..673185b3b 100644 --- a/src/twinkle/checkpoint_engine/__init__.py +++ b/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', diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py new file mode 100644 index 000000000..a4b62e1ee --- /dev/null +++ b/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'] diff --git a/src/twinkle/checkpoint_engine/delta/config.py b/src/twinkle/checkpoint_engine/delta/config.py new file mode 100644 index 000000000..3bc980f08 --- /dev/null +++ b/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 diff --git a/src/twinkle/checkpoint_engine/delta/platform.py b/src/twinkle/checkpoint_engine/delta/platform.py new file mode 100644 index 000000000..68022a6f6 --- /dev/null +++ b/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') diff --git a/src/twinkle/checkpoint_engine/delta/protocol/__init__.py b/src/twinkle/checkpoint_engine/delta/protocol/__init__.py new file mode 100644 index 000000000..d14647075 --- /dev/null +++ b/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.""" diff --git a/src/twinkle/checkpoint_engine/delta/protocol/decode.py b/src/twinkle/checkpoint_engine/delta/protocol/decode.py new file mode 100644 index 000000000..40b7f2d73 --- /dev/null +++ b/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 diff --git a/src/twinkle/checkpoint_engine/delta/protocol/wire.py b/src/twinkle/checkpoint_engine/delta/protocol/wire.py new file mode 100644 index 000000000..da52f0e17 --- /dev/null +++ b/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 diff --git a/src/twinkle/checkpoint_engine/delta/rollout/__init__.py b/src/twinkle/checkpoint_engine/delta/rollout/__init__.py new file mode 100644 index 000000000..92a91fa1e --- /dev/null +++ b/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.""" diff --git a/src/twinkle/checkpoint_engine/delta/rollout/cuda_receiver.py b/src/twinkle/checkpoint_engine/delta/rollout/cuda_receiver.py new file mode 100644 index 000000000..17092a180 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/rollout/cuda_receiver.py @@ -0,0 +1,79 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""CUDA rollout receiver: a Ray actor colocated with one vLLM worker. + +It joins the trainer's NCCL group, receives each delta flush (:meth:`DeltaShardedCheckpointEngine.receive_weights`) +and relays it to its vLLM worker over CUDA IPC. The NPU counterpart is :mod:`.hccl_receiver`, which runs inside +the vLLM worker itself because NPU has no device IPC. +""" +from __future__ import annotations + +import asyncio +import os +from collections.abc import Iterator + +import torch + +from ..config import DeltaReceiverPlacement, DeltaSyncConfig, DeltaTransportEndpoint +from ..transport.nccl_engine import DeltaShardedCheckpointEngine + +TensorFlush = tuple[list[tuple[str, torch.Tensor]], bool] + + +class CudaDeltaReceiver: + """Checkpoint worker colocated with one vLLM worker""" + + def __init__(self, config: DeltaSyncConfig, endpoint: DeltaTransportEndpoint, + placement: DeltaReceiverPlacement) -> None: + visible_device = placement.get('cuda_visible_device', placement['device_uuid']) + os.environ['CUDA_VISIBLE_DEVICES'] = visible_device + print(f'[delta receiver rank={placement["rank"]}] ' + f'CUDA_VISIBLE_DEVICES={visible_device!r}, expected_uuid={placement["device_uuid"]}', flush=True) + try: + torch.cuda.set_device(0) + except RuntimeError as error: + raise RuntimeError( + f'CUDA delta receiver initialization failed: CUDA_VISIBLE_DEVICES={visible_device!r}, ' + f'expected_uuid={placement["device_uuid"]}, ' + f'LD_PRELOAD={os.environ.get("LD_PRELOAD", "")!r}') from error + actual_uuid = str(torch.cuda.get_device_properties(0).uuid) + if actual_uuid != placement['device_uuid']: + raise ValueError(f'CUDA delta receiver GPU mismatch: expected {placement["device_uuid"]}, ' + f'got {actual_uuid}, CUDA_VISIBLE_DEVICES={visible_device!r}') + self.config = config + self.zmq_handle = placement['zmq_handle'] + self.engine = DeltaShardedCheckpointEngine(bucket_size=config.bucket_bytes) + self.engine.prepare() + self.engine.init_process_group(placement['rank'], endpoint['world_size'], endpoint['master_metadata']) + self._flushes: Iterator[TensorFlush] | None = None + self._current: TensorFlush | None = None + + def ready(self) -> int: + return self.engine.rank + + def next_flush(self) -> dict[str, bool] | None: + if self._flushes is None: + self._flushes = self.engine.receive_weights() + self._current = next(self._flushes, None) + if self._current is None: + self._flushes = None + return None + return {'is_last': self._current[1]} + + def send_flush(self) -> None: + from twinkle.sampler.vllm_sampler.delta.cuda_ipc import BucketedWeightSender + + tensors, _ = self._current + sender = BucketedWeightSender( + zmq_handle=self.zmq_handle, + bucket_size_mb=max(1, self.config.bucket_bytes >> 20), use_shm=False) + try: + asyncio.run(sender.async_send_weights(iter(tensors))) + finally: + self._current = None + + def close(self) -> None: + self._current = None + if self._flushes is not None: + self._flushes.close() + self._flushes = None + self.engine.close() diff --git a/src/twinkle/checkpoint_engine/delta/rollout/hccl_receiver.py b/src/twinkle/checkpoint_engine/delta/rollout/hccl_receiver.py new file mode 100644 index 000000000..9b8c92ff5 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/rollout/hccl_receiver.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Receiver-side HCCL delta transport (runs inside each vLLM worker process). + +Import-safe without ``torch_npu`` / ``vllm_ascend``: the communicator class is imported lazily by +``stateless_init_process_group`` and device modules are looked up by name at call time. + +The SUB socket MUST be connected before the communicator is built: the trainer builds its +communicator lazily at its first publish, which only completes once every worker joined, and a +zmq PUB drops messages sent before a subscriber connected. +""" +from __future__ import annotations + +import json +import os +import torch +from collections.abc import Callable + +from twinkle.utils import stateless_init_process_group + +_REQUIRED_KEYS = ('is_last', ) +_REQUIRED_FLUSH_KEYS = ('encoding', 'val_numel', 'val_dtype', 'spec') + + +def build_delta_hccl_communicator(dist_ip: str, dist_port: int, group_rank: int, world_size: int, device, + hccl_env: dict[str, str] | None = None): + """Join the trainer's HCCL broadcast group at ``group_rank``. + + ``hccl_env`` is trainer rank 0's HCCL socket env. It overrides this process's values: the trainer may + have derived them from its FSDP ``MASTER_PORT``, and ranks with different port layouts time out in + the communicator's first collective. + """ + if hccl_env: + os.environ.update(hccl_env) + if isinstance(device, str): + device = torch.device(device) + index = device.index + if index is None: + index = getattr(torch, device.type).current_device() + return stateless_init_process_group( + master_address=dist_ip, + master_port=dist_port, + rank=group_rank, + world_size=world_size, + device=index, + backend='hccl', + ) + + +def _validate_manifest(meta) -> None: + if not isinstance(meta, dict): + raise ValueError(f'Malformed delta manifest: expected dict, got {type(meta).__name__}') + if meta.get('terminal_empty'): + return + missing = [k for k in _REQUIRED_KEYS + _REQUIRED_FLUSH_KEYS if k not in meta] + if meta.get('encoding') != 'dense' and 'pos_numel' not in meta: + missing.append('pos_numel') + if missing: + raise ValueError(f'Malformed delta manifest: missing keys {sorted(missing)}') + if not isinstance(getattr(torch, str(meta['val_dtype']), None), torch.dtype): + raise ValueError(f"Malformed delta manifest: unknown val_dtype {meta['val_dtype']!r}") + + +def pump_hccl_delta_session( + sub, + comm, + device, + on_flush: Callable[[list[tuple[str, torch.Tensor]]], None], +) -> str: + """Pull one delta session (manifests + broadcasts) and hand each flush to ``on_flush``. + + Returns the session encoding (``'dense'`` | ``'indices'``); a terminal-only session returns + ``'indices'`` with zero ``on_flush`` calls. Broadcast order mirrors the trainer: positions + (non-dense only) then values. + """ + synchronize = getattr(torch, device.type).synchronize + encoding = 'indices' + while True: + sub.recv_string() + meta = sub.recv_pyobj() + _validate_manifest(meta) + if meta.get('terminal_empty'): + break + encoding = meta['encoding'] + dense = encoding == 'dense' + val_dtype = getattr(torch, meta['val_dtype']) + elem = torch.empty(0, dtype=val_dtype).element_size() + val_u8 = torch.empty(meta['val_numel'] * elem, dtype=torch.uint8, device=device) + pos = None + if not dense: + pos = torch.empty(meta['pos_numel'], dtype=torch.uint8, device=device) + comm.broadcast(pos, src=0) + comm.broadcast(val_u8, src=0) + synchronize() + val = val_u8.view(val_dtype) + spec_t = torch.frombuffer(bytearray(json.dumps(meta['spec']).encode()), dtype=torch.uint8).to(device) + named = [('__delta_spec__', spec_t), ('__values__', val)] + if pos is not None: + named.insert(1, ('__positions__', pos)) + on_flush(named) + del pos, val_u8, val, spec_t, named + if meta['is_last']: + break + return encoding diff --git a/src/twinkle/checkpoint_engine/delta/rollout/mixin.py b/src/twinkle/checkpoint_engine/delta/rollout/mixin.py new file mode 100644 index 000000000..799dca1f7 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/rollout/mixin.py @@ -0,0 +1,260 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rollout-side lifecycle for the delta checkpoint engine. + +CUDA: Ray receiver actors (:mod:`.cuda_receiver`) colocated with each vLLM worker feed it over IPC. Ascend NPU: +no receiver actors; the vLLM workers join the trainer's HCCL group themselves (:mod:`.hccl_receiver`) and the +sampler only drives vLLM's lifecycle RPCs. +""" +from __future__ import annotations + +import asyncio +from typing import Any, TYPE_CHECKING + +from twinkle import Platform, remote_function +from ..config import (DeltaReceiverPlacement, DeltaSyncConfig, DeltaTransportEndpoint, + DeltaWorkerCapabilities, DeltaWorkerMetadata) +from .cuda_receiver import CudaDeltaReceiver + +if TYPE_CHECKING: + from ray.actor import ActorHandle + + +class DeltaSamplerMixin: + """Sampler-side delta lifecycle RPCs (device-aware: CUDA uses receiver actors, NPU does not).""" + _delta_config: DeltaSyncConfig + _delta_job_failed: bool = False + _delta_weight_transfer_initialized: bool = False + _delta_receivers: list[ActorHandle] + _delta_device: str | None = None # 'cuda' / 'npu', remembered by validate_delta_sampler + _delta_tp: int | None = None # NPU: TP size of this replica + _delta_replica_rank: int | None = None # NPU: this actor's data-parallel replica index + _delta_endpoint: DeltaTransportEndpoint | None = None # NPU: HCCL rendezvous stored by init_delta_sampler_transport + + @property + def _delta_is_npu(self) -> bool: + return self._delta_device == 'npu' + + @staticmethod + def _delta_actor_replica_rank() -> int: + """Replica index of this sampler actor. + + ``RayHelper.create_workers`` starts actor ``i`` of ``sampler._actors`` with ``RANK=i`` and + ``WORLD_SIZE=len(_actors)``, and ``dispatch='all'`` results are collected in ``_actors`` order, so + ``Platform.get_rank()`` is the actor's position in the manager's replica list. ``DeviceMesh`` lookups + are not used: they index by GPU id and are wrong when an actor owns several GPUs (TP). + """ + rank = Platform.get_rank() + world_size = Platform.get_world_size() + if rank < 0 or rank >= world_size: + raise RuntimeError( + f'cannot determine the delta replica rank of this sampler actor (RANK={rank}, WORLD_SIZE={world_size}); ' + 'NPU delta sync requires sampler actors created by Ray with a per-actor RANK') + return rank + + async def _delta_rpc(self, method: str, kwargs: dict[str, Any] | None = None) -> list[Any]: + return await self.engine.engine.collective_rpc( + method, kwargs=kwargs or {}, timeout=self._delta_config.timeout_seconds) + + async def _reset_delta_prefix_cache(self) -> None: + reset = await self.engine.engine.reset_prefix_cache(reset_connector=True) + if reset is False: + raise RuntimeError('prefix cache reset failed during delta synchronization') + + async def _clear_delta_cache(self) -> None: + await self._reset_delta_prefix_cache() + await self.engine.engine.reset_mm_cache() + await self.engine.engine.reset_encoder_cache() + + @remote_function(dispatch='all', lazy_collect=False) + def begin_delta_sync(self) -> None: + async def begin(): + self.engine._delta_unavailable = True + await self.engine.pause_delta_generation(self._delta_config.timeout_seconds) + if self._delta_config.free_cache_engine: + # Level 1 preserves the baseline needed by sparse updates. + await self.engine.engine.sleep(level=1) + await self.engine.engine.wake_up(tags=['weights']) + self._run_in_loop(begin()) + + @remote_function(dispatch='all', lazy_collect=False) + def publish_delta_sampler(self) -> None: + async def publish(): + if self._delta_config.free_cache_engine: + await self.engine.engine.wake_up(tags=['kv_cache']) + await self._reset_delta_prefix_cache() + await self.engine.engine.resume_generation() + self.engine._delta_unavailable = False + self._run_in_loop(publish()) + + @remote_function(dispatch='all', lazy_collect=False) + def fail_delta_sampler(self) -> None: + self._delta_job_failed = True + self.engine._delta_unavailable = True + self._run_in_loop(self.engine.pause_delta_generation(self._delta_config.timeout_seconds)) + + @remote_function(dispatch='all', lazy_collect=False) + def validate_delta_sampler(self, config: DeltaSyncConfig) -> DeltaWorkerCapabilities: + self._delta_config = config + if self._delta_job_failed: + raise RuntimeError('delta sampler failed; restart the entire training/rollout job before reseeding') + if config.free_cache_engine and not self.engine.enable_sleep_mode: + raise ValueError('delta KV cache release requires enable_sleep_mode=True or free_cache_engine=False') + results = self._run_in_loop(self._delta_rpc('validate_delta_worker')) + if not results or any(result != results[0] for result in results): + raise ValueError('inconsistent delta worker capabilities') + if len(results) != results[0]['tp']: + raise ValueError('delta topology requires one vLLM worker per TP rank') + self._delta_device = results[0]['device'] + if self._delta_device == 'npu': + self._delta_tp = results[0]['tp'] + self._delta_replica_rank = self._delta_actor_replica_rank() + return dict(results[0], replica_rank=self._delta_replica_rank) + return results[0] + + @remote_function(dispatch='all', collect='flatten', lazy_collect=False) + def delta_sampler_metadata(self) -> list[DeltaWorkerMetadata]: + return self._run_in_loop(self._delta_rpc('delta_receiver_metadata')) + + def _delta_receiver_placements( + self, metadata: list[DeltaWorkerMetadata], endpoint: DeltaTransportEndpoint, + ) -> list[DeltaReceiverPlacement]: + import ray + + workers = endpoint['rollout_workers'] + ranks = {(item['node_ip'], item['device_uuid']): rank for rank, item in enumerate(workers, start=1)} + if not metadata or any(item not in workers for item in metadata): + raise ValueError('local vLLM workers are missing from the delta topology') + nodes = {node['NodeManagerAddress']: node['NodeID'] for node in ray.nodes() if node['Alive']} + return [dict(rank=ranks[item['node_ip'], item['device_uuid']], node_id=nodes[item['node_ip']], + device_uuid=item['device_uuid'], + cuda_visible_device=item.get('cuda_visible_device', item['device_uuid']), + zmq_handle=item['zmq_handle']) + for item in metadata] + + def _spawn_delta_receiver( + self, actor, endpoint: DeltaTransportEndpoint, placement: DeltaReceiverPlacement, + ) -> ActorHandle: + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + return actor.options( + scheduling_strategy=NodeAffinitySchedulingStrategy(placement['node_id'], soft=False), + runtime_env={'env_vars': { + 'CUDA_VISIBLE_DEVICES': placement.get('cuda_visible_device', placement['device_uuid']), + 'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES': '1', + }}).remote(self._delta_config, endpoint, placement) + + @remote_function(dispatch='all', lazy_collect=True) + def init_delta_sampler_transport(self, endpoint: DeltaTransportEndpoint) -> None: + if self._delta_is_npu: + # No receiver actors: workers join the trainer's HCCL group on their first receive. + self._delta_endpoint = endpoint + self._delta_weight_transfer_initialized = False + return + import ray + + metadata = self._run_in_loop(self._delta_rpc('delta_receiver_metadata')) + placements = self._delta_receiver_placements(metadata, endpoint) + actor = ray.remote(num_cpus=0, num_gpus=0)(CudaDeltaReceiver) + self._delta_receivers = [] + for placement in placements: + self._delta_receivers.append(self._spawn_delta_receiver(actor, endpoint, placement)) + ray.get([receiver.ready.remote() for receiver in self._delta_receivers], + timeout=self._delta_config.timeout_seconds) + self._delta_weight_transfer_initialized = False + + async def _next_delta_flush(self) -> dict[str, bool] | None: + flags = await asyncio.gather(*(receiver.next_flush.remote() for receiver in self._delta_receivers)) + if any(flag != flags[0] for flag in flags): + raise ValueError('delta checkpoint receivers disagree about stream completion') + return flags[0] + + async def _send_delta_flush(self) -> None: + receivers = asyncio.create_task(self._delta_rpc('update_twinkle_delta_weights')) + try: + await asyncio.gather(*(receiver.send_flush.remote() for receiver in self._delta_receivers)) + await receivers + finally: + if not receivers.done(): + receivers.cancel() + await asyncio.gather(receivers, return_exceptions=True) + + async def _receive_delta_weights_npu(self) -> dict[str, bool]: + """One NPU round: the vLLM workers pump the whole HCCL session inside ``update_twinkle_delta_weights``. + + That RPC blocks until the session is received and applied, and runs concurrently with the trainer's + send (the manager launches both before waiting). A zero-change round cannot be told apart here (the + worker pump just returns on the terminal manifest), so it is reported as changed; the cache clear is + harmless. + """ + endpoint = self._delta_endpoint + if endpoint is None: + raise RuntimeError('delta NPU transport is not initialized; run a full seed first') + if not self._delta_weight_transfer_initialized: + master = endpoint['master_metadata'] + init_info = { + 'dist_ip': master.dist_ip, 'dist_port': master.dist_port, + 'zmq_ip': master.zmq_ip, 'zmq_port': master.zmq_port, + 'topic': 'bucket_metadata', 'world_size': endpoint['world_size']} + if endpoint.get('hccl_env'): + init_info['hccl_env'] = dict(endpoint['hccl_env']) + await self._delta_rpc('init_weight_transfer_engine', {'init_info': init_info}) + self._delta_weight_transfer_initialized = True + await self._delta_rpc('start_weight_update') + # Group rank 0 is the trainer; workers are numbered replica-major, TP-minor. + base = 1 + self._delta_replica_rank * self._delta_tp + await self._delta_rpc('update_twinkle_delta_weights', { + 'update_info': [{'group_rank': base + index} for index in range(self._delta_tp)]}) + await self._delta_rpc('finish_weight_update') + await self._clear_delta_cache() + return {'changed': True} + + @remote_function(dispatch='all', lazy_collect=True) + def receive_delta_weights(self) -> dict[str, bool]: + if self._delta_is_npu: + return self._run_in_loop(self._receive_delta_weights_npu()) + + async def receive(): + first = await self._next_delta_flush() + if first is None: + return {'changed': False} + if not self._delta_weight_transfer_initialized: + await self._delta_rpc('init_weight_transfer_engine', {'init_info': {}}) + self._delta_weight_transfer_initialized = True + await self._delta_rpc('start_weight_update') + + saw_last = first['is_last'] + await self._send_delta_flush() + following = await self._next_delta_flush() + while following is not None: + if saw_last: + raise ValueError('DeltaFlush stream yielded data after is_last=True') + saw_last = following['is_last'] + await self._send_delta_flush() + following = await self._next_delta_flush() + if not saw_last: + raise ValueError('DeltaFlush stream ended without is_last=True') + await self._delta_rpc('finish_weight_update') + await self._clear_delta_cache() + return {'changed': True} + return self._run_in_loop(receive()) + + @remote_function(dispatch='all', lazy_collect=False) + def close_delta_sampler(self) -> None: + import ray + + self.engine._delta_unavailable = True + self._run_in_loop(self.engine.pause_delta_generation(self._delta_config.timeout_seconds)) + if self._delta_is_npu: + try: + self._run_in_loop(self._delta_rpc('shutdown_delta_weight_transfer')) + finally: + self._delta_weight_transfer_initialized = False + self._delta_endpoint = None + return + receivers = getattr(self, '_delta_receivers', []) + try: + ray.get([receiver.close.remote() for receiver in receivers], timeout=self._delta_config.timeout_seconds) + finally: + for receiver in receivers: + ray.kill(receiver, no_restart=True) + self._delta_receivers = [] diff --git a/src/twinkle/checkpoint_engine/delta/trainer/__init__.py b/src/twinkle/checkpoint_engine/delta/trainer/__init__.py new file mode 100644 index 000000000..2135334ee --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Trainer side: shard geometry, snapshot diff/export, training-backend hooks and the Ray mixin.""" diff --git a/src/twinkle/checkpoint_engine/delta/trainer/backends/__init__.py b/src/twinkle/checkpoint_engine/delta/trainer/backends/__init__.py new file mode 100644 index 000000000..ee3a14432 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer/backends/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import DeltaTrainerBackend +from .fsdp import FSDP2DeltaBackend + +__all__ = ['DeltaTrainerBackend', 'FSDP2DeltaBackend'] diff --git a/src/twinkle/checkpoint_engine/delta/trainer/backends/base.py b/src/twinkle/checkpoint_engine/delta/trainer/backends/base.py new file mode 100644 index 000000000..3cffb18fc --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer/backends/base.py @@ -0,0 +1,68 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Trainer-backend contract for sharded delta synchronization. + +The delta engine only needs a training object that can export (1) the full HF-keyed +weights for the seed, (2) its local parameter shards with placement metadata, and +(3) final HF-coordinate sparse deltas. A backend implements the first two plus a +config check and, when its local coordinates are not already HF coordinates, +overrides :meth:`DeltaTrainerBackend._hf_delta_entry`. Snapshots and diffing are +shared and need no backend code. +""" +from __future__ import annotations + +from collections.abc import Iterator + +import torch + +from ..export import _hf_entry_identity, hf_delta_export, prime_delta_snapshots +from ..spec import ShardSpec + +TensorWeights = Iterator[tuple[str, torch.Tensor]] +TensorShards = Iterator[tuple[str, torch.Tensor, ShardSpec]] +HFDeltaEntry = tuple[list[tuple[str, tuple[int, ...]]], str, torch.Tensor, + torch.Tensor, torch.Tensor, object] + + +class DeltaTrainerBackend: + """Export hooks a training backend provides to the delta engine. + + Required: :meth:`validate_delta_backend`, :meth:`get_per_tensor_param`, + :meth:`get_per_tensor_param_shard`. Optional: :meth:`_hf_delta_entry`. + Provided: :meth:`prime_delta_snapshots`, :meth:`get_per_tensor_param_delta_shard`. + """ + + _delta_device: str = 'cuda' + # Pinned snapshots speed up the diff read-back but compete with other pinned pools + # on the node; a backend with a tight host-memory profile can default this to False. + _delta_pin_snapshots: bool = True + + def validate_delta_backend(self) -> None: + """Raise ``ValueError`` when this rank's configuration cannot be delta-synced.""" + raise NotImplementedError + + def get_per_tensor_param(self) -> tuple[TensorWeights, None]: + """Full HF-keyed export for the seed. Collective: every rank must iterate it.""" + raise NotImplementedError + + def get_per_tensor_param_shard(self) -> tuple[TensorShards, None]: + """Yield ``(name, local flat shard, ShardSpec)``, in the same order on every rank.""" + raise NotImplementedError + + def _hf_delta_entry(self, name, spec, place, lidx, lval): + """Translate one shard-local delta into a final HF-coordinate entry. + + The default is the identity profile: the parameter is its own single HF slot.""" + return _hf_entry_identity(name, spec, place, lidx, lval) + + def prime_delta_snapshots(self) -> None: + """Snapshot the current shards as the diff base. Runs right after the seed sync.""" + self._delta_shard_snap = getattr(self, '_delta_shard_snap', {}) + shards, _ = self.get_per_tensor_param_shard() + pin = self._delta_device == 'cuda' and self._delta_pin_snapshots + prime_delta_snapshots(shards, self._delta_shard_snap, pin=pin) + + def get_per_tensor_param_delta_shard(self) -> tuple[Iterator[HFDeltaEntry], None]: + """Steady export: HF-coordinate deltas since the previous export. Needs a prior seed.""" + self._delta_shard_snap = getattr(self, '_delta_shard_snap', {}) + shards, _ = self.get_per_tensor_param_shard() + return hf_delta_export(shards, self._delta_shard_snap, self._hf_delta_entry), None diff --git a/src/twinkle/checkpoint_engine/delta/trainer/backends/fsdp.py b/src/twinkle/checkpoint_engine/delta/trainer/backends/fsdp.py new file mode 100644 index 000000000..35696e900 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer/backends/fsdp.py @@ -0,0 +1,100 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""FSDP2 implementation of the delta trainer backend.""" +from __future__ import annotations + +import re + +import torch + +from twinkle import Platform +from ..spec import ShardSpec, derive_dtensor_placement +from .base import DeltaTrainerBackend, TensorShards, TensorWeights + + +def convert_weight_keys(state_dict: dict[str, torch.Tensor], model): + # convert state dict keys: https://github.com/huggingface/transformers/pull/38385 + if not hasattr(model, "_checkpoint_conversion_mapping"): + return state_dict + + reverse_key_mapping = {v: k for k, v in model._checkpoint_conversion_mapping.items()} + original_weights = {} + for key, value in state_dict.items(): + for pattern, replacement in reverse_key_mapping.items(): + replacement = replacement.lstrip("^") # strip off un-needed chars and patterns + replacement = re.sub(r"\(.*\)", "", replacement) + key, n_replace = re.subn(pattern, replacement, key) + # Early exit of the loop + if n_replace > 0: + break + + original_weights[key] = value + + return original_weights + + +def validate_fsdp2(owner) -> None: + from torch.distributed.tensor import DTensor + from twinkle.model.transformers.strategy.accelerate import AccelerateStrategy + from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + + platform_name = Platform.get_platform().__name__ + if platform_name not in ('GPU', 'NPU'): + raise ValueError(f'delta_sharded requires GPU or NPU trainer actors, got platform {platform_name!r}') + if not isinstance(owner.strategy, (NativeFSDPStrategy, AccelerateStrategy)): + raise ValueError('delta_sharded requires native or Accelerate FSDP2') + model = owner.strategy.unwrap_model(owner.model) + if getattr(model, 'peft_config', None) or getattr(model.config, 'quantization_config', None): + raise ValueError('delta_sharded does not support adapters or quantization') + config = model.config.get_text_config() if hasattr(model.config, 'get_text_config') else model.config + if any(getattr(config, name, 0) for name in ('num_experts', 'moe_num_experts', + 'n_routed_experts', 'num_local_experts')): + raise ValueError('Twinkle FSDP delta export does not yet implement backend-specific MoE HF conversion') + if any((getattr(owner.device_mesh, dim, 1) or 1) != 1 + for dim in ('tp_world_size', 'pp_world_size', 'ep_world_size')): + raise ValueError('FSDP delta export requires training TP=PP=EP=1') + state = model.state_dict() + # FSDP leaves the model unwrapped on a single rank; its plain tensors export as unsharded parameters. + single_rank = getattr(owner.device_mesh, 'world_size', None) == 1 + if not single_rank and not any(isinstance(param, DTensor) for param in state.values()): + raise ValueError('delta_sharded requires a wrapped FSDP2 DTensor model') + if any(param.numel() >= 1 << 31 for param in state.values()): + raise ValueError('delta_sharded indices require fewer than 2^31 elements per tensor') + if any(not param.is_floating_point() for param in state.values()): + raise ValueError('FSDP delta checkpoint tensors must be floating point') + for param in state.values(): + derive_dtensor_placement(ShardSpec.from_param(param)) + + +class FSDP2DeltaBackend(DeltaTrainerBackend): + """Delta backend for native/Accelerate FSDP2 models (identity HF coordinates).""" + + def validate_delta_backend(self) -> None: + validate_fsdp2(self) + + def _delta_state_dict(self) -> dict[str, torch.Tensor]: + model = self.strategy.unwrap_model(self.model) + return convert_weight_keys(model.state_dict(), model) + + def get_per_tensor_param(self) -> tuple[TensorWeights, None]: + def generate(): + from torch.distributed.tensor import DTensor + + for name, param in self._delta_state_dict().items(): + param = param.to(self._delta_device, non_blocking=True) + if isinstance(param, DTensor): + param = param.full_tensor() + yield name, param + return generate(), None + + def get_per_tensor_param_shard(self) -> tuple[TensorShards, None]: + def generate(): + from torch.distributed.tensor import DTensor + + for name, param in self._delta_state_dict().items(): + spec = ShardSpec.from_param(param) + param = param.to(self._delta_device, non_blocking=True) + if param.is_floating_point(): + param = param.to(torch.bfloat16, non_blocking=True) + local = param.to_local() if isinstance(param, DTensor) else param + yield name, local.reshape(-1), spec + return generate(), None diff --git a/src/twinkle/checkpoint_engine/delta/trainer/export.py b/src/twinkle/checkpoint_engine/delta/trainer/export.py new file mode 100644 index 000000000..2fe96fae7 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer/export.py @@ -0,0 +1,91 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shard snapshots and HF delta export based on verl.""" +import torch + +_DTYPE_INT = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + +def shard_delta_indices( + local_new: torch.Tensor, + local_snap: torch.Tensor, + offset: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Byte-diff a local shard against its snapshot; return (global_positions, values). + + Positions are int64 indices into the *full flattened parameter* (offset + local index). + Dtype-agnostic, bytewise (view-as-int), no arithmetic -- matches ``bytewise_diff_mask``. + """ + es = local_new.element_size() + int_dtype = _DTYPE_INT.get(es) + if int_dtype is None: + raise ValueError(f"unsupported element size {es}") + mask = local_new.view(int_dtype) != local_snap.view(int_dtype) + local_idx = mask.nonzero(as_tuple=False).view(-1) + values = local_new[local_idx] + global_idx = local_idx.to(torch.int64) + offset + return global_idx, values + + +def _hf_entry_identity(name, spec, place, lidx, lval): + """Identity profile: the param IS its own single slot (weight name == HF name) + -- translate the shard-local delta to within-param coordinates. int32 + positions: the wire is int32 anyway and the engine asserts the range.""" + from .spec import translate_flat_indices + + gidx = (translate_flat_indices(lidx, place) if lidx.numel() else lidx).to(torch.int32) + counts = torch.zeros(1, dtype=torch.int64) + counts[0] = int(gidx.numel()) + return [(name, tuple(spec.full_shape))], str(lval.dtype).replace("torch.", ""), counts, gidx, lval + + +def hf_delta_export(gen, snaps: dict, entry_fn): + """STEADY export: wrap a raw ``(name, local_shard, spec)`` exporter into final + HF-coordinate delta entries ``(slots, dtype_str, counts, hf_idx, hf_val, + gather_group)`` -- diff against the pinned snapshot, refresh it, then hand the + shard-local delta to ``entry_fn(name, spec, place, lidx, lval)``, the engine's + per-param entry builder (FSDP: identity only; veomni adds the EP converter). + The delta engine consumes these entries verbatim (batch -> gather -> wire); no + spec, no placement and no conversion cross the boundary. Requires a prior seed + pass.""" + from .spec import derive_dtensor_placement + + for name, local, spec in gen: + local = local.detach().contiguous().view(-1) + snap = snaps.get(name) + assert snap is not None and snap.numel() == local.numel(), ( + f"{name}: no seed snapshot for this shard; run the seed export first" + ) + if spec.place is not None: + # explicit exporter override: the backend declared the whole triple + # (hybrid geometries are not derivable from DTensor facts alone). + place, contributes, pg = spec.place, spec.contributes, spec.gather_group + else: + place, contributes, pg = derive_dtensor_placement(spec) + if contributes: + base = snap.to(local.device, non_blocking=True) + lidx, lval = shard_delta_indices(local, base, 0) + else: + # replicated param owned by another rank; empty delta keeps lockstep. + lidx = torch.empty(0, dtype=torch.int64, device=local.device) + lval = torch.empty(0, dtype=local.dtype, device=local.device) + snap.copy_(local, non_blocking=True) + yield (*entry_fn(name, spec, place, lidx, lval), pg) + + +def prime_delta_snapshots(gen, snaps: dict, pin: bool) -> None: + """Snapshot each rank's current shards to CPU as the steady diff base. Run + right after the seed's full-weight sync: weights do not move during the + sync, so the snapshots equal exactly what the rollout side received. + + ``pin`` selects pinned vs pageable host memory and is the ENGINE's call + (``BaseEngine.delta_pin_snapshots``): pinning a whole shard set + (cudaHostAlloc) competes with everything else that pins on the node and its + failure surfaces as a CUDA out-of-memory; pageable costs a slower H2D on + the diff read-back but cannot OOM the device.""" + for name, local, _spec in gen: + local = local.detach().contiguous().view(-1) + snap = snaps.get(name) + if snap is None or snap.numel() != local.numel(): + snap = torch.empty_like(local, device="cpu", pin_memory=pin) + snaps[name] = snap + snap.copy_(local, non_blocking=True) diff --git a/src/twinkle/checkpoint_engine/delta/trainer/mixin.py b/src/twinkle/checkpoint_engine/delta/trainer/mixin.py new file mode 100644 index 000000000..34c106eb6 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer/mixin.py @@ -0,0 +1,95 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Trainer-side Ray lifecycle for the delta checkpoint engine (NCCL on GPU, HCCL on NPU).""" +from __future__ import annotations + +import asyncio +from datetime import timedelta +from time import perf_counter +from typing import Any + +import torch.distributed as dist + +from twinkle import Platform, remote_function +from ..config import DeltaMasterMetadata, DeltaSyncConfig, DeltaTrainerMetadata, DeltaTransportEndpoint +from ..platform import DELTA_DEVICES, resolve_delta_engine_cls + + +class DeltaTrainerMixin: + """Lifecycle RPCs only. Export hooks come from a DeltaTrainerBackend subclass on the same model.""" + _delta_engine: Any = None # transport chosen by platform (see resolve_delta_engine_cls) + _delta_metadata_group: dist.ProcessGroup | None = None + _delta_master_metadata: DeltaMasterMetadata | None = None + + @staticmethod + def _delta_rank() -> int: + # A single trainer rank never initializes the default process group. + return dist.get_rank() if dist.is_initialized() else 0 + + def _validate_delta_ranks(self) -> None: + error = None + try: + self.validate_delta_backend() + except Exception as validation_error: + error = str(validation_error) + if self._delta_metadata_group is None: + if error: + raise ValueError(f'Unsupported delta configuration: {error}') + return + errors = [None] * dist.get_world_size() + dist.all_gather_object(errors, error, group=self._delta_metadata_group) + if any(errors): + raise ValueError(f'Unsupported delta configuration: {errors}') + + def _initialize_delta_engine(self, config: DeltaSyncConfig) -> None: + # Resolve first so an unsupported platform fails before any process group is created. + platform_name = Platform.get_platform().__name__ + engine_cls = resolve_delta_engine_cls(platform_name) + if dist.is_initialized(): + self._delta_metadata_group = dist.new_group( + backend='gloo', timeout=timedelta(seconds=config.timeout_seconds)) + try: + self._validate_delta_ranks() + except Exception: + if self._delta_metadata_group is not None: + dist.destroy_process_group(self._delta_metadata_group) + self._delta_metadata_group = None + raise + self._delta_device = DELTA_DEVICES[platform_name] + self._delta_pin_snapshots = config.delta_pin_snapshots + self._delta_engine = engine_cls( + bucket_size=config.bucket_bytes, batch_gather=config.gather_batch_size) + self._delta_engine.is_master = self._delta_rank() == 0 + self._delta_master_metadata = self._delta_engine.prepare() + + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def prepare_delta_sync(self, config: DeltaSyncConfig) -> DeltaTrainerMetadata: + self._lazy_wrap_model() + if self._delta_engine is None: + self._initialize_delta_engine(config) + metadata: DeltaTrainerMetadata = {'device': self._delta_device, 'master_metadata': self._delta_master_metadata} + hccl_env = getattr(self._delta_engine, 'hccl_env', None) + if hccl_env: + metadata['hccl_env'] = dict(hccl_env) + return metadata + + @remote_function(dispatch='all', lazy_collect=True) + def init_delta_transport(self, endpoint: DeltaTransportEndpoint) -> None: + rank = 0 if self._delta_rank() == 0 else -1 + self._delta_engine.init_process_group(rank, endpoint['world_size'], endpoint['master_metadata']) + + @remote_function(dispatch='all', lazy_collect=True) + def send_delta_weights(self) -> dict[str, Any]: + started = perf_counter() + metrics = asyncio.run(self._delta_engine.send_weights(self)) + return {'send_seconds': perf_counter() - started, 'checkpoint_engine': metrics} + + @remote_function(dispatch='all', lazy_collect=False) + def close_delta_trainer(self) -> None: + if self._delta_engine is not None: + self._delta_engine.close() + self._delta_engine = None + snapshots = getattr(self, '_delta_shard_snap', {}) + snapshots.clear() + if self._delta_metadata_group is not None: + dist.destroy_process_group(self._delta_metadata_group) + self._delta_metadata_group = None diff --git a/src/twinkle/checkpoint_engine/delta/trainer/spec.py b/src/twinkle/checkpoint_engine/delta/trainer/spec.py new file mode 100644 index 000000000..ba365e230 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer/spec.py @@ -0,0 +1,229 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shard placement metadata""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch +from torch.distributed.tensor import DTensor +from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + from torch.distributed.device_mesh import DeviceMesh + +__all__ = ["BlockPlacement", "ShardSpec", "derive_dtensor_placement", "translate_flat_indices"] + + +@dataclass +class ShardSpec: + """Declarative placement descriptor for one exported local parameter shard.""" + + # Logical (full) tensor shape; the distribution facts below refer to it. + full_shape: tuple + # Distribution: torch DeviceMesh + per-mesh-dim Placement. None = unsharded. + mesh: Optional[DeviceMesh] = None + placements: Optional[tuple] = None + # Explicit placement override for trainers whose sharding is not fully captured + # by DTensor placements (e.g. veomni's manual expert-dim split): ``place`` is an + # int flat offset or a BlockPlacement in a *virtual* full tensor, and + # ``gather_group`` is the ProcessGroup covering every rank that holds a block + # (pass a real group object -- the engine treats ``None`` as "unsharded"). + # ``contributes=False`` marks a rank whose block is a replica owned by a + # peer (e.g. HSDP replicate dims): it keeps lockstep with an empty delta. + place: Optional[int | BlockPlacement] = None + gather_group: Optional[ProcessGroup] = None + contributes: bool = True + @classmethod + def from_param(cls, param: torch.Tensor) -> ShardSpec: + if isinstance(param, DTensor): + return cls(full_shape=tuple(param.shape), mesh=param.device_mesh, placements=tuple(param.placements)) + return cls(full_shape=tuple(param.shape)) + + +def _prod(xs: tuple | list) -> int: + n = 1 + for x in xs: + n *= int(x) + return n + + +def _row_major_strides(shape: tuple | list) -> tuple: + strides, acc = [], 1 + for d in reversed([int(x) for x in shape]): + strides.append(acc) + acc *= d + return tuple(reversed(strides)) + + +@dataclass(frozen=True) +class BlockPlacement: + """This rank's local shard is one hyper-rectangular block of the full tensor: + ``full[o0:o0+l0, o1:o1+l1, ...]`` with ``local_shape=(l0, l1, ...)`` and + ``global_offset=(o0, o1, ...)``. Produced by :func:`derive_dtensor_placement` for every + sharded geometry, including the dim-0 cut (FSDP2 ``Shard(0)``): that block is + flat-contiguous, and :func:`translate_flat_indices` detects it via + ``is_flat_contiguous`` and keeps the single-add fast path.""" + + local_shape: tuple + global_offset: tuple + full_shape: tuple + + @property + def local_strides(self) -> tuple: + return _row_major_strides(self.local_shape) + + @property + def full_strides(self) -> tuple: + return _row_major_strides(self.full_shape) + + @property + def is_flat_contiguous(self) -> bool: + """True when the block is one contiguous flat range (only dim 0 is cut). + + Then every trailing dim is whole, so the trailing offsets are all zero + and the block occupies flat positions ``[flat_offset, flat_offset+numel)``. + """ + return all(int(lo) == int(fu) for lo, fu in zip(self.local_shape[1:], self.full_shape[1:], strict=False)) + + @property + def flat_offset(self) -> int: + """Flat start of the block; only meaningful when ``is_flat_contiguous``.""" + return int(self.global_offset[0]) * _prod(self.full_shape[1:]) if self.full_shape else 0 + + +def _shard_dim(p) -> Optional[int]: + """Which tensor dim this placement cuts, or None. ``is_shard()`` misses ``_StridedShard`` on torch 2.13.""" + if p.is_shard(): + return int(p.dim) + return int(p.dim) if type(p).__name__ == "_StridedShard" else None + + +def _assert_even_strided(spec: ShardSpec, placements: tuple) -> None: + """Strided offsets assume equal shards, and switch every tensor dim to that rule -- so check them all.""" + if not any(type(p).__name__ == "_StridedShard" for p in placements): + return + cuts: dict[int, int] = {} + for mesh_dim, p in enumerate(placements): + tdim = _shard_dim(p) + if tdim is not None: + cuts[tdim] = cuts.get(tdim, 1) * spec.mesh.size(mesh_dim) + for tdim, total in cuts.items(): + if spec.full_shape[tdim] % total: + raise NotImplementedError( + f"sharded delta does not support uneven strided sharding: tensor dim {tdim} " + f"has size {spec.full_shape[tdim]}, which {total} shards do not divide evenly " + f"(placements={placements})" + ) + + +def translate_flat_indices(lidx: torch.Tensor, place: int | BlockPlacement) -> torch.Tensor: + """Map shard-local flat positions to full-tensor flat positions. + + ``place`` is what the caller's dispatch produced (``spec.place`` or + :func:`derive_dtensor_placement`): an ``int`` for the + identity cases (unsharded / replicated / explicit exporter overrides, translate + = add), or a :class:`BlockPlacement`. A flat-contiguous block (only dim 0 cut, + e.g. FSDP2 ``Shard(0)``) keeps the single-add fast path; any other block does a + mixed-radix decompose by the local shape, adds the per-dim offset, and recomposes + with the full-tensor strides -- a few divmods on the nnz tensor, no collectives. + """ + if isinstance(place, int): + return lidx + place if place else lidx + if place.is_flat_contiguous: + off = place.flat_offset + return lidx + off if off else lidx + out = torch.zeros_like(lidx) + rem = lidx + for lstride, off, fstride in zip(place.local_strides, place.global_offset, place.full_strides, strict=False): + coord = torch.div(rem, lstride, rounding_mode="floor") + rem = rem - coord * lstride + out = out + (coord + int(off)) * fstride + return out + + +def derive_dtensor_placement(spec: ShardSpec) -> tuple[int | BlockPlacement, bool, Optional[ProcessGroup]]: + """Derive ``(place, contributes, gather_group)`` for THIS rank from a spec + whose distribution is fully declared by ``mesh`` + ``placements`` (or is + unsharded). Specs carrying an explicit exporter override (``spec.place``) + never reach this function -- the caller dispatches on ``spec.place`` and + reads the exporter's own ``(place, contributes, gather_group)`` triple + verbatim, because a hybrid geometry (e.g. veomni's manual ep split) is not + derivable from the DTensor facts alone. + + ``place`` feeds :func:`translate_flat_indices`: + + * unsharded (``mesh is None``) or fully replicated: ``0``; no group (the local + tensor is already the full parameter). + * any sharded geometry: a :class:`BlockPlacement` computed from + ``compute_local_shape_and_global_offset`` (pure math, no collective). For a + single Shard dim, only ranks at coordinate 0 of every Replicate dim + contribute and the gather group is the Shard dim's subgroup; the FSDP2 + default ``Shard(0)`` yields a flat-contiguous block, which keeps the add + fast path in :func:`translate_flat_indices`. With several Shard dims (e.g. + automodel's EP x FSDP ``(Shard(0), Shard(1))`` expert mesh) every rank holds + a distinct block and the gather group spans every Shard dim at once; a + Replicate dim beside them is held fixed rather than spanned. + + ``_StridedShard`` is a block too: it permutes which block a rank gets, and + ``compute_local_shape_and_global_offset`` returns the permuted offset. + """ + import torch.distributed as dist + + assert spec.place is None, "explicit-place specs are dispatched by the caller, not derived" + + if spec.mesh is None: + return 0, (dist.get_rank() == 0 if dist.is_initialized() else True), None + + placements = spec.placements + shard_dims = [d for d, p in enumerate(placements) if _shard_dim(p) is not None] + _assert_even_strided(spec, placements) + + coord = spec.mesh.get_coordinate() + contributes = True + if coord is not None: + for d, p in enumerate(placements): + if p.is_replicate() and coord[d] != 0: + contributes = False + break + + if not shard_dims: + # replicated across every mesh dim: full tensor on each rank, no gather + return 0, contributes, None + + local_shape, global_offset = compute_local_shape_and_global_offset(spec.full_shape, spec.mesh, list(placements)) + place = BlockPlacement(tuple(local_shape), tuple(global_offset), tuple(spec.full_shape)) + + if len(shard_dims) == 1: + return place, contributes, spec.mesh.get_group(mesh_dim=shard_dims[0]) + return place, contributes, _shard_dims_group(spec.mesh, shard_dims) + + +_GATHER_GROUPS: dict[tuple, ProcessGroup] = {} + + +def _shard_dims_group(mesh: DeviceMesh, shard_dims: list[int]) -> ProcessGroup: + """The ranks reached by varying the Shard mesh dims, with this rank's coordinate fixed on the others.""" + import itertools + + import torch.distributed as dist + + coord = mesh.get_coordinate() + assert coord is not None, "cannot derive a gather group for a rank outside the mesh it exports from" + other = [d for d in range(mesh.ndim) if d not in shard_dims] + mine = None + for combo in itertools.product(*[range(mesh.size(d)) for d in other]): + index: list = [slice(None)] * mesh.ndim + for d, c in zip(other, combo, strict=True): + index[d] = c + ranks = tuple(sorted(int(r) for r in mesh.mesh[tuple(index)].flatten().tolist())) + got = _GATHER_GROUPS.get(ranks) + if got is None: + got = dist.new_group(list(ranks)) + _GATHER_GROUPS[ranks] = got + if all(coord[d] == c for d, c in zip(other, combo, strict=True)): + mine = got + assert mine is not None, f"no gather group covers coordinate {coord} of mesh {tuple(mesh.mesh.shape)}" + return mine diff --git a/src/twinkle/checkpoint_engine/delta/transport/__init__.py b/src/twinkle/checkpoint_engine/delta/transport/__init__.py new file mode 100644 index 000000000..0c7b963d4 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/transport/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sender pipeline and its transports: the transport-agnostic pipeline, sparse gather, NCCL and HCCL engines.""" diff --git a/src/twinkle/checkpoint_engine/delta/transport/gather.py b/src/twinkle/checkpoint_engine/delta/transport/gather.py new file mode 100644 index 000000000..dea7583f6 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/transport/gather.py @@ -0,0 +1,104 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sparse gather of per-rank deltas to rank 0, based on verl.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + + +def gather_slot_entries_to_rank0( + idx_concat: torch.Tensor, + val_concat: torch.Tensor, + counts: torch.Tensor, + group: dist.ProcessGroup | None = None, + max_round_bytes: int | None = None, +) -> list | None: + """Variable-length sparse gather, batched: one collective round for K parameters. + + Each rank passes its K per-parameter deltas concatenated (``idx_concat``, + ``val_concat``) plus the per-parameter length vector ``counts`` ([K] int64). + One all_gather exchanges the K x world count matrix; two padded gathers move + the blobs. Rank 0 slices per (rank, param) and returns K ``(idx, val)`` pairs + (None elsewhere) -- bit-identical to K individual gathers, ~K x fewer + collectives and host syncs. + """ + rank = dist.get_rank(group) + world = dist.get_world_size(group) + dst = dist.get_global_rank(group, 0) if group is not None else 0 + dev = idx_concat.device + k = int(counts.numel()) + + counts_all = [torch.zeros_like(counts) for _ in range(world)] + dist.all_gather(counts_all, counts.to(dev), group=group) + counts_cpu = torch.stack(counts_all).cpu().tolist() # one D2H sync instead of `world` + + if max_round_bytes is not None and k > 1: + # Deterministic sub-rounds: every rank sees the same counts matrix, so all + # ranks derive the SAME slot partition (no per-rank trigger asymmetry) and + # each padded round's largest blob stays within the byte budget. + per_elem = idx_concat.element_size() + val_concat.element_size() + budget = max(int(max_round_bytes) // per_elem, 1) # elements per rank per round + cuts = [0] + run = [0] * world + for i in range(k): + run = [run[r] + counts_cpu[r][i] for r in range(world)] + if max(run) > budget and cuts[-1] != i: + cuts.append(i) + run = [counts_cpu[r][i] for r in range(world)] + cuts.append(k) + if len(cuts) > 2: + out_all: list = [] + my_off = [0] + for i in range(k): + my_off.append(my_off[-1] + counts_cpu[rank][i]) + for lo, hi in zip(cuts[:-1], cuts[1:], strict=False): + sub_counts = torch.tensor(counts_cpu[rank][lo:hi], dtype=torch.int64, device=dev) + sub = gather_slot_entries_to_rank0( + idx_concat[my_off[lo] : my_off[hi]], + val_concat[my_off[lo] : my_off[hi]], + sub_counts, + group=group, + ) + if rank == 0: + out_all.extend(sub) + return out_all if rank == 0 else None + + totals = [sum(c) for c in counts_cpu] + max_n = max(totals) if totals else 0 + if max_n == 0: + if rank != 0: + return None + empty_i = torch.empty(0, dtype=idx_concat.dtype, device=dev) + empty_v = torch.empty(0, dtype=val_concat.dtype, device=dev) + return [(empty_i, empty_v) for _ in range(k)] + + idx_pad = torch.zeros(max_n, dtype=idx_concat.dtype, device=dev) + val_pad = torch.zeros(max_n, dtype=val_concat.dtype, device=dev) + n = int(idx_concat.numel()) + idx_pad[:n] = idx_concat + val_pad[:n] = val_concat + + idx_list = [torch.zeros(max_n, dtype=idx_pad.dtype, device=dev) for _ in range(world)] if rank == 0 else None + val_list = [torch.zeros(max_n, dtype=val_concat.dtype, device=dev) for _ in range(world)] if rank == 0 else None + dist.gather(idx_pad, idx_list, dst=dst, group=group) + dist.gather(val_pad, val_list, dst=dst, group=group) + if rank != 0: + return None + + # per-rank cumulative offsets into each blob, sliced per param then stitched across ranks + offs = [[0] * (k + 1) for _ in range(world)] + for r in range(world): + for i in range(k): + offs[r][i + 1] = offs[r][i] + counts_cpu[r][i] + out = [] + for i in range(k): + idx_pieces = [idx_list[r][offs[r][i] : offs[r][i + 1]] for r in range(world) if counts_cpu[r][i]] + val_pieces = [val_list[r][offs[r][i] : offs[r][i + 1]] for r in range(world) if counts_cpu[r][i]] + if idx_pieces: + out.append((torch.cat(idx_pieces), torch.cat(val_pieces))) + else: + out.append( + (torch.empty(0, dtype=idx_concat.dtype, device=dev), torch.empty(0, dtype=val_concat.dtype, device=dev)) + ) + return out diff --git a/src/twinkle/checkpoint_engine/delta/transport/hccl_engine.py b/src/twinkle/checkpoint_engine/delta/transport/hccl_engine.py new file mode 100644 index 000000000..5badd52bb --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/transport/hccl_engine.py @@ -0,0 +1,179 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""HCCL (Ascend NPU) transport for the delta pipeline (the pipeline itself lives in :mod:`.pipeline`). + +Group membership: trainer rank 0 is group rank 0 and the vLLM worker processes +themselves are group ranks 1..N. There are no rollout-side CheckpointEngine +workers on NPU, so this engine only SENDS; the receiver is a vLLM-side component. + +Lazy communicator +----------------- +``init_process_group`` only records rank, world size and the rendezvous metadata. +Trainer rank 0 builds the PyHccl communicator inside ``_ensure_wire()``, i.e. at +the first publish. The group completes only after every receiver has joined, and +every receiver connects its zmq SUB *before* joining the group, so "communicator +ready" implies "all receivers subscribed". A zmq PUB drops messages sent before a +subscriber has connected, so publishing the first manifest right after +``_ensure_wire()`` is lossless. The communicator is persistent across sync rounds +(HCCL group setup is expensive) and destroyed only in ``close()``. + +Wire order per flush: manifest (zmq) -> positions broadcast (absent for values-only +flushes) -> values broadcast (HCCL, src=0). + +All ``torch.npu`` / ``vllm_ascend`` usage is inside methods, so this module imports +on machines without torch_npu. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import torch + +from twinkle.utils import find_free_port, find_node_ip, stateless_init_process_group +from twinkle.utils.platforms import hccl_socket_env +from ...base import CheckpointEngine +from ...hccl_checkpoint_engine import MasterMetadata +from ...zmq_pubsub import start_publisher +from .pipeline import DeltaSyncMixin + +logger = logging.getLogger(__name__) + + +class DeltaHCCLCheckpointEngine(DeltaSyncMixin, CheckpointEngine): + """Sparse delta weight sync over HCCL; trainer (sender) side only.""" + + wire_format = 'delta_flush' + + def __init__( + self, + *args, + bucket_size: int = 3072 << 20, + group_name: str = 'twinkle_ckpt', + rebuild_group: bool = False, + rollout_dtype: torch.dtype = torch.bfloat16, + **kwargs, + ) -> None: + if rebuild_group: + raise NotImplementedError('delta_sharded on NPU keeps one PyHccl broadcast group across sync rounds; ' + 'rebuild_group=True is not supported') + super().__init__(*args, **kwargs) + self.bucket_size = bucket_size + self.group_name = group_name + self.rebuild_group = rebuild_group + self.rollout_dtype = rollout_dtype + + self.is_master = False + self.topic = 'bucket_metadata' + self.rank: int | None = None + self.world_size: int | None = None + self.socket = None + + self._comm = None + self._rendezvous: MasterMetadata | None = None + self._group_world_size: int | None = None + self.ip: str | None = None + self.zmq_port: int | None = None + self.dist_port: int | None = None + # HCCL socket env of trainer rank 0, published so vLLM workers join the group with the same ports. + self.hccl_env: dict[str, str] | None = None + + # ---- lifecycle ---- + def prepare(self) -> MasterMetadata | None: + """Master: bind the manifest PUB socket and allocate ports (idempotent). Others: None.""" + if not self.is_master: + return None + if self.socket is None: + ip = find_node_ip() + if not ip: + import ray + ip = ray.util.get_node_ip_address().strip('[]') + self.ip = ip + self.zmq_port = find_free_port() + self.dist_port = find_free_port() + self.socket = start_publisher(self.ip, self.zmq_port) + # The FSDP process group may already have pinned these from MASTER_PORT; workers must match. + self.hccl_env = hccl_socket_env(self.dist_port) + return MasterMetadata(zmq_ip=self.ip, zmq_port=self.zmq_port, dist_ip=self.ip, dist_port=self.dist_port) + + def init_process_group(self, rank: int, world_size: int, master_metadata: MasterMetadata): + """Record only; the communicator is built lazily by rank 0 (see module docstring).""" + self.rank = rank + self.world_size = world_size + if rank == 0: + self._group_world_size = world_size + self._rendezvous = master_metadata + + def _ensure_communicator(self) -> None: + if self._comm is not None: + return + assert self.is_master and self.rank == 0 and self._rendezvous is not None, ( + 'HCCL communicator is built by trainer rank 0 after init_process_group') + self._comm = stateless_init_process_group( + master_address=self._rendezvous.dist_ip, + master_port=self._rendezvous.dist_port, + rank=0, + world_size=self._group_world_size, + device=torch.npu.current_device(), + backend='hccl', + ) + logger.info('delta HCCL communicator ready: world_size=%s', self._group_world_size) + + def close(self): + """Destroy the communicator, close the PUB socket and reset state. Idempotent.""" + comm, self._comm = self._comm, None + if comm is not None: + try: + comm.destroyComm(comm.comm) + except Exception as e: + logger.warning('Error destroying HCCL communicator: %s', e) + socket, self.socket = self.socket, None + if socket is not None: + ctx = getattr(socket, 'context', None) + try: + socket.close(linger=0) + except Exception as e: + logger.warning('Error closing ZMQ socket: %s', e) + if ctx is not None: + try: + ctx.term() + except Exception as e: + logger.warning('Error terminating ZMQ context: %s', e) + self.rank = None + self.world_size = None + self._rendezvous = None + self._group_world_size = None + self.hccl_env = None + self.ip = self.zmq_port = self.dist_port = None + self._shard_seeded = False + + def finalize(self): + """Per-round finalize is a no-op: the group persists until ``close()``.""" + + # ---- transport hooks ---- + @property + def _wire_device(self): + return torch.device('npu', torch.npu.current_device()) + + def _ensure_wire(self) -> None: + self._ensure_communicator() + + def _broadcast_flush(self, pos_u8, val_u8) -> None: + comm = self._comm + if pos_u8 is not None: + comm.broadcast(pos_u8, src=0) + comm.broadcast(val_u8, src=0) + # PyHccl enqueues on the current stream and the caller frees the buffers right after. + torch.npu.synchronize() + + def _release_after_round(self, phase: str) -> None: + torch.npu.empty_cache() + logger.debug('delta HCCL %s round done, NPU cache released', phase) + + # ---- not provided on NPU ---- + @classmethod + def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metadata: list[dict]) -> Any: + raise NotImplementedError('the delta manager builds its own topology; vLLM workers are the group receivers') + + async def receive_weights(self, *args, **kwargs): + raise NotImplementedError('there is no rollout-side checkpoint engine on NPU; vLLM workers receive directly') diff --git a/src/twinkle/checkpoint_engine/delta/transport/nccl_engine.py b/src/twinkle/checkpoint_engine/delta/transport/nccl_engine.py new file mode 100644 index 000000000..ad69d9edd --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/transport/nccl_engine.py @@ -0,0 +1,111 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""NCCL transport for the delta pipeline (the pipeline itself lives in :mod:`.pipeline`).""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Iterator + +import torch + +from ...nccl_checkpoint_engine import MasterMetadata, NCCLCheckpointEngine, _pg_broadcast +from twinkle.utils.network import find_free_port +# Re-exported for callers that import the pipeline internals from this module. +from .pipeline import ( # noqa: F401 + DeltaSyncMixin, + _bucket_sliced, + _FlushBucket, + _FlushPiece, + _GatherQueue, + _prodshape, + _ValuesPiece, +) + +logger = logging.getLogger(__name__) + + +class DeltaShardedCheckpointEngine(DeltaSyncMixin, NCCLCheckpointEngine): + """Sparse delta weight sync over NCCL, diffed on each rank's local shard. + + The transport-agnostic pipeline (``send_weights`` state machine, bucketing, + gather batching, flush assembly) comes from :class:`DeltaSyncMixin`; this class + is the NCCL transport: it reuses NCCLCheckpointEngine's group/zmq machinery, + implements the four transport hooks on top of ``_pg_broadcast`` and holds the + rollout-side ``receive_weights``. + """ + + wire_format = "delta_flush" + + def prepare(self) -> MasterMetadata | None: + if not self._prepared: + if self.is_master: + self._start_zmq_server() + self._nccl_store_host = self.ip + self._nccl_store_port = find_free_port() + self._prepared = True + if self.is_master: + return MasterMetadata(self.ip, self.listen_port, self._nccl_store_host, self._nccl_store_port) + return None + + def close(self): + self.rebuild_group = True + self.finalize() + + def _broadcast(self, tensor): + if tensor.numel(): + _pg_broadcast(self._pg, tensor) + + # ---- transport hooks ---- + @property + def _wire_device(self): + return "cuda" + + def _broadcast_flush(self, pos_u8, val_u8) -> None: + if pos_u8 is not None: + self._broadcast(pos_u8) + self._broadcast(val_u8) + + def _release_after_round(self, phase: str) -> None: + """NCCL keeps no per-round scratch memory.""" + + # ---- rollout worker side ---- + def receive_weights(self, global_steps: int | None = None) -> Iterator[tuple[list[tuple[str, torch.Tensor]], bool]]: + """Yield delta flushes for the rollout adapter to consume. + + Each ``(named_tensors, is_last)`` item contains a JSON manifest, optional + int32 position bytes, and patch values. The generator keeps at most one + received GPU flush live at a time. + """ + assert self.rank > 0, "Rank 0 should not receive weights." + applied = 0 + while True: + self.socket.recv_string() + meta = self.socket.recv_pyobj() + if meta.get("terminal_empty"): + break + + dense = meta.get("encoding") == "dense" + val_dtype = getattr(torch, meta["val_dtype"]) + elem = torch.empty(0, dtype=val_dtype).element_size() + val_u8 = torch.empty(meta["val_numel"] * elem, dtype=torch.uint8, device=self._wire_device) + if dense: + pos = None + self._broadcast(val_u8) + else: + pos = torch.empty(meta["pos_numel"], dtype=torch.uint8, device=self._wire_device) + self._broadcast(pos) + self._broadcast(val_u8) + val = val_u8.view(val_dtype) + spec_bytes = json.dumps(meta["spec"]).encode() + spec_t = torch.frombuffer(bytearray(spec_bytes), dtype=torch.uint8).to(self._wire_device) + named = [("__delta_spec__", spec_t), ("__values__", val)] + if pos is not None: + named.insert(1, ("__positions__", pos)) + is_last = bool(meta["is_last"]) + yield named, is_last + applied += 1 + del pos, val_u8, val, spec_t + if is_last: + break + logger.info("delta recv v=%s flushes=%d (yielded to server adapter)", global_steps, applied) diff --git a/src/twinkle/checkpoint_engine/delta/transport/pipeline.py b/src/twinkle/checkpoint_engine/delta/transport/pipeline.py new file mode 100644 index 000000000..3b2efb393 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/transport/pipeline.py @@ -0,0 +1,536 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Transport-agnostic delta sync pipeline: bucketing, batching and the seed/steady state machine. + +Everything here is plain torch plus ``torch.distributed`` (the trainer-side gather +runs on the training engine's own process groups), so any transport can share it. +A concrete engine combines ``DeltaSyncMixin`` (FIRST in the MRO) with a transport +parent that contributes the zmq PUB/SUB manifest channel, rendezvous and topology +(e.g. ``NCCLCheckpointEngine``), and implements the transport hooks: + +* ``_wire_device`` -- device (str or torch.device) where flush tensors live +* ``_ensure_wire()`` -- transport lazy-init before each manifest (no-op default) +* ``_broadcast_flush(pos_u8, val_u8)`` -- put one flush's raw buffers on the wire +* ``_release_after_round(phase)`` -- release transport scratch memory after a round + +Data ladder (sender side, steady):: + + backend HF delta ENTRY (slots, dtype, counts, hf_idx, hf_val, group) + --_GatherQueue--> per-SLOT delta on rank 0 --_bucket_*--> _FlushPiece + --_FlushBucket--> FLUSH (DeltaFlush) --_publish_flush--> wire +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Generator +from dataclasses import dataclass + +import torch +import zmq + +from .gather import gather_slot_entries_to_rank0 +from ..protocol.wire import DeltaFlush, DeltaParam +from ..protocol.wire import checksum as _checksum + +logger = logging.getLogger(__name__) + + +def _prodshape(shape) -> int: + n = 1 + for x in shape: + n *= int(x) + return n + + +@dataclass(slots=True) +class _FlushPiece: + """One (possibly sliced) per-parameter piece buffered for a pending indices flush.""" + + name: str + dtype_str: str + shape: list + idx: torch.Tensor + val: torch.Tensor + + +@dataclass(slots=True) +class _ValuesPiece: + """One whole parameter's flat values buffered for the seed sync's values-only flush.""" + + name: str + dtype_str: str + shape: list + flat: torch.Tensor + + +class _FlushBucket: + """One-flush-lookahead bucket pipeline, shared by the steady loop and both + seed streams. Pieces accumulate until ``cap`` bytes; ``seal`` assembles them + into the single pending flush, first emitting the previous pending with + ``is_last=False`` (the lookahead: only the caller's finale knows which flush + is last and emits it with ``is_last=True``). ``assemble`` and ``publish`` + carry the only real differences between the streams -- the wire format + (indexed flush vs values-only flush) and the flush counters.""" + + __slots__ = ("cap", "pieces", "nbytes", "pending", "_assemble", "_publish") + + def __init__(self, cap: int, assemble, publish): + self.cap = int(cap) + self.pieces: list = [] + self.nbytes = 0 + self.pending = None + self._assemble = assemble + self._publish = publish + + def add(self, piece, nbytes: int) -> None: + self.pieces.append(piece) + self.nbytes += int(nbytes) + if self.nbytes >= self.cap: + self.seal() + + def seal(self) -> None: + if not self.pieces: + return + self.emit(is_last=False) + self.pending = self._assemble(self.pieces) + self.pieces, self.nbytes = [], 0 + + def emit(self, is_last: bool) -> None: + if self.pending is not None: + self._publish(self.pending, is_last) + self.pending = None + + +def _bucket_sliced(bkt: _FlushBucket, name: str, dtype_str: str, shape, aidx: torch.Tensor, aval: torch.Tensor) -> None: + """Slice one param's (idx, val) delta into <= MAX_ENTRY_ELEMS pieces and bucket + them (bounds the receiver-side decode transient; the masked apply is sequential, + so splitting is transparent). Bucket bytes = actual wire bytes (int32 positions + + values).""" + max_elems = DeltaSyncMixin.MAX_ENTRY_ELEMS + for s in range(0, aidx.numel(), max_elems): + e = min(s + max_elems, aidx.numel()) + bkt.add( + _FlushPiece(name, dtype_str, list(shape), aidx[s:e], aval[s:e]), + (e - s) * (4 + aval.element_size()), + ) + + +class _GatherQueue: + """Per-gather-group batching of slot-keyed queue entries + ``(slots, dtype_str, counts, idx, val)``. Entries carry FINAL-coordinate + payloads (identity specs: one slot = the param itself; converter specs: the + backend HF slots), so rank 0 never converts -- ``consume`` receives assembled + per-slot pieces straight from the gather. + + One queue per ProcessGroup: separate queues stop pg alternation (dense fsdp + group vs expert world group per layer) from shattering batches. The flush + trigger is COUNT-ONLY: entry counts are identical on every rank while byte + totals are not, so a count trigger is the only one that keeps the collective + sequence identical across ranks (a per-rank byte trigger desyncs the gathers + and deadlocks NCCL). Byte bounding happens INSIDE the batched gather via + ``max_round_bytes``, decided from the all-gathered counts every rank sees.""" + + __slots__ = ("batch_k", "max_round_bytes", "is_r0", "_consume", "_queues") + + def __init__(self, batch_k: int, max_round_bytes: int, is_r0: bool, consume): + self.batch_k = max(int(batch_k), 1) + self.max_round_bytes = int(max_round_bytes) + self.is_r0 = is_r0 + self._consume = consume + self._queues: dict[int, tuple] = {} # id(pg) -> (pg, [entries]) + + def put(self, pg, slots: list, dtype_str: str, counts: torch.Tensor, idx: torch.Tensor, val: torch.Tensor): + _pg, entries = self._queues.setdefault(id(pg), (pg, [])) + entries.append((slots, dtype_str, counts, idx, val)) + if len(entries) >= self.batch_k: + self._flush(pg, entries) + + def flush_all(self) -> None: + for pg, entries in self._queues.values(): + self._flush(pg, entries) + + def _flush(self, pg, entries: list) -> None: + """One gather round for one group's queue.""" + if not entries: + return + batch = list(entries) + entries.clear() + if pg is None: + # unsharded/replicated params: rank 0's local delta already is global + if self.is_r0: + for slots, dtype_str, counts, idx, val in batch: + off = 0 + for (name, shape), c in zip(slots, counts.tolist(), strict=True): + self._consume( + name, dtype_str, tuple(shape), _prodshape(shape), idx[off : off + c], val[off : off + c] + ) + off += c + return + dev = batch[0][3].device + counts_concat = torch.cat([c for _, _, c, _, _ in batch]).to(dev) + idx_concat = torch.cat([i for _, _, _, i, _ in batch]) + val_concat = torch.cat([v for _, _, _, _, v in batch]) + gathered = gather_slot_entries_to_rank0( + idx_concat, val_concat, counts_concat, group=pg, max_round_bytes=self.max_round_bytes + ) + if self.is_r0 and gathered is not None: + slot_i = 0 + for slots, dtype_str, _counts, _i, _v in batch: + for name, shape in slots: + aidx, aval = gathered[slot_i] + slot_i += 1 + self._consume(name, dtype_str, tuple(shape), _prodshape(shape), aidx, aval) + + +class DeltaSyncMixin: + """Transport-agnostic sparse delta sync pipeline (see module docstring for the hooks).""" + + # Cap on changed elements per DeltaParam entry. The receiver-side decode + # densifies per entry with an int64 index transient (8 B/element), so an + # uncapped entry (e.g. a 7B model's whole embedding on the full seed, ~545M + # elements) would spike several GiB at once. Oversized per-param deltas are + # sliced into multiple entries (the masked apply is sequential, so splitting + # is transparent); 64M elements bounds the transient to ~512 MiB. + MAX_ENTRY_ELEMS = 64 << 20 + + # ---- transport hooks (implemented by the concrete engine) ---- + @property + def _wire_device(self): + """Device (str or torch.device) where flush tensors are assembled and broadcast from.""" + raise NotImplementedError + + def _broadcast_flush(self, pos_u8: torch.Tensor | None, val_u8: torch.Tensor) -> None: + """Put one flush's raw uint8 buffers on the wire (positions are None for values-only flushes).""" + raise NotImplementedError + + def _ensure_wire(self) -> None: + """Transport lazy-init, invoked BEFORE each manifest is published. + + A lazily built communicator (e.g. HCCL) completes only once every receiver + has joined, and every receiver connects its zmq SUB before joining, so + "group ready" implies "subscribed" -- which is what makes publishing the + manifest after this hook lossless (a zmq PUB drops messages that predate + their subscribers). Transports that set the wire up eagerly (the NCCL group) + leave this as a no-op. + """ + + def _release_after_round(self, phase: str) -> None: + """Release transport scratch memory after a round (``phase``: 'seed' or 'steady', rank 0 only).""" + raise NotImplementedError + + # ---- trainer side ---- + # ---- shared STREAMING wire ---- + # Stream each flush as soon as it is produced so trainer peak memory stays near two buckets, + # including during the full seed. Each ZMQ manifest + NCCL broadcast carries ``is_last``; + # rollout adapters then forward that flush over same-GPU IPC to their backend-specific consumer. + def _publish_flush(self, flush: DeltaFlush, first: bool, is_last: bool) -> None: + meta = { + "is_full": first, + "encoding": self.encoding, + "is_last": is_last, + "terminal_empty": False, + "pos_numel": int(flush.positions_cpu.numel()), + "val_numel": int(flush.values_gpu.numel()), + "val_dtype": str(flush.values_gpu.dtype).replace("torch.", ""), + "spec": { + "encoding": self.encoding, + "params": [vars(p) for p in flush.params], + "checksum": int(flush.checksum), + }, + } + self._ensure_wire() + self.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + pos_u8 = flush.positions_cpu.to(self._wire_device, non_blocking=True).contiguous().view(torch.uint8) + val_u8 = flush.values_gpu.contiguous().view(torch.uint8) + self._broadcast_flush(pos_u8, val_u8) + + def _publish_values_flush( + self, params: list[DeltaParam], values: torch.Tensor, is_last: bool + ) -> None: + """Publish a values-only (full-coverage, positions-free) flush -- used by the first + sync. The wire encoding tag stays ``"dense"`` -- it is protocol, shared with the + receiver's delta_loader decode.""" + values = values.contiguous() + empty_pos = torch.empty(0, dtype=torch.uint8, device=values.device) + meta = { + "is_full": True, + "encoding": "dense", + "is_last": is_last, + "terminal_empty": False, + "pos_numel": 0, + "val_numel": int(values.numel()), + "val_dtype": str(values.dtype).replace("torch.", ""), + "spec": { + "encoding": "dense", + "verify": False, + "is_last": is_last, + "params": [vars(p) for p in params], + "checksum": int(_checksum(empty_pos, values)), + }, + } + self._ensure_wire() + self.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + self._broadcast_flush(None, values.view(torch.uint8)) + + def _publish_terminal(self, first: bool) -> None: + """End-of-stream marker when zero flushes were produced (no broadcast, just a signal).""" + meta = {"is_full": first, "encoding": self.encoding, "is_last": True, "terminal_empty": True} + self._ensure_wire() + self.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + + def __init__( + self, *args, encoding: str = "indices", batch_gather: int = 32, **kwargs + ) -> None: + super().__init__(*args, **kwargs) + assert encoding == "indices", f"delta_sharded ships only the 'indices' position encoding; got {encoding!r}" + self.encoding = encoding + self._shard_seeded = False + # Gather the per-param sparse deltas in groups of this many parameters + # (one count-matrix all_gather + two padded gathers per group instead of + # three collectives per parameter). + self.batch_gather = int(batch_gather) + + def _assemble_flush(self, per_param: list[_FlushPiece]) -> DeltaFlush: + """Build one DeltaFlush (indices encoding) from rank 0's gathered per-param deltas. + + ``per_param``: :class:`_FlushPiece` entries whose ``idx`` are within-parameter + flat positions (== what the receiver decodes). + + Positions stay on the GPU end to end (int32 pieces -> one cat -> uint8 view); + the wire broadcasts from the GPU anyway, and a host round-trip here + (``.cpu().numpy().tobytes()`` + join) dominated the whole send at scale + (~2.4s/sync at 7B steady state, ~83s on the full seed). + """ + idx_pieces: list[torch.Tensor] = [] + val_pieces: list[torch.Tensor] = [] + params: list[DeltaParam] = [] + pos_off = val_off = 0 + for piece in per_param: + nnz = int(piece.idx.numel()) + # positions ride the wire as int32 (pos_width=4); a parameter bigger than + # 2^31 elements would silently wrap, so fail loud instead. DeltaParam + # carries pos_width for a future 8-byte escalation if a model needs it. + assert _prodshape(piece.shape) < (1 << 31), ( + f"{piece.name}: {_prodshape(piece.shape)} elements exceeds the int32 position encoding" + ) + idx_pieces.append(piece.idx.to(torch.int32)) + val_pieces.append(piece.val) + params.append( + DeltaParam( + name=piece.name, + dtype=piece.dtype_str, + shape=list(piece.shape), + pos_start=pos_off, + pos_end=pos_off + nnz * 4, + pos_width=4, + val_start=val_off, + val_end=val_off + nnz, + ) + ) + pos_off += nnz * 4 + val_off += nnz + + values_gpu = torch.cat(val_pieces) if val_pieces else torch.empty(0, dtype=self.rollout_dtype, device=self._wire_device) + positions_u8 = ( + torch.cat(idx_pieces).contiguous().view(torch.uint8) + if idx_pieces + else torch.empty(0, dtype=torch.uint8, device=values_gpu.device) + ) + cks = _checksum(positions_u8, values_gpu) + return DeltaFlush( + encoding=self.encoding, params=params, positions_cpu=positions_u8, values_gpu=values_gpu, checksum=cks + ) + + def _send_full_seed( + self, + weights: Generator[tuple[str, torch.Tensor], None, None], + global_steps: int | None = None, + ) -> dict[str, float] | None: + """First sync: stream the backend's FULL HF export over the values-only wire. + + ``weights`` is ``get_per_tensor_param()`` -- every backend already knows how + to assemble and convert its own full tensors (FSDP all-gather, veomni expert + restack, Megatron TP/PP fusion), so the seed inherits all of that for free + and this engine only buckets and broadcasts. Every trainer rank iterates the + generator (the per-tensor assembly is collective); rank 0 buckets. Resume + works by construction: whatever the trainer restored is what ships.""" + is_r0 = self.is_master + t0 = time.time() + n_flushes = 0 + total_elems = 0 + wire_bytes = 0 + + def _assemble_values(pieces: list[_ValuesPiece]): + params = [] + val_off = 0 + for piece in pieces: + n = int(piece.flat.numel()) + params.append( + DeltaParam( + name=piece.name, + dtype=piece.dtype_str, + shape=list(piece.shape), + pos_start=0, + pos_end=0, + pos_width=4, + val_start=val_off, + val_end=val_off + n, + ) + ) + val_off += n + return params, torch.cat([piece.flat for piece in pieces]) + + def _publish_values(pending, is_last: bool) -> None: + nonlocal n_flushes, wire_bytes + params, values = pending + self._publish_values_flush(params, values, is_last=is_last) + n_flushes += 1 + wire_bytes += int(values.nbytes) + + bkt = _FlushBucket(self.bucket_size, _assemble_values, _publish_values) + + seen_names: set = set() + for name, tensor in weights: + # duplicate names in a full export mean two source params mapped to + # the same HF tensor -- the receiver would apply whichever came last + # and the delta diff base would silently disagree with the trainer + # (observed: NemotronH A_log). Fail loud instead. + assert name not in seen_names, f"full export yields duplicate HF tensor {name!r}" + seen_names.add(name) + tensor = tensor.detach() + if tensor.is_floating_point() and tensor.dtype != self.rollout_dtype: + tensor = tensor.to(self.rollout_dtype) + if not is_r0: + del tensor + continue + flat = tensor.contiguous().reshape(-1) + total_elems += int(flat.numel()) + bkt.add(_ValuesPiece(name, str(flat.dtype).replace("torch.", ""), list(tensor.shape), flat), flat.nbytes) + + if not is_r0: + return + bkt.seal() + if bkt.pending is not None: + bkt.emit(is_last=True) + else: + self._publish_terminal(True) + self._release_after_round("seed") + # warning level on purpose: worker default log level swallows info, and the + # one-off seed cost is the number people ask for when sizing a run. + logger.warning( + "delta-sharded FULL-SEED v=%s done in %.1fs (flushes=%d elems=%d wire=%.1fGB)", + global_steps, + time.time() - t0, + n_flushes, + total_elems, + wire_bytes / (1 << 30), + ) + if not total_elems: + return None + return { + "checkpoint_engine/changed_ratio": 1.0, + "checkpoint_engine/changed_elems": float(total_elems), + "checkpoint_engine/payload_mbytes": wire_bytes / (1 << 20), + "checkpoint_engine/flushes": float(n_flushes), + } + + async def send_weights( + self, + engine, + global_steps: int | None = None, + ) -> dict[str, float] | None: + """Drive one weight sync from the TRAINING ENGINE (unlike the full-sync + engines, which consume a weights iterator): the seed/steady phase choice + and the snapshot prime are this engine's own state machine, so the worker + stays delta-agnostic. The seed (first sync) streams the backend's full + ``get_per_tensor_param()`` export over the values-only wire, then pins the + diff base via ``engine.prime_delta_snapshots()``; every later sync consumes + ``get_per_tensor_param_delta_shard()`` (backend-computed final-HF deltas). + """ + # All actor ranks participate (gather-v is collective); only torch rank 0 broadcasts. + # rank 0 accumulates the gathered per-param deltas into bucket_size-sized flushes and streams + # each one as soon as it fills (then frees it), so peak memory is ~2 buckets rather than the + # whole model. + assert self.rank <= 0, "Trainer workers other than rank 0 should not send weights." + if not self._shard_seeded: + full, _ = engine.get_per_tensor_param() + metrics = self._send_full_seed(full, global_steps) + # weights do not move during the sync, so the snapshots equal exactly + # what the rollout just received. + engine.prime_delta_snapshots() + # The seed is complete only after its baseline snapshots are ready. + self._shard_seeded = True + return metrics + weights, _ = engine.get_per_tensor_param_delta_shard() + is_r0 = self.is_master + t0 = time.time() + n_flushes = 0 + changed_elems = 0 + total_elems = 0 + wire_bytes = 0 + + def _publish_steady(flush, is_last: bool) -> None: + nonlocal n_flushes + self._publish_flush(flush, first=False, is_last=is_last) + n_flushes += 1 + + bkt = _FlushBucket(self.bucket_size, self._assemble_flush, _publish_steady) + + batch_k = self.batch_gather + + def _bucket_slot_delta( + name: str, + dtype_str: str, + full_shape: tuple, + full_numel: int, + aidx: torch.Tensor | None, + aval: torch.Tensor | None, + ) -> None: + nonlocal total_elems, changed_elems, wire_bytes + total_elems += int(full_numel) + if aidx is None or aidx.numel() == 0: + return + changed_elems += int(aidx.numel()) + wire_bytes += int(aidx.numel()) * (4 + aval.element_size()) + _bucket_sliced(bkt, name, dtype_str, full_shape, aidx, aval) + + gq = _GatherQueue(batch_k, self.bucket_size, is_r0, _bucket_slot_delta) + + # ``weights`` is the BACKEND's HF delta stream (hf_delta_export): entries + # already carry final HF coordinates -- naming, conversion, diff and + # snapshot all happened on the backend side. This engine only batches, + # gathers and ships. + for slots, dtype_str, counts, hf_idx, hf_val, pg in weights: + gq.put(pg, slots, dtype_str, counts, hf_idx, hf_val) + gq.flush_all() + + if is_r0: + bkt.seal() + if bkt.pending is not None: + bkt.emit(is_last=True) + else: + self._publish_terminal(False) + if not is_r0: + return + self._release_after_round("steady") + elapsed = time.time() - t0 + mb = wire_bytes / (1 << 20) + logger.info( + "delta-sharded steady send v=%s: changed_ratio=%.4f wire_payload=%.1fMB wire_time=%.2fs " + "throughput=%.1fMB/s flushes=%d", + global_steps, + (changed_elems / total_elems) if total_elems else 0.0, + mb, elapsed, (mb / elapsed) if elapsed > 0 else 0.0, n_flushes, + ) + if not total_elems: + return None + return { + "checkpoint_engine/changed_ratio": changed_elems / total_elems, + "checkpoint_engine/changed_elems": float(changed_elems), + "checkpoint_engine/payload_mbytes": wire_bytes / (1 << 20), + "checkpoint_engine/flushes": float(n_flushes), + } diff --git a/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py b/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py index 314176ac9..d413721f5 100644 --- a/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py +++ b/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py @@ -419,8 +419,9 @@ def _flush(is_last: bool): elapsed = time.time() - start_time bandwidth = total_bytes / elapsed / (1024 * 1024 * 1024) if elapsed > 0 else 0.0 - logger.info(f'send_weights done: rank={self.rank}, params={total_params}, chunks={total_chunks}, ' - f'time={elapsed:.2f}s, bandwidth={bandwidth:.2f} GB/s') + logger.info(f'full-sync send_weights done: rank={self.rank}, params={total_params}, chunks={total_chunks}, ' + f'payload={total_bytes / (1 << 20):.1f}MB ({total_bytes / (1 << 30):.2f}GB), ' + f'time={elapsed:.2f}s, bandwidth={bandwidth:.2f} GB/s') @torch.no_grad() async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor]]: @@ -507,5 +508,6 @@ async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor]]: elapsed = time.time() - start_time bandwidth = total_bytes / elapsed / (1024 * 1024 * 1024) if elapsed > 0 else 0.0 - logger.info(f'receive_weights done: rank={self.rank}, params={total_params}, chunks={total_chunks}, ' - f'time={elapsed:.2f}s, bandwidth={bandwidth:.2f} GB/s') + logger.info(f'full-sync receive_weights done: rank={self.rank}, params={total_params}, chunks={total_chunks}, ' + f'payload={total_bytes / (1 << 20):.1f}MB ({total_bytes / (1 << 30):.2f}GB), ' + f'time={elapsed:.2f}s, bandwidth={bandwidth:.2f} GB/s') diff --git a/src/twinkle/checkpoint_engine/manager.py b/src/twinkle/checkpoint_engine/manager.py index c533437dc..51611e608 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -1,12 +1,15 @@ # Copyright (c) ModelScope Contributors. All rights reserved. # Adapted from https://github.com/volcengine/verl/blob/main/verl/checkpoint_engine/base.py import threading +from contextlib import contextmanager +from time import perf_counter from typing import List, Literal, Optional from twinkle.infra import get_worker_id from twinkle import Platform, get_logger from .base import CheckpointEngine +from .delta import DeltaSyncConfig, resolve_delta_engine_cls from .mixin import CheckpointEngineMixin from .zmq_auth import generate_auth_key @@ -28,6 +31,9 @@ class CheckpointEngineManager: ``auto`` resolves local objects to ``naive`` and Ray actor handlers to ``standalone``. It never guesses ``colocate`` because actor placement cannot be inferred reliably from the driver. + Delta synchronization follows verl's workflow ordering: finish training before syncing, + await sync before submitting more training, and do not overlap sync and close calls. + Architecture (following verl's CheckpointEngineManager): Trainer GPU(s) Rollout GPU(s) @@ -76,6 +82,8 @@ def __init__( sampler: 'CheckpointEngineMixin', platform: str = 'GPU', mode: CheckpointEngineMode = 'auto', + sync_backend: Literal['full', 'delta_sharded'] = 'full', + delta_config: Optional[DeltaSyncConfig] = None, channel: str = 'default', colocate_trainer_ranks: Optional[List[int]] = None, receiver_worker_ids: Optional[List[str]] = None, @@ -85,6 +93,22 @@ def __init__( self.requested_mode = mode self.mode = self._resolve_mode(mode, model, sampler) self.backend_cls = self.decide_backend_engine(platform, self.mode) + if sync_backend not in ('full', 'delta_sharded'): + raise ValueError(f'Unknown weight sync backend: {sync_backend}') + if sync_backend == 'delta_sharded' and self.mode != 'standalone': + raise ValueError('delta_sharded requires standalone trainer and rollout actors') + self.sync_backend = sync_backend + self.delta_config = delta_config or DeltaSyncConfig() + self.delta_engine_cls = None + self._delta_platform = None + if sync_backend == 'delta_sharded': + # Fail at construction, not mid-sync, when the platform has no delta transport. + self._delta_platform = Platform.get_platform(platform).__name__ + self.delta_engine_cls = resolve_delta_engine_cls(self._delta_platform) + self._delta_replica_count = None + self._delta_round = 0 + self._delta_failed = False + self.last_delta_metrics = None self._sync_lock = threading.RLock() self.receiver_worker_ids = self._validate_receiver_ids(receiver_worker_ids) self._receiver_snapshot = self.receiver_worker_ids @@ -292,6 +316,10 @@ def _sync_weights(self, merge_and_sync=True): Returns: None """ + if self.sync_backend == 'delta_sharded': + with self._delta_rpc_timeouts(): + self._sync_delta_weights(merge_and_sync) + return if self.mode == 'naive': self._sync_weights_naive(merge_and_sync) return @@ -343,6 +371,113 @@ def _sync_weights(self, merge_and_sync=True): if not merge_and_sync: logger.info('Base model sync completed, subsequent syncs will be LoRA-only') + def _prepare_delta_transport(self): + capabilities = self.sampler.validate_delta_sampler(self.delta_config) + replicas = capabilities if isinstance(capabilities, list) else [capabilities] + if self._delta_platform == 'NPU': + return self._prepare_delta_transport_npu(replicas) + workers = self.sampler.delta_sampler_metadata() + if len(workers) != sum(replica['tp'] for replica in replicas): + raise ValueError('delta topology does not include every rollout TP worker') + devices = {(worker['node_ip'], worker['device_uuid']) for worker in workers} + if len(devices) != len(workers): + raise ValueError('delta checkpoint receivers must occupy distinct rollout GPUs') + metadata = self.model.prepare_delta_sync(self.delta_config) + return dict(metadata, world_size=len(workers) + 1, rollout_workers=workers), len(replicas) + + def _prepare_delta_transport_npu(self, replicas): + """NPU rendezvous: workers join the trainer's HCCL group themselves, so there is no worker list. + + Group rank 0 is the trainer; rollout workers follow replica-major, TP-minor, so the world size is + every replica's TP size plus the trainer. + """ + metadata = self.model.prepare_delta_sync(self.delta_config) + for replica in replicas: + if replica['device'] != metadata['device']: + raise ValueError(f"delta device mismatch: trainer is {metadata['device']!r}, " + f"rollout worker is {replica['device']!r}") + ranks = [replica.get('replica_rank') for replica in replicas] + if sorted(ranks, key=lambda rank: (rank is None, rank)) != list(range(len(replicas))): + raise ValueError(f'delta replica ranks {ranks} must be exactly 0..{len(replicas) - 1} without duplicates') + world_size = sum(replica['tp'] for replica in replicas) + 1 + return dict(metadata, world_size=world_size), len(replicas) + + def _sync_delta_weights(self, merge_and_sync): + if not merge_and_sync: + raise ValueError('delta_sharded does not support LoRA synchronization') + if self._delta_failed: + raise RuntimeError('delta session failed; restart the entire training/rollout job before full seed') + started = perf_counter() + seed = self._delta_replica_count is None + if seed: + endpoint, replica_count = self._prepare_delta_transport() + try: + self.sampler.begin_delta_sync() + if seed: + trainer_init = self.model.init_delta_transport(endpoint) + sampler_init = self.sampler.init_delta_sampler_transport(endpoint) + trainer_init() + sampler_init() + self._delta_replica_count = replica_count + sender = self.model.send_delta_weights() + receiver = self.sampler.receive_delta_weights() + metrics = sender() + worker_metrics = receiver() + results = worker_metrics if isinstance(worker_metrics, list) else [worker_metrics] + if len(results) != self._delta_replica_count or any( + result['changed'] != results[0]['changed'] for result in results): + raise ValueError('delta rollout replicas disagree about update completion') + self.sampler.publish_delta_sampler() + self._delta_round += 1 + self.base_sync_done = True + self.last_delta_metrics = {'trainers': metrics, 'workers': worker_metrics, + 'total_seconds': perf_counter() - started, 'seed': seed} + except Exception: + self._delta_failed = True + try: + self.sampler.fail_delta_sampler() + except Exception: + logger.exception('Could not pause failed delta sampler; restart rollout actors before serving') + raise + + def close_sync_session(self): + """Release delta snapshots and communicators; keep rollout paused until a new full seed.""" + with self._delta_rpc_timeouts(): + if self.sync_backend == 'delta_sharded': + self._close_delta_session() + + @contextmanager + def _delta_rpc_timeouts(self): + missing = object() + previous = [(actor, getattr(actor, '_ray_get_timeout', missing)) for actor in (self.model, self.sampler)] + try: + for actor, _ in previous: + actor._ray_get_timeout = self.delta_config.timeout_seconds + yield + finally: + for actor, timeout in previous: + if timeout is missing: + # Ray driver proxies can inherit nn.Module without initializing its registries. + object.__delattr__(actor, '_ray_get_timeout') + else: + actor._ray_get_timeout = timeout + + def _close_delta_session(self): + errors = [] + close_methods = (self.sampler.close_delta_sampler, self.model.close_delta_trainer) + for close in close_methods: + try: + close() + except Exception as error: + errors.append(error) + if errors: + self._delta_failed = True + raise RuntimeError('delta cleanup failed; restart the affected actors') from errors[0] + self._delta_replica_count = None + self._delta_round = 0 + # Cleanup cannot make an unknown partially updated model safe for a new session. + self.base_sync_done = False + def _ensure_model_keys(self): if self._model_keys is not None: return diff --git a/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py b/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py index a4cdc0d7b..4cc45e8f0 100644 --- a/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py +++ b/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py @@ -10,9 +10,10 @@ from typing import Any, AsyncGenerator, Generator from twinkle import get_logger -from twinkle.utils.network import find_free_port, is_valid_ipv6_address +from twinkle.utils.network import find_free_port from .base import CheckpointEngine, TensorMeta from .zmq_auth import generate_auth_key, recv_auth_pyobj, send_auth_pyobj +from .zmq_pubsub import connect_subscriber, start_publisher logger = get_logger() @@ -156,29 +157,12 @@ def _start_zmq_server(self): self.listen_port = find_free_port() self.auth_key = generate_auth_key() # 256-bit shared secret - context = zmq.Context() - self.socket = context.socket(zmq.PUB) - if is_valid_ipv6_address(self.ip): - address = f'tcp://[{self.ip}]:{self.listen_port}' - self.socket.setsockopt(zmq.IPV6, 1) - else: - address = f'tcp://{self.ip}:{self.listen_port}' - - self.socket.bind(address) + self.socket = start_publisher(self.ip, self.listen_port) def _connect_zmq_client(self, metadata: MasterMetadata): """Connect to the ZMQ PUB server as a subscriber (receiver only).""" self.auth_key = metadata.auth_key # inherit the master's HMAC secret - context = zmq.Context() - self.socket = context.socket(zmq.SUB) - if is_valid_ipv6_address(metadata.zmq_ip): - address = f'tcp://[{metadata.zmq_ip}]:{metadata.zmq_port}' - self.socket.setsockopt(zmq.IPV6, 1) - else: - address = f'tcp://{metadata.zmq_ip}:{metadata.zmq_port}' - - self.socket.connect(address) - self.socket.setsockopt_string(zmq.SUBSCRIBE, self.topic) + self.socket = connect_subscriber(metadata.zmq_ip, metadata.zmq_port, self.topic) # ── Core lifecycle ─────────────────────────────────────────────────── diff --git a/src/twinkle/checkpoint_engine/zmq_pubsub.py b/src/twinkle/checkpoint_engine/zmq_pubsub.py new file mode 100644 index 000000000..a67ab0c28 --- /dev/null +++ b/src/twinkle/checkpoint_engine/zmq_pubsub.py @@ -0,0 +1,36 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""zmq PUB/SUB helpers for the manifest channel, shared by the NCCL and HCCL engines.""" + +from __future__ import annotations + +import zmq + +from twinkle.utils.network import is_valid_ipv6_address + + +def format_tcp_address(ip: str, port: int) -> tuple[str, bool]: + """Return ``(tcp address, is_ipv6)``; IPv6 literals are bracketed.""" + if is_valid_ipv6_address(ip): + return f'tcp://[{ip}]:{port}', True + return f'tcp://{ip}:{port}', False + + +def start_publisher(ip: str, port: int) -> zmq.Socket: + """Bind a PUB socket on ``ip:port`` (on its own context, reachable as ``socket.context``).""" + socket = zmq.Context().socket(zmq.PUB) + address, ipv6 = format_tcp_address(ip, port) + if ipv6: + socket.setsockopt(zmq.IPV6, 1) + socket.bind(address) + return socket + + +def connect_subscriber(ip: str, port: int, topic: str) -> zmq.Socket: + """Connect a SUB socket to ``ip:port`` and subscribe to ``topic``.""" + socket = zmq.Context().socket(zmq.SUB) + address, ipv6 = format_tcp_address(ip, port) + if ipv6: + socket.setsockopt(zmq.IPV6, 1) + socket.connect(address) + socket.setsockopt_string(zmq.SUBSCRIBE, topic) + return socket diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 61b1ead65..9b386f517 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -199,7 +199,10 @@ def initialize(nproc_per_node: int, ncpu_proc_per_node: int, device_groups: List import ray RayHelper.device_groups = device_groups if not RayHelper.ray_inited(): - resources = RayHelper._get_ray_custom_resources(device_groups) + # Connecting to an existing cluster (e.g. RAY_ADDRESS set for multi-node): custom + # resources are declared at ``ray start`` time, and Ray rejects ``resources`` on connect. + connecting = bool(os.environ.get('RAY_ADDRESS')) + resources = None if connecting else RayHelper._get_ray_custom_resources(device_groups) ray.init(ignore_reinit_error=True, resources=resources or None) if RayHelper.resource_manager is None: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 9c061ce7f..f66edb889 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -27,6 +27,8 @@ from twinkle import DeviceMesh, Platform, remote_class, remote_function from twinkle.checkpoint_engine import CheckpointEngine from twinkle.checkpoint_engine.mixin import CheckpointEngineMixin +from twinkle.checkpoint_engine.delta.trainer.backends import FSDP2DeltaBackend +from twinkle.checkpoint_engine.delta.trainer.mixin import DeltaTrainerMixin from twinkle.checkpoint_engine.weight_keys import (add_base_layer_suffix, is_lora_key, normalize_param_name, print_weight_example) from twinkle.data_format import InputFeature, ModelOutput, SampleResponse, SamplingParams, Trajectory @@ -301,7 +303,8 @@ def _read_hf_state_dict(checkpoint_dir: str) -> Dict[str, torch.Tensor]: @remote_class() -class TransformersModel(TrainableModel, PreTrainedModel, CheckpointEngineMixin): +class TransformersModel(TrainableModel, PreTrainedModel, CheckpointEngineMixin, DeltaTrainerMixin, + FSDP2DeltaBackend): """The transformers model wrapper. Args: @@ -697,7 +700,6 @@ def __repr__(self) -> str: # into the exception string and that dump swallows the real traceback. return f'{type(self).__name__}(model_id={getattr(self, "model_id", None)!r})' - @staticmethod def _lazy_wrap_model(self): if not self._model_wrapped: optimizer_groups = [og for og in self.optimizer_group.values() if og.optimizer is not None] @@ -2805,6 +2807,42 @@ def _get_weight_generator( # vLLM base_layer renaming) lives in twinkle.checkpoint_engine.weight_keys. model = self.strategy.unwrap_model(self.model) + def _checkpoint_tensor(tensor): + # Stage a (possibly CPU-offloaded) DTensor onto the device before the full_tensor all-gather. + from torch.distributed.tensor import DTensor + if isinstance(tensor, DTensor): + tensor = tensor.to(Torch.get_device(None), non_blocking=True) + return Torch.to_local_tensor(tensor) + + def _tied_duplicate_names() -> set: + """Names the HF state_dict repeats because they share one tensor, which must not be sent. + + With ``tie_word_embeddings`` the LM head and the input embedding are the SAME tensor, so + the state_dict lists it twice (``lm_head.weight`` and ``*.embed_tokens.weight``). vLLM + loads only the name its module traversal reaches first and records the other as a skipped + alias, then asserts the canonical one did load + (``AutoWeightsLoader._check_skipped_aliases``). Sending both makes that check fail, and + the duplicate is redundant anyway because one write covers both. + + Keyed off the config rather than runtime storage identity: the sharing can be broken + apart mid-run (an optimizer step or an FSDP all-gather may materialise the head as its + own tensor), which would make a storage probe stop matching after the first sync. The + flag can sit on a nested ``text_config`` for multimodal shells, so check both. + """ + config = getattr(self, 'hf_config', None) + if config is None: + return set() + tied = getattr(config, 'tie_word_embeddings', None) + if tied is None: + tied = getattr(getattr(config, 'text_config', None), 'tie_word_embeddings', None) + if not tied: + return set() + # Drop the head, keep the embedding: vLLM's traversal reaches the nested embedding first, + # making it the canonical name that ``_check_skipped_aliases`` wants to see loaded. + return {name for name in model.state_dict() if name.split('.')[-2:-1] == ['lm_head']} + + _tied_duplicates = _tied_duplicate_names() + if base_sync_done and adapter_name: if merge_and_sync: # LoRA Training and sync full model(merge_adapter) @@ -2815,9 +2853,9 @@ def weight_generator(): model.merge_adapter() names = [] for name, tensor in model.state_dict().items(): - if is_lora_key(name): + if is_lora_key(name) or name in _tied_duplicates: continue - tensor = Torch.to_local_tensor(tensor) + tensor = _checkpoint_tensor(tensor) name = normalize_param_name(name, keep_base_layer=False) names.append(name) yield name, tensor @@ -2833,7 +2871,7 @@ def weight_generator(): def weight_generator(): names = [] for name, tensor in lora_state_dict.items(): - tensor = Torch.to_local_tensor(tensor) + tensor = _checkpoint_tensor(tensor) name = normalize_param_name(name, keep_base_layer=True) names.append(name) yield name, tensor @@ -2844,15 +2882,17 @@ def weight_generator(): # ``enable_lora`` (``merge_and_sync=False``), rename base weights to # the ``.base_layer.`` form for every module vLLM has LoRA-wrapped # (detected via ``model_keys``); otherwise send canonical names. + from twinkle.checkpoint_engine.delta.trainer.backends.fsdp import convert_weight_keys + add_base_layer = not merge_and_sync - state_dict = model.state_dict() + state_dict = convert_weight_keys(model.state_dict(), model) def weight_generator(): names = [] for name, tensor in state_dict.items(): - if is_lora_key(name): + if is_lora_key(name) or name in _tied_duplicates: continue - tensor = Torch.to_local_tensor(tensor) + tensor = _checkpoint_tensor(tensor) name = normalize_param_name(name, keep_base_layer=False) if add_base_layer: name = add_base_layer_suffix(name, model_keys) diff --git a/src/twinkle/sampler/vllm_sampler/delta/__init__.py b/src/twinkle/sampler/vllm_sampler/delta/__init__.py new file mode 100644 index 000000000..dd2e3918d --- /dev/null +++ b/src/twinkle/sampler/vllm_sampler/delta/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""vLLM-side adapters for delta weight sync (bound to vLLM internals, so kept out of ``checkpoint_engine.delta``). + + transfer_engine.py vLLM ``WeightTransferEngine`` backends: ``twinkle_delta_ipc`` (GPU) and + ``twinkle_delta_hccl`` (NPU); decode each flush and apply it via vLLM's patch loaders + cuda_ipc.py bucketed zmq + CUDA-IPC channel between a CudaDeltaReceiver and its vLLM worker (GPU only) + weight_update_utils.py refresh of derived weights (tied aliases, packed conv layouts) after a sparse patch +""" diff --git a/src/twinkle/sampler/vllm_sampler/delta/cuda_ipc.py b/src/twinkle/sampler/vllm_sampler/delta/cuda_ipc.py new file mode 100644 index 000000000..42e7ea5f3 --- /dev/null +++ b/src/twinkle/sampler/vllm_sampler/delta/cuda_ipc.py @@ -0,0 +1,352 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Bucketed weight transfer based on verl.""" + +import logging +import os +from multiprocessing import shared_memory +from typing import Callable, TypedDict + +import torch +import zmq +from torch.multiprocessing.reductions import reduce_tensor + +def get_device_id(): + return torch.cuda.current_device() + + +def get_device_name(): + return 'cuda' + + +def get_torch_device(): + return torch.cuda + + +def is_support_ipc(): + return True + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) + + +class TensorMetadata(TypedDict): + name: str + shape: torch.Size + dtype: torch.dtype + offset: int + handle: tuple + + +# copy from https://github.com/vllm-project/vllm/blob/main/examples/offline_inference/rlhf_utils.py +def rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor: + func, args = handle + list_args = list(args) + if device_id is not None: + # the key is to change device id to the current device id + # in case two processes have different CUDA_VISIBLE_DEVICES + list_args[6] = device_id + buffer = func(*list_args) + return buffer + + +def create_shared_memory(size: int, name: str): + """Create shared memory for weight transfer. If already exists, attach to it.""" + try: + shm = shared_memory.SharedMemory(name=name, create=True, size=size) + except FileExistsError: + shm = shared_memory.SharedMemory(name=name) + assert shm.size >= size, f"Stale shm segment '{name}': expected {size} bytes, got {shm.size}" + return shm + + +def rebuild_shared_memory(name: str, size: int, dtype=torch.uint8): + """Rebuild tensor from shared memory.""" + shm = shared_memory.SharedMemory(name=name) + tensor = torch.frombuffer(shm.buf[:size], dtype=dtype) + + return tensor, shm + + +class BucketedWeightSender: + """ + Send model weights via bucketed IPC transfer over ZMQ. + + Packs weight tensors into a fixed-size communication buffer and sends them + in buckets to the receiver. Supports CUDA IPC and shared memory fallback. + + Args: + zmq_handle: ZMQ IPC socket path (e.g., "ipc:///tmp/rl-colocate-zmq-.sock") + bucket_size_mb: Communication buffer size in MB + use_shm: Use shared memory instead of CUDA IPC + """ + + def __init__( + self, + zmq_handle: str, + bucket_size_mb: int = 512, + use_shm: bool = False, + ): + self.zmq_handle = zmq_handle + self.bucket_size_mb = bucket_size_mb + self.bucket_size = int(bucket_size_mb) << 20 + self.use_shm = use_shm + + self.zmq_context = zmq.Context.instance() + self.socket = None + self.buffer = None + self.shm = None + + async def async_send_weights(self, weights): + """ + Send weights to the receiver. Accepts a sync generator or async iterator. + + Args: + weights: Generator or async iterator yielding (name, tensor) pairs + """ + try: + self._init_socket() + self._init_buffer() + + # send bucket weights + offset = 0 + bucket_meta: dict[str, TensorMetadata] = {} + async for name, weight in ensure_async_iterator(weights): + # Align each tensor before the receiver views the byte buffer + # using that tensor's dtype. + alignment = weight.element_size() + offset = (offset + alignment - 1) // alignment * alignment + + # fill the tensor bucket + if offset + weight.nbytes > self.bucket_size and len(bucket_meta) > 0: + get_torch_device().synchronize() + self.socket.send_pyobj({"bucket_meta": bucket_meta, "is_last": False}) + self.socket.recv() + bucket_meta = {} + offset = 0 + + if offset + weight.nbytes > self.bucket_size: + assert not self.use_shm, ( + f"Weight {name}({weight.shape}, {weight.dtype}) is too large to fit in the bucket." + f"Please increase rollout.update_weights_bucket_megabytes({self.bucket_size_mb} MB)." + ) + self._direct_send_large_weight(name, weight) + continue + + bucket_meta[name] = { + "name": name, + "shape": weight.shape, + "dtype": weight.dtype, + "offset": offset, + "handle": None, + } + self.buffer[offset : offset + weight.nbytes].view(dtype=weight.dtype).view(weight.shape).copy_( + weight, non_blocking=True + ) + offset += weight.nbytes + + # send the last bucket + name = weight = None + get_torch_device().synchronize() + self.socket.send_pyobj({"bucket_meta": bucket_meta, "is_last": True}) + self.socket.recv() + finally: + self._cleanup() + + def _init_socket(self): + """Initialize ZMQ REQ socket and bind.""" + if self.zmq_handle.startswith("ipc://"): + ipc_path = self.zmq_handle[len("ipc://") :] + try: + os.remove(ipc_path) + except OSError: + pass + self.socket = self.zmq_context.socket(zmq.REQ) + self.socket.bind(self.zmq_handle) + + def _init_buffer(self): + """build communication buffer""" + buffer, shm = None, None + if not self.use_shm: + buffer = torch.empty(self.bucket_size, dtype=torch.uint8, device=f"{get_device_name()}:{get_device_id()}") + handle = reduce_tensor(buffer) + self.socket.send_pyobj(handle) + else: + import uuid + + # Create unique name for shared memory + shm_name = f"verl_weights_{uuid.uuid4().hex}" + shm = create_shared_memory(self.bucket_size, shm_name) + buffer = torch.frombuffer(shm.buf, dtype=torch.uint8) + + comm_metadata = {"name": shm_name, "size": self.bucket_size} + self.socket.send_pyobj(comm_metadata) + + self.socket.recv() + self.buffer = buffer + self.shm = shm + + def _cleanup(self): + """clean up""" + if self.socket is not None: + self.socket.close() + self.socket = None + if self.zmq_handle.startswith("ipc://"): + ipc_path = self.zmq_handle[len("ipc://") :] + try: + os.remove(ipc_path) + except OSError: + pass + del self.buffer + self.buffer = None + if self.shm is not None: + self.shm.close() + self.shm.unlink() + del self.shm + self.shm = None + if is_support_ipc(): + get_torch_device().ipc_collect() + get_torch_device().empty_cache() + + def _direct_send_large_weight(self, name: str, weight: torch.Tensor): + """Send a weight larger than the bucket size via cuda ipc or share memory.""" + logger.debug(f"Direct sending large weight {name}({weight.shape}, {weight.dtype})") + # TODO: support fallback to shared memory + handle = reduce_tensor(weight) + bucket_meta: dict[str, TensorMetadata] = {} + bucket_meta[name] = { + "name": name, + "shape": weight.shape, + "dtype": weight.dtype, + "offset": 0, + "handle": handle, + } + self.socket.send_pyobj({"bucket_meta": bucket_meta, "is_last": False}) + self.socket.recv() + + +class BucketedWeightReceiver: + """ + Receive model weights via bucketed IPC transfer over ZMQ. + + Receives weight tensors from BucketedWeightSender and passes each + bucket to a callback for processing (e.g., loading into the model). + + Args: + zmq_handle: ZMQ IPC socket path (must match sender) + device: Target device for received tensors + use_shm: Use shared memory instead of CUDA IPC + """ + + def __init__( + self, + zmq_handle: str, + device: torch.device, + use_shm: bool = False, + ): + self.zmq_handle = zmq_handle + self.device = device + self.use_shm = use_shm + + self.zmq_context = zmq.Context.instance() + self.socket = None + self.buffer = None + self.shm = None + self._ack_pending = False + + def receive_weights(self, on_bucket_received: callable): + """ + Receive weights from sender and process each bucket via callback. + + Args: + on_bucket_received: Callback function(weights: list[(name, tensor)], + is_last: bool) called per bucket. ``is_last`` marks the final bucket + of a weight-sync round so consumers that need the complete tensor set + (e.g. vLLM ``add_lora``, which takes one adapter dict per call) can + defer their finalization until the whole adapter has arrived. + """ + try: + self._init_socket() + self._init_buffer() + + # receive bucket and update weights + while True: + metadata = self.socket.recv_pyobj() + weights, tensor = [], None + for name, meta in metadata["bucket_meta"].items(): + shape, dtype, offset, handle = meta["shape"], meta["dtype"], meta["offset"], meta["handle"] + if handle is not None: + tensor = rebuild_ipc(handle, self.device.index) + weights.append((name, tensor)) + continue + size = dtype.itemsize * shape.numel() + tensor = self.buffer[offset : offset + size].view(dtype=dtype).view(shape) + if self.use_shm: + tensor = tensor.to(self.device) + weights.append((name, tensor)) + is_last = metadata["is_last"] + on_bucket_received(weights, is_last) + get_torch_device().synchronize() + del weights, tensor + if not is_last: + self.socket.send(b"") + else: + self._ack_pending = True + break + finally: + self._cleanup() + + def _init_socket(self): + """Initialize ZMQ REP socket and connect.""" + self.socket = self.zmq_context.socket(zmq.REP) + self.socket.connect(self.zmq_handle) + + def _init_buffer(self): + """Receive and rebuild communication buffer from sender.""" + comm_metadata = self.socket.recv_pyobj() + buffer, shm = None, None + if not self.use_shm: + handle = comm_metadata + buffer = rebuild_ipc(handle, self.device.index) + assert buffer.dtype == torch.uint8 + else: + shm_name = comm_metadata["name"] + shm_size = comm_metadata["size"] + buffer, shm = rebuild_shared_memory(shm_name, shm_size, dtype=torch.uint8) + self.socket.send(b"") + self.buffer = buffer + self.shm = shm + + def _cleanup(self): + """clean up""" + # Synchronize before releasing the buffer to ensure all async ops + # referencing it (e.g. clone, .to()) have completed. + get_torch_device().synchronize() + del self.buffer + self.buffer = None + if self.shm is not None: + self.shm.close() + del self.shm + self.shm = None + if is_support_ipc(): + get_torch_device().ipc_collect() + get_torch_device().empty_cache() + # Ack last, after the buffer is released: the sender reclaims its bucket + # buffer the moment this ack returns, and CUDA IPC only frees the sender's + # block once our ref counter reaches zero. Acking earlier strands it in + # CudaIPCSentDataLimbo until some later round drains the limbo. + # The guard keeps the REP socket's strict recv/send alternation valid. + if self.socket is not None: + if self._ack_pending: + self.socket.send(b"") + self.socket.close() + self.socket = None + + +async def ensure_async_iterator(iterable): + """Convert an iterable to an async iterator.""" + if hasattr(iterable, '__aiter__'): + async for item in iterable: + yield item + else: + for item in iterable: + yield item diff --git a/src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py b/src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py new file mode 100644 index 000000000..69d2eaec2 --- /dev/null +++ b/src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py @@ -0,0 +1,441 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""vLLM delta weight transfer based on verl.""" + +from __future__ import annotations + +import inspect +import json +from collections.abc import Iterable, Iterator +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch +import vllm + +try: + from vllm.config.weight_transfer import WeightTransferConfig + from vllm.distributed.weight_transfer.base import ( + WeightTransferEngine, + WeightTransferInitInfo, + WeightTransferUpdateInfo, + ) +except ImportError as exc: + raise RuntimeError( + "vLLM delta_sharded support requires the vllm.config.weight_transfer " + "and vllm.distributed.weight_transfer modules" + ) from exc + +if TYPE_CHECKING: + from vllm.model_executor.model_loader.checkpoint_weight_patch import ( + CheckpointWeightPatch, + ) + +TWINKLE_DELTA_WEIGHT_TRANSFER_BACKEND = "twinkle_delta_ipc" +TWINKLE_DELTA_HCCL_WEIGHT_TRANSFER_BACKEND = "twinkle_delta_hccl" +_REGISTERED = False +_HCCL_REGISTERED = False + + +def _checkpoint_patch_api(): + try: + from vllm.model_executor.model_loader.checkpoint_weight_patch import ( + CheckpointWeightPatch, + load_checkpoint_weight_patches, + ) + except ImportError as exc: + if exc.name != "vllm.model_executor.model_loader.checkpoint_weight_patch": + raise + raise RuntimeError( + "sync_backend='delta_sharded' with vLLM requires the checkpoint patch API from " + f"vLLM #50723; installed vLLM {getattr(vllm, '__version__', 'unknown')} does not provide it. " + "Install a compatible vLLM build or select a different checkpoint engine backend." + ) from exc + return CheckpointWeightPatch, load_checkpoint_weight_patches + + +def require_vllm_delta_support() -> None: + """Verify that vLLM provides the required weight transfer and checkpoint patch APIs.""" + + init_parameters = inspect.signature(WeightTransferEngine.__init__).parameters + if not {"config", "vllm_config", "device", "model"}.issubset(init_parameters): + raise RuntimeError( + "sync_backend='delta_sharded' with vLLM requires the engine-owned WeightTransferEngine " + "API from vLLM #44353 and the checkpoint patch API from vLLM #50723; installed vLLM " + f"{getattr(vllm, '__version__', 'unknown')} does not provide the required WeightTransferEngine API. " + "Install a compatible vLLM build or select a different checkpoint engine backend." + ) + _checkpoint_patch_api() + + +def is_moe_model(hf_config) -> bool: + """Whether vLLM will treat this HF config as a MoE model. + + Delta updates require the Triton MoE backend, so this check must agree with + the engine's own MoE detection: expert counts may live under a different + attribute (Dbrx ``moe_num_experts``, DeepSeek ``n_routed_experts``, Mixtral + ``num_local_experts``) or on a nested text config (multimodal MoE). The + fallback mirrors vLLM's attribute list for configs its convertor cannot + process. + """ + try: + from vllm.transformers_utils.config import get_hf_text_config + from vllm.transformers_utils.model_arch_config_convertor import ( + MODEL_ARCH_CONFIG_CONVERTORS, + ModelArchConfigConvertorBase, + ) + + convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get( + getattr(hf_config, "model_type", ""), ModelArchConfigConvertorBase + ) + return int(convertor_cls(hf_config, get_hf_text_config(hf_config)).get_num_experts()) > 0 + except Exception: + text_config = hf_config.get_text_config() if hasattr(hf_config, "get_text_config") else hf_config + for name in ("num_experts", "moe_num_experts", "n_routed_experts", "num_local_experts"): + num_experts = getattr(text_config, name, 0) + if isinstance(num_experts, list): + num_experts = num_experts[0] if num_experts else 0 + if num_experts and int(num_experts) > 0: + return True + return False + + +@dataclass +class TwinkleDeltaIPCInitInfo(WeightTransferInitInfo): + """No initialization data is needed for the local IPC transport.""" + + +@dataclass +class TwinkleDeltaIPCUpdateInfo(WeightTransferUpdateInfo): + """Connection details for one flush.""" + + zmq_handle: str | None = None + + def __post_init__(self) -> None: + if not isinstance(self.zmq_handle, str) or not self.zmq_handle: + raise ValueError("Twinkle delta IPC updates require a worker-local zmq_handle") + + +def decode_delta_payload( + named_tensors: Iterable[tuple[str, torch.Tensor]], +) -> tuple[str, list[CheckpointWeightPatch]]: + """Validate and decode one DeltaFlush into vLLM checkpoint patches.""" + + from twinkle.checkpoint_engine.delta.protocol.decode import decode_delta_wire + + CheckpointWeightPatch, _ = _checkpoint_patch_api() + encoding, patches = decode_delta_wire(named_tensors) + return encoding, [ + CheckpointWeightPatch( + name=p.name, + shape=p.shape, + dtype=p.dtype, + values=p.values, + indices=p.positions, + ) + for p in patches + ] + + +class TwinkleDeltaIPCWeightTransferEngine(WeightTransferEngine): + """Consume Twinkle DeltaFlush payloads over the colocated CUDA-IPC channel.""" + + init_info_cls = TwinkleDeltaIPCInitInfo + update_info_cls = TwinkleDeltaIPCUpdateInfo + supports_draft_weight_update = False + required_device_type = "cuda" + + def __init__( + self, + config: WeightTransferConfig, + vllm_config, + device: torch.device, + model: torch.nn.Module, + ) -> None: + super().__init__(config, vllm_config, device, model) + self._session_encoding: str | None = None + self._update_failed = False + + def _device_error_message(self) -> str: + return "Twinkle delta weight transfer requires CUDA" + + def _validate_configuration(self) -> None: + require_vllm_delta_support() + if self.device.type != self.required_device_type: + raise NotImplementedError(self._device_error_message()) + if self.parallel_config.data_parallel_size != 1: + raise NotImplementedError("Twinkle delta weight transfer requires data_parallel_size=1") + if self.parallel_config.pipeline_parallel_size != 1: + raise NotImplementedError("Twinkle delta weight transfer requires pipeline_parallel_size=1") + if self.model_config.dtype != torch.bfloat16: + raise NotImplementedError("Twinkle delta weight transfer supports only BF16 vLLM models") + if getattr(self.vllm_config, "quant_config", None) is not None: + raise NotImplementedError("Twinkle delta weight transfer does not support quantized vLLM models") + if getattr(self.vllm_config, "speculative_config", None) is not None: + raise NotImplementedError("Twinkle delta weight transfer does not support speculative decoding") + + def init_transfer_engine(self, init_info: TwinkleDeltaIPCInitInfo) -> None: + self._validate_configuration() + + def start_weight_update(self) -> None: + if self._update_failed: + raise RuntimeError( + "A previous delta update did not complete; restart the full job so every rollout worker receives a " + "fresh dense seed. Restarting only the vLLM workers is unsafe because the producer snapshot may " + "already have advanced." + ) + self._session_encoding = None + + def _receive_payload(self, *, zmq_handle: str) -> list[tuple[str, torch.Tensor]]: + from twinkle.sampler.vllm_sampler.delta.cuda_ipc import ( + BucketedWeightReceiver, + ) + + receiver = BucketedWeightReceiver( + zmq_handle=zmq_handle, + device=self.device, + use_shm=False, + ) + payload: list[tuple[str, torch.Tensor]] = [] + callback_error: BaseException | None = None + + def retain_bucket(weights: list[tuple[str, torch.Tensor]], _is_last: bool) -> None: + nonlocal callback_error + if callback_error is not None: + return + try: + # The receiver reuses its IPC buffer for the next bucket, so + # retain a private copy until this flush is decoded and applied. + payload.extend((name, tensor.clone()) for name, tensor in weights) + except BaseException as exc: + # Finish the receiver handshake before reporting a callback error. + callback_error = exc + + receiver.receive_weights(on_bucket_received=retain_bucket) + if callback_error is not None: + raise callback_error + return payload + + def _decode_and_apply(self, payload: list[tuple[str, torch.Tensor]]) -> None: + """Decode one flush's named tensors into checkpoint patches and apply them.""" + from twinkle.sampler.vllm_sampler.delta.weight_update_utils import ( + drop_tied_alias_patches, + prepare_sparse_weight_loaders, + ) + + encoding, patches = decode_delta_payload(payload) + patches = drop_tied_alias_patches(self.model, patches) + + if self._session_encoding is None: + self._session_encoding = encoding + if encoding == "dense": + from vllm.model_executor.model_loader.reload import ( + initialize_layerwise_reload, + ) + + initialize_layerwise_reload(self.model) + elif encoding != self._session_encoding: + raise ValueError( + "One weight update cannot mix dense and sparse DeltaFlush payloads " + f"({self._session_encoding!r} then {encoding!r})" + ) + + if encoding == "indices": + prepare_sparse_weight_loaders(self.model) + _, load_checkpoint_weight_patches = _checkpoint_patch_api() + # The trusted producer emits every changed position at most once; + # skip vLLM's sort-based duplicate check on the rollout device. + load_checkpoint_weight_patches( + self.model, + patches, + validate_unique_indices=False, + ) + + def receive_weights(self, update_info: TwinkleDeltaIPCUpdateInfo) -> None: + assert update_info.zmq_handle is not None + try: + payload = self._receive_payload(zmq_handle=update_info.zmq_handle) + self._decode_and_apply(payload) + except BaseException: + # An earlier flush may already have changed runtime weights, so a + # later update would start from an unknown partial state. + self._update_failed = True + raise + + def finish_weight_update(self) -> None: + if self._session_encoding == "dense": + # Dense updates use vLLM's layerwise reload lifecycle. Finalize it + # once after the last flush to run model-specific post-load work. + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + ) + + try: + finalize_layerwise_reload(self.model, self.model_config) + except BaseException: + self._update_failed = True + raise + elif self._session_encoding == "indices": + # Sparse patches only write source parameters; re-derive runtime copies that + # process_weights_after_loading would have rebuilt (e.g. Ascend packed conv1d). + from twinkle.sampler.vllm_sampler.delta.weight_update_utils import ( + refresh_derived_weights, + ) + + try: + refresh_derived_weights(self.model) + except BaseException: + self._update_failed = True + raise + self._session_encoding = None + + def shutdown(self) -> None: + self._session_encoding = None + + @staticmethod + def trainer_send_weights( + iterator: Iterator[Any], + trainer_args: dict[str, Any] | Any, + ) -> None: + raise NotImplementedError("Twinkle sends DeltaFlush payloads; this adapter only receives them") + + +def register_twinkle_delta_weight_transfer_engine() -> None: + """Register the adapter before vLLM creates its worker-side engine.""" + + global _REGISTERED + if _REGISTERED: + return + from vllm.distributed.weight_transfer import WeightTransferEngineFactory + + WeightTransferEngineFactory.register_engine( + TWINKLE_DELTA_WEIGHT_TRANSFER_BACKEND, + TwinkleDeltaIPCWeightTransferEngine, + ) + _REGISTERED = True + + +@dataclass +class TwinkleDeltaHCCLInitInfo(WeightTransferInitInfo): + """Rendezvous endpoints of the trainer's HCCL delta broadcast group.""" + + dist_ip: str + dist_port: int + zmq_ip: str + zmq_port: int + topic: str = "bucket_metadata" + world_size: int = 1 + # Trainer rank 0's HCCL socket env (HCCL_IF_BASE_PORT etc.); applied before joining the group. + hccl_env: dict[str, str] = field(default_factory=dict) + + +@dataclass +class TwinkleDeltaHCCLUpdateInfo(WeightTransferUpdateInfo): + """This worker's rank in the trainer's HCCL delta broadcast group.""" + + group_rank: int + + +class TwinkleDeltaHCCLWeightTransferEngine(TwinkleDeltaIPCWeightTransferEngine): + """Consume Twinkle DeltaFlush payloads straight from the trainer's HCCL broadcasts (Ascend NPU). + + The vLLM worker joins the trainer's broadcast group itself: one manifest subscription plus + one or two broadcasts per flush. The SUB is connected before the communicator is built, and + both are created once and reused across sessions. + """ + + init_info_cls = TwinkleDeltaHCCLInitInfo + update_info_cls = TwinkleDeltaHCCLUpdateInfo + required_device_type = "npu" + + def __init__(self, config, vllm_config, device, model) -> None: + super().__init__(config, vllm_config, device, model) + self._rendezvous: TwinkleDeltaHCCLInitInfo | None = None + self._manifest_sub = None + self._comm = None + self._group_rank: int | None = None + + def _device_error_message(self) -> str: + return "Twinkle delta HCCL weight transfer requires Ascend NPU ('npu' device)" + + def init_transfer_engine(self, init_info: TwinkleDeltaHCCLInitInfo) -> None: + self._validate_configuration() + # A re-seed always builds a new trainer group: drop any channel from an earlier rendezvous + # (even an identical one) so _ensure_channel rebuilds against the new endpoint. + self.shutdown() + self._rendezvous = init_info + + def _ensure_channel(self, group_rank: int): + """Create (once) and return the (manifest SUB, communicator) pair, SUB first.""" + from twinkle.checkpoint_engine.delta.rollout.hccl_receiver import build_delta_hccl_communicator + from twinkle.checkpoint_engine.zmq_pubsub import connect_subscriber + + if self._comm is None: + if self._rendezvous is None: + raise RuntimeError("init_transfer_engine must run before receive_weights") + rv = self._rendezvous + self._manifest_sub = connect_subscriber(rv.zmq_ip, rv.zmq_port, rv.topic) + try: + self._comm = build_delta_hccl_communicator( + rv.dist_ip, rv.dist_port, group_rank, rv.world_size, self.device, hccl_env=rv.hccl_env + ) + except BaseException: + self._close_sub() + raise + self._group_rank = group_rank + elif group_rank != self._group_rank: + raise RuntimeError( + f"HCCL delta group_rank changed between sessions ({self._group_rank} -> {group_rank})" + ) + return self._manifest_sub, self._comm + + def receive_weights(self, update_info: TwinkleDeltaHCCLUpdateInfo) -> None: + from twinkle.checkpoint_engine.delta.rollout.hccl_receiver import pump_hccl_delta_session + + try: + sub, comm = self._ensure_channel(update_info.group_rank) + pump_hccl_delta_session(sub, comm, self.device, on_flush=self._decode_and_apply) + except BaseException: + self._update_failed = True + raise + + def _close_sub(self) -> None: + sub, self._manifest_sub = self._manifest_sub, None + if sub is None: + return + context = getattr(sub, "context", None) + try: + sub.close(linger=0) + except Exception: + pass + if context is not None: + try: + context.term() + except Exception: + pass + + def shutdown(self) -> None: + comm, self._comm = self._comm, None + if comm is not None: + try: + comm.destroyComm(comm.comm) + except Exception: + pass + self._close_sub() + self._group_rank = None + self._rendezvous = None + super().shutdown() + + +def register_twinkle_delta_hccl_weight_transfer_engine() -> None: + """Register the HCCL adapter before vLLM creates its worker-side engine.""" + + global _HCCL_REGISTERED + if _HCCL_REGISTERED: + return + from vllm.distributed.weight_transfer import WeightTransferEngineFactory + + WeightTransferEngineFactory.register_engine( + TWINKLE_DELTA_HCCL_WEIGHT_TRANSFER_BACKEND, + TwinkleDeltaHCCLWeightTransferEngine, + ) + _HCCL_REGISTERED = True diff --git a/src/twinkle/sampler/vllm_sampler/delta/weight_update_utils.py b/src/twinkle/sampler/vllm_sampler/delta/weight_update_utils.py new file mode 100644 index 000000000..56138c2c9 --- /dev/null +++ b/src/twinkle/sampler/vllm_sampler/delta/weight_update_utils.py @@ -0,0 +1,134 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright contributors to the vLLM project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Checkpoint update filtering from VERL and sparse loader adaptations for vLLM.""" +from __future__ import annotations + +import inspect +from functools import wraps +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + from vllm.model_executor.model_loader.checkpoint_weight_patch import CheckpointWeightPatch + + +def tied_embedding_aliases(model: torch.nn.Module) -> set[str]: + """Match the canonical embedding traversal used by vLLM's AutoWeightsLoader.""" + from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding + + canonical: dict[int, str] = {} + aliases: set[str] = set() + for prefix, submodule in model.named_modules(remove_duplicate=False): + if not isinstance(submodule, VocabParallelEmbedding): + continue + for name, parameter in submodule.named_parameters(remove_duplicate=False): + qualified_name = f'{prefix}.{name}' if prefix else name + if canonical.setdefault(id(parameter), qualified_name) != qualified_name: + aliases.add(qualified_name) + return aliases + + +def drop_tied_alias_patches( + model: torch.nn.Module, patches: list[CheckpointWeightPatch], +) -> list[CheckpointWeightPatch]: + """Drop redundant tied embedding names before each checkpoint patch batch. + + Canonical patches update the shared storage even when the alias arrives in + another flush. Map checkpoint names only for filtering; vLLM still receives + the original names and owns checkpoint-to-runtime loading. + """ + aliases = tied_embedding_aliases(model) + if not aliases: + return patches + mapper = getattr(model, 'hf_to_vllm_mapper', None) + retained = [] + for patch in patches: + mapped_names = mapper.apply_list([patch.name]) if mapper is not None else [patch.name] + if not mapped_names or any(name not in aliases for name in mapped_names): + retained.append(patch) + return retained + + +def prepare_sparse_weight_loaders(model: torch.nn.Module) -> None: + """Adapt vLLM's Mamba slice writes to the sparse patch copy_ contract. + + The native loader assigns slices directly, which bypasses vLLM's NaN mask + and overwrites unchanged positions. Keep its TP geometry, but write through + copy_ so the checkpoint patch API can preserve those positions. + """ + for _, parameter in model.named_parameters(): + loader = getattr(parameter, 'weight_loader', None) + if getattr(loader, '_twinkle_sparse_copy', False): + continue + if (getattr(loader, '__module__', '') == 'vllm.model_executor.layers.mamba.mamba_mixer2' + and getattr(loader, '__qualname__', '') == 'mamba_v2_sharded_weight_loader..loader'): + parameter.weight_loader = _mamba_copy_weight_loader(loader) + + +def _mamba_copy_weight_loader(original_loader): + geometry = inspect.getclosurevars(original_loader).nonlocals + if not {'shard_spec', 'tp_size', 'tp_rank'}.issubset(geometry): + raise RuntimeError('Unsupported vLLM Mamba loader closure; cannot safely apply sparse checkpoint patches') + shard_spec = geometry['shard_spec'] + tp_size = geometry['tp_size'] + tp_rank = geometry['tp_rank'] + + @wraps(original_loader) + def copy_weight_loader(param, loaded_weight): + boundary, loaded_boundary = 0, 0 + for full_dim, extra, duplicate_groups in shard_spec: + shard_size = full_dim // tp_size + rank = 0 if duplicate_groups else tp_rank + loaded_skip = rank * shard_size + loaded_start = loaded_boundary + loaded_skip + take = min(shard_size, full_dim - extra - loaded_skip) + destination = param.data[boundary:boundary + take, ...] + source = loaded_weight[loaded_start:loaded_start + take] + destination.copy_(source) + boundary += shard_size + loaded_boundary += full_dim - extra + + copy_weight_loader._twinkle_sparse_copy = True + return copy_weight_loader + + +# vllm-ascend (ops/gdn.py, ops/kimi_kda.py) keeps a kernel-layout copy of each linear-attention conv1d +# weight: ``[D, 1, W]`` source -> ``[W, D]`` contiguous parameter, repacked only by +# ``process_weights_after_loading``. Sparse patches never run that hook, so the copy must be refreshed here. +_ASCEND_PACKED_CONV_WEIGHT = 'ascend_conv1d_weight' + + +def refresh_derived_weights(model: torch.nn.Module) -> int: + """Re-derive runtime copies of source weights after a sparse checkpoint patch. + + Dense updates rebuild these in ``finalize_layerwise_reload``; sparse updates only write the source + parameters. Copies are refreshed in place so captured graphs keep their addresses. Returns the number + of refreshed tensors. + """ + refreshed = 0 + for name, module in model.named_modules(): + packed = (getattr(module, '_parameters', None) or {}).get(_ASCEND_PACKED_CONV_WEIGHT) + if packed is None: + continue + source = module.weight + if source.dim() != 3 or source.size(1) != 1 or tuple(packed.shape) != (source.size(2), source.size(0)): + raise RuntimeError( + f'{name}.{_ASCEND_PACKED_CONV_WEIGHT}: unsupported packed conv layout ' + f'{tuple(packed.shape)} for source {tuple(source.shape)}; expected [W, D] from [D, 1, W]') + # ``.data`` keeps the copy out of autograd; copy_ also casts to the packed dtype like vllm-ascend. + packed.data.copy_(source.data.view(source.size(0), source.size(2)).transpose(0, 1)) + refreshed += 1 + return refreshed diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index a81057b33..906d18c32 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -212,6 +212,37 @@ class VLLMEngine(BaseSamplerEngine): 2. Colocated with training (Ray): Single-tenant, weight sync via load_weights """ + def _check_delta_availability(self): + if getattr(self, '_delta_unavailable', False): + raise RuntimeError('rollout is unavailable until delta synchronization completes a full valid round') + + async def _delta_admitted_outputs(self, generator): + self._check_delta_availability() + # No await between the gate check and increment on the engine's event loop. + self._delta_admitting = getattr(self, '_delta_admitting', 0) + 1 + admitted = False + try: + async for output in generator: + if not admitted: + admitted = True + self._delta_admitting -= 1 + yield output + finally: + if not admitted: + self._delta_admitting -= 1 + await generator.aclose() + + async def pause_delta_generation(self, timeout: float): + self._delta_unavailable = True + + async def pause(): + while getattr(self, '_delta_admitting', 0): + await asyncio.sleep(0.01) + await self.engine.pause_generation(mode='abort') + + # A barrier timeout keeps the gate closed and prevents worker writes. + await asyncio.wait_for(pause(), timeout=timeout) + def __init__( self, model_id: str, @@ -342,10 +373,27 @@ def _create_engine(self): 'twinkle.sampler.vllm_sampler.vllm_worker_extension.TwinkleWorkerExtension') engine_config.update(self.engine_kwargs) + weight_transfer = engine_config.get('weight_transfer_config') + backend = (weight_transfer.get('backend') if isinstance(weight_transfer, dict) + else getattr(weight_transfer, 'backend', None)) + if backend == 'twinkle_delta_ipc': + engine_config.setdefault('enable_sleep_mode', True) + self.enable_sleep_mode = engine_config['enable_sleep_mode'] + from .delta.transfer_engine import register_twinkle_delta_weight_transfer_engine, require_vllm_delta_support + require_vllm_delta_support() + register_twinkle_delta_weight_transfer_engine() + elif backend == 'twinkle_delta_hccl': + # Ascend NPU: sleep mode stays opt-in (engine_args / enable_sleep_mode), it is not forced on here. + from .delta.transfer_engine import (register_twinkle_delta_hccl_weight_transfer_engine, + require_vllm_delta_support) + require_vllm_delta_support() + register_twinkle_delta_hccl_weight_transfer_engine() if self.enable_sampling_replay: engine_config['enable_return_sampling_mask'] = True engine_config['logprobs_mode'] = 'processed_logprobs' valid_args = inspect.signature(AsyncEngineArgs).parameters.keys() + if backend in ('twinkle_delta_ipc', 'twinkle_delta_hccl') and 'weight_transfer_config' not in valid_args: + raise RuntimeError('delta_sharded requires vLLM AsyncEngineArgs.weight_transfer_config') filtered_engine_config, invalid_args = _filter_engine_config( engine_config, valid_args, @@ -463,6 +511,7 @@ async def sample(self, # check to avoid per-request RPC overhead. lora_request = self._synced_lora_request + self._check_delta_availability() generator = self.engine.generate( prompt=prompt, sampling_params=vllm_params, @@ -473,7 +522,7 @@ async def sample(self, # Get final result result = None - async for output in generator: + async for output in self._delta_admitted_outputs(generator): result = output if result is None: @@ -606,6 +655,7 @@ async def encode(self, 'this pooling request') lora_request = None + self._check_delta_availability() generator = self.engine.encode( prompt=prompt_obj, pooling_params=vllm_pooling, @@ -615,7 +665,7 @@ async def encode(self, ) result = None - async for output in generator: + async for output in self._delta_admitted_outputs(generator): result = output if result is None: raise RuntimeError('Pooling encode did not produce a result') @@ -662,6 +712,7 @@ async def generate_stream(self, elif lora_request is None and self._synced_lora_request is not None: lora_request = self._synced_lora_request + self._check_delta_availability() generator = self.engine.generate( prompt=prompt, sampling_params=vllm_params, @@ -671,14 +722,15 @@ async def generate_stream(self, ) tokenizer = await self.get_tokenizer() - async for output in generator: - for comp_output in output.outputs: - delta_token_ids = list(comp_output.token_ids) - delta_text = tokenizer.decode(delta_token_ids, skip_special_tokens=True) if delta_token_ids else '' - finish_reason: Optional[str] = None - if comp_output.finish_reason is not None: - finish_reason = _map_finish_reason(comp_output.finish_reason) - yield delta_text, finish_reason + async with contextlib.aclosing(self._delta_admitted_outputs(generator)) as outputs: + async for output in outputs: + for comp_output in output.outputs: + delta_token_ids = list(comp_output.token_ids) + delta_text = tokenizer.decode(delta_token_ids, skip_special_tokens=True) if delta_token_ids else '' + finish_reason: Optional[str] = None + if comp_output.finish_reason is not None: + finish_reason = _map_finish_reason(comp_output.finish_reason) + yield delta_text, finish_reason # ----------------------------------------------------------------- # RL-training synced LoRA helpers diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index ab17f730f..48c4ad8ba 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -10,6 +10,7 @@ from twinkle import DeviceMesh, get_logger, remote_class, remote_function, requires from twinkle.checkpoint_engine import CheckpointEngineMixin +from twinkle.checkpoint_engine.delta.rollout.mixin import DeltaSamplerMixin from twinkle.data_format import (InputFeature, PoolingParams, PoolingResponse, SampledSequence, SampleResponse, SamplingParams, Trajectory) from twinkle.exceptions import FullWeightLoadError @@ -81,7 +82,7 @@ def _load_vllm_image(entry: Any) -> Any: @remote_class(max_concurrency=SAMPLER_MAX_CONCURRENCY) -class vLLMSampler(Sampler, CheckpointEngineMixin, PartialRolloutMixin): +class vLLMSampler(Sampler, CheckpointEngineMixin, PartialRolloutMixin, DeltaSamplerMixin): """A vLLM-based sampler using VLLMEngine (AsyncLLM). This sampler automatically configures vLLM based on available GPUs. diff --git a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py index 0a4ec3e56..b5343f1ea 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py @@ -18,13 +18,16 @@ import re import signal import torch -from typing import Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple from twinkle import get_logger from twinkle.utils import Platform from twinkle.utils.framework import Torch from twinkle.utils.zmq_utils import configure_zmq_socket, get_timeout_s_from_env +if TYPE_CHECKING: + from twinkle.checkpoint_engine.delta.config import DeltaWorkerCapabilities, DeltaWorkerMetadata + logger = get_logger() @@ -38,6 +41,17 @@ def set_death_signal(): os.kill(os.getpid(), signal.SIGKILL) +def _cuda_visible_device(device: torch.device) -> str: + device_index = device.index if device.index is not None else torch.cuda.current_device() + visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES') + if visible_devices is None: + return str(device_index) + selectors = [value.strip() for value in visible_devices.split(',')] + if not 0 <= device_index < len(selectors) or not selectors[device_index] or selectors[device_index] == '-1': + raise ValueError(f'Cannot map CUDA device {device_index} to CUDA_VISIBLE_DEVICES={visible_devices!r}') + return selectors[device_index] + + # Constants for the RL training LoRA adapter identity. VLLM_LORA_INT_ID = 111 VLLM_LORA_NAME = 'twinkle_lora' @@ -103,9 +117,78 @@ def __new__(cls, *args, **kwargs): from twinkle.patch.vllm_lora_weights import VLLMLoraWeights set_death_signal() VLLMLoraWeights()(None) + config = getattr(kwargs.get('vllm_config'), 'weight_transfer_config', None) + if getattr(config, 'backend', None) == 'twinkle_delta_ipc': + from .delta.transfer_engine import register_twinkle_delta_weight_transfer_engine + register_twinkle_delta_weight_transfer_engine() + elif getattr(config, 'backend', None) == 'twinkle_delta_hccl': + from .delta.transfer_engine import register_twinkle_delta_hccl_weight_transfer_engine + register_twinkle_delta_hccl_weight_transfer_engine() return super().__new__(cls) + def validate_delta_worker(self) -> 'DeltaWorkerCapabilities': + from .delta.transfer_engine import TwinkleDeltaIPCWeightTransferEngine, is_moe_model + engine = getattr(self, 'weight_transfer_engine', None) + if not isinstance(engine, TwinkleDeltaIPCWeightTransferEngine): + raise RuntimeError("Delta sync requires engine_args['weight_transfer_config'] when constructing the " + "sampler: {'backend': 'twinkle_delta_ipc'} on CUDA or " + "{'backend': 'twinkle_delta_hccl'} on Ascend NPU") + engine._validate_configuration() + if is_moe_model(engine.model_config.hf_config): + raise ValueError('Twinkle FSDP delta export does not yet implement backend-specific MoE HF conversion') + return {'tp': self.model_runner.parallel_config.tensor_parallel_size, 'device': engine.device.type} + + def delta_ipc_endpoint(self): + if not hasattr(self, '_delta_ipc_endpoint'): + import uuid + self._delta_ipc_endpoint = f'ipc:///tmp/twinkle-delta-{uuid.uuid4().hex}.sock' + return self._delta_ipc_endpoint + + def delta_receiver_metadata(self) -> 'DeltaWorkerMetadata': + from ray.util import get_node_ip_address + return { + 'node_ip': get_node_ip_address(), + 'device_uuid': str(torch.cuda.get_device_properties(self.device).uuid), + 'cuda_visible_device': _cuda_visible_device(self.device), + 'zmq_handle': self.delta_ipc_endpoint(), + } + + def _delta_update_info_for_this_worker(self, update_info: List[Dict]) -> Dict: + """Pick this worker's entry from a per-worker ``update_info`` list. + + Same indexing as vLLM's GPU ``Worker.update_weights``: ``data_parallel_index * world_size + rank`` + (a list indexed by global worker rank across data-parallel replicas). vllm-ascend's + ``NPUWorker.update_weights`` only accepts a dict, so the entry is selected here instead of relying + on the worker class to do it. + """ + parallel_config = getattr(getattr(self, 'vllm_config', None), 'parallel_config', None) + dp_index = int(getattr(parallel_config, 'data_parallel_index', 0) or 0) + world_size = int(getattr(parallel_config, 'world_size', 1) or 1) + index = dp_index * world_size + int(getattr(self, 'rank', 0) or 0) + if index >= len(update_info): + raise ValueError(f'delta update_info has {len(update_info)} entries but this worker needs entry {index}') + return update_info[index] + + def update_twinkle_delta_weights(self, update_info: Optional[List[Dict]] = None): + """Run one delta weight-update session. + + ``None``: CUDA IPC path, this worker's local zmq endpoint is used. A list of per-worker dicts + (HCCL path, e.g. ``[{'group_rank': 1}, {'group_rank': 2}]``): this worker's own entry is passed to + ``update_weights`` as a plain dict, so it works on both the GPU and the vllm-ascend worker. + """ + if update_info is None: + self.update_weights({'zmq_handle': self.delta_ipc_endpoint()}) + else: + self.update_weights(self._delta_update_info_for_this_worker(update_info)) + + def shutdown_delta_weight_transfer(self) -> None: + """Release the delta engine's transport channel (HCCL SUB + communicator); no-op for other engines.""" + from .delta.transfer_engine import TwinkleDeltaIPCWeightTransferEngine + engine = getattr(self, 'weight_transfer_engine', None) + if isinstance(engine, TwinkleDeltaIPCWeightTransferEngine): + engine.shutdown() + def monkey_patch_model(self): from twinkle.patch.vllm_moe_loader import VLLMMoEWeights VLLMMoEWeights()(self.model_runner.model) @@ -326,16 +409,19 @@ def _broadcast_obj(obj): weights.append((name, tensor)) del partial_tensors[name] - Torch.synchronize() - - if is_driver: - socket.send(b'') + # Wait for this worker's own clones out of the shared buffer. Pass the device + # explicitly: TP workers inherit the actor's LOCAL_RANK, so the default would + # point at the wrong card. + Torch.synchronize(self.device) - # Ensure all ranks finish reading the buffer before the driver - # proceeds to the next bucket (which overwrites the buffer). + # Every rank must finish reading before the driver acks: the ack lets the + # sender overwrite the buffer with the next bucket. if tp_size > 1: dist.barrier(group=cpu_group) + if is_driver: + socket.send(b'') + if lora_mode: lora_bucket_accum.extend(weights) else: diff --git a/src/twinkle/utils/framework.py b/src/twinkle/utils/framework.py index 5bd1c0a3d..eb97acdbf 100644 --- a/src/twinkle/utils/framework.py +++ b/src/twinkle/utils/framework.py @@ -204,13 +204,20 @@ def to_local_tensor(tensor: 'torch.Tensor') -> 'torch.Tensor': return tensor @staticmethod - def synchronize(): + def synchronize(device=None): + """Wait for all queued work on ``device``, defaulting to the LOCAL_RANK device. + + Pass ``device`` explicitly in subprocesses (e.g. vLLM TP workers) that inherit the + parent's LOCAL_RANK but are bound to a different device. + """ import torch + if device is None: + device = Platform.get_local_device() if Torch.is_gpu_available(): - torch.cuda.synchronize(Platform.get_local_device()) + torch.cuda.synchronize(device) elif Torch.is_npu_available(): import torch_npu - torch.npu.synchronize(Platform.get_local_device()) + torch.npu.synchronize(device) @staticmethod def contains_nan(*args, **kwargs) -> bool: diff --git a/src/twinkle/utils/platforms/__init__.py b/src/twinkle/utils/platforms/__init__.py index ea327f840..c0b533c3e 100644 --- a/src/twinkle/utils/platforms/__init__.py +++ b/src/twinkle/utils/platforms/__init__.py @@ -1,4 +1,4 @@ from .base import Platform from .gpu import GPU from .mps import MPS, is_mps_available -from .npu import NPU, ensure_hccl_socket_env, ensure_npu_backend +from .npu import NPU, ensure_hccl_socket_env, ensure_npu_backend, hccl_socket_env diff --git a/src/twinkle/utils/platforms/npu.py b/src/twinkle/utils/platforms/npu.py index 32e42538e..68552754b 100644 --- a/src/twinkle/utils/platforms/npu.py +++ b/src/twinkle/utils/platforms/npu.py @@ -48,6 +48,17 @@ def ensure_hccl_socket_env(master_port: int, environ: Optional[dict] = None) -> env.setdefault(key, value) +def hccl_socket_env(master_port: int, environ: Optional[dict] = None) -> dict: + """Return the HCCL socket env this process will use for a group rendezvoused on ``master_port``. + + Read-only counterpart of :func:`ensure_hccl_socket_env`: explicit or previously set values win + over the defaults derived from ``master_port``. Processes joining one HCCL group must agree on + these values, so a master can publish this dict and peers apply it before joining. + """ + env = os.environ if environ is None else environ + return {key: env.get(key, value) for key, value in _derive_hccl_socket_env_defaults(master_port).items()} + + def _resolve_ascend_physical_device_id(device_id: int) -> int: """Map local NPU device index to physical device id via visible devices.""" visible = os.environ.get('ASCEND_RT_VISIBLE_DEVICES', '').strip() diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 479798f63..076a8cc6a 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -381,13 +381,16 @@ def stateless_init_process_group( master_listen_fd=listen_fd, ) - pg = StatelessProcessGroup( - rank=rank, - world_size=world_size, - store=store, - socket=listen_socket, - data_expiration_seconds=3600, - ) + # Newer vLLM (e.g. 0.30) removed the ``socket`` field from StatelessProcessGroup; pass it only if present. + import dataclasses + pg_fields = {field.name for field in dataclasses.fields(StatelessProcessGroup)} + pg_kwargs = dict(rank=rank, world_size=world_size, store=store, data_expiration_seconds=3600) + if 'socket' in pg_fields: + pg_kwargs['socket'] = listen_socket + pg = StatelessProcessGroup(**pg_kwargs) + if 'socket' not in pg_fields and listen_socket is not None: + # TCPStore borrows this socket's fd via master_listen_fd; keep a reference so GC does not close it. + pg._twinkle_listen_socket = listen_socket communicator = Communicator(pg, device=device) return communicator diff --git a/tests/checkpoint_engine/_hccl_loopback.py b/tests/checkpoint_engine/_hccl_loopback.py new file mode 100644 index 000000000..6952f67b8 --- /dev/null +++ b/tests/checkpoint_engine/_hccl_loopback.py @@ -0,0 +1,371 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Multi-process loopback harness for the HCCL delta-sync protocol (CPU: gloo + zmq, no NPU). + +Everything the children need lives here (spawned processes must import it). Real code under test: +trainer ``DeltaHCCLCheckpointEngine`` (lazy communicator, zmq PUB, ``DeltaSyncMixin``) and receiver +``connect_subscriber`` + ``build_delta_hccl_communicator`` + ``pump_hccl_delta_session`` + +``decode_delta_wire``. Only ``PyHccl`` (-> gloo broadcast), ``torch.npu`` and the training engine +are stand-ins. +""" +from __future__ import annotations + +import asyncio +import multiprocessing as mp +import os +import queue as queue_mod +import random +import time +import traceback +from datetime import timedelta +from types import SimpleNamespace + +import torch + +ENGINE_KW = dict(batch_gather=4) +GROUP_TIMEOUT_S = 40 +BF16 = torch.bfloat16 + + +# --------------------------------------------------------------------------- model / rounds +def make_state(seed: int = 0, shapes=None) -> dict[str, torch.Tensor]: + shapes = shapes or {'layer.w1': (16, 32), 'layer.w2': (4000, ), 'layer.bias': (64, )} + g = torch.Generator().manual_seed(seed) + return {n: torch.randn(*s, generator=g).to(BF16) for n, s in shapes.items()} + + +def build_rounds(init_state, change_counts: list[dict[str, int]], seed: int = 1): + """Steady rounds from ``{name: n_changed}`` dicts. + + Returns ``(rounds, expected)``: ``rounds[i] = {name: (idx int64 sorted, val bf16)}`` and + ``expected[0]`` = seed state, ``expected[i+1]`` = state after steady round i (independently + computed here, so it does not depend on trainer/receiver code). + """ + g = torch.Generator().manual_seed(seed) + cur = {n: t.clone() for n, t in init_state.items()} + expected = [{n: t.clone() for n, t in cur.items()}] + rounds = [] + for counts in change_counts: + rnd = {} + for name, n in counts.items(): + numel = cur[name].numel() + idx = torch.randperm(numel, generator=g)[:n].sort().values + val = (cur[name].view(-1)[idx].float() + 1.0 + torch.rand(n, generator=g)).to(BF16) + cur[name].view(-1).index_copy_(0, idx, val) + rnd[name] = (idx, val) + rounds.append(rnd) + expected.append({n: t.clone() for n, t in cur.items()}) + return rounds, expected + + +def state_to_bytes(state) -> dict: + return {n: (tuple(t.shape), str(t.dtype), t.contiguous().view(torch.uint8).numpy().tobytes()) + for n, t in state.items()} + + +def wire_bytes(metas) -> int: + """Payload bytes on the broadcast wire for a list of (non-terminal) manifests.""" + return sum(m.get('pos_numel', 0) + m['val_numel'] * 2 for m in metas if not m.get('terminal_empty')) + + +# --------------------------------------------------------------------------- stand-ins +class LoopComm: + """PyHccl stand-in: broadcast over the gloo default group of this process.""" + + def __init__(self): + self.comm = 'handle' + + def broadcast(self, tensor, src): + import torch.distributed as dist + assert tensor.dtype == torch.uint8 + dist.broadcast(tensor, src=src) + + def destroyComm(self, handle): + import torch.distributed as dist + if dist.is_initialized(): + dist.destroy_process_group() + + +class CorruptingComm(LoopComm): + """Rank 0 only: flips one bit of the ``corrupt_call``-th broadcast buffer (after the checksum).""" + + def __init__(self, corrupt_call): + super().__init__() + self.corrupt_call = corrupt_call + self.calls = 0 + + def broadcast(self, tensor, src): + if self.calls == self.corrupt_call: + tensor[0] ^= 1 + self.calls += 1 + super().broadcast(tensor, src) + + +def _install_stand_ins(comm_factory, build_log): + """Patch torch.npu and stateless_init_process_group in BOTH modules (child process only).""" + import torch.distributed as dist + + from twinkle.checkpoint_engine.delta.rollout import hccl_receiver + from twinkle.checkpoint_engine.delta.transport import hccl_engine + + os.environ.setdefault('GLOO_SOCKET_IFNAME', 'lo') + torch.set_num_threads(1) + torch.npu = SimpleNamespace(current_device=lambda: 0, synchronize=lambda: None, empty_cache=lambda: None) + + def fake_init(master_address, master_port, rank, world_size, device=None, backend='nccl', **kw): + assert backend == 'hccl' + # blocks until every rank joined: the real HCCL rendezvous semantics + dist.init_process_group('gloo', init_method=f'tcp://{master_address}:{master_port}', rank=rank, + world_size=world_size, timeout=timedelta(seconds=GROUP_TIMEOUT_S)) + build_log.append(rank) + return comm_factory(rank) + + hccl_engine.stateless_init_process_group = fake_init + hccl_receiver.stateless_init_process_group = fake_init + hccl_engine.find_node_ip = lambda: '127.0.0.1' + + +class FakeTrainingEngine: + """get_per_tensor_param / prime_delta_snapshots / get_per_tensor_param_delta_shard stand-in.""" + + def __init__(self, init_state, rounds): + self.state = {n: t.clone() for n, t in init_state.items()} + self.rounds = rounds + self.next_round = 0 + self.primed = 0 + + def get_per_tensor_param(self): + return ((n, t.clone()) for n, t in self.state.items()), None + + def prime_delta_snapshots(self): + self.primed += 1 + + def get_per_tensor_param_delta_shard(self): + changes = self.rounds[self.next_round] + self.next_round += 1 + + def gen(): + for name, (idx, val) in changes.items(): + shape = tuple(self.state[name].shape) + yield [(name, shape)], 'bfloat16', torch.tensor([idx.numel()]), idx.clone(), val.clone(), None + + return gen(), None + + +# --------------------------------------------------------------------------- trainer child +def trainer_main(cfg, ep_qs, res_q, done_evs): + try: + build_log: list = [] + + def factory(rank): + if cfg.get('corrupt_call') is not None: + return CorruptingComm(cfg['corrupt_call']) + return LoopComm() + + _install_stand_ins(factory, build_log) + from twinkle.checkpoint_engine.delta.transport.hccl_engine import DeltaHCCLCheckpointEngine + from twinkle.checkpoint_engine.delta.transport.pipeline import DeltaSyncMixin + + class LoopbackEngine(DeltaHCCLCheckpointEngine): + + @property + def _wire_device(self): + return torch.device('cpu') + + if cfg.get('max_entry_elems'): + DeltaSyncMixin.MAX_ENTRY_ELEMS = cfg['max_entry_elems'] + world_size = 1 + cfg['n_recv'] + for s, (init_state, rounds) in enumerate(cfg['sessions']): + eng = LoopbackEngine(bucket_size=cfg['bucket_size'], **ENGINE_KW) + eng.is_master = True + md = eng.prepare() + assert eng.prepare() == md # idempotent + builds_before = len(build_log) + for _ in range(cfg['n_recv']): + ep_qs[s].put(dict(zmq_ip=md.zmq_ip, zmq_port=md.zmq_port, dist_ip=md.dist_ip, + dist_port=md.dist_port, world_size=world_size)) + eng.init_process_group(0, world_size, md) + assert eng._comm is None # lazy: nothing built before the first publish + train = FakeTrainingEngine(init_state, rounds) + for r in range(1 + len(rounds)): + metrics = asyncio.run(eng.send_weights(train, global_steps=r)) + res_q.put(('trainer_round', s, r, metrics, len(build_log) - builds_before)) + res_q.put(('trainer_sent', s, len(build_log) - builds_before, train.primed)) + if not done_evs[s].wait(60): + raise RuntimeError('receivers never finished') + eng.close() + closed = dict(comm=eng._comm, socket=eng.socket, seeded=eng._shard_seeded, rank=eng.rank) + res_q.put(('trainer_closed', s, closed)) + res_q.put(('trainer_exit', )) + except BaseException: + res_q.put(('error', 'trainer', traceback.format_exc())) + + +# --------------------------------------------------------------------------- receiver child +class RecordingSub: + + def __init__(self, sub): + self.sub = sub + self.metas: list = [] + + def recv_string(self): + return self.sub.recv_string() + + def recv_pyobj(self): + meta = self.sub.recv_pyobj() + self.metas.append(meta) + return meta + + +def apply_flush(state, named): + from twinkle.checkpoint_engine.delta.protocol.decode import decode_delta_wire + _enc, patches = decode_delta_wire(named) # includes the checksum check + for p in patches: + if p.positions is None: + n = 1 + for d in p.shape: + n *= d + assert p.values.numel() == n, (p.name, p.values.numel(), n) + state[p.name] = p.values.clone().reshape(p.shape) # dense: replace the whole param + else: + flat = state[p.name].view(-1) + flat.index_copy_(0, p.positions.long(), p.values) + + +def receiver_main(cfg, ep_q, res_q): + s, rank = cfg['session'], cfg['group_rank'] + try: + build_log: list = [] + _install_stand_ins(lambda r: LoopComm(), build_log) + import zmq + + from twinkle.checkpoint_engine.delta.rollout.hccl_receiver import (build_delta_hccl_communicator, + pump_hccl_delta_session) + from twinkle.checkpoint_engine.zmq_pubsub import connect_subscriber + time.sleep(cfg.get('start_delay', 0)) + ep = ep_q.get(timeout=60) + device = torch.device('cpu') + sub = None + if not cfg.get('sub_after_join'): + sub = connect_subscriber(ep['zmq_ip'], ep['zmq_port'], 'bucket_metadata') # SUB first ... + comm = build_delta_hccl_communicator(ep['dist_ip'], ep['dist_port'], rank, ep['world_size'], device) + if cfg.get('sub_after_join'): + time.sleep(cfg['sub_after_join']) # ... the wrong order, made deterministic + sub = connect_subscriber(ep['zmq_ip'], ep['zmq_port'], 'bucket_metadata') + sub.setsockopt(zmq.RCVTIMEO, cfg.get('pump_timeout_ms', 20000)) + rsub = RecordingSub(sub) + state: dict = {} + for r in range(cfg['n_rounds']): + del rsub.metas[:] + err = None + try: + enc = pump_hccl_delta_session(rsub, comm, device, lambda named: apply_flush(state, named)) + except Exception as e: # reported, then this receiver stops + enc, err = None, (type(e).__name__, str(e)) + res_q.put(('recv_round', s, rank, r, dict(state=state_to_bytes(state), metas=list(rsub.metas), + encoding=enc, error=err))) + if err: + break + res_q.put(('recv_finished', s, rank)) + try: + comm.destroyComm(comm.comm) + except Exception: + pass + except BaseException: + res_q.put(('error', f'receiver s{s} r{rank}', traceback.format_exc())) + + +# --------------------------------------------------------------------------- driver (parent) +class Result: + + def __init__(self): + self.recv: dict = {} # (session, rank) -> {round: payload} + self.trainer_rounds: dict = {} # (session, round) -> (metrics, comm builds so far) + self.trainer_sent: dict = {} # session -> (builds, primed) + self.trainer_closed: dict = {} + self.finished: set = set() + self.errors: list = [] + self.trainer_exited = False + self.elapsed = 0.0 + + +def run_loopback(sessions, n_recv=2, bucket_size=1 << 20, max_entry_elems=None, corrupt_call=None, + sub_after_join=0.0, pump_timeout_ms=20000, shuffle=False, timeout=60.0, allow_errors=False, + require_trainer_exit=True, seed=0) -> Result: + """Run ``sessions`` = [(init_state, rounds), ...] across 1 trainer + n_recv receiver processes. + + Session i>0 reuses the trainer process with a NEW engine (fresh ports/group) after ``close()``, + and gets NEW receiver processes. Always kills leftover children. + """ + ctx = mp.get_context('spawn') + res_q = ctx.Queue() + ep_qs = [ctx.Queue() for _ in sessions] + done_evs = [ctx.Event() for _ in sessions] + tcfg = dict(sessions=sessions, n_recv=n_recv, bucket_size=bucket_size, max_entry_elems=max_entry_elems, + corrupt_call=corrupt_call) + procs = [ctx.Process(target=trainer_main, args=(tcfg, ep_qs, res_q, done_evs), daemon=True)] + rng = random.Random(seed) + for s, (_st, rounds) in enumerate(sessions): + ranks = list(range(1, n_recv + 1)) + if shuffle: + rng.shuffle(ranks) + for rk in ranks: + rcfg = dict(session=s, group_rank=rk, n_rounds=1 + len(rounds), pump_timeout_ms=pump_timeout_ms, + sub_after_join=sub_after_join, start_delay=rng.random() * 0.8 if shuffle else 0.0) + procs.append(ctx.Process(target=receiver_main, args=(rcfg, ep_qs[s], res_q), daemon=True)) + res = Result() + t0 = time.time() + deadline = t0 + timeout + try: + for p in procs: + p.start() + while True: + all_recv_done = all((s, rk) in res.finished for s in range(len(sessions)) for rk in range(1, n_recv + 1)) + if all_recv_done and (res.trainer_exited or not require_trainer_exit): + break + try: + msg = res_q.get(timeout=max(0.05, min(1.0, deadline - time.time()))) + except queue_mod.Empty: + if time.time() >= deadline: + if all_recv_done and not require_trainer_exit: + break + raise AssertionError(f'loopback timed out after {timeout}s; got so far: finished={res.finished} ' + f'trainer_rounds={sorted(res.trainer_rounds)} errors={res.errors}') + continue + kind = msg[0] + if kind == 'error': + res.errors.append(msg[1:]) + if not allow_errors: + raise AssertionError(f'child failure in {msg[1]}:\n{msg[2]}') + elif kind == 'recv_round': + _, s, rk, r, payload = msg + res.recv.setdefault((s, rk), {})[r] = payload + elif kind == 'recv_finished': + res.finished.add((msg[1], msg[2])) + if all((msg[1], k) in res.finished for k in range(1, n_recv + 1)): + done_evs[msg[1]].set() + elif kind == 'trainer_round': + res.trainer_rounds[(msg[1], msg[2])] = (msg[3], msg[4]) + elif kind == 'trainer_sent': + res.trainer_sent[msg[1]] = (msg[2], msg[3]) + elif kind == 'trainer_closed': + res.trainer_closed[msg[1]] = msg[2] + elif kind == 'trainer_exit': + res.trainer_exited = True + finally: + res.elapsed = time.time() - t0 + for ev in done_evs: + ev.set() + for p in procs: + p.join(timeout=3 if res.trainer_exited or p is not procs[0] else 0.1) + for p in procs: + if p.is_alive(): + p.terminate() + for p in procs: + p.join(timeout=3) + if p.is_alive(): + p.kill() + p.join(timeout=3) + for q in [res_q, *ep_qs]: + q.close() + q.cancel_join_thread() + return res diff --git a/tests/checkpoint_engine/test_delta_checksum.py b/tests/checkpoint_engine/test_delta_checksum.py new file mode 100644 index 000000000..8f8d42477 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_checksum.py @@ -0,0 +1,171 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Device-agnostic XOR-tree delta wire checksum (CPU only).""" +import importlib.util +import json +import sys +from dataclasses import dataclass +from functools import reduce +from operator import xor +from pathlib import Path +from types import ModuleType + +import pytest +import torch + +from twinkle.checkpoint_engine.delta.protocol import wire +from twinkle.checkpoint_engine.delta.protocol.wire import checksum + +MASK = (1 << 64) - 1 + + +def ref_xor(t): + return reduce(xor, t.contiguous().view(torch.uint8).flatten().tolist(), 0) + + +def rand_bytes(n, seed=0): + g = torch.Generator().manual_seed(seed) + return torch.randint(0, 256, (n,), dtype=torch.uint8, generator=g) + + +def test_empty_is_zero(): + empty = torch.empty(0, dtype=torch.uint8) + assert wire._xor_tree(empty) == 0 + assert checksum(empty, empty) == 0 + assert checksum(empty, torch.empty(0, dtype=torch.bfloat16)) == 0 + + +def test_deterministic_and_dtype_view_invariant(): + values = torch.randn(33).to(torch.bfloat16) + pos = rand_bytes(12) + a = checksum(pos, values) + assert a == checksum(pos, values) + assert a == checksum(pos, values.view(torch.uint8)) + assert a == checksum(pos, values.view(torch.int16)) + assert a == checksum(pos.clone(), values.clone()) + + +@pytest.mark.parametrize('n', [1, 2, 3, 7, 8, 9, 1023]) +def test_matches_reference_for_lengths(n): + pos = rand_bytes(n, seed=n) + vals = rand_bytes(n, seed=n + 100) + assert wire._xor_tree(pos) == ref_xor(pos) + assert checksum(pos, vals) == (ref_xor(pos) ^ (ref_xor(vals) << 1)) & MASK + odd = rand_bytes(n * 2 + 2, seed=n).view(torch.bfloat16) # non-8-aligned byte size for most n + assert wire._xor_tree(odd) == ref_xor(odd) + + +def test_single_bit_flip_changes_value(): + pos = rand_bytes(13, seed=1) + vals = rand_bytes(40, seed=2) + base = checksum(pos, vals) + for i in range(vals.numel()): + for bit in range(8): + flipped = vals.clone() + flipped[i] ^= 1 << bit + assert checksum(pos, flipped) != base + for i in range(pos.numel()): + for bit in range(8): + flipped = pos.clone() + flipped[i] ^= 1 << bit + assert checksum(flipped, vals) != base + + +def test_input_not_mutated(): + for t in (rand_bytes(9), rand_bytes(16), torch.randn(5).to(torch.bfloat16)): + before = t.clone() + checksum(t, t) + assert torch.equal(t.view(torch.uint8), before.view(torch.uint8)) + + +def test_non_contiguous_input(): + base = rand_bytes(64).view(8, 8) + view = base.t() + assert wire._xor_tree(view) == ref_xor(view.contiguous()) + + +# ---------- end-to-end decode --------------------------------------------- + + +@pytest.fixture +def adapter(monkeypatch): + @dataclass + class CheckpointWeightPatch: + name: str + shape: tuple + dtype: torch.dtype + values: torch.Tensor + indices: torch.Tensor | None = None + + class Base: + def __init__(self, *args, **kwargs): + pass + + modules = { + 'vllm': {'__version__': 'test'}, + 'vllm.config.weight_transfer': {'WeightTransferConfig': type('WeightTransferConfig', (), {})}, + 'vllm.distributed.weight_transfer.base': { + 'WeightTransferEngine': Base, + 'WeightTransferInitInfo': type('WeightTransferInitInfo', (), {}), + 'WeightTransferUpdateInfo': type('WeightTransferUpdateInfo', (), {}), + }, + 'vllm.model_executor.model_loader.checkpoint_weight_patch': { + 'CheckpointWeightPatch': CheckpointWeightPatch, + 'load_checkpoint_weight_patches': lambda *a, **k: None, + }, + 'vllm.model_executor.layers.vocab_parallel_embedding': { + 'VocabParallelEmbedding': type('VocabParallelEmbedding', (), {}), + }, + 'vllm.model_executor.model_loader.reload': { + 'initialize_layerwise_reload': lambda *a, **k: None, + 'finalize_layerwise_reload': lambda *a, **k: None, + }, + } + for name, fields in modules.items(): + module = ModuleType(name) + module.__dict__.update(fields) + monkeypatch.setitem(sys.modules, name, module) + path = Path(__file__).parents[2] / 'src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py' + spec = importlib.util.spec_from_file_location('_twinkle_delta_checksum_adapter_test', path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + try: + spec.loader.exec_module(module) + module.decode_delta_payload # noqa: B018 + except Exception as exc: # pragma: no cover - environment dependent + pytest.skip(f'vLLM adapter unavailable: {exc!r}') + return module + + +def build_flush(encoding, corrupt=False): + values = torch.tensor([-0., 2.], dtype=torch.bfloat16) + positions = torch.tensor([0, 3], dtype=torch.int32).view(torch.uint8) + if encoding == 'dense': + positions = torch.empty(0, dtype=torch.uint8) + params = [dict(name='w', dtype='bfloat16', shape=[4 if encoding == 'indices' else 2], + pos_start=0, pos_end=positions.numel(), pos_width=4, val_start=0, val_end=2)] + spec = dict(encoding=encoding, params=params, checksum=checksum(positions, values)) + spec_tensor = torch.tensor(list(json.dumps(spec).encode()), dtype=torch.uint8) + if corrupt: + values[0] = 42 + return [('__delta_spec__', spec_tensor), ('__positions__', positions), ('__values__', values)] + + +@pytest.mark.parametrize('encoding', ['indices', 'dense']) +def test_decode_accepts_valid_payload(adapter, encoding): + kind, patches = adapter.decode_delta_payload(build_flush(encoding)) + assert kind == encoding and patches[0].name == 'w' + + +@pytest.mark.parametrize('encoding', ['indices', 'dense']) +def test_decode_rejects_corrupt_payload(adapter, encoding): + with pytest.raises(RuntimeError, match='checksum'): + adapter.decode_delta_payload(build_flush(encoding, corrupt=True)) + + +def test_decode_rejects_tampered_checksum_field(adapter): + flush = build_flush('indices') + spec = json.loads(bytes(flush[0][1].numpy().tobytes()).decode()) + spec['checksum'] ^= 1 + flush[0] = ('__delta_spec__', torch.tensor(list(json.dumps(spec).encode()), dtype=torch.uint8)) + with pytest.raises(RuntimeError, match='checksum'): + adapter.decode_delta_payload(flush) diff --git a/tests/checkpoint_engine/test_delta_hccl.py b/tests/checkpoint_engine/test_delta_hccl.py new file mode 100644 index 000000000..b339159eb --- /dev/null +++ b/tests/checkpoint_engine/test_delta_hccl.py @@ -0,0 +1,227 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeltaHCCLCheckpointEngine with fake torch.npu, communicator and sockets (CPU only).""" +from __future__ import annotations + +import asyncio +import os +import subprocess +import sys +from types import SimpleNamespace + +import pytest +import torch + +from twinkle.checkpoint_engine.delta.transport import hccl_engine as mod +from twinkle.checkpoint_engine.delta.transport.hccl_engine import DeltaHCCLCheckpointEngine +from twinkle.checkpoint_engine.hccl_checkpoint_engine import MasterMetadata + +from .test_delta_pipeline import FakeTrainingEngine + +META = MasterMetadata(zmq_ip='1.2.3.4', zmq_port=111, dist_ip='1.2.3.4', dist_port=222) + + +class Log(list): + pass + + +class FakeSocket: + + def __init__(self, log): + self.log = log + self.metas = [] + self.closed = False + self.context = SimpleNamespace(term=lambda: log.append('ctx_term')) + + def send_string(self, s, flags=0): + pass + + def send_pyobj(self, meta): + self.log.append('manifest') + self.metas.append(meta) + + def close(self, linger=None): + self.closed = True + + +class FakeComm: + + def __init__(self, log): + self.log = log + self.comm = 'handle' + self.destroyed = [] + + def broadcast(self, tensor, src): + assert src == 0 and tensor.dtype == torch.uint8 + self.log.append(('bcast', tensor.numel())) + + def destroyComm(self, handle): + self.destroyed.append(handle) + + +@pytest.fixture +def env(monkeypatch): + log = Log() + calls = [] + comms = [] + counters = {'empty_cache': 0, 'sync': 0} + + def fake_init(**kwargs): + calls.append(kwargs) + log.append('build_comm') + comm = FakeComm(log) + comms.append(comm) + return comm + + def sync(): + counters['sync'] += 1 + log.append('sync') + + def empty_cache(): + counters['empty_cache'] += 1 + + monkeypatch.setattr(torch, 'npu', + SimpleNamespace(current_device=lambda: 0, synchronize=sync, empty_cache=empty_cache), + raising=False) + monkeypatch.setattr(mod, 'stateless_init_process_group', fake_init) + monkeypatch.setattr(mod, 'find_node_ip', lambda: '1.2.3.4') + ports = iter([111, 222]) + monkeypatch.setattr(mod, 'find_free_port', lambda: next(ports)) + monkeypatch.setattr(mod, 'start_publisher', lambda ip, port: FakeSocket(log)) + # the pipeline broadcasts from `_wire_device`; keep tensors on CPU for the fake + monkeypatch.setattr(DeltaHCCLCheckpointEngine, '_wire_device', property(lambda self: 'cpu')) + return SimpleNamespace(log=log, calls=calls, comms=comms, counters=counters) + + +def make(master=True, rank=0, **kw): + e = DeltaHCCLCheckpointEngine(bucket_size=1 << 20, batch_gather=4, **kw) + e.is_master = master + return e + + +def run(e, eng): + return asyncio.run(e.send_weights(eng)) + + +def test_import_without_npu_modules(): + code = ("import sys, twinkle.checkpoint_engine.delta.transport.hccl_engine as m;" + "print('torch_npu' in sys.modules, 'vllm_ascend' in sys.modules)") + out = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, check=True).stdout + assert out.split() == ['False', 'False'] + + +def test_rebuild_group_rejected(): + with pytest.raises(NotImplementedError, match='rebuild_group=True is not supported'): + DeltaHCCLCheckpointEngine(rebuild_group=True) + e = DeltaHCCLCheckpointEngine() + assert e.rebuild_group is False and e.wire_format == 'delta_flush' and e.topic == 'bucket_metadata' + + +def test_wire_device_is_current_npu(monkeypatch): + monkeypatch.setattr(torch, 'npu', SimpleNamespace(current_device=lambda: 3), raising=False) + # `torch.device('npu', n)` only exists when the backend is registered; check the call shape + captured = [] + monkeypatch.setattr(mod.torch, 'device', lambda *a: captured.append(a) or a) + assert DeltaHCCLCheckpointEngine().__class__._wire_device.fget(None) == ('npu', 3) + assert DeltaHCCLCheckpointEngine.__mro__[1].__name__ == 'DeltaSyncMixin' + + +def test_prepare_master_idempotent(env): + e = make() + first = e.prepare() + sock = e.socket + assert first == META + assert e.prepare() == META and e.socket is sock + + +def test_prepare_non_master(env): + e = make(master=False, rank=-1) + assert e.prepare() is None and e.socket is None + assert e.hccl_env is None + + +def test_prepare_publishes_effective_hccl_env(env, monkeypatch): + from twinkle.utils.platforms.npu import _derive_hccl_socket_env_defaults + for key in _derive_hccl_socket_env_defaults(0): + monkeypatch.delenv(key, raising=False) + # Already pinned by the FSDP process group (from MASTER_PORT): must be published as-is. + monkeypatch.setenv('HCCL_IF_BASE_PORT', '12345') + e = make() + e.prepare() + derived = _derive_hccl_socket_env_defaults(222) + assert e.hccl_env == dict(derived, HCCL_IF_BASE_PORT='12345') + assert 'HCCL_HOST_SOCKET_PORT_RANGE' not in os.environ # read-only: prepare() must not mutate env + e.close() + assert e.hccl_env is None + + +def test_init_process_group_builds_no_communicator(env): + e = make() + e.prepare() + e.init_process_group(0, 3, META) + assert (e.rank, e._group_world_size, e._rendezvous) == (0, 3, META) + other = make(master=False) + other.init_process_group(-1, 3, META) + assert other.rank == -1 and other._rendezvous is None + assert env.calls == [] and e._comm is None + + +def test_end_to_end_rounds(env): + e = make() + e.prepare() + e.init_process_group(0, 3, META) + eng = FakeTrainingEngine([('a', (8,), [1, 3]), ('b', (4, 2), [0, 7, 5])]) + + run(e, eng) # seed: values-only flush, no positions + assert env.calls == [dict(master_address='1.2.3.4', master_port=222, rank=0, world_size=3, device=0, + backend='hccl')] + # comm built before the first manifest, then manifest -> values broadcast only -> synchronize + assert [x[0] if isinstance(x, tuple) else x for x in env.log] == ['build_comm', 'manifest', 'bcast', 'sync'] + assert env.log[2][1] == (8 + 8) * 2 # a(8) + b(4x2) bf16 values + assert env.counters['empty_cache'] == 1 + assert [m['encoding'] for m in e.socket.metas] == ['dense'] + + env.log.clear() + run(e, eng) # steady: positions then values + assert len(env.calls) == 1 # communicator reused + assert env.log[0] == 'manifest' + assert [x[0] if isinstance(x, tuple) else x for x in env.log] == ['manifest', 'bcast', 'bcast', 'sync'] + assert env.log[1][1] == 5 * 4 and env.log[2][1] == 5 * 2 # positions (uint8 bytes) then values + assert env.counters['empty_cache'] == 2 + assert env.counters['sync'] == 2 + + +def test_non_master_never_builds_or_publishes(env): + e = make(master=False, rank=-1) + e.prepare() + e.init_process_group(-1, 3, META) + eng = FakeTrainingEngine([('a', (8,), [1])]) + assert run(e, eng) is None and run(e, eng) is None + assert env.calls == [] and env.log == [] and env.counters['empty_cache'] == 0 + + +def test_terminal_only_round_builds_comm_but_broadcasts_nothing(env): + e = make() + e.prepare() + e.init_process_group(0, 2, META) + eng = FakeTrainingEngine([('a', (8,), [])]) + # skip the seed publish so the comm is still unbuilt at the terminal publish + e._shard_seeded = True + run(e, eng) + assert env.log == ['build_comm', 'manifest'] + assert e.socket.metas[0]['terminal_empty'] is True + assert len(env.calls) == 1 + + +def test_close(env): + e = make() + e.close() # safe before init + e.prepare() + e.init_process_group(0, 2, META) + run(e, FakeTrainingEngine([('a', (8,), [1])])) + comm, sock = env.comms[0], e.socket + e.close() + assert comm.destroyed == ['handle'] and sock.closed and 'ctx_term' in env.log + assert e._comm is None and e.socket is None and e.rank is None + e.close() # idempotent + assert comm.destroyed == ['handle'] + assert e._shard_seeded is False diff --git a/tests/checkpoint_engine/test_delta_hccl_loopback.py b/tests/checkpoint_engine/test_delta_hccl_loopback.py new file mode 100644 index 000000000..90bfebcfc --- /dev/null +++ b/tests/checkpoint_engine/test_delta_hccl_loopback.py @@ -0,0 +1,134 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Multi-process loopback of the HCCL delta-sync protocol (gloo stands in for PyHccl, zmq is real). + +Validates message ordering / rendezvous / decoding across real processes. It does NOT validate HCCL. +Every scenario is bounded by ``run_loopback(timeout=...)`` (children are always killed afterwards). +""" +from __future__ import annotations + +from ._hccl_loopback import build_rounds, make_state, run_loopback, state_to_bytes, wire_bytes + +W1, W2, B = 'layer.w1', 'layer.w2', 'layer.bias' + + +def _assert_states(res, expected, n_recv, session=0, rounds=None): + rounds = range(len(expected)) if rounds is None else rounds + for rk in range(1, n_recv + 1): + got = res.recv[(session, rk)] + for r in rounds: + assert got[r]['error'] is None, (rk, r, got[r]['error']) + assert got[r]['state'] == state_to_bytes(expected[r]), f'receiver {rk} differs after round {r}' + + +def _flags(payload): + return [m['is_last'] for m in payload['metas']] + + +def test_seed_only_values_flushes(): + init = make_state(0) + res = run_loopback([(init, [])], bucket_size=1024) + _assert_states(res, [init], 2) + for rk in (1, 2): + metas = res.recv[(0, rk)][0]['metas'] + assert len(metas) == 3 and all(m['encoding'] == 'dense' and m['pos_numel'] == 0 for m in metas) + assert _flags(res.recv[(0, rk)][0]) == [False, False, True] + assert res.recv[(0, rk)][0]['encoding'] == 'dense' + assert res.trainer_sent[0] == (1, 1) # communicator built once, diff base primed once + + +def test_seed_then_sparse_steady_rounds(): + init = make_state(1) + rounds, expected = build_rounds(init, [{W1: 0, W2: 500, B: 3}, {W1: 5, W2: 0, B: 64}], seed=2) + res = run_loopback([(init, rounds)], bucket_size=1024) + _assert_states(res, expected, 2) + seed_bytes = None + for rk in (1, 2): + got = res.recv[(0, rk)] + seed_bytes = wire_bytes(got[0]['metas']) + for r in (1, 2): + assert got[r]['encoding'] == 'indices' + flags = _flags(got[r]) + assert flags == [False] * (len(flags) - 1) + [True] + # round 1: w2 (3000 B) overflows the 1 KiB bucket -> 2 flushes; round 2 fits in one + assert len(flags) == (2 if r == 1 else 1) + assert wire_bytes(got[r]['metas']) < seed_bytes + assert len(res.recv[(0, 1)][1]['metas']) == len(res.recv[(0, 2)][1]['metas']) + assert res.trainer_rounds[(0, 1)][0]['checkpoint_engine/flushes'] == float(len(res.recv[(0, 1)][1]['metas'])) + + +def test_zero_change_round_is_terminal_only_and_comm_persists(): + init = make_state(3) + rounds, expected = build_rounds(init, [{W1: 7, W2: 20, B: 2}, {W1: 0, W2: 0, B: 0}, {W2: 33}], seed=4) + res = run_loopback([(init, rounds)], bucket_size=1024) + _assert_states(res, expected, 2) + assert all((expected[1][n] == expected[2][n]).all() for n in expected[1]) + for rk in (1, 2): + m = res.recv[(0, rk)][2]['metas'] + assert len(m) == 1 and m[0]['terminal_empty'] and m[0]['is_last'] + assert res.recv[(0, rk)][2]['state'] == res.recv[(0, rk)][1]['state'] + # the communicator was built exactly once, at the first publish, and persisted for all 4 rounds + assert res.trainer_sent[0][0] == 1 + assert all(res.trainer_rounds[(0, r)][1] == 1 for r in range(4)) + + +def test_tiny_bucket_slices_one_param_into_pieces(): + init = make_state(5) + rounds, expected = build_rounds(init, [{W2: 500, B: 10}, {W2: 301}], seed=6) + res = run_loopback([(init, rounds)], bucket_size=1024, max_entry_elems=64) + _assert_states(res, expected, 2) + for rk in (1, 2): + metas = res.recv[(0, rk)][1]['metas'] + assert len(metas) > 1 + entries = [p for m in metas for p in m['spec']['params'] if p['name'] == W2] + assert len(entries) >= 8 # 500 changes / 64 per entry + assert all(p['val_end'] - p['val_start'] <= 64 for p in entries) + + +def test_close_then_reseed_new_session(): + init0, init1 = make_state(7), make_state(8) + r0, e0 = build_rounds(init0, [{W2: 40, B: 4}], seed=9) + r1, e1 = build_rounds(init1, [{W1: 30, W2: 60}], seed=10) + res = run_loopback([(init0, r0), (init1, r1)], bucket_size=1024, timeout=90) + _assert_states(res, e0, 2, session=0) + _assert_states(res, e1, 2, session=1) + for s in (0, 1): + assert res.trainer_sent[s] == (1, 1) # fresh engine: one new communicator, one new seed + closed = res.trainer_closed[s] + assert closed == dict(comm=None, socket=None, seeded=False, rank=None) + # receivers of the second session saw a FULL seed (dense) first, not a delta + assert res.recv[(1, 1)][0]['encoding'] == 'dense' + + +def test_corrupted_payload_is_detected_end_to_end(): + init = make_state(11) + rounds, expected = build_rounds(init, [{W2: 50, B: 3}], seed=12) + # broadcast calls: 0 = seed values; 1 = steady positions; 2 = steady values (flipped after checksum) + res = run_loopback([(init, rounds)], bucket_size=1 << 20, corrupt_call=2) + _assert_states(res, expected, 2, rounds=[0]) + for rk in (1, 2): + err = res.recv[(0, rk)][1]['error'] + assert err is not None and err[0] == 'RuntimeError' and 'checksum mismatch' in err[1] + assert res.recv[(0, rk)][1]['state'] == state_to_bytes(expected[0]) # corrupt flush NOT applied + # single-flush round: the trainer is never blocked by the failed receivers and exits cleanly + assert res.trainer_exited and (0, 1) in res.trainer_rounds + + +def test_dp_layout_four_receivers_shuffled_start(): + init = make_state(13) + rounds, expected = build_rounds(init, [{W1: 9, W2: 200, B: 5}, {W2: 77}], seed=14) + res = run_loopback([(init, rounds)], n_recv=4, bucket_size=2048, shuffle=True, seed=15, timeout=90) + _assert_states(res, expected, 4) + assert res.trainer_sent[0][0] == 1 + + +def test_sub_after_join_loses_first_manifest(): + """Negative demo of the SUB-before-join rule. The late SUB connects 1s after the group formed, long + after the trainer published, so the loss is deterministic by construction (not a timing race).""" + init = make_state(16) + res = run_loopback([(init, [])], sub_after_join=1.0, pump_timeout_ms=3000, allow_errors=True, + require_trainer_exit=False, timeout=45) + for rk in (1, 2): + p = res.recv[(0, rk)][0] + assert p['error'] is not None and p['error'][0] == 'Again' # nothing ever arrived + assert p['metas'] == [] + assert (0, 0) not in res.trainer_rounds # trainer cannot finish: receivers never joined the broadcast diff --git a/tests/checkpoint_engine/test_delta_hccl_receiver.py b/tests/checkpoint_engine/test_delta_hccl_receiver.py new file mode 100644 index 000000000..3a433f2fc --- /dev/null +++ b/tests/checkpoint_engine/test_delta_hccl_receiver.py @@ -0,0 +1,623 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Worker-side HCCL delta receiver: decode, pump, vLLM engine, worker extension (CPU, fakes).""" +from __future__ import annotations + +import json +import subprocess +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from twinkle.checkpoint_engine.delta.rollout import hccl_receiver +from twinkle.checkpoint_engine.delta.protocol.decode import DeltaPatch, decode_delta_wire +from twinkle.checkpoint_engine.delta.protocol.wire import checksum + +from .test_delta_vllm import adapter, named_flush # noqa: F401 (adapter is a fixture) + + +# --------------------------------------------------------------------------- helpers +def _spec(named): + return json.loads(bytes(dict(named)['__delta_spec__'].numpy().tobytes()).decode()) + + +def _replace_spec(named, **changes): + spec = _spec(named) + spec.update(changes) + tensor = torch.tensor(list(json.dumps(spec).encode()), dtype=torch.uint8) + return [('__delta_spec__', tensor)] + [(k, v) for k, v in named if k != '__delta_spec__'] + + +def _mutate_param(named, **changes): + spec = _spec(named) + spec['params'][0].update(changes) + return _replace_spec(named, params=spec['params']) + + +def _manifest(named, is_last=True): + """Build the trainer-side manifest + the byte payloads for one flush.""" + t = dict(named) + spec = _spec(named) + pos = t.get('__positions__') if spec['encoding'] != 'dense' else None + val = t['__values__'] + meta = dict(is_full=spec['encoding'] == 'dense', encoding=spec['encoding'], is_last=is_last, + terminal_empty=False, val_numel=val.numel(), val_dtype=str(val.dtype).split('.')[-1], spec=spec) + if spec['encoding'] != 'dense': + meta['pos_numel'] = pos.numel() + return meta, (pos.clone() if pos is not None else None), val.view(torch.uint8).clone() + + +class FakeSub: + + def __init__(self, metas): + self.items = [('bucket_metadata', m) for m in metas] + self.recv_count = 0 + self.closed = False + self.close_kwargs = None + self.context = Mock() + + def recv_string(self): + self.recv_count += 1 + return self.items[0][0] + + def recv_pyobj(self): + return self.items.pop(0)[1] + + def close(self, **kwargs): + self.closed = True + self.close_kwargs = kwargs + + +class FakeComm: + """Fills broadcast targets from queued payloads, in call order.""" + + def __init__(self, payloads=()): + self.payloads = list(payloads) + self.calls = [] + self.comm = object() + self.destroyed = [] + + def broadcast(self, tensor, src=0): + assert src == 0 + data = self.payloads.pop(0) + assert data.numel() == tensor.numel() + tensor.copy_(data) + self.calls.append(tensor.numel()) + + def destroyComm(self, handle): + self.destroyed.append(handle) + + +@pytest.fixture +def sync(monkeypatch): + mock = Mock() + monkeypatch.setattr(torch.cpu, 'synchronize', mock) + return mock + + +def _terminal(): + return dict(is_full=False, encoding='indices', is_last=True, terminal_empty=True) + + +def _run_pump(flushes, sync_mock=None, extra_metas=()): + metas, payloads, applied = [], [], [] + for i, (named, is_last) in enumerate(flushes): + meta, pos, val = _manifest(named, is_last) + metas.append(meta) + if pos is not None: + payloads.append(pos) + payloads.append(val) + sub, comm = FakeSub(metas + list(extra_metas)), FakeComm(payloads) + enc = hccl_receiver.pump_hccl_delta_session(sub, comm, torch.device('cpu'), + on_flush=lambda n: applied.append( + [(k, v.clone()) for k, v in n])) + return enc, sub, comm, applied + + +# --------------------------------------------------------------------------- (a) decode +def test_decode_indices_and_dense(): + enc, patches = decode_delta_wire(named_flush()) + assert enc == 'indices' and isinstance(patches[0], DeltaPatch) + assert (patches[0].name, patches[0].shape, patches[0].dtype) == ('w', (4, ), torch.bfloat16) + assert patches[0].positions.tolist() == [0, 3] and patches[0].positions.dtype == torch.int32 + enc, patches = decode_delta_wire(named_flush('dense')) + assert enc == 'dense' and patches[0].positions is None and patches[0].values.numel() == 2 + + +def test_decode_corrupt_checksum(): + flush = named_flush() + flush[-1][1][0] = 42 + with pytest.raises(RuntimeError, match='checksum'): + decode_delta_wire(flush) + + +def test_decode_dense_with_positions(): + flush = named_flush('dense') + pos = torch.tensor([1], dtype=torch.int32).view(torch.uint8) + flush = [flush[0], ('__positions__', pos), flush[2]] + with pytest.raises(ValueError, match='must not contain positions'): + decode_delta_wire(flush) + + +def test_decode_indices_without_positions(): + flush = [x for x in named_flush() if x[0] != '__positions__'] + with pytest.raises(ValueError, match='requires __positions__'): + decode_delta_wire(flush) + + +def test_decode_bad_pos_width(): + with pytest.raises(ValueError, match='requires pos_width=4'): + decode_delta_wire(_mutate_param(named_flush(), pos_width=2)) + + +def test_decode_inconsistent_slice_lengths(): + with pytest.raises(ValueError, match='inconsistent lengths'): + decode_delta_wire(_mutate_param(named_flush(), val_end=1)) + + +def test_decode_unsupported_encoding(): + with pytest.raises(ValueError, match='Unsupported DeltaFlush encoding'): + decode_delta_wire(_replace_spec(named_flush(), encoding='bogus')) + + +def test_decode_missing_tensors_and_bad_json(): + with pytest.raises(ValueError, match='missing __values__'): + decode_delta_wire([x for x in named_flush() if x[0] != '__values__']) + with pytest.raises(ValueError, match='missing __delta_spec__'): + decode_delta_wire([x for x in named_flush() if x[0] != '__delta_spec__']) + bad = [('__delta_spec__', torch.tensor(list(b'{nope'), dtype=torch.uint8))] + named_flush()[1:] + with pytest.raises(ValueError, match='invalid JSON'): + decode_delta_wire(bad) + + +def test_decode_payload_wrapper_returns_vllm_patches(adapter): + encoding, patches = adapter.decode_delta_payload(named_flush()) + CheckpointWeightPatch = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch' + ].CheckpointWeightPatch + assert encoding == 'indices' and isinstance(patches[0], CheckpointWeightPatch) + assert patches[0].indices.tolist() == [0, 3] + _, dense = adapter.decode_delta_payload(named_flush('dense')) + assert dense[0].indices is None + + +# --------------------------------------------------------------------------- (b) pump +def test_pump_indices_flush_positions_before_values(sync): + named = named_flush() + enc, sub, comm, applied = _run_pump([(named, True)]) + assert enc == 'indices' + assert comm.calls == [8, 4] # positions (2 x int32 = 8 B) then values (2 x bf16 = 4 B) + assert [k for k, _ in applied[0]] == ['__delta_spec__', '__positions__', '__values__'] + got = dict(applied[0]) + assert got['__values__'].dtype == torch.bfloat16 + assert torch.equal(got['__values__'].view(torch.uint8), dict(named)['__values__'].view(torch.uint8)) + assert json.loads(bytes(got['__delta_spec__'].numpy().tobytes()).decode()) == _spec(named) + # the reassembled flush is a valid, checksummed payload + assert decode_delta_wire(applied[0])[0] == 'indices' + sync.assert_called() + + +def test_pump_dense_flush_has_no_positions_broadcast(sync): + enc, _, comm, applied = _run_pump([(named_flush('dense'), True)]) + assert enc == 'dense' and comm.calls == [4] + assert [k for k, _ in applied[0]] == ['__delta_spec__', '__values__'] + assert decode_delta_wire(applied[0])[0] == 'dense' + + +def test_pump_multi_flush_stops_after_is_last(sync): + enc, sub, comm, applied = _run_pump([(named_flush(), False), (named_flush(), True)], + extra_metas=[_terminal()]) + assert enc == 'indices' and len(applied) == 2 and comm.calls == [8, 4, 8, 4] + assert sub.recv_count == 2 and len(sub.items) == 1 # nothing consumed after is_last + + +def test_pump_terminal_only_session(sync): + enc, sub, comm, applied = _run_pump([], extra_metas=[_terminal()]) + assert enc == 'indices' and applied == [] and comm.calls == [] and sub.recv_count == 1 + + +def test_pump_terminal_after_flushes(sync): + enc, _, comm, applied = _run_pump([(named_flush('dense'), False)], extra_metas=[_terminal()]) + assert enc == 'dense' and len(applied) == 1 + + +@pytest.mark.parametrize('meta', [ + 'not-a-dict', + {'is_last': True}, + {'is_last': True, 'encoding': 'indices', 'val_numel': 1, 'val_dtype': 'bfloat16', 'spec': {}}, + {'encoding': 'dense', 'val_numel': 1, 'val_dtype': 'bfloat16', 'spec': {}}, + {'is_last': True, 'encoding': 'dense', 'val_numel': 1, 'val_dtype': 'nope', 'spec': {}}, +]) +def test_pump_malformed_manifest_raises_value_error(sync, meta): + sub, comm = FakeSub([meta]), FakeComm() + with pytest.raises(ValueError, match='Malformed delta manifest'): + hccl_receiver.pump_hccl_delta_session(sub, comm, torch.device('cpu'), on_flush=Mock()) + assert comm.calls == [] + + +def test_pump_synchronizes_device_type(monkeypatch): + npu = SimpleNamespace(synchronize=Mock()) + monkeypatch.setattr(torch, 'npu', npu, raising=False) + meta, pos, val = _manifest(named_flush('dense')) + dev = SimpleNamespace(type='npu') + # allocate on cpu while the "device type" resolves to the fake torch.npu module + real_empty = torch.empty + monkeypatch.setattr(torch, 'empty', lambda *a, **k: real_empty(*a, **{**k, 'device': 'cpu'})) + monkeypatch.setattr(torch.Tensor, 'to', lambda self, *a, **k: self) + hccl_receiver.pump_hccl_delta_session(FakeSub([meta]), FakeComm([val]), dev, on_flush=lambda n: None) + npu.synchronize.assert_called() + + +def test_build_communicator_args(monkeypatch): + stub = Mock(return_value='comm') + monkeypatch.setattr(hccl_receiver, 'stateless_init_process_group', stub) + assert hccl_receiver.build_delta_hccl_communicator('1.2.3.4', 99, 2, 3, torch.device('cpu', 1)) == 'comm' + stub.assert_called_once_with(master_address='1.2.3.4', master_port=99, rank=2, world_size=3, device=1, + backend='hccl') + + +def test_build_communicator_applies_trainer_hccl_env(monkeypatch): + seen = {} + monkeypatch.setenv('HCCL_IF_BASE_PORT', '11111') # worker's own (mismatched) value + monkeypatch.setattr(hccl_receiver, 'stateless_init_process_group', + lambda **_: seen.update(port=hccl_receiver.os.environ['HCCL_IF_BASE_PORT']) or 'comm') + env = {'HCCL_IF_BASE_PORT': '22222'} + assert hccl_receiver.build_delta_hccl_communicator('1.2.3.4', 99, 1, 2, torch.device('cpu', 0), + hccl_env=env) == 'comm' + assert seen['port'] == '22222' + + +# --------------------------------------------------------------------------- (c) engine +@pytest.fixture +def hccl(adapter, monkeypatch, sync): + order = [] + sub = FakeSub([]) + comm = FakeComm() + zmq_mod = sys.modules['twinkle.checkpoint_engine.zmq_pubsub'] + + def fake_connect(ip, port, topic): + order.append(('connect', ip, port, topic)) + return sub + + def fake_build(ip, port, rank, world, device, hccl_env=None): + order.append(('build', ip, port, rank, world)) + hccl_envs.append(hccl_env) + return comm + + hccl_envs = [] + monkeypatch.setattr(zmq_mod, 'connect_subscriber', fake_connect) + monkeypatch.setattr(hccl_receiver, 'build_delta_hccl_communicator', fake_build) + return SimpleNamespace(adapter=adapter, order=order, sub=sub, comm=comm, hccl_envs=hccl_envs) + + +def hccl_engine(adapter, device='cpu'): + config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size=1, pipeline_parallel_size=1), + model_config=SimpleNamespace(dtype=torch.bfloat16)) + model = SimpleNamespace(named_modules=lambda **_: [], named_parameters=lambda: []) + engine = adapter.TwinkleDeltaHCCLWeightTransferEngine(None, config, torch.device(device), model) + engine._validate_configuration = Mock() # device type is checked separately below + return engine + + +def init_info(adapter, world_size=3): + return adapter.TwinkleDeltaHCCLInitInfo(dist_ip='10.0.0.1', dist_port=5000, zmq_ip='10.0.0.2', zmq_port=6000, + world_size=world_size) + + +def feed(hccl_ns, flushes): + """Queue manifests on the fake SUB and payloads on the fake comm.""" + for named, is_last in flushes: + meta, pos, val = _manifest(named, is_last) + hccl_ns.sub.items.append(('bucket_metadata', meta)) + if pos is not None: + hccl_ns.comm.payloads.append(pos) + hccl_ns.comm.payloads.append(val) + + +def feed_terminal(hccl_ns): + hccl_ns.sub.items.append(('bucket_metadata', _terminal())) + + +def test_ensure_channel_order_reuse_and_rank_change(hccl): + engine = hccl_engine(hccl.adapter) + with pytest.raises(RuntimeError, match='init_transfer_engine must run'): + engine._ensure_channel(1) + engine.init_transfer_engine(init_info(hccl.adapter)) + sub, comm = engine._ensure_channel(1) + assert (sub, comm) == (hccl.sub, hccl.comm) + assert hccl.order == [('connect', '10.0.0.2', 6000, 'bucket_metadata'), ('build', '10.0.0.1', 5000, 1, 3)] + assert engine._ensure_channel(1) == (sub, comm) and len(hccl.order) == 2 + with pytest.raises(RuntimeError, match='group_rank changed'): + engine._ensure_channel(2) + assert hccl.hccl_envs == [{}] + + +def test_ensure_channel_forwards_hccl_env(hccl): + engine = hccl_engine(hccl.adapter) + env = {'HCCL_IF_BASE_PORT': '30000'} + engine.init_transfer_engine(hccl.adapter.TwinkleDeltaHCCLInitInfo( + dist_ip='10.0.0.1', dist_port=5000, zmq_ip='10.0.0.2', zmq_port=6000, world_size=3, hccl_env=env)) + engine._ensure_channel(1) + assert hccl.hccl_envs == [env] + + +def test_ensure_channel_build_failure_closes_sub(hccl, monkeypatch): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + monkeypatch.setattr(hccl_receiver, 'build_delta_hccl_communicator', Mock(side_effect=OSError('rdzv'))) + with pytest.raises(OSError): + engine._ensure_channel(1) + assert hccl.sub.closed and engine._comm is None and engine._manifest_sub is None + + +def _update(adapter, rank=1): + return adapter.TwinkleDeltaHCCLUpdateInfo(group_rank=rank) + + +def test_receive_sparse_session_prepares_sparse_loaders(hccl, monkeypatch): + utils = sys.modules.get('twinkle.sampler.vllm_sampler.delta.weight_update_utils') + prepare = Mock() + drop = Mock(side_effect=lambda model, patches: patches) + refresh = Mock(return_value=0) + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta.weight_update_utils', + SimpleNamespace(prepare_sparse_weight_loaders=prepare, drop_tied_alias_patches=drop, + refresh_derived_weights=refresh)) + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + feed(hccl, [(named_flush(), False), (named_flush(), True)]) + engine.start_weight_update() + engine.receive_weights(_update(hccl.adapter)) + engine.finish_weight_update() + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + reload = sys.modules['vllm.model_executor.model_loader.reload'] + assert prepare.call_count == 2 and drop.call_count == 2 + refresh.assert_called_once_with(engine.model) # once per session, after the last flush + assert loader.load_checkpoint_weight_patches.call_count == 2 + assert loader.load_checkpoint_weight_patches.call_args.kwargs == {'validate_unique_indices': False} + reload.initialize_layerwise_reload.assert_not_called() + reload.finalize_layerwise_reload.assert_not_called() + del utils + + +def test_receive_dense_session_runs_one_reload_lifecycle(hccl): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + feed(hccl, [(named_flush('dense'), False), (named_flush('dense'), True)]) + engine.start_weight_update() + engine.receive_weights(_update(hccl.adapter)) + engine.finish_weight_update() + reload = sys.modules['vllm.model_executor.model_loader.reload'] + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + assert reload.initialize_layerwise_reload.call_count == 1 + assert reload.finalize_layerwise_reload.call_count == 1 + assert loader.load_checkpoint_weight_patches.call_count == 2 + + +def test_receive_mixed_encodings_rejected_and_poisons(hccl): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + feed(hccl, [(named_flush('dense'), False), (named_flush(), True)]) + engine.start_weight_update() + with pytest.raises(ValueError, match='cannot mix dense and sparse'): + engine.receive_weights(_update(hccl.adapter)) + assert engine._update_failed + + +def test_terminal_only_session_is_noop(hccl): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + feed_terminal(hccl) + engine.start_weight_update() + engine.receive_weights(_update(hccl.adapter)) + engine.finish_weight_update() + reload = sys.modules['vllm.model_executor.model_loader.reload'] + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + loader.load_checkpoint_weight_patches.assert_not_called() + reload.initialize_layerwise_reload.assert_not_called() + reload.finalize_layerwise_reload.assert_not_called() + assert not engine._update_failed + + +def test_pump_failure_sets_update_failed(hccl): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + hccl.sub.items.append(('bucket_metadata', {'oops': 1})) + engine.start_weight_update() + with pytest.raises(ValueError, match='Malformed delta manifest'): + engine.receive_weights(_update(hccl.adapter)) + assert engine._update_failed + with pytest.raises(RuntimeError, match='previous delta update did not complete'): + engine.start_weight_update() + + +def test_channel_failure_sets_update_failed(hccl): + engine = hccl_engine(hccl.adapter) # init_transfer_engine never ran + with pytest.raises(RuntimeError, match='init_transfer_engine must run'): + engine.receive_weights(_update(hccl.adapter)) + assert engine._update_failed + + +def test_validate_configuration_device_types(adapter): + def build(cls, device): + config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size=1, pipeline_parallel_size=1), + model_config=SimpleNamespace(dtype=torch.bfloat16)) + model = SimpleNamespace() + engine = cls(None, config, SimpleNamespace(type=device), model) + return engine + + ok = build(adapter.TwinkleDeltaHCCLWeightTransferEngine, 'npu') + ok._validate_configuration() + with pytest.raises(NotImplementedError, match='NPU'): + build(adapter.TwinkleDeltaHCCLWeightTransferEngine, 'cuda')._validate_configuration() + build(adapter.TwinkleDeltaIPCWeightTransferEngine, 'cuda')._validate_configuration() + with pytest.raises(NotImplementedError, match='^Twinkle delta weight transfer requires CUDA$'): + build(adapter.TwinkleDeltaIPCWeightTransferEngine, 'npu')._validate_configuration() + assert adapter.TwinkleDeltaIPCWeightTransferEngine.required_device_type == 'cuda' + assert adapter.TwinkleDeltaHCCLWeightTransferEngine.required_device_type == 'npu' + + +def test_shutdown_destroys_once_and_is_idempotent(hccl): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + engine._ensure_channel(1) + engine.shutdown() + engine.shutdown() + assert hccl.comm.destroyed == [hccl.comm.comm] + assert hccl.sub.closed and hccl.sub.close_kwargs == {'linger': 0} + hccl.sub.context.term.assert_called_once() + assert (engine._comm, engine._manifest_sub, engine._group_rank, engine._rendezvous) == (None, None, None, None) + + +def test_register_hccl_engine_once(adapter, monkeypatch): + registry = [] + factory = SimpleNamespace(register_engine=lambda name, cls: registry.append((name, cls))) + module = type(sys)('vllm.distributed.weight_transfer') + module.WeightTransferEngineFactory = factory + monkeypatch.setitem(sys.modules, module.__name__, module) + adapter.register_twinkle_delta_hccl_weight_transfer_engine() + adapter.register_twinkle_delta_hccl_weight_transfer_engine() + assert registry == [('twinkle_delta_hccl', adapter.TwinkleDeltaHCCLWeightTransferEngine)] + + +# --------------------------------------------------------------------------- (d) worker extension +@pytest.fixture +def worker_cls(adapter, monkeypatch): + from twinkle.sampler.vllm_sampler import vllm_worker_extension as ext + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta.transfer_engine', adapter) + monkeypatch.setattr(ext, 'set_death_signal', Mock()) + import twinkle.patch.vllm_lora_weights as lora + monkeypatch.setattr(lora, 'VLLMLoraWeights', lambda: (lambda _: None)) + return ext.TwinkleWorkerExtension + + +@pytest.mark.parametrize('backend, ipc, hccl_calls', [('twinkle_delta_hccl', 0, 1), ('twinkle_delta_ipc', 1, 0), + ('other', 0, 0)]) +def test_worker_new_registers_by_backend(worker_cls, adapter, monkeypatch, backend, ipc, hccl_calls): + reg_ipc, reg_hccl = Mock(), Mock() + monkeypatch.setattr(adapter, 'register_twinkle_delta_weight_transfer_engine', reg_ipc) + monkeypatch.setattr(adapter, 'register_twinkle_delta_hccl_weight_transfer_engine', reg_hccl) + vllm_config = SimpleNamespace(weight_transfer_config=SimpleNamespace(backend=backend)) + worker_cls(vllm_config=vllm_config) + assert (reg_ipc.call_count, reg_hccl.call_count) == (ipc, hccl_calls) + + +def test_validate_delta_worker_reports_npu_and_names_both_backends(worker_cls, adapter): + engine = hccl_engine(adapter, 'cpu') + engine.device = SimpleNamespace(type='npu') + engine.model_config.hf_config = SimpleNamespace(model_type='llama') + worker = SimpleNamespace(weight_transfer_engine=engine, + model_runner=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=2))) + assert worker_cls.validate_delta_worker(worker) == {'tp': 2, 'device': 'npu'} + with pytest.raises(RuntimeError, match='weight_transfer_config') as err: + worker_cls.validate_delta_worker(SimpleNamespace(weight_transfer_engine=object())) + assert 'twinkle_delta_ipc' in str(err.value) and 'twinkle_delta_hccl' in str(err.value) + + +def fake_delta_worker(worker_cls, rank=0, dp_index=None, world_size=None): + """A worker stand-in whose update_weights records its argument; mirrors the bound helper like a real worker.""" + parallel = None if dp_index is None and world_size is None else SimpleNamespace( + data_parallel_index=dp_index or 0, world_size=world_size or 1) + worker = SimpleNamespace(update_weights=Mock(), delta_ipc_endpoint=lambda: 'ipc:///x', rank=rank, + vllm_config=SimpleNamespace(parallel_config=parallel)) + worker._delta_update_info_for_this_worker = lambda info: worker_cls._delta_update_info_for_this_worker( + worker, info) + return worker + + +def test_update_twinkle_delta_weights_dispatch(worker_cls): + worker = fake_delta_worker(worker_cls) + worker_cls.update_twinkle_delta_weights(worker) + worker.update_weights.assert_called_once_with({'zmq_handle': 'ipc:///x'}) + worker.update_weights.reset_mock() + info = [{'group_rank': 1}, {'group_rank': 2}] + worker_cls.update_twinkle_delta_weights(worker, info) + # a list is reduced to this worker's own entry and handed over as a plain dict (vllm-ascend's + # NPUWorker.update_weights does not accept lists) + worker.update_weights.assert_called_once_with({'group_rank': 1}) + + +@pytest.mark.parametrize('rank,expected', [(0, 5), (1, 6), (2, 7)]) +def test_update_twinkle_delta_weights_selects_entry_by_worker_rank(worker_cls, rank, expected): + worker = fake_delta_worker(worker_cls, rank=rank, dp_index=0, world_size=3) + worker_cls.update_twinkle_delta_weights(worker, [{'group_rank': 5}, {'group_rank': 6}, {'group_rank': 7}]) + worker.update_weights.assert_called_once_with({'group_rank': expected}) + + +def test_update_twinkle_delta_weights_indexes_across_data_parallel_replicas(worker_cls): + # same formula as vLLM's GPU Worker.update_weights: data_parallel_index * world_size + rank + worker = fake_delta_worker(worker_cls, rank=1, dp_index=1, world_size=2) + worker_cls.update_twinkle_delta_weights(worker, [{'group_rank': i} for i in range(4)]) + worker.update_weights.assert_called_once_with({'group_rank': 3}) + + +def test_update_twinkle_delta_weights_rejects_short_list(worker_cls): + worker = fake_delta_worker(worker_cls, rank=2, dp_index=0, world_size=3) + with pytest.raises(ValueError, match='needs entry 2'): + worker_cls.update_twinkle_delta_weights(worker, [{'group_rank': 1}, {'group_rank': 2}]) + worker.update_weights.assert_not_called() + + +# --------------------------------------------------------------------------- (e) import hygiene +def test_import_without_npu_modules(): + code = ("import sys, twinkle.checkpoint_engine.delta.rollout.hccl_receiver, twinkle.checkpoint_engine.delta.protocol.decode;" + "print('torch_npu' in sys.modules, 'vllm_ascend' in sys.modules)") + out = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, check=True).stdout + assert out.split() == ['False', 'False'] + + +def test_engine_pieces_do_not_import_npu_modules(adapter): + engine = hccl_engine(adapter) + engine.init_transfer_engine(init_info(adapter)) + assert 'torch_npu' not in sys.modules and 'vllm_ascend' not in sys.modules + assert checksum # keep import used + + +# --------------------------------------------------------------------------- (f) re-initialization / shutdown RPC +def test_init_transfer_engine_after_shutdown_rebuilds_against_new_rendezvous(hccl): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + engine._ensure_channel(1) + engine.shutdown() + new = hccl.adapter.TwinkleDeltaHCCLInitInfo(dist_ip='10.0.0.9', dist_port=5001, zmq_ip='10.0.0.8', + zmq_port=6001, world_size=5) + engine.init_transfer_engine(new) + engine._ensure_channel(3) + assert hccl.order[-2:] == [('connect', '10.0.0.8', 6001, 'bucket_metadata'), ('build', '10.0.0.9', 5001, 3, 5)] + assert hccl.comm.destroyed == [hccl.comm.comm] + + +@pytest.mark.parametrize('identical', [True, False]) +def test_init_transfer_engine_with_live_channel_releases_it_first(hccl, identical): + engine = hccl_engine(hccl.adapter) + engine.init_transfer_engine(init_info(hccl.adapter)) + engine._ensure_channel(1) + new = init_info(hccl.adapter) if identical else hccl.adapter.TwinkleDeltaHCCLInitInfo( + dist_ip='10.0.0.9', dist_port=5001, zmq_ip='10.0.0.8', zmq_port=6001, world_size=4) + engine.init_transfer_engine(new) + # the old communicator and SUB were released, nothing is live until the next receive + assert hccl.comm.destroyed == [hccl.comm.comm] and hccl.sub.closed + assert (engine._comm, engine._manifest_sub, engine._group_rank) == (None, None, None) + assert engine._rendezvous is new + # a new rendezvous may also use a different group rank + engine._ensure_channel(2) + assert hccl.order[-1][3] == 2 and len(hccl.order) == 4 + + +def test_shutdown_delta_weight_transfer_calls_engine_shutdown_for_delta_engines(worker_cls, adapter): + engine = hccl_engine(adapter, 'cpu') + engine.shutdown = Mock() + worker_cls.shutdown_delta_weight_transfer(SimpleNamespace(weight_transfer_engine=engine)) + engine.shutdown.assert_called_once() + ipc = adapter.TwinkleDeltaIPCWeightTransferEngine.__new__(adapter.TwinkleDeltaIPCWeightTransferEngine) + ipc.shutdown = Mock() + worker_cls.shutdown_delta_weight_transfer(SimpleNamespace(weight_transfer_engine=ipc)) + ipc.shutdown.assert_called_once() + + +def test_shutdown_delta_weight_transfer_is_noop_for_other_engines(worker_cls): + other = Mock() + worker_cls.shutdown_delta_weight_transfer(SimpleNamespace(weight_transfer_engine=other)) + other.shutdown.assert_not_called() + worker_cls.shutdown_delta_weight_transfer(SimpleNamespace(weight_transfer_engine=None)) + worker_cls.shutdown_delta_weight_transfer(SimpleNamespace()) diff --git a/tests/checkpoint_engine/test_delta_manager.py b/tests/checkpoint_engine/test_delta_manager.py new file mode 100644 index 000000000..8e8f30d42 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_manager.py @@ -0,0 +1,214 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Manager state machine for delta synchronization, driven by mock trainer and sampler actors.""" +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from torch import nn + +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.checkpoint_engine.delta import DeltaSyncConfig + + +def manager_pair(): + model = Mock(_actors=[1, 2], device_mesh=SimpleNamespace(world_size=2)) + sampler = Mock(_actors=[3], device_mesh=SimpleNamespace(data_world_size=1)) + model.prepare_delta_sync.return_value = dict(device='cuda', master_metadata=object()) + sampler.validate_delta_sampler.return_value = dict(device='cuda', tp=2) + sampler.delta_sampler_metadata.return_value = [ + dict(node_ip='10.0.0.1', device_uuid=f'GPU-{rank}', zmq_handle=f'ipc:///rank{rank}') + for rank in range(2)] + model.send_delta_weights.return_value = lambda: {'payload_bytes': 6} + sampler.receive_delta_weights.return_value = lambda: {'changed': True} + manager = CheckpointEngineManager(model, sampler, sync_backend='delta_sharded') + return manager, model, sampler + + +@pytest.fixture +def driver_proxy_manager(): + manager = CheckpointEngineManager.__new__(CheckpointEngineManager) + manager.model = nn.Module.__new__(nn.Module) + manager.sampler = SimpleNamespace() + manager.delta_config = SimpleNamespace(timeout_seconds=30) + return manager + + +@pytest.mark.parametrize('bucket_bytes', [0, 1, 5]) +def test_bucket_must_hold_one_patch_element(bucket_bytes): + with pytest.raises(ValueError): + DeltaSyncConfig(bucket_bytes=bucket_bytes) + + +@pytest.mark.parametrize('timeout', [0, -1, float('inf'), float('nan')]) +def test_timeout_is_finite_positive(timeout): + with pytest.raises(ValueError): + DeltaSyncConfig(timeout_seconds=timeout) + + +def test_delta_timeout_is_applied_to_driver_proxy(driver_proxy_manager): + with driver_proxy_manager._delta_rpc_timeouts(): + assert driver_proxy_manager.model._ray_get_timeout == 30 + + +def test_delta_timeout_cleanup_supports_uninitialized_module_proxy(driver_proxy_manager): + with driver_proxy_manager._delta_rpc_timeouts(): + pass + assert '_ray_get_timeout' not in vars(driver_proxy_manager.model) + + +@pytest.mark.parametrize('previous_timeout', [None, 10]) +def test_delta_timeout_restores_existing_proxy_value(driver_proxy_manager, previous_timeout): + driver_proxy_manager.model._ray_get_timeout = previous_timeout + with driver_proxy_manager._delta_rpc_timeouts(): + pass + assert driver_proxy_manager.model._ray_get_timeout == previous_timeout + + +def test_delta_rejects_local_deployment(): + with pytest.raises(ValueError, match='standalone'): + CheckpointEngineManager(SimpleNamespace(), SimpleNamespace(), sync_backend='delta_sharded') + + +def test_seed_then_sparse_reuses_transport(): + manager, model, sampler = manager_pair() + manager.sync_weights() + manager.sync_weights() + assert model.init_delta_transport.call_count == 1 + assert manager._delta_round == 2 + assert model.init_delta_transport.call_args.args[0]['world_size'] == 3 + sampler.validate_delta_sampler.assert_called_once() + sampler.delta_sampler_metadata.assert_called_once() + model.prepare_delta_sync.assert_called_once() + + +@pytest.mark.parametrize('stage', ['send_delta_weights', 'receive_delta_weights', 'publish_delta_sampler']) +def test_failure_does_not_publish_version(stage): + manager, model, sampler = manager_pair() + manager.sync_weights() + getattr(model if stage == 'send_delta_weights' else sampler, stage).side_effect = TimeoutError('injected') + with pytest.raises(TimeoutError): + manager.sync_weights() + assert manager._delta_round == 1 + sampler.fail_delta_sampler.assert_called_once() + with pytest.raises(RuntimeError, match='failed'): + manager.sync_weights() + + +def test_failed_round_requires_full_job_restart_even_after_close(): + manager, model, sampler = manager_pair() + manager.sync_weights() + sampler.receive_delta_weights.side_effect = ValueError('worker write failed') + with pytest.raises(ValueError): + manager.sync_weights() + sampler.receive_delta_weights.side_effect = None + manager.close_sync_session() + with pytest.raises(RuntimeError, match='restart the entire'): + manager.sync_weights() + assert manager._delta_failed + + +def test_cleanup_failure_does_not_clear_failed_state(): + manager, model, sampler = manager_pair() + manager.sync_weights() + model.close_delta_trainer.side_effect = RuntimeError('communicator hung') + with pytest.raises(RuntimeError, match='cleanup'): + manager.close_sync_session() + assert manager._delta_failed + + +def test_zero_change_still_resumes_manager_lifecycle(): + manager, model, sampler = manager_pair() + manager.sync_weights() + sampler.publish_delta_sampler.reset_mock() + sampler.receive_delta_weights.return_value = lambda: {'changed': False} + manager.sync_weights() + sampler.publish_delta_sampler.assert_called_once() + assert sampler.begin_delta_sync.call_count == 2 + + +def test_multiple_rollout_replicas_share_one_broadcast_group(): + manager, model, sampler = manager_pair() + sampler.validate_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 + sampler.delta_sampler_metadata.return_value = [ + dict(node_ip=f'10.0.0.{rank // 2 + 1}', device_uuid=f'GPU-{rank}', zmq_handle=f'ipc:///rank{rank}') + for rank in range(4)] + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}] * 2 + manager.sync_weights() + manager.sync_weights() + endpoint = model.init_delta_transport.call_args.args[0] + assert endpoint['world_size'] == 5 + assert endpoint['rollout_workers'] == sampler.delta_sampler_metadata.return_value + model.init_delta_transport.assert_called_once() + sampler.publish_delta_sampler.assert_called() + + +def test_replica_completion_disagreement_keeps_rollout_paused(): + manager, model, sampler = manager_pair() + sampler.validate_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 + sampler.delta_sampler_metadata.return_value = [ + dict(node_ip='10.0.0.1', device_uuid=f'GPU-{rank}', zmq_handle=f'ipc:///rank{rank}') + for rank in range(4)] + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}] * 2 + manager.sync_weights() + sampler.publish_delta_sampler.reset_mock() + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}, {'changed': False}] + with pytest.raises(ValueError, match='replicas disagree'): + manager.sync_weights() + sampler.publish_delta_sampler.assert_not_called() + sampler.fail_delta_sampler.assert_called_once() + + +def test_replica_metrics_may_differ_when_completion_agrees(): + manager, model, sampler = manager_pair() + sampler.validate_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 + sampler.delta_sampler_metadata.return_value = [ + dict(node_ip='10.0.0.1', device_uuid=f'GPU-{rank}', zmq_handle=f'ipc:///rank{rank}') + for rank in range(4)] + sampler.receive_delta_weights.return_value = lambda: [ + dict(changed=True, seconds=1), dict(changed=True, seconds=2)] + manager.sync_weights() + sampler.publish_delta_sampler.assert_called_once() + + +@pytest.mark.parametrize('invalid_topology', ['missing_worker', 'duplicate_gpu']) +def test_invalid_rollout_topology_fails_before_group_init(invalid_topology): + manager, model, sampler = manager_pair() + workers = sampler.delta_sampler_metadata.return_value + sampler.delta_sampler_metadata.return_value = workers[:1] if invalid_topology == 'missing_worker' else [workers[0]] * 2 + with pytest.raises(ValueError, match='topology|distinct'): + manager.sync_weights() + model.init_delta_transport.assert_not_called() + + +@pytest.mark.parametrize('stage', ['validate_delta_sampler', 'delta_sampler_metadata', 'prepare_delta_sync']) +def test_preflight_failure_can_be_corrected_and_retried(stage): + manager, model, sampler = manager_pair() + owner = model if stage == 'prepare_delta_sync' else sampler + getattr(owner, stage).side_effect = ValueError('invalid configuration') + with pytest.raises(ValueError, match='invalid configuration'): + manager.sync_weights() + sampler.fail_delta_sampler.assert_not_called() + sampler.begin_delta_sync.assert_not_called() + getattr(owner, stage).side_effect = None + manager.sync_weights() + assert manager.last_delta_metrics['seed'] is True + + +def test_manager_begin_pauses_before_receivers_or_sender_run(): + from twinkle.checkpoint_engine import CheckpointEngineManager + events = [] + model = Mock(_actors=[1]) + sampler = Mock(_actors=[2]) + model.prepare_delta_sync.return_value = dict(device='cuda') + sampler.validate_delta_sampler.return_value = {'tp': 2} + sampler.delta_sampler_metadata.return_value = [ + dict(node_ip='node', device_uuid=str(rank), zmq_handle=str(rank)) for rank in range(2)] + sampler.begin_delta_sync.side_effect = lambda: events.append('pause-release') + model.init_delta_transport.side_effect = lambda endpoint: (events.append('init'), lambda: None)[1] + sampler.init_delta_sampler_transport.return_value = lambda: None + model.send_delta_weights.side_effect = lambda: (events.append('send'), lambda: {})[1] + sampler.receive_delta_weights.return_value = lambda: {'changed': True} + sampler.publish_delta_sampler.side_effect = lambda: events.append('restore-resume') + manager = CheckpointEngineManager(model, sampler, sync_backend='delta_sharded') + manager.sync_weights() + assert events == ['pause-release', 'init', 'send', 'restore-resume'] diff --git a/tests/checkpoint_engine/test_delta_npu_orchestration.py b/tests/checkpoint_engine/test_delta_npu_orchestration.py new file mode 100644 index 000000000..966713bfc --- /dev/null +++ b/tests/checkpoint_engine/test_delta_npu_orchestration.py @@ -0,0 +1,425 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rollout-side NPU (HCCL) orchestration: manager rendezvous and DeltaSamplerMixin RPC sequence, all mocked.""" +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from twinkle import Platform +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.checkpoint_engine.delta import DeltaSyncConfig +from twinkle.checkpoint_engine.delta.rollout.mixin import DeltaSamplerMixin +from twinkle.checkpoint_engine.hccl_checkpoint_engine import MasterMetadata + +MASTER = MasterMetadata(zmq_ip='10.0.0.1', zmq_port=6000, dist_ip='10.0.0.1', dist_port=5000) + + +@pytest.fixture(autouse=True) +def fake_npu_platform(monkeypatch): + """No torch_npu here: resolve the platform names without importing the NPU backend.""" + classes = {'NPU': type('NPU', (), {}), 'GPU': type('GPU', (), {})} + monkeypatch.setattr(Platform, 'get_platform', staticmethod(lambda platform=None: classes[platform or 'GPU'])) + + +# --------------------------------------------------------------------------- manager +def npu_pair(tps=(2, ), trainer_device='npu', worker_device='npu'): + model = Mock(_actors=[1, 2], device_mesh=SimpleNamespace(world_size=2)) + sampler = Mock(_actors=list(range(len(tps))), device_mesh=SimpleNamespace(data_world_size=len(tps))) + model.prepare_delta_sync.return_value = dict(device=trainer_device, master_metadata=MASTER) + caps = [dict(tp=tp, device=worker_device, replica_rank=rank) for rank, tp in enumerate(tps)] + sampler.validate_delta_sampler.return_value = caps if len(caps) > 1 else caps[0] + model.send_delta_weights.return_value = lambda: {'payload_bytes': 6} + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}] * len(tps) + if len(tps) == 1: + sampler.receive_delta_weights.return_value = lambda: {'changed': True} + manager = CheckpointEngineManager(model, sampler, platform='NPU', sync_backend='delta_sharded') + return manager, model, sampler + + +def test_npu_preflight_order_and_no_worker_metadata(): + manager, model, sampler = npu_pair((2, 2)) + order = [] + sampler.validate_delta_sampler.side_effect = lambda cfg: (order.append('validate'), + [dict(tp=2, device='npu', replica_rank=0), + dict(tp=2, device='npu', replica_rank=1)])[1] + model.prepare_delta_sync.side_effect = lambda cfg: ( + order.append('prepare'), dict(device='npu', master_metadata=MASTER))[1] + manager.sync_weights() + assert order == ['validate', 'prepare'] + sampler.delta_sampler_metadata.assert_not_called() + + +def test_npu_endpoint_world_size_and_no_rollout_workers(): + manager, model, sampler = npu_pair((2, 2)) + manager.sync_weights() + manager.sync_weights() + endpoint = model.init_delta_transport.call_args.args[0] + assert endpoint['world_size'] == 5 and 'rollout_workers' not in endpoint + assert endpoint['master_metadata'] is MASTER + assert sampler.init_delta_sampler_transport.call_args.args[0] == endpoint + + +def test_npu_world_size_sums_tp_of_each_replica(): + manager, model, _ = npu_pair((1, )) + manager.sync_weights() + assert model.init_delta_transport.call_args.args[0]['world_size'] == 2 + + +def test_npu_begin_before_init_and_init_only_on_seed(): + events = [] + manager, model, sampler = npu_pair((2, )) + sampler.begin_delta_sync.side_effect = lambda: events.append('begin') + model.init_delta_transport.side_effect = lambda ep: (events.append('trainer-init'), lambda: None)[1] + sampler.init_delta_sampler_transport.side_effect = lambda ep: (events.append('sampler-init'), lambda: None)[1] + manager.sync_weights() + assert events == ['begin', 'trainer-init', 'sampler-init'] + manager.sync_weights() + assert events == ['begin', 'trainer-init', 'sampler-init', 'begin'] + assert manager.last_delta_metrics['seed'] is False + assert manager._delta_replica_count == 1 + + +def test_npu_device_mismatch_names_both_values(): + manager, model, sampler = npu_pair((2, ), trainer_device='npu', worker_device='cuda') + with pytest.raises(ValueError, match="'npu'.*'cuda'"): + manager.sync_weights() + sampler.begin_delta_sync.assert_not_called() + model.init_delta_transport.assert_not_called() + + +@pytest.mark.parametrize('ranks', [[0, 0], [1, 2], [0, 2], [None, 1], [0, None]]) +def test_npu_bad_replica_ranks_rejected(ranks): + manager, model, sampler = npu_pair((2, 2)) + sampler.validate_delta_sampler.return_value = [ + dict(tp=2, device='npu', **({} if rank is None else dict(replica_rank=rank))) for rank in ranks] + with pytest.raises(ValueError, match='delta replica ranks'): + manager.sync_weights() + sampler.begin_delta_sync.assert_not_called() + model.init_delta_transport.assert_not_called() + sampler.fail_delta_sampler.assert_not_called() + + +def test_npu_single_replica_must_be_rank_zero(): + manager, _, sampler = npu_pair((2, )) + sampler.validate_delta_sampler.return_value = dict(tp=2, device='npu', replica_rank=1) + with pytest.raises(ValueError, match='delta replica ranks'): + manager.sync_weights() + + +def test_npu_replica_order_does_not_matter(): + manager, model, sampler = npu_pair((2, 2)) + sampler.validate_delta_sampler.return_value = [dict(tp=2, device='npu', replica_rank=1), + dict(tp=2, device='npu', replica_rank=0)] + manager.sync_weights() + assert manager._delta_replica_count == 2 + + +def test_npu_replica_disagreement_keeps_rollout_paused(): + manager, model, sampler = npu_pair((2, 2)) + manager.sync_weights() + sampler.publish_delta_sampler.reset_mock() + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}, {'changed': False}] + with pytest.raises(ValueError, match='replicas disagree'): + manager.sync_weights() + sampler.publish_delta_sampler.assert_not_called() + sampler.fail_delta_sampler.assert_called_once() + + +def test_npu_replica_count_mismatch_fails(): + manager, _, sampler = npu_pair((2, 2)) + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}] + with pytest.raises(ValueError, match='replicas disagree'): + manager.sync_weights() + sampler.fail_delta_sampler.assert_called_once() + + +@pytest.mark.parametrize('stage', ['send_delta_weights', 'receive_delta_weights', 'init_delta_sampler_transport']) +def test_npu_failure_marks_failed_and_fails_sampler(stage): + manager, model, sampler = npu_pair((2, )) + getattr(model if stage == 'send_delta_weights' else sampler, stage).side_effect = TimeoutError('injected') + with pytest.raises(TimeoutError): + manager.sync_weights() + assert manager._delta_failed and manager._delta_round == 0 + sampler.fail_delta_sampler.assert_called_once() + sampler.publish_delta_sampler.assert_not_called() + with pytest.raises(RuntimeError, match='restart the entire'): + manager.sync_weights() + + +def test_npu_preflight_failure_is_retryable(): + manager, model, sampler = npu_pair((2, )) + model.prepare_delta_sync.side_effect = ValueError('bad config') + with pytest.raises(ValueError, match='bad config'): + manager.sync_weights() + sampler.begin_delta_sync.assert_not_called() + model.prepare_delta_sync.side_effect = None + manager.sync_weights() + assert manager.last_delta_metrics['seed'] is True + + +def test_npu_close_then_new_seed(): + manager, model, sampler = npu_pair((2, )) + manager.sync_weights() + manager.close_sync_session() + sampler.close_delta_sampler.assert_called_once() + manager.sync_weights() + assert model.init_delta_transport.call_count == 2 + assert sampler.validate_delta_sampler.call_count == 2 + + +def test_gpu_manager_still_uses_worker_metadata(): + model = Mock(_actors=[1], device_mesh=SimpleNamespace(world_size=1)) + sampler = Mock(_actors=[2], device_mesh=SimpleNamespace(data_world_size=1)) + model.prepare_delta_sync.return_value = dict(device='cuda', master_metadata=object()) + sampler.validate_delta_sampler.return_value = dict(device='cuda', tp=1) + sampler.delta_sampler_metadata.return_value = [dict(node_ip='n', device_uuid='g0', zmq_handle='h')] + sampler.receive_delta_weights.return_value = lambda: {'changed': True} + manager = CheckpointEngineManager(model, sampler, sync_backend='delta_sharded') + manager.sync_weights() + sampler.delta_sampler_metadata.assert_called_once() + assert 'rollout_workers' in model.init_delta_transport.call_args.args[0] + + +# --------------------------------------------------------------------------- sampler +class LocalSampler(DeltaSamplerMixin): + def _run_in_loop(self, coroutine): + return asyncio.run(coroutine) + + +def make_sampler(monkeypatch, *, rank=0, world=1, tp=2, device='npu', sleep=True): + if rank is None: + monkeypatch.delenv('RANK', raising=False) + else: + monkeypatch.setenv('RANK', str(rank)) + monkeypatch.setenv('WORLD_SIZE', str(world)) + sampler = LocalSampler() + sampler.engine = SimpleNamespace( + enable_sleep_mode=sleep, pause_delta_generation=AsyncMock(), + engine=SimpleNamespace(reset_prefix_cache=AsyncMock(return_value=True), reset_mm_cache=AsyncMock(), + reset_encoder_cache=AsyncMock())) + sampler._delta_rpc = AsyncMock(side_effect=lambda method, kwargs=None: ( + [dict(tp=tp, device=device)] * tp if method == 'validate_delta_worker' else [])) + return sampler + + +def rpc_calls(sampler): + return [(c.args[0], c.args[1] if len(c.args) > 1 else None) for c in sampler._delta_rpc.await_args_list] + + +def validated_npu_sampler(monkeypatch, *, rank=0, world=1, tp=2): + sampler = make_sampler(monkeypatch, rank=rank, world=world, tp=tp) + caps = sampler.validate_delta_sampler(DeltaSyncConfig(timeout_seconds=7)) + sampler.init_delta_sampler_transport({'world_size': 1 + world * tp, 'master_metadata': MASTER}) + sampler._delta_rpc.reset_mock() + return sampler, caps + + +def test_validate_returns_replica_rank_on_npu(monkeypatch): + sampler = make_sampler(monkeypatch, rank=1, world=2, tp=2) + assert sampler.validate_delta_sampler(DeltaSyncConfig()) == dict(tp=2, device='npu', replica_rank=1) + + +def test_validate_keeps_exact_legacy_dict_on_cuda(monkeypatch): + sampler = make_sampler(monkeypatch, rank=1, world=2, tp=2, device='cuda') + assert sampler.validate_delta_sampler(DeltaSyncConfig()) == dict(tp=2, device='cuda') + + +@pytest.mark.parametrize('rank, world', [(None, 1), (-1, 1), (2, 2), (3, 2)]) +def test_validate_npu_undeterminable_replica_rank_is_a_clear_error(monkeypatch, rank, world): + sampler = make_sampler(monkeypatch, rank=rank, world=world) + with pytest.raises(RuntimeError, match='cannot determine the delta replica rank'): + sampler.validate_delta_sampler(DeltaSyncConfig()) + + +def test_validate_npu_sleep_mode_message_covers_npu(monkeypatch): + sampler = make_sampler(monkeypatch, sleep=False) + with pytest.raises(ValueError, match='enable_sleep_mode=True or free_cache_engine=False'): + sampler.validate_delta_sampler(DeltaSyncConfig()) + sampler.validate_delta_sampler(DeltaSyncConfig(free_cache_engine=False)) + + +def test_npu_init_transport_spawns_no_actors_and_skips_receiver_metadata(monkeypatch): + import ray + remote = Mock() + monkeypatch.setattr(ray, 'remote', remote) + monkeypatch.setattr(ray, 'nodes', Mock(side_effect=AssertionError('must not touch ray'))) + sampler = make_sampler(monkeypatch) + sampler.validate_delta_sampler(DeltaSyncConfig()) + sampler._delta_rpc.reset_mock() + sampler._delta_weight_transfer_initialized = True + sampler.init_delta_sampler_transport({'world_size': 3, 'master_metadata': MASTER}) + remote.assert_not_called() + sampler._delta_rpc.assert_not_called() + assert not sampler._delta_weight_transfer_initialized + assert not getattr(sampler, '_delta_receivers', []) + + +def test_npu_round_rpc_sequence_and_kwargs(monkeypatch): + sampler, _ = validated_npu_sampler(monkeypatch, rank=0, world=1, tp=2) + assert sampler.receive_delta_weights() == {'changed': True} + assert rpc_calls(sampler) == [ + ('init_weight_transfer_engine', {'init_info': dict( + dist_ip='10.0.0.1', dist_port=5000, zmq_ip='10.0.0.1', zmq_port=6000, + topic='bucket_metadata', world_size=3)}), + ('start_weight_update', None), + ('update_twinkle_delta_weights', {'update_info': [{'group_rank': 1}, {'group_rank': 2}]}), + ('finish_weight_update', None), + ] + engine = sampler.engine.engine + engine.reset_prefix_cache.assert_awaited_once_with(reset_connector=True) + engine.reset_mm_cache.assert_awaited_once() + engine.reset_encoder_cache.assert_awaited_once() + + +def test_npu_init_info_forwards_trainer_hccl_env(monkeypatch): + sampler = make_sampler(monkeypatch, rank=0, world=1, tp=1) + sampler.validate_delta_sampler(DeltaSyncConfig()) + env = {'HCCL_IF_BASE_PORT': '30000'} + sampler.init_delta_sampler_transport({'world_size': 2, 'master_metadata': MASTER, 'hccl_env': env}) + sampler._delta_rpc.reset_mock() + sampler.receive_delta_weights() + method, kwargs = rpc_calls(sampler)[0] + assert method == 'init_weight_transfer_engine' and kwargs['init_info']['hccl_env'] == env + + +def test_npu_endpoint_carries_trainer_hccl_env(): + manager, model, sampler = npu_pair((2, )) + env = {'HCCL_IF_BASE_PORT': '30000'} + model.prepare_delta_sync.return_value = dict(device='npu', master_metadata=MASTER, hccl_env=env) + manager.sync_weights() + assert sampler.init_delta_sampler_transport.call_args.args[0]['hccl_env'] == env + + +@pytest.mark.parametrize('rank, world, tp, expected', [ + (0, 2, 2, [1, 2]), (1, 2, 2, [3, 4]), (0, 1, 1, [1]), (2, 3, 1, [3]), (1, 2, 4, [5, 6, 7, 8])]) +def test_npu_group_ranks_replica_major_tp_minor(monkeypatch, rank, world, tp, expected): + sampler, _ = validated_npu_sampler(monkeypatch, rank=rank, world=world, tp=tp) + sampler.receive_delta_weights() + update = [kw for name, kw in rpc_calls(sampler) if name == 'update_twinkle_delta_weights'][0] + assert [item['group_rank'] for item in update['update_info']] == expected + + +def test_npu_init_rpc_only_first_round_and_again_after_close(monkeypatch): + sampler, _ = validated_npu_sampler(monkeypatch) + sampler.receive_delta_weights() + sampler.receive_delta_weights() + names = [name for name, _ in rpc_calls(sampler)] + assert names.count('init_weight_transfer_engine') == 1 and names.count('update_twinkle_delta_weights') == 2 + sampler.close_delta_sampler() + sampler.validate_delta_sampler(DeltaSyncConfig()) + sampler.init_delta_sampler_transport({'world_size': 3, 'master_metadata': MASTER}) + sampler.receive_delta_weights() + names = [name for name, _ in rpc_calls(sampler)] + assert names.count('init_weight_transfer_engine') == 2 + + +def test_npu_close_shuts_workers_down_and_resets(monkeypatch): + sampler, _ = validated_npu_sampler(monkeypatch) + sampler.receive_delta_weights() + sampler._delta_rpc.reset_mock() + sampler.close_delta_sampler() + assert rpc_calls(sampler) == [('shutdown_delta_weight_transfer', None)] + sampler.engine.pause_delta_generation.assert_awaited_once_with(7) + assert sampler.engine._delta_unavailable is True + assert not sampler._delta_weight_transfer_initialized + + +def test_npu_close_resets_flag_even_when_worker_shutdown_fails(monkeypatch): + sampler, _ = validated_npu_sampler(monkeypatch) + sampler.receive_delta_weights() + sampler._delta_rpc.side_effect = RuntimeError('rpc down') + with pytest.raises(RuntimeError, match='rpc down'): + sampler.close_delta_sampler() + assert not sampler._delta_weight_transfer_initialized + + +@pytest.mark.parametrize('failing', ['init_weight_transfer_engine', 'update_twinkle_delta_weights']) +def test_npu_round_failure_propagates_without_cache_clear(monkeypatch, failing): + sampler, _ = validated_npu_sampler(monkeypatch) + + async def rpc(method, kwargs=None): + if method == failing: + raise TimeoutError('injected') + sampler._delta_rpc = AsyncMock(side_effect=rpc) + with pytest.raises(TimeoutError): + sampler.receive_delta_weights() + sampler.engine.engine.reset_prefix_cache.assert_not_called() + if failing == 'init_weight_transfer_engine': + assert not sampler._delta_weight_transfer_initialized + + +def test_npu_receive_without_transport_is_a_clear_error(monkeypatch): + sampler = make_sampler(monkeypatch) + sampler.validate_delta_sampler(DeltaSyncConfig()) + with pytest.raises(RuntimeError, match='not initialized'): + sampler.receive_delta_weights() + + +# --------------------------------------------------------------------------- vllm_engine backend wiring +@pytest.fixture +def engine_env(monkeypatch): + """Stub vLLM and the adapter module so VLLMEngine._create_engine runs up to AsyncEngineArgs.""" + import sys + from types import ModuleType + + captured = {} + + def args_cls(model=None, enable_sleep_mode=False, weight_transfer_config=None, **_): + captured['args'] = dict(enable_sleep_mode=enable_sleep_mode, weight_transfer_config=weight_transfer_config) + return SimpleNamespace(create_engine_config=lambda usage_context: 'vllm-config') + + modules = { + 'vllm.engine.arg_utils': SimpleNamespace(AsyncEngineArgs=args_cls), + 'vllm.usage.usage_lib': SimpleNamespace(UsageContext=SimpleNamespace(OPENAI_API_SERVER='api')), + 'vllm.v1.engine.async_llm': SimpleNamespace(AsyncLLM=SimpleNamespace(from_vllm_config=lambda **_: 'engine')), + } + for name, attrs in modules.items(): + module = ModuleType(name) + module.__dict__.update(vars(attrs)) + monkeypatch.setitem(sys.modules, name, module) + adapter = SimpleNamespace(require_vllm_delta_support=Mock(), register_twinkle_delta_weight_transfer_engine=Mock(), + register_twinkle_delta_hccl_weight_transfer_engine=Mock()) + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta.transfer_engine', adapter) + return SimpleNamespace(adapter=adapter, captured=captured) + + +def _bare_vllm_engine(**engine_kwargs): + from twinkle.sampler.vllm_sampler.vllm_engine import VLLMEngine + engine = VLLMEngine.__new__(VLLMEngine) + engine.__dict__.update( + model_id='m', tensor_parallel_size=1, gpu_memory_utilization=.5, max_num_seqs=1, trust_remote_code=False, + enforce_eager=True, dtype='bfloat16', load_format='auto', max_model_len=None, quantization=None, + enable_prefix_caching=False, enable_sleep_mode=False, runner=None, logprobs_mode=None, enable_lora=False, + enable_sampling_replay=False, engine_kwargs=engine_kwargs) + return engine + + +def test_vllm_engine_hccl_backend_registers_engine_without_forcing_sleep_mode(engine_env): + engine = _bare_vllm_engine(weight_transfer_config={'backend': 'twinkle_delta_hccl'}) + assert engine._create_engine() == 'engine' + engine_env.adapter.require_vllm_delta_support.assert_called_once() + engine_env.adapter.register_twinkle_delta_hccl_weight_transfer_engine.assert_called_once() + engine_env.adapter.register_twinkle_delta_weight_transfer_engine.assert_not_called() + assert engine.enable_sleep_mode is False and engine_env.captured['args']['enable_sleep_mode'] is False + assert engine_env.captured['args']['weight_transfer_config'] == {'backend': 'twinkle_delta_hccl'} + + +def test_vllm_engine_ipc_backend_still_enables_sleep_mode(engine_env): + engine = _bare_vllm_engine(weight_transfer_config={'backend': 'twinkle_delta_ipc'}) + engine._create_engine() + engine_env.adapter.register_twinkle_delta_weight_transfer_engine.assert_called_once() + engine_env.adapter.register_twinkle_delta_hccl_weight_transfer_engine.assert_not_called() + assert engine.enable_sleep_mode is True + + +def test_vllm_engine_hccl_backend_requires_weight_transfer_config_support(engine_env, monkeypatch): + import sys + + def old_args_cls(model=None, **_): + raise AssertionError('must fail before constructing args') + + sys.modules['vllm.engine.arg_utils'].AsyncEngineArgs = old_args_cls + engine = _bare_vllm_engine(weight_transfer_config={'backend': 'twinkle_delta_hccl'}) + with pytest.raises(RuntimeError, match='weight_transfer_config'): + engine._create_engine() diff --git a/tests/checkpoint_engine/test_delta_pipeline.py b/tests/checkpoint_engine/test_delta_pipeline.py new file mode 100644 index 000000000..ea1f69c7e --- /dev/null +++ b/tests/checkpoint_engine/test_delta_pipeline.py @@ -0,0 +1,205 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeltaSyncMixin driven by a fake CPU transport (no GPU, ray, zmq or NCCL).""" +from __future__ import annotations + +import asyncio + +import pytest +import torch + +from twinkle.checkpoint_engine.delta.transport.nccl_engine import DeltaShardedCheckpointEngine +from twinkle.checkpoint_engine.delta.transport.pipeline import DeltaSyncMixin +from twinkle.checkpoint_engine.nccl_checkpoint_engine import NCCLCheckpointEngine + +METRIC_KEYS = { + 'checkpoint_engine/changed_ratio', + 'checkpoint_engine/changed_elems', + 'checkpoint_engine/payload_mbytes', + 'checkpoint_engine/flushes', +} + + +class _Socket: + + def __init__(self, log): + self.log = log + self.metas = [] + + def send_string(self, s, flags=0): + self.log.append('manifest_topic') + + def send_pyobj(self, meta): + self.log.append('manifest') + self.metas.append(meta) + + +class _Base: + + def __init__(self, *args, bucket_size=1 << 20, is_master=True, rank=0, **kwargs): + self.bucket_size = bucket_size + self.rollout_dtype = torch.bfloat16 + self.is_master = is_master + self.rank = rank + self.topic = 'topic' + self.log = [] + self.socket = _Socket(self.log) + + +class FakeTransport(DeltaSyncMixin, _Base): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.broadcasts = [] + self.released = [] + + @property + def _wire_device(self): + return 'cpu' + + def _ensure_wire(self): + self.log.append('ensure_wire') + + def _broadcast_flush(self, pos_u8, val_u8): + self.log.append('broadcast') + self.broadcasts.append((pos_u8, val_u8)) + + def _release_after_round(self, phase): + self.released.append(phase) + + +class FakeTrainingEngine: + + def __init__(self, deltas): + self.deltas = deltas # list of (name, numel, changed idx) + self.primed = 0 + + def get_per_tensor_param(self): + def gen(): + yield 'a', torch.arange(8, dtype=torch.bfloat16) + yield 'b', torch.ones(4, 2, dtype=torch.bfloat16) + return gen(), None + + def prime_delta_snapshots(self): + self.primed += 1 + + def get_per_tensor_param_delta_shard(self): + def gen(): + for name, shape, idx in self.deltas: + idx_t = torch.tensor(idx, dtype=torch.int64) + val = torch.full((len(idx),), 2.0, dtype=torch.bfloat16) + yield [(name, shape)], 'bfloat16', torch.tensor([len(idx)]), idx_t, val, None + return gen(), None + + +def _run(t, eng): + return asyncio.run(t.send_weights(eng)) + + +def test_seed_then_steady(): + t = FakeTransport() + eng = FakeTrainingEngine([('a', (8,), [1, 3]), ('b', (4, 2), [0, 7, 5])]) + m = _run(t, eng) + assert set(m) == METRIC_KEYS + assert eng.primed == 1 and t.released == ['seed'] + assert t.socket.metas and all(x['encoding'] == 'dense' and x['is_full'] for x in t.socket.metas) + assert [x['is_last'] for x in t.socket.metas] == [False] * (len(t.socket.metas) - 1) + [True] + assert all(p is None for p, _ in t.broadcasts) + + t.log.clear() + t.socket.metas.clear() + t.broadcasts.clear() + m = _run(t, eng) + assert eng.primed == 1 and t.released == ['seed', 'steady'] + assert set(m) == METRIC_KEYS + assert m['checkpoint_engine/changed_elems'] == 5.0 + assert m['checkpoint_engine/changed_ratio'] == pytest.approx(5 / 16) + metas = t.socket.metas + assert len(metas) == 1 + assert metas[0]['encoding'] == 'indices' and metas[0]['is_full'] is False and metas[0]['is_last'] is True + pos, val = t.broadcasts[0] + assert pos.dtype == torch.uint8 and pos.numel() == 5 * 4 + assert pos.view(torch.int32).tolist() == [1, 3, 0, 7, 5] + assert val.numel() == 5 * 2 + assert t.log == ['ensure_wire', 'manifest_topic', 'manifest', 'broadcast'] + + +def test_small_bucket_multiple_flushes(): + t = FakeTransport(bucket_size=8) + eng = FakeTrainingEngine([('a', (8,), [1, 3]), ('b', (4, 2), [0, 7, 5])]) + _run(t, eng) + t.socket.metas.clear() + m = _run(t, eng) + flags = [x['is_last'] for x in t.socket.metas] + assert len(flags) > 1 and flags == [False] * (len(flags) - 1) + [True] + assert m['checkpoint_engine/flushes'] == float(len(flags)) + + +def test_seed_small_bucket_single_last(): + t = FakeTransport(bucket_size=8) + _run(t, FakeTrainingEngine([])) + flags = [x['is_last'] for x in t.socket.metas] + assert len(flags) > 1 and flags.count(True) == 1 and flags[-1] is True + + +def test_steady_no_change_terminal(): + t = FakeTransport() + eng = FakeTrainingEngine([('a', (8,), [])]) + _run(t, eng) + t.socket.metas.clear() + t.broadcasts.clear() + t.log.clear() + m = _run(t, eng) + assert len(t.socket.metas) == 1 + meta = t.socket.metas[0] + assert meta['terminal_empty'] is True and meta['is_last'] is True and meta['is_full'] is False + assert t.broadcasts == [] + assert t.log == ['ensure_wire', 'manifest_topic', 'manifest'] + assert m['checkpoint_engine/changed_elems'] == 0.0 + + +def test_non_master_publishes_nothing(): + t = FakeTransport(is_master=False, rank=-1) + eng = FakeTrainingEngine([('a', (8,), [1])]) + assert _run(t, eng) is None + assert _run(t, eng) is None + assert t.socket.metas == [] and t.broadcasts == [] and t.released == [] + assert eng.primed == 1 + + +def test_unimplemented_hooks(): + class Bare(DeltaSyncMixin, _Base): + pass + + b = Bare() + with pytest.raises(NotImplementedError): + b._wire_device + with pytest.raises(NotImplementedError): + b._broadcast_flush(None, torch.zeros(1, dtype=torch.uint8)) + with pytest.raises(NotImplementedError): + b._release_after_round('seed') + assert b._ensure_wire() is None + + +def test_encoding_must_be_indices(): + with pytest.raises(AssertionError): + FakeTransport(encoding='dense') + + +def test_nccl_engine_wiring(): + assert issubclass(DeltaShardedCheckpointEngine, DeltaSyncMixin) + assert issubclass(DeltaShardedCheckpointEngine, NCCLCheckpointEngine) + mro = DeltaShardedCheckpointEngine.__mro__ + assert mro.index(DeltaSyncMixin) < mro.index(NCCLCheckpointEngine) + assert DeltaShardedCheckpointEngine.MAX_ENTRY_ELEMS == DeltaSyncMixin.MAX_ENTRY_ELEMS + + e = object.__new__(DeltaShardedCheckpointEngine) + sent = [] + e._broadcast = sent.append + pos = torch.zeros(4, dtype=torch.uint8) + val = torch.zeros(2, dtype=torch.uint8) + e._broadcast_flush(None, val) + e._broadcast_flush(pos, val) + assert sent == [val, pos, val] + assert e._release_after_round('seed') is None + assert e._release_after_round('steady') is None + assert e._ensure_wire() is None diff --git a/tests/checkpoint_engine/test_delta_platform.py b/tests/checkpoint_engine/test_delta_platform.py new file mode 100644 index 000000000..a5d7ba293 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_platform.py @@ -0,0 +1,193 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Platform resolution of the delta transport and the entry points that use it.""" +import sys +import types +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from twinkle import DeviceMesh, Platform +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.checkpoint_engine.delta import DeltaSyncConfig, resolve_delta_engine_cls +from twinkle.checkpoint_engine.delta.trainer import mixin as delta_trainer +from twinkle.checkpoint_engine.delta.trainer.backends.fsdp import validate_fsdp2 +from twinkle.checkpoint_engine.delta.transport.nccl_engine import DeltaShardedCheckpointEngine +from twinkle.checkpoint_engine.delta.trainer.mixin import DeltaTrainerMixin +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + +HCCL_MODULE = 'twinkle.checkpoint_engine.delta.transport.hccl_engine' +NOT_IMPLEMENTED = 'not implemented yet' + + +def fake_platform(monkeypatch, name): + cls = type(name, (), {}) + monkeypatch.setattr(Platform, 'get_platform', staticmethod(lambda platform=None: cls)) + return cls + + +# ---- resolver ---- + +def test_resolver_gpu(): + assert resolve_delta_engine_cls('GPU') is DeltaShardedCheckpointEngine + + +def test_resolver_npu_real_engine(): + from twinkle.checkpoint_engine.delta.transport.hccl_engine import DeltaHCCLCheckpointEngine + assert resolve_delta_engine_cls('NPU') is DeltaHCCLCheckpointEngine + + +def test_resolver_npu_module_missing_not_implemented(monkeypatch): + # sys.modules[name] = None makes the import raise ModuleNotFoundError(name=HCCL_MODULE) + monkeypatch.setitem(sys.modules, HCCL_MODULE, None) + with pytest.raises(NotImplementedError, match=NOT_IMPLEMENTED): + resolve_delta_engine_cls('NPU') + + +def test_resolver_npu_returns_engine_class(monkeypatch): + module = types.ModuleType(HCCL_MODULE) + module.DeltaHCCLCheckpointEngine = type('DeltaHCCLCheckpointEngine', (), {}) + monkeypatch.setitem(sys.modules, HCCL_MODULE, module) + assert resolve_delta_engine_cls('NPU') is module.DeltaHCCLCheckpointEngine + + +def test_resolver_does_not_mask_other_import_errors(monkeypatch): + module = types.ModuleType(HCCL_MODULE) + + def boom(name): + raise ImportError('torch_npu is missing') + + module.__getattr__ = boom + monkeypatch.setitem(sys.modules, HCCL_MODULE, module) + with pytest.raises(ImportError, match='torch_npu is missing') as info: + resolve_delta_engine_cls('NPU') + assert not isinstance(info.value, NotImplementedError) + + +def test_resolver_unknown_platform(): + with pytest.raises(ValueError, match=r"'CPU'.*GPU, NPU"): + resolve_delta_engine_cls('CPU') + + +# ---- manager ---- + +def actors(): + return Mock(_actors=[1, 2]), Mock(_actors=[3]) + + +def test_manager_gpu_still_constructs(monkeypatch): + fake_platform(monkeypatch, 'GPU') + manager = CheckpointEngineManager(*actors(), platform='GPU', mode='standalone', sync_backend='delta_sharded') + assert manager.delta_engine_cls is DeltaShardedCheckpointEngine + + +def test_manager_full_sync_does_not_resolve_delta(monkeypatch): + fake_platform(monkeypatch, 'GPU') + manager = CheckpointEngineManager(*actors(), platform='GPU', mode='standalone') + assert manager.delta_engine_cls is None + + +def test_manager_npu_constructs_with_hccl_engine(monkeypatch): + fake_platform(monkeypatch, 'NPU') + from twinkle.checkpoint_engine.delta.transport.hccl_engine import DeltaHCCLCheckpointEngine + manager = CheckpointEngineManager(*actors(), platform='NPU', mode='standalone', sync_backend='delta_sharded') + assert manager.delta_engine_cls is DeltaHCCLCheckpointEngine + + +def test_manager_npu_module_missing_fails_early(monkeypatch): + fake_platform(monkeypatch, 'NPU') + monkeypatch.setitem(sys.modules, HCCL_MODULE, None) + with pytest.raises(NotImplementedError, match=NOT_IMPLEMENTED): + CheckpointEngineManager(*actors(), platform='NPU', mode='standalone', sync_backend='delta_sharded') + + +def test_manager_npu_resolves_when_engine_available(monkeypatch): + fake_platform(monkeypatch, 'NPU') + module = types.ModuleType(HCCL_MODULE) + module.DeltaHCCLCheckpointEngine = type('DeltaHCCLCheckpointEngine', (), {}) + monkeypatch.setitem(sys.modules, HCCL_MODULE, module) + manager = CheckpointEngineManager(*actors(), platform='NPU', mode='standalone', sync_backend='delta_sharded') + assert manager.delta_engine_cls is module.DeltaHCCLCheckpointEngine + + +def test_manager_non_standalone_still_rejected(monkeypatch): + fake_platform(monkeypatch, 'NPU') + with pytest.raises(ValueError, match='standalone'): + CheckpointEngineManager(SimpleNamespace(), SimpleNamespace(), platform='NPU', sync_backend='delta_sharded') + + +def test_manager_unknown_platform(monkeypatch): + fake_platform(monkeypatch, 'CPU') + # standalone backend selection rejects it first; either way construction fails early + with pytest.raises((ValueError, NotImplementedError)): + CheckpointEngineManager(*actors(), platform='CPU', mode='standalone', sync_backend='delta_sharded') + + +# ---- trainer ---- + +class Trainer(DeltaTrainerMixin): + + def validate_delta_backend(self): + pass + + +@pytest.mark.parametrize('name,device', [('GPU', 'cuda'), ('NPU', 'npu')]) +def test_trainer_device_follows_platform(monkeypatch, name, device): + fake_platform(monkeypatch, name) + created = {} + + class FakeEngine: + is_master = False + + def __init__(self, **kwargs): + created.update(kwargs) + + def prepare(self): + return 'metadata' + + monkeypatch.setattr(delta_trainer, 'resolve_delta_engine_cls', lambda platform: FakeEngine) + monkeypatch.setattr(delta_trainer.dist, 'is_initialized', lambda: False) + trainer = Trainer() + trainer._lazy_wrap_model = Mock() + result = trainer.prepare_delta_sync(DeltaSyncConfig()) + assert trainer._delta_device == device + assert result == {'device': device, 'master_metadata': 'metadata'} + assert isinstance(trainer._delta_engine, FakeEngine) + + +def test_trainer_npu_without_transport_fails_before_process_group(monkeypatch): + fake_platform(monkeypatch, 'NPU') + monkeypatch.setitem(sys.modules, HCCL_MODULE, None) + new_group = Mock() + monkeypatch.setattr(delta_trainer.dist, 'is_initialized', lambda: True) + monkeypatch.setattr(delta_trainer.dist, 'new_group', new_group) + with pytest.raises(NotImplementedError, match=NOT_IMPLEMENTED): + Trainer()._initialize_delta_engine(DeltaSyncConfig()) + new_group.assert_not_called() + + +# ---- validate_fsdp2 ---- + +def fsdp_owner(monkeypatch, name, **config): + fake_platform(monkeypatch, name) + model = SimpleNamespace(config=SimpleNamespace(**config), state_dict=lambda: {'w': torch.ones(4)}) + strategy = Mock(spec=NativeFSDPStrategy) + strategy.unwrap_model.return_value = model + mesh = DeviceMesh.from_sizes(world_size=1, fsdp_size=1, device_type='cuda') + return SimpleNamespace(strategy=strategy, model=model, device_mesh=mesh) + + +def test_validate_fsdp2_accepts_npu(monkeypatch): + validate_fsdp2(fsdp_owner(monkeypatch, 'NPU')) + + +def test_validate_fsdp2_rejects_unknown_platform(monkeypatch): + with pytest.raises(ValueError, match='GPU or NPU'): + validate_fsdp2(fsdp_owner(monkeypatch, 'CPU')) + + +@pytest.mark.parametrize('name', ['GPU', 'NPU']) +def test_validate_fsdp2_rejects_moe(monkeypatch, name): + with pytest.raises(ValueError, match='MoE'): + validate_fsdp2(fsdp_owner(monkeypatch, name, num_experts=8)) diff --git a/tests/checkpoint_engine/test_delta_rollout.py b/tests/checkpoint_engine/test_delta_rollout.py new file mode 100644 index 000000000..5b5826121 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_rollout.py @@ -0,0 +1,335 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rollout side: sampler coordination, GPU-local receivers, topology placement, request admission and device mapping.""" +import asyncio +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import AsyncMock, Mock, call + +import pytest +import torch + +from twinkle.checkpoint_engine.delta.rollout.cuda_receiver import CudaDeltaReceiver +from twinkle.checkpoint_engine.delta.rollout.mixin import DeltaSamplerMixin +from twinkle.sampler.vllm_sampler.vllm_engine import VLLMEngine +from twinkle.sampler.vllm_sampler.vllm_worker_extension import _cuda_visible_device + + +class LocalCudaSampler(DeltaSamplerMixin): + def _run_in_loop(self, coroutine): + return asyncio.run(coroutine) + + +def receiver_stub(flags): + return SimpleNamespace( + next_flush=SimpleNamespace(remote=AsyncMock(side_effect=flags)), + send_flush=SimpleNamespace(remote=AsyncMock()), + ) + + +def test_zero_change_skips_native_update_but_restores_kv_and_generation(): + sampler = LocalCudaSampler() + sampler._delta_receivers = [receiver_stub([None]), receiver_stub([None])] + sampler._delta_config = SimpleNamespace(timeout_seconds=1, free_cache_engine=True) + sampler.engine = SimpleNamespace(pause_delta_generation=AsyncMock(), engine=SimpleNamespace( + sleep=AsyncMock(), wake_up=AsyncMock(), reset_prefix_cache=AsyncMock(return_value=True), + resume_generation=AsyncMock())) + sampler._delta_rpc = AsyncMock() + sampler.begin_delta_sync() + assert sampler.receive_delta_weights() == {'changed': False} + sampler.publish_delta_sampler() + sampler._delta_rpc.assert_not_called() + sampler.engine.pause_delta_generation.assert_awaited_once_with(1) + sampler.engine.engine.sleep.assert_awaited_once_with(level=1) + assert sampler.engine.engine.wake_up.await_args_list == [ + call(tags=['weights']), call(tags=['kv_cache'])] + sampler.engine.engine.resume_generation.assert_awaited_once() + + +def test_sampler_coordinates_native_update_and_gpu_local_senders(): + sampler = LocalCudaSampler() + sampler._delta_receivers = [receiver_stub([{'is_last': True}, None]) for _ in range(2)] + sampler._delta_config = SimpleNamespace(timeout_seconds=1, free_cache_engine=False) + sampler._delta_weight_transfer_initialized = False + sampler.engine = SimpleNamespace(pause_delta_generation=AsyncMock(), engine=SimpleNamespace( + reset_prefix_cache=AsyncMock(return_value=True), reset_mm_cache=AsyncMock(), reset_encoder_cache=AsyncMock())) + sampler._delta_rpc = AsyncMock() + assert sampler.receive_delta_weights() == {'changed': True} + assert [call.args[0] for call in sampler._delta_rpc.await_args_list] == [ + 'init_weight_transfer_engine', 'start_weight_update', 'update_twinkle_delta_weights', 'finish_weight_update'] + for receiver in sampler._delta_receivers: + receiver.send_flush.remote.assert_awaited_once() + sampler.engine.engine.reset_prefix_cache.assert_awaited_once_with(reset_connector=True) + sampler.engine.engine.reset_mm_cache.assert_awaited_once() + sampler.engine.engine.reset_encoder_cache.assert_awaited_once() + + +def test_receiver_stream_disagreement_fails_before_native_update(): + sampler = LocalCudaSampler() + sampler._delta_receivers = [receiver_stub([None]), receiver_stub([{'is_last': True}])] + sampler._delta_rpc = AsyncMock() + with pytest.raises(ValueError, match='disagree'): + sampler.receive_delta_weights() + sampler._delta_rpc.assert_not_called() + + +def test_cache_reset_failure_does_not_resume_generation(): + sampler = LocalCudaSampler() + sampler._delta_receivers = [receiver_stub([{'is_last': True}, None])] + sampler._delta_weight_transfer_initialized = True + sampler._delta_rpc = AsyncMock() + sampler.engine = SimpleNamespace(engine=SimpleNamespace( + reset_prefix_cache=AsyncMock(return_value=False), resume_generation=AsyncMock())) + with pytest.raises(RuntimeError, match='prefix cache reset failed'): + sampler.receive_delta_weights() + sampler.engine.engine.resume_generation.assert_not_called() + + +def test_without_kv_release_pause_and_resume_still_run(): + sampler = LocalCudaSampler() + sampler._delta_config = SimpleNamespace(timeout_seconds=1, free_cache_engine=False) + sampler.engine = SimpleNamespace(pause_delta_generation=AsyncMock(), engine=SimpleNamespace( + sleep=AsyncMock(), wake_up=AsyncMock(), resume_generation=AsyncMock())) + sampler.begin_delta_sync() + sampler.publish_delta_sampler() + sampler.engine.pause_delta_generation.assert_awaited_once() + sampler.engine.engine.resume_generation.assert_awaited_once() + sampler.engine.engine.sleep.assert_not_called() + sampler.engine.engine.wake_up.assert_not_called() + + +def test_checkpoint_receiver_keeps_flush_on_local_gpu_until_ipc(monkeypatch): + from twinkle.checkpoint_engine.delta.rollout import cuda_receiver as delta_rollout + engine = Mock() + tensors = [('__values__', torch.zeros(2))] + engine.receive_weights.return_value = iter([(tensors, True)]) + monkeypatch.setattr(delta_rollout, 'DeltaShardedCheckpointEngine', Mock(return_value=engine)) + monkeypatch.setattr(torch.cuda, 'set_device', Mock()) + monkeypatch.setattr(torch.cuda, 'get_device_properties', Mock(return_value=SimpleNamespace(uuid='GPU-test'))) + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '3') + sender = SimpleNamespace(async_send_weights=AsyncMock()) + module = ModuleType('twinkle.sampler.vllm_sampler.delta.cuda_ipc') + module.BucketedWeightSender = Mock(return_value=sender) + monkeypatch.setitem(sys.modules, module.__name__, module) + config = SimpleNamespace(bucket_bytes=64 << 20) + receiver = CudaDeltaReceiver(config, {'world_size': 3, 'master_metadata': 'master'}, + {'rank': 2, 'device_uuid': 'GPU-test', 'cuda_visible_device': '3', + 'zmq_handle': 'ipc:///tmp/rank2.sock'}) + engine.init_process_group.assert_called_once_with(2, 3, 'master') + assert receiver.next_flush() == {'is_last': True} + receiver.send_flush() + assert list(sender.async_send_weights.call_args.args[0]) == tensors + assert receiver.next_flush() is None + assert receiver._current is None + + +def test_receiver_binds_gpu_before_cuda_initialization(monkeypatch): + from twinkle.checkpoint_engine.delta.rollout import cuda_receiver as delta_rollout + import os + + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '') + observed = [] + monkeypatch.setattr(torch.cuda, 'set_device', lambda device: observed.append(os.environ['CUDA_VISIBLE_DEVICES'])) + monkeypatch.setattr(torch.cuda, 'get_device_properties', Mock(return_value=SimpleNamespace(uuid='GPU-test'))) + monkeypatch.setattr(delta_rollout, 'DeltaShardedCheckpointEngine', Mock()) + CudaDeltaReceiver(SimpleNamespace(bucket_bytes=64), dict(world_size=2, master_metadata='master'), + dict(rank=1, device_uuid='GPU-test', cuda_visible_device='3', zmq_handle='ipc:///test')) + assert observed == ['3'] + + +def test_receiver_rejects_gpu_mismatch_before_joining_group(monkeypatch): + from twinkle.checkpoint_engine.delta.rollout import cuda_receiver as delta_rollout + + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '3') + monkeypatch.setattr(torch.cuda, 'set_device', Mock()) + monkeypatch.setattr(torch.cuda, 'get_device_properties', Mock(return_value=SimpleNamespace(uuid='GPU-wrong'))) + engine = Mock() + monkeypatch.setattr(delta_rollout, 'DeltaShardedCheckpointEngine', engine) + with pytest.raises(ValueError, match='GPU mismatch'): + CudaDeltaReceiver(SimpleNamespace(bucket_bytes=64), dict(world_size=2, master_metadata='master'), + dict(rank=1, device_uuid='GPU-test', cuda_visible_device='3', zmq_handle='ipc:///test')) + engine.assert_not_called() + + +@pytest.mark.parametrize('use_worker_selector', [False, True]) +def test_topology_places_one_receiver_on_each_vllm_gpu(monkeypatch, use_worker_selector): + import ray + metadata = [dict(node_ip='10.0.0.1', device_uuid=f'GPU-{rank}', zmq_handle=f'ipc:///tmp/rank{rank}.sock') + for rank in range(2)] + if use_worker_selector: + for rank, worker in enumerate(metadata): + worker['cuda_visible_device'] = str(rank + 2) + created = [] + + class Actor: + def options(self, **options): + self.options_for_rank = options + return self + + def remote(self, config, endpoint, placement): + created.append((self.options_for_rank, placement['rank'], placement['zmq_handle'])) + return SimpleNamespace(ready=SimpleNamespace(remote=lambda: placement['rank'])) + + remote = Mock(return_value=lambda cls: Actor()) + monkeypatch.setattr(ray, 'remote', remote) + monkeypatch.setattr(ray, 'nodes', lambda: [dict(Alive=True, NodeManagerAddress='10.0.0.1', NodeID='ab' * 28)]) + monkeypatch.setattr(ray, 'get', lambda refs, **kwargs: refs) + sampler = LocalCudaSampler() + sampler._delta_rpc = AsyncMock(return_value=metadata) + sampler._delta_config = SimpleNamespace(timeout_seconds=1) + sampler.init_delta_sampler_transport({'world_size': 3, 'master_metadata': 'master', + 'rollout_workers': metadata}) + assert [item[1:] for item in created] == [(1, 'ipc:///tmp/rank0.sock'), (2, 'ipc:///tmp/rank1.sock')] + expected_selectors = ['2', '3'] if use_worker_selector else ['GPU-0', 'GPU-1'] + assert [item[0]['runtime_env']['env_vars']['CUDA_VISIBLE_DEVICES'] for item in created] == expected_selectors + assert all(item[0]['scheduling_strategy'].node_id == 'ab' * 28 for item in created) + assert remote.call_args.kwargs == {'num_cpus': 0, 'num_gpus': 0} + + +def test_replica_receivers_use_global_ranks_across_nodes(monkeypatch): + import ray + workers = [dict(node_ip=f'10.0.0.{rank // 2 + 1}', device_uuid=f'GPU-{rank}', + zmq_handle=f'ipc:///rank{rank}') for rank in range(4)] + monkeypatch.setattr(ray, 'nodes', lambda: [ + dict(Alive=True, NodeManagerAddress=f'10.0.0.{rank + 1}', NodeID=str(rank)) for rank in range(2)]) + endpoint = dict(world_size=5, master_metadata='master', rollout_workers=workers) + sampler = LocalCudaSampler() + placements = [sampler._delta_receiver_placements(workers[:2], endpoint), + sampler._delta_receiver_placements(workers[2:], endpoint)] + assert [[item['rank'] for item in replica] for replica in placements] == [[1, 2], [3, 4]] + assert [[item['node_id'] for item in replica] for replica in placements] == [['0', '0'], ['1', '1']] + + +def test_sampler_dp_allows_independent_rollout_replicas(): + from twinkle.checkpoint_engine.delta import DeltaSyncConfig + sampler = LocalCudaSampler() + sampler.device_mesh = SimpleNamespace(data_world_size=2) + sampler.engine = SimpleNamespace(enable_sleep_mode=True) + sampler._delta_rpc = AsyncMock(return_value=[dict(tp=2, device='cuda')] * 2) + assert sampler.validate_delta_sampler(DeltaSyncConfig()) == dict(tp=2, device='cuda') + + +def test_submission_barrier_waits_for_request_already_past_gate(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + admitted = asyncio.Event() + release = asyncio.Event() + paused = asyncio.Event() + engine.engine = SimpleNamespace(pause_generation=AsyncMock(side_effect=lambda **_: paused.set())) + + async def source(): + admitted.set() + await release.wait() + yield 'first' + + async def consume(): + return [output async for output in engine._delta_admitted_outputs(source())] + + request = asyncio.create_task(consume()) + await admitted.wait() + update = asyncio.create_task(engine.pause_delta_generation(1)) + await asyncio.sleep(0) + assert engine._delta_unavailable + assert not paused.is_set() + with pytest.raises(RuntimeError, match='unavailable'): + await consume() + release.set() + assert await request == ['first'] + await update + assert paused.is_set() + assert engine._delta_admitting == 0 + asyncio.run(run()) + + +def test_admission_error_does_not_strand_barrier(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + engine.engine = SimpleNamespace(pause_generation=AsyncMock()) + async def source(): + raise ValueError('admission failed') + yield + with pytest.raises(ValueError, match='admission failed'): + async for _ in engine._delta_admitted_outputs(source()): + pass + await engine.pause_delta_generation(1) + engine.engine.pause_generation.assert_awaited_once() + assert engine._delta_admitting == 0 + asyncio.run(run()) + + +def test_admission_timeout_keeps_gate_closed_and_does_not_pause_or_write(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + engine._delta_admitting = 1 + engine.engine = SimpleNamespace(pause_generation=AsyncMock()) + with pytest.raises(TimeoutError): + await engine.pause_delta_generation(.01) + engine.engine.pause_generation.assert_not_awaited() + assert engine._delta_unavailable + asyncio.run(run()) + + +def test_stream_close_closes_admitted_engine_generator(monkeypatch): + inputs = ModuleType('vllm.inputs') + inputs.TextPrompt = lambda **fields: fields + inputs.TokensPrompt = lambda **fields: fields + sampling = ModuleType('vllm.sampling_params') + sampling.RequestOutputKind = SimpleNamespace(DELTA=object()) + monkeypatch.setitem(sys.modules, inputs.__name__, inputs) + monkeypatch.setitem(sys.modules, sampling.__name__, sampling) + + async def run(): + closed = asyncio.Event() + async def source(**kwargs): + try: + yield SimpleNamespace(outputs=[SimpleNamespace(token_ids=[1], finish_reason=None)]) + yield SimpleNamespace(outputs=[]) + finally: + closed.set() + engine = VLLMEngine.__new__(VLLMEngine) + engine.enable_lora = False + engine._synced_lora_request = None + engine.engine = SimpleNamespace(generate=source) + engine.get_tokenizer = AsyncMock(return_value=SimpleNamespace(decode=lambda *args, **kwargs: 'token')) + params = SimpleNamespace(to_vllm=lambda **kwargs: SimpleNamespace()) + stream = engine.generate_stream([1], params) + assert await anext(stream) == ('token', None) + await stream.aclose() + assert closed.is_set() + assert engine._delta_admitting == 0 + asyncio.run(run()) + + +@pytest.mark.parametrize(('visible_devices', 'expected'), [ + ('2,3', '3'), + ('3,2', '2'), + ('GPU-first,GPU-second', 'GPU-second'), + ('MIG-first,MIG-second', 'MIG-second'), + (' 2, 3 ', '3'), +]) +def test_selector_uses_worker_local_device_mapping(monkeypatch, visible_devices, expected): + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', visible_devices) + assert _cuda_visible_device(SimpleNamespace(index=1)) == expected + + +def test_single_card_worker_preserves_assigned_selector(monkeypatch): + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '3') + assert _cuda_visible_device(SimpleNamespace(index=0)) == '3' + + +def test_unrestricted_worker_uses_device_ordinal(monkeypatch): + monkeypatch.delenv('CUDA_VISIBLE_DEVICES', raising=False) + assert _cuda_visible_device(SimpleNamespace(index=2)) == '2' + + +@pytest.mark.parametrize('visible_devices', ['', '-1', '2,']) +def test_invalid_worker_visibility_is_rejected(monkeypatch, visible_devices): + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', visible_devices) + device = SimpleNamespace(index=1 if visible_devices == '2,' else 0) + with pytest.raises(ValueError, match='Cannot map CUDA device'): + _cuda_visible_device(device) + + +def test_device_ordinal_outside_visible_list_is_rejected(monkeypatch): + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '3') + with pytest.raises(ValueError, match='Cannot map CUDA device'): + _cuda_visible_device(SimpleNamespace(index=1)) diff --git a/tests/checkpoint_engine/test_delta_trainer.py b/tests/checkpoint_engine/test_delta_trainer.py new file mode 100644 index 000000000..097d96d64 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_trainer.py @@ -0,0 +1,438 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Trainer side: diff, placement and gather primitives, plus the FSDP2 delta backend and its lifecycle.""" +from __future__ import annotations + +import asyncio +import sys +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.device_mesh import DeviceMesh as TorchMesh +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +from twinkle import DeviceMesh, Platform +from twinkle.checkpoint_engine.delta.trainer.backends import FSDP2DeltaBackend +from twinkle.checkpoint_engine.delta.trainer.backends.fsdp import validate_fsdp2 +from twinkle.checkpoint_engine.delta.transport.nccl_engine import DeltaShardedCheckpointEngine, _GatherQueue +from twinkle.checkpoint_engine.delta.trainer.export import shard_delta_indices +from twinkle.checkpoint_engine.delta.trainer.mixin import DeltaTrainerMixin +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_shard_delta_indices_matches_bytewise_diff(dtype): + torch.manual_seed(0) + # A "shard" of some parameter, whose flat start in the full param is `offset`. + shard = torch.randn(1000, dtype=dtype) + new = shard.clone() + changed = torch.tensor([3, 17, 500, 999], dtype=torch.int64) + new[changed] = new[changed] + 0.5 + offset = 4096 # this shard begins at flat position 4096 within the full param + + gidx, gval = shard_delta_indices(new, shard, offset) + + # positions are (offset + local changed index), bytewise-exact values + assert torch.equal(gidx.sort().values, (changed + offset).sort().values) + order = torch.argsort(gidx) + got_pos = (gidx[order] - offset).to(torch.int64) + assert torch.equal( + gval[order].view(torch.int16 if dtype == torch.bfloat16 else torch.int32), + new[got_pos].view(torch.int16 if dtype == torch.bfloat16 else torch.int32), + ) + + +def test_shard_delta_indices_no_change_is_empty(): + shard = torch.randn(256, dtype=torch.bfloat16) + gidx, gval = shard_delta_indices(shard.clone(), shard, offset=0) + assert gidx.numel() == 0 + assert gval.numel() == 0 + + +def test_derive_dtensor_placement_unsharded(): + # A non-DTensor (replicated / unsharded) param: offset 0, no gather group, + # and outside a process group rank 0 is assumed -> contributes. + from twinkle.checkpoint_engine.delta.trainer.spec import ShardSpec, derive_dtensor_placement + + t = torch.randn(64, 8, dtype=torch.bfloat16) + spec = ShardSpec.from_param(t) + assert spec.mesh is None and spec.full_shape == (64, 8) + offset, contributes, group = derive_dtensor_placement(spec) + assert offset == 0 and contributes is True and group is None + + +def test_gather_slot_entries_sub_rounds_world1(): + """max_round_bytes splits the slot list into deterministic sub-rounds; the + reassembled output must equal the single-round result (world=1 mechanics).""" + import os + + import torch.distributed as dist + + from twinkle.checkpoint_engine.delta.transport.gather import gather_slot_entries_to_rank0 + + owns_pg = not dist.is_initialized() + if owns_pg: + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29512") + dist.init_process_group(backend="gloo", rank=0, world_size=1) + try: + torch.manual_seed(7) + k = 9 + counts = torch.tensor([5, 0, 3, 7, 1, 0, 4, 2, 6], dtype=torch.int64) + n = int(counts.sum()) + idx = torch.randint(0, 1000, (n,), dtype=torch.int32) + val = torch.randn(n, dtype=torch.bfloat16) + + ref = gather_slot_entries_to_rank0(idx, val, counts) + per_elem = idx.element_size() + val.element_size() + for budget_elems in (1, 4, 8): + got = gather_slot_entries_to_rank0(idx, val, counts, max_round_bytes=budget_elems * per_elem) + assert len(got) == k + for (ri, rv), (gi, gv) in zip(ref, got, strict=True): + assert torch.equal(ri, gi) + assert torch.equal(rv, gv) + finally: + # leaving the default pg alive leaks into whatever test runs next in the + # same session (repo convention: init in the test -> destroy in finally) + if owns_pg: + dist.destroy_process_group() + + +def test_prime_then_hf_delta_export_roundtrip(): + """The backend-side default delta strategy: prime_delta_snapshots pins the + shards; a later pass through hf_delta_export yields a final-HF-coordinate + entry with exactly the changed elements and refreshes the snapshot (a + second delta pass yields a zero-count entry).""" + from twinkle.checkpoint_engine.delta.trainer.spec import ShardSpec + from twinkle.checkpoint_engine.delta.trainer.export import _hf_entry_identity, hf_delta_export, prime_delta_snapshots + + w = torch.arange(12, dtype=torch.float32) + spec = ShardSpec(full_shape=(12,)) + snaps: dict = {} + + prime_delta_snapshots(iter([("w", w.clone(), spec)]), snaps, pin=False) + assert "w" in snaps and snaps["w"].numel() == 12 + + w2 = w.clone() + w2[3] = -1.0 + w2[7] = 42.0 + + ((slots, dtype_str, counts, hf_idx, hf_val, pg),) = list( + hf_delta_export(iter([("w", w2, spec)]), snaps, _hf_entry_identity) + ) + assert slots == [("w", (12,))] and dtype_str == "float32" and pg is None + assert counts.tolist() == [2] + assert hf_idx.dtype == torch.int32 and hf_idx.tolist() == [3, 7] + assert hf_val.tolist() == [-1.0, 42.0] + + ((_, _, counts2, hf_idx2, _, _),) = list( + hf_delta_export(iter([("w", w2.clone(), spec)]), snaps, _hf_entry_identity) + ) + assert counts2.tolist() == [0] and hf_idx2.numel() == 0 + + +def test_hf_delta_export_requires_seed(): + """A delta export without a prior prime must fail loud, not diff against + garbage.""" + import pytest + + from twinkle.checkpoint_engine.delta.trainer.spec import ShardSpec + from twinkle.checkpoint_engine.delta.trainer.export import _hf_entry_identity, hf_delta_export + + def raw(): + yield "w", torch.zeros(4), ShardSpec(full_shape=(4,)) + + with pytest.raises(AssertionError, match="seed snapshot"): + list(hf_delta_export(raw(), {}, _hf_entry_identity)) + + +class FSDPDeltaTrainer(DeltaTrainerMixin, FSDP2DeltaBackend): + """A model host: lifecycle mixin plus the FSDP2 export backend.""" + + +def make_delta_trainer(owner): + trainer = FSDPDeltaTrainer() + trainer.__dict__.update(vars(owner)) + return trainer + + +@pytest.fixture +def fsdp_owner(monkeypatch): + platform = type('GPU', (), {}) + monkeypatch.setattr(Platform, 'get_platform', staticmethod(lambda: platform)) + model = SimpleNamespace(config=SimpleNamespace(), state_dict=lambda: {}) + strategy = Mock(spec=NativeFSDPStrategy) + strategy.unwrap_model.return_value = model + return SimpleNamespace(strategy=strategy, model=model) + + +def test_default_two_card_fsdp_passes_parallelism_check(fsdp_owner): + fsdp_owner.device_mesh = DeviceMesh.from_sizes(world_size=2, fsdp_size=2, device_type='cuda') + with pytest.raises(ValueError, match='wrapped FSDP2 DTensor model'): + validate_fsdp2(fsdp_owner) + + +def test_single_rank_trainer_accepts_unsharded_model(fsdp_owner): + fsdp_owner.device_mesh = DeviceMesh.from_sizes(world_size=1, fsdp_size=1, device_type='cuda') + fsdp_owner.model.state_dict = lambda: {'w': torch.ones(4, dtype=torch.bfloat16)} + validate_fsdp2(fsdp_owner) + + +def test_single_rank_trainer_uses_rank_zero_without_process_group(monkeypatch): + from twinkle.checkpoint_engine.delta.trainer import mixin as delta_trainer + monkeypatch.setattr(delta_trainer.dist, 'is_initialized', lambda: False) + assert FSDPDeltaTrainer._delta_rank() == 0 + trainer = FSDPDeltaTrainer() + trainer.validate_delta_backend = Mock() + trainer._validate_delta_ranks() + assert trainer._delta_metadata_group is None + + +@pytest.mark.parametrize('size', [None, 0, 1]) +def test_disabled_parallel_dimensions_pass_parallelism_check(fsdp_owner, size): + fsdp_owner.device_mesh = SimpleNamespace(tp_world_size=size, pp_world_size=size, ep_world_size=size) + with pytest.raises(ValueError, match='wrapped FSDP2 DTensor model'): + validate_fsdp2(fsdp_owner) + + +@pytest.mark.parametrize('dimension', ['tp_world_size', 'pp_world_size', 'ep_world_size']) +def test_enabled_training_parallelism_is_rejected(fsdp_owner, dimension): + fsdp_owner.device_mesh = SimpleNamespace(tp_world_size=0, pp_world_size=0, ep_world_size=0) + setattr(fsdp_owner.device_mesh, dimension, 2) + with pytest.raises(ValueError, match='training TP=PP=EP=1'): + validate_fsdp2(fsdp_owner) + + +def test_full_seed_stages_dtensor_before_full_tensor(monkeypatch): + calls = [] + + class HostDTensor: + def to(self, device, non_blocking): + calls.append(('stage', device, non_blocking)) + return self + + def full_tensor(self): + calls.append(('full',)) + return torch.ones(2) + + monkeypatch.setattr(sys.modules['torch.distributed.tensor'], 'DTensor', HostDTensor) + exporter = make_delta_trainer(SimpleNamespace(_delta_device='cuda')) + exporter._delta_state_dict = lambda: {'w': HostDTensor()} + weights, _ = exporter.get_per_tensor_param() + list(weights) + assert calls == [('stage', 'cuda', True), ('full',)] + + +def test_full_sync_baseline_also_stages_cpu_offloaded_dtensor(monkeypatch): + from twinkle.model.transformers.transformers import TransformersModel + from twinkle.utils.framework import Torch + from torch.distributed import tensor as tensor_module + calls = [] + + class HostDTensor: + def to(self, device, non_blocking): + calls.append(('stage', device, non_blocking)) + return self + + def full_tensor(self): + calls.append(('full',)) + return torch.ones(2) + + monkeypatch.setattr(tensor_module, 'DTensor', HostDTensor) + monkeypatch.setattr(Torch, 'get_device', lambda rank: 'cuda') + model = SimpleNamespace(state_dict=lambda: {'w': HostDTensor()}) + owner = SimpleNamespace(model=model, strategy=SimpleNamespace(unwrap_model=lambda model: model)) + list(TransformersModel._get_weight_generator(owner, adapter_name='', merge_and_sync=True)) + assert calls == [('stage', 'cuda', True), ('full',)] + + +def test_hf_key_conversion_matches_seed_and_steady_export(): + weights = {'internal.weight': torch.zeros(2)} + model = SimpleNamespace(state_dict=lambda: weights, _checkpoint_conversion_mapping={'checkpoint': 'internal'}) + owner = SimpleNamespace(_delta_device='cpu', model=model, strategy=SimpleNamespace(unwrap_model=lambda m: m)) + exporter = make_delta_trainer(owner) + full, _ = exporter.get_per_tensor_param() + assert [name for name, _ in full] == ['checkpoint.weight'] + exporter.prime_delta_snapshots() + weights['internal.weight'][1] = 1 + delta, _ = exporter.get_per_tensor_param_delta_shard() + entries = list(delta) + assert entries[0][0] == [('checkpoint.weight', (2,))] + assert entries[0][3].tolist() == [1] + + +def test_full_sync_baseline_uses_same_hf_key_conversion(): + from twinkle.model.transformers.transformers import TransformersModel + model = SimpleNamespace(state_dict=lambda: {'internal.weight': torch.ones(1)}, + _checkpoint_conversion_mapping={'checkpoint': 'internal'}) + owner = SimpleNamespace(model=model, strategy=SimpleNamespace(unwrap_model=lambda m: m)) + weights = TransformersModel._get_weight_generator(owner, adapter_name='', merge_and_sync=True) + assert [name for name, _ in weights] == ['checkpoint.weight'] + + +def test_snapshot_pinning_can_be_disabled(monkeypatch): + from twinkle.checkpoint_engine.delta.trainer.backends import base as backend_base + prime = Mock() + monkeypatch.setattr(backend_base, 'prime_delta_snapshots', prime) + exporter = make_delta_trainer(SimpleNamespace(_delta_device='cuda', _delta_pin_snapshots=False)) + exporter.get_per_tensor_param_shard = lambda: (iter([]), None) + exporter.prime_delta_snapshots() + assert prime.call_args.kwargs == {'pin': False} + + +def test_seed_uses_whole_parameter_and_steady_uses_verl_schema(): + weights = {'w': torch.tensor([0., 1., 2., 3.]), 'v': torch.ones(3)} + model = SimpleNamespace(state_dict=lambda: weights) + owner = SimpleNamespace(_delta_device='cpu', model=model, strategy=SimpleNamespace(unwrap_model=lambda m: m)) + exporter = make_delta_trainer(owner) + engine = DeltaShardedCheckpointEngine(bucket_size=6, batch_gather=32) + engine.is_master, engine.rank = True, 0 + packets = [] + engine._publish_values_flush = lambda params, values, **flags: packets.append(('dense', params, values.clone(), flags)) + engine._publish_flush = lambda flush, **flags: packets.append(('indices', flush, flags)) + engine._publish_terminal = lambda first: packets.append(('terminal', first)) + asyncio.run(engine.send_weights(exporter)) + assert [packet[2].numel() for packet in packets] == [4, 3] + packets.clear() + weights['w'][0] = -0. + asyncio.run(engine.send_weights(exporter)) + flush = packets[0][1] + assert (flush.params[0].name, flush.positions_cpu.view(torch.int32).tolist(), + flush.values_gpu.view(torch.int16).tolist(), packets[0][2]['is_last']) == ('w', [0], [-32768], True) + packets.clear() + asyncio.run(engine.send_weights(exporter)) + assert packets == [('terminal', False)] + + +def test_eval_and_same_layout_restore_keep_snapshot_and_sparse_stream(): + model = torch.nn.Linear(3, 2, bias=False).to(torch.bfloat16) + trainer = make_delta_trainer(SimpleNamespace( + _delta_device='cpu', model=model, strategy=SimpleNamespace(unwrap_model=lambda m: m))) + engine = DeltaShardedCheckpointEngine(bucket_size=64 << 20) + engine.is_master, engine.rank = True, 0 + dense = Mock() + sparse = Mock() + terminal = Mock() + engine._publish_values_flush = dense + engine._publish_flush = sparse + engine._publish_terminal = terminal + asyncio.run(engine.send_weights(trainer)) + snapshot = trainer._delta_shard_snap['weight'] + baseline = snapshot.clone() + model.eval() + asyncio.run(engine.send_weights(trainer)) + terminal.assert_called_once_with(False) + model.train() + model.load_state_dict({'weight': torch.full_like(model.weight, 2)}) + asyncio.run(engine.send_weights(trainer)) + assert trainer._delta_shard_snap['weight'] is snapshot + assert dense.call_count == 1 + assert sparse.call_count == 1 + flush = sparse.call_args.args[0] + restored = baseline.clone() + restored[flush.positions_cpu.view(torch.int32).long()] = flush.values_gpu + assert torch.equal(restored.view(torch.int16), model.weight.detach().flatten().view(torch.int16)) + + +def test_trainer_prepare_reuses_engine_between_workflow_rounds(): + from twinkle.checkpoint_engine.delta import DeltaSyncConfig + + trainer = FSDPDeltaTrainer() + trainer._lazy_wrap_model = Mock() + engine = object() + master_metadata = object() + + def initialize(config): + trainer._delta_engine = engine + trainer._delta_master_metadata = master_metadata + + trainer._initialize_delta_engine = Mock(side_effect=initialize) + config = DeltaSyncConfig() + first = trainer.prepare_delta_sync(config) + second = trainer.prepare_delta_sync(config) + assert first == second == dict(device='cuda', master_metadata=master_metadata) + trainer._initialize_delta_engine.assert_called_once_with(config) + + +def test_trainer_validation_failure_releases_metadata_group(monkeypatch): + from twinkle.checkpoint_engine.delta import DeltaSyncConfig + from twinkle.checkpoint_engine.delta.trainer import mixin as delta_trainer + + trainer = FSDPDeltaTrainer() + group = object() + monkeypatch.setattr(delta_trainer.dist, 'is_initialized', lambda: True) + monkeypatch.setattr(delta_trainer.dist, 'new_group', Mock(return_value=group)) + destroy = Mock() + monkeypatch.setattr(delta_trainer.dist, 'destroy_process_group', destroy) + trainer._validate_delta_ranks = Mock(side_effect=ValueError('unsupported layout')) + with pytest.raises(ValueError, match='unsupported layout'): + trainer._initialize_delta_engine(DeltaSyncConfig()) + destroy.assert_called_once_with(group) + assert trainer._delta_metadata_group is None + assert trainer._delta_engine is None + + +def exercise_fsdp_export(rank, rendezvous): + dist.init_process_group('gloo', init_method=rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=30)) + try: + mesh = TorchMesh('cpu', [0, 1]) + replicated_mesh = TorchMesh('cpu', [[0, 1]]) + layouts = [(mesh, [Shard(0)]), (mesh, [Shard(1)]), + (replicated_mesh, [Shard(0), Replicate()])] + for parameter_mesh, placements in layouts: + baseline = torch.arange(15, dtype=torch.bfloat16).reshape(5, 3) + target = baseline.clone() + target.view(torch.int16)[0, 0] ^= -32768 + target[4, 2] = 42 + param = distribute_tensor(baseline, parameter_mesh, placements) + state = {'w': param} + model = SimpleNamespace(state_dict=lambda: state) + owner = SimpleNamespace(_delta_device='cpu', model=model, + strategy=SimpleNamespace(unwrap_model=lambda model: model)) + exporter = FSDPDeltaTrainer() + exporter.__dict__.update(vars(owner)) + exporter.prime_delta_snapshots() + changed = distribute_tensor(target, parameter_mesh, placements) + param.to_local().copy_(changed.to_local()) + restored = baseline.clone().flatten() + + def apply(name, dtype, shape, size, indices, values): + if indices is not None: + restored[indices.long()] = values + + queue = _GatherQueue(32, 30, rank == 0, apply) + original = DTensor.full_tensor + DTensor.full_tensor = lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError('steady export must not all-gather full tensors')) + try: + entries, _ = exporter.get_per_tensor_param_delta_shard() + for slots, dtype, counts, indices, values, group in entries: + queue.put(group, slots, dtype, counts, indices, values) + queue.flush_all() + finally: + DTensor.full_tensor = original + if rank == 0: + assert torch.equal(restored.view(torch.int16), target.flatten().view(torch.int16)) + entries, _ = exporter.get_per_tensor_param_delta_shard() + assert all(not entry[2].sum().item() for entry in entries) + # FSDP can produce an empty shard when there are more ranks than rows. + param = distribute_tensor(torch.ones(1, 3), mesh, [Shard(0)]) + state = {'w': param} + exporter = FSDPDeltaTrainer() + exporter.__dict__.update(vars(owner)) + exporter.prime_delta_snapshots() + param.to_local().add_(1) + entries, _ = exporter.get_per_tensor_param_delta_shard() + entry = next(entries) + assert entry[2].tolist() == [3 if rank == 0 else 0] + finally: + dist.destroy_process_group() + + +def test_fsdp_export_preserves_bytes_with_uneven_column_and_replica_shards(tmp_path): + mp.spawn(exercise_fsdp_export, args=(f'file://{tmp_path}/rendezvous',), nprocs=2, join=True) diff --git a/tests/checkpoint_engine/test_delta_vllm.py b/tests/checkpoint_engine/test_delta_vllm.py new file mode 100644 index 000000000..aee3f7315 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_vllm.py @@ -0,0 +1,508 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""vLLM weight-transfer adapter: flush decoding, native patch loading, tied weights and Mamba sparse loaders.""" +import importlib.util +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +from twinkle.checkpoint_engine.delta.protocol.wire import checksum + + +@pytest.fixture +def adapter(monkeypatch): + @dataclass + class CheckpointWeightPatch: + name: str + shape: tuple + dtype: torch.dtype + values: torch.Tensor + indices: torch.Tensor | None = None + + class WeightTransferEngine: + def __init__(self, config, vllm_config, device, model): + self.model = model + self.vllm_config = vllm_config + self.device = device + self.parallel_config = vllm_config.parallel_config + self.model_config = vllm_config.model_config + + modules = { + 'vllm': {'__version__': 'test'}, + 'vllm.config.weight_transfer': {'WeightTransferConfig': type('WeightTransferConfig', (), {})}, + 'vllm.distributed.weight_transfer.base': { + 'WeightTransferEngine': WeightTransferEngine, + 'WeightTransferInitInfo': type('WeightTransferInitInfo', (), {}), + 'WeightTransferUpdateInfo': type('WeightTransferUpdateInfo', (), {}), + }, + 'vllm.model_executor.model_loader.checkpoint_weight_patch': { + 'CheckpointWeightPatch': CheckpointWeightPatch, + 'load_checkpoint_weight_patches': Mock(), + }, + 'vllm.model_executor.layers.vocab_parallel_embedding': { + 'VocabParallelEmbedding': type('VocabParallelEmbedding', (), {}), + }, + 'vllm.model_executor.model_loader.reload': { + 'initialize_layerwise_reload': Mock(), + 'finalize_layerwise_reload': Mock(), + }, + } + for name, fields in modules.items(): + module = ModuleType(name) + module.__dict__.update(fields) + monkeypatch.setitem(sys.modules, name, module) + path = Path(__file__).parents[2] / 'src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py' + spec = importlib.util.spec_from_file_location('_twinkle_delta_adapter_test', path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +def named_flush(encoding='indices'): + values = torch.tensor([-0., 2.], dtype=torch.bfloat16) + positions = torch.tensor([0, 3], dtype=torch.int32).view(torch.uint8) + if encoding == 'dense': + positions = torch.empty(0, dtype=torch.uint8) + params = [dict(name='w', dtype='bfloat16', shape=[4 if encoding == 'indices' else 2], + pos_start=0, pos_end=positions.numel(), pos_width=4, val_start=0, val_end=2)] + spec = dict(encoding=encoding, params=params, checksum=checksum(positions, values)) + spec_tensor = torch.tensor(list(json.dumps(spec).encode()), dtype=torch.uint8) + return [('__delta_spec__', spec_tensor), ('__positions__', positions), ('__values__', values)] + + +def transfer_engine(adapter): + config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size=1, pipeline_parallel_size=1), + model_config=SimpleNamespace(dtype=torch.bfloat16)) + model = SimpleNamespace(named_modules=lambda **_: [], named_parameters=lambda: []) + return adapter.TwinkleDeltaIPCWeightTransferEngine(None, config, torch.device('cuda'), model) + + +def test_decoder_consumes_verl_flush_without_custom_envelope(adapter): + encoding, patches = adapter.decode_delta_payload(named_flush()) + assert (encoding, patches[0].name, patches[0].shape, patches[0].indices.tolist(), + patches[0].values.view(torch.int16).tolist()) == ('indices', 'w', (4,), [0, 3], [-32768, 16384]) + + +def test_corrupt_flush_is_rejected_before_loading(adapter): + flush = named_flush() + flush[-1][1][0] = 42 + with pytest.raises(RuntimeError, match='checksum'): + adapter.decode_delta_payload(flush) + + +def test_dense_update_runs_one_reload_lifecycle(adapter): + engine = transfer_engine(adapter) + engine._receive_payload = Mock(return_value=named_flush('dense')) + engine.start_weight_update() + update = adapter.TwinkleDeltaIPCUpdateInfo(zmq_handle='ipc:///tmp/test.sock') + engine.receive_weights(update) + engine.receive_weights(update) + engine.finish_weight_update() + reload = sys.modules['vllm.model_executor.model_loader.reload'] + assert (reload.initialize_layerwise_reload.call_count, reload.finalize_layerwise_reload.call_count) == (1, 1) + + +def test_sparse_update_uses_native_patch_loader_without_reload(adapter): + engine = transfer_engine(adapter) + engine._receive_payload = Mock(return_value=named_flush()) + engine.start_weight_update() + engine.receive_weights(adapter.TwinkleDeltaIPCUpdateInfo(zmq_handle='ipc:///tmp/test.sock')) + engine.finish_weight_update() + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + assert loader.load_checkpoint_weight_patches.call_args.kwargs == {'validate_unique_indices': False} + reload = sys.modules['vllm.model_executor.model_loader.reload'] + reload.initialize_layerwise_reload.assert_not_called() + reload.finalize_layerwise_reload.assert_not_called() + + +def ascend_gdn_model(dim=6, width=4): + """Mimic vllm-ascend ops/gdn.py: conv1d.weight [D, 1, W] plus a packed [W, D] copy.""" + conv1d = torch.nn.Module() + conv1d.weight = torch.nn.Parameter(torch.randn(dim, 1, width).to(torch.bfloat16), requires_grad=False) + packed = conv1d.weight.view(dim, width).transpose(0, 1).contiguous() + conv1d.register_parameter('ascend_conv1d_weight', torch.nn.Parameter(packed, requires_grad=False)) + model = torch.nn.Module() + model.layers = torch.nn.ModuleList([torch.nn.Module()]) + model.layers[0].conv1d = conv1d + return model, conv1d + + +def expected_packed(conv1d): + weight = conv1d.weight + return weight.view(weight.size(0), weight.size(2)).transpose(0, 1).contiguous() + + +@pytest.mark.parametrize('encoding, refreshed', [('indices', True), ('dense', False)]) +def test_sparse_update_refreshes_ascend_packed_conv_weight(adapter, encoding, refreshed): + model, conv1d = ascend_gdn_model() + stale = conv1d.ascend_conv1d_weight.detach().clone() + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + + def write_source(*_, **__): + # The checkpoint patch API only writes the source parameter. + flat = conv1d.weight.data.view(-1) + flat[[1, 9]] = -flat[[1, 9]] + 1 + + loader.load_checkpoint_weight_patches.side_effect = write_source + config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size=1, pipeline_parallel_size=1), + model_config=SimpleNamespace(dtype=torch.bfloat16)) + engine = adapter.TwinkleDeltaIPCWeightTransferEngine(None, config, torch.device('cuda'), model) + engine._receive_payload = Mock(return_value=named_flush(encoding)) + engine.start_weight_update() + engine.receive_weights(adapter.TwinkleDeltaIPCUpdateInfo(zmq_handle='ipc:///tmp/test.sock')) + assert not torch.equal(stale, expected_packed(conv1d)) # the source really changed + packed_ptr = conv1d.ascend_conv1d_weight.data_ptr() + engine.finish_weight_update() + assert conv1d.ascend_conv1d_weight.data_ptr() == packed_ptr # refreshed in place + if refreshed: + assert torch.equal(conv1d.ascend_conv1d_weight.view(torch.int16), expected_packed(conv1d).view(torch.int16)) + else: # dense: finalize_layerwise_reload (mocked here) owns the repack + assert torch.equal(conv1d.ascend_conv1d_weight, stale) + + +def test_refresh_derived_weights_rejects_unknown_packed_layout(adapter): + from twinkle.sampler.vllm_sampler.delta.weight_update_utils import refresh_derived_weights + model, conv1d = ascend_gdn_model() + conv1d.ascend_conv1d_weight = torch.nn.Parameter(torch.zeros(6, 4, dtype=torch.bfloat16), requires_grad=False) + with pytest.raises(RuntimeError, match='unsupported packed conv layout'): + refresh_derived_weights(model) + + +def test_refresh_derived_weights_is_noop_without_packed_weights(adapter): + from twinkle.sampler.vllm_sampler.delta.weight_update_utils import refresh_derived_weights + assert refresh_derived_weights(torch.nn.Linear(2, 2)) == 0 + + +def test_partial_update_failure_poisoned_across_sessions(adapter): + engine = transfer_engine(adapter) + engine._receive_payload = Mock(return_value=named_flush()) + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + loader.load_checkpoint_weight_patches.side_effect = ValueError('partial write') + engine.start_weight_update() + with pytest.raises(ValueError, match='partial write'): + engine.receive_weights(adapter.TwinkleDeltaIPCUpdateInfo(zmq_handle='ipc:///tmp/test.sock')) + with pytest.raises(RuntimeError, match='restart the full job'): + engine.start_weight_update() + + +def test_worker_requires_native_delta_backend(adapter, monkeypatch): + from twinkle.sampler.vllm_sampler.vllm_worker_extension import TwinkleWorkerExtension + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta.transfer_engine', adapter) + worker = SimpleNamespace(weight_transfer_engine=None) + with pytest.raises(RuntimeError, match='weight_transfer_config'): + TwinkleWorkerExtension.validate_delta_worker(worker) + + +def test_worker_allows_dense_model_and_tp_beyond_previous_allowlist(adapter, monkeypatch): + from twinkle.sampler.vllm_sampler.vllm_worker_extension import TwinkleWorkerExtension + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta.transfer_engine', adapter) + engine = transfer_engine(adapter) + engine.model_config.hf_config = SimpleNamespace(model_type='llama') + worker = SimpleNamespace(weight_transfer_engine=engine, + model_runner=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=4))) + assert TwinkleWorkerExtension.validate_delta_worker(worker) == {'tp': 4, 'device': 'cuda'} + + +class Embedding: + def __init__(self, weight): + self.weight = weight + + def named_parameters(self, remove_duplicate=False): + return [('weight', self.weight)] + + +class CheckpointMapper: + def apply_list(self, names): + return [f'language_model.{name}' for name in names] + + +def load_source(monkeypatch, filename): + path = Path(__file__).parents[2] / 'src/twinkle/sampler/vllm_sampler' / filename + spec = importlib.util.spec_from_file_location(f'_delta_test_{path.stem}', path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def tied_helpers(monkeypatch): + module = ModuleType('vllm.model_executor.layers.vocab_parallel_embedding') + module.VocabParallelEmbedding = Embedding + monkeypatch.setitem(sys.modules, module.__name__, module) + return load_source(monkeypatch, 'delta/weight_update_utils.py') + + +@pytest.fixture(params=[False, True], ids=['language-model', 'qwen-multimodal']) +def tied_model(request): + shared_weight = [1., 2., 3.] + prefix = 'language_model.' if request.param else '' + modules = [(f'{prefix}model.embed_tokens', Embedding(shared_weight)), + (f'{prefix}lm_head', Embedding(shared_weight))] + model = SimpleNamespace(named_modules=lambda **_: modules, named_parameters=lambda: [], shared_weight=shared_weight) + if request.param: + model.hf_to_vllm_mapper = CheckpointMapper() + return model + + +def test_separate_alias_bucket_is_dropped(tied_helpers, tied_model): + alias = SimpleNamespace(name='lm_head.weight', indices=[1], values=[7.]) + assert tied_helpers.drop_tied_alias_patches(tied_model, [alias]) == [] + + +def test_canonical_patch_preserves_checkpoint_name_and_payload(tied_helpers, tied_model): + canonical = SimpleNamespace(name='model.embed_tokens.weight', indices=[1], values=[7.]) + alias = SimpleNamespace(name='lm_head.weight', indices=[1], values=[7.]) + assert tied_helpers.drop_tied_alias_patches(tied_model, [canonical, alias]) == [canonical] + + +def test_untied_head_is_preserved(tied_helpers): + modules = [('model.embed_tokens', Embedding(object())), ('lm_head', Embedding(object()))] + model = SimpleNamespace(named_modules=lambda **_: modules) + alias = SimpleNamespace(name='lm_head.weight') + assert tied_helpers.drop_tied_alias_patches(model, [alias]) == [alias] + + +def test_shared_non_embedding_parameters_are_preserved(tied_helpers): + shared = object() + submodule = SimpleNamespace(named_parameters=lambda **_: [('weight', shared)]) + modules = [('first', submodule), ('second', submodule)] + model = SimpleNamespace(named_modules=lambda **_: modules) + patch = SimpleNamespace(name='second.weight') + assert tied_helpers.drop_tied_alias_patches(model, [patch]) == [patch] + + +def test_mapper_dropped_name_is_left_to_vllm(tied_helpers, tied_model): + tied_model.hf_to_vllm_mapper = SimpleNamespace(apply_list=lambda _: []) + patch = SimpleNamespace(name='unknown.weight') + assert tied_helpers.drop_tied_alias_patches(tied_model, [patch]) == [patch] + + +@pytest.fixture +def tied_adapter(monkeypatch, tied_helpers): + class WeightTransferEngine: + def __init__(self, config, vllm_config, device, model): + self.model = model + self.model_config = vllm_config.model_config + + modules = { + 'torch': {}, + 'vllm': {}, + 'vllm.config.weight_transfer': {'WeightTransferConfig': object}, + 'vllm.distributed.weight_transfer.base': { + 'WeightTransferEngine': WeightTransferEngine, + 'WeightTransferInitInfo': object, + 'WeightTransferUpdateInfo': object, + }, + 'vllm.model_executor.model_loader.checkpoint_weight_patch': { + 'CheckpointWeightPatch': object, + 'load_checkpoint_weight_patches': Mock(), + }, + 'vllm.model_executor.model_loader.reload': { + 'initialize_layerwise_reload': Mock(), + 'finalize_layerwise_reload': Mock(), + }, + } + for name, fields in modules.items(): + module = ModuleType(name) + module.__dict__.update(fields) + monkeypatch.setitem(sys.modules, name, module) + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta.weight_update_utils', tied_helpers) + return load_source(monkeypatch, 'delta/transfer_engine.py') + + +@pytest.mark.parametrize('encoding', ['indices', 'dense']) +def test_transfer_applies_shared_weight_across_separate_flushes(tied_adapter, tied_model, monkeypatch, encoding): + canonical = SimpleNamespace(name='model.embed_tokens.weight', indices=[1], values=[7.]) + alias = SimpleNamespace(name='lm_head.weight', indices=[1], values=[7.]) + config = SimpleNamespace(model_config=object()) + engine = tied_adapter.TwinkleDeltaIPCWeightTransferEngine(None, config, None, tied_model) + engine._receive_payload = Mock(side_effect=[[canonical], [alias]]) + monkeypatch.setattr(tied_adapter, 'decode_delta_payload', lambda patches: (encoding, patches)) + + def load_patches(model, patches, **_): + names = {patch.name for patch in patches} + if 'lm_head.weight' in names and 'model.embed_tokens.weight' not in names: + raise ValueError('tied weight is uninitialized') + for patch in patches: + for index, value in zip(patch.indices, patch.values): + model.shared_weight[index] = value + + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + loader.load_checkpoint_weight_patches.side_effect = load_patches + engine.start_weight_update() + update = tied_adapter.TwinkleDeltaIPCUpdateInfo(zmq_handle='ipc:///tmp/tied-test.sock') + engine.receive_weights(update) + engine.receive_weights(update) + engine.finish_weight_update() + assert (tied_model.shared_weight, loader.load_checkpoint_weight_patches.call_args.args[1]) == ([1., 7., 3.], []) + + +@pytest.mark.parametrize('encoding', ['indices', 'dense']) +def test_transfer_prepares_native_loaders_for_sparse_patches(tied_adapter, tied_model, monkeypatch, encoding): + tied_helpers = sys.modules['twinkle.sampler.vllm_sampler.delta.weight_update_utils'] + prepare = Mock() + monkeypatch.setattr(tied_helpers, 'prepare_sparse_weight_loaders', prepare) + engine = tied_adapter.TwinkleDeltaIPCWeightTransferEngine(None, SimpleNamespace(model_config=object()), None, tied_model) + engine._receive_payload = Mock(return_value=[]) + monkeypatch.setattr(tied_adapter, 'decode_delta_payload', lambda _: (encoding, [])) + loader = sys.modules['vllm.model_executor.model_loader.checkpoint_weight_patch'] + + def verify_preparation(*_, **__): + assert prepare.call_count == (1 if encoding == 'indices' else 0) + + loader.load_checkpoint_weight_patches.side_effect = verify_preparation + engine.start_weight_update() + engine.receive_weights(tied_adapter.TwinkleDeltaIPCUpdateInfo(zmq_handle='ipc:///tmp/tied-test.sock')) + engine.finish_weight_update() + + +class Tensor: + def __init__(self, values): + self.values = np.asarray(values, dtype=float) + + @property + def data(self): + return self + + def __getitem__(self, key): + return Tensor(self.values[key]) + + def __setitem__(self, key, source): + self.values[key] = source.values + + def copy_(self, source): + self.values[...] = source.values + + +def native_mamba_loader(shard_spec, tp_size, tp_rank): + """Reproduce the native slice assignment that bypasses Tensor.copy_.""" + def loader(param, loaded_weight): + boundary, loaded_boundary = 0, 0 + for full_dim, extra, duplicate_groups in shard_spec: + shard_size = full_dim // tp_size + rank = 0 if duplicate_groups else tp_rank + loaded_skip = rank * shard_size + loaded_start = loaded_boundary + loaded_skip + take = min(shard_size, full_dim - extra - loaded_skip) + param.data[boundary:boundary + take, ...] = loaded_weight[loaded_start:loaded_start + take] + boundary += shard_size + loaded_boundary += full_dim - extra + + loader.__module__ = 'vllm.model_executor.layers.mamba.mamba_mixer2' + loader.__qualname__ = 'mamba_v2_sharded_weight_loader..loader' + return loader + + +@pytest.fixture +def mamba_helpers(): + path = Path(__file__).parents[2] / 'src/twinkle/sampler/vllm_sampler/delta/weight_update_utils.py' + spec = importlib.util.spec_from_file_location('_delta_mamba_helpers', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def masked_copy(monkeypatch): + def copy_non_nan(destination, source): + destination.values[...] = np.where(np.isnan(source.values), destination.values, source.values) + + monkeypatch.setattr(Tensor, 'copy_', copy_non_nan) + + +def model_for(parameter): + return SimpleNamespace(named_parameters=lambda: [('conv1d.weight', parameter)]) + + +@pytest.mark.parametrize('tp_rank', [0, 1]) +@pytest.mark.parametrize('shard_spec', [ + [(4, 0, False), (4, 0, False), (4, 0, False)], + [(4, 0, False), (4, 0, True), (4, 1, False)], +]) +def test_sparse_updates_match_dense_tp_load(mamba_helpers, masked_copy, tp_rank, shard_spec): + checkpoint_size = sum(full_dim - extra for full_dim, extra, _ in shard_spec) + parameter_size = sum(full_dim // 2 for full_dim, _, _ in shard_spec) + checkpoint = Tensor(np.arange(checkpoint_size * 2).reshape(checkpoint_size, 1, 2)) + native = native_mamba_loader(shard_spec, 2, tp_rank) + parameter = Tensor(np.full((parameter_size, 1, 2), -1.)) + native(parameter, checkpoint) + expected = Tensor(parameter.values.copy()) + changed = np.arange(checkpoint_size) % 2 == 0 + updated = Tensor(checkpoint.values.copy()) + updated.values[changed] += 100. + native(expected, updated) + patch = Tensor(np.full_like(checkpoint.values, np.nan)) + patch.values[changed] = updated.values[changed] + parameter.weight_loader = native + mamba_helpers.prepare_sparse_weight_loaders(model_for(parameter)) + parameter.weight_loader(parameter, patch) + np.testing.assert_array_equal(parameter.values, expected.values) + + +def test_dense_updates_keep_native_geometry(mamba_helpers): + native = native_mamba_loader([(4, 0, False), (4, 0, True), (4, 1, False)], 2, 1) + checkpoint = Tensor(np.arange(22).reshape(11, 1, 2)) + expected = Tensor(np.full((6, 1, 2), -1.)) + native(expected, checkpoint) + parameter = Tensor(np.full((6, 1, 2), -1.)) + parameter.weight_loader = native + mamba_helpers.prepare_sparse_weight_loaders(model_for(parameter)) + parameter.weight_loader(parameter, checkpoint) + np.testing.assert_array_equal(parameter.values, expected.values) + + +def test_loader_adaptation_is_idempotent(mamba_helpers): + parameter = Tensor([0.]) + parameter.weight_loader = native_mamba_loader([(1, 0, False)], 1, 0) + model = model_for(parameter) + mamba_helpers.prepare_sparse_weight_loaders(model) + adapted = parameter.weight_loader + mamba_helpers.prepare_sparse_weight_loaders(model) + assert parameter.weight_loader is adapted + + +def test_unrelated_loaders_are_unchanged(mamba_helpers): + original = native_mamba_loader([(1, 0, False)], 1, 0) + original.__module__ = 'other_backend' + parameter = Tensor([0.]) + parameter.weight_loader = original + mamba_helpers.prepare_sparse_weight_loaders(model_for(parameter)) + assert parameter.weight_loader is original + + +def test_unknown_native_geometry_is_rejected(mamba_helpers): + def loader(param, loaded_weight): + pass + + loader.__module__ = 'vllm.model_executor.layers.mamba.mamba_mixer2' + loader.__qualname__ = 'mamba_v2_sharded_weight_loader..loader' + parameter = Tensor([0.]) + parameter.weight_loader = loader + with pytest.raises(RuntimeError, match='cannot safely apply sparse'): + mamba_helpers.prepare_sparse_weight_loaders(model_for(parameter)) + + +def test_sparse_conv_patch_preserves_torch_runtime_values(mamba_helpers, monkeypatch): + torch = pytest.importorskip('torch') + parameter = torch.nn.Parameter(torch.arange(12, dtype=torch.float32).reshape(6, 1, 2)) + parameter.weight_loader = native_mamba_loader([(2, 0, False)] * 3, 1, 0) + patch = torch.full_like(parameter, float('nan')) + patch[2, 0, 1] = 99. + expected = parameter.detach().clone() + expected[2, 0, 1] = 99. + + def copy_non_nan(destination, source): + return torch.where(torch.isnan(source), destination, source, out=destination) + + monkeypatch.setattr(torch.Tensor, 'copy_', copy_non_nan) + mamba_helpers.prepare_sparse_weight_loaders(model_for(parameter)) + parameter.weight_loader(parameter, patch) + assert torch.equal(parameter, expected) diff --git a/tests/checkpoint_engine/test_zmq_pubsub.py b/tests/checkpoint_engine/test_zmq_pubsub.py new file mode 100644 index 000000000..1313cea9e --- /dev/null +++ b/tests/checkpoint_engine/test_zmq_pubsub.py @@ -0,0 +1,41 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""zmq PUB/SUB manifest-channel helpers.""" +import time + +import pytest +import zmq + +from twinkle.checkpoint_engine.zmq_pubsub import connect_subscriber, format_tcp_address, start_publisher +from twinkle.utils.network import find_free_port + + +def test_format_tcp_address(): + assert format_tcp_address('10.0.0.1', 5555) == ('tcp://10.0.0.1:5555', False) + assert format_tcp_address('::1', 5555) == ('tcp://[::1]:5555', True) + + +def test_loopback_pub_sub(): + port = find_free_port() + try: + pub = start_publisher('127.0.0.1', port) + except zmq.ZMQError as e: + pytest.skip(f'zmq loopback unavailable: {e}') + sub = connect_subscriber('127.0.0.1', port, 'topic') + try: + sub.setsockopt(zmq.RCVTIMEO, 100) + got = None + deadline = time.time() + 5 + while got is None and time.time() < deadline: # tolerate slow joiner + pub.send_string('topic', flags=zmq.SNDMORE) + pub.send_pyobj({'x': 1}) + try: + sub.recv_string() + got = sub.recv_pyobj() + except zmq.error.Again: + pass + assert got == {'x': 1} + finally: + sub.close(linger=0) + pub.close(linger=0) + sub.context.term() + pub.context.term()