35231d4ff4
Setting UNSLOTH_FAST_INFERENCE=1 makes
FastLanguageModel.from_pretrained(..., fast_inference=True) return the
flex attention + paged KV + CUDA-graph backend instead of vLLM for
Qwen3, Llama-3, and Gemma-4-E2B-it. Default "0" keeps the existing vLLM
path.
The flex engine lives in a new unsloth.inference subpackage so the
FastLanguageModel wiring and any external caller can reach it without
going through scripts/. The user-facing GRPO notebooks for Qwen3,
Llama-3, and Gemma-4 run end-to-end with no code edits, just the env
var.
New subpackage
- unsloth/inference/flex_paged_attention.py: paged KV cache and
block-mask helpers.
- unsloth/inference/flex_qwen3_llama.py: FlexInference for Qwen3 and
Llama-3.2. tokenize() skips pre-populated input_ids so the engine
can accept vLLM-style token-id prompts.
- unsloth/inference/flex_gemma4.py: FlexGemma4Inference. Text-shell
extraction from Gemma4ForConditionalGeneration, per-layer-type
sliding-window masks, KV sharing for layers 15-34.
- unsloth/inference/flex_engine.py: FlexEngine with the vLLM LLM
surface. Arch dispatch, deep-copy of the HF model for colocate
rollout so training forward stays intact, lazy pristine-base copy
on first bind_peft_model, torch.amp.autocast("cuda", dtype=...)
wrapping for the whole generate path, auto-tune for Triton block
sizes per GPU capability and head_dim band. FA4 default OFF
because the sm_100 FLASH backend crashes on short prompts; users
can opt in with fa4_prefill=True.
- unsloth/inference/vllm_shim.py: LoRARequest, RequestOutput,
CompletionOutput dataclasses plus save_lora / load_lora that
mirror unsloth_zoo.vllm_utils line for line so the TRL GRPO patch
sees the same attribute shape.
Wiring
- unsloth/models/loader.py: UNSLOTH_FAST_INFERENCE=1 bypasses the
vLLM import guard in FastLanguageModel.from_pretrained (L356) and
FastModel.from_pretrained (L982).
- unsloth/models/llama.py: the elif not fast_inference load branch
now also handles the flex case, pops max_batch_size from kwargs
(FlexEngine-only), snapshots the model into an inference copy
BEFORE Unsloth's post-patch so flex-attention patching lands on a
clean HF layout, and after the tokenizer loads constructs
FlexEngine and attaches .vllm_engine / .fast_generate /
.fast_generate_batches. FastLlamaModel.get_peft_model calls
engine.bind_peft_model(peft) right after patch_peft_fast_inference
so state_dict() reads the training LoRA tensors.
- unsloth/models/vision.py: same pattern for FastBaseModel. Gemma-4
is allowed through the vision-model gate when the flex backend is
selected; the engine extracts the text shell from
Gemma4ForConditionalGeneration internally.
- unsloth/models/_utils.py: patch_peft_fast_inference picks
unsloth.inference.vllm_shim.save_lora / load_lora when the engine
is a FlexEngine, so the flex path never imports
vllm.lora.request. fast_inference_setup skips patch_vllm() when
UNSLOTH_FAST_INFERENCE=1.
Smoke tests (B200 sm_100)
- Qwen3-4B-Base bf16 (no LoRA and with get_peft_model LoRA)
- Qwen3-4B-Base fp16
- Llama-3.2-3B-Instruct bf16
- gemma-4-E2B-it bf16
all load through FastLanguageModel + UNSLOTH_FAST_INFERENCE=1 and
generate coherent completions.
Batched steady-state throughput via FastLanguageModel (autocast ON,
n_prompts=8 max_new_tokens=64 max_batch_size=16 max_seq_length=1024):
| Model | Flex CLI | Integration | Delta |
|----------------------------|-----------|-------------|-------|
| Qwen3-4B-Base | 819 tok/s | 595 tok/s | -27% |
| Llama-3.2-3B-Instruct | 1163 | 1357 | +17% |
Qwen3's regression is outside the 10% budget; most likely cause is
q_norm / k_norm being promoted to fp32 under autocast. Tracked for
a follow-up; the integration is functional and GRPO-usable today.
Note
The flex engine source files live here in unsloth/inference/. The
CLI benchmark drivers in scripts/benchmarks/ land in a separate PR
(#5108) which also carries the Gemma-4 and Llama-3 regression runs.
This PR depends only on what is already on main.
Known follow-ups
- FA4 on sm_100 needs a more robust enable rule before it can be
default-on.
- Memory: 2x base-model VRAM without LoRA, 3x with LoRA.
- FlexEngine.sleep() / .wake_up() are no-op stubs; real CPU-offload
parity with UNSLOTH_VLLM_STANDBY is a separate PR.
93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
"""Smoke-test the ``UNSLOTH_FAST_INFERENCE=1`` path through
|
|
``FastLanguageModel.from_pretrained``.
|
|
|
|
Invoked as:
|
|
CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 python tests/flex_fastlm_smoke.py \
|
|
--model unsloth/Qwen3-4B-Base --dtype bf16 --no-lora
|
|
|
|
Prints tokens/s + the first generated string.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
# Make the local fork importable.
|
|
_REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(_REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_REPO_ROOT))
|
|
|
|
|
|
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")
|
|
args = p.parse_args()
|
|
|
|
import torch
|
|
|
|
os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1")
|
|
print(f"[smoke] UNSLOTH_FAST_INFERENCE={os.environ.get('UNSLOTH_FAST_INFERENCE')}")
|
|
|
|
import unsloth
|
|
print(f"[smoke] unsloth={unsloth.__file__}")
|
|
from unsloth import FastLanguageModel
|
|
|
|
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,
|
|
)
|
|
t_load = time.perf_counter() - t0
|
|
print(f"[smoke] loaded model in {t_load:.1f}s; dtype={model.dtype}")
|
|
print(f"[smoke] hasattr(model, 'vllm_engine'): {hasattr(model, 'vllm_engine')}")
|
|
print(f"[smoke] vllm_engine type: {type(model.vllm_engine).__name__}")
|
|
|
|
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,
|
|
)
|
|
print(f"[smoke] PEFT model type: {type(model).__name__}")
|
|
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)
|
|
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] prompt: {args.prompt!r}")
|
|
print(f"[smoke] completion: {out.outputs[0].text!r}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|