fix pre-commit

Author: kip-cxj

tests/checkpoint_engine/test_correctness_on_gpu.py

@@ -147,7 +147,6 @@ async def test_kimi_checkpoint_engine(
             "env_vars": {
                 "NCCL_IB_HCA": "mlx5",
                 "VERL_LOGGING_LEVEL": "DEBUG",
-                "ASCEND_USE_SHORT_CONNECTION": "1",
             }
         }
     )

tests/checkpoint_engine/test_correctness_on_npu.py

@@ -77,25 +77,22 @@ async def test_hccl_checkpoint_engine(
 @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)])
+@pytest.mark.parametrize("num_trainer, num_rollout", [(4, 28)])
 async def test_kimi_checkpoint_engine(
     rebuild_group,
     num_trainer,
     num_rollout,
-    num_nodes=1,
-    num_gpus_per_node=8,
+    num_nodes=2,
+    num_gpus_per_node=16,
     check_allclose=True,
-    model_path="~/models/Qwen/Qwen3-8B-Base",
+    model_path="~/models/Qwen/Qwen3-32B",
 ):
     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",
             }
         }
     )

verl/checkpoint_engine/kimi_checkpoint_engine.py

@@ -270,19 +270,25 @@ def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metada
             "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,
+            "rollout_world_size": [rollout_world_size] * trainer_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,
+            "trainer_world_size": [trainer_world_size] * rollout_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):
+    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:
@@ -293,9 +299,8 @@ def init_process_group(self, rank: int, trainer_world_size: int, rollout_world_s
         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:
+
+        if not self.initialized:
             self.parameter_server = ParameterServer(
                 rank=rank,
                 world_size=self.world_size,
@@ -304,7 +309,7 @@ def init_process_group(self, rank: int, trainer_world_size: int, rollout_world_s
                 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()
 
@@ -345,7 +350,7 @@ def offload_cpu(named_tensors: dict[str, torch.Tensor], name: str, tensor: torch
 
         self.parameter_server.register_checkpoint(self.checkpoint_name, named_tensors=named_tensors)
         named_tensors = {}
-        torch.cuda.empty_cache()
+        get_torch_device().empty_cache()
         logger.info(f"Rank {self.rank} offload and register, time cost: {time.time() - start_time:.2f}s")
 
         self.parameter_server.gather_metas(self.checkpoint_name)