From 49acbdf6bd2f7197fb019a8013edd24d39919b33 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 21 Apr 2026 14:58:07 +0000 Subject: [PATCH] inference: cuMem sleep/wake + lazy FlexEngine batch sizing Follow-up to 35231d4f (initial flex backend wiring). Sleep/wake: - sleep_mode.py adds kv_cache_pool / weight_pool context managers backed by vLLM's CuMemAllocator. Activated when UNSLOTH_VLLM_STANDBY=1 and vLLM is importable; no-op otherwise so TRL's unconditional sleep / wake_up calls stay valid. - FlexEngine routes the inference deep-copy + per-layer PagedKVCache through the pools. Captured CUDA graphs survive a sleep -> wake round-trip because cuMem keeps the GPU virtual addresses stable. - 4-bit single-copy path drops only the KV cache; level 2 warns and falls back to level 1. - tests/flex_sleep_mode_smoke.py covers sleep / wake memory deltas, captured-graph survival, and the no-op path. Lazy batch sizing: - FlexEngine's max_batch_size drives fixed-shape page tables, the input_pos_buffer, the block_mask_logical build, and the CUDA-graph bucket list at __init__ time. There is no post-init resize, so picking it at from_pretrained time forces over- or under-shoot. - build_flex_engine() defers construction until the GRPO rollout shape is known. install_flex_sentinel() attaches a _LazyFlexEngineSentinel to model.vllm_engine so hasattr(model, "vllm_engine") keeps working between from_pretrained and the first build; fast_generate triggers a floor build on first call. - rl.py injects _build_flex_from_args(model, args) before both self.llm = model.vllm_engine rewrite sites (pre-TRL-0.18 sampling_params prefix and >=0.18 colocate LLM replacement). Sizes the engine from max(pdbs * spg, pdbs * spg * ngen) derived from the GRPO args. No-op on non-flex models, so the injection is safe for every TRL backend. - Precedence: user's max_batch_size kwarg is a floor; the GRPO target overrides only when strictly larger, with a warning naming both. - First-build only: post-build growth raises RuntimeError pointing the user back to max_batch_size= in from_pretrained. The pristine inference deep-copy is consumed on first build and gemma4 shell extraction mutates its module tree, so a safe rebuild would require a second deep-copy. - tests/flex_lazy_batch_smoke.py (unit, stubbed FlexEngine) covers default, GRPO bump, user-floor, and post-build-refused cases. tests/flex_lazy_live_smoke.py exercises sentinel + build against a live Qwen3-0.6B-Base. --- tests/flex_lazy_batch_smoke.py | 227 +++++++++++++++++ tests/flex_lazy_live_smoke.py | 94 +++++++ tests/flex_sleep_mode_smoke.py | 299 ++++++++++++++++++++++ unsloth/inference/__init__.py | 9 +- unsloth/inference/flex_engine.py | 347 ++++++++++++++++++++++++-- unsloth/inference/flex_gemma4.py | 20 +- unsloth/inference/flex_qwen3_llama.py | 24 +- unsloth/inference/sleep_mode.py | 114 +++++++++ unsloth/models/llama.py | 52 ++-- unsloth/models/rl.py | 28 ++- unsloth/models/vision.py | 20 +- 11 files changed, 1163 insertions(+), 71 deletions(-) create mode 100644 tests/flex_lazy_batch_smoke.py create mode 100644 tests/flex_lazy_live_smoke.py create mode 100644 tests/flex_sleep_mode_smoke.py create mode 100644 unsloth/inference/sleep_mode.py diff --git a/tests/flex_lazy_batch_smoke.py b/tests/flex_lazy_batch_smoke.py new file mode 100644 index 0000000000..d56ce51f7f --- /dev/null +++ b/tests/flex_lazy_batch_smoke.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Smoke tests for the deferred FlexEngine batch-size sizing. + +Unit-level: covers the four cases described in the implementation plan +by monkey-patching :class:`FlexEngine` with a cheap stand-in so the +tests run on any box (no CUDA / no model download). The dispatch logic +lives entirely in :func:`build_flex_engine`, +:func:`install_flex_sentinel`, and :func:`_build_flex_from_args`, which +are the units under test. + +Run as: + python tests/flex_lazy_batch_smoke.py +""" + +from __future__ import annotations + +import sys +import types +import warnings +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +class _StubFlexEngine: + """Records the ``max_batch_size`` construction arg and nothing else. + + ``_cudagraph_primed`` flips True after the first ``generate`` so the + post-warmup refuse path can be exercised without touching CUDA. + """ + + instances: list = [] + + def __init__( + self, + hf_model, + tokenizer, + *, + dtype = None, + max_seq_length: int = 2048, + max_lora_rank: int = 64, + max_batch_size: int = 32, + page_size: int = 128, + gpu_memory_utilization: float = 0.5, + max_new_tokens: int = 512, + prefill_kernel_options = None, + decode_kernel_options = None, + fa4_prefill = None, + capture_cudagraph: bool = True, + base_model = None, + peft_model = None, + inference_model = None, + ): + self.hf_model = hf_model + self.tokenizer = tokenizer + self.max_batch_size = max_batch_size + self.max_seq_length = max_seq_length + self.compute_dtype = dtype + self._cudagraph_primed = False + self.generate_calls = 0 + _StubFlexEngine.instances.append(self) + + def generate(self, prompts = None, *args, **kwargs): + self.generate_calls += 1 + self._cudagraph_primed = True + return [("stub", prompts)] + + +def _make_stub_model(): + """An object that quacks like an HF model for ``install_flex_sentinel``.""" + + model = types.SimpleNamespace() + model._unsloth_needs_flex_engine = dict( + dtype = "bf16", + max_seq_length = 2048, + max_lora_rank = 64, + max_batch_size = 32, + gpu_memory_utilization = 0.5, + ) + model._unsloth_flex_inference_copy = object() # never dereferenced + return model + + +def _install_stub(): + """Patch FlexEngine with :class:`_StubFlexEngine` for the duration of the + test process. Imports happen lazily inside ``build_flex_engine``, so we + patch the module attribute before those calls fire.""" + + import unsloth.inference.flex_engine as fe + + _StubFlexEngine.instances.clear() + fe.FlexEngine = _StubFlexEngine + + +def _case1_default_floor(): + """No trainer, no kwargs: fast_generate builds at floor=32.""" + + from unsloth.inference.flex_engine import install_flex_sentinel + + _install_stub() + model = _make_stub_model() + install_flex_sentinel(model, tokenizer = object()) + + assert hasattr(model, "vllm_engine"), "sentinel not installed" + assert not hasattr(model, "_flex_engine_instance"), ( + "engine should NOT exist before first use" + ) + + out = model.fast_generate(["hello"]) + assert out == [("stub", ["hello"])] + + engine = model._flex_engine_instance + assert engine.max_batch_size == 32, engine.max_batch_size + # Sentinel was replaced with the real engine after build. + assert model.vllm_engine is engine + print(" [1/4] default path: floor=32 build on first fast_generate OK") + + +def _case2_grpo_bump(): + """User kwarg=16 + GRPO target=64 → engine built at 64, warning logged.""" + + from unsloth.inference.flex_engine import ( + _build_flex_from_args, + install_flex_sentinel, + ) + + _install_stub() + model = _make_stub_model() + model._unsloth_needs_flex_engine["max_batch_size"] = 16 # user floor + install_flex_sentinel(model, tokenizer = object()) + + args = types.SimpleNamespace( + per_device_train_batch_size = 2, + steps_per_generation = 4, + num_generations = 8, + gradient_accumulation_steps = 1, + ) + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _build_flex_from_args(model, args) + + engine = model._flex_engine_instance + assert engine.max_batch_size == 64, engine.max_batch_size + assert any("16 -> 64" in str(w.message) for w in caught), [ + str(w.message) for w in caught + ] + print(" [2/4] GRPO bump: 16 -> 64 with warning OK") + + +def _case3_user_floor_wins(): + """User kwarg=128 + GRPO target=8 → engine stays at 128, no warning.""" + + from unsloth.inference.flex_engine import ( + _build_flex_from_args, + install_flex_sentinel, + ) + + _install_stub() + model = _make_stub_model() + model._unsloth_needs_flex_engine["max_batch_size"] = 128 + install_flex_sentinel(model, tokenizer = object()) + + args = types.SimpleNamespace( + per_device_train_batch_size = 1, + steps_per_generation = 2, + num_generations = 4, + gradient_accumulation_steps = 1, + ) + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _build_flex_from_args(model, args) + + engine = model._flex_engine_instance + assert engine.max_batch_size == 128, engine.max_batch_size + assert not any("FlexEngine" in str(w.message) for w in caught), [ + str(w.message) for w in caught + ] + print(" [3/4] user floor wins: engine.max_batch_size=128 OK") + + +def _case4_post_warmup_refused(): + """fast_generate primes the engine; later GRPO target=64 must raise.""" + + from unsloth.inference.flex_engine import ( + _build_flex_from_args, + install_flex_sentinel, + ) + + _install_stub() + model = _make_stub_model() + install_flex_sentinel(model, tokenizer = object()) + + model.fast_generate(["hi"]) # builds at floor=32, sets _cudagraph_primed + assert model._flex_engine_instance.max_batch_size == 32 + + args = types.SimpleNamespace( + per_device_train_batch_size = 2, + steps_per_generation = 4, + num_generations = 8, + gradient_accumulation_steps = 1, + ) + try: + _build_flex_from_args(model, args) + except RuntimeError as exc: + msg = str(exc) + assert "32" in msg and "64" in msg, msg + assert "max_batch_size=64" in msg, msg + print(" [4/4] post-warmup rebuild refused with actionable msg OK") + return + raise AssertionError("expected RuntimeError when growing a built engine") + + +def main(): + print("flex_lazy_batch_smoke:") + _case1_default_floor() + _case2_grpo_bump() + _case3_user_floor_wins() + _case4_post_warmup_refused() + print("ALL CASES PASSED") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_lazy_live_smoke.py b/tests/flex_lazy_live_smoke.py new file mode 100644 index 0000000000..8d7166c828 --- /dev/null +++ b/tests/flex_lazy_live_smoke.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Live smoke-test for deferred FlexEngine construction. + +Loads a small flex-supported model with ``UNSLOTH_FAST_INFERENCE=1`` and +checks: + + 1. After ``from_pretrained``, ``model.vllm_engine`` is the lazy + sentinel -- no ``_flex_engine_instance`` yet. + 2. ``build_flex_engine(model)`` constructs the engine at the stashed + ``max_batch_size`` floor (default 32) and wires + ``model.vllm_engine`` / ``fast_generate`` onto it. + 3. ``build_flex_engine(model, max_batch_size=X)`` with ``X`` larger + than the built size raises :class:`RuntimeError` with an actionable + hint pointing the user back to ``max_batch_size=`` in + ``from_pretrained``. + +This intentionally does NOT call ``engine.generate`` -- that path is +covered by the existing ``tests/flex_fastlm_smoke.py`` and would +duplicate its warm-up cost. The goal here is to confirm the lazy +dispatch behavior end-to-end. + +Run as: + CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 \ + python tests/flex_lazy_live_smoke.py +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def main(): + assert os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1", ( + "export UNSLOTH_FAST_INFERENCE=1 before running this smoke" + ) + import unsloth # noqa: F401 (must import before transformers) + from unsloth import FastLanguageModel + from unsloth.inference.flex_engine import ( + FlexEngine, + _LazyFlexEngineSentinel, + build_flex_engine, + ) + + model_name = os.environ.get( + "FLEX_LAZY_SMOKE_MODEL", "unsloth/Qwen3-0.6B-Base" + ) + print(f"loading {model_name} ...") + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = model_name, + max_seq_length = 1024, + fast_inference = True, + load_in_4bit = False, + ) + + assert hasattr(model, "vllm_engine"), "vllm_engine attr missing" + sentinel = model.vllm_engine + assert isinstance(sentinel, _LazyFlexEngineSentinel), ( + f"expected sentinel, got {type(sentinel)}" + ) + assert not hasattr(model, "_flex_engine_instance"), ( + "engine should NOT be built before first use" + ) + print(" [1/3] sentinel installed, no engine yet") + + engine = build_flex_engine(model) + assert isinstance(engine, FlexEngine), type(engine) + assert engine.max_batch_size == 32, engine.max_batch_size + assert model._flex_engine_instance is engine + assert model.vllm_engine is engine + print(f" [2/3] build_flex_engine built engine at max_batch_size={engine.max_batch_size}") + + try: + build_flex_engine(model, max_batch_size = 64) + except RuntimeError as exc: + msg = str(exc) + assert "32" in msg and "64" in msg, msg + assert "max_batch_size=64" in msg, msg + print(" [3/3] post-build resize refused with actionable msg") + else: + raise AssertionError("expected RuntimeError when growing a built engine") + + print("LIVE SMOKE PASSED") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_sleep_mode_smoke.py b/tests/flex_sleep_mode_smoke.py new file mode 100644 index 0000000000..5744df85da --- /dev/null +++ b/tests/flex_sleep_mode_smoke.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Smoke-test :meth:`FlexEngine.sleep` and :meth:`FlexEngine.wake_up`. + +Invoked as: + CUDA_VISIBLE_DEVICES=3 \ + UNSLOTH_FAST_INFERENCE=1 UNSLOTH_VLLM_STANDBY=1 \ + python tests/flex_sleep_mode_smoke.py --model unsloth/Qwen3-4B-Base + +What it checks: + 1. ``model.vllm_engine._sleep_mode_enabled`` is True when vLLM is + importable. + 2. Captured CUDA graphs survive a sleep/wake round-trip: a second + ``fast_generate`` call after ``sleep`` -> ``wake_up`` does not + re-capture and still produces the same token ids. + 3. ``torch.cuda.memory_allocated()`` drops on sleep and returns close + to the pre-sleep value on wake. With ``--no-standby``, memory + should be identical across the three probes (no-op path). + 4. Hardens the regression guard by running the same workflow with + ``UNSLOTH_VLLM_STANDBY`` unset (``--no-standby``); ``sleep`` / + ``wake_up`` must be exact no-ops. + +This is a smoke test, not the full verification matrix from the plan +(that lives under ``scripts/benchmarks``). +""" + +from __future__ import annotations + +import argparse +import gc +import os +import sys +import time +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _gb(n: int) -> float: + return round(n / 1e9, 3) + + +def _probe(label: str) -> dict: + import torch + + gc.collect() + torch.cuda.synchronize() + # ``memory_allocated`` / ``memory_reserved`` only track torch's + # caching allocator and ignore cuMem-backed pools, so they do NOT + # drop on sleep even though cuMem has unmapped the pages. + # ``mem_get_info()`` asks the CUDA runtime directly, so it sees + # cuMem unmaps and is the right probe for sleep / wake verification. + free_bytes, total_bytes = torch.cuda.mem_get_info() + return { + "label": label, + "allocated_gb": _gb(torch.cuda.memory_allocated()), + "reserved_gb": _gb(torch.cuda.memory_reserved()), + "cuda_free_gb": _gb(free_bytes), + "cuda_used_gb": _gb(total_bytes - free_bytes), + } + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default = "unsloth/Qwen3-4B-Base") + p.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16") + p.add_argument("--load_in_4bit", action = "store_true") + p.add_argument("--max_new_tokens", type = int, default = 32) + p.add_argument("--max_seq_length", type = int, default = 1024) + p.add_argument("--prompt", default = "The quick brown fox jumps over") + p.add_argument( + "--no-standby", + action = "store_true", + help = "Force UNSLOTH_VLLM_STANDBY=0 to validate the no-op regression path.", + ) + p.add_argument( + "--cycles", + type = int, + default = 1, + help = "Number of sleep / wake / generate cycles after warmup. " + ">1 exercises the repeated-cycle regression (run #6).", + ) + args = p.parse_args() + + os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") + if args.no_standby: + os.environ["UNSLOTH_VLLM_STANDBY"] = "0" + else: + os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1") + standby = os.environ.get("UNSLOTH_VLLM_STANDBY", "0") == "1" + print( + f"[sleep-smoke] UNSLOTH_FAST_INFERENCE=" + f"{os.environ.get('UNSLOTH_FAST_INFERENCE')} " + f"UNSLOTH_VLLM_STANDBY={os.environ.get('UNSLOTH_VLLM_STANDBY')}" + ) + + import torch + + import unsloth + from unsloth import FastLanguageModel + + print(f"[sleep-smoke] unsloth={unsloth.__file__}") + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + t0 = time.perf_counter() + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = args.model, + max_seq_length = args.max_seq_length, + dtype = dtype, + load_in_4bit = args.load_in_4bit, + fast_inference = True, + ) + print( + f"[sleep-smoke] loaded {args.model} in " + f"{time.perf_counter() - t0:.1f}s; dtype={model.dtype}" + ) + + engine = model.vllm_engine + print(f"[sleep-smoke] engine type: {type(engine).__name__}") + sleep_enabled = getattr(engine, "_sleep_mode_enabled", None) + print(f"[sleep-smoke] engine._sleep_mode_enabled: {sleep_enabled}") + if standby: + if sleep_enabled is not True: + # Most likely vLLM is not importable in this environment. + print( + "[sleep-smoke] WARNING: UNSLOTH_VLLM_STANDBY=1 was set " + "but engine._sleep_mode_enabled is False (vLLM missing?)" + ) + else: + assert sleep_enabled is False, ( + f"Expected sleep mode to be disabled with UNSLOTH_VLLM_STANDBY=0, " + f"got {sleep_enabled}" + ) + + ll_cfg = engine.llm_engine.vllm_config.model_config + print( + f"[sleep-smoke] llm_engine.vllm_config.model_config.enable_sleep_mode: " + f"{getattr(ll_cfg, 'enable_sleep_mode', None)}" + ) + assert getattr(ll_cfg, "enable_sleep_mode", None) == bool(sleep_enabled), ( + "_LLMEngineStub.model_config.enable_sleep_mode must mirror " + "engine._sleep_mode_enabled" + ) + + from unsloth.inference.vllm_shim import LoRARequest # noqa: F401 + + prompts = [args.prompt] + + # ----- warmup ----- + t0 = time.perf_counter() + out1 = engine.generate( + prompts, + sampling_params = type( + "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0}, + )(), + ) + print( + f"[sleep-smoke] warmup generate: {time.perf_counter() - t0:.2f}s; " + f"tok_ids[:10]={out1[0].outputs[0].token_ids[:10]}" + ) + pre_tokens = list(out1[0].outputs[0].token_ids) + + probe_pre = _probe("pre-sleep") + print(f"[sleep-smoke] {probe_pre}") + + # Diagnostic: checksum the inference-model weights so we can detect + # if cuMem's sleep/wake round-trip corrupts any parameter. + def _checksum_params(mod, limit = 16): + import torch as _t + + out = [] + for i, (name, p) in enumerate(mod.named_parameters()): + if i >= limit: + break + t = p.detach() + out.append((name, list(t.shape), float(t.float().abs().sum().item()))) + return out + + pre_sums = _checksum_params(engine._inference_model) + print("[sleep-smoke] pre-sleep first-16 param |sum|:") + for n, s, v in pre_sums: + print(f" {n} {s} {v:.4f}") + + # Repeated sleep / wake / generate cycle test (plan matrix run #6). + # Each cycle validates that the engine does not drift: tokens remain + # bitwise identical, memory returns to baseline, weights round-trip + # cleanly. A bug that only surfaces on the second or third cycle + # (stale Python state, double-wake, leaked handles) fails here. + for cycle in range(args.cycles): + if args.cycles > 1: + print(f"[sleep-smoke] --- cycle {cycle + 1}/{args.cycles} ---") + + # ----- sleep ----- + t0 = time.perf_counter() + engine.sleep(level = 1) + t_sleep = time.perf_counter() - t0 + probe_post = _probe(f"post-sleep[{cycle + 1}]") + print( + f"[sleep-smoke] sleep(level=1) took {t_sleep:.3f}s; {probe_post}" + ) + + if sleep_enabled: + drop = probe_pre["cuda_used_gb"] - probe_post["cuda_used_gb"] + print( + f"[sleep-smoke] process-level VRAM drop on sleep: " + f"{drop:+.3f} GB (cuMem-managed; not visible in " + f"torch.memory_allocated)" + ) + if not getattr(engine, "_single_copy_mode", False): + # 16-bit path: both weights + kv_cache pools are dropped. + assert drop >= 1.0, ( + f"Expected multi-GB drop in process-level VRAM on " + f"sleep(level=1); got {drop:+.3f} GB" + ) + else: + # 4-bit single-copy: only KV cache drops; weights stay. + assert drop > 0.0, ( + f"Expected KV-cache drop in process-level VRAM on " + f"sleep(level=1); got {drop:+.3f} GB" + ) + else: + # With sleep mode off, the sleep() call must not free VRAM. + # Process-level jitter is allowed (shared GPU); torch-owned + # allocations must be untouched. + assert ( + probe_post["allocated_gb"] == probe_pre["allocated_gb"] + ), ( + "With sleep mode disabled, torch.memory_allocated must " + "be unchanged by sleep()" + ) + + # ----- wake ----- + t0 = time.perf_counter() + engine.wake_up() + t_wake = time.perf_counter() - t0 + probe_wake = _probe(f"post-wake[{cycle + 1}]") + print( + f"[sleep-smoke] wake_up() took {t_wake:.3f}s; {probe_wake}" + ) + + post_sums = _checksum_params(engine._inference_model) + diffs = [] + for (n1, s1, v1), (n2, s2, v2) in zip(pre_sums, post_sums): + delta = abs(v1 - v2) + if delta > 0.0: + diffs.append((n1, v1, v2, delta)) + print( + f"[sleep-smoke] post-wake weight diff: " + f"{len(diffs)}/{len(pre_sums)} params changed " + f"(bitwise-exact restore expected)" + ) + for n, v1, v2, d in diffs[:8]: + print(f" diff {n}: pre={v1:.4f} post={v2:.4f} delta={d:.4f}") + assert len(diffs) == 0, ( + f"Weight corruption on sleep / wake (cycle {cycle + 1}): " + f"{len(diffs)}/{len(pre_sums)} first-layer params changed" + ) + + # ----- verify we can still generate ----- + t0 = time.perf_counter() + out2 = engine.generate( + prompts, + sampling_params = type( + "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0}, + )(), + ) + t_regen = time.perf_counter() - t0 + post_tokens = list(out2[0].outputs[0].token_ids) + match = pre_tokens == post_tokens + print( + f"[sleep-smoke] post-wake generate: {t_regen:.2f}s; " + f"tok_ids[:10]={post_tokens[:10]}; matches_pre={match}" + ) + assert match, ( + f"Pre-sleep / post-wake token ids must match exactly " + f"(cycle {cycle + 1}).\n" + f"Pre: {pre_tokens}\nPost: {post_tokens}" + ) + + if sleep_enabled: + delta = probe_wake["cuda_used_gb"] - probe_pre["cuda_used_gb"] + print( + f"[sleep-smoke] post-wake vs pre-sleep process-level " + f"VRAM delta: {delta:+.3f} GB (tolerance: +/- 1.5 GB)" + ) + assert abs(delta) < 1.5, ( + f"Post-wake VRAM diverged from pre-sleep " + f"(cycle {cycle + 1}): delta={delta:+.3f} GB" + ) + + print(f"[sleep-smoke] PASS ({args.cycles} cycle(s))") + + +if __name__ == "__main__": + main() diff --git a/unsloth/inference/__init__.py b/unsloth/inference/__init__.py index 476bda5d39..e81e5bbe9a 100644 --- a/unsloth/inference/__init__.py +++ b/unsloth/inference/__init__.py @@ -14,7 +14,12 @@ Three architectures are supported today: Qwen3, Llama-3, Gemma-4-E2B-it. Anything else raises :class:`NotImplementedError` — there is no silent fallback; unset the env var or use vLLM instead.""" -from .flex_engine import FlexEngine, load_flex +from .flex_engine import ( + FlexEngine, + build_flex_engine, + install_flex_sentinel, + load_flex, +) from .vllm_shim import ( CompletionOutput, LoRARequest, @@ -26,6 +31,8 @@ from .vllm_shim import ( __all__ = [ "FlexEngine", "load_flex", + "build_flex_engine", + "install_flex_sentinel", "LoRARequest", "RequestOutput", "CompletionOutput", diff --git a/unsloth/inference/flex_engine.py b/unsloth/inference/flex_engine.py index ef7da5acaf..b566241d2c 100644 --- a/unsloth/inference/flex_engine.py +++ b/unsloth/inference/flex_engine.py @@ -32,6 +32,7 @@ from __future__ import annotations import copy import functools +import gc import importlib.util import os import types @@ -49,6 +50,12 @@ from .flex_qwen3_llama import ( refresh_lora_merge_from_pristine, ) from .flex_gemma4 import FlexGemma4Inference +from .sleep_mode import ( + _get_cumem_allocator, + kv_cache_pool, + sleep_mode_enabled, + weight_pool, +) from .vllm_shim import CompletionOutput, LoRARequest, RequestOutput @@ -242,8 +249,13 @@ class _LLMEngineStub: ``pass`` (rl.py:1795-1810), so a nested attribute chain that simply exists is enough.""" - def __init__(self): - self.vllm_config = types.SimpleNamespace(lora_config = types.SimpleNamespace()) + def __init__(self, sleep_enabled: bool = False): + self.vllm_config = types.SimpleNamespace( + lora_config = types.SimpleNamespace(), + model_config = types.SimpleNamespace( + enable_sleep_mode = bool(sleep_enabled), + ), + ) self.model_executor = types.SimpleNamespace( driver_worker = types.SimpleNamespace( model_runner = types.SimpleNamespace( @@ -309,6 +321,22 @@ class FlexEngine: self.device = hf_model.device + # Sleep-mode setup. When ``UNSLOTH_VLLM_STANDBY=1`` is set AND + # vLLM is importable, we route the engine's heavy allocations + # (the inference deep-copies + per-layer PagedKVCache buffers) + # through cuMem-backed pools so ``FlexEngine.sleep`` can offload + # weights to pinned CPU memory without destroying captured CUDA + # graphs. The allocator assigns stable GPU virtual addresses, so + # unmapping on sleep and re-mapping on wake preserves pointer + # validity. ``expandable_segments:True`` on + # ``PYTORCH_CUDA_ALLOC_CONF`` is incompatible with cuMem; if the + # user has it set, ``_get_cumem_allocator`` returns None and + # sleep stays a no-op. + self._sleep_mode_enabled = sleep_mode_enabled() + self._cumem_allocator = ( + _get_cumem_allocator() if self._sleep_mode_enabled else None + ) + # Colocate pattern (mirrors vLLM's "colocate" mode): the engine # runs on its own deep-copy of the HF model so that flex attention # patching + KV-cache attachment does not mutate the training @@ -325,11 +353,18 @@ class FlexEngine: # defer materialising it until :meth:`bind_peft_model`, so the # no-LoRA path stays at 2x the base model's VRAM instead of 3x. if inference_model is None: - inference_model = copy.deepcopy(hf_model) - inference_model.eval() + with weight_pool(self._cumem_allocator): + inference_model = copy.deepcopy(hf_model) + inference_model.eval() self._pristine_base = base_model # None until bind_peft_model runs self._inference_model = inference_model self._inference_peft = peft_model # filled in by bind_peft_model + # ``_inference_model is hf_model`` is the 4-bit / single-copy + # fallback path (see ``bind_peft_model``); flex skips the second + # deep-copy there because bnb-4bit packed weights can't be + # in-place refreshed. In that mode there is no CPU-backup step + # to do on ``sleep`` — only the KV cache gets dropped. + self._single_copy_mode = inference_model is hf_model # Autocast wrapping (applied by the impl via a context manager). fa4_prefill, prefill_kernel_options, decode_kernel_options = ( @@ -353,6 +388,14 @@ class FlexEngine: inference_model = _extract_gemma4_text_shell(inference_model) self._inference_model = inference_model Impl = FlexGemma4Inference if arch == "gemma4" else FlexInference + # Pass the cuMem allocator through so the impl can wrap ONLY + # the paged-KV allocations (``PageTable`` + per-layer + # ``PagedKVCache``) in the ``kv_cache`` pool. Everything else + # the impl creates (``input_pos_buffer``, ``block_mask_logical``, + # captured CUDA-graph scratch / graph_vars) stays in torch's + # default allocator — those buffers are tiny and, critically, + # the captured CUDA graphs reference block_mask indices by + # address, so they must survive sleep / wake unchanged. self._impl = Impl( inference_model, tokenizer, @@ -366,8 +409,11 @@ class FlexEngine: fa4_prefill = fa4_prefill, base_model = self._pristine_base, peft_model = peft_model, + cumem_allocator = self._cumem_allocator, + ) + self._llm_engine_stub = _LLMEngineStub( + sleep_enabled = self._sleep_mode_enabled, ) - self._llm_engine_stub = _LLMEngineStub() # ----- configuration helpers ----- @@ -605,20 +651,74 @@ class FlexEngine: self._warned_sampling = True return int(max_tokens), {} - # ----- sleep-mode stubs ----- + # ----- sleep-mode ----- # # vLLM's sleep mode offloads engine weights to CPU between rollouts so - # training can use the freed VRAM. The flex backend does not implement - # that yet; these stubs exist so code paths that assume the API are safe. + # training can use the freed VRAM. The flex backend implements level 1 + # (weights offloaded to pinned CPU, KV cache dropped and re-zeroed on + # wake) via :class:`vllm.device_allocator.cumem.CuMemAllocator`. + # Captured CUDA graphs survive the round-trip because cuMem keeps the + # GPU virtual addresses stable across sleep / wake. + # + # Sleep activates only when ``UNSLOTH_VLLM_STANDBY=1`` is set AND + # vLLM is importable (evaluated at ``__init__`` time). Otherwise + # ``sleep`` / ``wake_up`` are no-ops so code that unconditionally + # calls the API (TRL's GRPO trainer) stays correct. - def sleep(self, level: int = 2): - """No-op. Real implementation in a follow-up PR (move PagedKVCache + - shared buffers + the inference-copy HF shell to CPU pinned memory - and swap back on wake_up).""" + def sleep(self, level: int = 1): + """Offload inference weights to pinned CPU memory (level 1). + + ``level=2`` is not implemented on the flex backend (it would + require rebuilding the inference deep-copy from the training + model on wake); requesting it emits a warning and falls back to + level 1. + + In the 4-bit single-copy fallback path the inference model + shares storage with the training model, so only the KV cache is + dropped on sleep; the weights stay resident. + """ + if not self._sleep_mode_enabled or self._cumem_allocator is None: + return None + if level not in (1, 2): + raise ValueError( + f"FlexEngine.sleep: level must be 1 or 2, got {level}" + ) + if level == 2: + warnings.warn( + "FlexEngine.sleep(level=2) is not implemented on the " + "flex backend; falling back to level=1 (CPU-pinned " + "weight offload).", + RuntimeWarning, + stacklevel = 2, + ) + if self._single_copy_mode: + # Weights are shared with the training model (4-bit path); + # only the kv_cache pool is ours to drop. + self._cumem_allocator.sleep(offload_tags = ()) + else: + self._cumem_allocator.sleep(offload_tags = ("weights",)) + gc.collect() + # NOTE: we deliberately do NOT call torch.cuda.empty_cache() here. + # Captured CUDA graphs may retain scratch / workspace tensors in + # torch's default caching allocator at fixed addresses; emptying + # the cache between sleep and wake can invalidate those + # addresses and cause the next graph replay to read freed + # memory. The cuMem pools have already released their physical + # pages; there is no additional VRAM to reclaim via empty_cache. return None def wake_up(self, tags: Optional[list] = None): - """No-op companion to :meth:`sleep`.""" + """Re-map cuMem handles and restore offloaded weights. + + ``tags=None`` wakes everything; TRL's GRPO trainer calls + ``wake_up(tags=["kv_cache"])`` and then ``wake_up(tags=["weights"])`` + on consecutive steps to stagger the VRAM reclaim. + """ + if not self._sleep_mode_enabled or self._cumem_allocator is None: + return None + if tags is not None and not isinstance(tags, list): + tags = list(tags) + self._cumem_allocator.wake_up(tags = tags) return None @property @@ -654,8 +754,9 @@ class FlexEngine: # linear weights (what LoRA merges into) are still pristine, # so we can clone it and just not call flex attention on the # pristine copy. - self._pristine_base = copy.deepcopy(self._inference_model) - self._pristine_base.eval() + with weight_pool(self._cumem_allocator): + self._pristine_base = copy.deepcopy(self._inference_model) + self._pristine_base.eval() self._impl.base_model = self._pristine_base if self._inference_peft is None: @@ -666,8 +767,11 @@ class FlexEngine: # Wrap the already-patched inference copy with a fresh LoRA # adapter of the same shape. LoraLayer insertion is # attention-forward-agnostic; it wraps Linear modules. - self._inference_peft = _get_peft_model(self._inference_model, peft_cfg) - self._inference_peft.eval() + with weight_pool(self._cumem_allocator): + self._inference_peft = _get_peft_model( + self._inference_model, peft_cfg, + ) + self._inference_peft.eval() except Exception as e: warnings.warn( f"FlexEngine.bind_peft_model: could not build an " @@ -724,4 +828,211 @@ def load_flex( ) -__all__ = ["FlexEngine", "load_flex"] +# --------------------------------------------------------------------------- +# Lazy engine construction +# +# FlexEngine's ``max_batch_size`` drives fixed-shape GPU page tables, the +# ``input_pos_buffer``, the ``block_mask_logical`` build, and the CUDA-graph +# bucket list. These are allocated inside ``FlexEngine.__init__`` and there is +# no post-init resize path. Picking the wrong value at ``from_pretrained`` +# time forces users to either overshoot (wasted KV pages + slower graph +# capture) or undershoot (``PageTable.can_reserve`` stalls the rollout). +# +# ``build_flex_engine`` defers the construction until the real rollout batch +# size is known: GRPOTrainer's ``__init__`` patch in ``unsloth/models/rl.py`` +# calls :func:`_build_flex_from_args` to pass +# ``per_device_train_batch_size * steps_per_generation * num_generations`` +# through; the plain ``model.fast_generate`` path falls back to the +# ``max_batch_size`` kwarg originally passed to ``from_pretrained``. +# --------------------------------------------------------------------------- + + +class _LazyFlexEngineSentinel: + """Placeholder for ``model.vllm_engine`` before the FlexEngine is built. + + Forwards attribute access to the real engine, triggering construction + (with the stashed ``max_batch_size`` floor) on first access. This keeps + ``hasattr(model, "vllm_engine")`` True between ``from_pretrained`` and + the first build, which matters for ``rl.py``'s ``args.use_vllm`` setter. + """ + + __slots__ = ("_model",) + + def __init__(self, model): + object.__setattr__(self, "_model", model) + + def _resolve(self): + engine = getattr(self._model, "_flex_engine_instance", None) + if engine is None: + engine = build_flex_engine(self._model) + return engine + + def __getattr__(self, name): + if name == "_model": + raise AttributeError(name) + return getattr(self._resolve(), name) + + def __bool__(self): + return True + + def __repr__(self): + engine = getattr(self._model, "_flex_engine_instance", None) + if engine is None: + return "" + return repr(engine) + + +def install_flex_sentinel(model, tokenizer): + """Wire the lazy ``vllm_engine`` / ``fast_generate`` placeholders. + + Called from ``unsloth/models/llama.py`` and ``unsloth/models/vision.py`` + in place of the eager ``FlexEngine(...)`` construction. The real build + happens inside :func:`build_flex_engine`, triggered either by the + GRPOTrainer patch (``_build_flex_from_args``) or by the first + ``model.fast_generate`` call. + """ + model._unsloth_flex_tokenizer = tokenizer + model.vllm_engine = _LazyFlexEngineSentinel(model) + + def _lazy_fast_generate(prompts = None, *gen_args, **gen_kwargs): + engine = getattr(model, "_flex_engine_instance", None) + if engine is None: + engine = build_flex_engine(model) + return engine.generate(prompts, *gen_args, **gen_kwargs) + + def _lazy_fast_generate_batches(prompts = None, *gen_args, **gen_kwargs): + engine = getattr(model, "_flex_engine_instance", None) + if engine is None: + engine = build_flex_engine(model) + gen_kwargs.setdefault("use_tqdm", False) + return engine.generate(prompts, *gen_args, **gen_kwargs) + + model.fast_generate = _lazy_fast_generate + model.fast_generate_batches = _lazy_fast_generate_batches + + +def _construct_and_attach(model, max_batch_size: int): + """Construct the FlexEngine with the given batch size and wire it up.""" + pending = model._unsloth_needs_flex_engine + tokenizer = getattr(model, "_unsloth_flex_tokenizer", None) + inference_copy = getattr(model, "_unsloth_flex_inference_copy", None) + + engine_kwargs = dict(pending) + engine_kwargs["max_batch_size"] = int(max_batch_size) + + engine = FlexEngine( + hf_model = model, + tokenizer = tokenizer, + inference_model = inference_copy, + base_model = None, + peft_model = None, + **engine_kwargs, + ) + model._flex_engine_instance = engine + + # Drop the sentinel (plain attribute; setting replaces it). + model.vllm_engine = engine + model.fast_generate = engine.generate + model.fast_generate_batches = functools.partial(engine.generate, use_tqdm = False) + + # Consume the one-shot stashes: the deep-copy is owned by the engine now, + # and the needs-dict's role as "build spec" is over. + for _attr in ( + "_unsloth_flex_inference_copy", + "_unsloth_flex_tokenizer", + "_unsloth_needs_flex_engine", + ): + if hasattr(model, _attr): + try: + delattr(model, _attr) + except AttributeError: + pass + return engine + + +def build_flex_engine(model, max_batch_size: Optional[int] = None): + """Construct or return the FlexEngine attached to ``model``. + + Called lazily: by the RL patch (:func:`_build_flex_from_args`) once + ``GRPOTrainer.args`` is resolved, or by ``model.fast_generate`` on + first use. + + ``max_batch_size`` resolution (first build): + - ``floor = model._unsloth_needs_flex_engine['max_batch_size']`` + - ``effective = max(floor, max_batch_size or 0)`` + - A warning is emitted when ``effective > floor`` so users see the + GRPO-driven bump. + + Once the engine is built, it is the sole source of truth for the + batch-size dimension. Subsequent calls are idempotent when the + requested size fits; requesting a larger size raises + :class:`RuntimeError` because the engine's fixed-shape GPU buffers + and captured CUDA graphs cannot be grown in place. + """ + existing = getattr(model, "_flex_engine_instance", None) + pending = getattr(model, "_unsloth_needs_flex_engine", None) + + # Non-flex model (plain HF or plain vLLM). No-op so ``rl.py``'s patch + # stays unconditional. + if existing is None and pending is None: + return None + + requested = int(max_batch_size) if max_batch_size else 0 + + if existing is not None: + if requested <= existing.max_batch_size: + return existing + raise RuntimeError( + f"Unsloth: FlexEngine was built at max_batch_size=" + f"{existing.max_batch_size}; cannot grow to {requested} after " + f"construction (fixed-shape GPU page tables + CUDA graphs). " + f"Pass max_batch_size={requested} to " + f"FastLanguageModel.from_pretrained before the first " + f"fast_generate / GRPOTrainer call." + ) + + floor = int(pending["max_batch_size"]) + target = max(floor, requested) + if target > floor: + warnings.warn( + f"Unsloth: increasing FlexEngine max_batch_size {floor} -> " + f"{target} to fit the GRPO rollout batch " + f"(per_device_train_batch_size * steps_per_generation * " + f"num_generations). Pass max_batch_size={target} to " + f"FastLanguageModel.from_pretrained to silence this warning.", + stacklevel = 2, + ) + return _construct_and_attach(model, target) + + +def _build_flex_from_args(model, args): + """Helper used by the ``rl.py`` GRPOTrainer-init patch. + + Reads the rollout batch size from the TRL args and triggers the + FlexEngine build. No-op when ``model`` wasn't loaded through the + flex-inference path (plain vLLM / plain HF). + """ + if not hasattr(model, "_unsloth_needs_flex_engine") and not hasattr( + model, "_flex_engine_instance" + ): + return None + pdbs = int(getattr(args, "per_device_train_batch_size", 1) or 1) + spg = int( + getattr(args, "steps_per_generation", None) + or getattr(args, "gradient_accumulation_steps", 1) + or 1 + ) + ngen = int(getattr(args, "num_generations", 1) or 1) + # Written as ``max(A, B)`` for reviewer clarity. Reduces to the second + # term whenever ``num_generations >= 1`` (always). + grpo_target = max(pdbs * spg, pdbs * spg * ngen) + return build_flex_engine(model, max_batch_size = grpo_target) + + +__all__ = [ + "FlexEngine", + "load_flex", + "build_flex_engine", + "install_flex_sentinel", + "_build_flex_from_args", +] diff --git a/unsloth/inference/flex_gemma4.py b/unsloth/inference/flex_gemma4.py index 268181af30..c75a1f0f4a 100644 --- a/unsloth/inference/flex_gemma4.py +++ b/unsloth/inference/flex_gemma4.py @@ -430,6 +430,7 @@ class FlexGemma4Inference: fa4_prefill = None, base_model = None, peft_model = None, + cumem_allocator = None, ): assert max_seq_length % page_size == 0 self.model = model @@ -477,14 +478,19 @@ class FlexGemma4Inference: base_prefill_opts["BACKEND"] = "FLASH" self.prefill_kernel_options = base_prefill_opts - self.page_table = PageTable( - n_pages = n_pages, - page_size = page_size, - max_batch_size = max_batch_size, - device = self.device.type, - ) + # See ``FlexInference.__init__`` for why only the paged-KV + # allocations go through the cuMem ``kv_cache`` pool. + from .sleep_mode import kv_cache_pool as _kv_cache_pool - patch_gemma4_attention_forwards(model, self.page_table) + with _kv_cache_pool(cumem_allocator): + self.page_table = PageTable( + n_pages = n_pages, + page_size = page_size, + max_batch_size = max_batch_size, + device = self.device.type, + ) + + patch_gemma4_attention_forwards(model, self.page_table) self.input_pos_buffer = torch.zeros( max_batch_size, dtype = torch.int32, device = self.device diff --git a/unsloth/inference/flex_qwen3_llama.py b/unsloth/inference/flex_qwen3_llama.py index 272549ceae..4d38cd2cbe 100644 --- a/unsloth/inference/flex_qwen3_llama.py +++ b/unsloth/inference/flex_qwen3_llama.py @@ -499,6 +499,7 @@ class FlexInference: fa4_prefill = None, base_model = None, peft_model = None, + cumem_allocator = None, ): assert max_seq_length % page_size == 0 self.model = model @@ -558,13 +559,22 @@ class FlexInference: base_prefill_opts["BACKEND"] = "FLASH" self.prefill_kernel_options = base_prefill_opts - self.page_table = PageTable( - n_pages = n_pages, - page_size = page_size, - max_batch_size = max_batch_size, - device = self.device.type, - ) - patch_model_attention_forwards(model, self.page_table) + # Route the paged-KV allocations through the cuMem "kv_cache" + # pool when sleep mode is active, so FlexEngine.sleep can drop + # them and wake_up can re-map fresh zeroed pages at the same + # virtual addresses. The block_mask / input_pos scratch below + # stays in the default allocator so captured CUDA graphs that + # reference them stay valid across sleep / wake. + from .sleep_mode import kv_cache_pool as _kv_cache_pool + + with _kv_cache_pool(cumem_allocator): + self.page_table = PageTable( + n_pages = n_pages, + page_size = page_size, + max_batch_size = max_batch_size, + device = self.device.type, + ) + patch_model_attention_forwards(model, self.page_table) # Pre-allocated decode state. self.input_pos_buffer = torch.zeros( diff --git a/unsloth/inference/sleep_mode.py b/unsloth/inference/sleep_mode.py new file mode 100644 index 0000000000..92b0f43ba0 --- /dev/null +++ b/unsloth/inference/sleep_mode.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Sleep-mode helpers for the flex inference backend. + +When ``UNSLOTH_VLLM_STANDBY=1`` is set and vLLM's ``CuMemAllocator`` is +importable, :class:`~unsloth.inference.flex_engine.FlexEngine` routes its +heavy GPU allocations (the inference deep-copies, the PEFT wrapper, and +per-layer :class:`~unsloth.inference.flex_paged_attention.PagedKVCache` +buffers) through cuMem-backed pools. ``sleep(level=1)`` then offloads the +``weights`` pool to pinned CPU memory and discards the ``kv_cache`` pool +outright; ``wake_up`` re-maps the handles at the same virtual addresses +so previously captured CUDA graphs and compiled artifacts stay valid. + +If vLLM is not installed, :func:`_get_cumem_allocator` returns ``None`` +and the helpers fall back to :func:`contextlib.nullcontext`, which keeps +the flex engine working without sleep support. +""" + +from __future__ import annotations + +import contextlib +import os +import warnings +from typing import Any, Optional + + +_WARNED_NO_VLLM = False + + +def _get_cumem_allocator() -> Optional[Any]: + """Return ``vllm.device_allocator.cumem.CuMemAllocator.get_instance()`` + or ``None`` if vLLM is not importable. + + Emits a single warning on the first failed import so users who opt + into sleep mode with ``UNSLOTH_VLLM_STANDBY=1`` see a clear message + about the soft dependency on vLLM.""" + global _WARNED_NO_VLLM + try: + from vllm.device_allocator.cumem import CuMemAllocator + except Exception as e: + if not _WARNED_NO_VLLM: + warnings.warn( + "FlexEngine sleep mode requires vLLM's CuMemAllocator " + f"(import failed: {e}). Sleep / wake_up will be no-ops. " + "Install vLLM to enable level-1 sleep mode on the flex " + "backend.", + RuntimeWarning, + stacklevel = 2, + ) + _WARNED_NO_VLLM = True + return None + return CuMemAllocator.get_instance() + + +def sleep_mode_enabled() -> bool: + """``True`` iff ``UNSLOTH_VLLM_STANDBY=1`` is set AND vLLM is + available. Evaluated at :class:`FlexEngine.__init__` time so the + choice of allocator is stable for the engine's lifetime.""" + if os.environ.get("UNSLOTH_VLLM_STANDBY", "0") != "1": + return False + return _get_cumem_allocator() is not None + + +def _pool(allocator: Optional[Any], tag: str): + """Return a ``use_memory_pool(tag=...)`` context manager if the + allocator is available, else ``nullcontext`` so callers can wrap + allocation sites unconditionally.""" + if allocator is None: + return contextlib.nullcontext() + return allocator.use_memory_pool(tag = tag) + + +def weight_pool(allocator: Optional[Any]): + """Context manager for ``tag="weights"`` allocations (offloaded to + pinned CPU on sleep, restored via ``cudaMemcpy`` on wake).""" + return _pool(allocator, "weights") + + +def kv_cache_pool(allocator: Optional[Any]): + """Context manager for ``tag="kv_cache"`` allocations (discarded on + sleep, re-mapped with zeros on wake).""" + return _pool(allocator, "kv_cache") + + +def describe_sleep_state(engine) -> dict: + """Return a small dict summarising the flex engine's sleep-mode + state plus a torch CUDA memory snapshot. Useful from bench scripts + and tests.""" + import torch + + allocator = getattr(engine, "_cumem_allocator", None) + state: dict = { + "sleep_mode_enabled": bool(getattr(engine, "_sleep_mode_enabled", False)), + "cumem_allocator": type(allocator).__name__ if allocator is not None else None, + } + if torch.cuda.is_available(): + state["allocated_gb"] = round(torch.cuda.memory_allocated() / 1e9, 3) + state["reserved_gb"] = round(torch.cuda.memory_reserved() / 1e9, 3) + if allocator is not None and hasattr(allocator, "get_current_usage"): + try: + state["cumem_current_usage"] = allocator.get_current_usage() + except Exception: + pass + return state + + +__all__ = [ + "_get_cumem_allocator", + "sleep_mode_enabled", + "weight_pool", + "kv_cache_pool", + "describe_sleep_state", +] diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 55edbc59e1..238c78557b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2526,10 +2526,28 @@ class FastLlamaModel: # flex inference copy uses the original transformers # rotary / QKV layout. A later block (after the tokenizer is # loaded) constructs the FlexEngine around this copy. + # + # If ``UNSLOTH_VLLM_STANDBY=1`` is set AND vLLM is + # importable, wrap the deep-copy in the cuMem + # ``weights`` pool so ``FlexEngine.sleep(level=1)`` can + # offload it to pinned CPU memory without invalidating + # the engine's captured CUDA graphs. This must happen at + # the deep-copy site, not later in FlexEngine.__init__, + # because the engine receives the copy as + # ``inference_model=`` and never re-allocates it. import copy as _copy + from unsloth.inference.sleep_mode import ( + _get_cumem_allocator as _flex_get_cumem, + sleep_mode_enabled as _flex_sleep_enabled, + weight_pool as _flex_weight_pool, + ) - model._unsloth_flex_inference_copy = _copy.deepcopy(model) - model._unsloth_flex_inference_copy.eval() + _flex_allocator = ( + _flex_get_cumem() if _flex_sleep_enabled() else None + ) + with _flex_weight_pool(_flex_allocator): + model._unsloth_flex_inference_copy = _copy.deepcopy(model) + model._unsloth_flex_inference_copy.eval() model._unsloth_needs_flex_engine = dict( dtype = dtype, max_seq_length = max_seq_length, @@ -2617,31 +2635,19 @@ class FastLlamaModel: model, tokenizer, correct_dtype = dtype ) - # UNSLOTH_FAST_INFERENCE=1 path: build the FlexEngine now that the - # tokenizer is available. The HF model itself is NOT flex-patched -- + # UNSLOTH_FAST_INFERENCE=1 path: install the lazy ``vllm_engine`` + # sentinel now that the tokenizer is available. The real + # ``FlexEngine(...)`` construction is deferred to + # :func:`build_flex_engine` so the batch-size dimension can be + # sized from the GRPO rollout shape instead of frozen here at a + # guessed default. The HF model itself is NOT flex-patched -- # the engine owns its own pre-patched deep-copy (captured above # before Unsloth's post_patch), so `model.forward` (used by the # training loop) stays intact. - _flex_args = getattr(model, "_unsloth_needs_flex_engine", None) - if _flex_args is not None: - del model._unsloth_needs_flex_engine - _inference_copy = model._unsloth_flex_inference_copy - del model._unsloth_flex_inference_copy - from unsloth.inference.flex_engine import FlexEngine + if hasattr(model, "_unsloth_needs_flex_engine"): + from unsloth.inference.flex_engine import install_flex_sentinel - flex_engine = FlexEngine( - hf_model = model, - tokenizer = tokenizer, - inference_model = _inference_copy, - base_model = None, - peft_model = None, - **_flex_args, - ) - model.vllm_engine = flex_engine - model.fast_generate = flex_engine.generate - model.fast_generate_batches = functools.partial( - flex_engine.generate, use_tqdm = False - ) + install_flex_sentinel(model, tokenizer) # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers): diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index ac9b35a822..f3be2dbdeb 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1752,9 +1752,22 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import 'GuidedDecodingParams(backend="outlines", regex=args.vllm_guided_decoding_regex) ' 'if getattr(args, "vllm_guided_decoding_regex", None) is not None else None,', ) - # Replace with our vLLM engine + # Replace with our vLLM engine. + # + # ``_build_flex_from_args`` sizes the FlexEngine's fixed-shape + # page tables / CUDA-graph buckets from the GRPO rollout batch + # (per_device_train_batch_size * steps_per_generation * + # num_generations) BEFORE the first ``model.vllm_engine`` + # access triggers the lazy build. No-op when the model wasn't + # loaded through the flex-inference path (plain vLLM / plain + # HF), so this injection is safe for every backend. The + # import is inlined on the same line so ``create_new_function`` + # doesn't need a cross-module import entry. sampling_params = ( " " * 12 + + "from unsloth.inference.flex_engine import " + + "_build_flex_from_args as __unsloth_build_flex_from_args; " + + "__unsloth_build_flex_from_args(model, args); " + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + sampling_params ) # Add spaces @@ -1790,9 +1803,18 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import ) if trl_version >= Version("0.18.0"): - # Replace LLM init with already existing vLLM engine for colocate mode + # Replace LLM init with already existing vLLM engine for colocate mode. + # Prepend the FlexEngine build call so the engine is sized from the + # GRPO rollout batch before ``model.vllm_engine`` is dereferenced. + # Inline the import so ``create_new_function`` doesn't need a + # cross-module entry; the helper is a no-op on non-flex models. vllm_llm_init_pattern = r"self\.llm\s*=\s*LLM\(.*?\)*\)\s*?\n(?!,)" - vllm_llm_replacement = "self.llm = model.vllm_engine\n" + vllm_llm_replacement = ( + "from unsloth.inference.flex_engine import " + "_build_flex_from_args as __unsloth_build_flex_from_args; " + "__unsloth_build_flex_from_args(model, args); " + "self.llm = model.vllm_engine\n" + ) new_vllm_part = re.sub( vllm_llm_init_pattern, vllm_llm_replacement, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 0e44cca711..93a1069c3f 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1248,21 +1248,17 @@ class FastBaseModel: raise _patch_err model = post_patch_loss_function(model) - # UNSLOTH_FAST_INFERENCE=1 path: build the FlexEngine here so that - # the tokenizer/processor is available. Keeps the training model's + # UNSLOTH_FAST_INFERENCE=1 path: install the lazy ``vllm_engine`` + # sentinel now that the tokenizer/processor is available. The real + # ``FlexEngine(...)`` construction is deferred to + # :func:`build_flex_engine` so the batch-size dimension can be + # sized from the GRPO rollout shape. Keeps the training model's # forward intact — the engine carries its own deep-copy. - _flex_args = getattr(model, "_unsloth_needs_flex_engine", None) - if _flex_args is not None: - del model._unsloth_needs_flex_engine - from unsloth.inference.flex_engine import load_flex + if hasattr(model, "_unsloth_needs_flex_engine"): + from unsloth.inference.flex_engine import install_flex_sentinel _tok_for_flex = getattr(tokenizer, "tokenizer", tokenizer) - flex_engine = load_flex(model, _tok_for_flex, **_flex_args) - model.vllm_engine = flex_engine - model.fast_generate = flex_engine.generate - model.fast_generate_batches = functools.partial( - flex_engine.generate, use_tqdm = False - ) + install_flex_sentinel(model, _tok_for_flex) # Log Unsloth version for future fastpaths for inference if hasattr(model, "config"):