ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-04-21 12:00:15 +00:00
commit 5e1ec3395a
6 changed files with 94 additions and 89 deletions
Download patch file Download diff file Expand all files Collapse all files

54
tests/flex_fastlm_bench.py
View file

@ -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__":

53
tests/flex_fastlm_smoke.py
View file

@ -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}")

30
unsloth/inference/flex_engine.py
View file

@ -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(

34
unsloth/inference/flex_gemma4.py
View file

@ -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))

4
unsloth/inference/vllm_shim.py
View file

@ -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)

8
unsloth/models/loader.py
View file

@ -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.