adapt new architecture

Author: kip-cxj

tests/checkpoint_engine/test_correctness_on_gpu.py

@@ -22,6 +22,7 @@
     RayResourcePool,
     split_resource_pool,
 )
+from verl.utils.device import get_device_name
 from verl.workers.config import CheckpointEngineConfig, HFModelConfig, RolloutConfig
 
 
@@ -127,6 +128,54 @@ async def test_nixl_checkpoint_engine(
     ray.shutdown()
 
 
+@pytest.mark.skip(reason="temporary skip since our ci environment is not ready")
+@pytest.mark.asyncio
+@pytest.mark.parametrize("rebuild_group", [False])
+@pytest.mark.parametrize("num_trainer, num_rollout", [(2, 6)])
+async def test_kimi_checkpoint_engine(
+    rebuild_group,
+    num_trainer,
+    num_rollout,
+    num_nodes=1,
+    num_gpus_per_node=8,
+    check_allclose=True,
+    model_path="~/models/Qwen/Qwen3-8B-Base",
+):
+    model_path = os.path.expanduser(model_path)
+    ray.init(
+        runtime_env={
+            "env_vars": {
+                "NCCL_IB_HCA": "mlx5",
+                "VERL_LOGGING_LEVEL": "DEBUG",
+                "ASCEND_USE_SHORT_CONNECTION": "1",
+            }
+        }
+    )
+
+    # initialize config
+    checkpoint_engine_config = CheckpointEngineConfig(
+        backend="kimi_ckpt_engine", engine_kwargs={"kimi_ckpt_engine": {"rebuild_group": rebuild_group}}
+    )
+    model_config = HFModelConfig(path=model_path, use_remove_padding=True)
+    rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config)
+
+    # create trainer and rollout worker group
+    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)
+    resource_pool.get_placement_groups(device_name=get_device_name())
+    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])
+    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)
+    trainer.reset()
+    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)
+
+    # create checkpoint engine manager
+    checkpoint_manager = CheckpointEngineManager(backend="kimi_ckpt_engine", trainer=trainer, replicas=replicas)
+    for _ in range(3):
+        await checkpoint_manager.update_weights()
+        rollout.check_weights()
+
+    ray.shutdown()
+
+
 if __name__ == "__main__":
     test_nccl_checkpoint_engine(
         rebuild_group=False,

tests/checkpoint_engine/test_correctness_on_npu.py

@@ -74,6 +74,56 @@ async def test_hccl_checkpoint_engine(
     ray.shutdown()
 
 
+@pytest.mark.skip(reason="temporary skip since our ci environment is not ready")
+@pytest.mark.asyncio
+@pytest.mark.parametrize("rebuild_group", [False])
+@pytest.mark.parametrize("num_trainer, num_rollout", [(2, 6)])
+async def test_kimi_checkpoint_engine(
+    rebuild_group,
+    num_trainer,
+    num_rollout,
+    num_nodes=1,
+    num_gpus_per_node=8,
+    check_allclose=True,
+    model_path="~/models/Qwen/Qwen3-8B-Base",
+):
+    model_path = os.path.expanduser(model_path)
+    ray.init(
+        runtime_env={
+            "env_vars": {
+                "HCCL_CONNECT_TIMEOUT": "1500",
+                "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050",
+                "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050",
+                "VERL_LOGGING_LEVEL": "DEBUG",
+                "ASCEND_USE_SHORT_CONNECTION": "1",
+            }
+        }
+    )
+
+    # initialize config
+    checkpoint_engine_config = CheckpointEngineConfig(
+        backend="kimi_ckpt_engine", engine_kwargs={"kimi_ckpt_engine": {"rebuild_group": rebuild_group}}
+    )
+    model_config = HFModelConfig(path=model_path, use_remove_padding=True)
+    rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config)
+
+    # create trainer and rollout worker group
+    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)
+    resource_pool.get_placement_groups(device_name=get_device_name())
+    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])
+    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)
+    trainer.reset()
+    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)
+
+    # create checkpoint engine manager
+    checkpoint_manager = CheckpointEngineManager(backend="kimi_ckpt_engine", trainer=trainer, replicas=replicas)
+    for _ in range(3):
+        await checkpoint_manager.update_weights()
+        rollout.check_weights()
+
+    ray.shutdown()
+
+
 if __name__ == "__main__":
     test_hccl_checkpoint_engine(
         rebuild_group=False,

tests/checkpoint_engine/test_kimi_checkpoint_engine.py

@@ -44,7 +44,6 @@
 except ImportError:
     HCCLCheckpointEngine = None
 
-
 try:
     from .nixl_checkpoint_engine import NIXLCheckpointEngine
 

verl/checkpoint_engine/__init__.py

@@ -174,8 +174,10 @@ async def receive_tensor(
 
 @dataclass
 class MasterMetadata:
-    ip: str
-    port: int
+    zmq_ip: str
+    zmq_port: int
+    dist_ip: str
+    dist_port: int
 
 
 class BroadcastOperation:
@@ -231,17 +233,11 @@ class KIMICheckpointEngine(CheckpointEngine):
 
     def __init__(
         self,
-        train_world_size: int,
-        rollout_world_size: int,
         bucket_size: int,
         rebuild_group: bool = False,
         is_master: bool = False,
         rollout_dtype: torch.dtype = torch.bfloat16,
     ) -> None:
-        self.train_world_size = train_world_size
-        self.rollout_world_size = rollout_world_size
-        self.world_size = train_world_size + rollout_world_size
-
         self.bucket_size = bucket_size
         self.rebuild_group = rebuild_group
         self.rollout_dtype = rollout_dtype
@@ -254,39 +250,65 @@ def prepare(self) -> MasterMetadata:
             self.ip = ray.util.get_node_ip_address().strip("[]")
             self.listen_port, _ = get_free_port(self.ip)
 
-        return MasterMetadata(ip=self.ip, port=self.listen_port) if self.is_master else None
+        return (
+            MasterMetadata(zmq_ip=None, zmq_port=None, dist_ip=self.ip, dist_port=self.listen_port)
+            if self.is_master
+            else None
+        )
 
-    def finish(self):
+    def finalize(self):
         """Destroy the ckpt engine process group if rebuild_group is True."""
         if self.rebuild_group:
             dist.destroy_process_group()
             self.rank = None
             self.world_size = None
             self.initialized = False
 
-    def init_process_group(self, rank: int, world_size: int, master_metadata: MasterMetadata):
+    @classmethod
+    def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metadata: list[dict]):
+        trainer_kwargs = {
+            "method": ["init_process_group"] * trainer_world_size,
+            "rank": list(range(0, trainer_world_size)),
+            "trainer_world_size": [trainer_world_size] * trainer_world_size,
+            "rollout_world_size": [rollout_world_size] * rollout_world_size,
+            "master_metadata": [metadata[0]] * trainer_world_size,
+        }
+        rollout_kwargs = {
+            "method": ["init_process_group"] * rollout_world_size,
+            "rank": list(range(trainer_world_size, trainer_world_size + rollout_world_size)),
+            "trainer_world_size": [trainer_world_size] * trainer_world_size,
+            "rollout_world_size": [rollout_world_size] * rollout_world_size,
+            "master_metadata": [metadata[0]] * rollout_world_size,
+        }
+        return trainer_kwargs, rollout_kwargs
+
+    def init_process_group(self, rank: int, trainer_world_size: int, rollout_world_size :int, master_metadata: MasterMetadata):
         """Initialize the ckpt engine process group.
 
         Args:
             rank (int): The rank of the current process.
             world_size (int): The total number of processes.
         """
         self.rank = rank
+        self.trainer_world_size = trainer_world_size
+        self.rollout_world_size = rollout_world_size
+        self.world_size = trainer_world_size + rollout_world_size
         # unregister_memory in transfer engine is not supported on NPU,
         # so we have to initialize ParameterServer each time
         if get_device_name() == "npu" or not self.initialized:
-            self.parameter_server = ParameterServer(rank=rank, world_size=world_size, auto_pg=False, custom_dist=True)
-            self.parameter_server.receive_tensor = types.MethodType(receive_tensor, self.parameter_server)
-        if not self.initialized:
-            dist.init_process_group(
-                host=master_metadata.ip,
-                port=master_metadata.port,
+            self.parameter_server = ParameterServer(
                 rank=rank,
-                world_size=world_size,
-                backend=get_nccl_backend(),
+                world_size=self.world_size,
+                auto_pg=False,
+                master_addr=master_metadata.dist_ip,
+                master_port=master_metadata.dist_port,
             )
+            self.parameter_server.receive_tensor = types.MethodType(receive_tensor, self.parameter_server)
+        if not self.initialized:
+            dist.use_backend(f"vllm_{get_nccl_backend()}")
+            self.parameter_server.init_process_group()
 
-            self.rollout_ranks = list(range(self.train_world_size, world_size))
+            self.rollout_ranks = list(range(self.trainer_world_size, self.world_size))
             self.rollout_group = dist.new_group(self.rollout_ranks)
             self.initialized = True
 
@@ -304,7 +326,7 @@ def offload_cpu(named_tensors: dict[str, torch.Tensor], name: str, tensor: torch
         start_time = time.time()
         named_tensors = {}
         for named_tensors_gpu in ckpt_get_named_tensor_buckets(
-            weights, self.bucket_size, self.train_world_size, self.rank, self.rollout_dtype
+            weights, self.bucket_size, self.trainer_world_size, self.rank, self.rollout_dtype
         ):
             with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
                 futures = [