diff --git a/scripts/benchmarks/cb_vs_vllm_generation.py b/scripts/benchmarks/cb_vs_vllm_generation.py index fd4fc715bb..eea9656f40 100644 --- a/scripts/benchmarks/cb_vs_vllm_generation.py +++ b/scripts/benchmarks/cb_vs_vllm_generation.py @@ -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) diff --git a/scripts/benchmarks/qwen3_flex_inference.py b/scripts/benchmarks/qwen3_flex_inference.py index 5c388bea74..43bd95f326 100644 --- a/scripts/benchmarks/qwen3_flex_inference.py +++ b/scripts/benchmarks/qwen3_flex_inference.py @@ -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, diff --git a/scripts/benchmarks/results/flex_vs_vllm.md b/scripts/benchmarks/results/flex_vs_vllm.md index 5543d63887..04be583084 100644 --- a/scripts/benchmarks/results/flex_vs_vllm.md +++ b/scripts/benchmarks/results/flex_vs_vllm.md @@ -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 | diff --git a/scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit.json b/scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit.json new file mode 100644 index 0000000000..f742fc3734 --- /dev/null +++ b/scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit.json @@ -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 + } +} \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit_tied.json b/scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit_tied.json new file mode 100644 index 0000000000..3007331af5 --- /dev/null +++ b/scripts/benchmarks/results/stats/cb_tpaged_64_lora_4bit_tied.json @@ -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 + } +} \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/flex_64_lora_4bit.json b/scripts/benchmarks/results/stats/flex_64_lora_4bit.json new file mode 100644 index 0000000000..9e17e5537a --- /dev/null +++ b/scripts/benchmarks/results/stats/flex_64_lora_4bit.json @@ -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\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." + ] +} \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/flex_64_lora_4bit_tied.json b/scripts/benchmarks/results/stats/flex_64_lora_4bit_tied.json new file mode 100644 index 0000000000..9ade340295 --- /dev/null +++ b/scripts/benchmarks/results/stats/flex_64_lora_4bit_tied.json @@ -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 +" + ] +} \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/unsloth_fi_true_64_lora_4bit.json b/scripts/benchmarks/results/stats/unsloth_fi_true_64_lora_4bit.json new file mode 100644 index 0000000000..1fa2185459 --- /dev/null +++ b/scripts/benchmarks/results/stats/unsloth_fi_true_64_lora_4bit.json @@ -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 + } +} \ No newline at end of file