Align GRPO vllm_enable_sleep_mode with the engine's actual sleep state (#6420)
The colocate GRPO setup set args.vllm_enable_sleep_mode from the raw UNSLOTH_VLLM_STANDBY env var (== '1'). When unsloth-zoo's vision standby gate disables sleep mode for a multimodal run, the engine is built with enable_sleep_mode=False but TRL was still told sleep mode is on, so it could drive sleep/wake against an engine that did not enable it. Read enable_sleep_mode from the colocated engine (model.vllm_engine.llm_engine.vllm_config.model_config), the same path check_sleep_mode uses, and fall back to the standby env var (!= '0', matching load_vllm/patch_vllm) only when the engine cannot be introspected. Pairs with unslothai/unsloth-zoo#768.
This commit is contained in:
parent
892d2983b0
commit
707641179f
1 changed files with 6 additions and 2 deletions
|
|
@ -1901,10 +1901,14 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
|
|||
# If model has vllm_engine, then use vllm in colocate mode. Donot wait for server
|
||||
vllm_setter += " " * 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
|
||||
# Align TRL sleep mode with the engine's actual enable_sleep_mode
|
||||
# (the vision standby gate may have disabled it); fall back to the
|
||||
# standby env var when the engine cannot be introspected.
|
||||
vllm_setter += (
|
||||
" " * 12
|
||||
+ "if os.environ.get('UNSLOTH_VLLM_STANDBY', '0') == '1':\n"
|
||||
+ "_unsloth_esm = getattr(getattr(getattr(getattr(model.vllm_engine, 'llm_engine', None), 'vllm_config', None), 'model_config', None), 'enable_sleep_mode', None)\n"
|
||||
+ " " * 12
|
||||
+ "if (_unsloth_esm if _unsloth_esm is not None else os.environ.get('UNSLOTH_VLLM_STANDBY', '0') != '0'):\n"
|
||||
+ " " * 16
|
||||
+ "args.vllm_enable_sleep_mode=True\n"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue