From 2bb393ec368354db6fac037e8f53e47bab10d44b Mon Sep 17 00:00:00 2001 From: kevssim Date: Tue, 13 Oct 2026 14:57:49 +0800 Subject: [PATCH 01/12] wip --- src/twinkle/checkpoint_engine/__init__.py | 3 + .../checkpoint_engine/delta/__init__.py | 18 + .../delta/backends/__init__.py | 5 + .../checkpoint_engine/delta/backends/base.py | 68 +++ .../checkpoint_engine/delta/backends/fsdp.py | 97 +++ src/twinkle/checkpoint_engine/delta/config.py | 59 ++ src/twinkle/checkpoint_engine/delta/engine.py | 563 ++++++++++++++++++ src/twinkle/checkpoint_engine/delta/export.py | 91 +++ src/twinkle/checkpoint_engine/delta/gather.py | 104 ++++ .../checkpoint_engine/delta/rollout.py | 245 ++++++++ src/twinkle/checkpoint_engine/delta/spec.py | 229 +++++++ .../checkpoint_engine/delta/trainer.py | 78 +++ src/twinkle/checkpoint_engine/delta/wire.py | 66 ++ src/twinkle/checkpoint_engine/manager.py | 112 ++++ .../model/transformers/transformers.py | 21 +- .../vllm_sampler/bucketed_weight_transfer.py | 352 +++++++++++ .../vllm_sampler/delta_weight_transfer.py | 337 +++++++++++ .../sampler/vllm_sampler/vllm_engine.py | 66 +- .../sampler/vllm_sampler/vllm_sampler.py | 3 +- .../vllm_sampler/vllm_worker_extension.py | 49 +- .../vllm_sampler/weight_update_utils.py | 105 ++++ tests/checkpoint_engine/test_delta_manager.py | 227 +++++++ tests/checkpoint_engine/test_delta_rollout.py | 334 +++++++++++ tests/checkpoint_engine/test_delta_trainer.py | 421 +++++++++++++ tests/checkpoint_engine/test_delta_vllm.py | 450 ++++++++++++++ 25 files changed, 4086 insertions(+), 17 deletions(-) create mode 100644 src/twinkle/checkpoint_engine/delta/__init__.py create mode 100644 src/twinkle/checkpoint_engine/delta/backends/__init__.py create mode 100644 src/twinkle/checkpoint_engine/delta/backends/base.py create mode 100644 src/twinkle/checkpoint_engine/delta/backends/fsdp.py create mode 100644 src/twinkle/checkpoint_engine/delta/config.py create mode 100644 src/twinkle/checkpoint_engine/delta/engine.py create mode 100644 src/twinkle/checkpoint_engine/delta/export.py create mode 100644 src/twinkle/checkpoint_engine/delta/gather.py create mode 100644 src/twinkle/checkpoint_engine/delta/rollout.py create mode 100644 src/twinkle/checkpoint_engine/delta/spec.py create mode 100644 src/twinkle/checkpoint_engine/delta/trainer.py create mode 100644 src/twinkle/checkpoint_engine/delta/wire.py create mode 100644 src/twinkle/sampler/vllm_sampler/bucketed_weight_transfer.py create mode 100644 src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py create mode 100644 src/twinkle/sampler/vllm_sampler/weight_update_utils.py create mode 100644 tests/checkpoint_engine/test_delta_manager.py create mode 100644 tests/checkpoint_engine/test_delta_rollout.py create mode 100644 tests/checkpoint_engine/test_delta_trainer.py create mode 100644 tests/checkpoint_engine/test_delta_vllm.py diff --git a/src/twinkle/checkpoint_engine/__init__.py b/src/twinkle/checkpoint_engine/__init__.py index 1b2cbd247..da1da3807 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 @@ -32,4 +33,6 @@ 'HCCLCheckpointEngine', 'IPCCheckpointEngine', 'TensorMeta', + 'DeltaSyncConfig', + 'DeltaShardedCheckpointEngine', ] diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py new file mode 100644 index 000000000..d783a4725 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sparse delta weight synchronization. + +Layout: + config.py DeltaSyncConfig and the TypedDicts exchanged between trainer and rollout + engine.py DeltaShardedCheckpointEngine: gather, bucket and broadcast the changed elements + wire.py DeltaFlush / DeltaParam wire structures and checksum + gather.py sparse gather of per-rank deltas to rank 0 + spec.py ShardSpec / BlockPlacement: where a local shard sits in the full tensor + export.py snapshot, diff and HF-coordinate delta export helpers + trainer.py CudaDeltaTrainerMixin: trainer-side Ray lifecycle + rollout.py CudaDeltaReceiver / CudaDeltaSamplerMixin: rollout-side lifecycle + backends/ per-training-backend export hooks (base.py contract, fsdp.py implementation) +""" +from .config import DeltaSyncConfig +from .engine import DeltaShardedCheckpointEngine + +__all__ = ['DeltaSyncConfig', 'DeltaShardedCheckpointEngine'] diff --git a/src/twinkle/checkpoint_engine/delta/backends/__init__.py b/src/twinkle/checkpoint_engine/delta/backends/__init__.py new file mode 100644 index 000000000..ee3a14432 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/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/backends/base.py b/src/twinkle/checkpoint_engine/delta/backends/base.py new file mode 100644 index 000000000..3cffb18fc --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/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/backends/fsdp.py b/src/twinkle/checkpoint_engine/delta/backends/fsdp.py new file mode 100644 index 000000000..d488a3d33 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/backends/fsdp.py @@ -0,0 +1,97 @@ +# 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 + + if Platform.get_platform().__name__ != 'GPU': + raise ValueError('CUDA delta requires CUDA trainer actors') + 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() + if 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/config.py b/src/twinkle/checkpoint_engine/delta/config.py new file mode 100644 index 000000000..f69b84e92 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/config.py @@ -0,0 +1,59 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Configuration for the verl CUDA delta checkpoint engine.""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import TypedDict +from typing_extensions import NotRequired + +from ..nccl_checkpoint_engine import MasterMetadata + +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: MasterMetadata + rollout_workers: list[DeltaWorkerMetadata] + + +class DeltaTrainerMetadata(TypedDict): + device: str + master_metadata: MasterMetadata | None + + +class DeltaWorkerCapabilities(TypedDict): + tp: int + device: str + + +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/engine.py b/src/twinkle/checkpoint_engine/delta/engine.py new file mode 100644 index 000000000..a0f73bef7 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/engine.py @@ -0,0 +1,563 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Delta checkpoint engine based on verl.""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import Generator, Iterator +from dataclasses import dataclass + +import torch +import zmq + +from ..nccl_checkpoint_engine import MasterMetadata, NCCLCheckpointEngine, _pg_broadcast +from twinkle.utils.network import find_free_port +from .gather import gather_slot_entries_to_rank0 +from .wire import DeltaFlush, DeltaParam +from .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 = DeltaShardedCheckpointEngine.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 DeltaShardedCheckpointEngine(NCCLCheckpointEngine): + """Sparse delta weight sync over NCCL, diffed on each rank's local shard. + + Reuses NCCLCheckpointEngine's group/zmq machinery but moves only changed + positions+values: each actor rank keeps a pinned-CPU snapshot of only *its* + shard, byte-diffs the shard, and only the changed ``(position, value)`` pairs + are gathered to rank 0 and streamed to the rollout side -- no rank ever holds + a full-model snapshot. + + ``send_weights`` takes the TRAINING ENGINE and drives the sync itself: the + seed (first sync) streams the backend's full ``get_per_tensor_param()`` + export values-only and pins the diff base (``prime_delta_snapshots``); every + steady sync consumes the backend's HF delta export + ``get_per_tensor_param_delta_shard()`` — per-parameter FINAL-HF-coordinate + entries ``(slots, dtype_str, counts, hf_idx, hf_val, gather_group)``. Naming, + to-HF conversion, diff and snapshot all live on the backend side (see + :mod:`twinkle.checkpoint_engine.delta.export`); this engine only batches, gathers, + buckets and ships, and so serves any backend that can produce HF deltas. + """ + + # 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 + + 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) + + # ---- 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.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + pos_u8 = flush.positions_cpu.to("cuda", non_blocking=True).contiguous().view(torch.uint8) + val_u8 = flush.values_gpu.contiguous().view(torch.uint8) + self._broadcast(pos_u8) + self._broadcast(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.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + val_u8 = values.view(torch.uint8) + self._broadcast(val_u8) + + 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.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + + # ---- 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="cuda") + if dense: + pos = None + self._broadcast(val_u8) + else: + pos = torch.empty(meta["pos_numel"], dtype=torch.uint8, device="cuda") + 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("cuda") + 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) + + 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="cuda") + 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) + # 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 + 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 + logger.info("delta-sharded send v=%s delta flushes=%d (streamed)", global_steps, 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/delta/export.py b/src/twinkle/checkpoint_engine/delta/export.py new file mode 100644 index 000000000..2fe96fae7 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/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/gather.py b/src/twinkle/checkpoint_engine/delta/gather.py new file mode 100644 index 000000000..dea7583f6 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/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/rollout.py b/src/twinkle/checkpoint_engine/delta/rollout.py new file mode 100644 index 000000000..0cf799b1c --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/rollout.py @@ -0,0 +1,245 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rollout-side Ray lifecycle and receivers for the CUDA delta checkpoint engine.""" +from __future__ import annotations + +import asyncio +import os +from collections.abc import Iterator +from typing import Any, TYPE_CHECKING + +import torch + +from twinkle import remote_function +from .config import (DeltaReceiverPlacement, DeltaSyncConfig, DeltaTransportEndpoint, + DeltaWorkerCapabilities, DeltaWorkerMetadata) +from .engine import DeltaShardedCheckpointEngine + +if TYPE_CHECKING: + from ray.actor import ActorHandle + +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.bucketed_weight_transfer 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() + + +class CudaDeltaSamplerMixin: + _delta_config: DeltaSyncConfig + _delta_job_failed: bool = False + _delta_weight_transfer_initialized: bool = False + _delta_receivers: list[ActorHandle] + + 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_cuda_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_cuda_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_cuda_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') + return results[0] + + @remote_function(dispatch='all', collect='flatten', lazy_collect=False) + def cuda_delta_sampler_metadata(self) -> list[DeltaWorkerMetadata]: + return self._run_in_loop(self._delta_rpc('cuda_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_cuda_delta_sampler_transport(self, endpoint: DeltaTransportEndpoint) -> None: + import ray + + metadata = self._run_in_loop(self._delta_rpc('cuda_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) + + @remote_function(dispatch='all', lazy_collect=True) + def receive_cuda_delta_weights(self) -> dict[str, bool]: + 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_cuda_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)) + 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/spec.py b/src/twinkle/checkpoint_engine/delta/spec.py new file mode 100644 index 000000000..ba365e230 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/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/trainer.py b/src/twinkle/checkpoint_engine/delta/trainer.py new file mode 100644 index 000000000..69d849302 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/trainer.py @@ -0,0 +1,78 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Trainer-side Ray lifecycle for the CUDA delta checkpoint engine.""" +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 remote_function +from ..nccl_checkpoint_engine import MasterMetadata +from .config import DeltaSyncConfig, DeltaTrainerMetadata, DeltaTransportEndpoint +from .engine import DeltaShardedCheckpointEngine + + +class CudaDeltaTrainerMixin: + """Lifecycle RPCs only. Export hooks come from a DeltaTrainerBackend subclass on the same model.""" + _delta_engine: DeltaShardedCheckpointEngine | None = None + _delta_metadata_group: dist.ProcessGroup | None = None + _delta_master_metadata: MasterMetadata | None = None + + def _validate_delta_ranks(self) -> None: + error = None + try: + self.validate_delta_backend() + except Exception as validation_error: + error = str(validation_error) + 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: + self._delta_metadata_group = dist.new_group( + backend='gloo', timeout=timedelta(seconds=config.timeout_seconds)) + try: + self._validate_delta_ranks() + except Exception: + dist.destroy_process_group(self._delta_metadata_group) + self._delta_metadata_group = None + raise + self._delta_device = 'cuda' + self._delta_pin_snapshots = config.delta_pin_snapshots + self._delta_engine = DeltaShardedCheckpointEngine( + bucket_size=config.bucket_bytes, batch_gather=config.gather_batch_size) + self._delta_engine.is_master = dist.get_rank() == 0 + self._delta_master_metadata = self._delta_engine.prepare() + + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def prepare_cuda_delta_sync(self, config: DeltaSyncConfig) -> DeltaTrainerMetadata: + self._lazy_wrap_model() + if self._delta_engine is None: + self._initialize_delta_engine(config) + return {'device': 'cuda', 'master_metadata': self._delta_master_metadata} + + @remote_function(dispatch='all', lazy_collect=True) + def init_cuda_delta_transport(self, endpoint: DeltaTransportEndpoint) -> None: + rank = 0 if dist.get_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_cuda_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_cuda_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/wire.py b/src/twinkle/checkpoint_engine/delta/wire.py new file mode 100644 index 000000000..1caeb991c --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/wire.py @@ -0,0 +1,66 @@ +# 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 checksum(positions: torch.Tensor, values: torch.Tensor) -> int: + """Wire-corruption check; sender computes pre-flush, receiver post-recv. + + Uses ``torch.hash_tensor`` (XOR-reduce over uint64 bitcast); one reduction + plus one ``.item()`` sync per argument. + """ + p = int(torch.hash_tensor(positions).item()) if positions.numel() else 0 + v = int(torch.hash_tensor(values).item()) if values.numel() else 0 + return p ^ (v << 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/manager.py b/src/twinkle/checkpoint_engine/manager.py index b5df00111..b947b6910 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -1,9 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. # Adapted from https://github.com/volcengine/verl/blob/main/verl/checkpoint_engine/base.py +from contextlib import contextmanager +from time import perf_counter from typing import List, Literal, Optional from twinkle import Platform, get_logger from .base import CheckpointEngine +from .delta import DeltaSyncConfig from .mixin import CheckpointEngineMixin logger = get_logger() @@ -24,6 +27,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) @@ -72,12 +78,26 @@ def __init__( sampler: 'CheckpointEngineMixin', platform: str = 'GPU', mode: CheckpointEngineMode = 'auto', + sync_backend: Literal['full', 'delta_sharded'] = 'full', + delta_config: Optional[DeltaSyncConfig] = None, ) -> None: self.model = model self.sampler = sampler 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() + if sync_backend == 'delta_sharded' and Platform.get_platform(platform).__name__ != 'GPU': + raise ValueError('delta_sharded currently supports CUDA only') + self._delta_replica_count = None + self._delta_round = 0 + self._delta_failed = False + self.last_delta_metrics = None if self.mode == 'colocate': # Each side builds its own engine inside its worker, so both have to be told which one. @@ -160,6 +180,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 @@ -202,6 +226,94 @@ 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_cuda_delta_sampler(self.delta_config) + replicas = capabilities if isinstance(capabilities, list) else [capabilities] + workers = self.sampler.cuda_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_cuda_delta_sync(self.delta_config) + return dict(metadata, world_size=len(workers) + 1, rollout_workers=workers), 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_cuda_delta_sync() + if seed: + trainer_init = self.model.init_cuda_delta_transport(endpoint) + sampler_init = self.sampler.init_cuda_delta_sampler_transport(endpoint) + trainer_init() + sampler_init() + self._delta_replica_count = replica_count + sender = self.model.send_cuda_delta_weights() + receiver = self.sampler.receive_cuda_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_cuda_delta_sampler, self.model.close_cuda_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/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 7d807fb2a..2cff791c6 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -28,6 +28,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.backends import FSDP2DeltaBackend +from twinkle.checkpoint_engine.delta.trainer import CudaDeltaTrainerMixin from twinkle.data_format import InputFeature, ModelOutput, SampleResponse, SamplingParams, Trajectory from twinkle.hub import HubOperation from twinkle.infra import collect_dp_and_flatten, collect_tensor_dict @@ -284,7 +286,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, CudaDeltaTrainerMixin, + FSDP2DeltaBackend): """The transformers model wrapper. Args: @@ -2187,6 +2190,12 @@ def _get_weight_generator( # Get state dict from unwrapped model model = self.strategy.unwrap_model(self.model) + def _checkpoint_tensor(tensor): + 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 _normalize(name: str, keep_base_layer: bool) -> str: name = name.replace('base_model.model.', '') if not keep_base_layer: @@ -2231,7 +2240,7 @@ def weight_generator(): for name, tensor in model.state_dict().items(): if _is_lora_key(name): continue - tensor = Torch.to_local_tensor(tensor) + tensor = _checkpoint_tensor(tensor) name = _normalize(name, keep_base_layer=False) names.append(name) yield name, tensor @@ -2247,7 +2256,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(name, keep_base_layer=True) names.append(name) yield name, tensor @@ -2258,15 +2267,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.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): continue - tensor = Torch.to_local_tensor(tensor) + tensor = _checkpoint_tensor(tensor) name = _normalize(name, keep_base_layer=False) if add_base_layer: name = _add_base_layer_suffix(name) diff --git a/src/twinkle/sampler/vllm_sampler/bucketed_weight_transfer.py b/src/twinkle/sampler/vllm_sampler/bucketed_weight_transfer.py new file mode 100644 index 000000000..42e7ea5f3 --- /dev/null +++ b/src/twinkle/sampler/vllm_sampler/bucketed_weight_transfer.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_weight_transfer.py b/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py new file mode 100644 index 000000000..ce7761d4b --- /dev/null +++ b/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py @@ -0,0 +1,337 @@ +# 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 +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" +_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.wire import checksum + + CheckpointWeightPatch, _ = _checkpoint_patch_api() + 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[CheckpointWeightPatch] = [] + 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_indices = 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_indices = pos_bytes.view(torch.int32) + + patches.append( + CheckpointWeightPatch( + name=name, + shape=tuple(param_spec["shape"]), + dtype=dtype, + values=patch_values, + indices=patch_indices, + ) + ) + return encoding, 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 + + 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 _validate_configuration(self) -> None: + require_vllm_delta_support() + if self.device.type != 'cuda': + raise NotImplementedError("Twinkle delta weight transfer requires CUDA") + 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.bucketed_weight_transfer 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 receive_weights(self, update_info: TwinkleDeltaIPCUpdateInfo) -> None: + from twinkle.sampler.vllm_sampler.weight_update_utils import ( + drop_tied_alias_patches, + prepare_sparse_weight_loaders, + ) + + assert update_info.zmq_handle is not None + try: + payload = self._receive_payload(zmq_handle=update_info.zmq_handle) + 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 GPU. + load_checkpoint_weight_patches( + self.model, + patches, + validate_unique_indices=False, + ) + 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 + # Sparse updates write to initialized runtime tensors, so there is + # nothing to rebuild after the last flush. + 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 + diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index a81057b33..fc8f0435c 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,21 @@ 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_weight_transfer import register_twinkle_delta_weight_transfer_engine, require_vllm_delta_support + require_vllm_delta_support() + register_twinkle_delta_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 == 'twinkle_delta_ipc' 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 +505,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 +516,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 +649,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 +659,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 +706,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 +716,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 d28d823b1..e4a2e95a7 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 import CudaDeltaSamplerMixin from twinkle.data_format import (InputFeature, PoolingParams, PoolingResponse, SampledSequence, SampleResponse, SamplingParams, Trajectory) from twinkle.hub import HubOperation @@ -76,7 +77,7 @@ def _load_vllm_image(entry: Any) -> Any: @remote_class(max_concurrency=_MAX_CONCURRENCY) -class vLLMSampler(Sampler, CheckpointEngineMixin, GenerationSubmissionMixin, PartialRolloutMixin): +class vLLMSampler(Sampler, CheckpointEngineMixin, GenerationSubmissionMixin, PartialRolloutMixin, CudaDeltaSamplerMixin): """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..66250db98 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,42 @@ 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_weight_transfer import register_twinkle_delta_weight_transfer_engine + register_twinkle_delta_weight_transfer_engine() return super().__new__(cls) + def validate_cuda_delta_worker(self) -> 'DeltaWorkerCapabilities': + from .delta_weight_transfer import TwinkleDeltaIPCWeightTransferEngine, is_moe_model + engine = getattr(self, 'weight_transfer_engine', None) + if not isinstance(engine, TwinkleDeltaIPCWeightTransferEngine): + raise RuntimeError("CUDA delta requires engine_args['weight_transfer_config']=" + "{'backend': 'twinkle_delta_ipc'} when constructing the sampler") + 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': 'cuda'} + + 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 cuda_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 update_twinkle_delta_weights(self): + self.update_weights({'zmq_handle': self.delta_ipc_endpoint()}) + def monkey_patch_model(self): from twinkle.patch.vllm_moe_loader import VLLMMoEWeights VLLMMoEWeights()(self.model_runner.model) diff --git a/src/twinkle/sampler/vllm_sampler/weight_update_utils.py b/src/twinkle/sampler/vllm_sampler/weight_update_utils.py new file mode 100644 index 000000000..4469ec401 --- /dev/null +++ b/src/twinkle/sampler/vllm_sampler/weight_update_utils.py @@ -0,0 +1,105 @@ +# 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 diff --git a/tests/checkpoint_engine/test_delta_manager.py b/tests/checkpoint_engine/test_delta_manager.py new file mode 100644 index 000000000..a6db26441 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_manager.py @@ -0,0 +1,227 @@ +# 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_cuda_delta_sync.return_value = dict(device='cuda', master_metadata=object()) + sampler.validate_cuda_delta_sampler.return_value = dict(device='cuda', tp=2) + sampler.cuda_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_cuda_delta_weights.return_value = lambda: {'payload_bytes': 6} + sampler.receive_cuda_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_delta_rejects_npu_deployment(monkeypatch): + from twinkle import Platform + + class NPU: + pass + + monkeypatch.setattr(Platform, 'get_platform', staticmethod(lambda platform=None: NPU)) + model = Mock(_actors=[1, 2]) + sampler = Mock(_actors=[3]) + with pytest.raises(ValueError, match='CUDA only'): + CheckpointEngineManager(model, sampler, platform='NPU', 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_cuda_delta_transport.call_count == 1 + assert manager._delta_round == 2 + assert model.init_cuda_delta_transport.call_args.args[0]['world_size'] == 3 + sampler.validate_cuda_delta_sampler.assert_called_once() + sampler.cuda_delta_sampler_metadata.assert_called_once() + model.prepare_cuda_delta_sync.assert_called_once() + + +@pytest.mark.parametrize('stage', ['send_cuda_delta_weights', 'receive_cuda_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_cuda_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_cuda_delta_weights.side_effect = ValueError('worker write failed') + with pytest.raises(ValueError): + manager.sync_weights() + sampler.receive_cuda_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_cuda_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_cuda_delta_weights.return_value = lambda: {'changed': False} + manager.sync_weights() + sampler.publish_delta_sampler.assert_called_once() + assert sampler.begin_cuda_delta_sync.call_count == 2 + + +def test_multiple_rollout_replicas_share_one_broadcast_group(): + manager, model, sampler = manager_pair() + sampler.validate_cuda_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 + sampler.cuda_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_cuda_delta_weights.return_value = lambda: [{'changed': True}] * 2 + manager.sync_weights() + manager.sync_weights() + endpoint = model.init_cuda_delta_transport.call_args.args[0] + assert endpoint['world_size'] == 5 + assert endpoint['rollout_workers'] == sampler.cuda_delta_sampler_metadata.return_value + model.init_cuda_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_cuda_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 + sampler.cuda_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_cuda_delta_weights.return_value = lambda: [{'changed': True}] * 2 + manager.sync_weights() + sampler.publish_delta_sampler.reset_mock() + sampler.receive_cuda_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_cuda_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 + sampler.cuda_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_cuda_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.cuda_delta_sampler_metadata.return_value + sampler.cuda_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_cuda_delta_transport.assert_not_called() + + +@pytest.mark.parametrize('stage', ['validate_cuda_delta_sampler', 'cuda_delta_sampler_metadata', 'prepare_cuda_delta_sync']) +def test_preflight_failure_can_be_corrected_and_retried(stage): + manager, model, sampler = manager_pair() + owner = model if stage == 'prepare_cuda_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_cuda_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_cuda_delta_sync.return_value = dict(device='cuda') + sampler.validate_cuda_delta_sampler.return_value = {'tp': 2} + sampler.cuda_delta_sampler_metadata.return_value = [ + dict(node_ip='node', device_uuid=str(rank), zmq_handle=str(rank)) for rank in range(2)] + sampler.begin_cuda_delta_sync.side_effect = lambda: events.append('pause-release') + model.init_cuda_delta_transport.side_effect = lambda endpoint: (events.append('init'), lambda: None)[1] + sampler.init_cuda_delta_sampler_transport.return_value = lambda: None + model.send_cuda_delta_weights.side_effect = lambda: (events.append('send'), lambda: {})[1] + sampler.receive_cuda_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_rollout.py b/tests/checkpoint_engine/test_delta_rollout.py new file mode 100644 index 000000000..68e792d0d --- /dev/null +++ b/tests/checkpoint_engine/test_delta_rollout.py @@ -0,0 +1,334 @@ +# 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 import CudaDeltaReceiver, CudaDeltaSamplerMixin +from twinkle.sampler.vllm_sampler.vllm_engine import VLLMEngine +from twinkle.sampler.vllm_sampler.vllm_worker_extension import _cuda_visible_device + + +class LocalCudaSampler(CudaDeltaSamplerMixin): + 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_cuda_delta_sync() + assert sampler.receive_cuda_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_cuda_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_cuda_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_cuda_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_cuda_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 import rollout 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.bucketed_weight_transfer') + 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 import rollout 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 import rollout 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_cuda_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_cuda_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..dc1d185e1 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_trainer.py @@ -0,0 +1,421 @@ +# 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.backends import FSDP2DeltaBackend +from twinkle.checkpoint_engine.delta.backends.fsdp import validate_fsdp2 +from twinkle.checkpoint_engine.delta.engine import DeltaShardedCheckpointEngine, _GatherQueue +from twinkle.checkpoint_engine.delta.export import shard_delta_indices +from twinkle.checkpoint_engine.delta.trainer import CudaDeltaTrainerMixin +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.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.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.spec import ShardSpec + from twinkle.checkpoint_engine.delta.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.spec import ShardSpec + from twinkle.checkpoint_engine.delta.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(CudaDeltaTrainerMixin, 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) + + +@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.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_cuda_delta_sync(config) + second = trainer.prepare_cuda_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 import trainer as delta_trainer + + trainer = FSDPDeltaTrainer() + group = object() + 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..26b22993d --- /dev/null +++ b/tests/checkpoint_engine/test_delta_vllm.py @@ -0,0 +1,450 @@ +# 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.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_weight_transfer.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 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_weight_transfer', adapter) + worker = SimpleNamespace(weight_transfer_engine=None) + with pytest.raises(RuntimeError, match='weight_transfer_config'): + TwinkleWorkerExtension.validate_cuda_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_weight_transfer', 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_cuda_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, '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.weight_update_utils', tied_helpers) + return load_source(monkeypatch, 'delta_weight_transfer.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.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/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) From 788ac698ce0afe312b5698651bed81e29cfb5d02 Mon Sep 17 00:00:00 2001 From: kevssim Date: Tue, 13 Oct 2026 17:32:57 +0800 Subject: [PATCH 02/12] wip --- .../checkpoint_engine/delta/backends/fsdp.py | 4 +++- .../checkpoint_engine/delta/trainer.py | 23 ++++++++++++++----- .../vllm_sampler/vllm_worker_extension.py | 15 +++++++----- src/twinkle/utils/framework.py | 13 ++++++++--- tests/checkpoint_engine/test_delta_trainer.py | 17 ++++++++++++++ 5 files changed, 56 insertions(+), 16 deletions(-) diff --git a/src/twinkle/checkpoint_engine/delta/backends/fsdp.py b/src/twinkle/checkpoint_engine/delta/backends/fsdp.py index d488a3d33..b9f580dff 100644 --- a/src/twinkle/checkpoint_engine/delta/backends/fsdp.py +++ b/src/twinkle/checkpoint_engine/delta/backends/fsdp.py @@ -52,7 +52,9 @@ def validate_fsdp2(owner) -> None: 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() - if not any(isinstance(param, DTensor) for param in state.values()): + # 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') diff --git a/src/twinkle/checkpoint_engine/delta/trainer.py b/src/twinkle/checkpoint_engine/delta/trainer.py index 69d849302..9d3b9ba3e 100644 --- a/src/twinkle/checkpoint_engine/delta/trainer.py +++ b/src/twinkle/checkpoint_engine/delta/trainer.py @@ -21,31 +21,42 @@ class CudaDeltaTrainerMixin: _delta_metadata_group: dist.ProcessGroup | None = None _delta_master_metadata: MasterMetadata | 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: - self._delta_metadata_group = dist.new_group( - backend='gloo', timeout=timedelta(seconds=config.timeout_seconds)) + 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: - dist.destroy_process_group(self._delta_metadata_group) - self._delta_metadata_group = None + 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 = 'cuda' self._delta_pin_snapshots = config.delta_pin_snapshots self._delta_engine = DeltaShardedCheckpointEngine( bucket_size=config.bucket_bytes, batch_gather=config.gather_batch_size) - self._delta_engine.is_master = dist.get_rank() == 0 + 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) @@ -57,7 +68,7 @@ def prepare_cuda_delta_sync(self, config: DeltaSyncConfig) -> DeltaTrainerMetada @remote_function(dispatch='all', lazy_collect=True) def init_cuda_delta_transport(self, endpoint: DeltaTransportEndpoint) -> None: - rank = 0 if dist.get_rank() == 0 else -1 + 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) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py index 66250db98..708887097 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py @@ -373,16 +373,19 @@ def _broadcast_obj(obj): weights.append((name, tensor)) del partial_tensors[name] - Torch.synchronize() + # 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) - if is_driver: - socket.send(b'') - - # 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/tests/checkpoint_engine/test_delta_trainer.py b/tests/checkpoint_engine/test_delta_trainer.py index dc1d185e1..d53687a5b 100644 --- a/tests/checkpoint_engine/test_delta_trainer.py +++ b/tests/checkpoint_engine/test_delta_trainer.py @@ -176,6 +176,22 @@ def test_default_two_card_fsdp_passes_parallelism_check(fsdp_owner): 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 import trainer 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) @@ -348,6 +364,7 @@ def test_trainer_validation_failure_releases_metadata_group(monkeypatch): 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) From a0f3299c974bcdca02fa1494e340af131b3d09ca Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 14 Oct 2026 10:17:43 +0800 Subject: [PATCH 03/12] refactor(delta): use platform-neutral names for delta sync RPCs and mixins Rename *_cuda_delta_* methods to *_delta_* and CudaDelta{Trainer,Sampler}Mixin to Delta{Trainer,Sampler}Mixin. CudaDeltaReceiver keeps its name since it is CUDA-specific. No behavior change. --- .../checkpoint_engine/delta/__init__.py | 4 +- .../checkpoint_engine/delta/rollout.py | 20 ++--- .../checkpoint_engine/delta/trainer.py | 10 +-- src/twinkle/checkpoint_engine/manager.py | 18 ++-- .../model/transformers/transformers.py | 4 +- .../sampler/vllm_sampler/vllm_sampler.py | 4 +- .../vllm_sampler/vllm_worker_extension.py | 4 +- tests/checkpoint_engine/test_delta_manager.py | 88 +++++++++---------- tests/checkpoint_engine/test_delta_rollout.py | 20 ++--- tests/checkpoint_engine/test_delta_trainer.py | 8 +- tests/checkpoint_engine/test_delta_vllm.py | 4 +- 11 files changed, 92 insertions(+), 92 deletions(-) diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py index d783a4725..cb81409b3 100644 --- a/src/twinkle/checkpoint_engine/delta/__init__.py +++ b/src/twinkle/checkpoint_engine/delta/__init__.py @@ -8,8 +8,8 @@ gather.py sparse gather of per-rank deltas to rank 0 spec.py ShardSpec / BlockPlacement: where a local shard sits in the full tensor export.py snapshot, diff and HF-coordinate delta export helpers - trainer.py CudaDeltaTrainerMixin: trainer-side Ray lifecycle - rollout.py CudaDeltaReceiver / CudaDeltaSamplerMixin: rollout-side lifecycle + trainer.py DeltaTrainerMixin: trainer-side Ray lifecycle + rollout.py CudaDeltaReceiver / DeltaSamplerMixin: rollout-side lifecycle backends/ per-training-backend export hooks (base.py contract, fsdp.py implementation) """ from .config import DeltaSyncConfig diff --git a/src/twinkle/checkpoint_engine/delta/rollout.py b/src/twinkle/checkpoint_engine/delta/rollout.py index 0cf799b1c..ad57eb7c0 100644 --- a/src/twinkle/checkpoint_engine/delta/rollout.py +++ b/src/twinkle/checkpoint_engine/delta/rollout.py @@ -80,7 +80,7 @@ def close(self) -> None: self.engine.close() -class CudaDeltaSamplerMixin: +class DeltaSamplerMixin: _delta_config: DeltaSyncConfig _delta_job_failed: bool = False _delta_weight_transfer_initialized: bool = False @@ -101,7 +101,7 @@ async def _clear_delta_cache(self) -> None: await self.engine.engine.reset_encoder_cache() @remote_function(dispatch='all', lazy_collect=False) - def begin_cuda_delta_sync(self) -> None: + 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) @@ -128,13 +128,13 @@ def fail_delta_sampler(self) -> None: self._run_in_loop(self.engine.pause_delta_generation(self._delta_config.timeout_seconds)) @remote_function(dispatch='all', lazy_collect=False) - def validate_cuda_delta_sampler(self, config: DeltaSyncConfig) -> DeltaWorkerCapabilities: + 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_cuda_delta_worker')) + 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']: @@ -142,8 +142,8 @@ def validate_cuda_delta_sampler(self, config: DeltaSyncConfig) -> DeltaWorkerCap return results[0] @remote_function(dispatch='all', collect='flatten', lazy_collect=False) - def cuda_delta_sampler_metadata(self) -> list[DeltaWorkerMetadata]: - return self._run_in_loop(self._delta_rpc('cuda_delta_receiver_metadata')) + 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, @@ -174,10 +174,10 @@ def _spawn_delta_receiver( }}).remote(self._delta_config, endpoint, placement) @remote_function(dispatch='all', lazy_collect=True) - def init_cuda_delta_sampler_transport(self, endpoint: DeltaTransportEndpoint) -> None: + def init_delta_sampler_transport(self, endpoint: DeltaTransportEndpoint) -> None: import ray - metadata = self._run_in_loop(self._delta_rpc('cuda_delta_receiver_metadata')) + 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 = [] @@ -204,7 +204,7 @@ async def _send_delta_flush(self) -> None: await asyncio.gather(receivers, return_exceptions=True) @remote_function(dispatch='all', lazy_collect=True) - def receive_cuda_delta_weights(self) -> dict[str, bool]: + def receive_delta_weights(self) -> dict[str, bool]: async def receive(): first = await self._next_delta_flush() if first is None: @@ -231,7 +231,7 @@ async def receive(): return self._run_in_loop(receive()) @remote_function(dispatch='all', lazy_collect=False) - def close_cuda_delta_sampler(self) -> None: + def close_delta_sampler(self) -> None: import ray self.engine._delta_unavailable = True diff --git a/src/twinkle/checkpoint_engine/delta/trainer.py b/src/twinkle/checkpoint_engine/delta/trainer.py index 9d3b9ba3e..8b9eea427 100644 --- a/src/twinkle/checkpoint_engine/delta/trainer.py +++ b/src/twinkle/checkpoint_engine/delta/trainer.py @@ -15,7 +15,7 @@ from .engine import DeltaShardedCheckpointEngine -class CudaDeltaTrainerMixin: +class DeltaTrainerMixin: """Lifecycle RPCs only. Export hooks come from a DeltaTrainerBackend subclass on the same model.""" _delta_engine: DeltaShardedCheckpointEngine | None = None _delta_metadata_group: dist.ProcessGroup | None = None @@ -60,25 +60,25 @@ def _initialize_delta_engine(self, config: DeltaSyncConfig) -> None: self._delta_master_metadata = self._delta_engine.prepare() @remote_function(dispatch='all', collect='first', lazy_collect=False) - def prepare_cuda_delta_sync(self, config: DeltaSyncConfig) -> DeltaTrainerMetadata: + def prepare_delta_sync(self, config: DeltaSyncConfig) -> DeltaTrainerMetadata: self._lazy_wrap_model() if self._delta_engine is None: self._initialize_delta_engine(config) return {'device': 'cuda', 'master_metadata': self._delta_master_metadata} @remote_function(dispatch='all', lazy_collect=True) - def init_cuda_delta_transport(self, endpoint: DeltaTransportEndpoint) -> None: + 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_cuda_delta_weights(self) -> dict[str, Any]: + 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_cuda_delta_trainer(self) -> None: + def close_delta_trainer(self) -> None: if self._delta_engine is not None: self._delta_engine.close() self._delta_engine = None diff --git a/src/twinkle/checkpoint_engine/manager.py b/src/twinkle/checkpoint_engine/manager.py index b947b6910..ec75249e6 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -227,15 +227,15 @@ def sync_weights(self, merge_and_sync=True): logger.info('Base model sync completed, subsequent syncs will be LoRA-only') def _prepare_delta_transport(self): - capabilities = self.sampler.validate_cuda_delta_sampler(self.delta_config) + capabilities = self.sampler.validate_delta_sampler(self.delta_config) replicas = capabilities if isinstance(capabilities, list) else [capabilities] - workers = self.sampler.cuda_delta_sampler_metadata() + 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_cuda_delta_sync(self.delta_config) + metadata = self.model.prepare_delta_sync(self.delta_config) return dict(metadata, world_size=len(workers) + 1, rollout_workers=workers), len(replicas) def _sync_delta_weights(self, merge_and_sync): @@ -248,15 +248,15 @@ def _sync_delta_weights(self, merge_and_sync): if seed: endpoint, replica_count = self._prepare_delta_transport() try: - self.sampler.begin_cuda_delta_sync() + self.sampler.begin_delta_sync() if seed: - trainer_init = self.model.init_cuda_delta_transport(endpoint) - sampler_init = self.sampler.init_cuda_delta_sampler_transport(endpoint) + 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_cuda_delta_weights() - receiver = self.sampler.receive_cuda_delta_weights() + 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] @@ -300,7 +300,7 @@ def _delta_rpc_timeouts(self): def _close_delta_session(self): errors = [] - close_methods = (self.sampler.close_cuda_delta_sampler, self.model.close_cuda_delta_trainer) + close_methods = (self.sampler.close_delta_sampler, self.model.close_delta_trainer) for close in close_methods: try: close() diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 2cff791c6..8fd6530e6 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -29,7 +29,7 @@ from twinkle.checkpoint_engine import CheckpointEngine from twinkle.checkpoint_engine.mixin import CheckpointEngineMixin from twinkle.checkpoint_engine.delta.backends import FSDP2DeltaBackend -from twinkle.checkpoint_engine.delta.trainer import CudaDeltaTrainerMixin +from twinkle.checkpoint_engine.delta.trainer import DeltaTrainerMixin from twinkle.data_format import InputFeature, ModelOutput, SampleResponse, SamplingParams, Trajectory from twinkle.hub import HubOperation from twinkle.infra import collect_dp_and_flatten, collect_tensor_dict @@ -286,7 +286,7 @@ def _read_hf_state_dict(checkpoint_dir: str) -> Dict[str, torch.Tensor]: @remote_class() -class TransformersModel(TrainableModel, PreTrainedModel, CheckpointEngineMixin, CudaDeltaTrainerMixin, +class TransformersModel(TrainableModel, PreTrainedModel, CheckpointEngineMixin, DeltaTrainerMixin, FSDP2DeltaBackend): """The transformers model wrapper. diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index e4a2e95a7..2daf21151 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -10,7 +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 import CudaDeltaSamplerMixin +from twinkle.checkpoint_engine.delta.rollout import DeltaSamplerMixin from twinkle.data_format import (InputFeature, PoolingParams, PoolingResponse, SampledSequence, SampleResponse, SamplingParams, Trajectory) from twinkle.hub import HubOperation @@ -77,7 +77,7 @@ def _load_vllm_image(entry: Any) -> Any: @remote_class(max_concurrency=_MAX_CONCURRENCY) -class vLLMSampler(Sampler, CheckpointEngineMixin, GenerationSubmissionMixin, PartialRolloutMixin, CudaDeltaSamplerMixin): +class vLLMSampler(Sampler, CheckpointEngineMixin, GenerationSubmissionMixin, 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 708887097..222cdea23 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py @@ -124,7 +124,7 @@ def __new__(cls, *args, **kwargs): return super().__new__(cls) - def validate_cuda_delta_worker(self) -> 'DeltaWorkerCapabilities': + def validate_delta_worker(self) -> 'DeltaWorkerCapabilities': from .delta_weight_transfer import TwinkleDeltaIPCWeightTransferEngine, is_moe_model engine = getattr(self, 'weight_transfer_engine', None) if not isinstance(engine, TwinkleDeltaIPCWeightTransferEngine): @@ -141,7 +141,7 @@ def delta_ipc_endpoint(self): self._delta_ipc_endpoint = f'ipc:///tmp/twinkle-delta-{uuid.uuid4().hex}.sock' return self._delta_ipc_endpoint - def cuda_delta_receiver_metadata(self) -> 'DeltaWorkerMetadata': + def delta_receiver_metadata(self) -> 'DeltaWorkerMetadata': from ray.util import get_node_ip_address return { 'node_ip': get_node_ip_address(), diff --git a/tests/checkpoint_engine/test_delta_manager.py b/tests/checkpoint_engine/test_delta_manager.py index a6db26441..10a24a166 100644 --- a/tests/checkpoint_engine/test_delta_manager.py +++ b/tests/checkpoint_engine/test_delta_manager.py @@ -13,13 +13,13 @@ 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_cuda_delta_sync.return_value = dict(device='cuda', master_metadata=object()) - sampler.validate_cuda_delta_sampler.return_value = dict(device='cuda', tp=2) - sampler.cuda_delta_sampler_metadata.return_value = [ + 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_cuda_delta_weights.return_value = lambda: {'payload_bytes': 6} - sampler.receive_cuda_delta_weights.return_value = lambda: {'changed': True} + 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 @@ -86,19 +86,19 @@ def test_seed_then_sparse_reuses_transport(): manager, model, sampler = manager_pair() manager.sync_weights() manager.sync_weights() - assert model.init_cuda_delta_transport.call_count == 1 + assert model.init_delta_transport.call_count == 1 assert manager._delta_round == 2 - assert model.init_cuda_delta_transport.call_args.args[0]['world_size'] == 3 - sampler.validate_cuda_delta_sampler.assert_called_once() - sampler.cuda_delta_sampler_metadata.assert_called_once() - model.prepare_cuda_delta_sync.assert_called_once() + 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_cuda_delta_weights', 'receive_cuda_delta_weights', 'publish_delta_sampler']) +@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_cuda_delta_weights' else sampler, stage).side_effect = TimeoutError('injected') + 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 @@ -110,10 +110,10 @@ def test_failure_does_not_publish_version(stage): def test_failed_round_requires_full_job_restart_even_after_close(): manager, model, sampler = manager_pair() manager.sync_weights() - sampler.receive_cuda_delta_weights.side_effect = ValueError('worker write failed') + sampler.receive_delta_weights.side_effect = ValueError('worker write failed') with pytest.raises(ValueError): manager.sync_weights() - sampler.receive_cuda_delta_weights.side_effect = None + sampler.receive_delta_weights.side_effect = None manager.close_sync_session() with pytest.raises(RuntimeError, match='restart the entire'): manager.sync_weights() @@ -123,7 +123,7 @@ def test_failed_round_requires_full_job_restart_even_after_close(): def test_cleanup_failure_does_not_clear_failed_state(): manager, model, sampler = manager_pair() manager.sync_weights() - model.close_cuda_delta_trainer.side_effect = RuntimeError('communicator hung') + model.close_delta_trainer.side_effect = RuntimeError('communicator hung') with pytest.raises(RuntimeError, match='cleanup'): manager.close_sync_session() assert manager._delta_failed @@ -133,38 +133,38 @@ def test_zero_change_still_resumes_manager_lifecycle(): manager, model, sampler = manager_pair() manager.sync_weights() sampler.publish_delta_sampler.reset_mock() - sampler.receive_cuda_delta_weights.return_value = lambda: {'changed': False} + sampler.receive_delta_weights.return_value = lambda: {'changed': False} manager.sync_weights() sampler.publish_delta_sampler.assert_called_once() - assert sampler.begin_cuda_delta_sync.call_count == 2 + assert sampler.begin_delta_sync.call_count == 2 def test_multiple_rollout_replicas_share_one_broadcast_group(): manager, model, sampler = manager_pair() - sampler.validate_cuda_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 - sampler.cuda_delta_sampler_metadata.return_value = [ + 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_cuda_delta_weights.return_value = lambda: [{'changed': True}] * 2 + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}] * 2 manager.sync_weights() manager.sync_weights() - endpoint = model.init_cuda_delta_transport.call_args.args[0] + endpoint = model.init_delta_transport.call_args.args[0] assert endpoint['world_size'] == 5 - assert endpoint['rollout_workers'] == sampler.cuda_delta_sampler_metadata.return_value - model.init_cuda_delta_transport.assert_called_once() + 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_cuda_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 - sampler.cuda_delta_sampler_metadata.return_value = [ + 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_cuda_delta_weights.return_value = lambda: [{'changed': True}] * 2 + sampler.receive_delta_weights.return_value = lambda: [{'changed': True}] * 2 manager.sync_weights() sampler.publish_delta_sampler.reset_mock() - sampler.receive_cuda_delta_weights.return_value = lambda: [{'changed': True}, {'changed': False}] + 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() @@ -173,11 +173,11 @@ def test_replica_completion_disagreement_keeps_rollout_paused(): def test_replica_metrics_may_differ_when_completion_agrees(): manager, model, sampler = manager_pair() - sampler.validate_cuda_delta_sampler.return_value = [dict(device='cuda', tp=2)] * 2 - sampler.cuda_delta_sampler_metadata.return_value = [ + 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_cuda_delta_weights.return_value = lambda: [ + 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() @@ -186,22 +186,22 @@ def test_replica_metrics_may_differ_when_completion_agrees(): @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.cuda_delta_sampler_metadata.return_value - sampler.cuda_delta_sampler_metadata.return_value = workers[:1] if invalid_topology == 'missing_worker' else [workers[0]] * 2 + 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_cuda_delta_transport.assert_not_called() + model.init_delta_transport.assert_not_called() -@pytest.mark.parametrize('stage', ['validate_cuda_delta_sampler', 'cuda_delta_sampler_metadata', 'prepare_cuda_delta_sync']) +@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_cuda_delta_sync' else sampler + 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_cuda_delta_sync.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 @@ -212,15 +212,15 @@ def test_manager_begin_pauses_before_receivers_or_sender_run(): events = [] model = Mock(_actors=[1]) sampler = Mock(_actors=[2]) - model.prepare_cuda_delta_sync.return_value = dict(device='cuda') - sampler.validate_cuda_delta_sampler.return_value = {'tp': 2} - sampler.cuda_delta_sampler_metadata.return_value = [ + 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_cuda_delta_sync.side_effect = lambda: events.append('pause-release') - model.init_cuda_delta_transport.side_effect = lambda endpoint: (events.append('init'), lambda: None)[1] - sampler.init_cuda_delta_sampler_transport.return_value = lambda: None - model.send_cuda_delta_weights.side_effect = lambda: (events.append('send'), lambda: {})[1] - sampler.receive_cuda_delta_weights.return_value = lambda: {'changed': True} + 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() diff --git a/tests/checkpoint_engine/test_delta_rollout.py b/tests/checkpoint_engine/test_delta_rollout.py index 68e792d0d..059c0e54c 100644 --- a/tests/checkpoint_engine/test_delta_rollout.py +++ b/tests/checkpoint_engine/test_delta_rollout.py @@ -8,12 +8,12 @@ import pytest import torch -from twinkle.checkpoint_engine.delta.rollout import CudaDeltaReceiver, CudaDeltaSamplerMixin +from twinkle.checkpoint_engine.delta.rollout import CudaDeltaReceiver, DeltaSamplerMixin from twinkle.sampler.vllm_sampler.vllm_engine import VLLMEngine from twinkle.sampler.vllm_sampler.vllm_worker_extension import _cuda_visible_device -class LocalCudaSampler(CudaDeltaSamplerMixin): +class LocalCudaSampler(DeltaSamplerMixin): def _run_in_loop(self, coroutine): return asyncio.run(coroutine) @@ -33,8 +33,8 @@ def test_zero_change_skips_native_update_but_restores_kv_and_generation(): sleep=AsyncMock(), wake_up=AsyncMock(), reset_prefix_cache=AsyncMock(return_value=True), resume_generation=AsyncMock())) sampler._delta_rpc = AsyncMock() - sampler.begin_cuda_delta_sync() - assert sampler.receive_cuda_delta_weights() == {'changed': False} + 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) @@ -52,7 +52,7 @@ def test_sampler_coordinates_native_update_and_gpu_local_senders(): 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_cuda_delta_weights() == {'changed': True} + 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: @@ -67,7 +67,7 @@ def test_receiver_stream_disagreement_fails_before_native_update(): sampler._delta_receivers = [receiver_stub([None]), receiver_stub([{'is_last': True}])] sampler._delta_rpc = AsyncMock() with pytest.raises(ValueError, match='disagree'): - sampler.receive_cuda_delta_weights() + sampler.receive_delta_weights() sampler._delta_rpc.assert_not_called() @@ -79,7 +79,7 @@ def test_cache_reset_failure_does_not_resume_generation(): 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_cuda_delta_weights() + sampler.receive_delta_weights() sampler.engine.engine.resume_generation.assert_not_called() @@ -88,7 +88,7 @@ def test_without_kv_release_pause_and_resume_still_run(): 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_cuda_delta_sync() + sampler.begin_delta_sync() sampler.publish_delta_sampler() sampler.engine.pause_delta_generation.assert_awaited_once() sampler.engine.engine.resume_generation.assert_awaited_once() @@ -175,7 +175,7 @@ def remote(self, config, endpoint, placement): sampler = LocalCudaSampler() sampler._delta_rpc = AsyncMock(return_value=metadata) sampler._delta_config = SimpleNamespace(timeout_seconds=1) - sampler.init_cuda_delta_sampler_transport({'world_size': 3, 'master_metadata': 'master', + 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'] @@ -204,7 +204,7 @@ def test_sampler_dp_allows_independent_rollout_replicas(): 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_cuda_delta_sampler(DeltaSyncConfig()) == dict(tp=2, device='cuda') + assert sampler.validate_delta_sampler(DeltaSyncConfig()) == dict(tp=2, device='cuda') def test_submission_barrier_waits_for_request_already_past_gate(): diff --git a/tests/checkpoint_engine/test_delta_trainer.py b/tests/checkpoint_engine/test_delta_trainer.py index d53687a5b..33d0a3c3b 100644 --- a/tests/checkpoint_engine/test_delta_trainer.py +++ b/tests/checkpoint_engine/test_delta_trainer.py @@ -20,7 +20,7 @@ from twinkle.checkpoint_engine.delta.backends.fsdp import validate_fsdp2 from twinkle.checkpoint_engine.delta.engine import DeltaShardedCheckpointEngine, _GatherQueue from twinkle.checkpoint_engine.delta.export import shard_delta_indices -from twinkle.checkpoint_engine.delta.trainer import CudaDeltaTrainerMixin +from twinkle.checkpoint_engine.delta.trainer import DeltaTrainerMixin from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy @@ -150,7 +150,7 @@ def raw(): list(hf_delta_export(raw(), {}, _hf_entry_identity)) -class FSDPDeltaTrainer(CudaDeltaTrainerMixin, FSDP2DeltaBackend): +class FSDPDeltaTrainer(DeltaTrainerMixin, FSDP2DeltaBackend): """A model host: lifecycle mixin plus the FSDP2 export backend.""" @@ -352,8 +352,8 @@ def initialize(config): trainer._initialize_delta_engine = Mock(side_effect=initialize) config = DeltaSyncConfig() - first = trainer.prepare_cuda_delta_sync(config) - second = trainer.prepare_cuda_delta_sync(config) + 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) diff --git a/tests/checkpoint_engine/test_delta_vllm.py b/tests/checkpoint_engine/test_delta_vllm.py index 26b22993d..4a25b20e3 100644 --- a/tests/checkpoint_engine/test_delta_vllm.py +++ b/tests/checkpoint_engine/test_delta_vllm.py @@ -139,7 +139,7 @@ def test_worker_requires_native_delta_backend(adapter, monkeypatch): monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta_weight_transfer', adapter) worker = SimpleNamespace(weight_transfer_engine=None) with pytest.raises(RuntimeError, match='weight_transfer_config'): - TwinkleWorkerExtension.validate_cuda_delta_worker(worker) + TwinkleWorkerExtension.validate_delta_worker(worker) def test_worker_allows_dense_model_and_tp_beyond_previous_allowlist(adapter, monkeypatch): @@ -149,7 +149,7 @@ def test_worker_allows_dense_model_and_tp_beyond_previous_allowlist(adapter, mon 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_cuda_delta_worker(worker) == {'tp': 4, 'device': 'cuda'} + assert TwinkleWorkerExtension.validate_delta_worker(worker) == {'tp': 4, 'device': 'cuda'} class Embedding: From 78a296f80f64c8617e9db268ad09db472d986ff6 Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 14 Oct 2026 10:17:43 +0800 Subject: [PATCH 04/12] feat(delta): use device-agnostic XOR-tree wire checksum Replace torch.hash_tensor (CPU/CUDA kernels only) with a byte-wise XOR fold that runs on any device, including NPU. Sender and receiver must be upgraded together. --- src/twinkle/checkpoint_engine/delta/wire.py | 34 +++- .../checkpoint_engine/test_delta_checksum.py | 171 ++++++++++++++++++ 2 files changed, 200 insertions(+), 5 deletions(-) create mode 100644 tests/checkpoint_engine/test_delta_checksum.py diff --git a/src/twinkle/checkpoint_engine/delta/wire.py b/src/twinkle/checkpoint_engine/delta/wire.py index 1caeb991c..da52f0e17 100644 --- a/src/twinkle/checkpoint_engine/delta/wire.py +++ b/src/twinkle/checkpoint_engine/delta/wire.py @@ -32,15 +32,39 @@ class DeltaParam: 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. - Uses ``torch.hash_tensor`` (XOR-reduce over uint64 bitcast); one reduction - plus one ``.item()`` sync per argument. + 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 = int(torch.hash_tensor(positions).item()) if positions.numel() else 0 - v = int(torch.hash_tensor(values).item()) if values.numel() else 0 - return p ^ (v << 1) + p = _xor_tree(positions) + v = _xor_tree(values) + return (p ^ (v << 1)) & ((1 << 64) - 1) # ---------- encode -------------------------------------------------------- diff --git a/tests/checkpoint_engine/test_delta_checksum.py b/tests/checkpoint_engine/test_delta_checksum.py new file mode 100644 index 000000000..de3e89786 --- /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 import wire +from twinkle.checkpoint_engine.delta.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_weight_transfer.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) From 84988370454e3f24de73ea132cf2b3fde92d3fd1 Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 14 Oct 2026 10:29:29 +0800 Subject: [PATCH 05/12] refactor(delta): extract transport-agnostic DeltaSyncMixin from the NCCL engine Move the seed/steady state machine, gather batching, bucketing and flush/manifest assembly into delta/pipeline.py. The NCCL engine keeps only transport concerns and implements four hooks: _wire_device, _ensure_wire, _broadcast_flush and _release_after_round. No behavior change for CUDA. --- .../checkpoint_engine/delta/__init__.py | 3 +- src/twinkle/checkpoint_engine/delta/engine.py | 512 +---------------- .../checkpoint_engine/delta/pipeline.py | 527 ++++++++++++++++++ .../checkpoint_engine/test_delta_pipeline.py | 205 +++++++ 4 files changed, 764 insertions(+), 483 deletions(-) create mode 100644 src/twinkle/checkpoint_engine/delta/pipeline.py create mode 100644 tests/checkpoint_engine/test_delta_pipeline.py diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py index cb81409b3..01ced82dd 100644 --- a/src/twinkle/checkpoint_engine/delta/__init__.py +++ b/src/twinkle/checkpoint_engine/delta/__init__.py @@ -3,7 +3,8 @@ Layout: config.py DeltaSyncConfig and the TypedDicts exchanged between trainer and rollout - engine.py DeltaShardedCheckpointEngine: gather, bucket and broadcast the changed elements + pipeline.py DeltaSyncMixin: transport-agnostic gather/bucket/flush pipeline and seed/steady state machine + engine.py DeltaShardedCheckpointEngine: NCCL transport (hooks + rollout receive) over the pipeline wire.py DeltaFlush / DeltaParam wire structures and checksum gather.py sparse gather of per-rank deltas to rank 0 spec.py ShardSpec / BlockPlacement: where a local shard sits in the full tensor diff --git a/src/twinkle/checkpoint_engine/delta/engine.py b/src/twinkle/checkpoint_engine/delta/engine.py index a0f73bef7..1481886f9 100644 --- a/src/twinkle/checkpoint_engine/delta/engine.py +++ b/src/twinkle/checkpoint_engine/delta/engine.py @@ -1,201 +1,40 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Delta checkpoint engine based on verl.""" +"""NCCL transport for the delta pipeline (the pipeline itself lives in :mod:`.pipeline`).""" from __future__ import annotations import json import logging -import time -from collections.abc import Generator, Iterator -from dataclasses import dataclass +from collections.abc import Iterator import torch -import zmq from ..nccl_checkpoint_engine import MasterMetadata, NCCLCheckpointEngine, _pg_broadcast from twinkle.utils.network import find_free_port -from .gather import gather_slot_entries_to_rank0 -from .wire import DeltaFlush, DeltaParam -from .wire import checksum as _checksum +# 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__) -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 = DeltaShardedCheckpointEngine.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 DeltaShardedCheckpointEngine(NCCLCheckpointEngine): +class DeltaShardedCheckpointEngine(DeltaSyncMixin, NCCLCheckpointEngine): """Sparse delta weight sync over NCCL, diffed on each rank's local shard. - Reuses NCCLCheckpointEngine's group/zmq machinery but moves only changed - positions+values: each actor rank keeps a pinned-CPU snapshot of only *its* - shard, byte-diffs the shard, and only the changed ``(position, value)`` pairs - are gathered to rank 0 and streamed to the rollout side -- no rank ever holds - a full-model snapshot. - - ``send_weights`` takes the TRAINING ENGINE and drives the sync itself: the - seed (first sync) streams the backend's full ``get_per_tensor_param()`` - export values-only and pins the diff base (``prime_delta_snapshots``); every - steady sync consumes the backend's HF delta export - ``get_per_tensor_param_delta_shard()`` — per-parameter FINAL-HF-coordinate - entries ``(slots, dtype_str, counts, hf_idx, hf_val, gather_group)``. Naming, - to-HF conversion, diff and snapshot all live on the backend side (see - :mod:`twinkle.checkpoint_engine.delta.export`); this engine only batches, gathers, - buckets and ships, and so serves any backend that can produce HF deltas. + 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``. """ - # 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 - wire_format = "delta_flush" def prepare(self) -> MasterMetadata | None: @@ -217,67 +56,18 @@ def _broadcast(self, tensor): if tensor.numel(): _pg_broadcast(self._pg, tensor) - # ---- 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.socket.send_string(self.topic, flags=zmq.SNDMORE) - self.socket.send_pyobj(meta) - pos_u8 = flush.positions_cpu.to("cuda", non_blocking=True).contiguous().view(torch.uint8) - val_u8 = flush.values_gpu.contiguous().view(torch.uint8) - self._broadcast(pos_u8) - self._broadcast(val_u8) + # ---- transport hooks ---- + @property + def _wire_device(self): + return "cuda" - 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.socket.send_string(self.topic, flags=zmq.SNDMORE) - self.socket.send_pyobj(meta) - val_u8 = values.view(torch.uint8) + def _broadcast_flush(self, pos_u8, val_u8) -> None: + if pos_u8 is not None: + self._broadcast(pos_u8) self._broadcast(val_u8) - 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.socket.send_string(self.topic, flags=zmq.SNDMORE) - self.socket.send_pyobj(meta) + 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]]: @@ -298,17 +88,17 @@ def receive_weights(self, global_steps: int | None = None) -> Iterator[tuple[lis 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="cuda") + 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="cuda") + 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("cuda") + 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)) @@ -319,245 +109,3 @@ def receive_weights(self, global_steps: int | None = None) -> Iterator[tuple[lis if is_last: break logger.info("delta recv v=%s flushes=%d (yielded to server adapter)", global_steps, applied) - - 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="cuda") - 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) - # 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 - 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 - logger.info("delta-sharded send v=%s delta flushes=%d (streamed)", global_steps, 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/delta/pipeline.py b/src/twinkle/checkpoint_engine/delta/pipeline.py new file mode 100644 index 000000000..f25b3c31e --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/pipeline.py @@ -0,0 +1,527 @@ +# 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 .wire import DeltaFlush, DeltaParam +from .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 + 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") + logger.info("delta-sharded send v=%s delta flushes=%d (streamed)", global_steps, 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/tests/checkpoint_engine/test_delta_pipeline.py b/tests/checkpoint_engine/test_delta_pipeline.py new file mode 100644 index 000000000..2dd2d919d --- /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.engine import DeltaShardedCheckpointEngine +from twinkle.checkpoint_engine.delta.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 From e12c0bc0f507eb9a4b610255baa25136feccf163 Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 14 Oct 2026 10:44:30 +0800 Subject: [PATCH 06/12] feat(delta): resolve the delta transport by platform and accept NPU trainers Add resolve_delta_engine_cls (GPU -> NCCL engine, NPU -> HCCL engine, imported lazily) and use it in the manager and trainer. The manager no longer rejects non-GPU platforms and fails at construction when no transport exists. The trainer derives its device from the platform and validate_fsdp2 accepts NPU. The NPU transport itself is not implemented yet, so selecting it raises NotImplementedError. --- .../checkpoint_engine/delta/__init__.py | 4 +- .../checkpoint_engine/delta/backends/fsdp.py | 5 +- src/twinkle/checkpoint_engine/delta/config.py | 16 +- .../checkpoint_engine/delta/platform.py | 28 +++ .../checkpoint_engine/delta/trainer.py | 22 ++- src/twinkle/checkpoint_engine/manager.py | 8 +- tests/checkpoint_engine/test_delta_manager.py | 13 -- .../checkpoint_engine/test_delta_platform.py | 187 ++++++++++++++++++ 8 files changed, 249 insertions(+), 34 deletions(-) create mode 100644 src/twinkle/checkpoint_engine/delta/platform.py create mode 100644 tests/checkpoint_engine/test_delta_platform.py diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py index 01ced82dd..74af7abfc 100644 --- a/src/twinkle/checkpoint_engine/delta/__init__.py +++ b/src/twinkle/checkpoint_engine/delta/__init__.py @@ -5,6 +5,7 @@ config.py DeltaSyncConfig and the TypedDicts exchanged between trainer and rollout pipeline.py DeltaSyncMixin: transport-agnostic gather/bucket/flush pipeline and seed/steady state machine engine.py DeltaShardedCheckpointEngine: NCCL transport (hooks + rollout receive) over the pipeline + platform.py resolve_delta_engine_cls: GPU -> NCCL engine, NPU -> HCCL engine (lazy import) wire.py DeltaFlush / DeltaParam wire structures and checksum gather.py sparse gather of per-rank deltas to rank 0 spec.py ShardSpec / BlockPlacement: where a local shard sits in the full tensor @@ -15,5 +16,6 @@ """ from .config import DeltaSyncConfig from .engine import DeltaShardedCheckpointEngine +from .platform import resolve_delta_engine_cls -__all__ = ['DeltaSyncConfig', 'DeltaShardedCheckpointEngine'] +__all__ = ['DeltaSyncConfig', 'DeltaShardedCheckpointEngine', 'resolve_delta_engine_cls'] diff --git a/src/twinkle/checkpoint_engine/delta/backends/fsdp.py b/src/twinkle/checkpoint_engine/delta/backends/fsdp.py index b9f580dff..35696e900 100644 --- a/src/twinkle/checkpoint_engine/delta/backends/fsdp.py +++ b/src/twinkle/checkpoint_engine/delta/backends/fsdp.py @@ -37,8 +37,9 @@ def validate_fsdp2(owner) -> None: from twinkle.model.transformers.strategy.accelerate import AccelerateStrategy from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy - if Platform.get_platform().__name__ != 'GPU': - raise ValueError('CUDA delta requires CUDA trainer actors') + 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) diff --git a/src/twinkle/checkpoint_engine/delta/config.py b/src/twinkle/checkpoint_engine/delta/config.py index f69b84e92..ded758a6f 100644 --- a/src/twinkle/checkpoint_engine/delta/config.py +++ b/src/twinkle/checkpoint_engine/delta/config.py @@ -1,12 +1,18 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Configuration for the verl CUDA delta checkpoint engine.""" +"""Configuration for the delta checkpoint engine.""" from __future__ import annotations from dataclasses import dataclass -from typing import TypedDict +from typing import TYPE_CHECKING, TypedDict, Union from typing_extensions import NotRequired -from ..nccl_checkpoint_engine import MasterMetadata +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 @@ -30,13 +36,13 @@ def __post_init__(self): class DeltaTransportEndpoint(TypedDict): world_size: int - master_metadata: MasterMetadata + master_metadata: DeltaMasterMetadata rollout_workers: list[DeltaWorkerMetadata] class DeltaTrainerMetadata(TypedDict): device: str - master_metadata: MasterMetadata | None + master_metadata: DeltaMasterMetadata | None class DeltaWorkerCapabilities(TypedDict): diff --git a/src/twinkle/checkpoint_engine/delta/platform.py b/src/twinkle/checkpoint_engine/delta/platform.py new file mode 100644 index 000000000..8e6169354 --- /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.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 .engine import DeltaShardedCheckpointEngine + return DeltaShardedCheckpointEngine + if platform_name == 'NPU': + try: + from .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/trainer.py b/src/twinkle/checkpoint_engine/delta/trainer.py index 8b9eea427..58028c77f 100644 --- a/src/twinkle/checkpoint_engine/delta/trainer.py +++ b/src/twinkle/checkpoint_engine/delta/trainer.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Trainer-side Ray lifecycle for the CUDA delta checkpoint engine.""" +"""Trainer-side Ray lifecycle for the delta checkpoint engine (NCCL on GPU, HCCL on NPU).""" from __future__ import annotations import asyncio @@ -9,17 +9,16 @@ import torch.distributed as dist -from twinkle import remote_function -from ..nccl_checkpoint_engine import MasterMetadata -from .config import DeltaSyncConfig, DeltaTrainerMetadata, DeltaTransportEndpoint -from .engine import DeltaShardedCheckpointEngine +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: DeltaShardedCheckpointEngine | None = None + _delta_engine: Any = None # transport chosen by platform (see resolve_delta_engine_cls) _delta_metadata_group: dist.ProcessGroup | None = None - _delta_master_metadata: MasterMetadata | None = None + _delta_master_metadata: DeltaMasterMetadata | None = None @staticmethod def _delta_rank() -> int: @@ -42,6 +41,9 @@ def _validate_delta_ranks(self) -> None: 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)) @@ -52,9 +54,9 @@ def _initialize_delta_engine(self, config: DeltaSyncConfig) -> None: dist.destroy_process_group(self._delta_metadata_group) self._delta_metadata_group = None raise - self._delta_device = 'cuda' + self._delta_device = DELTA_DEVICES[platform_name] self._delta_pin_snapshots = config.delta_pin_snapshots - self._delta_engine = DeltaShardedCheckpointEngine( + 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() @@ -64,7 +66,7 @@ def prepare_delta_sync(self, config: DeltaSyncConfig) -> DeltaTrainerMetadata: self._lazy_wrap_model() if self._delta_engine is None: self._initialize_delta_engine(config) - return {'device': 'cuda', 'master_metadata': self._delta_master_metadata} + return {'device': self._delta_device, 'master_metadata': self._delta_master_metadata} @remote_function(dispatch='all', lazy_collect=True) def init_delta_transport(self, endpoint: DeltaTransportEndpoint) -> None: diff --git a/src/twinkle/checkpoint_engine/manager.py b/src/twinkle/checkpoint_engine/manager.py index ec75249e6..a3191f4bb 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -6,7 +6,7 @@ from twinkle import Platform, get_logger from .base import CheckpointEngine -from .delta import DeltaSyncConfig +from .delta import DeltaSyncConfig, resolve_delta_engine_cls from .mixin import CheckpointEngineMixin logger = get_logger() @@ -92,8 +92,10 @@ def __init__( raise ValueError('delta_sharded requires standalone trainer and rollout actors') self.sync_backend = sync_backend self.delta_config = delta_config or DeltaSyncConfig() - if sync_backend == 'delta_sharded' and Platform.get_platform(platform).__name__ != 'GPU': - raise ValueError('delta_sharded currently supports CUDA only') + self.delta_engine_cls = None + if sync_backend == 'delta_sharded': + # Fail at construction, not mid-sync, when the platform has no delta transport. + self.delta_engine_cls = resolve_delta_engine_cls(Platform.get_platform(platform).__name__) self._delta_replica_count = None self._delta_round = 0 self._delta_failed = False diff --git a/tests/checkpoint_engine/test_delta_manager.py b/tests/checkpoint_engine/test_delta_manager.py index 10a24a166..8e8f30d42 100644 --- a/tests/checkpoint_engine/test_delta_manager.py +++ b/tests/checkpoint_engine/test_delta_manager.py @@ -69,19 +69,6 @@ def test_delta_rejects_local_deployment(): CheckpointEngineManager(SimpleNamespace(), SimpleNamespace(), sync_backend='delta_sharded') -def test_delta_rejects_npu_deployment(monkeypatch): - from twinkle import Platform - - class NPU: - pass - - monkeypatch.setattr(Platform, 'get_platform', staticmethod(lambda platform=None: NPU)) - model = Mock(_actors=[1, 2]) - sampler = Mock(_actors=[3]) - with pytest.raises(ValueError, match='CUDA only'): - CheckpointEngineManager(model, sampler, platform='NPU', sync_backend='delta_sharded') - - def test_seed_then_sparse_reuses_transport(): manager, model, sampler = manager_pair() manager.sync_weights() diff --git a/tests/checkpoint_engine/test_delta_platform.py b/tests/checkpoint_engine/test_delta_platform.py new file mode 100644 index 000000000..fb1b9cf9e --- /dev/null +++ b/tests/checkpoint_engine/test_delta_platform.py @@ -0,0 +1,187 @@ +# 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 import trainer as delta_trainer +from twinkle.checkpoint_engine.delta.backends.fsdp import validate_fsdp2 +from twinkle.checkpoint_engine.delta.engine import DeltaShardedCheckpointEngine +from twinkle.checkpoint_engine.delta.trainer import DeltaTrainerMixin +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + +HCCL_MODULE = 'twinkle.checkpoint_engine.delta.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_not_implemented(monkeypatch): + monkeypatch.delitem(sys.modules, HCCL_MODULE, raising=False) + try: + import importlib.util + exists = importlib.util.find_spec(HCCL_MODULE) is not None + except ModuleNotFoundError: + exists = False + if exists: + pytest.skip('HCCL delta module exists') + 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_fails_early_without_cuda_only_error(monkeypatch): + fake_platform(monkeypatch, 'NPU') + monkeypatch.delitem(sys.modules, HCCL_MODULE, raising=False) + 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.delitem(sys.modules, HCCL_MODULE, raising=False) + 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)) From 55914ee0f5011d72083a8770267d4ef9cdea687d Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 14 Oct 2026 10:54:23 +0800 Subject: [PATCH 07/12] feat(delta): add the HCCL delta transport for NPU trainers DeltaHCCLCheckpointEngine implements the four DeltaSyncMixin hooks over a PyHccl broadcast group whose other members are the vLLM workers. Rank 0 builds the communicator lazily at the first publish, so a ready group implies every receiver has subscribed to the zmq manifest channel. The group persists across rounds; rebuild_group=True is rejected. The zmq PUB/SUB helpers move to zmq_pubsub.py and are shared with the NCCL engine. --- .../checkpoint_engine/delta/__init__.py | 2 + .../checkpoint_engine/delta/hccl_engine.py | 173 +++++++++++++++ .../nccl_checkpoint_engine.py | 24 +- src/twinkle/checkpoint_engine/zmq_pubsub.py | 36 +++ tests/checkpoint_engine/test_delta_hccl.py | 210 ++++++++++++++++++ .../checkpoint_engine/test_delta_platform.py | 30 ++- tests/checkpoint_engine/test_zmq_pubsub.py | 41 ++++ 7 files changed, 484 insertions(+), 32 deletions(-) create mode 100644 src/twinkle/checkpoint_engine/delta/hccl_engine.py create mode 100644 src/twinkle/checkpoint_engine/zmq_pubsub.py create mode 100644 tests/checkpoint_engine/test_delta_hccl.py create mode 100644 tests/checkpoint_engine/test_zmq_pubsub.py diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py index 74af7abfc..f1e259b26 100644 --- a/src/twinkle/checkpoint_engine/delta/__init__.py +++ b/src/twinkle/checkpoint_engine/delta/__init__.py @@ -5,6 +5,8 @@ config.py DeltaSyncConfig and the TypedDicts exchanged between trainer and rollout pipeline.py DeltaSyncMixin: transport-agnostic gather/bucket/flush pipeline and seed/steady state machine engine.py DeltaShardedCheckpointEngine: NCCL transport (hooks + rollout receive) over the pipeline + hccl_engine.py DeltaHCCLCheckpointEngine: HCCL (Ascend NPU) sender transport over the pipeline + ../zmq_pubsub.py zmq PUB/SUB manifest-channel helpers shared by the NCCL and HCCL engines platform.py resolve_delta_engine_cls: GPU -> NCCL engine, NPU -> HCCL engine (lazy import) wire.py DeltaFlush / DeltaParam wire structures and checksum gather.py sparse gather of per-rank deltas to rank 0 diff --git a/src/twinkle/checkpoint_engine/delta/hccl_engine.py b/src/twinkle/checkpoint_engine/delta/hccl_engine.py new file mode 100644 index 000000000..607dcfdcc --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/hccl_engine.py @@ -0,0 +1,173 @@ +# 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 ..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 + + # ---- 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) + 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.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/nccl_checkpoint_engine.py b/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py index d6209c58b..dee851f98 100644 --- a/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py +++ b/src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py @@ -10,8 +10,9 @@ 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_pubsub import connect_subscriber, start_publisher logger = get_logger() @@ -138,28 +139,11 @@ def _start_zmq_server(self): self.ip = ray.util.get_node_ip_address().strip('[]') self.listen_port = find_free_port() - 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).""" - 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/tests/checkpoint_engine/test_delta_hccl.py b/tests/checkpoint_engine/test_delta_hccl.py new file mode 100644 index 000000000..5aff17948 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_hccl.py @@ -0,0 +1,210 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeltaHCCLCheckpointEngine with fake torch.npu, communicator and sockets (CPU only).""" +from __future__ import annotations + +import asyncio +import subprocess +import sys +from types import SimpleNamespace + +import pytest +import torch + +from twinkle.checkpoint_engine.delta import hccl_engine as mod +from twinkle.checkpoint_engine.delta.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.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 + + +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_platform.py b/tests/checkpoint_engine/test_delta_platform.py index fb1b9cf9e..6cd0b5e1d 100644 --- a/tests/checkpoint_engine/test_delta_platform.py +++ b/tests/checkpoint_engine/test_delta_platform.py @@ -33,15 +33,14 @@ def test_resolver_gpu(): assert resolve_delta_engine_cls('GPU') is DeltaShardedCheckpointEngine -def test_resolver_npu_not_implemented(monkeypatch): - monkeypatch.delitem(sys.modules, HCCL_MODULE, raising=False) - try: - import importlib.util - exists = importlib.util.find_spec(HCCL_MODULE) is not None - except ModuleNotFoundError: - exists = False - if exists: - pytest.skip('HCCL delta module exists') +def test_resolver_npu_real_engine(): + from twinkle.checkpoint_engine.delta.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') @@ -89,9 +88,16 @@ def test_manager_full_sync_does_not_resolve_delta(monkeypatch): assert manager.delta_engine_cls is None -def test_manager_npu_fails_early_without_cuda_only_error(monkeypatch): +def test_manager_npu_constructs_with_hccl_engine(monkeypatch): + fake_platform(monkeypatch, 'NPU') + from twinkle.checkpoint_engine.delta.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.delitem(sys.modules, HCCL_MODULE, raising=False) + monkeypatch.setitem(sys.modules, HCCL_MODULE, None) with pytest.raises(NotImplementedError, match=NOT_IMPLEMENTED): CheckpointEngineManager(*actors(), platform='NPU', mode='standalone', sync_backend='delta_sharded') @@ -152,7 +158,7 @@ def prepare(self): def test_trainer_npu_without_transport_fails_before_process_group(monkeypatch): fake_platform(monkeypatch, 'NPU') - monkeypatch.delitem(sys.modules, HCCL_MODULE, raising=False) + 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) 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() From 57259d873d34ea634e7621e8def848bfbe60d233 Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 14 Oct 2026 11:20:38 +0800 Subject: [PATCH 08/12] feat(delta): add the vLLM-side HCCL delta receiver for NPU Extract the device-agnostic flush decode into delta/decode.py and share the decode-and-apply path between the IPC and HCCL weight-transfer engines. The new twinkle_delta_hccl engine runs inside each vLLM worker: it connects the zmq manifest SUB before joining the trainer's PyHccl group, then pumps manifests and broadcasts into vLLM checkpoint patches. The worker extension registers the engine, reports the real device type and forwards per-worker group ranks. --- src/twinkle/checkpoint_engine/delta/decode.py | 80 +++ .../checkpoint_engine/delta/hccl_receiver.py | 94 ++++ .../vllm_sampler/delta_weight_transfer.py | 255 ++++++--- .../vllm_sampler/vllm_worker_extension.py | 23 +- .../test_delta_hccl_receiver.py | 513 ++++++++++++++++++ 5 files changed, 877 insertions(+), 88 deletions(-) create mode 100644 src/twinkle/checkpoint_engine/delta/decode.py create mode 100644 src/twinkle/checkpoint_engine/delta/hccl_receiver.py create mode 100644 tests/checkpoint_engine/test_delta_hccl_receiver.py diff --git a/src/twinkle/checkpoint_engine/delta/decode.py b/src/twinkle/checkpoint_engine/delta/decode.py new file mode 100644 index 000000000..40b7f2d73 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/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/hccl_receiver.py b/src/twinkle/checkpoint_engine/delta/hccl_receiver.py new file mode 100644 index 000000000..37bc2c741 --- /dev/null +++ b/src/twinkle/checkpoint_engine/delta/hccl_receiver.py @@ -0,0 +1,94 @@ +# 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 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): + """Join the trainer's HCCL broadcast group at ``group_rank``.""" + 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/sampler/vllm_sampler/delta_weight_transfer.py b/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py index ce7761d4b..8c06736d7 100644 --- a/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py +++ b/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py @@ -31,7 +31,9 @@ ) 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(): @@ -118,62 +120,20 @@ def decode_delta_payload( ) -> tuple[str, list[CheckpointWeightPatch]]: """Validate and decode one DeltaFlush into vLLM checkpoint patches.""" - from twinkle.checkpoint_engine.delta.wire import checksum + from twinkle.checkpoint_engine.delta.decode import decode_delta_wire CheckpointWeightPatch, _ = _checkpoint_patch_api() - 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[CheckpointWeightPatch] = [] - 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_indices = 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_indices = pos_bytes.view(torch.int32) - - patches.append( - CheckpointWeightPatch( - name=name, - shape=tuple(param_spec["shape"]), - dtype=dtype, - values=patch_values, - indices=patch_indices, - ) + 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, ) - return encoding, patches + for p in patches + ] class TwinkleDeltaIPCWeightTransferEngine(WeightTransferEngine): @@ -182,6 +142,7 @@ class TwinkleDeltaIPCWeightTransferEngine(WeightTransferEngine): init_info_cls = TwinkleDeltaIPCInitInfo update_info_cls = TwinkleDeltaIPCUpdateInfo supports_draft_weight_update = False + required_device_type = "cuda" def __init__( self, @@ -194,10 +155,13 @@ def __init__( 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 != 'cuda': - raise NotImplementedError("Twinkle delta weight transfer requires CUDA") + 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: @@ -251,42 +215,46 @@ def retain_bucket(weights: list[tuple[str, torch.Tensor]], _is_last: bool) -> No raise callback_error return payload - def receive_weights(self, update_info: TwinkleDeltaIPCUpdateInfo) -> None: + 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.weight_update_utils import ( drop_tied_alias_patches, prepare_sparse_weight_loaders, ) - assert update_info.zmq_handle is not None - try: - payload = self._receive_payload(zmq_handle=update_info.zmq_handle) - 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})" + 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, ) - 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 GPU. - load_checkpoint_weight_patches( - self.model, - patches, - validate_unique_indices=False, + 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. @@ -335,3 +303,124 @@ def register_twinkle_delta_weight_transfer_engine() -> None: ) _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 + + +@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() + 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.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 + ) + 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.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/vllm_worker_extension.py b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py index 222cdea23..ab9c3495c 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py @@ -121,6 +121,9 @@ def __new__(cls, *args, **kwargs): if getattr(config, 'backend', None) == 'twinkle_delta_ipc': from .delta_weight_transfer import register_twinkle_delta_weight_transfer_engine register_twinkle_delta_weight_transfer_engine() + elif getattr(config, 'backend', None) == 'twinkle_delta_hccl': + from .delta_weight_transfer import register_twinkle_delta_hccl_weight_transfer_engine + register_twinkle_delta_hccl_weight_transfer_engine() return super().__new__(cls) @@ -128,12 +131,13 @@ def validate_delta_worker(self) -> 'DeltaWorkerCapabilities': from .delta_weight_transfer import TwinkleDeltaIPCWeightTransferEngine, is_moe_model engine = getattr(self, 'weight_transfer_engine', None) if not isinstance(engine, TwinkleDeltaIPCWeightTransferEngine): - raise RuntimeError("CUDA delta requires engine_args['weight_transfer_config']=" - "{'backend': 'twinkle_delta_ipc'} when constructing the sampler") + 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': 'cuda'} + 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'): @@ -150,8 +154,17 @@ def delta_receiver_metadata(self) -> 'DeltaWorkerMetadata': 'zmq_handle': self.delta_ipc_endpoint(), } - def update_twinkle_delta_weights(self): - self.update_weights({'zmq_handle': self.delta_ipc_endpoint()}) + 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}]``) is forwarded as is; + ``WorkerBase.update_weights`` is expected to pick the entry for this worker's rank. + """ + if update_info is None: + self.update_weights({'zmq_handle': self.delta_ipc_endpoint()}) + else: + self.update_weights(update_info) def monkey_patch_model(self): from twinkle.patch.vllm_moe_loader import VLLMMoEWeights 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..62f40e7a1 --- /dev/null +++ b/tests/checkpoint_engine/test_delta_hccl_receiver.py @@ -0,0 +1,513 @@ +# 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 import hccl_receiver +from twinkle.checkpoint_engine.delta.decode import DeltaPatch, decode_delta_wire +from twinkle.checkpoint_engine.delta.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') + + +# --------------------------------------------------------------------------- (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): + order.append(('build', ip, port, rank, world)) + return comm + + 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) + + +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) + + +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.weight_update_utils') + prepare = Mock() + drop = Mock(side_effect=lambda model, patches: patches) + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.weight_update_utils', + SimpleNamespace(prepare_sparse_weight_loaders=prepare, drop_tied_alias_patches=drop)) + 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 + 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_weight_transfer', 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 test_update_twinkle_delta_weights_dispatch(worker_cls): + worker = SimpleNamespace(update_weights=Mock(), delta_ipc_endpoint=lambda: 'ipc:///x') + 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) + worker.update_weights.assert_called_once_with(info) + + +# --------------------------------------------------------------------------- (e) import hygiene +def test_import_without_npu_modules(): + code = ("import sys, twinkle.checkpoint_engine.delta.hccl_receiver, twinkle.checkpoint_engine.delta.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 From b8c610ed5ddb74d4c143e77d19aeaec85455f2aa Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 21 Oct 2026 15:53:16 +0800 Subject: [PATCH 09/12] wip --- cookbook/checkpoint_engine/check_npu_env.py | 199 ++++++++ .../checkpoint_engine/check_weight_names.py | 64 +++ .../checkpoint_engine/crossnode_comm_check.py | 247 ++++++++++ cookbook/checkpoint_engine/delta_benchmark.py | 83 ++++ .../delta_benchmark_notes.md | 144 ++++++ .../checkpoint_engine/delta_validation.py | 393 ++++++++++++++++ cookbook/checkpoint_engine/device_utils.py | 232 ++++++++++ .../checkpoint_engine/e2e_step_benchmark.py | 257 +++++++++++ delta_perf.md | 69 +++ .../checkpoint_engine/delta/__init__.py | 3 +- src/twinkle/checkpoint_engine/delta/config.py | 8 +- .../checkpoint_engine/delta/hccl_engine.py | 6 + .../checkpoint_engine/delta/hccl_receiver.py | 13 +- .../checkpoint_engine/delta/pipeline.py | 11 +- .../checkpoint_engine/delta/rollout.py | 84 +++- .../checkpoint_engine/delta/trainer.py | 6 +- .../hccl_checkpoint_engine.py | 10 +- src/twinkle/checkpoint_engine/manager.py | 23 +- src/twinkle/infra/_ray/ray_helper.py | 5 +- .../model/transformers/transformers.py | 33 +- .../vllm_sampler/delta_weight_transfer.py | 23 +- .../sampler/vllm_sampler/vllm_engine.py | 8 +- .../vllm_sampler/vllm_worker_extension.py | 31 +- .../vllm_sampler/weight_update_utils.py | 29 ++ src/twinkle/utils/platforms/__init__.py | 2 +- src/twinkle/utils/platforms/npu.py | 11 + src/twinkle/utils/torch_utils.py | 17 +- tests/checkpoint_engine/_hccl_loopback.py | 369 +++++++++++++++ tests/checkpoint_engine/test_delta_hccl.py | 17 + .../test_delta_hccl_loopback.py | 134 ++++++ .../test_delta_hccl_receiver.py | 120 ++++- .../test_delta_npu_orchestration.py | 425 ++++++++++++++++++ tests/checkpoint_engine/test_delta_vllm.py | 58 +++ 33 files changed, 3096 insertions(+), 38 deletions(-) create mode 100644 cookbook/checkpoint_engine/check_npu_env.py create mode 100644 cookbook/checkpoint_engine/check_weight_names.py create mode 100644 cookbook/checkpoint_engine/crossnode_comm_check.py create mode 100644 cookbook/checkpoint_engine/delta_benchmark.py create mode 100644 cookbook/checkpoint_engine/delta_benchmark_notes.md create mode 100644 cookbook/checkpoint_engine/delta_validation.py create mode 100644 cookbook/checkpoint_engine/device_utils.py create mode 100644 cookbook/checkpoint_engine/e2e_step_benchmark.py create mode 100644 delta_perf.md create mode 100644 tests/checkpoint_engine/_hccl_loopback.py create mode 100644 tests/checkpoint_engine/test_delta_hccl_loopback.py create mode 100644 tests/checkpoint_engine/test_delta_npu_orchestration.py diff --git a/cookbook/checkpoint_engine/check_npu_env.py b/cookbook/checkpoint_engine/check_npu_env.py new file mode 100644 index 000000000..088b230ad --- /dev/null +++ b/cookbook/checkpoint_engine/check_npu_env.py @@ -0,0 +1,199 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""检查当前环境是否满足 delta_sharded（HCCL/NPU 或 IPC/CUDA）增量权重同步的依赖。 + +只做导入和签名检查，不启动 Ray、不接触加速卡，可在任一节点单独运行： + + python -m cookbook.checkpoint_engine.check_npu_env # 自动探测平台 + python -m cookbook.checkpoint_engine.check_npu_env --platform NPU + +每项打印 [PASS]/[FAIL]/[WARN]，退出码 0 表示所有必需项通过。换机器或换镜像后先跑它， +可以在拉起几分钟的 vLLM 之前就发现 vLLM / vllm-ascend 版本不兼容的问题。 +""" +from __future__ import annotations + +import argparse +import importlib +import importlib.metadata +import inspect + + +class Checker: + def __init__(self) -> None: + self.failed = 0 + self.warned = 0 + + def ok(self, name: str, detail: str = '') -> None: + print(f'[PASS] {name}' + (f'：{detail}' if detail else '')) + + def fail(self, name: str, detail: str) -> None: + self.failed += 1 + print(f'[FAIL] {name}：{detail}') + + def warn(self, name: str, detail: str) -> None: + self.warned += 1 + print(f'[WARN] {name}：{detail}') + + def check(self, name: str, fn, *, required: bool = True) -> object: + """运行 fn；返回 (detail_or_value)。约定：fn 返回字符串即 PASS 详情，抛异常即失败。""" + try: + detail = fn() + except Exception as error: # noqa: BLE001 - 汇总所有失败，不在第一项中止 + (self.fail if required else self.warn)(name, f'{type(error).__name__}: {error}') + return None + self.ok(name, detail if isinstance(detail, str) else '') + return detail + + +def _version(package: str) -> str: + try: + return importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + return 'not-installed' + + +def check_versions(checker: Checker) -> None: + for package in ('torch', 'torch_npu', 'vllm', 'vllm_ascend', 'transformers'): + version = _version(package.replace('_', '-')) + print(f' {package}: {version}') + + +def check_weight_transfer_engine(checker: Checker) -> None: + def fn(): + from vllm.distributed.weight_transfer.base import WeightTransferEngine + params = inspect.signature(WeightTransferEngine.__init__).parameters + missing = {'config', 'vllm_config', 'device', 'model'} - set(params) + if missing: + raise RuntimeError(f'WeightTransferEngine.__init__ 缺少参数 {sorted(missing)}（需要 vLLM #44353）') + return 'engine-owned WeightTransferEngine API 可用（vLLM #44353）' + + checker.check('WeightTransferEngine API (#44353)', fn) + + +def check_checkpoint_patch_api(checker: Checker) -> None: + def fn(): + module = importlib.import_module('vllm.model_executor.model_loader.checkpoint_weight_patch') + for attr in ('CheckpointWeightPatch', 'load_checkpoint_weight_patches'): + if not hasattr(module, attr): + raise RuntimeError(f'缺少 {attr}（需要 vLLM #50723）') + sig = inspect.signature(module.load_checkpoint_weight_patches) + if 'validate_unique_indices' not in sig.parameters: + raise RuntimeError('load_checkpoint_weight_patches 缺少 validate_unique_indices 参数') + return 'checkpoint patch API 可用（vLLM #50723）' + + checker.check('checkpoint patch API (#50723)', fn) + + +def check_layerwise_reload(checker: Checker) -> None: + def fn(): + module = importlib.import_module('vllm.model_executor.model_loader.reload') + for attr in ('initialize_layerwise_reload', 'finalize_layerwise_reload'): + if not hasattr(module, attr): + raise RuntimeError(f'缺少 {attr}（dense/full 同步依赖 layerwise reload）') + return 'layerwise reload 生命周期可用' + + checker.check('layerwise reload API', fn) + + +def check_stateless_pg(checker: Checker) -> None: + def fn(): + import dataclasses + from vllm.distributed.utils import StatelessProcessGroup + fields = {f.name for f in dataclasses.fields(StatelessProcessGroup)} + # torch_utils.stateless_init_process_group 已兼容有/无 socket 两种版本，这里只提示。 + return f"StatelessProcessGroup 可导入（socket 字段：{'有' if 'socket' in fields else '无，走兼容分支'}）" + + checker.check('StatelessProcessGroup', fn) + + +def check_npu_hccl(checker: Checker) -> None: + def torch_npu_available(): + import torch + import torch_npu # noqa: F401 + if not torch.npu.is_available(): + raise RuntimeError('torch.npu.is_available() 为 False（本节点看不到 NPU）') + return f'torch.npu 可用，device_count={torch.npu.device_count()}' + + checker.check('torch_npu 设备', torch_npu_available) + + def pyhccl(): + from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator # noqa: F401 + return 'vllm_ascend PyHcclCommunicator 可导入（delta HCCL broadcast 依赖）' + + checker.check('vllm-ascend PyHccl', pyhccl) + + def rl_config(): + # RL 场景必须支持 additional_config={'rl_config': {'enabled': True}}；仅检查配置类存在。 + from vllm_ascend import ascend_config # noqa: F401 + return 'vllm_ascend.ascend_config 可导入（rl_config 开关所在模块）' + + checker.check('vllm-ascend rl_config 模块', rl_config, required=False) + + +def check_cuda_ipc(checker: Checker) -> None: + def fn(): + import torch + if not torch.cuda.is_available(): + raise RuntimeError('torch.cuda.is_available() 为 False') + return f'torch.cuda 可用，device_count={torch.cuda.device_count()}' + + checker.check('torch.cuda 设备', fn) + + +def check_twinkle_delta(checker: Checker) -> None: + def fn(): + from twinkle.sampler.vllm_sampler.delta_weight_transfer import require_vllm_delta_support + require_vllm_delta_support() + return 'twinkle require_vllm_delta_support() 通过' + + checker.check('twinkle delta 依赖自检', fn) + + def refresh(): + from twinkle.sampler.vllm_sampler.weight_update_utils import refresh_derived_weights # noqa: F401 + return 'refresh_derived_weights 可用（sparse 同步后刷新 Ascend 打包权重）' + + checker.check('twinkle refresh_derived_weights', refresh, required=False) + + +def detect_platform() -> str: + try: + import torch + import torch_npu # noqa: F401 + if torch.npu.is_available(): + return 'NPU' + except Exception: # noqa: BLE001 + pass + return 'GPU' + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--platform', choices=('GPU', 'NPU'), help='默认自动探测（能导入 torch_npu 且 NPU 可用则 NPU）') + args = parser.parse_args() + platform = args.platform or detect_platform() + + print(f'== 依赖检查 platform={platform} ==') + checker = Checker() + check_versions(checker) + print('-- vLLM 通用 delta 依赖 --') + check_weight_transfer_engine(checker) + check_checkpoint_patch_api(checker) + check_layerwise_reload(checker) + check_stateless_pg(checker) + print(f'-- {platform} 专用依赖 --') + if platform == 'NPU': + check_npu_hccl(checker) + else: + check_cuda_ipc(checker) + print('-- twinkle 侧 --') + check_twinkle_delta(checker) + + print('== 汇总 ==') + if checker.failed: + print(f'{checker.failed} 项必需检查未通过，{checker.warned} 项告警。请先解决 [FAIL] 再运行 delta_validation。') + return 1 + print(f'全部必需检查通过，{checker.warned} 项告警。') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/cookbook/checkpoint_engine/check_weight_names.py b/cookbook/checkpoint_engine/check_weight_names.py new file mode 100644 index 000000000..61867ca9a --- /dev/null +++ b/cookbook/checkpoint_engine/check_weight_names.py @@ -0,0 +1,64 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""打印 Qwen3.5 的权重名映射，用于定位 full 同步的 tied embedding 报错。 + +不加载权重（meta device 建骨架），不启动 Ray，只看名字： + python -m cookbook.checkpoint_engine.check_weight_names --model-path /nas/disk1/Qwen3.5-4B +""" +import argparse + +import torch +import transformers +from transformers import AutoConfig, AutoModelForCausalLM + +from twinkle.checkpoint_engine.delta.backends.fsdp import convert_weight_keys + +KEYWORDS = ('embed_tokens', 'lm_head') + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument('--model-path', required=True) + args = parser.parse_args() + + config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True) + architectures = getattr(config, 'architectures', None) + print(f'architectures : {architectures}') + text_config = getattr(config, 'text_config', config) + print(f'tie_word_embeddings : {getattr(config, "tie_word_embeddings", None)} ' + f'(text_config: {getattr(text_config, "tie_word_embeddings", None)})') + + # 与 TransformersModel 一致：优先用 config.architectures[0]，而不是 AutoModelForCausalLM + # （两者命名空间不同：ForConditionalGeneration 带 language_model. 前缀）。 + model_cls = AutoModelForCausalLM + if architectures: + model_cls = getattr(transformers, architectures[0], AutoModelForCausalLM) + print(f'model_cls (同 trainer): {getattr(model_cls, "__name__", model_cls)}') + + with torch.device('meta'): + if hasattr(model_cls, 'from_config'): + model = model_cls.from_config(config, trust_remote_code=True) + else: + model = model_cls._from_config(config) + print(f'model class : {type(model).__name__}') + mapping = getattr(model, '_checkpoint_conversion_mapping', None) + print(f'_checkpoint_conversion_mapping: {mapping}') + + state_dict = model.state_dict() + print(f'\nstate_dict 里的相关键（trainer 侧原始名，共 {len(state_dict)} 个权重）：') + for name in state_dict: + if any(k in name for k in KEYWORDS): + print(f' {name}') + + converted = convert_weight_keys(state_dict, model) + print(f'\nconvert_weight_keys 之后（实际发给 vLLM 的名字）：') + for name in converted: + if any(k in name for k in KEYWORDS): + print(f' {name}') + + print('\nvLLM 报错期望的是 language_model.model.embed_tokens.weight；' + '若上面没有这个名字，就是导出侧少了 language_model. 前缀。') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/cookbook/checkpoint_engine/crossnode_comm_check.py b/cookbook/checkpoint_engine/crossnode_comm_check.py new file mode 100644 index 000000000..b24cd023f --- /dev/null +++ b/cookbook/checkpoint_engine/crossnode_comm_check.py @@ -0,0 +1,247 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""最小化的两机通信连通性测试，不依赖 twinkle / vLLM / Ray。 + +用途：当 delta_validation 跨机卡在 HCCL 建链时，用它判断是“网络本身不通”还是“HCCL 跨机不通”。 +分两个后端测： + - gloo：纯 TCP，验证两机三层网络 + MASTER_PORT 能否 rendezvous（不碰 NPU）。 + - hccl：昇腾集合通信，验证 NPU 跨机能否真正建链并做 all_reduce。 + +两个 Pod 各跑一条命令（先 gloo 再 hccl）。MASTER 必须是 rank 0 所在 Pod 的 IP。 + + # Pod1（rank 0，172.61.8.99） + python -m cookbook.checkpoint_engine.crossnode_comm_check \ + --backend gloo --rank 0 --world-size 2 --master-addr 172.61.8.99 --master-port 29555 + # Pod2（rank 1，172.61.10.44） + python -m cookbook.checkpoint_engine.crossnode_comm_check \ + --backend gloo --rank 1 --world-size 2 --master-addr 172.61.8.99 --master-port 29555 + +gloo 通过后把 --backend 换成 hccl，--master-port 换一个新端口（如 29556）再各跑一次。 +gloo 通、hccl 不通，基本可判定：该环境的 NPU 不具备跨机 HCCL（通常是缺 RDMA/RoCE）。 +""" +from __future__ import annotations + +import argparse +import datetime +import os +import socket + + +def log(message: str) -> None: + print(f'[{datetime.datetime.now():%H:%M:%S}] [{socket.gethostname()}] {message}', flush=True) + + +def _run_stateless(args, torch, device) -> int: + """用 delta 完全相同的 pyhccl 路径（stateless_init_process_group）建链并做 all_reduce。 + + 这正是 delta 超时的那条调用栈： + stateless_init_process_group -> PyHcclCommunicator(pg, device) -> __init__ 里的 all_reduce。 + 若此处也超时，说明问题在 pyhccl 初始化路径本身（与并发组/时间错位无关）； + 若此处 PASS，说明裸 pyhccl 单组可通，delta 的超时来自并发的第二个 HCCL 组或 join 时间错位。 + """ + from twinkle.utils import stateless_init_process_group + + log(f'开始 stateless_init_process_group（pyhccl）：rank={args.rank}/{args.world_size}, ' + f'master={args.master_addr}:{args.master_port}, device={device.index}') + log('这条路径 = delta 的 build_delta_hccl_communicator；PyHcclCommunicator.__init__ 内部会做一次 all_reduce') + try: + comm = stateless_init_process_group( + master_address=args.master_addr, + master_port=args.master_port, + rank=args.rank, + world_size=args.world_size, + device=device.index, + backend='hccl', + ) + except Exception as error: # noqa: BLE001 + import traceback + log(f'FAIL：stateless_init_process_group 抛异常：{type(error).__name__}: {error}') + traceback.print_exc() + return 1 + log('stateless_init_process_group 返回成功（communicator 已建立，内部 all_reduce 已通过）') + + try: + tensor = torch.ones(1, device=device) + # PyHcclCommunicator.all_reduce 是 out-of-place：结果在返回值里，不改入参。 + out = comm.all_reduce(tensor) + torch.npu.synchronize() + value = int(out.item()) + log(f'comm.all_reduce 完成，结果={value}（期望={args.world_size}）') + if value != args.world_size: + log('FAIL：all_reduce 结果不等于 world_size，通信异常') + return 1 + log('PASS：pyhccl 跨机单组通信正常') + except Exception as error: # noqa: BLE001 + import traceback + log(f'FAIL：comm.all_reduce 抛异常：{type(error).__name__}: {error}') + traceback.print_exc() + return 1 + return 0 + + +def _build_pyhccl_group(args, torch, device, master_port, label): + """用 delta 的 pyhccl 路径建一个组并做一次 all_reduce，返回 (comm, 是否正常)。""" + from twinkle.utils import stateless_init_process_group + + log(f'[{label}] stateless_init_process_group：master_port={master_port}, ' + f"HCCL_IF_BASE_PORT={os.environ.get('HCCL_IF_BASE_PORT')}, " + f"NPU_RANGE={os.environ.get('HCCL_NPU_SOCKET_PORT_RANGE')}") + comm = stateless_init_process_group( + master_address=args.master_addr, master_port=master_port, + rank=args.rank, world_size=args.world_size, device=device.index, backend='hccl') + out = comm.all_reduce(torch.ones(1, device=device)) + torch.npu.synchronize() + value = int(out.item()) + ok = value == args.world_size + log(f'[{label}] all_reduce 结果={value}（期望={args.world_size}）{"OK" if ok else "异常"}') + return comm, ok + + +def _run_two_groups(args, torch, device) -> int: + """在同一张卡上先后建两个并发 HCCL 组，复现 delta 的「第二个并发组」场景。 + + delta 真实情况：trainer 卡上已有 FSDP 组、Pod2 卡上已有 vLLM TP 组，delta 组是第二个； + 且 trainer 侧 delta 组与 FSDP 组共用同一个 HCCL_IF_BASE_PORT。 + + 组 1 建好后不销毁（占住设备），再建组 2。 + - 默认：组 2 复用相同 HCCL 端口（= delta 与 FSDP 共享 base port 的情形）。 + - 指定 --second-base-port：组 2 换一套端口（验证「换端口即可解决」的修复方向）。 + """ + log('==== 建立组 1（占住设备，模拟已有的 FSDP / TP 通信域）====') + try: + comm1, ok1 = _build_pyhccl_group(args, torch, device, args.master_port, 'group1') + except Exception as error: # noqa: BLE001 + import traceback + log(f'FAIL：组 1 建链失败：{type(error).__name__}: {error}') + traceback.print_exc() + return 1 + if not ok1: + log('FAIL：组 1 本身就不正常，无法继续') + return 1 + log('组 1 保持存活（不销毁），现在在同一张卡上建组 2') + + if args.second_base_port: + base = args.second_base_port + os.environ['HCCL_IF_BASE_PORT'] = str(base) + os.environ['HCCL_HOST_SOCKET_PORT_RANGE'] = f'{base + 1000}-{base + 1511}' + os.environ['HCCL_NPU_SOCKET_PORT_RANGE'] = f'{base + 2000}-{base + 2511}' + log(f'==== 建立组 2（换一套 HCCL 端口 base={base}）====') + else: + log('==== 建立组 2（复用与组 1 相同的 HCCL 端口 = delta 与 FSDP 共享 base port 的情形）====') + + try: + comm2, ok2 = _build_pyhccl_group(args, torch, device, args.master_port + 1, 'group2') + except Exception as error: # noqa: BLE001 + import traceback + log(f'FAIL：组 2 建链失败（这就是 delta 的超时点）：{type(error).__name__}: {error}') + traceback.print_exc() + return 1 + if not ok2: + log('FAIL：组 2 并发通信异常') + return 1 + log('PASS：同一张卡上第二个并发 HCCL 组也正常') + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--backend', choices=('gloo', 'hccl'), required=True) + parser.add_argument('--rank', type=int, required=True, help='本进程的全局 rank（rank 0 即 master）') + parser.add_argument('--world-size', type=int, default=2) + parser.add_argument('--master-addr', required=True, help='rank 0 所在 Pod 的 IP') + parser.add_argument('--master-port', type=int, default=29555) + parser.add_argument('--device', type=int, default=0, help='本机使用的本地 NPU 序号（hccl 后端用）') + parser.add_argument('--timeout', type=int, default=120, help='rendezvous/集合通信超时秒数') + parser.add_argument('--iface', help='指定网卡名，会设置 GLOO_SOCKET_IFNAME / HCCL_SOCKET_IFNAME') + parser.add_argument( + '--stateless', action='store_true', + help='hccl 后端用 delta 完全相同的 vllm_ascend pyhccl 路径（stateless_init_process_group）建链，' + '而非 torch.distributed。用于判断 delta 超时是否来自这条 pyhccl 初始化路径本身。') + parser.add_argument( + '--startup-delay', type=int, default=0, + help='建链前本进程先睡眠的秒数。用于复现 delta 里 trainer（rank 0）很晚才 join 的时间错位，' + '例如给 rank 0 设一个大的值，看非 0 rank 是否因等待超时。') + parser.add_argument( + '--two-groups', action='store_true', + help='在同一张卡上先后建两个并发 HCCL 组（pyhccl 路径），复现 delta 的第二个并发组场景。') + parser.add_argument( + '--second-base-port', type=int, default=0, + help='配合 --two-groups：给第二个组换一套 HCCL_IF_BASE_PORT（及派生端口段）。' + '不指定则第二个组复用与第一个组相同的端口（= delta 与 FSDP 共享 base port）。') + args = parser.parse_args() + + if args.iface: + os.environ['GLOO_SOCKET_IFNAME'] = args.iface + os.environ['HCCL_SOCKET_IFNAME'] = args.iface + os.environ.setdefault('TP_SOCKET_IFNAME', args.iface) + + os.environ['MASTER_ADDR'] = args.master_addr + os.environ['MASTER_PORT'] = str(args.master_port) + os.environ['RANK'] = str(args.rank) + os.environ['WORLD_SIZE'] = str(args.world_size) + + import torch + import torch.distributed as dist + + if args.backend == 'hccl': + # 必须在创建进程组前导入并选卡，否则 HCCL 不知道用哪张 NPU。 + import torch_npu # noqa: F401 + if not torch.npu.is_available(): + log('FAIL：torch.npu.is_available() 为 False，本机看不到 NPU') + return 1 + torch.npu.set_device(args.device) + device = torch.device('npu', args.device) + log(f'使用 NPU device {args.device}，device_count={torch.npu.device_count()}') + else: + device = torch.device('cpu') + + if args.startup_delay > 0: + import time + log(f'建链前先睡眠 {args.startup_delay}s（模拟 join 时间错位）') + time.sleep(args.startup_delay) + + if args.stateless: + if args.backend != 'hccl': + log('FAIL：--stateless 仅支持 hccl 后端') + return 1 + return _run_stateless(args, torch, device) + + if args.two_groups: + if args.backend != 'hccl': + log('FAIL：--two-groups 仅支持 hccl 后端') + return 1 + return _run_two_groups(args, torch, device) + + log(f'开始 init_process_group：backend={args.backend}, rank={args.rank}/{args.world_size}, ' + f'master={args.master_addr}:{args.master_port}, timeout={args.timeout}s') + log('若进程在此长时间停住，说明两机无法在该 backend 上 rendezvous（gloo=TCP 不通；hccl=NPU 跨机建链失败）') + try: + dist.init_process_group( + backend=args.backend, init_method='env://', rank=args.rank, world_size=args.world_size, + timeout=datetime.timedelta(seconds=args.timeout)) + except Exception as error: # noqa: BLE001 + log(f'FAIL：init_process_group 抛异常：{type(error).__name__}: {error}') + return 1 + log('init_process_group 成功，进程组已建立') + + try: + tensor = torch.ones(1, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + if args.backend == 'hccl': + torch.npu.synchronize() + value = int(tensor.item()) + log(f'all_reduce 完成，结果={value}（期望={args.world_size}）') + if value != args.world_size: + log('FAIL：all_reduce 结果不等于 world_size，通信异常') + return 1 + log('PASS：跨机通信正常') + except Exception as error: # noqa: BLE001 + log(f'FAIL：all_reduce 抛异常：{type(error).__name__}: {error}') + return 1 + finally: + dist.barrier() + dist.destroy_process_group() + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/cookbook/checkpoint_engine/delta_benchmark.py b/cookbook/checkpoint_engine/delta_benchmark.py new file mode 100644 index 000000000..170c1f7b1 --- /dev/null +++ b/cookbook/checkpoint_engine/delta_benchmark.py @@ -0,0 +1,83 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""权重同步耗时基准，口径对齐 verl 的 tests/checkpoint_engine/bench_delta_vllm.py。 + +和 delta_validation.py 不同，这里不做逐字节比较，只计时 manager.sync_weights()： +先做一次 seed 同步，再对每个变化比例各同步 --repeats 次。一个进程只测一个后端。 + +python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend delta_sharded +python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend full +""" +import argparse +from statistics import median +from time import perf_counter + +from cookbook.checkpoint_engine.delta_validation import build_components, log_progress +from twinkle.checkpoint_engine import CheckpointEngineManager, DeltaSyncConfig + + +def timed_sync(manager) -> float: + started = perf_counter() + manager.sync_weights() + return perf_counter() - started + + +def run_benchmark(args): + trainer, sampler = build_components(args) + if args.backend == 'delta_sharded': + manager = CheckpointEngineManager( + trainer, sampler, platform='GPU', mode='standalone', sync_backend='delta_sharded', + delta_config=DeltaSyncConfig(bucket_bytes=args.bucket_bytes, gather_batch_size=args.gather_batch_size)) + else: + manager = CheckpointEngineManager(trainer, sampler, platform='GPU', mode='standalone') + + def metrics(): + return manager.last_delta_metrics if args.backend == 'delta_sharded' else {} + + try: + seconds = timed_sync(manager) + log_progress(f'[seed] backend={args.backend} {seconds:.3f}s metrics={metrics()}') + + rows = [] + step = 0 + for ratio in args.ratios: + samples = [] + for _ in range(args.repeats): + step += 1 + trainer.change_fraction(ratio) + seconds = timed_sync(manager) + samples.append(seconds) + log_progress(f'[ratio={ratio}] step={step} sync={seconds:.3f}s metrics={metrics()}') + rows.append((ratio, samples, metrics())) + finally: + if args.backend == 'delta_sharded': + manager.close_sync_session() + + print(f'\nbackend={args.backend} model={args.model} trainer={args.trainer_gpus} rollout_tp={args.tp} dp={args.dp} ' + f'bucket={args.bucket_bytes >> 20}MB') + print(f"{'ratio':>8} {'median(s)':>10} {'min(s)':>8} {'max(s)':>8} last metrics") + for ratio, samples, last in rows: + print(f'{ratio:>8} {median(samples):>10.3f} {min(samples):>8.3f} {max(samples):>8.3f} {last}') + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--model', required=True, help='训练与 rollout 使用的同一模型，例如 Qwen3-8B') + parser.add_argument('--backend', choices=('delta_sharded', 'full'), default='delta_sharded') + parser.add_argument('--trainer-gpus', type=int, choices=(1, 2), default=2, help='FSDP2 训练卡数') + parser.add_argument('--tp', type=int, choices=(1, 2), default=2, help='rollout 副本的 GPU 数') + parser.add_argument('--dp', type=int, default=1, help='独立 rollout 副本数') + parser.add_argument('--strategy', choices=('native_fsdp', 'accelerate'), default='native_fsdp') + parser.add_argument('--cpu-offload', action='store_true', help='使用 Accelerate FSDP2 CPU offload') + parser.add_argument('--ratios', type=float, nargs='+', default=[0.005, 0.01, 0.02, 0.03]) + parser.add_argument('--repeats', type=int, default=2, help='每个比例的同步次数') + parser.add_argument('--bucket-bytes', type=int, default=2048 << 20, help='verl 默认 2048MB') + parser.add_argument('--gather-batch-size', type=int, default=32) + parser.add_argument('--devices-per-node', type=int, help='每个节点的 GPU 数，默认单节点运行') + args = parser.parse_args() + if args.cpu_offload and args.strategy != 'accelerate': + parser.error('--cpu-offload 需要 --strategy accelerate') + run_benchmark(args) + + +if __name__ == '__main__': + main() diff --git a/cookbook/checkpoint_engine/delta_benchmark_notes.md b/cookbook/checkpoint_engine/delta_benchmark_notes.md new file mode 100644 index 000000000..035bea797 --- /dev/null +++ b/cookbook/checkpoint_engine/delta_benchmark_notes.md @@ -0,0 +1,144 @@ +# Delta 权重同步实验记录 + +## 背景 + +- 起因：`delta_validation.py` 在 `--tp 2` 时崩溃，`--dp 2` 正常。 +- 修复后验证通过，但 delta 同步比 full 同步慢，甚至是负优化。 +- 目标：先在 verl 上跑出 baseline，再把 twinkle 的脚本对齐 verl，比较两边。 + +## 环境 + +- HAMi vGPU（libvgpu 拦截 `cuMemCreate`），所以绝对耗时不一定能外推到裸机。 +- trainer 2 卡（FSDP2），rollout 2 卡（vLLM TP=2）。 +- 模型：Qwen3-8B（`/nas/disk1/Qwen3-8B`）。 + +## 1. tp=2 崩溃 + +- 现象：NCCL 通信组初始化时报 `Illegal device id` 并 segfault。 +- 原因：HAMi 环境下 NCCL 的 cuMem 路径（`NCCL_CUMEM_ENABLE`、`NCCL_CUMEM_HOST_ENABLE`）不可用，tp>1 时 vLLM 才会初始化 NCCL，所以 dp=2 不受影响。 +- 处理：在 `delta_validation.py` 开头用 `os.environ.setdefault` 把两个变量设为 `0`，与 verl 的做法一致，外部可覆盖。 +- 结果：重跑后 status=passed。 + +## 2. twinkle 早期观察（Qwen3.5-0.8B，bucket 64MB） + +- delta 耗时随 flush 数增长。每个 flush 都要新建 sender/receiver、两次 NCCL broadcast、checksum 同步、Ray RPC 和 vLLM `collective_rpc`。 +- 把 bucket 调大后：seed 10.1s 降到 1.9s；50% 和 100% 变化的用例从 11.1s、16.7s 降到约 2.7s。 +- 但 1% 到 10% 变化仍然是 3.2 到 3.5s，full 只要 1.0 到 1.6s。0% 约 1.0s。 +- 结论：存在 2 到 3 秒的固定开销，与 flush 数无关。来源未定位。 + +## 3. 数据量原理 + +- 稀疏条目每个变化元素 6 字节（int32 位置 + bf16 值），full 同步每个元素 2 字节。 +- 变化比例超过约 1/3 时，delta 的传输量就比 full 大。 +- 真实 RL 单步变化约 1% 到 3%（verl 文档在 Qwen3-8B 上测得 0.884% 到 0.895%），所以比例改成 0.5%、1%、2%、3%。 + +## 4. 脚本 + +| 脚本 | 说明 | +|---|---| +| `verl/tests/checkpoint_engine/bench_delta_vllm.py` | verl baseline。FSDP2 trainer（`forward_only=True`，`model_dtype=bfloat16`），随机翻转 bf16 最低位模拟变化，计时整个 `update_weights()` | +| `cookbook/checkpoint_engine/delta_benchmark.py` | twinkle 对齐脚本。先 seed，再按比例同步，计时整个 `sync_weights()`，输出 median/min/max 表 | + +运行： + +``` +# verl +python tests/checkpoint_engine/bench_delta_vllm.py --model /nas/disk1/Qwen3-8B --backend delta_sharded +python tests/checkpoint_engine/bench_delta_vllm.py --model /nas/disk1/Qwen3-8B --backend nccl + +# twinkle +python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend delta_sharded +python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend full +``` + +verl 脚本调试中修过的问题： + +- 缺少 flash-attn：改用 `attn_implementation=sdpa`，并设置 `use_remove_padding=False`。 +- `perturb` 断言 bf16 失败：verl 的 `model_dtype` 默认是 `fp32`，需要显式设为 `bfloat16`。 +- `perturb` 是 blocking 调用，不能再套 `ray.get`。 + +## 5. verl baseline（Qwen3-8B，bucket 2048MB，每档 2 次） + +seed（delta_sharded）：256.9s，传输 15622MB，7 个 flush。 + +| 比例 | delta 中位数 (min~max) | nccl 中位数 (min~max) | delta 传输量 | +|---|---|---|---| +| 0.5% | 38.3s (33.9~42.6) | 21.9s (14.8~29.1) | 234MB | +| 1% | 27.6s (23.7~31.5) | 28.8s (27.4~30.2) | 469MB | +| 2% | 17.8s (15.7~19.8) | 7.3s (6.2~8.4) | 937MB | +| 3% | 34.8s (27.9~41.8) | 6.5s (5.7~7.3) | 1406MB | + +日志字段说明（以 0.5% 为例）： + +- `flipped` 和 `changed_elems` 相等，说明 delta 检测正确。 +- `payload_mbytes` = 变化元素数 × 6 字节，与预期吻合。 +- 变化量都小于 2048MB，所以每次只有 1 个 flush。 + +## 6. 结论 + +能确定： + +- verl 自己的 delta 在这个环境里也不比 nccl 快，负优化不是 twinkle 特有的。 +- nccl 每次传输量固定约 15.6GB，耗时却在 5.7s 到 30s 之间，环境本身波动约 5 倍。 +- delta 的耗时与传输量无关：传 937MB 最快，传 234MB 反而最慢。传 234MB 要 34 到 42s，这部分不是带宽造成的。 +- 2% 和 3% 档，nccl 比 delta 快 2.4 到 5 倍。 + +不能确定： + +- 每档只有 2 个样本，没有统计意义。 +- 固定开销花在哪一步未知。候选：HAMi 对 CUDA IPC 的影响、ZMQ 和 Ray RPC、暂停和恢复生成、释放和恢复 KV cache。均未测量。 +- 机器上是否有其他任务占用 GPU 或 NAS，未排查。 + +## 7. 已知的口径差异（twinkle 与 verl） + +- 变化位置：twinkle 的 `change_fraction` 翻转每个参数前 N% 元素的符号位，verl 随机翻转最低位。变化数量相同，位置分布不同。 +- CPU offload：verl 的 `forward_only=True` 会启用，twinkle 默认不启用。对齐需加 `--strategy accelerate --cpu-offload`。 +- 分片 dtype：verl 现在是 bf16；twinkle 取决于 transformers 版本的默认加载 dtype，未核实。传输字节数不受影响。 +- `delta_benchmark.py` 不做 `forbid_full_tensor` 检查。 + +## 8. twinkle 结果（Qwen3-8B，bucket 2048MB，tp=2，每档 2 次） + +| 比例 | delta 中位数 (min~max) | trainer `send_seconds` | full 中位数 (min~max) | delta 传输量 | +|---|---|---|---|---| +| 0.5% | 8.4s (7.2~9.6) | 0.44s | 1.3s (0.96~1.6) | 234MB | +| 1% | 9.8s (9.1~10.6) | 0.44s | 0.9s (0.83~1.0) | 469MB | +| 2% | 11.6s (11.2~12.1) | 0.44s | 1.2s (0.86~1.5) | 937MB | +| 3% | 10.0s (9.0~11.1) | 0.71s | 1.3s (0.80~1.7) | 1406MB | + +- 传输量和 verl 完全一致（0.5% 档 234.34MB、1 个 flush），说明编码和分桶等价。 +- trainer 侧（diff、gather、broadcast）只要 0.44~0.71s，不是瓶颈。 +- delta 比 full 多出的 7~11s 都在 trainer `send_seconds` 之外，且随传输量有增长趋势（0.5%→2% 约 4.7s/GB）；3% 比 2% 低，属于 2 样本噪声。 +- 嫌疑（均未测量）：`begin_delta_sync` 里 `sleep(level=1)` + `wake_up(['weights'])`（full 路径没有）、每个 flush 新建 ZMQ/IPC、接收端 `clone` + checksum `.item()` + 稀疏 patch、`_clear_delta_cache` 与 `wake_up(['kv_cache'])`。 + +## 9. verl 与 twinkle 对比 + +| 比例 | verl delta | verl nccl | twinkle delta | twinkle full | +|---|---|---|---|---| +| 0.5% | 38.3s | 21.9s | 8.4s | 1.3s | +| 1% | 27.6s | 28.8s | 9.8s | 0.9s | +| 2% | 17.8s | 7.3s | 11.6s | 1.2s | +| 3% | 34.8s | 6.5s | 10.0s | 1.3s | + +- 两边 delta 都比各自的全量同步慢，负优化不是 twinkle 特有的。 +- twinkle 的 delta 绝对耗时反而更低（8~12s 对 18~38s）；twinkle 的 delta/full 比例更差，主要因为 twinkle full 基线只有约 1s。 +- 计时范围不一致，不能直接下结论：verl `update_weights()` 包含终止请求、释放/恢复 KV cache、重建进程组、恢复生成（来自代码阅读，未逐行核对）；twinkle delta 含暂停生成、sleep/wake_up、清缓存、恢复生成；twinkle full 只有 `sync_weights()`，没有 sleep/wake_up，所以最轻。 +- 其他口径差异见第 7 节。 + +## 10. verl 官方文档的测法（`docs/advance/delta_weight_sync.md`，"Measured results"） + +- 没有独立 benchmark 脚本，数据来自真实 GSM8K GRPO 训练中每步稳态权重同步的计时，`delta_sharded` 与 `nccl` 各跑一遍。 +- 环境：H100 或 A100/A800 多节点，训练与 rollout 分离（多为 2+2 节点），SGLang rollout。部分行为 50 步中位数或 50 步持续运行。 +- 变化比例由训练自然产生（稠密模型约 0.884%~0.895%），不是人为设定。 +- 正确性另测：200 步 GRPO 奖励曲线对比、perturb/apply/revert/apply 位级一致。 +- 官方 Qwen3-8B 数据：2+2 节点 FSDP2 delta 2.24~3.40s，nccl 34~38s；1+1 节点 delta 3.00s，nccl 11.00s。Qwen3-0.6B 单机 NVLink 为 0.59s 对 0.61s（持平）。 +- 文档自己说明：快链路加小模型时，全量广播本来就便宜，delta 的优势随模型规模和网络距离增长。 +- 与我们测试的差异：官方是多节点、走网络、真实训练循环；我们是单机 4 张 HAMi 虚拟卡、合成扰动、微基准。因此单机上 delta 不占优不与文档矛盾。 +- 官方 verl delta 在 8B 上约 2.2~3.4s，我们的 verl delta 是 18~38s，twinkle 是 8~12s；这一差距的原因未查明（HAMi、单机部署、计时范围都有可能）。 + +## 11. 待办 + +- 决定：改为参照官方做法，用真实 GRPO 训练做 delta 与 full 的对比（用户已确认方向，具体方案待定）。 +- 跑之前检查 `nvidia-smi`，确认没有其他任务占卡。 +- 在 delta 路径加分阶段计时（trainer 的 diff/gather、broadcast、接收端 IPC 与 vLLM patch、前后的暂停恢复），定位 7~11s 的来源。涉及 `checkpoint_engine/delta/engine.py`、`delta/trainer.py` 和 `delta/rollout.py`，需先确认。 +- 可选：给 `delta_benchmark.py` 加 `--no-free-cache`（`free_cache_engine=False`），看 sleep/wake_up 占多少。 +- 可选：把 NCCL 环境变量移到库里（`vLLMSampler.__init__`），让所有使用者受益。 diff --git a/cookbook/checkpoint_engine/delta_validation.py b/cookbook/checkpoint_engine/delta_validation.py new file mode 100644 index 000000000..ae6d42c83 --- /dev/null +++ b/cookbook/checkpoint_engine/delta_validation.py @@ -0,0 +1,393 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""在 GPU(CUDA) 或 Ascend NPU 上验证 FSDP2 增量同步：逐字节比较 delta 与 full 同步后的 vLLM 权重。 + +示例（2 张训练卡 + 4 张 rollout 卡，GPU）： +python -m cookbook.checkpoint_engine.delta_validation --dp 2 --tp 2 --rounds 10 +示例（只有 2 张卡：1 张训练、1 张 rollout，FSDP 不分片）： +python -m cookbook.checkpoint_engine.delta_validation --trainer-gpus 1 --dp 1 --tp 1 --rounds 2 + +NPU（单机 4 卡、不联网，模型必须是本地目录；NPU 默认 --offline 且 --num-devices 4）。先用 --dry-run 检查计划： +(a) 1 训练 + 1 rollout(tp=1)： +python -m cookbook.checkpoint_engine.delta_validation --platform NPU --model-path /path/to/model \\ + --trainer-gpus 1 --dp 1 --tp 1 --rounds 2 +(b) 1 训练 + 1 rollout(tp=2)： --trainer-gpus 1 --dp 1 --tp 2 +(c) 2 训练 + 2 个 rollout 副本(tp=1)： --trainer-gpus 2 --dp 2 --tp 1 +(d) 2 训练 + 1 rollout(tp=2)： --trainer-gpus 2 --dp 1 --tp 2 +“多副本且 tp>1”至少需要 5 张卡，4 卡机器无法验证。NPU 路径尚未在真实 NPU 上跑过。 +""" +import argparse +import importlib.metadata +import json +import os +import subprocess +from contextlib import contextmanager +from datetime import datetime +from pathlib import Path +from time import perf_counter + +import torch +from torch.distributed.tensor import DTensor + +import twinkle +from cookbook.checkpoint_engine.device_utils import (add_platform_arguments, build_device_groups, build_engine_args, + delta_config_kwargs, device_peak_bytes, device_synchronize, + nproc_per_node, run_preflight_or_exit) +from twinkle import DeviceMesh, Platform, remote_class, remote_function +from twinkle.checkpoint_engine import CheckpointEngineManager, DeltaSyncConfig +from twinkle.checkpoint_engine.delta.spec import ShardSpec, derive_dtensor_placement, translate_flat_indices +from twinkle.data_format.sampling import SamplingParams +from twinkle.model import TransformersModel +from twinkle.sampler import vLLMSampler +from twinkle.sampler.vllm_sampler.vllm_worker_extension import TwinkleWorkerExtension + +# tp>1 时 vLLM 会初始化 NCCL 通信组；在 HAMi 等 vGPU 环境下，NCCL 的 cuMem* 路径会触发 +# "Illegal device id" 并 segfault。与 verl 一致，在 Ray actor 继承环境前关闭（可由外部覆盖）。 +os.environ.setdefault('NCCL_CUMEM_ENABLE', '0') +os.environ.setdefault('NCCL_CUMEM_HOST_ENABLE', '0') + +# 0% 用来测固定开销；其余贴近真实 RL 单步变化（稠密模型约 1%~3%）。 +CHANGE_RATIOS = (0., .005, .01, .02, .03) + + +def log_progress(message: str) -> None: + print(f'[{datetime.now():%Y-%m-%d %H:%M:%S}] {message}', flush=True) + + +@contextmanager +def validation_stage(name: str): + log_progress(f'开始：{name}') + started = perf_counter() + try: + yield + except Exception as error: + log_progress(f'失败：{name}，耗时 {perf_counter() - started:.3f}s，{error!r}') + raise + else: + log_progress(f'完成：{name}，耗时 {perf_counter() - started:.3f}s') + + +def environment_report(): + """记录软件版本和驱动信息，方便复现设备上的结果。""" + report = {} + for package in ('torch', 'torch_npu', 'vllm', 'vllm_ascend', 'transformers', 'accelerate', 'ray'): + try: + report[package] = importlib.metadata.version(package.replace('_', '-')) + except importlib.metadata.PackageNotFoundError: + report[package] = None + # GPU 机器有 nvidia-smi，Ascend 机器有 npu-smi；哪个存在就记录哪个。 + for tool, command in (('nvidia-smi', ['nvidia-smi']), ('npu-smi', ['npu-smi', 'info'])): + try: + report[tool] = subprocess.run(command, capture_output=True, text=True, timeout=15).stdout + except (FileNotFoundError, subprocess.TimeoutExpired): + report[tool] = None + return report + + +@remote_class() +class ValidationTrainer(TransformersModel): + def __init__(self, device_mesh: DeviceMesh = None, **kwargs): + TransformersModel.__init__.__wrapped__(self, device_mesh=device_mesh, **kwargs) + + @remote_function(dispatch='all', lazy_collect=False) + def change_fraction(self, ratio: float): + """翻转每个参数前 ratio 比例元素的符号位，产生可控的 BF16 字节变化。""" + self._lazy_wrap_model() + with torch.no_grad(): + for tensor in self.strategy.unwrap_model(self.model).parameters(): + place, _, _ = derive_dtensor_placement(ShardSpec.from_param(tensor)) + local = tensor.to_local() if isinstance(tensor, DTensor) else tensor + values = local.detach().to(torch.bfloat16).contiguous().reshape(-1) + indices = torch.arange(values.numel(), dtype=torch.int64, device=values.device) + # 使用完整参数的坐标，使各 FSDP shard 合起来恰好修改指定比例。 + changed = translate_flat_indices(indices, place) < int(tensor.numel() * ratio) + bits = values.view(torch.int16) + bits[changed] = bits[changed].bitwise_xor(-32768) + local.copy_(values.reshape(local.shape).to(device=local.device, dtype=local.dtype)) + + @remote_function(dispatch='all', lazy_collect=False) + def forbid_full_tensor(self): + # 在 trainer 进程里临时禁用全量聚合，检查稳态 delta 没有退回 full 路径。 + self._validation_full_tensor = DTensor.full_tensor + + def forbidden(*args, **kwargs): + raise AssertionError('稳态 delta 调用了 DTensor.full_tensor()') + + DTensor.full_tensor = forbidden + + @remote_function(dispatch='all', lazy_collect=False) + def allow_full_tensor(self): + DTensor.full_tensor = self._validation_full_tensor + + @remote_function(dispatch='all', lazy_collect=False) + def training_step(self): + """执行一次真实的前向、反向和 SGD 更新，补充人为修改之外的验证。""" + optimizer = torch.optim.SGD(self.model.parameters(), lr=0.001) + tokens = torch.tensor([[1, 2, 3, 4]], device=self._delta_device) + self.model.train() + optimizer.zero_grad() + loss = self.model(input_ids=tokens, labels=tokens).loss + loss.backward() + optimizer.step() + optimizer.zero_grad() + return float(loss.detach().cpu()) + + @remote_function(dispatch='all', lazy_collect=False) + def sync_device(self): + """等待本进程设备上已排队的 kernel 全部完成。 + + 端到端耗时基准用它隔开“训练计算”和“权重同步”：clip_grad_and_step 返回时优化器 kernel 可能还没跑完， + 不同步的话这部分时间会被算进随后的 sync。 + """ + device_synchronize(Platform.get_platform().device_prefix()) + + @remote_function(dispatch='all', lazy_collect=False) + def validation_environment(self): + return environment_report() + + +class ValidationWorker(TwinkleWorkerExtension): + """在每个 vLLM TP worker 内保存和比较真正参与推理的 runtime 权重。""" + + def validation_environment(self): + return environment_report() + + def capture_weights(self): + self._validation_weights = { + name: tensor.detach().cpu().clone() + for name, tensor in self.model_runner.model.state_dict().items()} + # self.device 是 torch.device('cuda:N') 或 torch.device('npu:N')，不要写死 torch.cuda。 + return {'device_peak_bytes': device_peak_bytes(torch.device(self.device).type)} + + def begin_full_reload(self): + # vLLM 可能打包或转置权重；全量基准也必须走完整 reload 生命周期。 + from vllm.model_executor.model_loader.reload.layerwise import ( + initialize_layerwise_reload, record_metadata_for_reloading) + with torch.device(self.device): + record_metadata_for_reloading(self.model_runner.model) + initialize_layerwise_reload(self.model_runner.model) + + def compare_weights(self): + from vllm.model_executor.model_loader.reload.layerwise import finalize_layerwise_reload + finalize_layerwise_reload(self.model_runner.model, self.model_runner.model_config) + state = self.model_runner.model.state_dict() + if state.keys() != self._validation_weights.keys(): + raise AssertionError('全量 reload 改变了 runtime 权重键') + mismatches = [] + for name, tensor in state.items(): + # 比较字节而非浮点近似值，包含正负零等位级差异。 + actual = tensor.detach().cpu().contiguous() + expected = self._validation_weights[name].contiguous() + if actual.shape != expected.shape or not torch.equal( + actual.reshape(-1).view(torch.uint8), expected.reshape(-1).view(torch.uint8)): + mismatches.append(self._describe_mismatch(name, actual, expected)) + self._validation_weights = None + if mismatches: + # 汇总全部不一致的 tensor，而不是在第一个处中止，便于判断是整体错误还是局部错误。 + shown = '\n '.join(mismatches[:20]) + more = f'\n ...另有 {len(mismatches) - 20} 个' if len(mismatches) > 20 else '' + raise AssertionError( + f'delta 与 full 权重不一致：rank={self.rank}, {len(mismatches)}/{len(state)} 个 tensor 不一致\n' + f' {shown}{more}') + return True + + @staticmethod + def _describe_mismatch(name, actual, expected): + """描述一个不一致 tensor：不一致元素占比、行范围，以及首个差异处的 full/delta 取值。 + + actual 是 full reload 后的运行时权重，expected 是 delta 同步后保存的快照。 + """ + if actual.shape != expected.shape: + return f'{name}: shape 不同 full={tuple(actual.shape)} delta={tuple(expected.shape)}' + int_dtype = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}[actual.element_size()] + differs = actual.reshape(-1).view(int_dtype) != expected.reshape(-1).view(int_dtype) + total = differs.numel() + count = int(differs.sum()) + first = int(differs.nonzero()[0]) + last = int(differs.nonzero()[-1]) + text = (f'{name}: shape={tuple(actual.shape)} 不一致元素 {count}/{total} ({count / total:.4%}), ' + f'flat 范围 [{first}, {last}]') + if actual.dim() >= 2: + rows = differs.reshape(actual.shape[0], -1).any(dim=1).nonzero().flatten() + text += (f', 不一致行 {rows.numel()}/{actual.shape[0]} ' + f'(行范围 [{int(rows[0])}, {int(rows[-1])}])') + full_value = actual.reshape(-1)[first].item() + delta_value = expected.reshape(-1)[first].item() + return f'{text}, 首个差异 flat={first}: full={full_value} delta={delta_value}' + + +@remote_class() +class ValidationSampler(vLLMSampler): + def __init__(self, device_mesh: DeviceMesh = None, **kwargs): + vLLMSampler.__init__.__wrapped__(self, device_mesh=device_mesh, **kwargs) + + @remote_function(dispatch='all', lazy_collect=False, collect='flatten') + def worker_rpc(self, method: str): + # 遍历全部 DP 副本及其 TP workers，任何 worker 的比较失败都会向上传递。 + return self._run_in_loop(self.engine.engine.collective_rpc(method)) + + @remote_function(dispatch='all', lazy_collect=False) + def pause_validation(self): + self._run_in_loop(self.engine.pause_delta_generation(self._delta_config.timeout_seconds)) + + @remote_function(dispatch='all', lazy_collect=False) + def generate_one_token(self): + return self._run_in_loop(self.engine.sample( + prompt=[1, 2, 3], sampling_params=SamplingParams(max_tokens=1, temperature=0.))) + + +WORKER_EXTENSION_CLS = 'cookbook.checkpoint_engine.delta_validation.ValidationWorker' + + +def build_components(args): + """训练使用 trainer_gpus 张卡的 FSDP2；rollout 使用 dp 个独立副本，每个副本占 tp 张卡。 + + GPU 与 NPU 共用同一套流程，差别（device_type、vLLM 的 weight_transfer 后端、NPU 的 rl_config 等） + 都在 device_utils 里按 --platform 解析。 + """ + trainer_gpus = args.trainer_gpus + rollout_gpus = args.dp * args.tp + with validation_stage(f'初始化 Ray 和 {args.platform} 资源组'): + twinkle.initialize(mode='ray', nproc_per_node=nproc_per_node(args), + groups=build_device_groups(args), lazy_collect=False) + with validation_stage(f'创建 {trainer_gpus} 卡 FSDP trainer'): + trainer = ValidationTrainer( + model_id=args.model, remote_group='trainer', strategy=args.strategy, mixed_precision='bf16', + fsdp_config={'cpu_offload': True} if args.cpu_offload else {}, + device_mesh=DeviceMesh.from_sizes(world_size=trainer_gpus, fsdp_size=trainer_gpus)) + with validation_stage(f'创建 rollout：{args.dp} 个副本，每个 {args.tp} 张 GPU'): + sampler = ValidationSampler( + model_id=args.model, remote_group='rollout', + device_mesh=DeviceMesh.from_sizes(world_size=rollout_gpus, tp_size=args.tp, dp_size=args.dp), + engine_args=build_engine_args(args, WORKER_EXTENSION_CLS)) + return trainer, sampler + + +def steady_sync(trainer, delta): + """稳态同步必须只传增量；无论是否成功，都恢复 full_tensor 方法。""" + with validation_stage('禁用 trainer 的全量权重聚合'): + trainer.forbid_full_tensor() + try: + started = perf_counter() + with validation_stage('稳态 delta 权重同步'): + delta.sync_weights() + if delta.last_delta_metrics['seed']: + raise AssertionError('稳态同步意外重新执行了全量 seed') + return perf_counter() - started + finally: + with validation_stage('恢复 trainer 的全量权重聚合'): + trainer.allow_full_tensor() + + +def compare_with_full(sampler, full): + """保存 delta 结果，再全量加载相同 trainer 权重，比较两条路径的最终结果。""" + with validation_stage('rollout 生成一个 token，检查同步后推理'): + sampler.generate_one_token() + with validation_stage('暂停 rollout 推理'): + sampler.pause_validation() + with validation_stage('保存所有 rollout worker 的 delta 权重快照'): + workers = sampler.worker_rpc('capture_weights') + with validation_stage('准备全量权重 reload'): + sampler.worker_rpc('begin_full_reload') + started = perf_counter() + with validation_stage('full 权重同步'): + full.sync_weights() + seconds = perf_counter() - started + with validation_stage('所有 rollout worker 逐字节比较 delta 与 full 权重'): + sampler.worker_rpc('compare_weights') + return {'full_seconds': seconds, 'workers': workers} + + +def run_validation(args): + log_progress( + f'启动验证：platform={args.platform}，model={args.model}，strategy={args.strategy}，' + f'trainer={args.trainer_gpus} 张卡，rollout dp={args.dp}、tp={args.tp}，' + f'rounds={args.rounds}，report={Path(args.report).resolve()}') + trainer, sampler = build_components(args) + with validation_stage('创建 delta 和 full checkpoint manager'): + delta = CheckpointEngineManager( + trainer, sampler, platform=args.platform, mode='standalone', sync_backend='delta_sharded', + delta_config=DeltaSyncConfig(**delta_config_kwargs(args))) + full = CheckpointEngineManager(trainer, sampler, platform=args.platform, mode='standalone') + report = {'status': 'running', 'arguments': vars(args), 'rounds': []} + + def save_report(): + Path(args.report).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding='utf-8') + + def check_round(stage): + with validation_stage(stage): + seconds = steady_sync(trainer, delta) + result = compare_with_full(sampler, full) + result.update(stage=stage, delta_seconds=seconds, metrics=delta.last_delta_metrics) + report['rounds'].append(result) + save_report() + log_progress(f'{stage}：字节比较通过，delta={seconds:.3f}s，full={result["full_seconds"]:.3f}s') + + try: + with validation_stage('收集 driver 和 trainer 环境信息'): + report.update(driver=environment_report(), trainers=trainer.validation_environment()) + # 1. 首次发送 dense seed 并建立 CPU shard 快照，后续各轮复用同一个通信组。 + with validation_stage('首次全量 seed 同步，建立增量快照和通信组'): + delta.sync_weights() + with validation_stage('收集 rollout worker 环境信息'): + report['workers'] = sampler.worker_rpc('validation_environment') + with validation_stage('首次 seed 权重验证'): + report['seed'] = compare_with_full(sampler, full) + + # 2. 轮流覆盖零变化、稀疏变化和全部变化；full 基准不修改 trainer 快照。 + for step in range(args.rounds): + ratio = CHANGE_RATIOS[step % len(CHANGE_RATIOS)] + with validation_stage(f'第 {step + 1}/{args.rounds} 轮：修改 {ratio:.1%} 权重'): + trainer.change_fraction(ratio) + check_round(f'第 {step + 1}/{args.rounds} 轮验证，修改比例 {ratio:.1%}') + + # 3. 真实训练更新后的 delta 仍必须与 full 逐字节一致。 + with validation_stage('真实训练：前向、反向和 SGD 更新'): + report['training_loss'] = trainer.training_step() + log_progress(f'训练 loss={report["training_loss"]}') + check_round('真实训练更新') + + report['status'] = 'passed' + except Exception as error: + report.update(status='failed', error=repr(error)) + log_progress(f'验证失败：{error!r}') + raise + finally: + # 4. 关闭接收 actors、通信组和快照；即使验证失败也保存报告。 + try: + with validation_stage('关闭增量同步会话，释放通信组和快照'): + delta.close_sync_session() + except Exception as error: + report.update(status='failed', cleanup_error=repr(error)) + raise + finally: + with validation_stage(f'保存验证报告：{args.report}'): + save_report() + log_progress(f'验证结束：status={report["status"]}') + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_platform_arguments(parser) # --platform {GPU,NPU}、--model-path、--offline、--num-devices、--dry-run 等 + parser.add_argument('--model', default='Qwen/Qwen3-0.6B', help='训练与 rollout 使用的同一模型') + parser.add_argument('--trainer-gpus', type=int, choices=(1, 2), default=2, help='FSDP2 训练卡数；1 表示不分片') + parser.add_argument('--tp', type=int, choices=(1, 2), default=1, help='每个 rollout 副本的卡数') + parser.add_argument('--dp', type=int, default=1, help='独立 rollout 副本数') + parser.add_argument('--strategy', choices=('native_fsdp', 'accelerate'), default='native_fsdp') + parser.add_argument('--cpu-offload', action='store_true', help='使用 Accelerate FSDP2 CPU offload') + parser.add_argument('--rounds', type=int, default=1, help='不同变化比例的验证轮数') + parser.add_argument('--gather-batch-size', type=int, default=32) + parser.add_argument('--bucket-bytes', type=int, default=64 << 20) + parser.add_argument('--devices-per-node', type=int, help='每个节点的卡数，默认单节点运行') + parser.add_argument('--report', default='delta-validation.json', help='结果报告路径') + args = parser.parse_args() + if args.dp < 1 or args.rounds < 1: + parser.error('--dp 和 --rounds 必须为正数') + if args.cpu_offload and args.strategy != 'accelerate': + parser.error('--cpu-offload 需要 --strategy accelerate') + # 解析模型/离线设置、预检卡数预算并打印计划；--dry-run 在这里退出，不启动 Ray。 + run_preflight_or_exit(args, parser, WORKER_EXTENSION_CLS) + run_validation(args) + + +if __name__ == '__main__': + main() diff --git a/cookbook/checkpoint_engine/device_utils.py b/cookbook/checkpoint_engine/device_utils.py new file mode 100644 index 000000000..e166f0ded --- /dev/null +++ b/cookbook/checkpoint_engine/device_utils.py @@ -0,0 +1,232 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""GPU/NPU 通用的 delta 同步脚本辅助函数（delta_validation.py / delta_benchmark.py 共用）。 + +这里的大部分函数是纯 Python（不启动 Ray、不接触设备），因此可以在没有加速卡的机器上做 --dry-run 和单测。 +""" +from __future__ import annotations + +import argparse +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +DEFAULT_NPU_CARDS = 4 + +# vLLM-Ascend 在 RL 场景下必须带 rl_config（它会关闭 FRACTAL_NZ，否则 start_weight_update / wake_up 报 weight_nz_mode 错误）。 +NPU_ADDITIONAL_CONFIG = {'rl_config': {'enabled': True}} + +TOKENIZER_FILES = ('tokenizer.json', 'tokenizer_config.json', 'tokenizer.model', 'vocab.json', 'merges.txt') + +OFFLINE_ENV = { + 'HF_HUB_OFFLINE': '1', + 'TRANSFORMERS_OFFLINE': '1', + 'HF_DATASETS_OFFLINE': '1', + 'MODELSCOPE_OFFLINE': '1', # 尽力而为：不同 modelscope 版本的离线开关不一致，真正的保护是本地目录检查 +} + +NPU_REQUIREMENTS = ( + "vLLM-Ascend 需要 additional_config={'rl_config': {'enabled': True}}（脚本已自动加入）；否则 " + 'start_weight_update/wake_up 会抛 weight_nz_mode 的 ValueError。', + '分离部署（trainer 与 rollout 使用不同的卡）通常不需要 sleep mode：NPU 默认 free_cache_engine=False；' + '若用 --free-cache-engine 打开，需要从源码编译的 vllm-ascend。', + "sampler 的 engine_args 必须带 weight_transfer_config={'backend': 'twinkle_delta_hccl'}（CUDA 为 twinkle_delta_ipc）。", + 'vLLM 必须提供 engine-owned WeightTransferEngine API（vLLM #44353）和 checkpoint patch API（vLLM #50723）；' + '可先运行 cookbook/checkpoint_engine/check_npu_env.py 检查。', + '支持范围：仅 vLLM、BF16 稠密（非 MoE）、非量化、无推测解码、vLLM 内部 data_parallel_size=1、' + 'trainer FSDP2 且 TP=PP=EP=1、无 LoRA、不支持 rebuild_group；rollout 可有多个副本（replica 为主、tp 为次）。', + 'GPU 与 NPU 的其余配置接口一致。以上内容尚未在真实 NPU 上验证。', +) + + +@dataclass(frozen=True) +class PlatformSpec: + name: str # twinkle Platform 名称，同时用作 DeviceGroup.device_type 与 CheckpointEngineManager(platform=) + torch_device: str # torch 设备模块名：torch.cuda / torch.npu + backend: str # vLLM weight_transfer_config 的 backend + + +PLATFORM_SPECS = { + 'GPU': PlatformSpec('GPU', 'cuda', 'twinkle_delta_ipc'), + 'NPU': PlatformSpec('NPU', 'npu', 'twinkle_delta_hccl'), +} +_PLATFORM_ALIASES = {'GPU': 'GPU', 'CUDA': 'GPU', 'NPU': 'NPU', 'ASCEND': 'NPU'} + + +def normalize_platform(value: str) -> str: + """把 gpu/cuda/npu/ascend（不区分大小写）归一化为 'GPU' 或 'NPU'。""" + try: + return _PLATFORM_ALIASES[str(value).upper()] + except KeyError: + raise argparse.ArgumentTypeError(f'未知平台 {value!r}，可选 GPU 或 NPU') from None + + +def add_platform_arguments(parser: argparse.ArgumentParser) -> None: + """两个脚本共用的 GPU/NPU、离线模型与 dry-run 参数。""" + parser.add_argument('--platform', type=normalize_platform, default='GPU', metavar='{GPU,NPU}', + help='设备平台（默认 GPU；cuda 是 GPU 的别名）') + parser.add_argument('--model-path', help='本地模型目录（不联网）；给出后覆盖 --model') + parser.add_argument('--offline', action=argparse.BooleanOptionalAction, default=None, + help='禁止任何联网下载（设置 HF_HUB_OFFLINE 等）；NPU 默认开启，GPU 默认关闭') + parser.add_argument('--num-devices', type=int, + help='可用卡数，用于预检卡数预算；NPU 默认 4，GPU 默认不限制') + parser.add_argument('--free-cache-engine', action=argparse.BooleanOptionalAction, default=None, + help='DeltaSyncConfig.free_cache_engine；GPU 默认开启，NPU 默认关闭（需要 sleep mode）') + parser.add_argument('--use-token-ids', action=argparse.BooleanOptionalAction, default=None, + help='让 vLLM 使用 skip_tokenizer_init，只用 token id 生成并比较；模型目录没有 tokenizer 时自动开启') + parser.add_argument('--dry-run', action='store_true', + help='只解析参数、构造 DeviceGroup/配置并做卡数预检，打印执行计划后退出（不启动 Ray、不接触设备）') + + +def has_tokenizer(model_dir: str) -> bool: + return any((Path(model_dir) / name).is_file() for name in TOKENIZER_FILES) + + +def resolve_model(args) -> None: + """解析 --model-path / --offline / --use-token-ids，原地更新 args。""" + if args.model_path: + args.model = args.model_path + if args.offline is None: + args.offline = args.platform == 'NPU' + if args.num_devices is None and args.platform == 'NPU': + args.num_devices = DEFAULT_NPU_CARDS + if args.free_cache_engine is None: + args.free_cache_engine = args.platform == 'GPU' + if args.use_token_ids is None: + local = Path(args.model).is_dir() + args.use_token_ids = bool(local and not has_tokenizer(args.model)) + + +def apply_offline_env(args) -> dict[str, str]: + """--offline 时设置离线环境变量（必须在启动 Ray 之前调用，Ray 子进程会继承）。""" + if not args.offline: + return {} + for key, value in OFFLINE_ENV.items(): + os.environ[key] = value + return dict(OFFLINE_ENV) + + +def preflight(args) -> tuple[list[str], list[str]]: + """返回 (errors, warnings)。不接触设备。""" + errors: list[str] = [] + warnings: list[str] = [] + trainer, dp, tp = args.trainer_gpus, args.dp, args.tp + if dp < 1 or tp < 1 or trainer < 1: + errors.append('--trainer-gpus、--dp、--tp 必须为正数') + return errors, warnings + needed = trainer + dp * tp + if args.num_devices is not None and needed > args.num_devices: + errors.append( + f'卡数预算不足：trainer {trainer} + rollout dp*tp={dp}*{tp}={dp * tp} = {needed} 张，' + f'超过 --num-devices={args.num_devices}。4 卡可验证的布局：' + '(a) --trainer-gpus 1 --dp 1 --tp 1；(b) --trainer-gpus 1 --dp 1 --tp 2；' + '(c) --trainer-gpus 2 --dp 2 --tp 1；(d) --trainer-gpus 2 --dp 1 --tp 2') + if args.platform == 'NPU' and dp > 1 and tp > 1 and (args.num_devices or DEFAULT_NPU_CARDS) <= DEFAULT_NPU_CARDS: + warnings.append(f'rollout 副本数>1 且 tp>1（dp={dp}, tp={tp}）至少需要 {1 + dp * tp} 张卡，' + f'4 卡机器无法验证“多副本 + tp>1”') + model_dir = Path(args.model) + if args.offline and not model_dir.is_dir(): + errors.append(f'--offline 下模型必须是本地目录（--model-path），当前为 {args.model!r}') + if model_dir.is_dir(): + if not (model_dir / 'config.json').is_file(): + errors.append(f'模型目录 {args.model} 缺少 config.json') + if not has_tokenizer(args.model) and not args.use_token_ids: + errors.append(f'模型目录 {args.model} 没有 tokenizer 文件，请加 --use-token-ids') + if args.platform == 'NPU' and args.free_cache_engine: + warnings.append('NPU 上启用 free_cache_engine 需要 sleep mode（已加入 enable_sleep_mode=True），' + '且需要从源码编译的 vllm-ascend') + if args.platform == 'NPU' and args.strategy == 'accelerate' and args.cpu_offload: + warnings.append('NPU 上的 Accelerate FSDP2 CPU offload 未验证') + return errors, warnings + + +def build_engine_args(args, worker_extension_cls: str) -> dict[str, Any]: + spec = PLATFORM_SPECS[args.platform] + engine_args: dict[str, Any] = dict( + tensor_parallel_size=args.tp, dtype='bfloat16', max_model_len=256, enforce_eager=True, + gpu_memory_utilization=0.4, worker_extension_cls=worker_extension_cls, + weight_transfer_config={'backend': spec.backend}) + if args.platform == 'NPU': + engine_args['additional_config'] = {'rl_config': dict(NPU_ADDITIONAL_CONFIG['rl_config'])} + if args.free_cache_engine: + engine_args['enable_sleep_mode'] = True + if args.use_token_ids: + engine_args['skip_tokenizer_init'] = True + return engine_args + + +def delta_config_kwargs(args) -> dict[str, Any]: + return dict(bucket_bytes=args.bucket_bytes, gather_batch_size=args.gather_batch_size, + free_cache_engine=bool(args.free_cache_engine)) + + +def build_device_groups(args) -> list: + from twinkle import DeviceGroup + spec = PLATFORM_SPECS[args.platform] + trainer, rollout = args.trainer_gpus, args.dp * args.tp + return [ + DeviceGroup(name='trainer', ranks=list(range(trainer)), device_type=spec.name), + DeviceGroup(name='rollout', ranks=list(range(trainer, trainer + rollout)), device_type=spec.name, + gpus_per_worker=args.tp), + ] + + +def nproc_per_node(args) -> int: + return args.devices_per_node or args.trainer_gpus + args.dp * args.tp + + +def print_plan(args, worker_extension_cls: str) -> None: + """打印解析后的执行计划（--dry-run 以及真实运行前都会打印）。""" + from twinkle.checkpoint_engine import DeltaSyncConfig + spec = PLATFORM_SPECS[args.platform] + print(f'== 执行计划 platform={args.platform} (torch.{spec.torch_device}) ==') + print(f'model : {args.model} (offline={args.offline}, use_token_ids={args.use_token_ids})') + print(f'cards : trainer={args.trainer_gpus} + rollout dp*tp={args.dp}*{args.tp}=' + f'{args.dp * args.tp} / num_devices={args.num_devices}') + for group in build_device_groups(args): + print(f'device group : {group}') + print(f'nproc_per_node : {nproc_per_node(args)}') + print(f'manager : CheckpointEngineManager(platform={args.platform!r}, mode=\'standalone\', ' + f"sync_backend='delta_sharded')") + print(f'backend : {spec.backend}') + print(f'engine_args : {build_engine_args(args, worker_extension_cls)}') + print(f'DeltaSyncConfig: {DeltaSyncConfig(**delta_config_kwargs(args))}') + if args.offline: + print(f'offline env : {OFFLINE_ENV}') + + +def run_preflight_or_exit(args, parser: argparse.ArgumentParser, worker_extension_cls: str) -> None: + """解析 -> 预检 -> 打印计划；--dry-run 时在此退出（返回前不接触设备）。""" + resolve_model(args) + errors, warnings = preflight(args) + if args.platform == 'NPU': + print('== NPU 使用须知（均未在真实 NPU 上验证）==') + for line in NPU_REQUIREMENTS: + print(f' - {line}') + for warning in warnings: + print(f'[警告] {warning}') + if errors: + parser.error('预检失败：\n ' + '\n '.join(errors)) + print_plan(args, worker_extension_cls) + if args.dry_run: + print('--dry-run：未启动 Ray，未接触设备，退出。') + raise SystemExit(0) + apply_offline_env(args) + + +def torch_device_module(platform_or_type: str): + """返回 torch.cuda 或 torch.npu（按 'GPU'/'NPU'/'cuda'/'npu' 选择）；NPU 会先 import torch_npu。""" + import torch + name = platform_or_type.lower() + name = PLATFORM_SPECS[normalize_platform(name)].torch_device if name in ('gpu', 'npu', 'ascend') else name + if name == 'npu': + import torch_npu # noqa: F401 + return getattr(torch, name) + + +def device_peak_bytes(device_type: str) -> int: + return int(torch_device_module(device_type).max_memory_allocated()) + + +def device_synchronize(device_type: str) -> None: + torch_device_module(device_type).synchronize() diff --git a/cookbook/checkpoint_engine/e2e_step_benchmark.py b/cookbook/checkpoint_engine/e2e_step_benchmark.py new file mode 100644 index 000000000..d9f92d94e --- /dev/null +++ b/cookbook/checkpoint_engine/e2e_step_benchmark.py @@ -0,0 +1,257 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""端到端“一个训练 step”耗时对比：delta_sharded vs full。 + +一个 step = 训练计算（前向 + 反向 + AdamW，forward_backward + clip_grad_and_step）+ 权重同步（manager.sync_weights，阻塞到 rollout 就绪）。 +一个进程只测一个后端，各自写 JSON 报告，最后用 --compare 合并成对比表（不启动 Ray）。 + + # 1. 分别跑两个后端（两次之间重启 Ray 集群/actors，保证干净） + python -m cookbook.checkpoint_engine.e2e_step_benchmark --platform NPU --model-path /nas/disk1/Qwen3.5-9B \\ + --backend delta_sharded --trainer-gpus 2 --dp 1 --tp 2 --devices-per-node 2 \\ + --bucket-bytes 2147483648 --seq-len 128 --report e2e-delta.json + python -m cookbook.checkpoint_engine.e2e_step_benchmark --platform NPU --model-path /nas/disk1/Qwen3.5-9B \\ + --backend full --trainer-gpus 2 --dp 1 --tp 2 --devices-per-node 2 \\ + --bucket-bytes 2147483648 --seq-len 128 --report e2e-full.json + # 2. 合并对比 + python -m cookbook.checkpoint_engine.e2e_step_benchmark --compare e2e-delta.json e2e-full.json --md e2e.md + +默认只测真实 step：训练一步（前向 + 反向 + AdamW）后立刻 sync，变化比例由 AdamW 自然产生 +（报告里记录 delta 的实际 changed_ratio）。首次 seed 同步和 --warmup-steps 指定的预热步只记录、不进统计。 +""" +import argparse +import json +from datetime import datetime +from pathlib import Path +from statistics import median +from time import perf_counter + +from cookbook.checkpoint_engine.device_utils import add_platform_arguments, delta_config_kwargs, run_preflight_or_exit + + +def log_progress(message: str) -> None: + print(f'[{datetime.now():%Y-%m-%d %H:%M:%S}] {message}', flush=True) + + +def delta_info(manager) -> dict: + """从 manager.last_delta_metrics 里取 delta 的实际变化比例和 payload；full 后端返回 {}。""" + metrics = getattr(manager, 'last_delta_metrics', None) + if not metrics: + return {} + info = {'manager_seconds': metrics.get('total_seconds')} + trainers = metrics.get('trainers') + for item in trainers if isinstance(trainers, list) else [trainers]: + # 只有 trainer rank 0 带 checkpoint_engine 指标，其余 rank 为 None。 + checkpoint = item.get('checkpoint_engine') if isinstance(item, dict) else None + if isinstance(checkpoint, dict): + info['changed_ratio'] = checkpoint.get('checkpoint_engine/changed_ratio') + info['payload_mbytes'] = checkpoint.get('checkpoint_engine/payload_mbytes') + info['send_seconds'] = item.get('send_seconds') + break + return info + + +def timed_sync(manager) -> tuple[float, dict]: + started = perf_counter() + manager.sync_weights() + seconds = perf_counter() - started + return seconds, delta_info(manager) + + +def build_batch(batch_size: int, seq_len: int) -> list[dict]: + """构造已编码的最小训练 batch（InputFeature 列表）：全部 token 参与 loss，无需 tokenizer / 数据集。""" + row = [(i % 1000) + 1 for i in range(seq_len)] + return [{'input_ids': list(row), 'attention_mask': [1] * seq_len, 'labels': list(row)} + for _ in range(batch_size)] + + +def resolve(result): + """lazy_collect 打开时远程调用返回 callable，需要调用一次才会真正等待结果。""" + return result() if callable(result) else result + + +def train_step(trainer, batch, max_grad_norm: float) -> float: + """一次真实训练：前向 + 反向 + AdamW 更新（接口与 cookbook/transformers 一致）。返回耗时（秒）。 + + 结尾的 sync_device 用来等设备上的优化器 kernel 跑完，避免这部分时间漏进随后的 sync。 + """ + started = perf_counter() + resolve(trainer.forward_backward(inputs=batch)) + resolve(trainer.clip_grad_and_step(max_grad_norm=max_grad_norm)) + resolve(trainer.sync_device()) + return perf_counter() - started + + +def run_benchmark(args) -> None: + # 延迟导入：--compare 不需要 torch / Ray。 + from cookbook.checkpoint_engine.delta_validation import build_components, validation_stage + from twinkle.checkpoint_engine import CheckpointEngineManager, DeltaSyncConfig + + log_progress(f'e2e step benchmark：backend={args.backend}, platform={args.platform}, model={args.model}, ' + f'trainer={args.trainer_gpus}, rollout dp={args.dp} tp={args.tp}, ' + f'batch={args.batch_size} seq={args.seq_len}') + trainer, sampler = build_components(args) + if args.backend == 'delta_sharded': + manager = CheckpointEngineManager( + trainer, sampler, platform=args.platform, mode='standalone', sync_backend='delta_sharded', + delta_config=DeltaSyncConfig(**delta_config_kwargs(args))) + else: + manager = CheckpointEngineManager(trainer, sampler, platform=args.platform, mode='standalone') + + report = {'status': 'running', 'backend': args.backend, 'arguments': vars(args), 'seed': None, 'real_steps': []} + + def save_report(): + Path(args.report).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding='utf-8') + + try: + # 必须先 set_optimizer 再做首次同步：同步会触发 FSDP 懒包装，包装时才会把优化器绑到分片参数上。 + trainer.set_optimizer('AdamW', lr=args.lr, weight_decay=args.weight_decay, foreach=False) + batch = build_batch(args.batch_size, args.seq_len) + + # 首次 sync：delta 为 dense seed（建立快照和通信组），full 为首次全量；只记录，不进入统计。 + with validation_stage('首次 seed 同步（不计入统计）'): + seconds, info = timed_sync(manager) + report['seed'] = {'sync_seconds': seconds, **info} + log_progress(f'[seed] sync={seconds:.3f}s {info}') + save_report() + + # ---- real：训练计算 + 同步，变化比例由 AdamW 自然产生 ---- + total_real = args.warmup_steps + args.real_steps + for index in range(total_real): + warmup = index < args.warmup_steps + train_seconds = train_step(trainer, batch, args.max_grad_norm) + sync_seconds, info = timed_sync(manager) + row = {'index': index, 'warmup': warmup, 'train_seconds': train_seconds, 'sync_seconds': sync_seconds, + 'step_seconds': train_seconds + sync_seconds, **info} + report['real_steps'].append(row) + tag = 'warmup' if warmup else 'real' + log_progress(f'[{tag} {index + 1}/{total_real}] train={train_seconds:.3f}s sync={sync_seconds:.3f}s ' + f'step={row["step_seconds"]:.3f}s {info}') + save_report() + + report['status'] = 'passed' + except Exception as error: + report.update(status='failed', error=repr(error)) + log_progress(f'基准失败：{error!r}') + raise + finally: + try: + if args.backend == 'delta_sharded': + with validation_stage('关闭增量同步会话，释放通信组和快照'): + manager.close_sync_session() + finally: + save_report() + log_progress(f'报告已保存：{Path(args.report).resolve()}（status={report["status"]}）') + + print_single_summary(report) + + +def _med(rows, key): + values = [row[key] for row in rows if row.get(key) is not None] + return median(values) if values else None + + +def _fmt(value, spec='.3f', suffix=''): + return '-' if value is None else f'{value:{spec}}{suffix}' + + +def print_single_summary(report) -> None: + measured = [row for row in report['real_steps'] if not row['warmup']] + print(f"\nbackend={report['backend']} real steps（中位数，{len(measured)} 步，已去掉 warmup）") + print(f" train={_fmt(_med(measured, 'train_seconds'))}s sync={_fmt(_med(measured, 'sync_seconds'))}s " + f"step={_fmt(_med(measured, 'step_seconds'))}s changed_ratio={_fmt(_med(measured, 'changed_ratio'), '.4f')} " + f"payload={_fmt(_med(measured, 'payload_mbytes'), '.1f', 'MB')}") + + +def compare_reports(paths, md_path=None) -> None: + reports = {} + for path in paths: + data = json.loads(Path(path).read_text(encoding='utf-8')) + if data.get('backend') in reports: + raise SystemExit(f'--compare 需要一份 delta_sharded 和一份 full 报告，但 {path} 的 backend 重复') + if data.get('status') != 'passed': + log_progress(f'警告：{path} 的 status={data.get("status")}，结果可能不完整') + reports[data['backend']] = data + if set(reports) != {'delta_sharded', 'full'}: + raise SystemExit(f'--compare 需要 delta_sharded 与 full 各一份报告，实际: {sorted(reports)}') + delta, full = reports['delta_sharded'], reports['full'] + + lines = ['# 端到端训练 step 耗时：delta vs full', ''] + args = delta['arguments'] + lines += [f"环境：platform={args['platform']}，model={args['model']}，trainer={args['trainer_gpus']} 卡，" + f"rollout dp={args['dp']} × tp={args['tp']}，batch={args['batch_size']} × seq={args['seq_len']}，" + f"bucket={args['bucket_bytes'] >> 20}MB。", + '一个 step = 训练计算（前向+反向+AdamW）+ 权重同步（阻塞到 rollout 就绪）。`speedup = full step / delta step`，' + '>1 表示 delta 更快。', ''] + + def real_rows(report): + return [row for row in report['real_steps'] if not row['warmup']] + + d_real, f_real = real_rows(delta), real_rows(full) + lines += ['## 真实 step（AdamW 自然产生的变化比例）', '', + '| backend | train (s) | sync (s) | step (s) | changed_ratio | payload (MB) |', + '| --- | --- | --- | --- | --- | --- |'] + for name, rows in (('delta_sharded', d_real), ('full', f_real)): + lines.append(f"| {name} | {_fmt(_med(rows, 'train_seconds'))} | {_fmt(_med(rows, 'sync_seconds'))} | " + f"{_fmt(_med(rows, 'step_seconds'))} | {_fmt(_med(rows, 'changed_ratio'), '.4f')} | " + f"{_fmt(_med(rows, 'payload_mbytes'), '.1f')} |") + d_step, f_step = _med(d_real, 'step_seconds'), _med(f_real, 'step_seconds') + if d_step and f_step: + lines += ['', f'speedup（real）= {f_step / d_step:.2f}x（full {f_step:.3f}s / delta {d_step:.3f}s）'] + + text = '\n'.join(lines) + '\n' + print(text) + if md_path: + Path(md_path).write_text(text, encoding='utf-8') + log_progress(f'对比表已写入：{Path(md_path).resolve()}') + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--compare', nargs=2, metavar=('REPORT_A', 'REPORT_B'), + help='合并 delta_sharded 与 full 两份 JSON 报告并输出对比表（不启动 Ray）') + parser.add_argument('--md', help='--compare 时把对比表另存为 markdown') + add_platform_arguments(parser) + parser.add_argument('--model', default='Qwen/Qwen3-0.6B', help='训练与 rollout 使用的同一模型') + parser.add_argument('--backend', choices=('delta_sharded', 'full'), default='delta_sharded') + parser.add_argument('--trainer-gpus', type=int, choices=(1, 2), default=2, help='FSDP2 训练卡数；1 表示不分片') + parser.add_argument('--tp', type=int, choices=(1, 2), default=1, help='每个 rollout 副本的卡数') + parser.add_argument('--dp', type=int, default=1, help='独立 rollout 副本数') + parser.add_argument('--strategy', choices=('native_fsdp', 'accelerate'), default='native_fsdp') + parser.add_argument('--cpu-offload', action='store_true', help='使用 Accelerate FSDP2 CPU offload') + parser.add_argument('--gather-batch-size', type=int, default=32) + parser.add_argument('--bucket-bytes', type=int, default=2048 << 20, help='默认 2048MB，与 verl 一致') + parser.add_argument('--devices-per-node', type=int, help='每个节点的卡数，默认单节点运行') + parser.add_argument('--batch-size', type=int, default=None, + help='全局 batch size，按 DP 切分给 trainer 各卡，必须能被 --trainer-gpus 整除；默认等于 --trainer-gpus') + parser.add_argument('--seq-len', type=int, default=128, + help='每条样本的 token 数。全参数 AdamW 的优化器状态很占显存（9B 模型每卡约 30~50GB），' + '默认取较小值，OOM 时调小或加 --strategy accelerate --cpu-offload') + parser.add_argument('--lr', type=float, default=1e-5, help='AdamW 学习率；它决定 real step 自然产生的变化比例') + parser.add_argument('--weight-decay', type=float, default=0.01) + parser.add_argument('--max-grad-norm', type=float, default=1.0) + parser.add_argument('--warmup-steps', type=int, default=1, help='real step 的预热步数（不进入统计）') + parser.add_argument('--real-steps', type=int, default=3, help='参与统计的 real step 数') + parser.add_argument('--report', default=None, help='JSON 报告路径，默认 e2e-.json') + args = parser.parse_args() + + if args.compare: + compare_reports(args.compare, args.md) + return + if args.report is None: + args.report = f'e2e-{args.backend}.json' + if args.dp < 1 or args.real_steps < 0 or args.warmup_steps < 0: + parser.error('--dp 必须为正数，--real-steps/--warmup-steps 不能为负') + if args.batch_size is None: + args.batch_size = args.trainer_gpus + if args.batch_size < 1 or args.seq_len < 1: + parser.error('--batch-size 和 --seq-len 必须为正数') + if args.batch_size % args.trainer_gpus: + parser.error(f'--batch-size={args.batch_size} 必须能被 --trainer-gpus={args.trainer_gpus} 整除（forward_backward 按 DP 切分）') + if args.cpu_offload and args.strategy != 'accelerate': + parser.error('--cpu-offload 需要 --strategy accelerate') + # 复用 delta_validation 的 worker extension 做预检（延迟导入常量所在模块会拉起 torch，所以直接写字符串）。 + run_preflight_or_exit(args, parser, 'cookbook.checkpoint_engine.delta_validation.ValidationWorker') + run_benchmark(args) + + +if __name__ == '__main__': + main() diff --git a/delta_perf.md b/delta_perf.md new file mode 100644 index 000000000..6a8c8496e --- /dev/null +++ b/delta_perf.md @@ -0,0 +1,69 @@ +# Delta vs Full 权重同步性能 + +测试环境：Qwen3.5-9B（BF16，整模型 17.53 GB / 760 个张量），2 卡 FSDP2 trainer + 2 卡 rollout（TP=2），跨机 HCCL，bucket=2048MB。 +`delta time` 为发送侧传输耗时（算 diff + 跨 rank gather-v + 编码 + HCCL 广播），`full payload` 为每轮全量同步实际广播的字节数。 + +| changed_ratio | delta payload | delta time | full payload | reduction vs full | +| ------------- | ------------- | ---------- | ------------ | ----------------- | +| 0.5% | 269.2 MB | 3.59 s | 17.53 GB | 66.7x | +| 1.0% | 538.4 MB | 4.31 s | 17.53 GB | 33.3x | +| 2.0% | 1076.9 MB | 5.62 s | 17.53 GB | 16.7x | +| 3.0% | 1615.3 MB | 6.76 s | 17.53 GB | 11.1x | + +说明： +- `reduction vs full` = full payload / delta payload，即 delta 相比全量同步减少的网络传输量。 +- full payload 每轮恒为整模型 17.53 GB（本环境链路 ~9.5 GB/s，全量广播约 1.8 s）。 +- delta payload 每个变化元素 6 字节（int32 位置 + bf16 值），随改动比例线性增长。 +- delta time 可近似分解为「固定开销 ~2.8 s（diff 全扫 + gather + 逐参数）+ payload / ~400 MB/s」。 + +# 端到端单 step 对比（Qwen3.5-4B） + +环境：2 卡 FSDP2 trainer + 2 卡 rollout（TP=2），batch=2 × seq=2048，AdamW lr=1e-5，bucket=2048MB。 +一个 step = 训练计算（前向 + 反向 + AdamW）+ 权重同步（阻塞到 rollout 就绪）。 + +| backend | train (s) | sync (s) | step (s) | changed_ratio | payload (MB) | +| ------------- | --------- | --------- | --------- | ------------- | ------------ | +| delta_sharded | 3.670 | 36.998 | 40.927 | 14.11% | 4179.3 | +| full | 3.677 | 3.098 | 6.775 | - | 9870.5 | + +此工作点下 delta 慢于 full（speedup 0.17x）：payload 只降 2.4x，而稀疏路径吞吐约为 dense 广播的 1/10 ~ 1/20。 +改动比例由优化器配置与训练阶段决定，不是 delta 实现的属性；上游参考实现在稳态 RL 下报的是 1 ~ 2.3%。 + +## 耗时分解（取上表那一步，sync 37.0 s） + +发送端 wire_time 与接收端 apply 是流水重叠的（flushes=2），因此各项之和大于 37 s。 + +| 阶段 | 耗时 (s) | 占 sync | +| ---------------------------- | -------- | ------- | +| 接收端 apply（decode+scatter） | 27.0 | 73% | +| 发送端 broadcast | 11.0 | 30% | +| assemble（编码 + checksum） | 5.2 | 14% | +| produce（diff 全扫 + 提取） | 4.7 | 13% | +| gather（跨 rank） | 2.5 | 7% | +| 接收端 broadcast_recv | 0.1 | ~0% | + +两个大头： +- **接收端稀疏 scatter**：4179MB / 27.0s ≈ 155 MB/s。按索引散写是随机访存，9B 上实测约 110 MB/s，同一量级，属稀疏路径的固有吞吐。 +- **发送端 broadcast**：4179MB / 11.0s ≈ 380 MB/s，而 full 的连续广播为 9870MB / 3.1s ≈ 3200 MB/s，慢约 8 倍。delta 把 positions 与 values 分两次广播，且每个 flush 后 `torch.npu.synchronize()`；full 是单块连续 buffer 一次推完。payload 小且单 flush 时（9B 的 1GB 场景）broadcast 可忽略，payload 大且多 flush 时才显现。 + +## 改动比例随训练下降 + +同一份固定数据反复训练，梯度收敛，更新量递减： + +| step | changed_ratio | payload (MB) | sync (s) | step (s) | +| ------ | ------------- | ------------ | -------- | -------- | +| warmup | 21.64% | 6408.5 | 56.1 | 64.2 | +| real 2 | 15.83% | 4686.8 | 42.1 | 45.6 | +| real 3 | 14.11% | 4179.3 | 37.0 | 40.9 | +| real 4 | 12.65% | 3744.4 | 31.2 | 34.9 | + +14.11% 不是稳态值，只是该合成训练设置的早期状态。 + +## 优化优先级 + +1. 接收端 scatter（73%）：合并 kernel、连续段走整段 copy、避免 int64 索引中转。 +2. 发送端 broadcast（30%）：positions 与 values 合并为一次广播、去掉每 flush 的强制同步、增大 bucket 以减少 flush 数。 +3. assemble 的 checksum（14%）：verify 关闭时跳过。 +4. produce 的 diff 全扫（13%）：按 optimizer 标记的 dirty 参数增量化，避免每次 O(模型) 扫描。 +5. gather（7%）：此规模无需改动。 + diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py index f1e259b26..c0caea509 100644 --- a/src/twinkle/checkpoint_engine/delta/__init__.py +++ b/src/twinkle/checkpoint_engine/delta/__init__.py @@ -13,7 +13,8 @@ spec.py ShardSpec / BlockPlacement: where a local shard sits in the full tensor export.py snapshot, diff and HF-coordinate delta export helpers trainer.py DeltaTrainerMixin: trainer-side Ray lifecycle - rollout.py CudaDeltaReceiver / DeltaSamplerMixin: rollout-side lifecycle + rollout.py CudaDeltaReceiver / DeltaSamplerMixin: rollout-side lifecycle (CUDA: Ray receiver actors + IPC; + NPU: no actors, vLLM workers join the trainer's HCCL group via the worker extension) backends/ per-training-backend export hooks (base.py contract, fsdp.py implementation) """ from .config import DeltaSyncConfig diff --git a/src/twinkle/checkpoint_engine/delta/config.py b/src/twinkle/checkpoint_engine/delta/config.py index ded758a6f..3bc980f08 100644 --- a/src/twinkle/checkpoint_engine/delta/config.py +++ b/src/twinkle/checkpoint_engine/delta/config.py @@ -37,17 +37,23 @@ def __post_init__(self): class DeltaTransportEndpoint(TypedDict): world_size: int master_metadata: DeltaMasterMetadata - rollout_workers: list[DeltaWorkerMetadata] + # 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): diff --git a/src/twinkle/checkpoint_engine/delta/hccl_engine.py b/src/twinkle/checkpoint_engine/delta/hccl_engine.py index 607dcfdcc..02b8ec455 100644 --- a/src/twinkle/checkpoint_engine/delta/hccl_engine.py +++ b/src/twinkle/checkpoint_engine/delta/hccl_engine.py @@ -31,6 +31,7 @@ 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 @@ -74,6 +75,8 @@ def __init__( 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: @@ -89,6 +92,8 @@ def prepare(self) -> MasterMetadata | None: 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): @@ -138,6 +143,7 @@ def close(self): 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 diff --git a/src/twinkle/checkpoint_engine/delta/hccl_receiver.py b/src/twinkle/checkpoint_engine/delta/hccl_receiver.py index 37bc2c741..9b8c92ff5 100644 --- a/src/twinkle/checkpoint_engine/delta/hccl_receiver.py +++ b/src/twinkle/checkpoint_engine/delta/hccl_receiver.py @@ -11,6 +11,7 @@ from __future__ import annotations import json +import os import torch from collections.abc import Callable @@ -20,8 +21,16 @@ _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): - """Join the trainer's HCCL broadcast group at ``group_rank``.""" +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 diff --git a/src/twinkle/checkpoint_engine/delta/pipeline.py b/src/twinkle/checkpoint_engine/delta/pipeline.py index f25b3c31e..05154be9f 100644 --- a/src/twinkle/checkpoint_engine/delta/pipeline.py +++ b/src/twinkle/checkpoint_engine/delta/pipeline.py @@ -467,6 +467,7 @@ async def send_weights( 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 @@ -516,7 +517,15 @@ def _bucket_slot_delta( if not is_r0: return self._release_after_round("steady") - logger.info("delta-sharded send v=%s delta flushes=%d (streamed)", global_steps, n_flushes) + 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 { diff --git a/src/twinkle/checkpoint_engine/delta/rollout.py b/src/twinkle/checkpoint_engine/delta/rollout.py index ad57eb7c0..16c748e48 100644 --- a/src/twinkle/checkpoint_engine/delta/rollout.py +++ b/src/twinkle/checkpoint_engine/delta/rollout.py @@ -1,5 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Rollout-side Ray lifecycle and receivers for the CUDA delta checkpoint engine.""" +"""Rollout-side lifecycle for the delta checkpoint engine. + +CUDA: Ray receiver actors colocated with each vLLM worker feed it over IPC. Ascend NPU: no receiver actors; +the vLLM workers join the trainer's HCCL group themselves and the sampler only drives vLLM's lifecycle RPCs. +""" from __future__ import annotations import asyncio @@ -9,7 +13,7 @@ import torch -from twinkle import remote_function +from twinkle import Platform, remote_function from .config import (DeltaReceiverPlacement, DeltaSyncConfig, DeltaTransportEndpoint, DeltaWorkerCapabilities, DeltaWorkerMetadata) from .engine import DeltaShardedCheckpointEngine @@ -81,10 +85,36 @@ def close(self) -> None: 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( @@ -139,6 +169,11 @@ def validate_delta_sampler(self, config: DeltaSyncConfig) -> DeltaWorkerCapabili 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) @@ -175,6 +210,11 @@ def _spawn_delta_receiver( @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')) @@ -203,8 +243,41 @@ async def _send_delta_flush(self) -> None: 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: @@ -236,6 +309,13 @@ def close_delta_sampler(self) -> None: 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) diff --git a/src/twinkle/checkpoint_engine/delta/trainer.py b/src/twinkle/checkpoint_engine/delta/trainer.py index 58028c77f..3140ddfdc 100644 --- a/src/twinkle/checkpoint_engine/delta/trainer.py +++ b/src/twinkle/checkpoint_engine/delta/trainer.py @@ -66,7 +66,11 @@ def prepare_delta_sync(self, config: DeltaSyncConfig) -> DeltaTrainerMetadata: self._lazy_wrap_model() if self._delta_engine is None: self._initialize_delta_engine(config) - return {'device': self._delta_device, 'master_metadata': self._delta_master_metadata} + 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: diff --git a/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py b/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py index 33db57640..7cd9180f9 100644 --- a/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py +++ b/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py @@ -386,8 +386,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]]: @@ -474,5 +475,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 a3191f4bb..c27ca438d 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -93,9 +93,11 @@ def __init__( 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_engine_cls = resolve_delta_engine_cls(Platform.get_platform(platform).__name__) + 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 @@ -231,6 +233,8 @@ def sync_weights(self, merge_and_sync=True): 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') @@ -240,6 +244,23 @@ def _prepare_delta_transport(self): 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') diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index ffd4e1a42..8031d74ec 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -96,7 +96,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 8fd6530e6..820273c44 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -2202,6 +2202,35 @@ def _normalize(name: str, keep_base_layer: bool) -> str: name = name.replace('.base_layer', '') return name + 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() + def _add_base_layer_suffix(name): # vLLM (enable_lora) wraps some modules as ``*WithLoRA`` and exposes # only their ``.base_layer.*`` param, even when PEFT does not target @@ -2238,7 +2267,7 @@ 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 = _checkpoint_tensor(tensor) name = _normalize(name, keep_base_layer=False) @@ -2275,7 +2304,7 @@ def weight_generator(): 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 = _checkpoint_tensor(tensor) name = _normalize(name, keep_base_layer=False) diff --git a/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py b/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py index 8c06736d7..a9c0ae78c 100644 --- a/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py +++ b/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py @@ -6,7 +6,7 @@ import inspect import json from collections.abc import Iterable, Iterator -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import torch @@ -274,8 +274,18 @@ def finish_weight_update(self) -> None: except BaseException: self._update_failed = True raise - # Sparse updates write to initialized runtime tensors, so there is - # nothing to rebuild after the last flush. + 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.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: @@ -314,6 +324,8 @@ class TwinkleDeltaHCCLInitInfo(WeightTransferInitInfo): 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 @@ -347,6 +359,9 @@ def _device_error_message(self) -> str: 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): @@ -361,7 +376,7 @@ def _ensure_channel(self, group_rank: int): 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 + rv.dist_ip, rv.dist_port, group_rank, rv.world_size, self.device, hccl_env=rv.hccl_env ) except BaseException: self._close_sub() diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index fc8f0435c..3998a2df6 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -382,11 +382,17 @@ def _create_engine(self): from .delta_weight_transfer 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_weight_transfer 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 == 'twinkle_delta_ipc' and 'weight_transfer_config' not in valid_args: + 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, diff --git a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py index ab9c3495c..8486f68f3 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py @@ -154,17 +154,40 @@ def delta_receiver_metadata(self) -> 'DeltaWorkerMetadata': '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}]``) is forwarded as is; - ``WorkerBase.update_weights`` is expected to pick the entry for this worker's rank. + ``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(update_info) + 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_weight_transfer 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 diff --git a/src/twinkle/sampler/vllm_sampler/weight_update_utils.py b/src/twinkle/sampler/vllm_sampler/weight_update_utils.py index 4469ec401..56138c2c9 100644 --- a/src/twinkle/sampler/vllm_sampler/weight_update_utils.py +++ b/src/twinkle/sampler/vllm_sampler/weight_update_utils.py @@ -103,3 +103,32 @@ def copy_weight_loader(param, loaded_weight): 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/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 fb76f1293..6f0bb5da6 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 fd3aa15d3..1d4a2d33f 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..6219eb55f --- /dev/null +++ b/tests/checkpoint_engine/_hccl_loopback.py @@ -0,0 +1,369 @@ +# 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 import hccl_engine, hccl_receiver + 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.hccl_engine import DeltaHCCLCheckpointEngine + from twinkle.checkpoint_engine.delta.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.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.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_hccl.py b/tests/checkpoint_engine/test_delta_hccl.py index 5aff17948..deb6c6d52 100644 --- a/tests/checkpoint_engine/test_delta_hccl.py +++ b/tests/checkpoint_engine/test_delta_hccl.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import os import subprocess import sys from types import SimpleNamespace @@ -135,6 +136,22 @@ def test_prepare_master_idempotent(env): 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): 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 index 62f40e7a1..f641f128a 100644 --- a/tests/checkpoint_engine/test_delta_hccl_receiver.py +++ b/tests/checkpoint_engine/test_delta_hccl_receiver.py @@ -257,6 +257,17 @@ def test_build_communicator_args(monkeypatch): 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): @@ -269,13 +280,15 @@ def fake_connect(ip, port, topic): order.append(('connect', ip, port, topic)) return sub - def fake_build(ip, port, rank, world, device): + 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) + return SimpleNamespace(adapter=adapter, order=order, sub=sub, comm=comm, hccl_envs=hccl_envs) def hccl_engine(adapter, device='cpu'): @@ -317,6 +330,16 @@ def test_ensure_channel_order_reuse_and_rank_change(hccl): 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): @@ -336,8 +359,10 @@ def test_receive_sparse_session_prepares_sparse_loaders(hccl, monkeypatch): utils = sys.modules.get('twinkle.sampler.vllm_sampler.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.weight_update_utils', - SimpleNamespace(prepare_sparse_weight_loaders=prepare, drop_tied_alias_patches=drop)) + 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)]) @@ -347,6 +372,7 @@ def test_receive_sparse_session_prepares_sparse_loaders(hccl, monkeypatch): 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() @@ -488,14 +514,48 @@ def test_validate_delta_worker_reports_npu_and_names_both_backends(worker_cls, a 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 = SimpleNamespace(update_weights=Mock(), delta_ipc_endpoint=lambda: 'ipc:///x') + 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) - worker.update_weights.assert_called_once_with(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 @@ -511,3 +571,53 @@ def test_engine_pieces_do_not_import_npu_modules(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_npu_orchestration.py b/tests/checkpoint_engine/test_delta_npu_orchestration.py new file mode 100644 index 000000000..d0658622f --- /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 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_weight_transfer', 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_vllm.py b/tests/checkpoint_engine/test_delta_vllm.py index 4a25b20e3..556b1124e 100644 --- a/tests/checkpoint_engine/test_delta_vllm.py +++ b/tests/checkpoint_engine/test_delta_vllm.py @@ -122,6 +122,64 @@ def test_sparse_update_uses_native_patch_loader_without_reload(adapter): 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.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.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()) From f8237bf47e7ea6cfca761c9442eff86fb0c3edd1 Mon Sep 17 00:00:00 2001 From: kevssim Date: Wed, 21 Oct 2026 15:53:27 +0800 Subject: [PATCH 10/12] wip: drop cookbook and perf notes --- cookbook/checkpoint_engine/check_npu_env.py | 199 --------- .../checkpoint_engine/check_weight_names.py | 64 --- .../checkpoint_engine/crossnode_comm_check.py | 247 ----------- cookbook/checkpoint_engine/delta_benchmark.py | 83 ---- .../delta_benchmark_notes.md | 144 ------- .../checkpoint_engine/delta_validation.py | 393 ------------------ cookbook/checkpoint_engine/device_utils.py | 232 ----------- .../checkpoint_engine/e2e_step_benchmark.py | 257 ------------ delta_perf.md | 69 --- 9 files changed, 1688 deletions(-) delete mode 100644 cookbook/checkpoint_engine/check_npu_env.py delete mode 100644 cookbook/checkpoint_engine/check_weight_names.py delete mode 100644 cookbook/checkpoint_engine/crossnode_comm_check.py delete mode 100644 cookbook/checkpoint_engine/delta_benchmark.py delete mode 100644 cookbook/checkpoint_engine/delta_benchmark_notes.md delete mode 100644 cookbook/checkpoint_engine/delta_validation.py delete mode 100644 cookbook/checkpoint_engine/device_utils.py delete mode 100644 cookbook/checkpoint_engine/e2e_step_benchmark.py delete mode 100644 delta_perf.md diff --git a/cookbook/checkpoint_engine/check_npu_env.py b/cookbook/checkpoint_engine/check_npu_env.py deleted file mode 100644 index 088b230ad..000000000 --- a/cookbook/checkpoint_engine/check_npu_env.py +++ /dev/null @@ -1,199 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""检查当前环境是否满足 delta_sharded（HCCL/NPU 或 IPC/CUDA）增量权重同步的依赖。 - -只做导入和签名检查，不启动 Ray、不接触加速卡，可在任一节点单独运行： - - python -m cookbook.checkpoint_engine.check_npu_env # 自动探测平台 - python -m cookbook.checkpoint_engine.check_npu_env --platform NPU - -每项打印 [PASS]/[FAIL]/[WARN]，退出码 0 表示所有必需项通过。换机器或换镜像后先跑它， -可以在拉起几分钟的 vLLM 之前就发现 vLLM / vllm-ascend 版本不兼容的问题。 -""" -from __future__ import annotations - -import argparse -import importlib -import importlib.metadata -import inspect - - -class Checker: - def __init__(self) -> None: - self.failed = 0 - self.warned = 0 - - def ok(self, name: str, detail: str = '') -> None: - print(f'[PASS] {name}' + (f'：{detail}' if detail else '')) - - def fail(self, name: str, detail: str) -> None: - self.failed += 1 - print(f'[FAIL] {name}：{detail}') - - def warn(self, name: str, detail: str) -> None: - self.warned += 1 - print(f'[WARN] {name}：{detail}') - - def check(self, name: str, fn, *, required: bool = True) -> object: - """运行 fn；返回 (detail_or_value)。约定：fn 返回字符串即 PASS 详情，抛异常即失败。""" - try: - detail = fn() - except Exception as error: # noqa: BLE001 - 汇总所有失败，不在第一项中止 - (self.fail if required else self.warn)(name, f'{type(error).__name__}: {error}') - return None - self.ok(name, detail if isinstance(detail, str) else '') - return detail - - -def _version(package: str) -> str: - try: - return importlib.metadata.version(package) - except importlib.metadata.PackageNotFoundError: - return 'not-installed' - - -def check_versions(checker: Checker) -> None: - for package in ('torch', 'torch_npu', 'vllm', 'vllm_ascend', 'transformers'): - version = _version(package.replace('_', '-')) - print(f' {package}: {version}') - - -def check_weight_transfer_engine(checker: Checker) -> None: - def fn(): - from vllm.distributed.weight_transfer.base import WeightTransferEngine - params = inspect.signature(WeightTransferEngine.__init__).parameters - missing = {'config', 'vllm_config', 'device', 'model'} - set(params) - if missing: - raise RuntimeError(f'WeightTransferEngine.__init__ 缺少参数 {sorted(missing)}（需要 vLLM #44353）') - return 'engine-owned WeightTransferEngine API 可用（vLLM #44353）' - - checker.check('WeightTransferEngine API (#44353)', fn) - - -def check_checkpoint_patch_api(checker: Checker) -> None: - def fn(): - module = importlib.import_module('vllm.model_executor.model_loader.checkpoint_weight_patch') - for attr in ('CheckpointWeightPatch', 'load_checkpoint_weight_patches'): - if not hasattr(module, attr): - raise RuntimeError(f'缺少 {attr}（需要 vLLM #50723）') - sig = inspect.signature(module.load_checkpoint_weight_patches) - if 'validate_unique_indices' not in sig.parameters: - raise RuntimeError('load_checkpoint_weight_patches 缺少 validate_unique_indices 参数') - return 'checkpoint patch API 可用（vLLM #50723）' - - checker.check('checkpoint patch API (#50723)', fn) - - -def check_layerwise_reload(checker: Checker) -> None: - def fn(): - module = importlib.import_module('vllm.model_executor.model_loader.reload') - for attr in ('initialize_layerwise_reload', 'finalize_layerwise_reload'): - if not hasattr(module, attr): - raise RuntimeError(f'缺少 {attr}（dense/full 同步依赖 layerwise reload）') - return 'layerwise reload 生命周期可用' - - checker.check('layerwise reload API', fn) - - -def check_stateless_pg(checker: Checker) -> None: - def fn(): - import dataclasses - from vllm.distributed.utils import StatelessProcessGroup - fields = {f.name for f in dataclasses.fields(StatelessProcessGroup)} - # torch_utils.stateless_init_process_group 已兼容有/无 socket 两种版本，这里只提示。 - return f"StatelessProcessGroup 可导入（socket 字段：{'有' if 'socket' in fields else '无，走兼容分支'}）" - - checker.check('StatelessProcessGroup', fn) - - -def check_npu_hccl(checker: Checker) -> None: - def torch_npu_available(): - import torch - import torch_npu # noqa: F401 - if not torch.npu.is_available(): - raise RuntimeError('torch.npu.is_available() 为 False（本节点看不到 NPU）') - return f'torch.npu 可用，device_count={torch.npu.device_count()}' - - checker.check('torch_npu 设备', torch_npu_available) - - def pyhccl(): - from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator # noqa: F401 - return 'vllm_ascend PyHcclCommunicator 可导入（delta HCCL broadcast 依赖）' - - checker.check('vllm-ascend PyHccl', pyhccl) - - def rl_config(): - # RL 场景必须支持 additional_config={'rl_config': {'enabled': True}}；仅检查配置类存在。 - from vllm_ascend import ascend_config # noqa: F401 - return 'vllm_ascend.ascend_config 可导入（rl_config 开关所在模块）' - - checker.check('vllm-ascend rl_config 模块', rl_config, required=False) - - -def check_cuda_ipc(checker: Checker) -> None: - def fn(): - import torch - if not torch.cuda.is_available(): - raise RuntimeError('torch.cuda.is_available() 为 False') - return f'torch.cuda 可用，device_count={torch.cuda.device_count()}' - - checker.check('torch.cuda 设备', fn) - - -def check_twinkle_delta(checker: Checker) -> None: - def fn(): - from twinkle.sampler.vllm_sampler.delta_weight_transfer import require_vllm_delta_support - require_vllm_delta_support() - return 'twinkle require_vllm_delta_support() 通过' - - checker.check('twinkle delta 依赖自检', fn) - - def refresh(): - from twinkle.sampler.vllm_sampler.weight_update_utils import refresh_derived_weights # noqa: F401 - return 'refresh_derived_weights 可用（sparse 同步后刷新 Ascend 打包权重）' - - checker.check('twinkle refresh_derived_weights', refresh, required=False) - - -def detect_platform() -> str: - try: - import torch - import torch_npu # noqa: F401 - if torch.npu.is_available(): - return 'NPU' - except Exception: # noqa: BLE001 - pass - return 'GPU' - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('--platform', choices=('GPU', 'NPU'), help='默认自动探测（能导入 torch_npu 且 NPU 可用则 NPU）') - args = parser.parse_args() - platform = args.platform or detect_platform() - - print(f'== 依赖检查 platform={platform} ==') - checker = Checker() - check_versions(checker) - print('-- vLLM 通用 delta 依赖 --') - check_weight_transfer_engine(checker) - check_checkpoint_patch_api(checker) - check_layerwise_reload(checker) - check_stateless_pg(checker) - print(f'-- {platform} 专用依赖 --') - if platform == 'NPU': - check_npu_hccl(checker) - else: - check_cuda_ipc(checker) - print('-- twinkle 侧 --') - check_twinkle_delta(checker) - - print('== 汇总 ==') - if checker.failed: - print(f'{checker.failed} 项必需检查未通过，{checker.warned} 项告警。请先解决 [FAIL] 再运行 delta_validation。') - return 1 - print(f'全部必需检查通过，{checker.warned} 项告警。') - return 0 - - -if __name__ == '__main__': - raise SystemExit(main()) diff --git a/cookbook/checkpoint_engine/check_weight_names.py b/cookbook/checkpoint_engine/check_weight_names.py deleted file mode 100644 index 61867ca9a..000000000 --- a/cookbook/checkpoint_engine/check_weight_names.py +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""打印 Qwen3.5 的权重名映射，用于定位 full 同步的 tied embedding 报错。 - -不加载权重（meta device 建骨架），不启动 Ray，只看名字： - python -m cookbook.checkpoint_engine.check_weight_names --model-path /nas/disk1/Qwen3.5-4B -""" -import argparse - -import torch -import transformers -from transformers import AutoConfig, AutoModelForCausalLM - -from twinkle.checkpoint_engine.delta.backends.fsdp import convert_weight_keys - -KEYWORDS = ('embed_tokens', 'lm_head') - - -def main() -> int: - parser = argparse.ArgumentParser() - parser.add_argument('--model-path', required=True) - args = parser.parse_args() - - config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True) - architectures = getattr(config, 'architectures', None) - print(f'architectures : {architectures}') - text_config = getattr(config, 'text_config', config) - print(f'tie_word_embeddings : {getattr(config, "tie_word_embeddings", None)} ' - f'(text_config: {getattr(text_config, "tie_word_embeddings", None)})') - - # 与 TransformersModel 一致：优先用 config.architectures[0]，而不是 AutoModelForCausalLM - # （两者命名空间不同：ForConditionalGeneration 带 language_model. 前缀）。 - model_cls = AutoModelForCausalLM - if architectures: - model_cls = getattr(transformers, architectures[0], AutoModelForCausalLM) - print(f'model_cls (同 trainer): {getattr(model_cls, "__name__", model_cls)}') - - with torch.device('meta'): - if hasattr(model_cls, 'from_config'): - model = model_cls.from_config(config, trust_remote_code=True) - else: - model = model_cls._from_config(config) - print(f'model class : {type(model).__name__}') - mapping = getattr(model, '_checkpoint_conversion_mapping', None) - print(f'_checkpoint_conversion_mapping: {mapping}') - - state_dict = model.state_dict() - print(f'\nstate_dict 里的相关键（trainer 侧原始名，共 {len(state_dict)} 个权重）：') - for name in state_dict: - if any(k in name for k in KEYWORDS): - print(f' {name}') - - converted = convert_weight_keys(state_dict, model) - print(f'\nconvert_weight_keys 之后（实际发给 vLLM 的名字）：') - for name in converted: - if any(k in name for k in KEYWORDS): - print(f' {name}') - - print('\nvLLM 报错期望的是 language_model.model.embed_tokens.weight；' - '若上面没有这个名字，就是导出侧少了 language_model. 前缀。') - return 0 - - -if __name__ == '__main__': - raise SystemExit(main()) diff --git a/cookbook/checkpoint_engine/crossnode_comm_check.py b/cookbook/checkpoint_engine/crossnode_comm_check.py deleted file mode 100644 index b24cd023f..000000000 --- a/cookbook/checkpoint_engine/crossnode_comm_check.py +++ /dev/null @@ -1,247 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""最小化的两机通信连通性测试，不依赖 twinkle / vLLM / Ray。 - -用途：当 delta_validation 跨机卡在 HCCL 建链时，用它判断是“网络本身不通”还是“HCCL 跨机不通”。 -分两个后端测： - - gloo：纯 TCP，验证两机三层网络 + MASTER_PORT 能否 rendezvous（不碰 NPU）。 - - hccl：昇腾集合通信，验证 NPU 跨机能否真正建链并做 all_reduce。 - -两个 Pod 各跑一条命令（先 gloo 再 hccl）。MASTER 必须是 rank 0 所在 Pod 的 IP。 - - # Pod1（rank 0，172.61.8.99） - python -m cookbook.checkpoint_engine.crossnode_comm_check \ - --backend gloo --rank 0 --world-size 2 --master-addr 172.61.8.99 --master-port 29555 - # Pod2（rank 1，172.61.10.44） - python -m cookbook.checkpoint_engine.crossnode_comm_check \ - --backend gloo --rank 1 --world-size 2 --master-addr 172.61.8.99 --master-port 29555 - -gloo 通过后把 --backend 换成 hccl，--master-port 换一个新端口（如 29556）再各跑一次。 -gloo 通、hccl 不通，基本可判定：该环境的 NPU 不具备跨机 HCCL（通常是缺 RDMA/RoCE）。 -""" -from __future__ import annotations - -import argparse -import datetime -import os -import socket - - -def log(message: str) -> None: - print(f'[{datetime.datetime.now():%H:%M:%S}] [{socket.gethostname()}] {message}', flush=True) - - -def _run_stateless(args, torch, device) -> int: - """用 delta 完全相同的 pyhccl 路径（stateless_init_process_group）建链并做 all_reduce。 - - 这正是 delta 超时的那条调用栈： - stateless_init_process_group -> PyHcclCommunicator(pg, device) -> __init__ 里的 all_reduce。 - 若此处也超时，说明问题在 pyhccl 初始化路径本身（与并发组/时间错位无关）； - 若此处 PASS，说明裸 pyhccl 单组可通，delta 的超时来自并发的第二个 HCCL 组或 join 时间错位。 - """ - from twinkle.utils import stateless_init_process_group - - log(f'开始 stateless_init_process_group（pyhccl）：rank={args.rank}/{args.world_size}, ' - f'master={args.master_addr}:{args.master_port}, device={device.index}') - log('这条路径 = delta 的 build_delta_hccl_communicator；PyHcclCommunicator.__init__ 内部会做一次 all_reduce') - try: - comm = stateless_init_process_group( - master_address=args.master_addr, - master_port=args.master_port, - rank=args.rank, - world_size=args.world_size, - device=device.index, - backend='hccl', - ) - except Exception as error: # noqa: BLE001 - import traceback - log(f'FAIL：stateless_init_process_group 抛异常：{type(error).__name__}: {error}') - traceback.print_exc() - return 1 - log('stateless_init_process_group 返回成功（communicator 已建立，内部 all_reduce 已通过）') - - try: - tensor = torch.ones(1, device=device) - # PyHcclCommunicator.all_reduce 是 out-of-place：结果在返回值里，不改入参。 - out = comm.all_reduce(tensor) - torch.npu.synchronize() - value = int(out.item()) - log(f'comm.all_reduce 完成，结果={value}（期望={args.world_size}）') - if value != args.world_size: - log('FAIL：all_reduce 结果不等于 world_size，通信异常') - return 1 - log('PASS：pyhccl 跨机单组通信正常') - except Exception as error: # noqa: BLE001 - import traceback - log(f'FAIL：comm.all_reduce 抛异常：{type(error).__name__}: {error}') - traceback.print_exc() - return 1 - return 0 - - -def _build_pyhccl_group(args, torch, device, master_port, label): - """用 delta 的 pyhccl 路径建一个组并做一次 all_reduce，返回 (comm, 是否正常)。""" - from twinkle.utils import stateless_init_process_group - - log(f'[{label}] stateless_init_process_group：master_port={master_port}, ' - f"HCCL_IF_BASE_PORT={os.environ.get('HCCL_IF_BASE_PORT')}, " - f"NPU_RANGE={os.environ.get('HCCL_NPU_SOCKET_PORT_RANGE')}") - comm = stateless_init_process_group( - master_address=args.master_addr, master_port=master_port, - rank=args.rank, world_size=args.world_size, device=device.index, backend='hccl') - out = comm.all_reduce(torch.ones(1, device=device)) - torch.npu.synchronize() - value = int(out.item()) - ok = value == args.world_size - log(f'[{label}] all_reduce 结果={value}（期望={args.world_size}）{"OK" if ok else "异常"}') - return comm, ok - - -def _run_two_groups(args, torch, device) -> int: - """在同一张卡上先后建两个并发 HCCL 组，复现 delta 的「第二个并发组」场景。 - - delta 真实情况：trainer 卡上已有 FSDP 组、Pod2 卡上已有 vLLM TP 组，delta 组是第二个； - 且 trainer 侧 delta 组与 FSDP 组共用同一个 HCCL_IF_BASE_PORT。 - - 组 1 建好后不销毁（占住设备），再建组 2。 - - 默认：组 2 复用相同 HCCL 端口（= delta 与 FSDP 共享 base port 的情形）。 - - 指定 --second-base-port：组 2 换一套端口（验证「换端口即可解决」的修复方向）。 - """ - log('==== 建立组 1（占住设备，模拟已有的 FSDP / TP 通信域）====') - try: - comm1, ok1 = _build_pyhccl_group(args, torch, device, args.master_port, 'group1') - except Exception as error: # noqa: BLE001 - import traceback - log(f'FAIL：组 1 建链失败：{type(error).__name__}: {error}') - traceback.print_exc() - return 1 - if not ok1: - log('FAIL：组 1 本身就不正常，无法继续') - return 1 - log('组 1 保持存活（不销毁），现在在同一张卡上建组 2') - - if args.second_base_port: - base = args.second_base_port - os.environ['HCCL_IF_BASE_PORT'] = str(base) - os.environ['HCCL_HOST_SOCKET_PORT_RANGE'] = f'{base + 1000}-{base + 1511}' - os.environ['HCCL_NPU_SOCKET_PORT_RANGE'] = f'{base + 2000}-{base + 2511}' - log(f'==== 建立组 2（换一套 HCCL 端口 base={base}）====') - else: - log('==== 建立组 2（复用与组 1 相同的 HCCL 端口 = delta 与 FSDP 共享 base port 的情形）====') - - try: - comm2, ok2 = _build_pyhccl_group(args, torch, device, args.master_port + 1, 'group2') - except Exception as error: # noqa: BLE001 - import traceback - log(f'FAIL：组 2 建链失败（这就是 delta 的超时点）：{type(error).__name__}: {error}') - traceback.print_exc() - return 1 - if not ok2: - log('FAIL：组 2 并发通信异常') - return 1 - log('PASS：同一张卡上第二个并发 HCCL 组也正常') - return 0 - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('--backend', choices=('gloo', 'hccl'), required=True) - parser.add_argument('--rank', type=int, required=True, help='本进程的全局 rank（rank 0 即 master）') - parser.add_argument('--world-size', type=int, default=2) - parser.add_argument('--master-addr', required=True, help='rank 0 所在 Pod 的 IP') - parser.add_argument('--master-port', type=int, default=29555) - parser.add_argument('--device', type=int, default=0, help='本机使用的本地 NPU 序号（hccl 后端用）') - parser.add_argument('--timeout', type=int, default=120, help='rendezvous/集合通信超时秒数') - parser.add_argument('--iface', help='指定网卡名，会设置 GLOO_SOCKET_IFNAME / HCCL_SOCKET_IFNAME') - parser.add_argument( - '--stateless', action='store_true', - help='hccl 后端用 delta 完全相同的 vllm_ascend pyhccl 路径（stateless_init_process_group）建链，' - '而非 torch.distributed。用于判断 delta 超时是否来自这条 pyhccl 初始化路径本身。') - parser.add_argument( - '--startup-delay', type=int, default=0, - help='建链前本进程先睡眠的秒数。用于复现 delta 里 trainer（rank 0）很晚才 join 的时间错位，' - '例如给 rank 0 设一个大的值，看非 0 rank 是否因等待超时。') - parser.add_argument( - '--two-groups', action='store_true', - help='在同一张卡上先后建两个并发 HCCL 组（pyhccl 路径），复现 delta 的第二个并发组场景。') - parser.add_argument( - '--second-base-port', type=int, default=0, - help='配合 --two-groups：给第二个组换一套 HCCL_IF_BASE_PORT（及派生端口段）。' - '不指定则第二个组复用与第一个组相同的端口（= delta 与 FSDP 共享 base port）。') - args = parser.parse_args() - - if args.iface: - os.environ['GLOO_SOCKET_IFNAME'] = args.iface - os.environ['HCCL_SOCKET_IFNAME'] = args.iface - os.environ.setdefault('TP_SOCKET_IFNAME', args.iface) - - os.environ['MASTER_ADDR'] = args.master_addr - os.environ['MASTER_PORT'] = str(args.master_port) - os.environ['RANK'] = str(args.rank) - os.environ['WORLD_SIZE'] = str(args.world_size) - - import torch - import torch.distributed as dist - - if args.backend == 'hccl': - # 必须在创建进程组前导入并选卡，否则 HCCL 不知道用哪张 NPU。 - import torch_npu # noqa: F401 - if not torch.npu.is_available(): - log('FAIL：torch.npu.is_available() 为 False，本机看不到 NPU') - return 1 - torch.npu.set_device(args.device) - device = torch.device('npu', args.device) - log(f'使用 NPU device {args.device}，device_count={torch.npu.device_count()}') - else: - device = torch.device('cpu') - - if args.startup_delay > 0: - import time - log(f'建链前先睡眠 {args.startup_delay}s（模拟 join 时间错位）') - time.sleep(args.startup_delay) - - if args.stateless: - if args.backend != 'hccl': - log('FAIL：--stateless 仅支持 hccl 后端') - return 1 - return _run_stateless(args, torch, device) - - if args.two_groups: - if args.backend != 'hccl': - log('FAIL：--two-groups 仅支持 hccl 后端') - return 1 - return _run_two_groups(args, torch, device) - - log(f'开始 init_process_group：backend={args.backend}, rank={args.rank}/{args.world_size}, ' - f'master={args.master_addr}:{args.master_port}, timeout={args.timeout}s') - log('若进程在此长时间停住，说明两机无法在该 backend 上 rendezvous（gloo=TCP 不通；hccl=NPU 跨机建链失败）') - try: - dist.init_process_group( - backend=args.backend, init_method='env://', rank=args.rank, world_size=args.world_size, - timeout=datetime.timedelta(seconds=args.timeout)) - except Exception as error: # noqa: BLE001 - log(f'FAIL：init_process_group 抛异常：{type(error).__name__}: {error}') - return 1 - log('init_process_group 成功，进程组已建立') - - try: - tensor = torch.ones(1, device=device) - dist.all_reduce(tensor, op=dist.ReduceOp.SUM) - if args.backend == 'hccl': - torch.npu.synchronize() - value = int(tensor.item()) - log(f'all_reduce 完成，结果={value}（期望={args.world_size}）') - if value != args.world_size: - log('FAIL：all_reduce 结果不等于 world_size，通信异常') - return 1 - log('PASS：跨机通信正常') - except Exception as error: # noqa: BLE001 - log(f'FAIL：all_reduce 抛异常：{type(error).__name__}: {error}') - return 1 - finally: - dist.barrier() - dist.destroy_process_group() - return 0 - - -if __name__ == '__main__': - raise SystemExit(main()) diff --git a/cookbook/checkpoint_engine/delta_benchmark.py b/cookbook/checkpoint_engine/delta_benchmark.py deleted file mode 100644 index 170c1f7b1..000000000 --- a/cookbook/checkpoint_engine/delta_benchmark.py +++ /dev/null @@ -1,83 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""权重同步耗时基准，口径对齐 verl 的 tests/checkpoint_engine/bench_delta_vllm.py。 - -和 delta_validation.py 不同，这里不做逐字节比较，只计时 manager.sync_weights()： -先做一次 seed 同步，再对每个变化比例各同步 --repeats 次。一个进程只测一个后端。 - -python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend delta_sharded -python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend full -""" -import argparse -from statistics import median -from time import perf_counter - -from cookbook.checkpoint_engine.delta_validation import build_components, log_progress -from twinkle.checkpoint_engine import CheckpointEngineManager, DeltaSyncConfig - - -def timed_sync(manager) -> float: - started = perf_counter() - manager.sync_weights() - return perf_counter() - started - - -def run_benchmark(args): - trainer, sampler = build_components(args) - if args.backend == 'delta_sharded': - manager = CheckpointEngineManager( - trainer, sampler, platform='GPU', mode='standalone', sync_backend='delta_sharded', - delta_config=DeltaSyncConfig(bucket_bytes=args.bucket_bytes, gather_batch_size=args.gather_batch_size)) - else: - manager = CheckpointEngineManager(trainer, sampler, platform='GPU', mode='standalone') - - def metrics(): - return manager.last_delta_metrics if args.backend == 'delta_sharded' else {} - - try: - seconds = timed_sync(manager) - log_progress(f'[seed] backend={args.backend} {seconds:.3f}s metrics={metrics()}') - - rows = [] - step = 0 - for ratio in args.ratios: - samples = [] - for _ in range(args.repeats): - step += 1 - trainer.change_fraction(ratio) - seconds = timed_sync(manager) - samples.append(seconds) - log_progress(f'[ratio={ratio}] step={step} sync={seconds:.3f}s metrics={metrics()}') - rows.append((ratio, samples, metrics())) - finally: - if args.backend == 'delta_sharded': - manager.close_sync_session() - - print(f'\nbackend={args.backend} model={args.model} trainer={args.trainer_gpus} rollout_tp={args.tp} dp={args.dp} ' - f'bucket={args.bucket_bytes >> 20}MB') - print(f"{'ratio':>8} {'median(s)':>10} {'min(s)':>8} {'max(s)':>8} last metrics") - for ratio, samples, last in rows: - print(f'{ratio:>8} {median(samples):>10.3f} {min(samples):>8.3f} {max(samples):>8.3f} {last}') - - -def main(): - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('--model', required=True, help='训练与 rollout 使用的同一模型，例如 Qwen3-8B') - parser.add_argument('--backend', choices=('delta_sharded', 'full'), default='delta_sharded') - parser.add_argument('--trainer-gpus', type=int, choices=(1, 2), default=2, help='FSDP2 训练卡数') - parser.add_argument('--tp', type=int, choices=(1, 2), default=2, help='rollout 副本的 GPU 数') - parser.add_argument('--dp', type=int, default=1, help='独立 rollout 副本数') - parser.add_argument('--strategy', choices=('native_fsdp', 'accelerate'), default='native_fsdp') - parser.add_argument('--cpu-offload', action='store_true', help='使用 Accelerate FSDP2 CPU offload') - parser.add_argument('--ratios', type=float, nargs='+', default=[0.005, 0.01, 0.02, 0.03]) - parser.add_argument('--repeats', type=int, default=2, help='每个比例的同步次数') - parser.add_argument('--bucket-bytes', type=int, default=2048 << 20, help='verl 默认 2048MB') - parser.add_argument('--gather-batch-size', type=int, default=32) - parser.add_argument('--devices-per-node', type=int, help='每个节点的 GPU 数，默认单节点运行') - args = parser.parse_args() - if args.cpu_offload and args.strategy != 'accelerate': - parser.error('--cpu-offload 需要 --strategy accelerate') - run_benchmark(args) - - -if __name__ == '__main__': - main() diff --git a/cookbook/checkpoint_engine/delta_benchmark_notes.md b/cookbook/checkpoint_engine/delta_benchmark_notes.md deleted file mode 100644 index 035bea797..000000000 --- a/cookbook/checkpoint_engine/delta_benchmark_notes.md +++ /dev/null @@ -1,144 +0,0 @@ -# Delta 权重同步实验记录 - -## 背景 - -- 起因：`delta_validation.py` 在 `--tp 2` 时崩溃，`--dp 2` 正常。 -- 修复后验证通过，但 delta 同步比 full 同步慢，甚至是负优化。 -- 目标：先在 verl 上跑出 baseline，再把 twinkle 的脚本对齐 verl，比较两边。 - -## 环境 - -- HAMi vGPU（libvgpu 拦截 `cuMemCreate`），所以绝对耗时不一定能外推到裸机。 -- trainer 2 卡（FSDP2），rollout 2 卡（vLLM TP=2）。 -- 模型：Qwen3-8B（`/nas/disk1/Qwen3-8B`）。 - -## 1. tp=2 崩溃 - -- 现象：NCCL 通信组初始化时报 `Illegal device id` 并 segfault。 -- 原因：HAMi 环境下 NCCL 的 cuMem 路径（`NCCL_CUMEM_ENABLE`、`NCCL_CUMEM_HOST_ENABLE`）不可用，tp>1 时 vLLM 才会初始化 NCCL，所以 dp=2 不受影响。 -- 处理：在 `delta_validation.py` 开头用 `os.environ.setdefault` 把两个变量设为 `0`，与 verl 的做法一致，外部可覆盖。 -- 结果：重跑后 status=passed。 - -## 2. twinkle 早期观察（Qwen3.5-0.8B，bucket 64MB） - -- delta 耗时随 flush 数增长。每个 flush 都要新建 sender/receiver、两次 NCCL broadcast、checksum 同步、Ray RPC 和 vLLM `collective_rpc`。 -- 把 bucket 调大后：seed 10.1s 降到 1.9s；50% 和 100% 变化的用例从 11.1s、16.7s 降到约 2.7s。 -- 但 1% 到 10% 变化仍然是 3.2 到 3.5s，full 只要 1.0 到 1.6s。0% 约 1.0s。 -- 结论：存在 2 到 3 秒的固定开销，与 flush 数无关。来源未定位。 - -## 3. 数据量原理 - -- 稀疏条目每个变化元素 6 字节（int32 位置 + bf16 值），full 同步每个元素 2 字节。 -- 变化比例超过约 1/3 时，delta 的传输量就比 full 大。 -- 真实 RL 单步变化约 1% 到 3%（verl 文档在 Qwen3-8B 上测得 0.884% 到 0.895%），所以比例改成 0.5%、1%、2%、3%。 - -## 4. 脚本 - -| 脚本 | 说明 | -|---|---| -| `verl/tests/checkpoint_engine/bench_delta_vllm.py` | verl baseline。FSDP2 trainer（`forward_only=True`，`model_dtype=bfloat16`），随机翻转 bf16 最低位模拟变化，计时整个 `update_weights()` | -| `cookbook/checkpoint_engine/delta_benchmark.py` | twinkle 对齐脚本。先 seed，再按比例同步，计时整个 `sync_weights()`，输出 median/min/max 表 | - -运行： - -``` -# verl -python tests/checkpoint_engine/bench_delta_vllm.py --model /nas/disk1/Qwen3-8B --backend delta_sharded -python tests/checkpoint_engine/bench_delta_vllm.py --model /nas/disk1/Qwen3-8B --backend nccl - -# twinkle -python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend delta_sharded -python -m cookbook.checkpoint_engine.delta_benchmark --model /nas/disk1/Qwen3-8B --backend full -``` - -verl 脚本调试中修过的问题： - -- 缺少 flash-attn：改用 `attn_implementation=sdpa`，并设置 `use_remove_padding=False`。 -- `perturb` 断言 bf16 失败：verl 的 `model_dtype` 默认是 `fp32`，需要显式设为 `bfloat16`。 -- `perturb` 是 blocking 调用，不能再套 `ray.get`。 - -## 5. verl baseline（Qwen3-8B，bucket 2048MB，每档 2 次） - -seed（delta_sharded）：256.9s，传输 15622MB，7 个 flush。 - -| 比例 | delta 中位数 (min~max) | nccl 中位数 (min~max) | delta 传输量 | -|---|---|---|---| -| 0.5% | 38.3s (33.9~42.6) | 21.9s (14.8~29.1) | 234MB | -| 1% | 27.6s (23.7~31.5) | 28.8s (27.4~30.2) | 469MB | -| 2% | 17.8s (15.7~19.8) | 7.3s (6.2~8.4) | 937MB | -| 3% | 34.8s (27.9~41.8) | 6.5s (5.7~7.3) | 1406MB | - -日志字段说明（以 0.5% 为例）： - -- `flipped` 和 `changed_elems` 相等，说明 delta 检测正确。 -- `payload_mbytes` = 变化元素数 × 6 字节，与预期吻合。 -- 变化量都小于 2048MB，所以每次只有 1 个 flush。 - -## 6. 结论 - -能确定： - -- verl 自己的 delta 在这个环境里也不比 nccl 快，负优化不是 twinkle 特有的。 -- nccl 每次传输量固定约 15.6GB，耗时却在 5.7s 到 30s 之间，环境本身波动约 5 倍。 -- delta 的耗时与传输量无关：传 937MB 最快，传 234MB 反而最慢。传 234MB 要 34 到 42s，这部分不是带宽造成的。 -- 2% 和 3% 档，nccl 比 delta 快 2.4 到 5 倍。 - -不能确定： - -- 每档只有 2 个样本，没有统计意义。 -- 固定开销花在哪一步未知。候选：HAMi 对 CUDA IPC 的影响、ZMQ 和 Ray RPC、暂停和恢复生成、释放和恢复 KV cache。均未测量。 -- 机器上是否有其他任务占用 GPU 或 NAS，未排查。 - -## 7. 已知的口径差异（twinkle 与 verl） - -- 变化位置：twinkle 的 `change_fraction` 翻转每个参数前 N% 元素的符号位，verl 随机翻转最低位。变化数量相同，位置分布不同。 -- CPU offload：verl 的 `forward_only=True` 会启用，twinkle 默认不启用。对齐需加 `--strategy accelerate --cpu-offload`。 -- 分片 dtype：verl 现在是 bf16；twinkle 取决于 transformers 版本的默认加载 dtype，未核实。传输字节数不受影响。 -- `delta_benchmark.py` 不做 `forbid_full_tensor` 检查。 - -## 8. twinkle 结果（Qwen3-8B，bucket 2048MB，tp=2，每档 2 次） - -| 比例 | delta 中位数 (min~max) | trainer `send_seconds` | full 中位数 (min~max) | delta 传输量 | -|---|---|---|---|---| -| 0.5% | 8.4s (7.2~9.6) | 0.44s | 1.3s (0.96~1.6) | 234MB | -| 1% | 9.8s (9.1~10.6) | 0.44s | 0.9s (0.83~1.0) | 469MB | -| 2% | 11.6s (11.2~12.1) | 0.44s | 1.2s (0.86~1.5) | 937MB | -| 3% | 10.0s (9.0~11.1) | 0.71s | 1.3s (0.80~1.7) | 1406MB | - -- 传输量和 verl 完全一致（0.5% 档 234.34MB、1 个 flush），说明编码和分桶等价。 -- trainer 侧（diff、gather、broadcast）只要 0.44~0.71s，不是瓶颈。 -- delta 比 full 多出的 7~11s 都在 trainer `send_seconds` 之外，且随传输量有增长趋势（0.5%→2% 约 4.7s/GB）；3% 比 2% 低，属于 2 样本噪声。 -- 嫌疑（均未测量）：`begin_delta_sync` 里 `sleep(level=1)` + `wake_up(['weights'])`（full 路径没有）、每个 flush 新建 ZMQ/IPC、接收端 `clone` + checksum `.item()` + 稀疏 patch、`_clear_delta_cache` 与 `wake_up(['kv_cache'])`。 - -## 9. verl 与 twinkle 对比 - -| 比例 | verl delta | verl nccl | twinkle delta | twinkle full | -|---|---|---|---|---| -| 0.5% | 38.3s | 21.9s | 8.4s | 1.3s | -| 1% | 27.6s | 28.8s | 9.8s | 0.9s | -| 2% | 17.8s | 7.3s | 11.6s | 1.2s | -| 3% | 34.8s | 6.5s | 10.0s | 1.3s | - -- 两边 delta 都比各自的全量同步慢，负优化不是 twinkle 特有的。 -- twinkle 的 delta 绝对耗时反而更低（8~12s 对 18~38s）；twinkle 的 delta/full 比例更差，主要因为 twinkle full 基线只有约 1s。 -- 计时范围不一致，不能直接下结论：verl `update_weights()` 包含终止请求、释放/恢复 KV cache、重建进程组、恢复生成（来自代码阅读，未逐行核对）；twinkle delta 含暂停生成、sleep/wake_up、清缓存、恢复生成；twinkle full 只有 `sync_weights()`，没有 sleep/wake_up，所以最轻。 -- 其他口径差异见第 7 节。 - -## 10. verl 官方文档的测法（`docs/advance/delta_weight_sync.md`，"Measured results"） - -- 没有独立 benchmark 脚本，数据来自真实 GSM8K GRPO 训练中每步稳态权重同步的计时，`delta_sharded` 与 `nccl` 各跑一遍。 -- 环境：H100 或 A100/A800 多节点，训练与 rollout 分离（多为 2+2 节点），SGLang rollout。部分行为 50 步中位数或 50 步持续运行。 -- 变化比例由训练自然产生（稠密模型约 0.884%~0.895%），不是人为设定。 -- 正确性另测：200 步 GRPO 奖励曲线对比、perturb/apply/revert/apply 位级一致。 -- 官方 Qwen3-8B 数据：2+2 节点 FSDP2 delta 2.24~3.40s，nccl 34~38s；1+1 节点 delta 3.00s，nccl 11.00s。Qwen3-0.6B 单机 NVLink 为 0.59s 对 0.61s（持平）。 -- 文档自己说明：快链路加小模型时，全量广播本来就便宜，delta 的优势随模型规模和网络距离增长。 -- 与我们测试的差异：官方是多节点、走网络、真实训练循环；我们是单机 4 张 HAMi 虚拟卡、合成扰动、微基准。因此单机上 delta 不占优不与文档矛盾。 -- 官方 verl delta 在 8B 上约 2.2~3.4s，我们的 verl delta 是 18~38s，twinkle 是 8~12s；这一差距的原因未查明（HAMi、单机部署、计时范围都有可能）。 - -## 11. 待办 - -- 决定：改为参照官方做法，用真实 GRPO 训练做 delta 与 full 的对比（用户已确认方向，具体方案待定）。 -- 跑之前检查 `nvidia-smi`，确认没有其他任务占卡。 -- 在 delta 路径加分阶段计时（trainer 的 diff/gather、broadcast、接收端 IPC 与 vLLM patch、前后的暂停恢复），定位 7~11s 的来源。涉及 `checkpoint_engine/delta/engine.py`、`delta/trainer.py` 和 `delta/rollout.py`，需先确认。 -- 可选：给 `delta_benchmark.py` 加 `--no-free-cache`（`free_cache_engine=False`），看 sleep/wake_up 占多少。 -- 可选：把 NCCL 环境变量移到库里（`vLLMSampler.__init__`），让所有使用者受益。 diff --git a/cookbook/checkpoint_engine/delta_validation.py b/cookbook/checkpoint_engine/delta_validation.py deleted file mode 100644 index ae6d42c83..000000000 --- a/cookbook/checkpoint_engine/delta_validation.py +++ /dev/null @@ -1,393 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""在 GPU(CUDA) 或 Ascend NPU 上验证 FSDP2 增量同步：逐字节比较 delta 与 full 同步后的 vLLM 权重。 - -示例（2 张训练卡 + 4 张 rollout 卡，GPU）： -python -m cookbook.checkpoint_engine.delta_validation --dp 2 --tp 2 --rounds 10 -示例（只有 2 张卡：1 张训练、1 张 rollout，FSDP 不分片）： -python -m cookbook.checkpoint_engine.delta_validation --trainer-gpus 1 --dp 1 --tp 1 --rounds 2 - -NPU（单机 4 卡、不联网，模型必须是本地目录；NPU 默认 --offline 且 --num-devices 4）。先用 --dry-run 检查计划： -(a) 1 训练 + 1 rollout(tp=1)： -python -m cookbook.checkpoint_engine.delta_validation --platform NPU --model-path /path/to/model \\ - --trainer-gpus 1 --dp 1 --tp 1 --rounds 2 -(b) 1 训练 + 1 rollout(tp=2)： --trainer-gpus 1 --dp 1 --tp 2 -(c) 2 训练 + 2 个 rollout 副本(tp=1)： --trainer-gpus 2 --dp 2 --tp 1 -(d) 2 训练 + 1 rollout(tp=2)： --trainer-gpus 2 --dp 1 --tp 2 -“多副本且 tp>1”至少需要 5 张卡，4 卡机器无法验证。NPU 路径尚未在真实 NPU 上跑过。 -""" -import argparse -import importlib.metadata -import json -import os -import subprocess -from contextlib import contextmanager -from datetime import datetime -from pathlib import Path -from time import perf_counter - -import torch -from torch.distributed.tensor import DTensor - -import twinkle -from cookbook.checkpoint_engine.device_utils import (add_platform_arguments, build_device_groups, build_engine_args, - delta_config_kwargs, device_peak_bytes, device_synchronize, - nproc_per_node, run_preflight_or_exit) -from twinkle import DeviceMesh, Platform, remote_class, remote_function -from twinkle.checkpoint_engine import CheckpointEngineManager, DeltaSyncConfig -from twinkle.checkpoint_engine.delta.spec import ShardSpec, derive_dtensor_placement, translate_flat_indices -from twinkle.data_format.sampling import SamplingParams -from twinkle.model import TransformersModel -from twinkle.sampler import vLLMSampler -from twinkle.sampler.vllm_sampler.vllm_worker_extension import TwinkleWorkerExtension - -# tp>1 时 vLLM 会初始化 NCCL 通信组；在 HAMi 等 vGPU 环境下，NCCL 的 cuMem* 路径会触发 -# "Illegal device id" 并 segfault。与 verl 一致，在 Ray actor 继承环境前关闭（可由外部覆盖）。 -os.environ.setdefault('NCCL_CUMEM_ENABLE', '0') -os.environ.setdefault('NCCL_CUMEM_HOST_ENABLE', '0') - -# 0% 用来测固定开销；其余贴近真实 RL 单步变化（稠密模型约 1%~3%）。 -CHANGE_RATIOS = (0., .005, .01, .02, .03) - - -def log_progress(message: str) -> None: - print(f'[{datetime.now():%Y-%m-%d %H:%M:%S}] {message}', flush=True) - - -@contextmanager -def validation_stage(name: str): - log_progress(f'开始：{name}') - started = perf_counter() - try: - yield - except Exception as error: - log_progress(f'失败：{name}，耗时 {perf_counter() - started:.3f}s，{error!r}') - raise - else: - log_progress(f'完成：{name}，耗时 {perf_counter() - started:.3f}s') - - -def environment_report(): - """记录软件版本和驱动信息，方便复现设备上的结果。""" - report = {} - for package in ('torch', 'torch_npu', 'vllm', 'vllm_ascend', 'transformers', 'accelerate', 'ray'): - try: - report[package] = importlib.metadata.version(package.replace('_', '-')) - except importlib.metadata.PackageNotFoundError: - report[package] = None - # GPU 机器有 nvidia-smi，Ascend 机器有 npu-smi；哪个存在就记录哪个。 - for tool, command in (('nvidia-smi', ['nvidia-smi']), ('npu-smi', ['npu-smi', 'info'])): - try: - report[tool] = subprocess.run(command, capture_output=True, text=True, timeout=15).stdout - except (FileNotFoundError, subprocess.TimeoutExpired): - report[tool] = None - return report - - -@remote_class() -class ValidationTrainer(TransformersModel): - def __init__(self, device_mesh: DeviceMesh = None, **kwargs): - TransformersModel.__init__.__wrapped__(self, device_mesh=device_mesh, **kwargs) - - @remote_function(dispatch='all', lazy_collect=False) - def change_fraction(self, ratio: float): - """翻转每个参数前 ratio 比例元素的符号位，产生可控的 BF16 字节变化。""" - self._lazy_wrap_model() - with torch.no_grad(): - for tensor in self.strategy.unwrap_model(self.model).parameters(): - place, _, _ = derive_dtensor_placement(ShardSpec.from_param(tensor)) - local = tensor.to_local() if isinstance(tensor, DTensor) else tensor - values = local.detach().to(torch.bfloat16).contiguous().reshape(-1) - indices = torch.arange(values.numel(), dtype=torch.int64, device=values.device) - # 使用完整参数的坐标，使各 FSDP shard 合起来恰好修改指定比例。 - changed = translate_flat_indices(indices, place) < int(tensor.numel() * ratio) - bits = values.view(torch.int16) - bits[changed] = bits[changed].bitwise_xor(-32768) - local.copy_(values.reshape(local.shape).to(device=local.device, dtype=local.dtype)) - - @remote_function(dispatch='all', lazy_collect=False) - def forbid_full_tensor(self): - # 在 trainer 进程里临时禁用全量聚合，检查稳态 delta 没有退回 full 路径。 - self._validation_full_tensor = DTensor.full_tensor - - def forbidden(*args, **kwargs): - raise AssertionError('稳态 delta 调用了 DTensor.full_tensor()') - - DTensor.full_tensor = forbidden - - @remote_function(dispatch='all', lazy_collect=False) - def allow_full_tensor(self): - DTensor.full_tensor = self._validation_full_tensor - - @remote_function(dispatch='all', lazy_collect=False) - def training_step(self): - """执行一次真实的前向、反向和 SGD 更新，补充人为修改之外的验证。""" - optimizer = torch.optim.SGD(self.model.parameters(), lr=0.001) - tokens = torch.tensor([[1, 2, 3, 4]], device=self._delta_device) - self.model.train() - optimizer.zero_grad() - loss = self.model(input_ids=tokens, labels=tokens).loss - loss.backward() - optimizer.step() - optimizer.zero_grad() - return float(loss.detach().cpu()) - - @remote_function(dispatch='all', lazy_collect=False) - def sync_device(self): - """等待本进程设备上已排队的 kernel 全部完成。 - - 端到端耗时基准用它隔开“训练计算”和“权重同步”：clip_grad_and_step 返回时优化器 kernel 可能还没跑完， - 不同步的话这部分时间会被算进随后的 sync。 - """ - device_synchronize(Platform.get_platform().device_prefix()) - - @remote_function(dispatch='all', lazy_collect=False) - def validation_environment(self): - return environment_report() - - -class ValidationWorker(TwinkleWorkerExtension): - """在每个 vLLM TP worker 内保存和比较真正参与推理的 runtime 权重。""" - - def validation_environment(self): - return environment_report() - - def capture_weights(self): - self._validation_weights = { - name: tensor.detach().cpu().clone() - for name, tensor in self.model_runner.model.state_dict().items()} - # self.device 是 torch.device('cuda:N') 或 torch.device('npu:N')，不要写死 torch.cuda。 - return {'device_peak_bytes': device_peak_bytes(torch.device(self.device).type)} - - def begin_full_reload(self): - # vLLM 可能打包或转置权重；全量基准也必须走完整 reload 生命周期。 - from vllm.model_executor.model_loader.reload.layerwise import ( - initialize_layerwise_reload, record_metadata_for_reloading) - with torch.device(self.device): - record_metadata_for_reloading(self.model_runner.model) - initialize_layerwise_reload(self.model_runner.model) - - def compare_weights(self): - from vllm.model_executor.model_loader.reload.layerwise import finalize_layerwise_reload - finalize_layerwise_reload(self.model_runner.model, self.model_runner.model_config) - state = self.model_runner.model.state_dict() - if state.keys() != self._validation_weights.keys(): - raise AssertionError('全量 reload 改变了 runtime 权重键') - mismatches = [] - for name, tensor in state.items(): - # 比较字节而非浮点近似值，包含正负零等位级差异。 - actual = tensor.detach().cpu().contiguous() - expected = self._validation_weights[name].contiguous() - if actual.shape != expected.shape or not torch.equal( - actual.reshape(-1).view(torch.uint8), expected.reshape(-1).view(torch.uint8)): - mismatches.append(self._describe_mismatch(name, actual, expected)) - self._validation_weights = None - if mismatches: - # 汇总全部不一致的 tensor，而不是在第一个处中止，便于判断是整体错误还是局部错误。 - shown = '\n '.join(mismatches[:20]) - more = f'\n ...另有 {len(mismatches) - 20} 个' if len(mismatches) > 20 else '' - raise AssertionError( - f'delta 与 full 权重不一致：rank={self.rank}, {len(mismatches)}/{len(state)} 个 tensor 不一致\n' - f' {shown}{more}') - return True - - @staticmethod - def _describe_mismatch(name, actual, expected): - """描述一个不一致 tensor：不一致元素占比、行范围，以及首个差异处的 full/delta 取值。 - - actual 是 full reload 后的运行时权重，expected 是 delta 同步后保存的快照。 - """ - if actual.shape != expected.shape: - return f'{name}: shape 不同 full={tuple(actual.shape)} delta={tuple(expected.shape)}' - int_dtype = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}[actual.element_size()] - differs = actual.reshape(-1).view(int_dtype) != expected.reshape(-1).view(int_dtype) - total = differs.numel() - count = int(differs.sum()) - first = int(differs.nonzero()[0]) - last = int(differs.nonzero()[-1]) - text = (f'{name}: shape={tuple(actual.shape)} 不一致元素 {count}/{total} ({count / total:.4%}), ' - f'flat 范围 [{first}, {last}]') - if actual.dim() >= 2: - rows = differs.reshape(actual.shape[0], -1).any(dim=1).nonzero().flatten() - text += (f', 不一致行 {rows.numel()}/{actual.shape[0]} ' - f'(行范围 [{int(rows[0])}, {int(rows[-1])}])') - full_value = actual.reshape(-1)[first].item() - delta_value = expected.reshape(-1)[first].item() - return f'{text}, 首个差异 flat={first}: full={full_value} delta={delta_value}' - - -@remote_class() -class ValidationSampler(vLLMSampler): - def __init__(self, device_mesh: DeviceMesh = None, **kwargs): - vLLMSampler.__init__.__wrapped__(self, device_mesh=device_mesh, **kwargs) - - @remote_function(dispatch='all', lazy_collect=False, collect='flatten') - def worker_rpc(self, method: str): - # 遍历全部 DP 副本及其 TP workers，任何 worker 的比较失败都会向上传递。 - return self._run_in_loop(self.engine.engine.collective_rpc(method)) - - @remote_function(dispatch='all', lazy_collect=False) - def pause_validation(self): - self._run_in_loop(self.engine.pause_delta_generation(self._delta_config.timeout_seconds)) - - @remote_function(dispatch='all', lazy_collect=False) - def generate_one_token(self): - return self._run_in_loop(self.engine.sample( - prompt=[1, 2, 3], sampling_params=SamplingParams(max_tokens=1, temperature=0.))) - - -WORKER_EXTENSION_CLS = 'cookbook.checkpoint_engine.delta_validation.ValidationWorker' - - -def build_components(args): - """训练使用 trainer_gpus 张卡的 FSDP2；rollout 使用 dp 个独立副本，每个副本占 tp 张卡。 - - GPU 与 NPU 共用同一套流程，差别（device_type、vLLM 的 weight_transfer 后端、NPU 的 rl_config 等） - 都在 device_utils 里按 --platform 解析。 - """ - trainer_gpus = args.trainer_gpus - rollout_gpus = args.dp * args.tp - with validation_stage(f'初始化 Ray 和 {args.platform} 资源组'): - twinkle.initialize(mode='ray', nproc_per_node=nproc_per_node(args), - groups=build_device_groups(args), lazy_collect=False) - with validation_stage(f'创建 {trainer_gpus} 卡 FSDP trainer'): - trainer = ValidationTrainer( - model_id=args.model, remote_group='trainer', strategy=args.strategy, mixed_precision='bf16', - fsdp_config={'cpu_offload': True} if args.cpu_offload else {}, - device_mesh=DeviceMesh.from_sizes(world_size=trainer_gpus, fsdp_size=trainer_gpus)) - with validation_stage(f'创建 rollout：{args.dp} 个副本，每个 {args.tp} 张 GPU'): - sampler = ValidationSampler( - model_id=args.model, remote_group='rollout', - device_mesh=DeviceMesh.from_sizes(world_size=rollout_gpus, tp_size=args.tp, dp_size=args.dp), - engine_args=build_engine_args(args, WORKER_EXTENSION_CLS)) - return trainer, sampler - - -def steady_sync(trainer, delta): - """稳态同步必须只传增量；无论是否成功，都恢复 full_tensor 方法。""" - with validation_stage('禁用 trainer 的全量权重聚合'): - trainer.forbid_full_tensor() - try: - started = perf_counter() - with validation_stage('稳态 delta 权重同步'): - delta.sync_weights() - if delta.last_delta_metrics['seed']: - raise AssertionError('稳态同步意外重新执行了全量 seed') - return perf_counter() - started - finally: - with validation_stage('恢复 trainer 的全量权重聚合'): - trainer.allow_full_tensor() - - -def compare_with_full(sampler, full): - """保存 delta 结果，再全量加载相同 trainer 权重，比较两条路径的最终结果。""" - with validation_stage('rollout 生成一个 token，检查同步后推理'): - sampler.generate_one_token() - with validation_stage('暂停 rollout 推理'): - sampler.pause_validation() - with validation_stage('保存所有 rollout worker 的 delta 权重快照'): - workers = sampler.worker_rpc('capture_weights') - with validation_stage('准备全量权重 reload'): - sampler.worker_rpc('begin_full_reload') - started = perf_counter() - with validation_stage('full 权重同步'): - full.sync_weights() - seconds = perf_counter() - started - with validation_stage('所有 rollout worker 逐字节比较 delta 与 full 权重'): - sampler.worker_rpc('compare_weights') - return {'full_seconds': seconds, 'workers': workers} - - -def run_validation(args): - log_progress( - f'启动验证：platform={args.platform}，model={args.model}，strategy={args.strategy}，' - f'trainer={args.trainer_gpus} 张卡，rollout dp={args.dp}、tp={args.tp}，' - f'rounds={args.rounds}，report={Path(args.report).resolve()}') - trainer, sampler = build_components(args) - with validation_stage('创建 delta 和 full checkpoint manager'): - delta = CheckpointEngineManager( - trainer, sampler, platform=args.platform, mode='standalone', sync_backend='delta_sharded', - delta_config=DeltaSyncConfig(**delta_config_kwargs(args))) - full = CheckpointEngineManager(trainer, sampler, platform=args.platform, mode='standalone') - report = {'status': 'running', 'arguments': vars(args), 'rounds': []} - - def save_report(): - Path(args.report).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding='utf-8') - - def check_round(stage): - with validation_stage(stage): - seconds = steady_sync(trainer, delta) - result = compare_with_full(sampler, full) - result.update(stage=stage, delta_seconds=seconds, metrics=delta.last_delta_metrics) - report['rounds'].append(result) - save_report() - log_progress(f'{stage}：字节比较通过，delta={seconds:.3f}s，full={result["full_seconds"]:.3f}s') - - try: - with validation_stage('收集 driver 和 trainer 环境信息'): - report.update(driver=environment_report(), trainers=trainer.validation_environment()) - # 1. 首次发送 dense seed 并建立 CPU shard 快照，后续各轮复用同一个通信组。 - with validation_stage('首次全量 seed 同步，建立增量快照和通信组'): - delta.sync_weights() - with validation_stage('收集 rollout worker 环境信息'): - report['workers'] = sampler.worker_rpc('validation_environment') - with validation_stage('首次 seed 权重验证'): - report['seed'] = compare_with_full(sampler, full) - - # 2. 轮流覆盖零变化、稀疏变化和全部变化；full 基准不修改 trainer 快照。 - for step in range(args.rounds): - ratio = CHANGE_RATIOS[step % len(CHANGE_RATIOS)] - with validation_stage(f'第 {step + 1}/{args.rounds} 轮：修改 {ratio:.1%} 权重'): - trainer.change_fraction(ratio) - check_round(f'第 {step + 1}/{args.rounds} 轮验证，修改比例 {ratio:.1%}') - - # 3. 真实训练更新后的 delta 仍必须与 full 逐字节一致。 - with validation_stage('真实训练：前向、反向和 SGD 更新'): - report['training_loss'] = trainer.training_step() - log_progress(f'训练 loss={report["training_loss"]}') - check_round('真实训练更新') - - report['status'] = 'passed' - except Exception as error: - report.update(status='failed', error=repr(error)) - log_progress(f'验证失败：{error!r}') - raise - finally: - # 4. 关闭接收 actors、通信组和快照；即使验证失败也保存报告。 - try: - with validation_stage('关闭增量同步会话，释放通信组和快照'): - delta.close_sync_session() - except Exception as error: - report.update(status='failed', cleanup_error=repr(error)) - raise - finally: - with validation_stage(f'保存验证报告：{args.report}'): - save_report() - log_progress(f'验证结束：status={report["status"]}') - - -def main(): - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - add_platform_arguments(parser) # --platform {GPU,NPU}、--model-path、--offline、--num-devices、--dry-run 等 - parser.add_argument('--model', default='Qwen/Qwen3-0.6B', help='训练与 rollout 使用的同一模型') - parser.add_argument('--trainer-gpus', type=int, choices=(1, 2), default=2, help='FSDP2 训练卡数；1 表示不分片') - parser.add_argument('--tp', type=int, choices=(1, 2), default=1, help='每个 rollout 副本的卡数') - parser.add_argument('--dp', type=int, default=1, help='独立 rollout 副本数') - parser.add_argument('--strategy', choices=('native_fsdp', 'accelerate'), default='native_fsdp') - parser.add_argument('--cpu-offload', action='store_true', help='使用 Accelerate FSDP2 CPU offload') - parser.add_argument('--rounds', type=int, default=1, help='不同变化比例的验证轮数') - parser.add_argument('--gather-batch-size', type=int, default=32) - parser.add_argument('--bucket-bytes', type=int, default=64 << 20) - parser.add_argument('--devices-per-node', type=int, help='每个节点的卡数，默认单节点运行') - parser.add_argument('--report', default='delta-validation.json', help='结果报告路径') - args = parser.parse_args() - if args.dp < 1 or args.rounds < 1: - parser.error('--dp 和 --rounds 必须为正数') - if args.cpu_offload and args.strategy != 'accelerate': - parser.error('--cpu-offload 需要 --strategy accelerate') - # 解析模型/离线设置、预检卡数预算并打印计划；--dry-run 在这里退出，不启动 Ray。 - run_preflight_or_exit(args, parser, WORKER_EXTENSION_CLS) - run_validation(args) - - -if __name__ == '__main__': - main() diff --git a/cookbook/checkpoint_engine/device_utils.py b/cookbook/checkpoint_engine/device_utils.py deleted file mode 100644 index e166f0ded..000000000 --- a/cookbook/checkpoint_engine/device_utils.py +++ /dev/null @@ -1,232 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""GPU/NPU 通用的 delta 同步脚本辅助函数（delta_validation.py / delta_benchmark.py 共用）。 - -这里的大部分函数是纯 Python（不启动 Ray、不接触设备），因此可以在没有加速卡的机器上做 --dry-run 和单测。 -""" -from __future__ import annotations - -import argparse -import os -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -DEFAULT_NPU_CARDS = 4 - -# vLLM-Ascend 在 RL 场景下必须带 rl_config（它会关闭 FRACTAL_NZ，否则 start_weight_update / wake_up 报 weight_nz_mode 错误）。 -NPU_ADDITIONAL_CONFIG = {'rl_config': {'enabled': True}} - -TOKENIZER_FILES = ('tokenizer.json', 'tokenizer_config.json', 'tokenizer.model', 'vocab.json', 'merges.txt') - -OFFLINE_ENV = { - 'HF_HUB_OFFLINE': '1', - 'TRANSFORMERS_OFFLINE': '1', - 'HF_DATASETS_OFFLINE': '1', - 'MODELSCOPE_OFFLINE': '1', # 尽力而为：不同 modelscope 版本的离线开关不一致，真正的保护是本地目录检查 -} - -NPU_REQUIREMENTS = ( - "vLLM-Ascend 需要 additional_config={'rl_config': {'enabled': True}}（脚本已自动加入）；否则 " - 'start_weight_update/wake_up 会抛 weight_nz_mode 的 ValueError。', - '分离部署（trainer 与 rollout 使用不同的卡）通常不需要 sleep mode：NPU 默认 free_cache_engine=False；' - '若用 --free-cache-engine 打开，需要从源码编译的 vllm-ascend。', - "sampler 的 engine_args 必须带 weight_transfer_config={'backend': 'twinkle_delta_hccl'}（CUDA 为 twinkle_delta_ipc）。", - 'vLLM 必须提供 engine-owned WeightTransferEngine API（vLLM #44353）和 checkpoint patch API（vLLM #50723）；' - '可先运行 cookbook/checkpoint_engine/check_npu_env.py 检查。', - '支持范围：仅 vLLM、BF16 稠密（非 MoE）、非量化、无推测解码、vLLM 内部 data_parallel_size=1、' - 'trainer FSDP2 且 TP=PP=EP=1、无 LoRA、不支持 rebuild_group；rollout 可有多个副本（replica 为主、tp 为次）。', - 'GPU 与 NPU 的其余配置接口一致。以上内容尚未在真实 NPU 上验证。', -) - - -@dataclass(frozen=True) -class PlatformSpec: - name: str # twinkle Platform 名称，同时用作 DeviceGroup.device_type 与 CheckpointEngineManager(platform=) - torch_device: str # torch 设备模块名：torch.cuda / torch.npu - backend: str # vLLM weight_transfer_config 的 backend - - -PLATFORM_SPECS = { - 'GPU': PlatformSpec('GPU', 'cuda', 'twinkle_delta_ipc'), - 'NPU': PlatformSpec('NPU', 'npu', 'twinkle_delta_hccl'), -} -_PLATFORM_ALIASES = {'GPU': 'GPU', 'CUDA': 'GPU', 'NPU': 'NPU', 'ASCEND': 'NPU'} - - -def normalize_platform(value: str) -> str: - """把 gpu/cuda/npu/ascend（不区分大小写）归一化为 'GPU' 或 'NPU'。""" - try: - return _PLATFORM_ALIASES[str(value).upper()] - except KeyError: - raise argparse.ArgumentTypeError(f'未知平台 {value!r}，可选 GPU 或 NPU') from None - - -def add_platform_arguments(parser: argparse.ArgumentParser) -> None: - """两个脚本共用的 GPU/NPU、离线模型与 dry-run 参数。""" - parser.add_argument('--platform', type=normalize_platform, default='GPU', metavar='{GPU,NPU}', - help='设备平台（默认 GPU；cuda 是 GPU 的别名）') - parser.add_argument('--model-path', help='本地模型目录（不联网）；给出后覆盖 --model') - parser.add_argument('--offline', action=argparse.BooleanOptionalAction, default=None, - help='禁止任何联网下载（设置 HF_HUB_OFFLINE 等）；NPU 默认开启，GPU 默认关闭') - parser.add_argument('--num-devices', type=int, - help='可用卡数，用于预检卡数预算；NPU 默认 4，GPU 默认不限制') - parser.add_argument('--free-cache-engine', action=argparse.BooleanOptionalAction, default=None, - help='DeltaSyncConfig.free_cache_engine；GPU 默认开启，NPU 默认关闭（需要 sleep mode）') - parser.add_argument('--use-token-ids', action=argparse.BooleanOptionalAction, default=None, - help='让 vLLM 使用 skip_tokenizer_init，只用 token id 生成并比较；模型目录没有 tokenizer 时自动开启') - parser.add_argument('--dry-run', action='store_true', - help='只解析参数、构造 DeviceGroup/配置并做卡数预检，打印执行计划后退出（不启动 Ray、不接触设备）') - - -def has_tokenizer(model_dir: str) -> bool: - return any((Path(model_dir) / name).is_file() for name in TOKENIZER_FILES) - - -def resolve_model(args) -> None: - """解析 --model-path / --offline / --use-token-ids，原地更新 args。""" - if args.model_path: - args.model = args.model_path - if args.offline is None: - args.offline = args.platform == 'NPU' - if args.num_devices is None and args.platform == 'NPU': - args.num_devices = DEFAULT_NPU_CARDS - if args.free_cache_engine is None: - args.free_cache_engine = args.platform == 'GPU' - if args.use_token_ids is None: - local = Path(args.model).is_dir() - args.use_token_ids = bool(local and not has_tokenizer(args.model)) - - -def apply_offline_env(args) -> dict[str, str]: - """--offline 时设置离线环境变量（必须在启动 Ray 之前调用，Ray 子进程会继承）。""" - if not args.offline: - return {} - for key, value in OFFLINE_ENV.items(): - os.environ[key] = value - return dict(OFFLINE_ENV) - - -def preflight(args) -> tuple[list[str], list[str]]: - """返回 (errors, warnings)。不接触设备。""" - errors: list[str] = [] - warnings: list[str] = [] - trainer, dp, tp = args.trainer_gpus, args.dp, args.tp - if dp < 1 or tp < 1 or trainer < 1: - errors.append('--trainer-gpus、--dp、--tp 必须为正数') - return errors, warnings - needed = trainer + dp * tp - if args.num_devices is not None and needed > args.num_devices: - errors.append( - f'卡数预算不足：trainer {trainer} + rollout dp*tp={dp}*{tp}={dp * tp} = {needed} 张，' - f'超过 --num-devices={args.num_devices}。4 卡可验证的布局：' - '(a) --trainer-gpus 1 --dp 1 --tp 1；(b) --trainer-gpus 1 --dp 1 --tp 2；' - '(c) --trainer-gpus 2 --dp 2 --tp 1；(d) --trainer-gpus 2 --dp 1 --tp 2') - if args.platform == 'NPU' and dp > 1 and tp > 1 and (args.num_devices or DEFAULT_NPU_CARDS) <= DEFAULT_NPU_CARDS: - warnings.append(f'rollout 副本数>1 且 tp>1（dp={dp}, tp={tp}）至少需要 {1 + dp * tp} 张卡，' - f'4 卡机器无法验证“多副本 + tp>1”') - model_dir = Path(args.model) - if args.offline and not model_dir.is_dir(): - errors.append(f'--offline 下模型必须是本地目录（--model-path），当前为 {args.model!r}') - if model_dir.is_dir(): - if not (model_dir / 'config.json').is_file(): - errors.append(f'模型目录 {args.model} 缺少 config.json') - if not has_tokenizer(args.model) and not args.use_token_ids: - errors.append(f'模型目录 {args.model} 没有 tokenizer 文件，请加 --use-token-ids') - if args.platform == 'NPU' and args.free_cache_engine: - warnings.append('NPU 上启用 free_cache_engine 需要 sleep mode（已加入 enable_sleep_mode=True），' - '且需要从源码编译的 vllm-ascend') - if args.platform == 'NPU' and args.strategy == 'accelerate' and args.cpu_offload: - warnings.append('NPU 上的 Accelerate FSDP2 CPU offload 未验证') - return errors, warnings - - -def build_engine_args(args, worker_extension_cls: str) -> dict[str, Any]: - spec = PLATFORM_SPECS[args.platform] - engine_args: dict[str, Any] = dict( - tensor_parallel_size=args.tp, dtype='bfloat16', max_model_len=256, enforce_eager=True, - gpu_memory_utilization=0.4, worker_extension_cls=worker_extension_cls, - weight_transfer_config={'backend': spec.backend}) - if args.platform == 'NPU': - engine_args['additional_config'] = {'rl_config': dict(NPU_ADDITIONAL_CONFIG['rl_config'])} - if args.free_cache_engine: - engine_args['enable_sleep_mode'] = True - if args.use_token_ids: - engine_args['skip_tokenizer_init'] = True - return engine_args - - -def delta_config_kwargs(args) -> dict[str, Any]: - return dict(bucket_bytes=args.bucket_bytes, gather_batch_size=args.gather_batch_size, - free_cache_engine=bool(args.free_cache_engine)) - - -def build_device_groups(args) -> list: - from twinkle import DeviceGroup - spec = PLATFORM_SPECS[args.platform] - trainer, rollout = args.trainer_gpus, args.dp * args.tp - return [ - DeviceGroup(name='trainer', ranks=list(range(trainer)), device_type=spec.name), - DeviceGroup(name='rollout', ranks=list(range(trainer, trainer + rollout)), device_type=spec.name, - gpus_per_worker=args.tp), - ] - - -def nproc_per_node(args) -> int: - return args.devices_per_node or args.trainer_gpus + args.dp * args.tp - - -def print_plan(args, worker_extension_cls: str) -> None: - """打印解析后的执行计划（--dry-run 以及真实运行前都会打印）。""" - from twinkle.checkpoint_engine import DeltaSyncConfig - spec = PLATFORM_SPECS[args.platform] - print(f'== 执行计划 platform={args.platform} (torch.{spec.torch_device}) ==') - print(f'model : {args.model} (offline={args.offline}, use_token_ids={args.use_token_ids})') - print(f'cards : trainer={args.trainer_gpus} + rollout dp*tp={args.dp}*{args.tp}=' - f'{args.dp * args.tp} / num_devices={args.num_devices}') - for group in build_device_groups(args): - print(f'device group : {group}') - print(f'nproc_per_node : {nproc_per_node(args)}') - print(f'manager : CheckpointEngineManager(platform={args.platform!r}, mode=\'standalone\', ' - f"sync_backend='delta_sharded')") - print(f'backend : {spec.backend}') - print(f'engine_args : {build_engine_args(args, worker_extension_cls)}') - print(f'DeltaSyncConfig: {DeltaSyncConfig(**delta_config_kwargs(args))}') - if args.offline: - print(f'offline env : {OFFLINE_ENV}') - - -def run_preflight_or_exit(args, parser: argparse.ArgumentParser, worker_extension_cls: str) -> None: - """解析 -> 预检 -> 打印计划；--dry-run 时在此退出（返回前不接触设备）。""" - resolve_model(args) - errors, warnings = preflight(args) - if args.platform == 'NPU': - print('== NPU 使用须知（均未在真实 NPU 上验证）==') - for line in NPU_REQUIREMENTS: - print(f' - {line}') - for warning in warnings: - print(f'[警告] {warning}') - if errors: - parser.error('预检失败：\n ' + '\n '.join(errors)) - print_plan(args, worker_extension_cls) - if args.dry_run: - print('--dry-run：未启动 Ray，未接触设备，退出。') - raise SystemExit(0) - apply_offline_env(args) - - -def torch_device_module(platform_or_type: str): - """返回 torch.cuda 或 torch.npu（按 'GPU'/'NPU'/'cuda'/'npu' 选择）；NPU 会先 import torch_npu。""" - import torch - name = platform_or_type.lower() - name = PLATFORM_SPECS[normalize_platform(name)].torch_device if name in ('gpu', 'npu', 'ascend') else name - if name == 'npu': - import torch_npu # noqa: F401 - return getattr(torch, name) - - -def device_peak_bytes(device_type: str) -> int: - return int(torch_device_module(device_type).max_memory_allocated()) - - -def device_synchronize(device_type: str) -> None: - torch_device_module(device_type).synchronize() diff --git a/cookbook/checkpoint_engine/e2e_step_benchmark.py b/cookbook/checkpoint_engine/e2e_step_benchmark.py deleted file mode 100644 index d9f92d94e..000000000 --- a/cookbook/checkpoint_engine/e2e_step_benchmark.py +++ /dev/null @@ -1,257 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""端到端“一个训练 step”耗时对比：delta_sharded vs full。 - -一个 step = 训练计算（前向 + 反向 + AdamW，forward_backward + clip_grad_and_step）+ 权重同步（manager.sync_weights，阻塞到 rollout 就绪）。 -一个进程只测一个后端，各自写 JSON 报告，最后用 --compare 合并成对比表（不启动 Ray）。 - - # 1. 分别跑两个后端（两次之间重启 Ray 集群/actors，保证干净） - python -m cookbook.checkpoint_engine.e2e_step_benchmark --platform NPU --model-path /nas/disk1/Qwen3.5-9B \\ - --backend delta_sharded --trainer-gpus 2 --dp 1 --tp 2 --devices-per-node 2 \\ - --bucket-bytes 2147483648 --seq-len 128 --report e2e-delta.json - python -m cookbook.checkpoint_engine.e2e_step_benchmark --platform NPU --model-path /nas/disk1/Qwen3.5-9B \\ - --backend full --trainer-gpus 2 --dp 1 --tp 2 --devices-per-node 2 \\ - --bucket-bytes 2147483648 --seq-len 128 --report e2e-full.json - # 2. 合并对比 - python -m cookbook.checkpoint_engine.e2e_step_benchmark --compare e2e-delta.json e2e-full.json --md e2e.md - -默认只测真实 step：训练一步（前向 + 反向 + AdamW）后立刻 sync，变化比例由 AdamW 自然产生 -（报告里记录 delta 的实际 changed_ratio）。首次 seed 同步和 --warmup-steps 指定的预热步只记录、不进统计。 -""" -import argparse -import json -from datetime import datetime -from pathlib import Path -from statistics import median -from time import perf_counter - -from cookbook.checkpoint_engine.device_utils import add_platform_arguments, delta_config_kwargs, run_preflight_or_exit - - -def log_progress(message: str) -> None: - print(f'[{datetime.now():%Y-%m-%d %H:%M:%S}] {message}', flush=True) - - -def delta_info(manager) -> dict: - """从 manager.last_delta_metrics 里取 delta 的实际变化比例和 payload；full 后端返回 {}。""" - metrics = getattr(manager, 'last_delta_metrics', None) - if not metrics: - return {} - info = {'manager_seconds': metrics.get('total_seconds')} - trainers = metrics.get('trainers') - for item in trainers if isinstance(trainers, list) else [trainers]: - # 只有 trainer rank 0 带 checkpoint_engine 指标，其余 rank 为 None。 - checkpoint = item.get('checkpoint_engine') if isinstance(item, dict) else None - if isinstance(checkpoint, dict): - info['changed_ratio'] = checkpoint.get('checkpoint_engine/changed_ratio') - info['payload_mbytes'] = checkpoint.get('checkpoint_engine/payload_mbytes') - info['send_seconds'] = item.get('send_seconds') - break - return info - - -def timed_sync(manager) -> tuple[float, dict]: - started = perf_counter() - manager.sync_weights() - seconds = perf_counter() - started - return seconds, delta_info(manager) - - -def build_batch(batch_size: int, seq_len: int) -> list[dict]: - """构造已编码的最小训练 batch（InputFeature 列表）：全部 token 参与 loss，无需 tokenizer / 数据集。""" - row = [(i % 1000) + 1 for i in range(seq_len)] - return [{'input_ids': list(row), 'attention_mask': [1] * seq_len, 'labels': list(row)} - for _ in range(batch_size)] - - -def resolve(result): - """lazy_collect 打开时远程调用返回 callable，需要调用一次才会真正等待结果。""" - return result() if callable(result) else result - - -def train_step(trainer, batch, max_grad_norm: float) -> float: - """一次真实训练：前向 + 反向 + AdamW 更新（接口与 cookbook/transformers 一致）。返回耗时（秒）。 - - 结尾的 sync_device 用来等设备上的优化器 kernel 跑完，避免这部分时间漏进随后的 sync。 - """ - started = perf_counter() - resolve(trainer.forward_backward(inputs=batch)) - resolve(trainer.clip_grad_and_step(max_grad_norm=max_grad_norm)) - resolve(trainer.sync_device()) - return perf_counter() - started - - -def run_benchmark(args) -> None: - # 延迟导入：--compare 不需要 torch / Ray。 - from cookbook.checkpoint_engine.delta_validation import build_components, validation_stage - from twinkle.checkpoint_engine import CheckpointEngineManager, DeltaSyncConfig - - log_progress(f'e2e step benchmark：backend={args.backend}, platform={args.platform}, model={args.model}, ' - f'trainer={args.trainer_gpus}, rollout dp={args.dp} tp={args.tp}, ' - f'batch={args.batch_size} seq={args.seq_len}') - trainer, sampler = build_components(args) - if args.backend == 'delta_sharded': - manager = CheckpointEngineManager( - trainer, sampler, platform=args.platform, mode='standalone', sync_backend='delta_sharded', - delta_config=DeltaSyncConfig(**delta_config_kwargs(args))) - else: - manager = CheckpointEngineManager(trainer, sampler, platform=args.platform, mode='standalone') - - report = {'status': 'running', 'backend': args.backend, 'arguments': vars(args), 'seed': None, 'real_steps': []} - - def save_report(): - Path(args.report).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding='utf-8') - - try: - # 必须先 set_optimizer 再做首次同步：同步会触发 FSDP 懒包装，包装时才会把优化器绑到分片参数上。 - trainer.set_optimizer('AdamW', lr=args.lr, weight_decay=args.weight_decay, foreach=False) - batch = build_batch(args.batch_size, args.seq_len) - - # 首次 sync：delta 为 dense seed（建立快照和通信组），full 为首次全量；只记录，不进入统计。 - with validation_stage('首次 seed 同步（不计入统计）'): - seconds, info = timed_sync(manager) - report['seed'] = {'sync_seconds': seconds, **info} - log_progress(f'[seed] sync={seconds:.3f}s {info}') - save_report() - - # ---- real：训练计算 + 同步，变化比例由 AdamW 自然产生 ---- - total_real = args.warmup_steps + args.real_steps - for index in range(total_real): - warmup = index < args.warmup_steps - train_seconds = train_step(trainer, batch, args.max_grad_norm) - sync_seconds, info = timed_sync(manager) - row = {'index': index, 'warmup': warmup, 'train_seconds': train_seconds, 'sync_seconds': sync_seconds, - 'step_seconds': train_seconds + sync_seconds, **info} - report['real_steps'].append(row) - tag = 'warmup' if warmup else 'real' - log_progress(f'[{tag} {index + 1}/{total_real}] train={train_seconds:.3f}s sync={sync_seconds:.3f}s ' - f'step={row["step_seconds"]:.3f}s {info}') - save_report() - - report['status'] = 'passed' - except Exception as error: - report.update(status='failed', error=repr(error)) - log_progress(f'基准失败：{error!r}') - raise - finally: - try: - if args.backend == 'delta_sharded': - with validation_stage('关闭增量同步会话，释放通信组和快照'): - manager.close_sync_session() - finally: - save_report() - log_progress(f'报告已保存：{Path(args.report).resolve()}（status={report["status"]}）') - - print_single_summary(report) - - -def _med(rows, key): - values = [row[key] for row in rows if row.get(key) is not None] - return median(values) if values else None - - -def _fmt(value, spec='.3f', suffix=''): - return '-' if value is None else f'{value:{spec}}{suffix}' - - -def print_single_summary(report) -> None: - measured = [row for row in report['real_steps'] if not row['warmup']] - print(f"\nbackend={report['backend']} real steps（中位数，{len(measured)} 步，已去掉 warmup）") - print(f" train={_fmt(_med(measured, 'train_seconds'))}s sync={_fmt(_med(measured, 'sync_seconds'))}s " - f"step={_fmt(_med(measured, 'step_seconds'))}s changed_ratio={_fmt(_med(measured, 'changed_ratio'), '.4f')} " - f"payload={_fmt(_med(measured, 'payload_mbytes'), '.1f', 'MB')}") - - -def compare_reports(paths, md_path=None) -> None: - reports = {} - for path in paths: - data = json.loads(Path(path).read_text(encoding='utf-8')) - if data.get('backend') in reports: - raise SystemExit(f'--compare 需要一份 delta_sharded 和一份 full 报告，但 {path} 的 backend 重复') - if data.get('status') != 'passed': - log_progress(f'警告：{path} 的 status={data.get("status")}，结果可能不完整') - reports[data['backend']] = data - if set(reports) != {'delta_sharded', 'full'}: - raise SystemExit(f'--compare 需要 delta_sharded 与 full 各一份报告，实际: {sorted(reports)}') - delta, full = reports['delta_sharded'], reports['full'] - - lines = ['# 端到端训练 step 耗时：delta vs full', ''] - args = delta['arguments'] - lines += [f"环境：platform={args['platform']}，model={args['model']}，trainer={args['trainer_gpus']} 卡，" - f"rollout dp={args['dp']} × tp={args['tp']}，batch={args['batch_size']} × seq={args['seq_len']}，" - f"bucket={args['bucket_bytes'] >> 20}MB。", - '一个 step = 训练计算（前向+反向+AdamW）+ 权重同步（阻塞到 rollout 就绪）。`speedup = full step / delta step`，' - '>1 表示 delta 更快。', ''] - - def real_rows(report): - return [row for row in report['real_steps'] if not row['warmup']] - - d_real, f_real = real_rows(delta), real_rows(full) - lines += ['## 真实 step（AdamW 自然产生的变化比例）', '', - '| backend | train (s) | sync (s) | step (s) | changed_ratio | payload (MB) |', - '| --- | --- | --- | --- | --- | --- |'] - for name, rows in (('delta_sharded', d_real), ('full', f_real)): - lines.append(f"| {name} | {_fmt(_med(rows, 'train_seconds'))} | {_fmt(_med(rows, 'sync_seconds'))} | " - f"{_fmt(_med(rows, 'step_seconds'))} | {_fmt(_med(rows, 'changed_ratio'), '.4f')} | " - f"{_fmt(_med(rows, 'payload_mbytes'), '.1f')} |") - d_step, f_step = _med(d_real, 'step_seconds'), _med(f_real, 'step_seconds') - if d_step and f_step: - lines += ['', f'speedup（real）= {f_step / d_step:.2f}x（full {f_step:.3f}s / delta {d_step:.3f}s）'] - - text = '\n'.join(lines) + '\n' - print(text) - if md_path: - Path(md_path).write_text(text, encoding='utf-8') - log_progress(f'对比表已写入：{Path(md_path).resolve()}') - - -def main(): - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('--compare', nargs=2, metavar=('REPORT_A', 'REPORT_B'), - help='合并 delta_sharded 与 full 两份 JSON 报告并输出对比表（不启动 Ray）') - parser.add_argument('--md', help='--compare 时把对比表另存为 markdown') - add_platform_arguments(parser) - parser.add_argument('--model', default='Qwen/Qwen3-0.6B', help='训练与 rollout 使用的同一模型') - parser.add_argument('--backend', choices=('delta_sharded', 'full'), default='delta_sharded') - parser.add_argument('--trainer-gpus', type=int, choices=(1, 2), default=2, help='FSDP2 训练卡数；1 表示不分片') - parser.add_argument('--tp', type=int, choices=(1, 2), default=1, help='每个 rollout 副本的卡数') - parser.add_argument('--dp', type=int, default=1, help='独立 rollout 副本数') - parser.add_argument('--strategy', choices=('native_fsdp', 'accelerate'), default='native_fsdp') - parser.add_argument('--cpu-offload', action='store_true', help='使用 Accelerate FSDP2 CPU offload') - parser.add_argument('--gather-batch-size', type=int, default=32) - parser.add_argument('--bucket-bytes', type=int, default=2048 << 20, help='默认 2048MB，与 verl 一致') - parser.add_argument('--devices-per-node', type=int, help='每个节点的卡数，默认单节点运行') - parser.add_argument('--batch-size', type=int, default=None, - help='全局 batch size，按 DP 切分给 trainer 各卡，必须能被 --trainer-gpus 整除；默认等于 --trainer-gpus') - parser.add_argument('--seq-len', type=int, default=128, - help='每条样本的 token 数。全参数 AdamW 的优化器状态很占显存（9B 模型每卡约 30~50GB），' - '默认取较小值，OOM 时调小或加 --strategy accelerate --cpu-offload') - parser.add_argument('--lr', type=float, default=1e-5, help='AdamW 学习率；它决定 real step 自然产生的变化比例') - parser.add_argument('--weight-decay', type=float, default=0.01) - parser.add_argument('--max-grad-norm', type=float, default=1.0) - parser.add_argument('--warmup-steps', type=int, default=1, help='real step 的预热步数（不进入统计）') - parser.add_argument('--real-steps', type=int, default=3, help='参与统计的 real step 数') - parser.add_argument('--report', default=None, help='JSON 报告路径，默认 e2e-.json') - args = parser.parse_args() - - if args.compare: - compare_reports(args.compare, args.md) - return - if args.report is None: - args.report = f'e2e-{args.backend}.json' - if args.dp < 1 or args.real_steps < 0 or args.warmup_steps < 0: - parser.error('--dp 必须为正数，--real-steps/--warmup-steps 不能为负') - if args.batch_size is None: - args.batch_size = args.trainer_gpus - if args.batch_size < 1 or args.seq_len < 1: - parser.error('--batch-size 和 --seq-len 必须为正数') - if args.batch_size % args.trainer_gpus: - parser.error(f'--batch-size={args.batch_size} 必须能被 --trainer-gpus={args.trainer_gpus} 整除（forward_backward 按 DP 切分）') - if args.cpu_offload and args.strategy != 'accelerate': - parser.error('--cpu-offload 需要 --strategy accelerate') - # 复用 delta_validation 的 worker extension 做预检（延迟导入常量所在模块会拉起 torch，所以直接写字符串）。 - run_preflight_or_exit(args, parser, 'cookbook.checkpoint_engine.delta_validation.ValidationWorker') - run_benchmark(args) - - -if __name__ == '__main__': - main() diff --git a/delta_perf.md b/delta_perf.md deleted file mode 100644 index 6a8c8496e..000000000 --- a/delta_perf.md +++ /dev/null @@ -1,69 +0,0 @@ -# Delta vs Full 权重同步性能 - -测试环境：Qwen3.5-9B（BF16，整模型 17.53 GB / 760 个张量），2 卡 FSDP2 trainer + 2 卡 rollout（TP=2），跨机 HCCL，bucket=2048MB。 -`delta time` 为发送侧传输耗时（算 diff + 跨 rank gather-v + 编码 + HCCL 广播），`full payload` 为每轮全量同步实际广播的字节数。 - -| changed_ratio | delta payload | delta time | full payload | reduction vs full | -| ------------- | ------------- | ---------- | ------------ | ----------------- | -| 0.5% | 269.2 MB | 3.59 s | 17.53 GB | 66.7x | -| 1.0% | 538.4 MB | 4.31 s | 17.53 GB | 33.3x | -| 2.0% | 1076.9 MB | 5.62 s | 17.53 GB | 16.7x | -| 3.0% | 1615.3 MB | 6.76 s | 17.53 GB | 11.1x | - -说明： -- `reduction vs full` = full payload / delta payload，即 delta 相比全量同步减少的网络传输量。 -- full payload 每轮恒为整模型 17.53 GB（本环境链路 ~9.5 GB/s，全量广播约 1.8 s）。 -- delta payload 每个变化元素 6 字节（int32 位置 + bf16 值），随改动比例线性增长。 -- delta time 可近似分解为「固定开销 ~2.8 s（diff 全扫 + gather + 逐参数）+ payload / ~400 MB/s」。 - -# 端到端单 step 对比（Qwen3.5-4B） - -环境：2 卡 FSDP2 trainer + 2 卡 rollout（TP=2），batch=2 × seq=2048，AdamW lr=1e-5，bucket=2048MB。 -一个 step = 训练计算（前向 + 反向 + AdamW）+ 权重同步（阻塞到 rollout 就绪）。 - -| backend | train (s) | sync (s) | step (s) | changed_ratio | payload (MB) | -| ------------- | --------- | --------- | --------- | ------------- | ------------ | -| delta_sharded | 3.670 | 36.998 | 40.927 | 14.11% | 4179.3 | -| full | 3.677 | 3.098 | 6.775 | - | 9870.5 | - -此工作点下 delta 慢于 full（speedup 0.17x）：payload 只降 2.4x，而稀疏路径吞吐约为 dense 广播的 1/10 ~ 1/20。 -改动比例由优化器配置与训练阶段决定，不是 delta 实现的属性；上游参考实现在稳态 RL 下报的是 1 ~ 2.3%。 - -## 耗时分解（取上表那一步，sync 37.0 s） - -发送端 wire_time 与接收端 apply 是流水重叠的（flushes=2），因此各项之和大于 37 s。 - -| 阶段 | 耗时 (s) | 占 sync | -| ---------------------------- | -------- | ------- | -| 接收端 apply（decode+scatter） | 27.0 | 73% | -| 发送端 broadcast | 11.0 | 30% | -| assemble（编码 + checksum） | 5.2 | 14% | -| produce（diff 全扫 + 提取） | 4.7 | 13% | -| gather（跨 rank） | 2.5 | 7% | -| 接收端 broadcast_recv | 0.1 | ~0% | - -两个大头： -- **接收端稀疏 scatter**：4179MB / 27.0s ≈ 155 MB/s。按索引散写是随机访存，9B 上实测约 110 MB/s，同一量级，属稀疏路径的固有吞吐。 -- **发送端 broadcast**：4179MB / 11.0s ≈ 380 MB/s，而 full 的连续广播为 9870MB / 3.1s ≈ 3200 MB/s，慢约 8 倍。delta 把 positions 与 values 分两次广播，且每个 flush 后 `torch.npu.synchronize()`；full 是单块连续 buffer 一次推完。payload 小且单 flush 时（9B 的 1GB 场景）broadcast 可忽略，payload 大且多 flush 时才显现。 - -## 改动比例随训练下降 - -同一份固定数据反复训练，梯度收敛，更新量递减： - -| step | changed_ratio | payload (MB) | sync (s) | step (s) | -| ------ | ------------- | ------------ | -------- | -------- | -| warmup | 21.64% | 6408.5 | 56.1 | 64.2 | -| real 2 | 15.83% | 4686.8 | 42.1 | 45.6 | -| real 3 | 14.11% | 4179.3 | 37.0 | 40.9 | -| real 4 | 12.65% | 3744.4 | 31.2 | 34.9 | - -14.11% 不是稳态值，只是该合成训练设置的早期状态。 - -## 优化优先级 - -1. 接收端 scatter（73%）：合并 kernel、连续段走整段 copy、避免 int64 索引中转。 -2. 发送端 broadcast（30%）：positions 与 values 合并为一次广播、去掉每 flush 的强制同步、增大 bucket 以减少 flush 数。 -3. assemble 的 checksum（14%）：verify 关闭时跳过。 -4. produce 的 diff 全扫（13%）：按 optimizer 标记的 dirty 参数增量化，避免每次 O(模型) 扫描。 -5. gather（7%）：此规模无需改动。 - From 2eb12f432d63ff706670236985e8f8c2f1281510 Mon Sep 17 00:00:00 2001 From: kevssim Date: Thu, 22 Oct 2026 11:41:46 +0800 Subject: [PATCH 11/12] refactor(delta): split delta sync into subpackages - checkpoint_engine/delta: protocol/ (wire, decode), trainer/ (mixin, spec, export, backends), transport/ (pipeline, gather, nccl_engine, hccl_engine), rollout/ (mixin, cuda_receiver, hccl_receiver); CudaDeltaReceiver split out of the rollout mixin. - sampler/vllm_sampler/delta: transfer_engine (was delta_weight_transfer), cuda_ipc (was bucketed_weight_transfer), weight_update_utils. - transformers: restore _checkpoint_tensor so CPU-offloaded DTensors are staged on device before full_tensor (dropped in the upstream merge). --- .../checkpoint_engine/delta/__init__.py | 36 +++++---- .../checkpoint_engine/delta/platform.py | 6 +- .../delta/protocol/__init__.py | 2 + .../delta/{ => protocol}/decode.py | 0 .../delta/{ => protocol}/wire.py | 0 .../delta/rollout/__init__.py | 2 + .../delta/rollout/cuda_receiver.py | 79 +++++++++++++++++++ .../delta/{ => rollout}/hccl_receiver.py | 0 .../delta/{rollout.py => rollout/mixin.py} | 75 ++---------------- .../delta/trainer/__init__.py | 2 + .../delta/{ => trainer}/backends/__init__.py | 0 .../delta/{ => trainer}/backends/base.py | 0 .../delta/{ => trainer}/backends/fsdp.py | 0 .../delta/{ => trainer}/export.py | 0 .../delta/{trainer.py => trainer/mixin.py} | 4 +- .../delta/{ => trainer}/spec.py | 0 .../delta/transport/__init__.py | 2 + .../delta/{ => transport}/gather.py | 0 .../delta/{ => transport}/hccl_engine.py | 6 +- .../{engine.py => transport/nccl_engine.py} | 2 +- .../delta/{ => transport}/pipeline.py | 4 +- .../model/transformers/transformers.py | 19 +++-- .../sampler/vllm_sampler/delta/__init__.py | 8 ++ .../cuda_ipc.py} | 0 .../transfer_engine.py} | 12 +-- .../{ => delta}/weight_update_utils.py | 0 .../sampler/vllm_sampler/vllm_engine.py | 4 +- .../sampler/vllm_sampler/vllm_sampler.py | 2 +- .../vllm_sampler/vllm_worker_extension.py | 8 +- tests/checkpoint_engine/_hccl_loopback.py | 12 +-- .../checkpoint_engine/test_delta_checksum.py | 6 +- tests/checkpoint_engine/test_delta_hccl.py | 6 +- .../test_delta_hccl_receiver.py | 14 ++-- .../test_delta_npu_orchestration.py | 4 +- .../checkpoint_engine/test_delta_pipeline.py | 4 +- .../checkpoint_engine/test_delta_platform.py | 14 ++-- tests/checkpoint_engine/test_delta_rollout.py | 11 +-- tests/checkpoint_engine/test_delta_trainer.py | 28 +++---- tests/checkpoint_engine/test_delta_vllm.py | 22 +++--- 39 files changed, 220 insertions(+), 174 deletions(-) create mode 100644 src/twinkle/checkpoint_engine/delta/protocol/__init__.py rename src/twinkle/checkpoint_engine/delta/{ => protocol}/decode.py (100%) rename src/twinkle/checkpoint_engine/delta/{ => protocol}/wire.py (100%) create mode 100644 src/twinkle/checkpoint_engine/delta/rollout/__init__.py create mode 100644 src/twinkle/checkpoint_engine/delta/rollout/cuda_receiver.py rename src/twinkle/checkpoint_engine/delta/{ => rollout}/hccl_receiver.py (100%) rename src/twinkle/checkpoint_engine/delta/{rollout.py => rollout/mixin.py} (80%) create mode 100644 src/twinkle/checkpoint_engine/delta/trainer/__init__.py rename src/twinkle/checkpoint_engine/delta/{ => trainer}/backends/__init__.py (100%) rename src/twinkle/checkpoint_engine/delta/{ => trainer}/backends/base.py (100%) rename src/twinkle/checkpoint_engine/delta/{ => trainer}/backends/fsdp.py (100%) rename src/twinkle/checkpoint_engine/delta/{ => trainer}/export.py (100%) rename src/twinkle/checkpoint_engine/delta/{trainer.py => trainer/mixin.py} (96%) rename src/twinkle/checkpoint_engine/delta/{ => trainer}/spec.py (100%) create mode 100644 src/twinkle/checkpoint_engine/delta/transport/__init__.py rename src/twinkle/checkpoint_engine/delta/{ => transport}/gather.py (100%) rename src/twinkle/checkpoint_engine/delta/{ => transport}/hccl_engine.py (98%) rename src/twinkle/checkpoint_engine/delta/{engine.py => transport/nccl_engine.py} (97%) rename src/twinkle/checkpoint_engine/delta/{ => transport}/pipeline.py (99%) create mode 100644 src/twinkle/sampler/vllm_sampler/delta/__init__.py rename src/twinkle/sampler/vllm_sampler/{bucketed_weight_transfer.py => delta/cuda_ipc.py} (100%) rename src/twinkle/sampler/vllm_sampler/{delta_weight_transfer.py => delta/transfer_engine.py} (97%) rename src/twinkle/sampler/vllm_sampler/{ => delta}/weight_update_utils.py (100%) diff --git a/src/twinkle/checkpoint_engine/delta/__init__.py b/src/twinkle/checkpoint_engine/delta/__init__.py index c0caea509..a4b62e1ee 100644 --- a/src/twinkle/checkpoint_engine/delta/__init__.py +++ b/src/twinkle/checkpoint_engine/delta/__init__.py @@ -2,23 +2,29 @@ """Sparse delta weight synchronization. Layout: - config.py DeltaSyncConfig and the TypedDicts exchanged between trainer and rollout - pipeline.py DeltaSyncMixin: transport-agnostic gather/bucket/flush pipeline and seed/steady state machine - engine.py DeltaShardedCheckpointEngine: NCCL transport (hooks + rollout receive) over the pipeline - hccl_engine.py DeltaHCCLCheckpointEngine: HCCL (Ascend NPU) sender transport over the pipeline - ../zmq_pubsub.py zmq PUB/SUB manifest-channel helpers shared by the NCCL and HCCL engines - platform.py resolve_delta_engine_cls: GPU -> NCCL engine, NPU -> HCCL engine (lazy import) - wire.py DeltaFlush / DeltaParam wire structures and checksum - gather.py sparse gather of per-rank deltas to rank 0 - spec.py ShardSpec / BlockPlacement: where a local shard sits in the full tensor - export.py snapshot, diff and HF-coordinate delta export helpers - trainer.py DeltaTrainerMixin: trainer-side Ray lifecycle - rollout.py CudaDeltaReceiver / DeltaSamplerMixin: rollout-side lifecycle (CUDA: Ray receiver actors + IPC; - NPU: no actors, vLLM workers join the trainer's HCCL group via the worker extension) - backends/ per-training-backend export hooks (base.py contract, fsdp.py implementation) + 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 .engine import DeltaShardedCheckpointEngine +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/platform.py b/src/twinkle/checkpoint_engine/delta/platform.py index 8e6169354..68022a6f6 100644 --- a/src/twinkle/checkpoint_engine/delta/platform.py +++ b/src/twinkle/checkpoint_engine/delta/platform.py @@ -2,7 +2,7 @@ """Platform -> delta transport resolution.""" from __future__ import annotations -_HCCL_MODULE = 'twinkle.checkpoint_engine.delta.hccl_engine' +_HCCL_MODULE = 'twinkle.checkpoint_engine.delta.transport.hccl_engine' # Trainer-side torch device type per supported platform. DELTA_DEVICES = {'GPU': 'cuda', 'NPU': 'npu'} @@ -14,11 +14,11 @@ def resolve_delta_engine_cls(platform_name: str) -> type: Imports are lazy so the GPU path never touches NPU-specific modules. """ if platform_name == 'GPU': - from .engine import DeltaShardedCheckpointEngine + from .transport.nccl_engine import DeltaShardedCheckpointEngine return DeltaShardedCheckpointEngine if platform_name == 'NPU': try: - from .hccl_engine import DeltaHCCLCheckpointEngine + from .transport.hccl_engine import DeltaHCCLCheckpointEngine except ModuleNotFoundError as error: if error.name != _HCCL_MODULE: raise 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/decode.py b/src/twinkle/checkpoint_engine/delta/protocol/decode.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/decode.py rename to src/twinkle/checkpoint_engine/delta/protocol/decode.py diff --git a/src/twinkle/checkpoint_engine/delta/wire.py b/src/twinkle/checkpoint_engine/delta/protocol/wire.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/wire.py rename to src/twinkle/checkpoint_engine/delta/protocol/wire.py 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/hccl_receiver.py b/src/twinkle/checkpoint_engine/delta/rollout/hccl_receiver.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/hccl_receiver.py rename to src/twinkle/checkpoint_engine/delta/rollout/hccl_receiver.py diff --git a/src/twinkle/checkpoint_engine/delta/rollout.py b/src/twinkle/checkpoint_engine/delta/rollout/mixin.py similarity index 80% rename from src/twinkle/checkpoint_engine/delta/rollout.py rename to src/twinkle/checkpoint_engine/delta/rollout/mixin.py index 16c748e48..799dca1f7 100644 --- a/src/twinkle/checkpoint_engine/delta/rollout.py +++ b/src/twinkle/checkpoint_engine/delta/rollout/mixin.py @@ -1,88 +1,23 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Rollout-side lifecycle for the delta checkpoint engine. -CUDA: Ray receiver actors colocated with each vLLM worker feed it over IPC. Ascend NPU: no receiver actors; -the vLLM workers join the trainer's HCCL group themselves and the sampler only drives vLLM's lifecycle RPCs. +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 -import os -from collections.abc import Iterator from typing import Any, TYPE_CHECKING -import torch - from twinkle import Platform, remote_function -from .config import (DeltaReceiverPlacement, DeltaSyncConfig, DeltaTransportEndpoint, +from ..config import (DeltaReceiverPlacement, DeltaSyncConfig, DeltaTransportEndpoint, DeltaWorkerCapabilities, DeltaWorkerMetadata) -from .engine import DeltaShardedCheckpointEngine +from .cuda_receiver import CudaDeltaReceiver if TYPE_CHECKING: from ray.actor import ActorHandle -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.bucketed_weight_transfer 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() - class DeltaSamplerMixin: """Sampler-side delta lifecycle RPCs (device-aware: CUDA uses receiver actors, NPU does not).""" 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/backends/__init__.py b/src/twinkle/checkpoint_engine/delta/trainer/backends/__init__.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/backends/__init__.py rename to src/twinkle/checkpoint_engine/delta/trainer/backends/__init__.py diff --git a/src/twinkle/checkpoint_engine/delta/backends/base.py b/src/twinkle/checkpoint_engine/delta/trainer/backends/base.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/backends/base.py rename to src/twinkle/checkpoint_engine/delta/trainer/backends/base.py diff --git a/src/twinkle/checkpoint_engine/delta/backends/fsdp.py b/src/twinkle/checkpoint_engine/delta/trainer/backends/fsdp.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/backends/fsdp.py rename to src/twinkle/checkpoint_engine/delta/trainer/backends/fsdp.py diff --git a/src/twinkle/checkpoint_engine/delta/export.py b/src/twinkle/checkpoint_engine/delta/trainer/export.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/export.py rename to src/twinkle/checkpoint_engine/delta/trainer/export.py diff --git a/src/twinkle/checkpoint_engine/delta/trainer.py b/src/twinkle/checkpoint_engine/delta/trainer/mixin.py similarity index 96% rename from src/twinkle/checkpoint_engine/delta/trainer.py rename to src/twinkle/checkpoint_engine/delta/trainer/mixin.py index 3140ddfdc..34c106eb6 100644 --- a/src/twinkle/checkpoint_engine/delta/trainer.py +++ b/src/twinkle/checkpoint_engine/delta/trainer/mixin.py @@ -10,8 +10,8 @@ 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 +from ..config import DeltaMasterMetadata, DeltaSyncConfig, DeltaTrainerMetadata, DeltaTransportEndpoint +from ..platform import DELTA_DEVICES, resolve_delta_engine_cls class DeltaTrainerMixin: diff --git a/src/twinkle/checkpoint_engine/delta/spec.py b/src/twinkle/checkpoint_engine/delta/trainer/spec.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/spec.py rename to src/twinkle/checkpoint_engine/delta/trainer/spec.py 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/gather.py b/src/twinkle/checkpoint_engine/delta/transport/gather.py similarity index 100% rename from src/twinkle/checkpoint_engine/delta/gather.py rename to src/twinkle/checkpoint_engine/delta/transport/gather.py diff --git a/src/twinkle/checkpoint_engine/delta/hccl_engine.py b/src/twinkle/checkpoint_engine/delta/transport/hccl_engine.py similarity index 98% rename from src/twinkle/checkpoint_engine/delta/hccl_engine.py rename to src/twinkle/checkpoint_engine/delta/transport/hccl_engine.py index 02b8ec455..5badd52bb 100644 --- a/src/twinkle/checkpoint_engine/delta/hccl_engine.py +++ b/src/twinkle/checkpoint_engine/delta/transport/hccl_engine.py @@ -32,9 +32,9 @@ 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 ...base import CheckpointEngine +from ...hccl_checkpoint_engine import MasterMetadata +from ...zmq_pubsub import start_publisher from .pipeline import DeltaSyncMixin logger = logging.getLogger(__name__) diff --git a/src/twinkle/checkpoint_engine/delta/engine.py b/src/twinkle/checkpoint_engine/delta/transport/nccl_engine.py similarity index 97% rename from src/twinkle/checkpoint_engine/delta/engine.py rename to src/twinkle/checkpoint_engine/delta/transport/nccl_engine.py index 1481886f9..ad69d9edd 100644 --- a/src/twinkle/checkpoint_engine/delta/engine.py +++ b/src/twinkle/checkpoint_engine/delta/transport/nccl_engine.py @@ -9,7 +9,7 @@ import torch -from ..nccl_checkpoint_engine import MasterMetadata, NCCLCheckpointEngine, _pg_broadcast +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 diff --git a/src/twinkle/checkpoint_engine/delta/pipeline.py b/src/twinkle/checkpoint_engine/delta/transport/pipeline.py similarity index 99% rename from src/twinkle/checkpoint_engine/delta/pipeline.py rename to src/twinkle/checkpoint_engine/delta/transport/pipeline.py index 05154be9f..3b2efb393 100644 --- a/src/twinkle/checkpoint_engine/delta/pipeline.py +++ b/src/twinkle/checkpoint_engine/delta/transport/pipeline.py @@ -30,8 +30,8 @@ import zmq from .gather import gather_slot_entries_to_rank0 -from .wire import DeltaFlush, DeltaParam -from .wire import checksum as _checksum +from ..protocol.wire import DeltaFlush, DeltaParam +from ..protocol.wire import checksum as _checksum logger = logging.getLogger(__name__) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 9f68d8d25..a70ca903e 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -27,8 +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.backends import FSDP2DeltaBackend -from twinkle.checkpoint_engine.delta.trainer import DeltaTrainerMixin +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 @@ -2789,6 +2789,13 @@ 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. @@ -2830,7 +2837,7 @@ def weight_generator(): for name, tensor in model.state_dict().items(): 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 @@ -2846,7 +2853,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 @@ -2857,7 +2864,7 @@ 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.backends.fsdp import convert_weight_keys + from twinkle.checkpoint_engine.delta.trainer.backends.fsdp import convert_weight_keys add_base_layer = not merge_and_sync state_dict = convert_weight_keys(model.state_dict(), model) @@ -2867,7 +2874,7 @@ def weight_generator(): for name, tensor in state_dict.items(): 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/bucketed_weight_transfer.py b/src/twinkle/sampler/vllm_sampler/delta/cuda_ipc.py similarity index 100% rename from src/twinkle/sampler/vllm_sampler/bucketed_weight_transfer.py rename to src/twinkle/sampler/vllm_sampler/delta/cuda_ipc.py diff --git a/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py b/src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py similarity index 97% rename from src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py rename to src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py index a9c0ae78c..69d2eaec2 100644 --- a/src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py +++ b/src/twinkle/sampler/vllm_sampler/delta/transfer_engine.py @@ -120,7 +120,7 @@ def decode_delta_payload( ) -> tuple[str, list[CheckpointWeightPatch]]: """Validate and decode one DeltaFlush into vLLM checkpoint patches.""" - from twinkle.checkpoint_engine.delta.decode import decode_delta_wire + from twinkle.checkpoint_engine.delta.protocol.decode import decode_delta_wire CheckpointWeightPatch, _ = _checkpoint_patch_api() encoding, patches = decode_delta_wire(named_tensors) @@ -186,7 +186,7 @@ def start_weight_update(self) -> None: self._session_encoding = None def _receive_payload(self, *, zmq_handle: str) -> list[tuple[str, torch.Tensor]]: - from twinkle.sampler.vllm_sampler.bucketed_weight_transfer import ( + from twinkle.sampler.vllm_sampler.delta.cuda_ipc import ( BucketedWeightReceiver, ) @@ -217,7 +217,7 @@ def retain_bucket(weights: list[tuple[str, torch.Tensor]], _is_last: bool) -> No 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.weight_update_utils import ( + from twinkle.sampler.vllm_sampler.delta.weight_update_utils import ( drop_tied_alias_patches, prepare_sparse_weight_loaders, ) @@ -277,7 +277,7 @@ def finish_weight_update(self) -> None: 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.weight_update_utils import ( + from twinkle.sampler.vllm_sampler.delta.weight_update_utils import ( refresh_derived_weights, ) @@ -366,7 +366,7 @@ def init_transfer_engine(self, init_info: TwinkleDeltaHCCLInitInfo) -> None: def _ensure_channel(self, group_rank: int): """Create (once) and return the (manifest SUB, communicator) pair, SUB first.""" - from twinkle.checkpoint_engine.delta.hccl_receiver import build_delta_hccl_communicator + 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: @@ -389,7 +389,7 @@ def _ensure_channel(self, group_rank: int): return self._manifest_sub, self._comm def receive_weights(self, update_info: TwinkleDeltaHCCLUpdateInfo) -> None: - from twinkle.checkpoint_engine.delta.hccl_receiver import pump_hccl_delta_session + from twinkle.checkpoint_engine.delta.rollout.hccl_receiver import pump_hccl_delta_session try: sub, comm = self._ensure_channel(update_info.group_rank) diff --git a/src/twinkle/sampler/vllm_sampler/weight_update_utils.py b/src/twinkle/sampler/vllm_sampler/delta/weight_update_utils.py similarity index 100% rename from src/twinkle/sampler/vllm_sampler/weight_update_utils.py rename to src/twinkle/sampler/vllm_sampler/delta/weight_update_utils.py diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index 3998a2df6..906d18c32 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -379,12 +379,12 @@ def _create_engine(self): if backend == 'twinkle_delta_ipc': engine_config.setdefault('enable_sleep_mode', True) self.enable_sleep_mode = engine_config['enable_sleep_mode'] - from .delta_weight_transfer import register_twinkle_delta_weight_transfer_engine, require_vllm_delta_support + 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_weight_transfer import (register_twinkle_delta_hccl_weight_transfer_engine, + 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() diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 1bd20dd30..c48585a42 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -10,7 +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 import DeltaSamplerMixin +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 diff --git a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py index 8486f68f3..b5343f1ea 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_worker_extension.py @@ -119,16 +119,16 @@ def __new__(cls, *args, **kwargs): VLLMLoraWeights()(None) config = getattr(kwargs.get('vllm_config'), 'weight_transfer_config', None) if getattr(config, 'backend', None) == 'twinkle_delta_ipc': - from .delta_weight_transfer import register_twinkle_delta_weight_transfer_engine + 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_weight_transfer import register_twinkle_delta_hccl_weight_transfer_engine + 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_weight_transfer import TwinkleDeltaIPCWeightTransferEngine, is_moe_model + 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 " @@ -184,7 +184,7 @@ def update_twinkle_delta_weights(self, update_info: Optional[List[Dict]] = None) def shutdown_delta_weight_transfer(self) -> None: """Release the delta engine's transport channel (HCCL SUB + communicator); no-op for other engines.""" - from .delta_weight_transfer import TwinkleDeltaIPCWeightTransferEngine + from .delta.transfer_engine import TwinkleDeltaIPCWeightTransferEngine engine = getattr(self, 'weight_transfer_engine', None) if isinstance(engine, TwinkleDeltaIPCWeightTransferEngine): engine.shutdown() diff --git a/tests/checkpoint_engine/_hccl_loopback.py b/tests/checkpoint_engine/_hccl_loopback.py index 6219eb55f..6952f67b8 100644 --- a/tests/checkpoint_engine/_hccl_loopback.py +++ b/tests/checkpoint_engine/_hccl_loopback.py @@ -104,7 +104,9 @@ 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 import hccl_engine, hccl_receiver + 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) @@ -160,8 +162,8 @@ def factory(rank): return LoopComm() _install_stand_ins(factory, build_log) - from twinkle.checkpoint_engine.delta.hccl_engine import DeltaHCCLCheckpointEngine - from twinkle.checkpoint_engine.delta.pipeline import DeltaSyncMixin + from twinkle.checkpoint_engine.delta.transport.hccl_engine import DeltaHCCLCheckpointEngine + from twinkle.checkpoint_engine.delta.transport.pipeline import DeltaSyncMixin class LoopbackEngine(DeltaHCCLCheckpointEngine): @@ -215,7 +217,7 @@ def recv_pyobj(self): def apply_flush(state, named): - from twinkle.checkpoint_engine.delta.decode import decode_delta_wire + 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: @@ -236,7 +238,7 @@ def receiver_main(cfg, ep_q, res_q): _install_stand_ins(lambda r: LoopComm(), build_log) import zmq - from twinkle.checkpoint_engine.delta.hccl_receiver import (build_delta_hccl_communicator, + 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)) diff --git a/tests/checkpoint_engine/test_delta_checksum.py b/tests/checkpoint_engine/test_delta_checksum.py index de3e89786..8f8d42477 100644 --- a/tests/checkpoint_engine/test_delta_checksum.py +++ b/tests/checkpoint_engine/test_delta_checksum.py @@ -12,8 +12,8 @@ import pytest import torch -from twinkle.checkpoint_engine.delta import wire -from twinkle.checkpoint_engine.delta.wire import checksum +from twinkle.checkpoint_engine.delta.protocol import wire +from twinkle.checkpoint_engine.delta.protocol.wire import checksum MASK = (1 << 64) - 1 @@ -124,7 +124,7 @@ def __init__(self, *args, **kwargs): module = ModuleType(name) module.__dict__.update(fields) monkeypatch.setitem(sys.modules, name, module) - path = Path(__file__).parents[2] / 'src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py' + 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) diff --git a/tests/checkpoint_engine/test_delta_hccl.py b/tests/checkpoint_engine/test_delta_hccl.py index deb6c6d52..b339159eb 100644 --- a/tests/checkpoint_engine/test_delta_hccl.py +++ b/tests/checkpoint_engine/test_delta_hccl.py @@ -11,8 +11,8 @@ import pytest import torch -from twinkle.checkpoint_engine.delta import hccl_engine as mod -from twinkle.checkpoint_engine.delta.hccl_engine import DeltaHCCLCheckpointEngine +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 @@ -103,7 +103,7 @@ def run(e, eng): def test_import_without_npu_modules(): - code = ("import sys, twinkle.checkpoint_engine.delta.hccl_engine as m;" + 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'] diff --git a/tests/checkpoint_engine/test_delta_hccl_receiver.py b/tests/checkpoint_engine/test_delta_hccl_receiver.py index f641f128a..3a433f2fc 100644 --- a/tests/checkpoint_engine/test_delta_hccl_receiver.py +++ b/tests/checkpoint_engine/test_delta_hccl_receiver.py @@ -11,9 +11,9 @@ import pytest import torch -from twinkle.checkpoint_engine.delta import hccl_receiver -from twinkle.checkpoint_engine.delta.decode import DeltaPatch, decode_delta_wire -from twinkle.checkpoint_engine.delta.wire import checksum +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) @@ -356,11 +356,11 @@ def _update(adapter, rank=1): def test_receive_sparse_session_prepares_sparse_loaders(hccl, monkeypatch): - utils = sys.modules.get('twinkle.sampler.vllm_sampler.weight_update_utils') + 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.weight_update_utils', + 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) @@ -484,7 +484,7 @@ def test_register_hccl_engine_once(adapter, monkeypatch): @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_weight_transfer', adapter) + 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)) @@ -560,7 +560,7 @@ def test_update_twinkle_delta_weights_rejects_short_list(worker_cls): # --------------------------------------------------------------------------- (e) import hygiene def test_import_without_npu_modules(): - code = ("import sys, twinkle.checkpoint_engine.delta.hccl_receiver, twinkle.checkpoint_engine.delta.decode;" + 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'] diff --git a/tests/checkpoint_engine/test_delta_npu_orchestration.py b/tests/checkpoint_engine/test_delta_npu_orchestration.py index d0658622f..966713bfc 100644 --- a/tests/checkpoint_engine/test_delta_npu_orchestration.py +++ b/tests/checkpoint_engine/test_delta_npu_orchestration.py @@ -9,7 +9,7 @@ from twinkle import Platform from twinkle.checkpoint_engine import CheckpointEngineManager from twinkle.checkpoint_engine.delta import DeltaSyncConfig -from twinkle.checkpoint_engine.delta.rollout import DeltaSamplerMixin +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) @@ -380,7 +380,7 @@ def args_cls(model=None, enable_sleep_mode=False, weight_transfer_config=None, * 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_weight_transfer', adapter) + monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.delta.transfer_engine', adapter) return SimpleNamespace(adapter=adapter, captured=captured) diff --git a/tests/checkpoint_engine/test_delta_pipeline.py b/tests/checkpoint_engine/test_delta_pipeline.py index 2dd2d919d..ea1f69c7e 100644 --- a/tests/checkpoint_engine/test_delta_pipeline.py +++ b/tests/checkpoint_engine/test_delta_pipeline.py @@ -7,8 +7,8 @@ import pytest import torch -from twinkle.checkpoint_engine.delta.engine import DeltaShardedCheckpointEngine -from twinkle.checkpoint_engine.delta.pipeline import DeltaSyncMixin +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 = { diff --git a/tests/checkpoint_engine/test_delta_platform.py b/tests/checkpoint_engine/test_delta_platform.py index 6cd0b5e1d..a5d7ba293 100644 --- a/tests/checkpoint_engine/test_delta_platform.py +++ b/tests/checkpoint_engine/test_delta_platform.py @@ -11,13 +11,13 @@ 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 import trainer as delta_trainer -from twinkle.checkpoint_engine.delta.backends.fsdp import validate_fsdp2 -from twinkle.checkpoint_engine.delta.engine import DeltaShardedCheckpointEngine -from twinkle.checkpoint_engine.delta.trainer import DeltaTrainerMixin +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.hccl_engine' +HCCL_MODULE = 'twinkle.checkpoint_engine.delta.transport.hccl_engine' NOT_IMPLEMENTED = 'not implemented yet' @@ -34,7 +34,7 @@ def test_resolver_gpu(): def test_resolver_npu_real_engine(): - from twinkle.checkpoint_engine.delta.hccl_engine import DeltaHCCLCheckpointEngine + from twinkle.checkpoint_engine.delta.transport.hccl_engine import DeltaHCCLCheckpointEngine assert resolve_delta_engine_cls('NPU') is DeltaHCCLCheckpointEngine @@ -90,7 +90,7 @@ def test_manager_full_sync_does_not_resolve_delta(monkeypatch): def test_manager_npu_constructs_with_hccl_engine(monkeypatch): fake_platform(monkeypatch, 'NPU') - from twinkle.checkpoint_engine.delta.hccl_engine import DeltaHCCLCheckpointEngine + 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 diff --git a/tests/checkpoint_engine/test_delta_rollout.py b/tests/checkpoint_engine/test_delta_rollout.py index 059c0e54c..5b5826121 100644 --- a/tests/checkpoint_engine/test_delta_rollout.py +++ b/tests/checkpoint_engine/test_delta_rollout.py @@ -8,7 +8,8 @@ import pytest import torch -from twinkle.checkpoint_engine.delta.rollout import CudaDeltaReceiver, DeltaSamplerMixin +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 @@ -97,7 +98,7 @@ def test_without_kv_release_pause_and_resume_still_run(): def test_checkpoint_receiver_keeps_flush_on_local_gpu_until_ipc(monkeypatch): - from twinkle.checkpoint_engine.delta import rollout as delta_rollout + 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)]) @@ -106,7 +107,7 @@ def test_checkpoint_receiver_keeps_flush_on_local_gpu_until_ipc(monkeypatch): 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.bucketed_weight_transfer') + 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) @@ -122,7 +123,7 @@ def test_checkpoint_receiver_keeps_flush_on_local_gpu_until_ipc(monkeypatch): def test_receiver_binds_gpu_before_cuda_initialization(monkeypatch): - from twinkle.checkpoint_engine.delta import rollout as delta_rollout + from twinkle.checkpoint_engine.delta.rollout import cuda_receiver as delta_rollout import os monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '') @@ -136,7 +137,7 @@ def test_receiver_binds_gpu_before_cuda_initialization(monkeypatch): def test_receiver_rejects_gpu_mismatch_before_joining_group(monkeypatch): - from twinkle.checkpoint_engine.delta import rollout as delta_rollout + 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()) diff --git a/tests/checkpoint_engine/test_delta_trainer.py b/tests/checkpoint_engine/test_delta_trainer.py index 33d0a3c3b..097d96d64 100644 --- a/tests/checkpoint_engine/test_delta_trainer.py +++ b/tests/checkpoint_engine/test_delta_trainer.py @@ -16,11 +16,11 @@ from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor from twinkle import DeviceMesh, Platform -from twinkle.checkpoint_engine.delta.backends import FSDP2DeltaBackend -from twinkle.checkpoint_engine.delta.backends.fsdp import validate_fsdp2 -from twinkle.checkpoint_engine.delta.engine import DeltaShardedCheckpointEngine, _GatherQueue -from twinkle.checkpoint_engine.delta.export import shard_delta_indices -from twinkle.checkpoint_engine.delta.trainer import DeltaTrainerMixin +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 @@ -56,7 +56,7 @@ def test_shard_delta_indices_no_change_is_empty(): 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.spec import ShardSpec, derive_dtensor_placement + 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) @@ -72,7 +72,7 @@ def test_gather_slot_entries_sub_rounds_world1(): import torch.distributed as dist - from twinkle.checkpoint_engine.delta.gather import gather_slot_entries_to_rank0 + from twinkle.checkpoint_engine.delta.transport.gather import gather_slot_entries_to_rank0 owns_pg = not dist.is_initialized() if owns_pg: @@ -107,8 +107,8 @@ def test_prime_then_hf_delta_export_roundtrip(): 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.spec import ShardSpec - from twinkle.checkpoint_engine.delta.export import _hf_entry_identity, hf_delta_export, prime_delta_snapshots + 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,)) @@ -140,8 +140,8 @@ def test_hf_delta_export_requires_seed(): garbage.""" import pytest - from twinkle.checkpoint_engine.delta.spec import ShardSpec - from twinkle.checkpoint_engine.delta.export import _hf_entry_identity, hf_delta_export + 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,)) @@ -183,7 +183,7 @@ def test_single_rank_trainer_accepts_unsharded_model(fsdp_owner): def test_single_rank_trainer_uses_rank_zero_without_process_group(monkeypatch): - from twinkle.checkpoint_engine.delta import trainer as delta_trainer + 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() @@ -275,7 +275,7 @@ def test_full_sync_baseline_uses_same_hf_key_conversion(): def test_snapshot_pinning_can_be_disabled(monkeypatch): - from twinkle.checkpoint_engine.delta.backends import base as backend_base + 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)) @@ -360,7 +360,7 @@ def initialize(config): def test_trainer_validation_failure_releases_metadata_group(monkeypatch): from twinkle.checkpoint_engine.delta import DeltaSyncConfig - from twinkle.checkpoint_engine.delta import trainer as delta_trainer + from twinkle.checkpoint_engine.delta.trainer import mixin as delta_trainer trainer = FSDPDeltaTrainer() group = object() diff --git a/tests/checkpoint_engine/test_delta_vllm.py b/tests/checkpoint_engine/test_delta_vllm.py index 556b1124e..aee3f7315 100644 --- a/tests/checkpoint_engine/test_delta_vllm.py +++ b/tests/checkpoint_engine/test_delta_vllm.py @@ -12,7 +12,7 @@ import pytest import torch -from twinkle.checkpoint_engine.delta.wire import checksum +from twinkle.checkpoint_engine.delta.protocol.wire import checksum @pytest.fixture @@ -57,7 +57,7 @@ def __init__(self, config, vllm_config, device, model): module = ModuleType(name) module.__dict__.update(fields) monkeypatch.setitem(sys.modules, name, module) - path = Path(__file__).parents[2] / 'src/twinkle/sampler/vllm_sampler/delta_weight_transfer.py' + 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) @@ -168,7 +168,7 @@ def write_source(*_, **__): def test_refresh_derived_weights_rejects_unknown_packed_layout(adapter): - from twinkle.sampler.vllm_sampler.weight_update_utils import refresh_derived_weights + 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'): @@ -176,7 +176,7 @@ def test_refresh_derived_weights_rejects_unknown_packed_layout(adapter): def test_refresh_derived_weights_is_noop_without_packed_weights(adapter): - from twinkle.sampler.vllm_sampler.weight_update_utils import refresh_derived_weights + from twinkle.sampler.vllm_sampler.delta.weight_update_utils import refresh_derived_weights assert refresh_derived_weights(torch.nn.Linear(2, 2)) == 0 @@ -194,7 +194,7 @@ def test_partial_update_failure_poisoned_across_sessions(adapter): 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_weight_transfer', adapter) + 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) @@ -202,7 +202,7 @@ def test_worker_requires_native_delta_backend(adapter, monkeypatch): 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_weight_transfer', adapter) + 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, @@ -237,7 +237,7 @@ 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, 'weight_update_utils.py') + return load_source(monkeypatch, 'delta/weight_update_utils.py') @pytest.fixture(params=[False, True], ids=['language-model', 'qwen-multimodal']) @@ -314,8 +314,8 @@ def __init__(self, config, vllm_config, device, model): module = ModuleType(name) module.__dict__.update(fields) monkeypatch.setitem(sys.modules, name, module) - monkeypatch.setitem(sys.modules, 'twinkle.sampler.vllm_sampler.weight_update_utils', tied_helpers) - return load_source(monkeypatch, 'delta_weight_transfer.py') + 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']) @@ -347,7 +347,7 @@ def load_patches(model, patches, **_): @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.weight_update_utils'] + 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) @@ -403,7 +403,7 @@ def loader(param, loaded_weight): @pytest.fixture def mamba_helpers(): - path = Path(__file__).parents[2] / 'src/twinkle/sampler/vllm_sampler/weight_update_utils.py' + 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) From 9badf353fbf4cd972e9ab01d9c8bbba822dec67b Mon Sep 17 00:00:00 2001 From: kevssim Date: Thu, 22 Oct 2026 15:48:33 +0800 Subject: [PATCH 12/12] fix(transformers): drop stray @staticmethod on _lazy_wrap_model A refactor removed _not_encoded but left its @staticmethod behind, which then decorated _lazy_wrap_model(self) and broke every self._lazy_wrap_model() call with a missing 'self' argument. --- src/twinkle/model/transformers/transformers.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 2d0b6fe83..9b20cf0a5 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -699,7 +699,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]