fix

hebiao064 · hebiao064 · commit 503e19fd0da9 · 2025-07-01T04:09:37.000Z
diff --git a/verl/workers/sharding_manager/fsdp_sglang.py b/verl/workers/sharding_manager/fsdp_sglang.py
@@ -103,7 +103,7 @@ def __enter__(self):
 
             loop = asyncio.get_event_loop()
 
-            if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine:
+            if self.device_mesh["infer_tp"].get_local_rank() == 0:
                 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)
@@ -130,7 +130,7 @@ def __enter__(self):
             get_torch_device().empty_cache()
             log_gpu_memory_usage("After del state_dict and empty_cache in sharding manager", logger=logger)
 
-            if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.multi_stage_wake_up and self.rollout_config.free_cache_engine:
+            if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.multi_stage_wake_up:
                 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)
 
@@ -141,11 +141,10 @@ def __enter__(self):
 
     @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)
+        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()
 
@@ -188,7 +187,7 @@ async def update_weights(self, params):
                 )
 
     async def release_memory(self):
-        if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine:
+        if self.device_mesh["infer_tp"].get_local_rank() == 0:
             await self.inference_engine.release_memory_occupation()
 
     @GPUMemoryLogger(role="FSDPSGLangShardingManager enter", logger=logger)
@@ -218,10 +217,9 @@ async def wake_up(self):
 
     @GPUMemoryLogger(role="FSDPSGLangShardingManager 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)
-            await self.release_memory()
-            log_gpu_memory_usage("After SGLang offload in sharding manager", logger=logger)
+        log_gpu_memory_usage("Before SGLang offload in sharding manager", logger=logger)
+        await self.release_memory()
+        log_gpu_memory_usage("After SGLang offload in sharding manager", logger=logger)
 
         self.module.train()
 
