feat: add local weight synchronization to CheckpointEngineManager

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.

docs/source_zh/组件/检查点引擎/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，或使用多租户时，

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',

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)

src/twinkle/checkpoint_engine/manager.py

@@ -1,21 +1,28 @@
 # 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
 from .mixin import CheckpointEngineMixin
 
 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')