add checkpoint engine for one step off policy

Author: kip-cxj

recipe/one_step_off_policy/ckpt_engine_worker.py

@@ -0,0 +1,95 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan Ltd. and/or its affiliates
+# Copyright 2025 Huawei Ltd. and/or its affiliates
+#
+# 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.
+
+import logging
+import os
+import time
+
+import httpx
+import torch
+import torch.distributed
+from checkpoint_engine.ps import ParameterServer, request_inference_to_update
+from omegaconf import DictConfig, OmegaConf
+
+from verl.single_controller.base import Worker
+from verl.single_controller.base.decorator import Dispatch, register
+from verl.utils.device import (
+    get_device_name,
+)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+device_name = get_device_name()
+
+
+class CkptEngineWorker(Worker):
+    def __init__(self, rank_offset, ps_world_size, inference_parallel_size):
+        super().__init__()
+        rank = self.rank + rank_offset
+        self.ps_rank = rank
+        self.ps_rank_offset = rank_offset
+        self.ps_world_size = ps_world_size
+        self.inference_parallel_size = inference_parallel_size
+        self.ps = ParameterServer(rank=rank, world_size=ps_world_size)
+        self.index = 0
+
+    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+    def init_process_group(self):
+        os.environ["HCCL_NPU_SOCKET_PORT_RANGE"] = "61020-61050"
+        self.ps.init_process_group(device_index=0, master_port=60010)
+        del os.environ["HCCL_NPU_SOCKET_PORT_RANGE"]
+
+    def check_vllm_ready(self, uds: str | None = None):
+        if self.ps_rank != self.ps_rank // self.inference_parallel_size * self.inference_parallel_size:
+            return
+        retry_num = 0
+        transport = None
+        if uds is not None:
+            transport = httpx.HTTPTransport(uds=uds)
+        while True:
+            try:
+                response = httpx.Client(transport=transport).get(f"{self.endpoint}/health", timeout=10)
+                response.raise_for_status()
+                break
+            except (httpx.ConnectError, httpx.HTTPStatusError) as e:
+                retry_num += 1
+                logger.warning(f"fail to check vllm ready, retry {retry_num} times, error: {e}")
+                time.sleep(5)
+
+    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+    def set_server_addresses(self, server_addresses: list[str]):
+        # todo support multiple api server
+        self.endpoint = f"http://{server_addresses[0]}"
+        self.check_vllm_ready()
+
+    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+    def sync_rollout_weights_by_ckpt_engine(self):
+        rank = self.rank
+        src = rank // self.inference_parallel_size * self.inference_parallel_size
+
+        def req_func(socket_paths: list[tuple[str, str]]) -> None:
+            if rank == src:
+                request_inference_to_update(
+                    url=f"{self.endpoint}/collective_rpc",
+                    socket_paths=dict(socket_paths),
+                )
+
+        checkpoint_name = f"sync_{self.index}"
+        self.ps.register_checkpoint(checkpoint_name=checkpoint_name)
+        self.ps.gather_metas(checkpoint_name)
+        ranks = list(range(self.ps_rank_offset, self.ps_world_size))
+        self.ps.update(checkpoint_name, req_func, ranks=ranks)

recipe/one_step_off_policy/config/one_step_off_ppo_trainer.yaml

@@ -20,9 +20,12 @@ actor_rollout_ref:
     free_cache_engine: False
     # Must be enabled! Otherwise, log_probs cannot be calculated.
     calculate_log_probs: True
+    engine_kwargs:
+      vllm:
+        worker_extension_cls: checkpoint_engine.worker.VllmColocateWorkerExtension
 
 # Only then will the use of log probs be correct.
 # And it can be used in conjunction with other rollout_correction algorithms.
 algorithm:
   rollout_correction:
-    bypass_mode: True
+    bypass_mode: True

recipe/one_step_off_policy/fsdp_workers.py

@@ -18,11 +18,11 @@
 
 import torch
 import torch.distributed
+from checkpoint_engine.ps import ParameterServer
 from omegaconf import DictConfig
 from ray.util.collective import collective
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
-from recipe.one_step_off_policy.distributed_util import vllm_stateless_init_process_group
 from verl.single_controller.base.decorator import Dispatch, register
 from verl.utils.device import (
     get_device_name,
@@ -53,17 +53,6 @@ class DetachSync(AsyncActorRolloutRefWorker):
     def _get_actor_params(self):
         pass
 
-    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
-    def create_weight_sync_group(self, master_address, master_port, rank_offset, world_size):
-        rank = torch.distributed.get_rank() + rank_offset
-        self._weight_sync_group = vllm_stateless_init_process_group(
-            master_address,
-            master_port,
-            rank,
-            world_size,
-            get_torch_device().current_device(),
-        )
-
     @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
     def sync_rollout_weights(self):
         assert (self._is_actor or self._is_rollout) and not self.config.hybrid_engine
@@ -127,6 +116,61 @@ async def update_weights(self, inference_engine, params):
 
 
 class DetachActorWorker(DetachSync):
+    def __init__(self, config: DictConfig, role: str, **kwargs):
+        ActorRolloutRefWorker.__init__(self, config, role)
+
+        if role == "actor":
+            self.ps_rank_offset = kwargs.get("rank_offset", self.rank)
+            self.ps_world_size = kwargs.get("ps_world_size", self.world_size)
+            self.ps = ParameterServer(rank=self.rank, world_size=self.ps_world_size)
+
+            self.index = 0
+
+    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+    def init_process_group(self):
+        os.environ["HCCL_NPU_SOCKET_PORT_RANGE"] = "61020"
+        self.ps.init_process_group(device_index=0, master_port=60010)
+        del os.environ["HCCL_NPU_SOCKET_PORT_RANGE"]
+
+    def split_tensors(self) -> dict[str, torch.Tensor]:
+        assert self._is_actor and not self.config.hybrid_engine
+        assert hasattr(self, "_weights_info") and self._weights_info is not None
+
+        if self._is_actor and self._is_offload_param:
+            load_fsdp_model_to_gpu(self.actor_module_fsdp)
+        params = self._get_actor_params()
+
+        named_tensors = {}
+
+        world_size = self.world_size
+        rank = self.rank
+
+        weights_per_rank = (len(self._weights_info) + world_size - 1) // world_size
+        for index, (key, _, _) in enumerate(self._weights_info):
+            assert key in params
+            tensor = params[key].full_tensor()
+            if index >= rank * weights_per_rank and index < (rank + 1) * weights_per_rank:
+                named_tensors[key] = tensor.to("cpu", non_blocking=True)
+
+        get_torch_device().synchronize()
+
+        return named_tensors
+
+    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+    def sync_rollout_weights_by_ckpt_engine(self):
+        def req_func(socket_paths: list[tuple[str, str]]):
+            return
+
+        named_tensors = self.split_tensors()
+
+        checkpoint_name = f"sync_{self.index}"
+
+        self.ps.register_checkpoint(checkpoint_name=checkpoint_name, named_tensors=named_tensors)
+        self.ps.gather_metas(checkpoint_name)
+        ranks = list(range(self.ps_rank_offset, self.ps_world_size))
+
+        self.ps.update(checkpoint_name, req_func, ranks=ranks)
+
     def _get_actor_params(self):
         assert self._is_actor
         params = self.actor_module_fsdp.state_dict()
@@ -159,8 +203,7 @@ def get_actor_weights_info(self):
 
 
 class DetachAsyncRolloutWorker(DetachSync):
-    def __init__(self, config: DictConfig, role: str):
-        print(f"[DetachAsyncRolloutWorker] {DetachAsyncRolloutWorker.__mro__}")
+    def __init__(self, config: DictConfig, role: str, **kwargs):
         ActorRolloutRefWorker.__init__(self, config, role)
 
     @register(dispatch_mode=Dispatch.ONE_TO_ALL)

recipe/one_step_off_policy/main_ppo.py

@@ -32,6 +32,8 @@
 from verl.utils.config import validate_config
 from verl.utils.device import auto_set_ascend_device_name
 
+from .ckpt_engine_worker import CkptEngineWorker
+
 
 def create_resource_pool_manager(config, roles: list) -> ResourcePoolManager:
     """
@@ -69,6 +71,14 @@ def create_resource_pool_manager(config, roles: list) -> ResourcePoolManager:
         resource_pool_spec["rollout_pool"] = rollout_pool
         mapping[Role.Rollout] = "rollout_pool"
 
+    if Role.CkptEngine in roles:
+        assert config.rollout.n_gpus_per_node > 0, "ckpt_engine config.rollout.n_gpus_per_node must be greater than 0"
+        assert config.rollout.nnodes > 0, "ckpt_engine config.rollout.nnodes must be greater than 0"
+        # the same as rollout pool
+        ckpt_engine_pool = [config.rollout.n_gpus_per_node] * config.rollout.nnodes
+        resource_pool_spec["ckpt_engine_pool"] = ckpt_engine_pool
+        mapping[Role.CkptEngine] = "ckpt_engine_pool"
+
     return ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
 
 
@@ -111,6 +121,7 @@ def create_role_worker_mapping(config):
         Role.Actor: ray.remote(DetachActorWorker),
         Role.Rollout: ray.remote(DetachAsyncRolloutWorker),
         Role.Critic: ray.remote(CriticWorker),
+        Role.CkptEngine: ray.remote(CkptEngineWorker),
     }
 
     if config.reward_model.enable:
@@ -140,6 +151,9 @@ def run(self, config):
 
         from verl.utils.fs import copy_to_local
 
+        if os.environ.get("ASCEND_RT_VISIBLE_DEVICES", None) is not None:
+            del os.environ["ASCEND_RT_VISIBLE_DEVICES"]
+
         print(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
 
         pprint(OmegaConf.to_container(config, resolve=True))

recipe/one_step_off_policy/ray_trainer.py

@@ -127,6 +127,9 @@ def __init__(
         if config.algorithm.use_kl_in_reward:
             self.kl_ctrl_in_reward = core_algos.get_kl_controller(config.algorithm.kl_ctrl)
 
+        self.rank_offset = config.trainer.n_gpus_per_node * config.trainer.nnodes
+        self.ps_world_size = self.rank_offset + config.rollout.n_gpus_per_node * config.rollout.nnodes
+
         self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)
 
     def _validate(self):
@@ -149,7 +152,8 @@ def init_workers(self):
         self._init_async_rollout_manager()
 
     def _init_resource_pools(self):
-        self.resource_pool_manager.create_resource_pool()
+        additional = {"ckpt_engine_pool": {"CPU": 1, "NPU": 0.2}, "rollout_pool": {"CPU": 1, "NPU": 0.8}}
+        self.resource_pool_manager.create_resource_pool(additional=additional)
 
         self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
 
@@ -158,6 +162,19 @@ def _create_worker_classes(self):
         self._create_critic_class()
         self._create_reference_policy_class()
         self._create_reward_model_class()
+        self._create_ckpt_engine_class()
+
+    def _create_ckpt_engine_class(self):
+        # create ckpt engine
+        if True:
+            resource_pool = self.resource_pool_manager.get_resource_pool(Role.CkptEngine)
+            ckpt_engine_cls = RayClassWithInitArgs(
+                cls=self.role_worker_mapping[Role.CkptEngine],
+                rank_offset=self.rank_offset,
+                ps_world_size=self.ps_world_size,
+                inference_parallel_size=self.config.actor_rollout_ref.rollout.tensor_model_parallel_size,
+            )
+            self.resource_pool_to_cls[resource_pool][str(Role.CkptEngine)] = ckpt_engine_cls
 
     def _create_actor_rollout_classes(self):
         for role in [Role.Actor, Role.Rollout]:
@@ -166,6 +183,8 @@ def _create_actor_rollout_classes(self):
                 cls=self.role_worker_mapping[role],
                 config=self.config.actor_rollout_ref,
                 role=str(role),
+                rank_offset=self.rank_offset,
+                ps_world_size=self.ps_world_size,
             )
             self.resource_pool_to_cls[resource_pool][str(role)] = role_cls
 
@@ -249,26 +268,15 @@ def _init_models(self):
         self.rollout_wg = self.all_wg[str(Role.Rollout)]
         self.actor_wg.init_model()
         self.rollout_wg.init_model()
+        self.ckpt_engine_wg = self.all_wg[str(Role.CkptEngine)]
         self.actor_rollout_wg = self.actor_wg
         weights_info = self.actor_wg.get_actor_weights_info()[0]
         self.rollout_wg.set_actor_weights_info(weights_info)
         self._create_weight_sync_group()
 
     def _create_weight_sync_group(self):
-        # TODO: NPU support
-        from verl.utils.device import get_nccl_backend
-
-        actor_rollout_workers = self.actor_wg.workers + self.rollout_wg.workers
-        n_workers = len(actor_rollout_workers)
-
-        # Create Ray collective group for fallback communication
-        collective.create_collective_group(
-            actor_rollout_workers,
-            n_workers,
-            list(range(0, n_workers)),
-            backend=get_nccl_backend(),
-            group_name="actor_rollout",
-        )
+        self.actor_wg.init_process_group()
+        ray.get(self.ckpt_engine_wg.init_process_group())
 
     def _init_async_rollout_manager(self):
         # create async rollout manager and request scheduler
@@ -286,9 +294,11 @@ def _init_async_rollout_manager(self):
             config=self.config, worker_group=self.rollout_wg, rm_resource_pool=rm_resource_pool
         )
 
+        ray.get(self.ckpt_engine_wg.set_server_addresses(self.async_rollout_manager.server_addresses))
+
     def sync_rollout_weights(self):
-        self.actor_wg.sync_rollout_weights()
-        ray.get(self.rollout_wg.sync_rollout_weights())
+        self.actor_wg.sync_rollout_weights_by_ckpt_engine()
+        ray.get(self.ckpt_engine_wg.sync_rollout_weights_by_ckpt_engine())
 
     def _create_continuous_iterator(self):
         """