Fix vLLM sleep mode: add collective RPC call to reload weights in vLLM wake-up process

trl/experimental/gold/gold_trainer.py

@@ -1845,6 +1845,7 @@ def _move_model_to_vllm(self):
         if self.vllm_mode == "colocate" and self.vllm_enable_sleep_mode:
             empty_cache()
             self.vllm_engine.wake_up(tags=["weights"])
+            self.vllm_engine.collective_rpc("reload_weights")
 
         if is_peft_model(self.model):
             # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as

trl/experimental/openenv/utils.py

@@ -101,6 +101,7 @@ def generate_rollout_completions(
 
     if trainer.args.vllm_enable_sleep_mode:
         trainer.llm.wake_up(tags=["kv_cache"])
+        trainer.llm.collective_rpc("reload_weights")
 
     with profiling_context(trainer, "vLLM.generate_rollout"):
         if as_chat:

trl/trainer/grpo_trainer.py

@@ -1165,6 +1165,7 @@ def _generate_single_turn(self, prompts: list):
                 # wake up colocated vLLM instances if needed
                 torch.cuda.empty_cache()  # required to avoid OOM in some cases
                 self.llm.wake_up(tags=["weights"])
+                self.llm.collective_rpc("reload_weights")
 
             # First, update the vLLM weights if needed
             if self.state.global_step != self._last_loaded_step:

trl/trainer/rloo_trainer.py

@@ -1006,6 +1006,7 @@ def _generate_single_turn(self, prompts: list):
                 # wake up colocated vLLM instances if needed
                 torch.cuda.empty_cache()  # required to avoid OOM in some cases
                 self.llm.wake_up(tags=["weights"])
+                self.llm.collective_rpc("reload_weights")
 
             # First, update the vLLM weights if needed
             if self.state.global_step != self._last_loaded_step: