From db64e71508651b535883eb5e43371642336071f2 Mon Sep 17 00:00:00 2001 From: weikaiwen <34648228+kevssim@users.noreply.github.com> Date: Fri, 28 Aug 2026 11:22:09 +0800 Subject: [PATCH] wip --- .../Checkpoint Engine/CheckpointEngine.md | 19 +- .../CheckpointEngine.md" | 18 +- src/twinkle/checkpoint_engine/__init__.py | 9 +- src/twinkle/checkpoint_engine/base.py | 6 +- src/twinkle/checkpoint_engine/manager.py | 186 +++++++--- src/twinkle/model/megatron/megatron.py | 27 +- .../model/transformers/transformers.py | 28 +- .../sampler/sglang_sampler/sglang_sampler.py | 22 +- .../sampler/vllm_sampler/vllm_sampler.py | 22 +- tests/checkpoint_engine/__init__.py | 1 + tests/checkpoint_engine/test_manager_naive.py | 340 ++++++++++++++++++ tests/sampler/test_ipc_checkpoint_engine.py | 2 +- 12 files changed, 585 insertions(+), 95 deletions(-) create mode 100644 tests/checkpoint_engine/__init__.py create mode 100644 tests/checkpoint_engine/test_manager_naive.py diff --git a/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md b/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md index 1a7c39bfa..cd4f371e6 100644 --- a/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md +++ b/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md @@ -2,6 +2,15 @@ CheckpointEngine is a component used to synchronize model weights between trainer and inference processes, primarily used in RLHF training to synchronize weights between Actor models and Rollout samplers. +`CheckpointEngineManager` exposes four modes: + +- `auto`: local objects use `naive`; Ray actor handlers use `standalone`. +- `naive`: stream the model's weight generator directly into a local sampler without creating a checkpoint engine. +- `colocate`: synchronize Ray actors sharing GPUs through CUDA IPC. +- `standalone`: synchronize disaggregated Ray actors through NCCL on GPU or HCCL on NPU. + +`auto` never infers `colocate`, because actor placement cannot be determined reliably from the driver. + ## Basic Interface ```python @@ -39,7 +48,7 @@ class CheckpointEngine(ABC): ## Available Checkpoint Engines -Twinkle provides two checkpoint engine implementations: +Twinkle provides three cross-process checkpoint engine implementations. `naive` mode bypasses them. ### NCCLCheckpointEngine @@ -61,10 +70,18 @@ A checkpoint engine that uses HCCL for weight transfer between Ascend NPUs. See: [HCCLCheckpointEngine](HCCLCheckpointEngine.md) +### IPCCheckpointEngine + +A CUDA IPC engine for model and sampler Ray actors placed on the same physical GPUs. NCCL cannot be +used for this topology because it rejects multiple ranks bound to one GPU. Weight buckets are mapped +between the actor processes rather than broadcast between devices. + ## How to Choose - **NCCLCheckpointEngine**: Suitable for GPU environments, provides the highest transfer performance - **HCCLCheckpointEngine**: Suitable for Ascend NPU environments +- **IPCCheckpointEngine**: Required for colocated Ray actors sharing physical GPUs +- **No engine (`naive`)**: Local model and sampler objects in the same process > Checkpoint engine is a key component of RLHF training infrastructure, ensuring that trainers and samplers use consistent model weights. > Currently, synchronization is divided into two cases based on merge_and_sync=True/False. When set to True, the LoRA is merged into the base model and then synchronized. diff --git "a/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" index 338be10db..24e90947b 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" @@ -2,6 +2,15 @@ CheckpointEngine (检查点引擎) 是用于在训练器和推理进程之间同步模型权重的组件,主要用于 RLHF 训练中 Actor 模型和 Rollout 采样器之间的权重同步。 +`CheckpointEngineManager` 提供四种模式： + +- `auto`：本地对象使用 `naive`；Ray actor handler 使用 `standalone`。 +- `naive`：模型的权重生成器直接流式传入本地 sampler，不创建 CheckpointEngine。 +- `colocate`：共享 GPU 的 Ray actors 通过 CUDA IPC 同步。 +- `standalone`：分离部署的 Ray actors 在 GPU 上使用 NCCL，在 NPU 上使用 HCCL。 + +`auto` 不会推断 `colocate`，因为 driver 无法可靠判断 actor 的实际设备放置。 + ## 基本接口 ```python @@ -39,7 +48,7 @@ class CheckpointEngine(ABC): ## 可用的检查点引擎 -Twinkle 提供了两种检查点引擎实现: +Twinkle 提供三种跨进程检查点引擎实现；`naive` 模式会绕过这些引擎。 ### NCCLCheckpointEngine @@ -61,10 +70,17 @@ Twinkle 提供了两种检查点引擎实现: 详见: [HCCLCheckpointEngine](HCCLCheckpointEngine.md) +### IPCCheckpointEngine + +适用于模型和 sampler Ray actors 被放置在同一组物理 GPU 上的 CUDA IPC 引擎。NCCL 会拒绝多个 +rank 绑定同一张 GPU，因此该拓扑必须通过 CUDA IPC 在进程间映射权重 bucket，而不是跨设备广播。 + ## 如何选择 - **NCCLCheckpointEngine**: 适用于 GPU 环境,提供最高的传输性能 - **HCCLCheckpointEngine**: 适用于昇腾 NPU 环境 +- **IPCCheckpointEngine**: 适用于共享物理 GPU 的 colocated Ray actors +- **不创建引擎 (`naive`)**: 适用于同一进程内的本地 model 和 sampler > 检查点引擎是 RLHF 训练基础设施的关键组件,确保训练器和采样器使用一致的模型权重。 > 目前的同步分为merge_and_sync=True/False两种情况，为True时将lora合并仅基模并同步，为False时仅同步lora权重。另外，多租户直接附加lora文件到vLLM上，在merge_and_sync=False，或使用多租户时， diff --git a/src/twinkle/checkpoint_engine/__init__.py b/src/twinkle/checkpoint_engine/__init__.py index 37fba79ba..1b2cbd247 100644 --- a/src/twinkle/checkpoint_engine/__init__.py +++ b/src/twinkle/checkpoint_engine/__init__.py @@ -1,8 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Checkpoint Engine for weight synchronization between trainer and rollout. -Provides NCCL/HCCL-based weight broadcast from training model workers to -inference sampler workers in STANDALONE (disaggregated) deployment mode. +``CheckpointEngineManager`` supports three synchronization modes: direct +generator streaming for local objects (``naive``), CUDA IPC for colocated Ray +actors (``colocate``), and NCCL/HCCL for disaggregated Ray actors +(``standalone``). Reference: https://github.com/volcengine/verl/tree/main/verl/checkpoint_engine @@ -16,7 +18,7 @@ from .base import CheckpointEngine, TensorMeta from .hccl_checkpoint_engine import HCCLCheckpointEngine from .ipc_checkpoint_engine import IPCCheckpointEngine -from .manager import CheckpointEngineManager +from .manager import CheckpointEngineManager, CheckpointEngineMode from .mixin import CheckpointEngineMixin # Import backend implementations to register them from .nccl_checkpoint_engine import NCCLCheckpointEngine @@ -25,6 +27,7 @@ 'CheckpointEngine', 'CheckpointEngineMixin', 'CheckpointEngineManager', + 'CheckpointEngineMode', 'NCCLCheckpointEngine', 'HCCLCheckpointEngine', 'IPCCheckpointEngine', diff --git a/src/twinkle/checkpoint_engine/base.py b/src/twinkle/checkpoint_engine/base.py index f3a1d8918..e5f00c7f0 100644 --- a/src/twinkle/checkpoint_engine/base.py +++ b/src/twinkle/checkpoint_engine/base.py @@ -16,10 +16,12 @@ class TensorMeta(TypedDict): class CheckpointEngine(ABC): - """Abstract base class for checkpoint engines. + """Abstract base class for cross-process checkpoint engines. A checkpoint engine handles weight synchronization between trainer and rollout - processes. The typical workflow is: + processes. Local ``naive`` synchronization bypasses this interface and streams + the model's weight generator directly into the sampler. The typical cross-process + workflow is: In trainer process (rank 0): >>> engine = CheckpointEngineRegistry.new('nccl', bucket_size=512<<20) diff --git a/src/twinkle/checkpoint_engine/manager.py b/src/twinkle/checkpoint_engine/manager.py index 355a9b58b..9f7a24666 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -1,6 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. # Adapted from https://github.com/volcengine/verl/blob/main/verl/checkpoint_engine/base.py -from typing import List, Optional +from typing import List, Literal, Optional from twinkle import Platform, get_logger from .base import CheckpointEngine @@ -8,14 +8,21 @@ logger = get_logger() +CheckpointEngineMode = Literal['auto', 'naive', 'colocate', 'standalone'] +_VALID_MODES = {'auto', 'naive', 'colocate', 'standalone'} + class CheckpointEngineManager: - """Weight synchronization manager for Twinkle. + """Weight synchronization manager for local and Ray deployments. + + ``mode`` selects one of three synchronization paths: + + * ``naive`` streams a local model's weight generator directly into a local sampler. + * ``colocate`` connects Ray model and sampler actors sharing GPUs through CUDA IPC. + * ``standalone`` connects disaggregated Ray actors through NCCL/HCCL. - Coordinates weight synchronization between training model and inference sampler, either when they - reside on **different GPUs** (disaggregated / standalone deployment, the default) or when they - **share** one (``colocate=True``). Colocation replaces the NCCL broadcast drawn below with a CUDA - IPC handover per GPU -- not as an optimisation, but because NCCL refuses two ranks on one device. + ``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. Architecture (following verl's CheckpointEngineManager): @@ -25,7 +32,7 @@ class CheckpointEngineManager: │ (Ray actors) │ │ (Ray actors) │ │ │ │ │ │ │ │ ▼ │ │ ▼ │ - │ CheckpointEngine │ NCCL broadcast │ CheckpointEngine │ + │ CheckpointEngine │ NCCL/HCCL/CUDA IPC │ CheckpointEngine │ │ send_weights() │ ─────────────────► │ receive_weights()│ │ │ │ │ │ │ │ │ ▼ │ @@ -42,11 +49,11 @@ class CheckpointEngineManager: >>> manager = CheckpointEngineManager(model=model, sampler=sampler) >>> manager.sync_weights() # Call after each training step - Colocated, the caller also owns the memory schedule, because only it knows where in the loop the - device is free. The sampler must have its weights resident to be written into -- ``sleep(1)`` puts - them on the host -- and the trainer has to step aside before a rollout: + With colocated Ray actors, the caller also owns the memory schedule, because only it knows where + in the loop the device is free. The sampler must have its weights resident to be written into -- + ``sleep(1)`` puts them on the host -- and the trainer has to step aside before a rollout: - >>> manager = CheckpointEngineManager(model=model, sampler=sampler, colocate=True) + >>> manager = CheckpointEngineManager(model=model, sampler=sampler, mode='colocate') >>> sampler.wake_up(tags=['weights']) # able to receive, still without a KV cache >>> manager.sync_weights() >>> model.offload_to_cpu() # the trainer's turn is over @@ -64,20 +71,15 @@ def __init__( model: 'CheckpointEngineMixin', sampler: 'CheckpointEngineMixin', platform: str = 'GPU', - colocate: bool = False, + mode: CheckpointEngineMode = 'auto', ) -> None: self.model = model self.sampler = sampler - self.colocate = colocate - self.backend_cls = self.decide_backend_engine(platform, colocate) + self.requested_mode = mode + self.mode = self._resolve_mode(mode, model, sampler) + self.backend_cls = self.decide_backend_engine(platform, self.mode) - # Validate Ray actors - assert hasattr(model, '_actors') and model._actors, \ - 'CheckpointEngineManager requires model to be deployed as Ray actors' - assert hasattr(sampler, '_actors') and sampler._actors, \ - 'CheckpointEngineManager requires sampler to be deployed as Ray actors' - - if colocate: + if self.mode == 'colocate': # Each side builds its own engine inside its worker, so both have to be told which one. self.model.set_checkpoint_engine_backend('ipc') self.sampler.set_checkpoint_engine_backend('ipc') @@ -91,14 +93,50 @@ def __init__( self._model_keys: Optional[List[str]] = None @staticmethod - def decide_backend_engine(platform: Optional[str] = None, colocate: bool = False) -> 'CheckpointEngine': - if colocate: + def _resolve_mode( + mode: CheckpointEngineMode, + model: 'CheckpointEngineMixin', + sampler: 'CheckpointEngineMixin', + ) -> Literal['naive', 'colocate', 'standalone']: + if mode not in _VALID_MODES: + valid = ', '.join(sorted(_VALID_MODES)) + raise ValueError(f'Unknown checkpoint engine mode {mode!r}; expected one of: {valid}.') + + model_has_actors = bool(getattr(model, '_actors', None)) + sampler_has_actors = bool(getattr(sampler, '_actors', None)) + if model_has_actors != sampler_has_actors: + raise ValueError( + 'CheckpointEngineManager requires model and sampler to use the same deployment shape: ' + 'both must be local objects or both must be Ray actor handlers.') + + if mode == 'auto': + return 'standalone' if model_has_actors else 'naive' + if mode == 'naive' and model_has_actors: + raise ValueError("mode='naive' requires local model and sampler objects without Ray actors.") + if mode in ('colocate', 'standalone') and not model_has_actors: + raise ValueError(f"mode={mode!r} requires model and sampler to be Ray actor handlers.") + return mode + + @staticmethod + def decide_backend_engine( + platform: Optional[str] = None, + mode: Literal['naive', 'colocate', 'standalone'] = 'standalone', + ) -> Optional['CheckpointEngine']: + if mode == 'naive': + return None + + platform_name = Platform.get_platform(platform).__name__ + if mode == 'colocate': + if platform_name != 'GPU': + raise NotImplementedError("mode='colocate' currently requires the GPU platform.") from twinkle.checkpoint_engine import IPCCheckpointEngine return IPCCheckpointEngine - if Platform.get_platform(platform).__name__ == 'GPU': + if mode != 'standalone': + raise ValueError(f'Cannot select a backend for unresolved mode {mode!r}.') + if platform_name == 'GPU': from twinkle.checkpoint_engine import NCCLCheckpointEngine return NCCLCheckpointEngine - elif Platform.get_platform(platform).__name__ == 'NPU': + elif platform_name == 'NPU': from twinkle.checkpoint_engine import HCCLCheckpointEngine return HCCLCheckpointEngine else: @@ -124,8 +162,12 @@ def sync_weights(self, merge_and_sync=True): Returns: None """ - model_metadata = self.model.prepare_checkpoint_engine([True] - + [False] * (self.model.device_mesh.world_size - 1)) + if self.mode == 'naive': + self._sync_weights_naive(merge_and_sync) + return + + is_master = [True] + [False] * (self.model.device_mesh.world_size - 1) + model_metadata = self.model.prepare_checkpoint_engine(is_master) self.sampler.prepare_checkpoint_engine(False) model_kwargs, sampler_kwargs = self.backend_cls.build_topology( self.model.device_mesh.world_size, @@ -146,36 +188,7 @@ def sync_weights(self, merge_and_sync=True): self._peft_config = self.model.get_peft_config_dict() peft_config = self._peft_config - if self._model_keys is None: - if hasattr(self.sampler, 'get_state_keys'): - self._model_keys = self.sampler.get_state_keys() - - if self._model_keys is None: - self._model_keys = [] - - # vLLM may have grouped params - use word boundaries to avoid substring matches - import re - _STACKED_MAPPINGS = [ - (re.compile(r'\bqkv_proj\b'), ('q_proj', 'k_proj', 'v_proj', 'q', 'k', 'v')), - (re.compile(r'\bgate_up_proj\b'), ('gate_proj', 'up_proj')), - (re.compile(r'\bin_proj_ba\b'), ('in_proj_b', 'in_proj_a')), - (re.compile(r'\blanguage_model\.model\b'), ('model.language_model', )), - (re.compile(r'^visual\.'), ('model.visual.', )), - ] - - def _expand_keys(keys): - result = set(keys) - for key in keys: - for pattern, individuals in _STACKED_MAPPINGS: - if pattern.search(key): - for ind in individuals: - result.add(pattern.sub(ind, key)) - return result - - # Two passes for chain expansion (e.g., language_model.model + qkv_proj) - expanded = _expand_keys(self._model_keys) - expanded = _expand_keys(expanded) - self._model_keys = list(expanded) + self._ensure_model_keys() model_result = self.model.send_weights( base_sync_done=self.base_sync_done, merge_and_sync=merge_and_sync, model_keys=self._model_keys) @@ -190,3 +203,62 @@ def _expand_keys(keys): self.base_sync_done = True if not merge_and_sync: logger.info('Base model sync completed, subsequent syncs will be LoRA-only') + + def _ensure_model_keys(self): + if self._model_keys is not None: + return + + if hasattr(self.sampler, 'get_state_keys'): + self._model_keys = self.sampler.get_state_keys() + + if self._model_keys is None: + self._model_keys = [] + + # vLLM may have grouped params - use word boundaries to avoid substring matches + import re + _STACKED_MAPPINGS = [ + (re.compile(r'\bqkv_proj\b'), ('q_proj', 'k_proj', 'v_proj', 'q', 'k', 'v')), + (re.compile(r'\bgate_up_proj\b'), ('gate_proj', 'up_proj')), + (re.compile(r'\bin_proj_ba\b'), ('in_proj_b', 'in_proj_a')), + (re.compile(r'\blanguage_model\.model\b'), ('model.language_model', )), + (re.compile(r'^visual\.'), ('model.visual.', )), + ] + + def _expand_keys(keys): + result = set(keys) + for key in keys: + for pattern, individuals in _STACKED_MAPPINGS: + if pattern.search(key): + for ind in individuals: + result.add(pattern.sub(ind, key)) + return result + + # Two passes for chain expansion (e.g., language_model.model + qkv_proj) + expanded = _expand_keys(self._model_keys) + expanded = _expand_keys(expanded) + self._model_keys = list(expanded) + + def _sync_weights_naive(self, merge_and_sync): + """Stream model weights directly into a local sampler.""" + peft_config = None + if self.base_sync_done and not merge_and_sync: + if self._peft_config is None: + self._peft_config = self.model.get_peft_config_dict() + peft_config = self._peft_config + + self._ensure_model_keys() + weights = self.model._get_weight_generator( + base_sync_done=self.base_sync_done, + merge_and_sync=merge_and_sync, + model_keys=self._model_keys, + ) + self.sampler.receive_weights( + weights=weights, + base_sync_done=self.base_sync_done, + peft_config=peft_config, + ) + + if not self.base_sync_done: + self.base_sync_done = True + if not merge_and_sync: + logger.info('Base model sync completed, subsequent syncs will be LoRA-only') diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 37a3ee22e..214e99be0 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -1784,6 +1784,7 @@ def get_train_configs(self, **kwargs): # prepare_checkpoint_engine, init_checkpoint_process_group, and # finalize_checkpoint_engine are inherited from CheckpointEngineMixin. # + # The weight generator is shared by direct and checkpoint-engine sync. # Key difference from TransformersModel: Megatron uses TP/PP, so # get_hf_state_dict() internally performs TP allgather and handles PP # layer distribution. All model ranks MUST execute the weight generator @@ -1791,8 +1792,7 @@ def get_train_configs(self, **kwargs): # model_actor[0] (rank=0 in the checkpoint engine) actually broadcasts # via NCCL; others consume the generator silently (rank=-1). - @remote_function(dispatch='all', lazy_collect=True) - def send_weights( + def _get_weight_generator( self, adapter_name: str = None, base_sync_done: bool = False, @@ -1801,7 +1801,6 @@ def send_weights( ): if adapter_name is None: adapter_name = self._get_default_group() - engine = self._get_or_create_checkpoint_engine() @contextmanager def merge_lora(): @@ -1893,15 +1892,33 @@ def weight_generator(): else: yield from _raw_weights(False) + return weight_generator() + + @remote_function(dispatch='all', lazy_collect=True) + def send_weights( + self, + adapter_name: str = None, + base_sync_done: bool = False, + merge_and_sync: bool = False, + model_keys: List[str] = None, + ): + engine = self._get_or_create_checkpoint_engine() + weight_generator = self._get_weight_generator( + adapter_name=adapter_name, + base_sync_done=base_sync_done, + merge_and_sync=merge_and_sync, + model_keys=model_keys, + ) + is_sender = (engine.rank is not None and engine.rank == 0) if not is_sender: - for _name, _tensor in weight_generator(): + for _name, _tensor in weight_generator: pass return async def _send(): - await engine.send_weights(weight_generator()) + await engine.send_weights(weight_generator) result_container = {'error': None} diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 5b4a2eef1..715ae5649 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1576,10 +1576,9 @@ def get_train_configs(self, **kwargs) -> str: # ========================================================================= # prepare_checkpoint_engine, init_checkpoint_process_group, and # finalize_checkpoint_engine are inherited from CheckpointEngineMixin. - # Only send_weights_via_checkpoint_engine is model-specific. + # The weight generator is shared by direct and checkpoint-engine sync. - @remote_function(dispatch='all', lazy_collect=True) - def send_weights( + def _get_weight_generator( self, adapter_name: str = None, base_sync_done: bool = False, @@ -1589,7 +1588,6 @@ def send_weights( ): if adapter_name is None: adapter_name = self._get_default_group() - engine = self._get_or_create_checkpoint_engine() # Get state dict from unwrapped model model = self.strategy.unwrap_model(self.model) @@ -1663,11 +1661,31 @@ def weight_generator(): yield name, tensor _print_weight_example(names) + return weight_generator() + + @remote_function(dispatch='all', lazy_collect=True) + def send_weights( + self, + adapter_name: str = None, + base_sync_done: bool = False, + merge_and_sync: bool = False, + model_keys: List[str] = None, + **kwargs, + ): + engine = self._get_or_create_checkpoint_engine() + weight_generator = self._get_weight_generator( + adapter_name=adapter_name, + base_sync_done=base_sync_done, + merge_and_sync=merge_and_sync, + model_keys=model_keys, + **kwargs, + ) + # Run async send_weights in a dedicated event loop thread. # We cannot use the Ray worker's event loop because it may already # be occupied, and send_weights uses run_in_executor internally. async def _send(): - await engine.send_weights(weight_generator()) + await engine.send_weights(weight_generator) result_container = {'error': None} diff --git a/src/twinkle/sampler/sglang_sampler/sglang_sampler.py b/src/twinkle/sampler/sglang_sampler/sglang_sampler.py index 7903bd42a..230bda272 100644 --- a/src/twinkle/sampler/sglang_sampler/sglang_sampler.py +++ b/src/twinkle/sampler/sglang_sampler/sglang_sampler.py @@ -395,30 +395,32 @@ def receive_weights( self, base_sync_done: bool = False, peft_config: dict = None, + weights=None, ): """Receive weights from the trainer and stream them into sglang. - Which transport delivers them is the checkpoint engine's business, not this method's: NCCL - broadcast when the trainer is on other GPUs, CUDA IPC when it shares this one, where NCCL - cannot be used at all. - - The checkpoint engine's ``receive_weights()`` async generator is handed straight to - :meth:`SGLangEngine.update_weights`, which consumes it one tensor at a time into a bucket and - forwards each full bucket to sglang. Peak extra memory is one bucket rather than a second copy - of the model, the same reason the vLLM path streams. + With no ``weights`` argument, the checkpoint engine supplies an async generator from its + NCCL/HCCL/CUDA IPC transport. A local naive caller can instead provide the model's synchronous + generator directly. Either iterator is handed straight to :meth:`SGLangEngine.update_weights`, + which consumes it one tensor at a time into a bucket and forwards each full bucket to sglang. + Peak extra memory is one bucket rather than a second copy of the model. Args: base_sync_done: If True, this would be a LoRA-only sync. peft_config: PEFT config dict for LoRA adapter loading. + weights: Optional synchronous/asynchronous weight iterator. If + omitted, weights are received from the checkpoint engine. Raises: NotImplementedError: For a LoRA-only sync; see the class docstring. """ - engine = self._get_or_create_checkpoint_engine() + if weights is None: + engine = self._get_or_create_checkpoint_engine() + weights = engine.receive_weights() async def _receive_and_load(): await self.engine.update_weights( - engine.receive_weights(), # async generator — not materialised + weights, # async/sync generator — not materialised peft_config=peft_config, base_sync_done=base_sync_done, ) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index d470a4b0c..d2c9520a1 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -467,19 +467,17 @@ def receive_weights( self, base_sync_done: bool = False, peft_config: dict = None, + weights=None, ): """Receive weights from the trainer and stream them into vLLM. - Which transport delivers them is the checkpoint engine's business, not this method's: NCCL - broadcast when the trainer is on other GPUs, CUDA IPC when it shares this one, where NCCL - cannot be used at all. Either way what arrives here is the same async generator. - Uses a **streaming pipeline** to avoid accumulating a full model-weight copy on GPU: - 1. ``CheckpointEngine.receive_weights()`` yields tensors from - the engine's buckets (async generator, GPU tensors). - 2. The async generator is passed **directly** to + 1. With no ``weights`` argument, ``CheckpointEngine.receive_weights()`` + yields tensors from NCCL/HCCL/CUDA IPC buckets. A local naive + caller can instead provide the model's synchronous generator. + 2. The weight iterator is passed **directly** to ``VLLMEngine.update_weights()`` which consumes it one tensor at a time, copying each into a GPU IPC bucket and flushing to the vLLM worker subprocess when the bucket is full. @@ -490,18 +488,22 @@ def receive_weights( Args: base_sync_done: If True, this is a LoRA-only sync. peft_config: PEFT config dict for LoRA adapter loading. + weights: Optional synchronous/asynchronous weight iterator. If + omitted, weights are received from the checkpoint engine. Returns: Number of weights loaded (approximate, from engine log). """ - engine = self._get_or_create_checkpoint_engine() + if weights is None: + engine = self._get_or_create_checkpoint_engine() + weights = engine.receive_weights() async def _receive_and_load(): - # Stream the received tensors directly into vLLM via IPC. + # Stream model/checkpoint-engine tensors directly into vLLM via IPC. # VLLMEngine.update_weights accepts an async generator and # handles bucket packing + ZMQ transfer internally. await self.engine.update_weights( - engine.receive_weights(), # async generator — not materialised + weights, # async/sync generator — not materialised peft_config=peft_config, base_sync_done=base_sync_done, ) diff --git a/tests/checkpoint_engine/__init__.py b/tests/checkpoint_engine/__init__.py new file mode 100644 index 000000000..85b3e739d --- /dev/null +++ b/tests/checkpoint_engine/__init__.py @@ -0,0 +1 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. diff --git a/tests/checkpoint_engine/test_manager_naive.py b/tests/checkpoint_engine/test_manager_naive.py new file mode 100644 index 000000000..eb2c041b1 --- /dev/null +++ b/tests/checkpoint_engine/test_manager_naive.py @@ -0,0 +1,340 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""CPU-only tests for CheckpointEngineManager mode selection and direct sync.""" + +import asyncio + +import pytest + +from twinkle.checkpoint_engine.manager import CheckpointEngineManager + + +class _Mesh: + world_size = 1 + data_world_size = 1 + + +class _Model: + + def __init__(self, weights): + self.device_mesh = _Mesh() + self._weights = weights + self.generator_calls = [] + self.peft_config_calls = 0 + self._checkpoint_engine = None + + def _get_weight_generator(self, **kwargs): + self.generator_calls.append(kwargs) + weights = self._weights(kwargs) if callable(self._weights) else self._weights + + def _weights(): + yield from weights + + return _weights() + + def get_peft_config_dict(self): + self.peft_config_calls += 1 + return {'r': 8, 'target_modules': ['q_proj']} + + +class _Sampler: + + def __init__(self, fail=None): + self.device_mesh = _Mesh() + self.calls = [] + self.loaded = [] + self.fail = fail + self._checkpoint_engine = None + + def get_state_keys(self): + return ['q_proj.weight'] + + def receive_weights(self, **kwargs): + self.calls.append(kwargs) + if self.fail is not None: + raise self.fail + self.loaded.append(list(kwargs['weights'])) + + +@pytest.mark.parametrize('requested_mode', ['auto', 'naive']) +def test_local_sync_streams_generator_without_checkpoint_engine(requested_mode): + model = _Model([('q_proj.weight', 'base')]) + sampler = _Sampler() + + manager = CheckpointEngineManager(model, sampler, platform='CPU', mode=requested_mode) + manager.sync_weights(merge_and_sync=False) + + assert manager.requested_mode == requested_mode + assert manager.mode == 'naive' + assert manager.backend_cls is None + assert manager.base_sync_done is True + assert sampler.loaded == [[('q_proj.weight', 'base')]] + assert model._checkpoint_engine is None + assert sampler._checkpoint_engine is None + assert model.generator_calls == [{ + 'base_sync_done': False, + 'merge_and_sync': False, + 'model_keys': ['q_proj.weight'], + }] + + +def test_local_lora_sync_reuses_peft_config_and_sends_only_incremental_weights(): + model = _Model(lambda call: ([('q_proj.lora_A', 'adapter')] + if call['base_sync_done'] else [('q_proj.weight', 'base')])) + sampler = _Sampler() + manager = CheckpointEngineManager(model, sampler, platform='CPU') + + manager.sync_weights(merge_and_sync=False) + manager.sync_weights(merge_and_sync=False) + + assert sampler.loaded == [[('q_proj.weight', 'base')], [('q_proj.lora_A', 'adapter')]] + assert model.peft_config_calls == 1 + assert sampler.calls[0]['peft_config'] is None + assert sampler.calls[1]['peft_config'] == {'r': 8, 'target_modules': ['q_proj']} + assert sampler.calls[0]['base_sync_done'] is False + assert sampler.calls[1]['base_sync_done'] is True + assert model.generator_calls[1]['base_sync_done'] is True + + +def test_local_merge_sync_generates_a_full_weight_set_each_time(): + model = _Model([('q_proj.weight', 'merged')]) + sampler = _Sampler() + manager = CheckpointEngineManager(model, sampler, platform='CPU') + + manager.sync_weights(merge_and_sync=True) + manager.sync_weights(merge_and_sync=True) + + assert sampler.loaded == [[('q_proj.weight', 'merged')], [('q_proj.weight', 'merged')]] + assert all(call['merge_and_sync'] is True for call in model.generator_calls) + assert [call['base_sync_done'] for call in model.generator_calls] == [False, True] + + +def test_local_failure_does_not_mark_base_sync_done(): + model = _Model([('q_proj.weight', 'base')]) + error = RuntimeError('sampler failed') + sampler = _Sampler(fail=error) + manager = CheckpointEngineManager(model, sampler, platform='CPU') + + with pytest.raises(RuntimeError, match='sampler failed') as exc_info: + manager.sync_weights(merge_and_sync=False) + + assert exc_info.value is error + assert manager.base_sync_done is False + + sampler.fail = None + manager.sync_weights(merge_and_sync=False) + assert model.generator_calls[-1]['base_sync_done'] is False + + +def test_local_weight_generator_failure_does_not_mark_base_sync_done(): + error = ValueError('weight generation failed') + + def broken_weights(): + yield 'q_proj.weight', 'base' + raise error + + model = _Model(broken_weights()) + sampler = _Sampler() + manager = CheckpointEngineManager(model, sampler, platform='CPU') + + with pytest.raises(ValueError, match='weight generation failed') as exc_info: + manager.sync_weights(merge_and_sync=False) + + assert exc_info.value is error + assert manager.base_sync_done is False + + +def test_mixed_deployment_shape_fails_at_initialization(): + model = _Model([]) + sampler = _Sampler() + sampler._actors = [object()] + + with pytest.raises(ValueError, match='same deployment shape'): + CheckpointEngineManager(model, sampler, platform='CPU') + + +@pytest.mark.parametrize( + ('mode', 'use_actors', 'match'), + [ + ('naive', True, "mode='naive' requires local"), + ('colocate', False, "mode='colocate' requires"), + ('standalone', False, "mode='standalone' requires"), + ('unknown', False, 'Unknown checkpoint engine mode'), + ], +) +def test_explicit_mode_validates_deployment_shape(mode, use_actors, match): + model = _Model([]) + sampler = _Sampler() + if use_actors: + model._actors = [object()] + sampler._actors = [object()] + + with pytest.raises(ValueError, match=match): + CheckpointEngineManager(model, sampler, platform='CPU', mode=mode) + + +def test_backend_selection_uses_resolved_mode(): + from twinkle.checkpoint_engine import IPCCheckpointEngine, NCCLCheckpointEngine + + assert CheckpointEngineManager.decide_backend_engine('GPU', mode='naive') is None + assert CheckpointEngineManager.decide_backend_engine('GPU', mode='colocate') is IPCCheckpointEngine + assert CheckpointEngineManager.decide_backend_engine('GPU', mode='standalone') is NCCLCheckpointEngine + + +@pytest.mark.parametrize('sampler_backend', ['vllm', 'sglang']) +@pytest.mark.parametrize('provide_weights', [True, False]) +def test_sampler_receive_weights_selects_direct_or_checkpoint_stream(sampler_backend, provide_weights): + if sampler_backend == 'vllm': + from twinkle.sampler.vllm_sampler.vllm_sampler import vLLMSampler as sampler_cls + else: + from twinkle.sampler.sglang_sampler.sglang_sampler import SGLangSampler as sampler_cls + + class _InferenceEngine: + + def __init__(self): + self.loaded = None + self.invalidated = False + + async def update_weights(self, weights, **kwargs): + self.loaded = (list(weights), kwargs) + + def invalidate_synced_lora(self): + self.invalidated = True + + class _CheckpointEngine: + + def receive_weights(self): + return iter([('checkpoint.weight', 'checkpoint')]) + + sampler = object.__new__(sampler_cls) + sampler.engine = _InferenceEngine() + sampler._run_in_loop = asyncio.run + checkpoint_engine = _CheckpointEngine() + checkpoint_engine_calls = [] + + def get_checkpoint_engine(): + checkpoint_engine_calls.append(True) + return checkpoint_engine + + sampler._get_or_create_checkpoint_engine = get_checkpoint_engine + direct_weights = iter([('direct.weight', 'direct')]) if provide_weights else None + + sampler_cls.receive_weights.__wrapped__(sampler, weights=direct_weights) + + expected_weights = [('direct.weight', 'direct')] if provide_weights else [('checkpoint.weight', 'checkpoint')] + assert sampler.engine.loaded == (expected_weights, {'peft_config': None, 'base_sync_done': False}) + assert len(checkpoint_engine_calls) == (0 if provide_weights else 1) + assert sampler.engine.invalidated is (sampler_backend == 'vllm') + + +def test_auto_ray_actor_sync_keeps_standalone_checkpoint_engine_lifecycle(monkeypatch): + events = [] + + class _Backend: + + @classmethod + def build_topology(cls, trainer_world_size, rollout_world_size, metadata): + events.append(('build_topology', trainer_world_size, rollout_world_size)) + return ({'rank': [0], 'world_size': [2], 'master_metadata': [metadata[0]]}, + {'rank': [1], 'world_size': [2], 'master_metadata': [metadata[0]]}) + + class _ActorModel(_Model): + + def __init__(self): + super().__init__([('q_proj.weight', 'base')]) + self._actors = [object()] + + def prepare_checkpoint_engine(self, is_master): + events.append(('model_prepare', is_master)) + return {'zmq_ip': '127.0.0.1', 'zmq_port': 1} + + def init_checkpoint_process_group(self, **kwargs): + events.append(('model_init_submitted', kwargs)) + return lambda: events.append(('model_init_waited', kwargs)) + + def send_weights(self, **kwargs): + events.append(('send_submitted', kwargs)) + return lambda: events.append(('send_waited', kwargs)) + + def finalize_checkpoint_engine(self): + events.append('model_finalize') + + class _ActorSampler(_Sampler): + + def __init__(self): + super().__init__() + self._actors = [object()] + + def prepare_checkpoint_engine(self, is_master): + events.append(('sampler_prepare', is_master)) + + def init_checkpoint_process_group(self, **kwargs): + events.append(('sampler_init_submitted', kwargs)) + return lambda: events.append(('sampler_init_waited', kwargs)) + + def receive_weights(self, **kwargs): + events.append(('receive_submitted', kwargs)) + return lambda: events.append(('receive_waited', kwargs)) + + def finalize_checkpoint_engine(self): + events.append('sampler_finalize') + + model = _ActorModel() + sampler = _ActorSampler() + + def decide_backend(platform=None, mode='standalone'): + assert mode == 'standalone' + return _Backend + + monkeypatch.setattr(CheckpointEngineManager, 'decide_backend_engine', staticmethod(decide_backend)) + + manager = CheckpointEngineManager(model, sampler, platform='GPU') + manager.sync_weights() + + assert manager.requested_mode == 'auto' + assert manager.mode == 'standalone' + assert [event if isinstance(event, str) else event[0] for event in events] == [ + 'model_prepare', + 'sampler_prepare', + 'build_topology', + 'model_init_submitted', + 'sampler_init_submitted', + 'model_init_waited', + 'sampler_init_waited', + 'send_submitted', + 'receive_submitted', + 'send_waited', + 'receive_waited', + 'model_finalize', + 'sampler_finalize', + ] + receive_kwargs = next( + event[1] + for event in events + if isinstance(event, tuple) and event[0] == 'receive_submitted' + ) + assert 'weights' not in receive_kwargs + assert manager.base_sync_done is True + + +def test_colocate_mode_configures_actor_backends(monkeypatch): + configured = [] + backend = object() + + class _ActorRole: + _actors = [object()] + + def set_checkpoint_engine_backend(self, name): + configured.append(name) + + def decide_backend(platform=None, mode='standalone'): + assert platform == 'GPU' + assert mode == 'colocate' + return backend + + monkeypatch.setattr(CheckpointEngineManager, 'decide_backend_engine', staticmethod(decide_backend)) + + manager = CheckpointEngineManager(_ActorRole(), _ActorRole(), platform='GPU', mode='colocate') + + assert manager.mode == 'colocate' + assert manager.backend_cls is backend + assert configured == ['ipc', 'ipc'] diff --git a/tests/sampler/test_ipc_checkpoint_engine.py b/tests/sampler/test_ipc_checkpoint_engine.py index 0e7532b26..70158a8dc 100644 --- a/tests/sampler/test_ipc_checkpoint_engine.py +++ b/tests/sampler/test_ipc_checkpoint_engine.py @@ -234,7 +234,7 @@ def weights(): def test_the_manager_picks_this_engine_only_when_colocating(): from twinkle.checkpoint_engine import CheckpointEngineManager, IPCCheckpointEngine, NCCLCheckpointEngine - assert CheckpointEngineManager.decide_backend_engine('GPU', colocate=True) is IPCCheckpointEngine + assert CheckpointEngineManager.decide_backend_engine('GPU', mode='colocate') is IPCCheckpointEngine assert CheckpointEngineManager.decide_backend_engine('GPU') is NCCLCheckpointEngine