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

flex: support --load_in_4bit with PEFT adapter (bnb-4bit shard)

Browse source 
Adds --load_in_4bit (+ --model_name_4bit override) to both the flex
benchmark script and the vllm/tpaged benchmark script. When set, loads the
pre-quantized Unsloth bnb-4bit shard (e.g.
unsloth/Qwen3-4B-Base-unsloth-bnb-4bit) and keeps the LoRA adapter as a
PEFT wrapper instead of merging, because merging into 4-bit weights is not
supported.

Ties lm_head.weight to model.embed_tokens.weight post-load in both scripts,
because the bnb-4bit shards ship without an lm_head parameter even though
tie_word_embeddings is True in the config, so transformers leaves it
randomly initialised otherwise (garbage generations).

Results at batch 64 + LoRA rank 32:

| Backend                       | tok/s | peak mem | output    |
|-------------------------------|------:|---------:|-----------|
| Unsloth fast_inference (vLLM) |  4515 | 159 GB   | coherent  |
| flex (this PR)                |  1738 |  40.6 GB | coherent  |
| transformers CB (sdpa)        |   504 | 124 GB   | gibberish |

4-bit costs ~40 % throughput on the vLLM path vs bf16 and ~70 % on flex.
flex regresses worse because PEFT-without-merge doubles the matmuls per
projection (base + LoRA add) on top of bnb dequant, whereas bf16 flex
merges LoRA into the base. Peak memory barely moves for vLLM because KV
cache at gpu_memory_utilization=0.8 dominates regardless of base size.

transformers CB (generate_batch) at 4-bit + LoRA produces garbage even
with lm_head tied. Likely PEFT-over-bnb + batched CB interaction; not
debugged further -- it was always the 10 % reference path.

Writeup updated in scripts/benchmarks/results/flex_vs_vllm.md with a new
"Same workload at load_in_4bit=True" section.
This commit is contained in:
Daniel Han 2026-04-21 02:13:06 +00:00
commit 4717bce97e
8 changed files with 248 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

37
scripts/benchmarks/cb_vs_vllm_generation.py
View file

@ -77,7 +77,7 @@ def run_vllm(args):
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = args.model_name,
max_seq_length = args.max_seq_length,
load_in_4bit = False,
load_in_4bit = args.load_in_4bit,
fast_inference = True,
max_lora_rank = 32,
gpu_memory_utilization = args.gpu_memory_utilization,
@ -155,11 +155,26 @@ def run_tpaged(args):
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
args.model_name,
dtype = torch.bfloat16,
attn_implementation = args.attn_impl,
).to("cuda")
if args.load_in_4bit:
bnb_model_name = args.model_name_4bit or f"{args.model_name}-unsloth-bnb-4bit"
print(f"[tpaged] loading 4-bit base: {bnb_model_name}")
model = AutoModelForCausalLM.from_pretrained(
bnb_model_name,
attn_implementation = args.attn_impl,
device_map = "cuda:0",
)
# HF transformers logs "lm_head.weight newly initialized" for
# bnb-4bit shards of tied-embedding models. tie_word_embeddings is
# True in the config but the dequant path leaves lm_head unbound.
# Tie manually so we don't generate gibberish.
if getattr(model.config, "tie_word_embeddings", False):
model.lm_head.weight = model.model.embed_tokens.weight
else:
model = AutoModelForCausalLM.from_pretrained(
args.model_name,
dtype = torch.bfloat16,
attn_implementation = args.attn_impl,
).to("cuda")
model.eval()
if args.lora_adapter:
@ -422,6 +437,16 @@ def parse_args():
default = None,
help = "Path to a PEFT adapter (rank 32) applied in every backend.",
)
p.add_argument(
"--load_in_4bit",
action = "store_true",
help = "Load base as bitsandbytes 4-bit (Unsloth shard).",
)
p.add_argument(
"--model_name_4bit",
default = None,
help = "Override 4-bit shard name. Default `{model_name}-unsloth-bnb-4bit`.",
)
p.add_argument("--temperature", type = float, default = 0.1)
p.add_argument("--top_p", type = float, default = 0.97)
p.add_argument("--min_p", type = float, default = 0.5)

63
scripts/benchmarks/qwen3_flex_inference.py
View file

@ -632,6 +632,23 @@ def main():
"CuTeDSL FA4 kernel on Blackwell (SM100)."
),
)
p.add_argument(
"--load_in_4bit",
action = "store_true",
help = (
"Load the base model as bitsandbytes 4-bit. When set with "
"--lora_adapter, the LoRA is kept as a PEFT wrapper (no merge) "
"because merging into 4-bit weights is not supported."
),
)
p.add_argument(
"--model_name_4bit",
default = None,
help = (
"Override the 4-bit shard name. Defaults to "
"`{model_name}-unsloth-bnb-4bit`."
),
)
p.add_argument("--stats_path", required = True)
args = p.parse_args()
@ -645,26 +662,50 @@ def main():
tok = AutoTokenizer.from_pretrained(args.model_name)
if tok.pad_token is None:
tok.pad_token = tok.eos_token
# Load eager; we swap attention forward below.
model = AutoModelForCausalLM.from_pretrained(
args.model_name,
dtype = torch.bfloat16,
attn_implementation = "eager",
).to("cuda")
if args.load_in_4bit:
# Load the pre-quantized Unsloth 4-bit shard. Compute dtype comes
# from the packaged config (bf16 for these shards).
bnb_model_name = args.model_name_4bit or f"{args.model_name}-unsloth-bnb-4bit"
print(f"[flex] loading 4-bit base: {bnb_model_name}")
model = AutoModelForCausalLM.from_pretrained(
bnb_model_name,
attn_implementation = "eager",
device_map = "cuda:0",
)
# See note in cb_vs_vllm_generation.py: tie lm_head to embed_tokens
# for bnb-4bit shards of tied-embedding models.
if getattr(model.config, "tie_word_embeddings", False):
model.lm_head.weight = model.model.embed_tokens.weight
else:
model = AutoModelForCausalLM.from_pretrained(
args.model_name,
dtype = torch.bfloat16,
attn_implementation = "eager",
).to("cuda")
model.eval()
if args.lora_adapter:
from peft import PeftModel
model = PeftModel.from_pretrained(
peft_model = PeftModel.from_pretrained(
model,
str(Path(args.lora_adapter).resolve()),
is_trainable = False,
)
# Merge so attention forward below sees merged weights without the
# PEFT wrapper mangling `self.q_proj` etc.
model = model.merge_and_unload()
model.eval()
if args.load_in_4bit:
# Can't merge LoRA into 4-bit base. Keep PEFT wrapping active --
# `q_proj`/etc on each layer are now LoraLayer(base_layer=Linear4bit,
# lora_A=..., lora_B=...). The monkey-patched attention forward
# calls `self.q_proj(hidden_states)` which routes through LoRA.
# For `patch_qwen3_model` / `call_model_with_flex_kwargs` we pass
# the underlying Qwen3ForCausalLM that PEFT has already modified
# in-place.
model = peft_model.base_model.model
else:
# bf16 path: merge LoRA so there's no PEFT wrapper at call time.
model = peft_model.merge_and_unload()
model.eval()
from unsloth_grpo_common import (
SYSTEM_PROMPT,

25
scripts/benchmarks/results/flex_vs_vllm.md
View file

@ -56,6 +56,31 @@ At the GRPO workload flex reaches **72 % of vLLM throughput at
3.5 × less memory**. Up from 9 % with transformers CB at the start of this
work.
### Same workload at `load_in_4bit=True` (Unsloth bnb-4bit shard)
Loading base as bitsandbytes 4-bit (`unsloth/Qwen3-4B-Base-unsloth-bnb-4bit`,
compute dtype bf16). LoRA kept as PEFT wrapper (can't merge into 4-bit).
lm_head is tied to embed_tokens post-load because the 4-bit shard ships
without an lm_head parameter.
| Backend | tok/s | peak mem | output |
|---------------------------------|------:|---------:|:------------|
| Unsloth fast_inference (vLLM) | 4515 | 159 GB | coherent |
| **flex** (this PR) |**1738**| **40.6 GB** | coherent |
| transformers CB (sdpa) | 504 | 124 GB | **gibberish** |
4-bit costs throughput on every backend (vLLM-path 4515 vs bf16 7775 = 58 %;
flex 1738 vs bf16 5744 = 30 %). The regression is worse for flex because
PEFT-without-merge doubles the number of matmuls per projection (base + LoRA
add, separately) on top of the bnb dequant cost; the bf16 path merges LoRA
into the base and skips both. Peak memory barely moves for vLLM because KV
cache at `gpu_memory_utilization=0.8` dominates regardless of base size.
Transformers CB at 4-bit + LoRA produces garbage tokens even with
`model.lm_head.weight = model.model.embed_tokens.weight` tied explicitly.
Likely a PEFT-over-bnb + batched `generate_batch` interaction bug; did not
debug further.
## What each option did (batch 64, no LoRA, after CUDA graph capture)
| Config | tok/s | vs baseline |

30
scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit.json Normal file
View file

@ -0,0 +1,30 @@
{
"backend": "tpaged",
"lora_adapter": "outputs/lora_rank32_fresh",
"attn_impl": "sdpa",
"persistent_cb": false,
"n_prompts": 64,
"n_prompt_tokens": 9129,
"n_decoded_tokens": 32768,
"wall_times_s": [
60.9150581190479,
56.917108469991945,
58.08906611002749
],
"median_wall_s": 58.08906611002749,
"prompt_tps": 157.15522061774251,
"decode_tps": 564.0992736556235,
"max_new_tokens": 512,
"sample_completions": [
"FirstFirst??? ? ",
"Let 11111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
"First.$^ "
],
"peak_memory_gb": 124.05089378356934,
"sampling": {
"temperature": 0.1,
"top_p": 0.97,
"min_p": 0.5,
"top_k": 5
}
}

30
scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit_tied.json Normal file
View file

@ -0,0 +1,30 @@
{
"backend": "tpaged",
"lora_adapter": "outputs/lora_rank32_fresh",
"attn_impl": "sdpa",
"persistent_cb": false,
"n_prompts": 64,
"n_prompt_tokens": 9129,
"n_decoded_tokens": 32275,
"wall_times_s": [
61.436025188013446,
63.98777190799592,
65.29973084997619
],
"median_wall_s": 63.98777190799592,
"prompt_tps": 142.66788368762124,
"decode_tps": 504.39324635973,
"max_new_tokens": 512,
"sample_completions": [
"First list list list list list list list list list list list list list list list list list<|endoftext|>",
"FirstFirst's???. \n,?.. and and and2222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222",
"First. and. and and2222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222"
],
"peak_memory_gb": 124.05089378356934,
"sampling": {
"temperature": 0.1,
"top_p": 0.97,
"min_p": 0.5,
"top_k": 5
}
}

25
scripts/benchmarks/results/stats/flex_64_lora_4bit.json Normal file
View file

@ -0,0 +1,25 @@
{
"backend": "qwen3_flex",
"capture_cudagraph": true,
"lora_adapter": "outputs/lora_rank32_fresh",
"n_prompts": 64,
"n_decoded_tokens": 28172,
"wall_times_s": [
19.392819664964918,
14.716534855018836,
14.195494230021723,
15.050612487946637,
16.052564749028534
],
"median_wall_s": 15.050612487946637,
"best_wall_s": 14.195494230021723,
"decode_tps_median": 1871.8175105871403,
"decode_tps_best": 1984.5733824765107,
"max_new_tokens": 512,
"peak_memory_gb": 40.59123468399048,
"sample_completions": [
"Let's denote the length of segment $DB$ as $ units.\n<start_working_out>\nTo solve this problem, we can use the Power of a Point theorem, which states that for a point P inside a circle, the product of the lengths of the segments of any two intersecting chords through P is constant. In this case, we have two intersecting chords: AB and CD. Let",
"Let $P(x)$ be a monic polynomial of degree $2023$ such that $P(k) = k^{2023}P(1-\\frac{1}{k})$ for every positive integer $1 \\leq k \\leq 2023$. We want to find $P(-1)$.\n\nFirst, we can rewrite the given equation",
"First, let's consider the cube's edges. A cube has 12 edges. Each edge is parallel to 3 other edges. However, we need to be careful not to double-count the pairs.\n\nLet's count the pairs of parallel edges:\n\n1. Each edge is parallel to 3 other edges, so there are 12 * 3 = 36 pairs.\n2."
]
}

25
scripts/benchmarks/results/stats/flex_64_lora_4bit_tied.json Normal file
View file

@ -0,0 +1,25 @@
{
"backend": "qwen3_flex",
"capture_cudagraph": true,
"lora_adapter": "outputs/lora_rank32_fresh",
"n_prompts": 64,
"n_decoded_tokens": 27691,
"wall_times_s": [
19.849357629951555,
16.3931127290125,
15.936416806012858,
17.375963038997725,
17.99713112297468
],
"median_wall_s": 17.375963038997725,
"best_wall_s": 15.936416806012858,
"decode_tps_median": 1593.638288585889,
"decode_tps_best": 1737.59260548156,
"max_new_tokens": 512,
"peak_memory_gb": 40.59123468399048,
"sample_completions": [
"First, let's consider the cube's edges. A cube has 12 edges. Each edge is parallel to 3 other edges. However, we need to be careful not to double-count the pairs.\n\nLet's count the pairs of parallel edges:\n\n1. Each edge is parallel to 3 other edges, so there are 12 * 3 = 36 pairs.\n2.",
"Let the common ratio of the geometric sequence be $r$. Then the second term is $\\frac{3}{4}r=15$, so $r=20$. The $n$th term of the sequence is $\\frac{3}{4}r^{n-1}$. We want to find the smallest $n$ such that $\\frac{3}{4}r^{",
"Let $n = 20k + r$ and $n = 16m + s$, where $0 \\leq r < 20$ and $0 \\leq s < 16$. We want to find the number of integers $n$ such that $r < s$.\n\nSince $n$ is an integer, we have $20k +"
]
}

30
scripts/benchmarks/results/stats/unsloth_fi_true_64_lora_4bit.json Normal file
View file

@ -0,0 +1,30 @@
{
"backend": "vllm",
"lora_adapter": "outputs/lora_rank32_fresh",
"n_prompts": 64,
"n_prompt_tokens": 9129,
"n_decoded_tokens": 30050,
"wall_times_s": [
6.68179759796476,
6.6636974540306255,
6.6514031600090675,
6.653628617990762,
6.65495745599037
],
"median_wall_s": 6.65495745599037,
"prompt_tps": 1371.7593328538342,
"decode_tps": 4515.430819614166,
"max_new_tokens": 512,
"sample_completions": [
"First, we need to find the length of the legs of the trapezoid. Since the trapezoid is isosceles, the legs are equal in length. Let's call the length of each leg $x$. We can use the Pythagorean theore",
"Let $P(x)$ be a monic polynomial of degree $2023$ such that $P(k) = k^{2023}P(1-\\frac{1}{k})$ for every positive integer $1 \\leq k \\leq 2023$. We want to find $P(-1)$.\n\nFirst, we can rewrite the given",
"First, we need to find the value of $a$ such that the graph of $y = mx + 2$ passes through no lattice point with $0 < x \\leq 100$ for all $m$ such that $\\frac{1}{2} < m < a$.\n\nLet's consider the equat"
],
"peak_memory_gb": 159.28503799438477,
"sampling": {
"temperature": 0.1,
"top_p": 0.97,
"min_p": 0.5,
"top_k": 5
}
}