From 377c510e955c4db01b9e26858ed0945d92dc6e8d Mon Sep 17 00:00:00 2001 From: VED <146507396+ved1beta@users.noreply.github.com> Date: Wed, 8 Oct 2025 17:09:21 +0530 Subject: [PATCH] sleep model support (#3135) Co-authored-by: salman --- src/axolotl/core/trainers/grpo/__init__.py | 1 + src/axolotl/utils/schemas/trl.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/src/axolotl/core/trainers/grpo/__init__.py b/src/axolotl/core/trainers/grpo/__init__.py index 7eda7a0ba..d1a6b7fd9 100644 --- a/src/axolotl/core/trainers/grpo/__init__.py +++ b/src/axolotl/core/trainers/grpo/__init__.py @@ -52,6 +52,7 @@ class GRPOStrategy: if trl.vllm_mode: grpo_args_kwargs["vllm_mode"] = trl.vllm_mode if trl.vllm_mode == "colocate": + grpo_args_kwargs["enable_sleep_mode"] = trl.vllm_enable_sleep_mode # type: ignore[attr-defined] grpo_args_kwargs["vllm_gpu_memory_utilization"] = ( vllm_cfg.gpu_memory_utilization ) diff --git a/src/axolotl/utils/schemas/trl.py b/src/axolotl/utils/schemas/trl.py index 980474e87..624f7663e 100644 --- a/src/axolotl/utils/schemas/trl.py +++ b/src/axolotl/utils/schemas/trl.py @@ -167,3 +167,9 @@ class TRLConfig(BaseModel): "description": "Whether to exclude truncated completions from loss calculation." }, ) + vllm_enable_sleep_mode: bool | None = Field( + default=None, + json_schema_extra={ + "description": "Enable sleep mode for vLLM to offload VRAM when idle" + }, + )