From 26ddbb5b8e745d4906a0fda4aeb8276193a2e495 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 22 Oct 2025 17:30:52 +0530 Subject: [PATCH] Patch sleep mode properly for trl (#3492) --- unsloth/models/rl_replacements.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 7ecffdf63d..996c4f62b0 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -303,6 +303,34 @@ def grpo_trainer__generate_and_score_completions(function_name, function): if self.use_vllm:""" function = function.replace(replace_part, new_replacement) + if 'wake_up()' not in function: + # Sleep functionality has been added to trl in v0.23.0. We do not want to redo this. + # https://github.com/huggingface/trl/commit/edbe8234bc7e528f72ac76607de9d3e4753e2709 + + pattern = re.compile(r'.*self\.llm\.generate\(.*\).*', re.MULTILINE) + matches = list(pattern.finditer(function)) + patched = function + + # Generally there's only one match. But this is just to make sure we don't miss any. + for match in reversed(matches): + line = match.group(0) + indent_match = re.match(r'(\s*)', line) + indent = indent_match.group(1) if indent_match else '' + + wrapped = ( + f"{indent}if hasattr(self, 'llm'):\n" + f"{indent} if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n" + f"{indent} self.llm.wake_up()\n\n" + f"{line}\n\n" + f"{indent}if hasattr(self, 'llm'):\n" + f"{indent} if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n" + f"{indent} self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1))" + ) + + patched = patched[:match.start()] + wrapped + patched[match.end():] + + function = patched + return function pass RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_and_score_completions)