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.
49acbdf6bd