Discard weight when finish generation in the main loop

Signed-off-by: Guyue Huang <guyueh@login-lyris01.lyris.clusters.nvidia.com>

Author: Guyue Huang

nemo_rl/algorithms/grpo.py

@@ -1586,7 +1586,11 @@ def grpo_train(
                             max_rollout_turns=master_config.grpo["max_rollout_turns"],
                             greedy=False,
                         )
-                    policy_generation.finish_generation()
+                    policy_generation.finish_generation(
+                        discard_weights=colocated_inference
+                    )
+                    if colocated_inference:
+                        POLICY_GENERATION_STALE = True
                     # Collect generation logger metrics for performance reporting after each generation step
                     # inflight batch sizes and num pending samples are collected from each worker
                     if policy_generation is not None:

nemo_rl/models/generation/vllm/vllm_generation.py

@@ -732,10 +732,12 @@ def finish_generation(self, *args: Any, **kwargs: Any) -> bool:
                     if self.cfg["vllm_cfg"]["async_engine"]
                     else "reset_prefix_cache"
                 )
+                kwargs = {}
             # Use run_all_workers_single_data for methods that don't need data
             futures = self.worker_group.run_all_workers_single_data(
                 method_name,
                 run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
+                **kwargs,
             )
             # Wait for all futures to complete
             results = ray.get(futures)

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -986,7 +986,7 @@ def reset_prefix_cache(self):
         gc.collect()
         torch.cuda.empty_cache()
 
-    def sleep(self):
+    def sleep(self, discard_weights: bool = False):
         """Put the vLLM engine to sleep."""
         assert self.llm is not None, (
             "Attempting to sleep with either an uninitialized vLLM or non-model-owner"
@@ -1009,7 +1009,7 @@ def sleep(self):
             self.llm.renderer, "clear_mm_cache"
         ):
             self.llm.renderer.clear_mm_cache()
-        self.llm.sleep(level=1)
+        self.llm.sleep(level=2 if discard_weights else 1)
 
         gc.collect()
         torch.cuda.empty_cache()

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -1129,7 +1129,7 @@ async def reset_prefix_cache_async(self):
         gc.collect()
         torch.cuda.empty_cache()
 
-    async def sleep_async(self):
+    async def sleep_async(self, discard_weights: bool = False):
         """Async version of sleep."""
         assert self.llm is not None, (
             "Attempting to sleep with either an uninitialized vLLM or non-model-owner"
@@ -1148,7 +1148,7 @@ async def sleep_async(self):
         # the receiver and sends data=None, causing an assertion error.
         if hasattr(self.llm, "reset_mm_cache"):
             await self.llm.reset_mm_cache()
-        await self.llm.sleep(level=1)
+        await self.llm.sleep(level=2 if discard_weights else 1)
 
         gc.collect()
         torch.cuda.empty_cache()