From b71a3b355bca7ff6731ff32c640863b365d6b291 Mon Sep 17 00:00:00 2001
From: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Date: Tue, 4 Nov 2025 12:30:54 +0530
Subject: [PATCH] Sleep trl patch (#3517)

* Patch sleep mode properly for trl

* empty cache after sleep/wakeup

* no extra wakeups

* Do not redo wakeups

* cleanup

* post trl 0.23 sleep patch
---
 unsloth/models/rl.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 67b407075c..6d4edb09f8 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -950,8 +950,13 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
             if "grpo" in trainer_file and trl_version >= Version("0.18.0"):
                 # If model has vllm_engine, then use vllm in colocate mode. Donot wait for server
                 vllm_setter += \
-                " " * 12 + "args.vllm_mode='colocate'\n"               
-    
+                " " * 12 + "args.vllm_mode='colocate'\n"
+                if trl_version >= Version("0.23.0"):
+                    # We need to set this flag for sleep mode auto working with trl update
+                    vllm_setter += \
+                    " " * 12 + "if os.environ.get('UNSLOTH_VLLM_STANDBY', '0') == '1':\n" + \
+                    " " * 16 + "args.vllm_enable_sleep_mode=True\n"
+
             init = init.replace(replacer, replacer + vllm_setter)
         pass
     pass