[perf] feat: clear megatron global buffer memory (verl-project#5173)

### What does this PR do?

- [x] clear megatron global buffer memory
- [x] synchronize device before empty_cache in vllm worker
- [ ] clear vllm encoder cache memory
vllm-project/vllm#33452

Author: wuxibin89

verl/utils/megatron_utils.py

@@ -31,6 +31,7 @@
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.enums import ModelType
 from megatron.core.optimizer import ChainedOptimizer
+from megatron.core.parallel_state import get_global_memory_buffer
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
 from megatron.core.transformer.module import Float16Module
 from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
@@ -598,7 +599,7 @@ def _iter_opts(opt):
             pass
 
         # Free Megatron-LM's global memory buffer
-        # get_global_memory_buffer().buffer.clear()
+        get_global_memory_buffer().buffer.clear()
 
         gc.collect()
         get_torch_device().empty_cache()

verl/workers/engine_workers.py

@@ -596,6 +596,8 @@ async def update_weights(self):
             return
 
         set_expandable_segments(False)
+        log_gpu_memory_usage("Before resume weights", logger=logger)
+
         # 1. resume weights and update weights
         if self.config.rollout.free_cache_engine:
             await self.rollout.resume(tags=["weights"])

verl/workers/rollout/vllm_rollout/utils.py

@@ -209,7 +209,7 @@ def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False
         # receive bucket and update weights
         while True:
             metadata = socket.recv_pyobj()
-            weights = []
+            weights, tensor = [], None
             for name, meta in metadata["bucket_meta"].items():
                 shape, dtype, offset = meta["shape"], meta["dtype"], meta["offset"]
                 size = dtype.itemsize * shape.numel()
@@ -225,7 +225,7 @@ def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False
             get_torch_device().synchronize()
             socket.send(b"")
             self._update_weights(weights, peft_config=peft_config, base_sync_done=base_sync_done)
-            del weights
+            del weights, tensor
             if metadata["is_last"]:
                 break
 
@@ -235,6 +235,7 @@ def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False
         if shm is not None:
             shm.close()
             del shm
+        get_torch_device().synchronize()
         gc.collect()
         get_torch_device().ipc_collect()
         get_torch_device().empty_cache()

verl/workers/rollout/vllm_rollout/vllm_async_server.py

@@ -580,6 +580,9 @@ async def sleep(self):
         if self.rollout_mode == RolloutMode.HYBRID:
             # Don't use engine.sleep(level=2) here
             await self.engine.collective_rpc("sleep", kwargs={"level": 2})
+
+            # clear encoder cache: https://github.com/vllm-project/vllm/pull/33452
+            # await self.engine.reset_encoder_cache()
         elif self.rollout_mode == RolloutMode.COLOCATED:
             await self.engine.sleep(level=1)
         elif self.rollout_mode == RolloutMode.STANDALONE: