diff --git a/tests/flex_lazy_batch_smoke.py b/tests/flex_lazy_batch_smoke.py index d56ce51f7f..87b3fcf917 100644 --- a/tests/flex_lazy_batch_smoke.py +++ b/tests/flex_lazy_batch_smoke.py @@ -106,9 +106,9 @@ def _case1_default_floor(): 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" - ) + assert not hasattr( + model, "_flex_engine_instance" + ), "engine should NOT exist before first use" out = model.fast_generate(["hello"]) assert out == [("stub", ["hello"])] diff --git a/tests/flex_lazy_live_smoke.py b/tests/flex_lazy_live_smoke.py index 8d7166c828..093aceb64c 100644 --- a/tests/flex_lazy_live_smoke.py +++ b/tests/flex_lazy_live_smoke.py @@ -38,9 +38,9 @@ if str(_REPO_ROOT) not in sys.path: def main(): - assert os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1", ( - "export UNSLOTH_FAST_INFERENCE=1 before running this smoke" - ) + 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 ( @@ -49,9 +49,7 @@ def main(): build_flex_engine, ) - model_name = os.environ.get( - "FLEX_LAZY_SMOKE_MODEL", "unsloth/Qwen3-0.6B-Base" - ) + 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, @@ -62,12 +60,12 @@ def main(): 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" - ) + 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) @@ -75,7 +73,9 @@ def main(): 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}") + 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) diff --git a/tests/flex_sleep_mode_smoke.py b/tests/flex_sleep_mode_smoke.py index 5744df85da..d206eae773 100644 --- a/tests/flex_sleep_mode_smoke.py +++ b/tests/flex_sleep_mode_smoke.py @@ -155,7 +155,9 @@ def main(): out1 = engine.generate( prompts, sampling_params = type( - "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0}, + "SP", + (), + {"max_tokens": args.max_new_tokens, "temperature": 0.0}, )(), ) print( @@ -199,9 +201,7 @@ def main(): 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}" - ) + 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"] @@ -226,9 +226,7 @@ def main(): # 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"] - ), ( + assert probe_post["allocated_gb"] == probe_pre["allocated_gb"], ( "With sleep mode disabled, torch.memory_allocated must " "be unchanged by sleep()" ) @@ -238,9 +236,7 @@ def main(): 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}" - ) + print(f"[sleep-smoke] wake_up() took {t_wake:.3f}s; {probe_wake}") post_sums = _checksum_params(engine._inference_model) diffs = [] @@ -265,7 +261,9 @@ def main(): out2 = engine.generate( prompts, sampling_params = type( - "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0}, + "SP", + (), + {"max_tokens": args.max_new_tokens, "temperature": 0.0}, )(), ) t_regen = time.perf_counter() - t0 diff --git a/unsloth/inference/flex_engine.py b/unsloth/inference/flex_engine.py index b566241d2c..8c220cfd26 100644 --- a/unsloth/inference/flex_engine.py +++ b/unsloth/inference/flex_engine.py @@ -680,9 +680,7 @@ class FlexEngine: 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}" - ) + 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 " @@ -769,7 +767,8 @@ class FlexEngine: # attention-forward-agnostic; it wraps Linear modules. with weight_pool(self._cumem_allocator): self._inference_peft = _get_peft_model( - self._inference_model, peft_cfg, + self._inference_model, + peft_cfg, ) self._inference_peft.eval() except Exception as e: diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 238c78557b..93a1592a5e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2542,9 +2542,7 @@ class FastLlamaModel: weight_pool as _flex_weight_pool, ) - _flex_allocator = ( - _flex_get_cumem() if _flex_sleep_enabled() else None - ) + _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()