fix pre_commit

Author: kip-cxj

verl/trainer/config/_generated_ppo_megatron_trainer.yaml

@@ -272,6 +272,7 @@ actor_rollout_ref:
     skip_dump_dir: /tmp/rollout_dump
     skip_tokenizer_init: true
     enable_rollout_routing_replay: false
+    enable_checkpoint_engine: false
     profiler:
       _target_: verl.utils.profiler.ProfilerConfig
       tool: ${oc.select:global_profiler.tool,null}

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -261,6 +261,7 @@ actor_rollout_ref:
     skip_dump_dir: /tmp/rollout_dump
     skip_tokenizer_init: true
     enable_rollout_routing_replay: false
+    enable_checkpoint_engine: false
     profiler:
       _target_: verl.utils.profiler.ProfilerConfig
       tool: ${oc.select:global_profiler.tool,null}

verl/workers/fsdp_workers.py

@@ -581,22 +581,21 @@ def _build_model_optimizer(
     def update_weighs_by_checkpoint_engine(
         self,
         weights: Generator[tuple[str, torch.Tensor], None, None],
-        req_func: Callable[[list[tuple[str, str]]], None]
+        req_func: Callable[[list[tuple[str, str]]], None],
     ):
         named_tensors = {}
         for tensor_idx, (name, tensor) in enumerate(weights):
             if tensor_idx % self.world_size == self.rank:
                 named_tensors[name] = tensor
 
-        checkpoint_name = f"checkpoint_engine"
+        checkpoint_name = "checkpoint_engine"
         self.parameter_server.register_checkpoint(checkpoint_name, named_tensors=named_tensors)
         named_tensors = {}
         dist.barrier()
         self.parameter_server.gather_metas(checkpoint_name)
         self.parameter_server.update(checkpoint_name, req_func)
         self.parameter_server.unregister_checkpoint(checkpoint_name)
 
-
     def _build_rollout(self, trust_remote_code=False):
         from torch.distributed.device_mesh import init_device_mesh
 
@@ -744,16 +743,18 @@ async def rollout_mode(self):
             )
             if self.config.rollout.enable_checkpoint_engine:
                 req_func = await self.rollout.checkpoint_engine_req_func(self.infer_world_size)
-                self.update_weighs_by_checkpoint_engine(per_tensor_param, req_func)
+                self.update_weighs_by_checkpoint_engine(per_tensor_base_params, req_func)
             else:
                 await self.rollout.update_weights(per_tensor_base_params, base_sync_done=False)
             del base_model_params, per_tensor_base_params
-        
+
         if self.config.rollout.enable_checkpoint_engine:
             req_func = await self.rollout.checkpoint_engine_req_func(self.infer_world_size)
             self.update_weighs_by_checkpoint_engine(per_tensor_param, req_func)
         else:
-            await self.rollout.update_weights(per_tensor_param, peft_config=peft_config, base_sync_done=self.base_sync_done)
+            await self.rollout.update_weights(
+                per_tensor_param, peft_config=peft_config, base_sync_done=self.base_sync_done
+            )
         log_gpu_memory_usage("After update_weights", logger=logger)
         del params, per_tensor_param
         aggressive_empty_cache(force_sync=True)

verl/workers/megatron_workers.py

@@ -19,12 +19,12 @@
 import logging
 import os
 import time
+from collections.abc import Callable
 from typing import Any, Generator, Optional
 
 import psutil
 import torch
 import torch.distributed
-from collections.abc import Callable
 from codetiming import Timer
 from omegaconf import DictConfig, OmegaConf
 
@@ -487,14 +487,14 @@ def _build_model_optimizer(
     def update_weighs_by_checkpoint_engine(
         self,
         weights: Generator[tuple[str, torch.Tensor], None, None],
-        req_func: Callable[[list[tuple[str, str]]], None]
+        req_func: Callable[[list[tuple[str, str]]], None],
     ):
         named_tensors = {}
         for tensor_idx, (name, tensor) in enumerate(weights):
             if tensor_idx % self.world_size == self.rank:
                 named_tensors[name] = tensor.to("cpu", non_blocking=True)
 
-        checkpoint_name = f"checkpoint_engine"
+        checkpoint_name = "checkpoint_engine"
         self.parameter_server.register_checkpoint(checkpoint_name, named_tensors=named_tensors)
         named_tensors = {}
         torch.distributed.barrier()

verl/workers/rollout/sglang_rollout/async_sglang_server.py

@@ -77,8 +77,10 @@ def __init__(
         cuda_visible_devices: str,
     ):
         print(f"SGLang http server: {rollout_mode=}, {replica_rank=}, {node_rank=}, {nnodes=}, {cuda_visible_devices=}")
-        os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
-        assert torch.cuda.is_available(), "SGLang http server should run on GPU node"
+        os.environ["CUDA_VISIBLE_DEVICES" if torch.cuda.is_avilable else "ASCEND_RT_VISIBLE_DEVICES"] = (
+            cuda_visible_devices
+        )
+        assert torch.cuda.is_available() or torch.npu.is_available(), "SGLang http server should run on GPU/NPU node"
 
         self.config: RolloutConfig = omega_conf_to_dataclass(config)
         self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)
@@ -337,7 +339,13 @@ async def launch_servers(self):
                     node_id=node_id,
                     soft=False,
                 ),
-                runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}},
+                runtime_env={
+                    "env_vars": {
+                        "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"
+                        if torch.cuda.is_available()
+                        else "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1"
+                    }
+                },
                 name=name,
             ).remote(
                 config=self.config,

verl/workers/rollout/sglang_rollout/sglang_rollout.py

@@ -18,12 +18,13 @@
 import logging
 import multiprocessing as mp
 import os
+from collections.abc import Callable
 from typing import Generator
 
 import ray
 import sglang.srt.entrypoints.engine
 import torch
-from collections.abc import Callable
+from sglang.srt.checkpoint_engine.update import req_inference
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
     assert_pkg_version,
@@ -32,7 +33,6 @@
     set_ulimit,
 )
 from sglang.srt.weight_sync.utils import update_weights as sgl_update_weights
-from sglang.srt.checkpoint_engine.update import req_inference
 from torch.distributed.device_mesh import DeviceMesh
 
 from verl.workers.config import HFModelConfig, RolloutConfig
@@ -203,4 +203,4 @@ async def checkpoint_engine_req_func(self, inference_parallel_size: int) -> Call
 
         req_func = req_inference(endpoint=endpoint, inference_parallel_size=inference_parallel_size)
 
-        return req_func
+        return req_func

verl/workers/rollout/vllm_rollout/vllm_async_server.py

@@ -290,7 +290,9 @@ async def launch_server(self, master_address: str = None, master_port: int = Non
             "override_generation_config": json.dumps(override_generation_config),
             "quantization": quantization,
             "hf_overrides": {"quantization_config": fp8_block_quant_kwargs} if quantization == "fp8" else None,
-            "worker_extension_cls": "checkpoint_engine.worker.VllmColocateWorkerExtension" if self.config.enable_checkpoint_engine else None,
+            "worker_extension_cls": "checkpoint_engine.worker.VllmColocateWorkerExtension" 
+            if self.config.enable_checkpoint_engine
+            else None,
             **engine_kwargs,
         }
 
@@ -691,7 +693,9 @@ async def launch_servers(self):
                     soft=False,
                 ),
                 name=name,
-                runtime_env={"env_vars": {"VLLM_SERVER_DEV_MODE": "1"}} if self.config.enable_checkpoint_engine else None,
+                runtime_env={"env_vars": {"VLLM_SERVER_DEV_MODE": "1"}}
+                if self.config.enable_checkpoint_engine
+                else None,
             ).remote(
                 config=self.config,
                 model_config=self.model_config,

verl/workers/rollout/vllm_rollout/vllm_rollout.py

@@ -30,8 +30,8 @@
 import getpass
 import logging
 import os
-from dataclasses import asdict
 from collections.abc import Callable
+from dataclasses import asdict
 from types import MethodType
 from typing import Any, Generator
 
@@ -271,19 +271,7 @@ async def update_weights(self, weights: Generator[tuple[str, torch.Tensor], None
                 model.load_weights(weights)
 
     async def checkpoint_engine_req_func(self, inference_parallel_size: int) -> Callable[[list[tuple[str, str]]], None]:
-        from checkpoint_engine.ps import request_inference_to_update
-        rank = int(os.getenv("RANK", None))
-        src = rank // inference_parallel_size * inference_parallel_size
-
-        server_actor = ray.get_actor(f"vllm_server_{self.replica_rank}_{self.node_rank}")
-        server_address, server_port = await server_actor.get_server_address.remote()
-        def req_func(socket_paths: list[tuple[str, str]]):
-            if rank == src:
-                request_inference_to_update(
-                    f"http://{server_address}:{server_port}/collective_rpc",
-                    dict(socket_paths[src : src + inference_parallel_size]),
-                )
-        return req_func
+        raise NotImplementedError
 
     def generate_sequences(self, prompts: DataProto) -> DataProto:
         """Batch generate sequences in sync mode."""