support megatron and sglang

Author: x1314aq

recipe/one_step_off_policy/ckpt_engine_worker.py

@@ -37,19 +37,19 @@
 
 
 class CkptEngineWorker(Worker):
-    def __init__(self, rank_offset, ps_world_size, inference_parallel_size):
+    def __init__(self, rank_offset, ps_world_size, inference_parallel_size, rollout_name):
         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.rollout_name = rollout_name
         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"
+    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"]
 
@@ -70,6 +70,27 @@ def check_vllm_ready(self, uds: str | None = None):
                 logger.warning(f"fail to check vllm ready, retry {retry_num} times, error: {e}")
                 time.sleep(5)
 
+    def check_sglang_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)
+        with httpx.Client(transport=transport) as client:
+            while True:
+                try:
+                    response = client.get(f"{self.endpoint}/ping", timeout=10)
+                    response.raise_for_status()
+                    break
+                except (httpx.ConnectError, httpx.HTTPStatusError) as e:
+                    if retry_num % 10 == 0:
+                        logger.warning(
+                            f"fail to check sglang ready, retry {retry_num} times, error: {e}"
+                        )
+                    retry_num += 1
+                    time.sleep(0.1)
+
     @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
     def set_server_addresses(self, server_addresses: list[str]):
         # todo support multiple api server
@@ -81,15 +102,36 @@ 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:
+        def vllm_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),
                 )
 
+        def vllm_req_func(socket_paths: list[tuple[str, str]]) -> None:
+            if rank == src:
+                with httpx.Client(transport=httpx.HTTPTransport()) as client:
+                    resp = client.post(
+                        f"{self.endpoint}/update_weights_from_ipc",
+                        json={
+                            "zmq_handles": dict(socket_paths),
+                            "flush_cache": True,
+                            "weight_version": None,
+                        },
+                        timeout=300.0,
+                    )
+                    resp.raise_for_status()
+            pass
+
+        if self.rollout_name == "sglang":
+            req_func = sglang_req_func
+        elif self.rollout_name == "vllm":
+            req_func = vllm_req_func
+
+        self._init_process_group()
         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)
+        self.ps.update(checkpoint_name, req_func, ranks=list(range(self.ps_rank_offset, self.ps_world_size)))
+        self.index += 1

recipe/one_step_off_policy/fsdp_workers.py

@@ -123,10 +123,8 @@ def __init__(self, config: DictConfig, role: str, **kwargs):
             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)
@@ -161,15 +159,15 @@ def sync_rollout_weights_by_ckpt_engine(self):
         def req_func(socket_paths: list[tuple[str, str]]):
             return
 
+        self.init_process_group()
         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=list(range(self.ps_rank_offset, self.ps_world_size)))
 
-        self.ps.update(checkpoint_name, req_func, ranks=ranks)
+        self.index += 1
 
     def _get_actor_params(self):
         assert self._is_actor

recipe/one_step_off_policy/megatron_workers.py

@@ -18,6 +18,7 @@
 
 import torch
 import torch.distributed
+from checkpoint_engine.ps import ParameterServer
 from omegaconf import DictConfig
 from ray.util.collective import collective
 
@@ -120,6 +121,58 @@ 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
+
+    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
+
+        params_generator = self._get_actor_params_generator() if self._is_actor else None
+
+        if self._is_actor and self._is_offload_param:
+            load_megatron_model_to_gpu(self.actor_module)
+
+        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, tensor) in enumerate(params_generator):
+            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
+
+        self.init_process_group()
+        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)
+        self.ps.update(checkpoint_name, req_func, ranks=list(range(self.ps_rank_offset, self.ps_world_size)))
+
+        self.index += 1
+
     @register(dispatch_mode=Dispatch.ONE_TO_ALL)
     def _get_actor_params_generator(self):
         assert self._is_actor
@@ -160,7 +213,7 @@ def get_actor_weights_info(self):
 
 
 class DetachAsyncRolloutWorker(DetachSync):
-    def __init__(self, config: DictConfig, role: str):
+    def __init__(self, config: DictConfig, role: str, **kwargs):
         print(f"[DetachAsyncRolloutWorker] {DetachAsyncRolloutWorker.__mro__}")
         ActorRolloutRefWorker.__init__(self, config, role)
 

recipe/one_step_off_policy/ray_trainer.py

@@ -173,6 +173,7 @@ def _create_ckpt_engine_class(self):
                 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,
+                rollout_name=self.config.actor_rollout_ref.rollout.name
             )
             self.resource_pool_to_cls[resource_pool][str(Role.CkptEngine)] = ckpt_engine_cls
 
@@ -272,11 +273,6 @@ def _init_models(self):
         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):
-        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