[rollout] fix: sglang async fail with Multi-stage Awake feature (#2365)

paolo328 · paolo328 · commit cd734a4a71ef · 2025-07-07T13:56:07.000+08:00
### What does this PR do? Fix a regression from verl-project/verl#1911, because the PR did not change the sglang async branch. CI did not catch this error because it only run 1 step, but this error happen in the second test. So I update the testcases to run 2 steps. To reproduce the bug, run test: TOTAL_TRAIN_STEPS=2 ENGINE=sglang ROLLOUT_MODE=async bash tests/special_e2e/ppo_trainer/run_function_reward.sh It fail with: ``` (WorkerDict pid=1257286) Total steps: 2, num_warmup_steps: 0 (WorkerDict pid=1257286) Actor use_remove_padding=True (WorkerDict pid=1257286) Actor use_fused_kernels=False (AsyncSglangServer pid=1260392) FastAPI listen on [192.168.111.48:40451](http://192.168.111.48:40451/) (WorkerDict pid=1257286) terminate called after throwing an instance of 'c10::Error' (WorkerDict pid=1257286) what(): CUDA error: an illegal memory access was encountered (WorkerDict pid=1257286) CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect. (WorkerDict pid=1257286) For debugging consider passing CUDA_LAUNCH_BLOCKING=1 (WorkerDict pid=1257286) Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions. (WorkerDict pid=1257286) (WorkerDict pid=1257286) Exception raised from c10_cuda_check_implementation at /pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first): (WorkerDict pid=1257286) frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x96 (0x7fbf6036c1b6 in /usr/local/lib/python3.10/dist-packages/torch/lib/[libc10.so](http://libc10.so/)) (WorkerDict pid=1257286) frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7fbf60315a76 in /usr/local/lib/python3.10/dist-packages/torch/lib/[libc10.so](http://libc10.so/)) (WorkerDict pid=1257286) frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x118 (0x7fbf6080d918 in ``` ### Checklist Before Starting - [X] Search for similar PRs. Paste at least one query link here: https://github.com/volcengine/verl/issues?q=is%3Aissue%20state%3Aopen%20an%20illegal%20memory%20access%20was%20encountered - [X] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) - `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data` - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` ### Test ``` (TaskRunner pid=1647269) step:2 - global_seqlen/min:13075 - global_seqlen/max:14837 - global_seqlen/minmax_diff:1762 - global_seqlen/balanced_min:14231 - global_seqlen/balanced_max:14232 - global_seqlen/mean:14231.5 - actor/entropy:2.0606913566589355 - critic/vf_loss:8.7157882153 ``` ### API and Usage Example > Demonstrate how the API changes if any, and provide usage example(s) if possible. ```python # Add code snippet or script demonstrating how to use this ``` ### High-Level Design > Demonstrate the high-level design if this PR is complex. ### Specific Changes > List the specific changes. ### Checklist Before Submitting > [!IMPORTANT] > Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. - [X] Read the [Contribute Guide](https://github.com/volcengine/verl?tab=readme-ov-file#contribution-guide). - [ X] Apply [pre-commit checks](https://github.com/volcengine/verl?tab=readme-ov-file#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` - [X] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs). - [X] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... - [X] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
diff --git a/.github/workflows/e2e_ppo_trainer.yml b/.github/workflows/e2e_ppo_trainer.yml
@@ -292,13 +292,13 @@ jobs:
       - name: Running GSM8K E2E training tests on sglang async
         run: |
           ray stop --force
-          ENGINE=sglang ROLLOUT_MODE=async bash tests/special_e2e/ppo_trainer/run_function_reward.sh
+          TOTAL_TRAIN_STEPS=2 ENGINE=sglang ROLLOUT_MODE=async bash tests/special_e2e/ppo_trainer/run_function_reward.sh
       - name: Running GSM8K E2E training tests on vllm async
         run: |
           ray stop --force
           export VLLM_USE_V1=1
           ray start --head
-          ENGINE=vllm ROLLOUT_MODE=async bash tests/special_e2e/ppo_trainer/run_function_reward.sh
+          TOTAL_TRAIN_STEPS=2 ENGINE=vllm ROLLOUT_MODE=async bash tests/special_e2e/ppo_trainer/run_function_reward.sh
 
   e2e_ppo_trainer_sglang_multiturn_with_tool:
     runs-on: [L20x8]
diff --git a/verl/workers/megatron_workers.py b/verl/workers/megatron_workers.py
@@ -19,7 +19,7 @@
 import logging
 import os
 import time
-from typing import Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import psutil
 import torch
@@ -692,6 +692,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:
diff --git a/verl/workers/sharding_manager/fsdp_sglang.py b/verl/workers/sharding_manager/fsdp_sglang.py
@@ -32,7 +32,6 @@
 from verl.protocol import all_gather_data_proto
 from verl.utils.device import get_device_id, get_torch_device
 from verl.utils.fsdp_utils import fsdp_version, load_fsdp_model_to_gpu, offload_fsdp_model_to_cpu
-from verl.utils.model import convert_weight_keys
 from verl.utils.profiler import GPUMemoryLogger, log_gpu_memory_usage, simple_timer
 from verl.utils.torch_functional import check_device_is_available
 
@@ -101,65 +100,13 @@ def __init__(
     def __enter__(self):
         self.timing = {}
         with simple_timer("reshard", self.timing):
-            get_torch_device().empty_cache()
-
             loop = asyncio.get_event_loop()
-
-            if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine:
-                if self.multi_stage_wake_up:
-                    loop.run_until_complete(self.inference_engine.resume_memory_occupation(tags=["weights"]))
-                    log_gpu_memory_usage("Before resume SGLang weights in sharding manager", logger=logger)
-                else:
-                    loop.run_until_complete(self.inference_engine.resume_memory_occupation())
-                    log_gpu_memory_usage("Before resume SGLang weights + kv_cache in sharding manager", logger=logger)
-            get_torch_device().empty_cache()
-
-            log_gpu_memory_usage("Before state_dict() in sharding manager memory", logger=logger)
-            if self.offload_param:
-                load_fsdp_model_to_gpu(self.module)
-            params = self.module.state_dict()
-            log_gpu_memory_usage("After state_dict() in sharding manager memory", logger=logger)
-            device = get_device_id()  # used when fsdp2 set cpu_offload_policy
-            params = {
-                k: v.to(device, non_blocking=True) if fsdp_version(self.module) == 2 else v for k, v in params.items()
-            }
-            params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module))
-            # Copy, not share memory
-            loop.run_until_complete(self.update_weights(params))
-            log_gpu_memory_usage("After sync model weights in sharding manager", logger=logger)
-
-            del params
-            if self.offload_param:
-                offload_fsdp_model_to_cpu(self.module)
-            get_torch_device().empty_cache()
-            log_gpu_memory_usage("After del state_dict and empty_cache in sharding manager", logger=logger)
-
-            if self.multi_stage_wake_up and self.rollout_config.free_cache_engine:
-                loop.run_until_complete(self.inference_engine.resume_memory_occupation(tags=["kv_cache"]))
-                log_gpu_memory_usage("After resume SGLang kv_cache in sharding manager", logger=logger)
-
-            # important: need to manually set the random states of each tp to be identical.
-            if self.device_mesh is not None:
-                self.torch_random_states = get_torch_device().get_rng_state()
-                get_torch_device().set_rng_state(self.gen_random_states)
+            loop.run_until_complete(self.wake_up())
 
     @GPUMemoryLogger(role="FSDPSGLangShardingManager exit", logger=logger)
     def __exit__(self, exc_type, exc_value, traceback):
-        if self.rollout_config.free_cache_engine:
-            log_gpu_memory_usage("Before SGLang offload in sharding manager", logger=logger)
-            loop = asyncio.get_event_loop()
-            loop.run_until_complete(self.release_memory())
-            log_gpu_memory_usage("After SGLang offload in sharding manager", logger=logger)
-
-        self.module.train()
-
-        # add empty cache after each compute
-        get_torch_device().empty_cache()
-
-        # restore random states
-        if self.device_mesh is not None:
-            self.gen_random_states = get_torch_device().get_rng_state()
-            get_torch_device().set_rng_state(self.torch_random_states)
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self.sleep())
 
     async def update_weights(self, params):
         # Most naive implementation, can optimize a lot if it is bottleneck from sglang Engine weight update
@@ -207,6 +154,15 @@ async def wake_up(self):
         params = {
             k: v.to(device, non_blocking=True) if fsdp_version(self.module) == 2 else v for k, v in params.items()
         }
+
+        if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine:
+            if self.multi_stage_wake_up:
+                await self.inference_engine.resume_memory_occupation(tags=["weights"])
+                log_gpu_memory_usage("Before resume SGLang weights in sharding manager", logger=logger)
+            else:
+                await self.inference_engine.resume_memory_occupation()
+                log_gpu_memory_usage("Before resume SGLang weights + kv_cache in sharding manager", logger=logger)
+
         # Copy, not share memory
         await self.update_weights(params)
         log_gpu_memory_usage("After sync model weights in sharding manager", logger=logger)
@@ -217,6 +173,10 @@ async def wake_up(self):
         get_torch_device().empty_cache()
         log_gpu_memory_usage("After del state_dict and empty_cache in sharding manager", logger=logger)
 
+        if self.multi_stage_wake_up and self.rollout_config.free_cache_engine:
+            await self.inference_engine.resume_memory_occupation(tags=["kv_cache"])
+            log_gpu_memory_usage("After resume SGLang kv_cache in sharding manager", logger=logger)
+
         # important: need to manually set the random states of each tp to be identical.
         if self.device_mesh is not None:
             self.torch_random_states = get_torch_device().get_rng_state()
diff --git a/verl/workers/sharding_manager/megatron_sglang.py b/verl/workers/sharding_manager/megatron_sglang.py
@@ -114,45 +114,13 @@ def __init__(
     def __enter__(self):
         self.timing = {}
         with simple_timer("reshard", self.timing):
-            if self.offload_param:
-                load_megatron_model_to_gpu(self.actor_module)
-            if self.bridge is not None:
-                per_tensor_param = self.bridge.export_weights(self.actor_module)
-            else:
-                per_tensor_param = per_tensor_generator(
-                    self.actor_module,
-                    self.model_config,
-                    self.weight_converter,
-                    self.transformer_config,
-                    self.layer_name_mapping,
-                )
             loop = asyncio.get_event_loop()
-            loop.run_until_complete(self.update_weights(per_tensor_param))
-            if self.offload_param:
-                offload_megatron_model_to_cpu(self.actor_module)
-            get_torch_device().empty_cache()
-            # important: need to manually set the random states of each tp to be identical.
-            if self.device_mesh is not None:
-                self.torch_random_states = get_torch_device().get_rng_state()
-                get_torch_device().set_rng_state(self.gen_random_states)
+            loop.run_until_complete(self.wake_up())
 
     @GPUMemoryLogger(role="MegatronSGLangShardingManager exit", logger=logger)
     def __exit__(self, exc_type, exc_value, traceback):
-        if self.rollout_config.free_cache_engine:
-            log_gpu_memory_usage("Before SGLang offload in sharding manager", logger=logger)
-            loop = asyncio.get_event_loop()
-            loop.run_until_complete(self.release_memory())
-            log_gpu_memory_usage("After SGLang offload in sharding manager", logger=logger)
-
-        for model in self.actor_module:
-            model.train()
-        # add empty cache after each compute
-        get_torch_device().empty_cache()
-
-        # restore random states
-        if self.device_mesh is not None:
-            self.gen_random_states = get_torch_device().get_rng_state()
-            get_torch_device().set_rng_state(self.torch_random_states)
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self.sleep())
 
     async def update_weights(self, params):
         if self.device_mesh["tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine:
@@ -182,8 +150,10 @@ async def release_memory(self):
         if self.device_mesh["tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine:
             await self.inference_engine.release_memory_occupation()
 
-    @GPUMemoryLogger(role="FSDPSGLangShardingManager enter", logger=logger)
+    @GPUMemoryLogger(role="MegatronSGLangShardingManager enter", logger=logger)
     async def wake_up(self):
+        if self.offload_param:
+            load_megatron_model_to_gpu(self.actor_module)
         if self.bridge is not None:
             per_tensor_param = self.bridge.export_weights(self.actor_module)
         else:
@@ -195,12 +165,15 @@ async def wake_up(self):
                 self.layer_name_mapping,
             )
         await self.update_weights(per_tensor_param)
+        if self.offload_param:
+            offload_megatron_model_to_cpu(self.actor_module)
+        get_torch_device().empty_cache()
         # important: need to manually set the random states of each tp to be identical.
         if self.device_mesh is not None:
             self.torch_random_states = get_torch_device().get_rng_state()
             get_torch_device().set_rng_state(self.gen_random_states)
 
-    @GPUMemoryLogger(role="FSDPSGLangShardingManager exit", logger=logger)
+    @GPUMemoryLogger(role="MegatronSGLangShardingManager exit", logger=logger)
     async def sleep(self):
         if self.rollout_config.free_cache_engine:
             log_gpu_memory_usage("Before SGLang offload in sharding manager", logger=logger)
