[rollout] fix: sglang megatron backend missing generate function

Author: wuxibin89

.github/workflows/sgl.yml

@@ -131,4 +131,8 @@ jobs:
         run: |
           cd tests/workers/rollout
           pytest -s test_sglang_async_rollout_mcp_tools.py
+      - name: Test the latest SGLang Rollout async with agent loop
+        run: |
+          BACKEND=fsdp ROLLOUT_NAME=sglang pytest -svvv tests/experimental/agent_loop/test_basic_agent_loop.py
+          BACKEND=megatron ROLLOUT_NAME=sglang pytest -svvv tests/experimental/agent_loop/test_basic_agent_loop.py
       # Note(haibin.lin): for any new test, please update gpu_unit_tests.yaml to avoid repeated tests

.github/workflows/vllm.yml

@@ -121,6 +121,6 @@ jobs:
       - name: Running multi-turn rollout tests on 8 L20 GPUs
         run: |
           pip3 install --upgrade vllm==0.8.3 tensordict==0.7.2
-          pytest -svvv tests/workers/rollout/rollout_vllm/test_vllm_chat_scheduler.py
-          ROLLOUT_NAME=vllm pytest -svvv tests/experimental/agent_loop/test_basic_agent_loop.py
+          BACKEND=fsdp ROLLOUT_NAME=vllm pytest -svvv tests/experimental/agent_loop/test_basic_agent_loop.py
+          BACKEND=megatron ROLLOUT_NAME=vllm pytest -svvv tests/experimental/agent_loop/test_basic_agent_loop.py
       # Note(haibin.lin): for any new test, please update gpu_unit_tests.yaml to avoid repeated tests

tests/experimental/agent_loop/agent_utils.py

@@ -20,14 +20,23 @@
 from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup
 from verl.single_controller.ray.base import create_colocated_worker_cls
 from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role
-from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker
 
 
 def init_agent_loop_manager(config: DictConfig) -> Union[AgentLoopManager, RayWorkerGroup]:
     # =========================== 1. Create hybrid ActorRollout workers ===========================
-    actor_rollout_cls = (
-        AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker
-    )
+    if config.actor_rollout_ref.actor.strategy in ["fsdp", "fsdp2"]:
+        from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker
+
+        actor_rollout_cls = (
+            AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker
+        )
+    elif config.actor_rollout_ref.actor.strategy == "megatron":
+        from verl.workers.megatron_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker
+
+        actor_rollout_cls = (
+            AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker
+        )
+
     role_worker_mapping = {
         Role.ActorRollout: ray.remote(actor_rollout_cls),
     }

tests/experimental/agent_loop/test_basic_agent_loop.py

@@ -27,9 +27,29 @@
 from verl.utils import hf_tokenizer
 
 
+def init_fsdp_config() -> DictConfig:
+    config = OmegaConf.load("verl/trainer/config/ppo_trainer.yaml")
+    # test sleep/wake_up with fsdp offload
+    config.actor_rollout_ref.actor.fsdp_config.param_offload = True
+    config.actor_rollout_ref.actor.fsdp_config.optimizer_offload = True
+    return config
+
+
+def init_megatron_config() -> DictConfig:
+    config = OmegaConf.load("verl/trainer/config/ppo_megatron_trainer.yaml")
+    config.actor_rollout_ref.actor.megatron.tensor_model_parallel_size = 2
+    config.actor_rollout_ref.actor.megatron.pipeline_model_parallel_size = 2
+
+    # FIXME: sglang with megatron param_offload got error:
+    # "CUDA error: an illegal memory access was encountered"
+    config.actor_rollout_ref.actor.megatron.param_offload = False
+    config.actor_rollout_ref.actor.megatron.optimizer_offload = True
+    return config
+
+
 @pytest.fixture
 def init_config() -> DictConfig:
-    config = OmegaConf.load("verl/trainer/config/ppo_trainer.yaml")
+    config = init_fsdp_config() if os.getenv("BACKEND", "fsdp") == "fsdp" else init_megatron_config()
     model_path = "Qwen/Qwen2.5-1.5B-Instruct"
     config.actor_rollout_ref.model.path = model_path
     config.actor_rollout_ref.rollout.name = os.getenv("ROLLOUT_NAME", "vllm")
@@ -38,10 +58,7 @@ def init_config() -> DictConfig:
     config.actor_rollout_ref.rollout.response_length = 4096
     config.actor_rollout_ref.rollout.n = 4
     config.actor_rollout_ref.rollout.agent.num_workers = 2
-
-    # test sleep/wake_up with fsdp offload
-    config.actor_rollout_ref.actor.fsdp_config.param_offload = True
-    config.actor_rollout_ref.actor.fsdp_config.optimizer_offload = True
+    config.actor_rollout_ref.actor.optim.total_training_steps = 100
 
     return config
 

verl/workers/megatron_workers.py

@@ -19,7 +19,7 @@
 import logging
 import os
 import time
-from typing import Union
+from typing import Any, Dict, List, Union
 
 import psutil
 import torch
@@ -700,6 +700,11 @@ async def chat_completion(self, json_request):
         ret = await self.rollout.chat_completion(json_request)
         return ret
 
+    @register(dispatch_mode=Dispatch.DIRECT_ROLLOUT_METHOD, blocking=False)
+    async def generate(self, prompt_ids: List[int], sampling_params: Dict[str, Any], request_id: str) -> List[int]:
+        ret = await self.rollout.generate(prompt_ids, sampling_params, request_id)
+        return ret
+
     @register(dispatch_mode=Dispatch.DIRECT_ROLLOUT_METHOD)
     async def wake_up(self):
         if self.config.rollout.free_cache_engine: