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,8 @@ 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"])
+            # Work around for https://github.com/vllm-project/vllm/issues/29341
+            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,8 @@ def generate_rollout_completions(
 
     if trainer.args.vllm_enable_sleep_mode:
         trainer.llm.wake_up(tags=["kv_cache"])
+        # Work around for https://github.com/vllm-project/vllm/issues/29341
+        trainer.llm.collective_rpc("reload_weights")
 
     with profiling_context(trainer, "vLLM.generate_rollout"):
         if as_chat:

trl/trainer/grpo_trainer.py

@@ -1165,6 +1165,8 @@ 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"])
+                # Work around for https://github.com/vllm-project/vllm/issues/29341
+                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,8 @@ 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"])
+                # Work around for https://github.com/vllm-project/vllm/issues/29341
+                self.llm.collective_rpc("reload_weights")
 
             # First, update the vLLM weights if needed
             if self.state.global_step != self._last_loaded_step: