ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Sleep trl patch (#3517)

Browse source 
* 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
This commit is contained in:
Datta Nimmaturi 2025-11-04 12:30:54 +05:30 • committed by GitHub
commit b71a3b355b
1 changed files with 7 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/models/rl.py
View file

@ -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