diff --git a/tests/flex_fastlm_bench.py b/tests/flex_fastlm_bench.py index 1f1ea7fa3f..c7c3d152dd 100644 --- a/tests/flex_fastlm_bench.py +++ b/tests/flex_fastlm_bench.py @@ -5,6 +5,7 @@ numbers for the same workload.""" import os, sys, 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)) @@ -15,13 +16,13 @@ import torch 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("--n_prompts", type=int, default=8) - p.add_argument("--max_new_tokens", type=int, default=64) - p.add_argument("--max_batch_size", type=int, default=16) - p.add_argument("--max_seq_length", type=int, default=1024) - p.add_argument("--n_rounds", type=int, default=3) + p.add_argument("--model", default = "unsloth/Qwen3-4B-Base") + p.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16") + p.add_argument("--n_prompts", type = int, default = 8) + p.add_argument("--max_new_tokens", type = int, default = 64) + p.add_argument("--max_batch_size", type = int, default = 16) + p.add_argument("--max_seq_length", type = int, default = 1024) + p.add_argument("--n_rounds", type = int, default = 3) args = p.parse_args() os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") @@ -31,18 +32,27 @@ def main(): dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 model, tok = FastLanguageModel.from_pretrained( - model_name=args.model, - max_seq_length=args.max_seq_length, - dtype=dtype, - load_in_4bit=False, - fast_inference=True, - max_batch_size=args.max_batch_size, + model_name = args.model, + max_seq_length = args.max_seq_length, + dtype = dtype, + load_in_4bit = False, + fast_inference = True, + max_batch_size = args.max_batch_size, ) print(f"[bench] model={args.model} dtype={args.dtype}") - prompts = [f"In one sentence, a fact about {t} is" for t in [ - "the moon", "gravity", "the ocean", "the sun", "honey", - "rain", "trees", "mountains" - ][:args.n_prompts]] + prompts = [ + f"In one sentence, a fact about {t} is" + for t in [ + "the moon", + "gravity", + "the ocean", + "the sun", + "honey", + "rain", + "trees", + "mountains", + ][: args.n_prompts] + ] class _SP: max_tokens = args.max_new_tokens @@ -51,7 +61,7 @@ def main(): # Warmup round — captures CUDA graphs. print("[bench] warmup (CUDA graph capture)...") t0 = time.perf_counter() - _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + _ = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) print(f"[bench] warmup wall: {time.perf_counter() - t0:.2f}s") walls = [] @@ -59,7 +69,7 @@ def main(): for r in range(args.n_rounds): torch.cuda.synchronize() t1 = time.perf_counter() - outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + outs = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) torch.cuda.synchronize() dt = time.perf_counter() - t1 n_tok = sum(len(o.outputs[0].token_ids) for o in outs) @@ -70,8 +80,10 @@ def main(): if walls: wall_med = sorted(walls)[len(walls) // 2] tok_med = tok_counts[len(walls) // 2] - print(f"[bench] median: {tok_med} toks in {wall_med:.2f}s " - f"=> {tok_med / wall_med:.1f} tok/s") + print( + f"[bench] median: {tok_med} toks in {wall_med:.2f}s " + f"=> {tok_med / wall_med:.1f} tok/s" + ) if __name__ == "__main__": diff --git a/tests/flex_fastlm_smoke.py b/tests/flex_fastlm_smoke.py index 05b620552c..1bea61dc52 100644 --- a/tests/flex_fastlm_smoke.py +++ b/tests/flex_fastlm_smoke.py @@ -24,13 +24,13 @@ if str(_REPO_ROOT) not in sys.path: 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("--with_lora", 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("--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("--with_lora", 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") args = p.parse_args() import torch @@ -39,6 +39,7 @@ def main(): print(f"[smoke] UNSLOTH_FAST_INFERENCE={os.environ.get('UNSLOTH_FAST_INFERENCE')}") import unsloth + print(f"[smoke] unsloth={unsloth.__file__}") from unsloth import FastLanguageModel @@ -46,11 +47,11 @@ def main(): 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, + model_name = args.model, + max_seq_length = args.max_seq_length, + dtype = dtype, + load_in_4bit = args.load_in_4bit, + fast_inference = True, ) t_load = time.perf_counter() - t0 print(f"[smoke] loaded model in {t_load:.1f}s; dtype={model.dtype}") @@ -60,31 +61,35 @@ def main(): if args.with_lora: model = FastLanguageModel.get_peft_model( model, - r=16, - target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], - lora_alpha=16, - lora_dropout=0.0, - bias="none", - use_gradient_checkpointing="unsloth", - random_state=3407, + r = 16, + target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"], + lora_alpha = 16, + lora_dropout = 0.0, + bias = "none", + use_gradient_checkpointing = "unsloth", + random_state = 3407, ) print(f"[smoke] PEFT model type: {type(model).__name__}") - print(f"[smoke] model.vllm_engine bound to PEFT: " - f"{hasattr(model, 'vllm_engine')}") + print( + f"[smoke] model.vllm_engine bound to PEFT: " + f"{hasattr(model, 'vllm_engine')}" + ) from unsloth.inference.vllm_shim import LoRARequest + prompts = [args.prompt] + # Minimal SamplingParams stand-in class _SP: max_tokens = args.max_new_tokens temperature = 0.0 + t1 = time.perf_counter() - outputs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + outputs = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) dt = time.perf_counter() - t1 out = outputs[0] n_tok = len(out.outputs[0].token_ids) - print(f"[smoke] generated {n_tok} tokens in {dt:.2f}s " - f"({n_tok / dt:.1f} tok/s)") + print(f"[smoke] generated {n_tok} tokens in {dt:.2f}s " f"({n_tok / dt:.1f} tok/s)") print(f"[smoke] prompt: {args.prompt!r}") print(f"[smoke] completion: {out.outputs[0].text!r}") diff --git a/unsloth/inference/flex_engine.py b/unsloth/inference/flex_engine.py index b17eb39d9c..5c83735460 100644 --- a/unsloth/inference/flex_engine.py +++ b/unsloth/inference/flex_engine.py @@ -72,7 +72,7 @@ def _fa4_ok_for_head_dim(head_dim: int, device_cap: tuple[int, int]) -> bool: return False if major == 9: return 8 <= head_dim <= 256 # Hopper - return 8 <= head_dim <= 128 # Blackwell (sm_100 / sm_120) + return 8 <= head_dim <= 128 # Blackwell (sm_100 / sm_120) def _triton_block_defaults( @@ -131,7 +131,11 @@ def _auto_kernel_options( ``'NoneType' object is not subscriptable``). Users who want FA4 can pass ``fa4_prefill=True`` explicitly to the engine after confirming it works on their workload.""" - cap = torch.cuda.get_device_capability(device) if torch.cuda.is_available() else (0, 0) + cap = ( + torch.cuda.get_device_capability(device) + if torch.cuda.is_available() + else (0, 0) + ) head_dims = _collect_head_dims(hf_model) or [128] head_dim_max = max(head_dims) @@ -285,9 +289,10 @@ class FlexEngine: peft_model = None, inference_model = None, ): - assert dtype in (torch.bfloat16, torch.float16), ( - f"FlexEngine requires bf16 or fp16 dtype; got {dtype}." - ) + assert dtype in ( + torch.bfloat16, + torch.float16, + ), f"FlexEngine requires bf16 or fp16 dtype; got {dtype}." self.hf_model = hf_model self.tokenizer = tokenizer self.compute_dtype = dtype @@ -335,7 +340,9 @@ class FlexEngine: ) # Size n_pages from available VRAM so big prompts don't OOM. - n_pages = self._compute_n_pages(gpu_memory_utilization, max_batch_size, page_size) + n_pages = self._compute_n_pages( + gpu_memory_utilization, max_batch_size, page_size + ) arch = _detect_arch(inference_model) self.arch = arch @@ -506,9 +513,7 @@ class FlexEngine: if candidate in target_sd: renamed[candidate] = v if renamed: - missing, unexpected = peft_model.load_state_dict( - renamed, strict = False - ) + missing, unexpected = peft_model.load_state_dict(renamed, strict = False) # ``missing`` will be every non-LoRA param; that's fine. refresh_lora_merge_from_pristine(base_model, peft_model) self._current_lora_int_id = getattr(lora_request, "lora_int_id", None) @@ -545,7 +550,8 @@ class FlexEngine: ) elif "prompt" in p: seq = Sequence( - text = p["prompt"], max_new_tokens = max_new_tokens, + text = p["prompt"], + max_new_tokens = max_new_tokens, ) else: raise ValueError( @@ -657,9 +663,7 @@ 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 = _get_peft_model(self._inference_model, peft_cfg) self._inference_peft.eval() except Exception as e: warnings.warn( diff --git a/unsloth/inference/flex_gemma4.py b/unsloth/inference/flex_gemma4.py index 42664c1469..87fa9d2ec7 100644 --- a/unsloth/inference/flex_gemma4.py +++ b/unsloth/inference/flex_gemma4.py @@ -105,7 +105,9 @@ from torch.nn.attention.flex_attention import create_block_mask as _create_block # (local to this file so that file is untouched). -def _causal_blockmask_with_window(B: int, L: int, block_size: int, window: int, device: str): +def _causal_blockmask_with_window( + B: int, L: int, block_size: int, window: int, device: str +): def causal_windowed(b, h, q_idx, kv_idx): return (q_idx >= kv_idx) & (q_idx - kv_idx < window) @@ -297,9 +299,7 @@ def make_flex_gemma4_attention_forward(page_table: PageTable): return forward -def patch_gemma4_attention_forwards( - model: torch.nn.Module, page_table: PageTable -): +def patch_gemma4_attention_forwards(model: torch.nn.Module, page_table: PageTable): """Attach a PagedKVCache to every non-shared attention layer, link every shared attention layer to its store layer's cache, and swap in the flex_attention forward above. @@ -341,9 +341,7 @@ def patch_gemma4_attention_forwards( # --- model forward walker -------------------------------------------------- -def call_gemma4_model_with_flex_kwargs( - model, input_ids, position_ids, flex_kwargs -): +def call_gemma4_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs): """Walk the Gemma-4 text model manually so we can inject flex_* kwargs into each attention call. Mirrors `call_model_with_flex_kwargs` in `qwen3_flex_inference.py` but: @@ -609,9 +607,9 @@ class FlexGemma4Inference: kv_num_blocks = block_mask.kv_num_blocks[ batch_idx, :, input_block_idx ].view(B, 1, 1) - kv_indices = block_mask.kv_indices[ - batch_idx, :, input_block_idx - ].view(B, 1, 1, -1) + kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view( + B, 1, 1, -1 + ) full_num = full_idx = None if block_mask.full_kv_num_blocks is not None: full_num = block_mask.full_kv_num_blocks[ @@ -640,9 +638,7 @@ class FlexGemma4Inference: def causal_offset_windowed(off, window): def m(b, h, q_idx, kv_idx): - return (q_idx + off[b] >= kv_idx) & ( - q_idx + off[b] - kv_idx < window - ) + return (q_idx + off[b] >= kv_idx) & (q_idx + off[b] - kv_idx < window) return m @@ -711,9 +707,7 @@ class FlexGemma4Inference: allocated = self.page_table.allocate() self.page_table.reserve( allocated, - torch.tensor( - [allocated], device = self.device, dtype = torch.long - ), + torch.tensor([allocated], device = self.device, dtype = torch.long), self.page_size, ) reserved_batches.append(allocated) @@ -1041,9 +1035,7 @@ def main(): "--verify_no_drift only applies to the bf16 double-copy path." ) if args.no_merge_lora: - raise SystemExit( - "--verify_no_drift is incompatible with --no_merge_lora." - ) + raise SystemExit("--verify_no_drift is incompatible with --no_merge_lora.") if base_model is None or peft_model is None: raise SystemExit( "--verify_no_drift requires --lora_adapter against the bf16 path." @@ -1160,9 +1152,7 @@ def main(): "peak_memory_gb": peak, "sample_completions": sample_completions, } - os.makedirs( - os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True - ) + os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True) with open(args.stats_path, "w") as f: json.dump(res, f, indent = 2) print(json.dumps(res, indent = 2)) diff --git a/unsloth/inference/vllm_shim.py b/unsloth/inference/vllm_shim.py index 16b7025d3b..19a37b51f4 100644 --- a/unsloth/inference/vllm_shim.py +++ b/unsloth/inference/vllm_shim.py @@ -129,9 +129,7 @@ def load_lora( if load_tensors: model.peft_config["default"].save_pretrained(save_directory) elif not os.path.exists(save_directory): - raise OSError( - f"Unsloth: LoRA filepath = {save_directory} does not exist!" - ) + raise OSError(f"Unsloth: LoRA filepath = {save_directory} does not exist!") if load_tensors: peft_config = _get_peft_config(save_directory) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index ebd3ae8f9e..3165855534 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -353,9 +353,7 @@ class FastLanguageModel(FastLlamaModel): or dtype == torch.float32 ) - _use_flex_fast_inference = ( - os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" - ) + _use_flex_fast_inference = os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" if fast_inference and _use_flex_fast_inference: # Flex backend path: skip the vLLM import gate entirely. The # actual engine is attached further down in @@ -987,9 +985,7 @@ class FastModel(FastBaseModel): ) load_in_4bit = False - _use_flex_fast_inference = ( - os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" - ) + _use_flex_fast_inference = os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" if fast_inference and _use_flex_fast_inference: # Flex backend path: skip the vLLM import gate. The engine is # attached further down in the ``FastBaseModel`` loader.