diff --git a/tests/gemma4_flex_bench.py b/tests/gemma4_flex_bench.py new file mode 100644 index 0000000000..f896005263 --- /dev/null +++ b/tests/gemma4_flex_bench.py @@ -0,0 +1,132 @@ +"""Throughput bench for Unsloth flex fast-inference on Gemma 4. + +Runs ``FastLanguageModel(fast_inference=True)`` through the flex engine +(``UNSLOTH_FAST_INFERENCE=1``) and measures tok/s at a list of batch sizes. +Captures CUDA graphs at each bucket; the bench reports the post-warmup pass. + +Example: + CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 \\ + UNSLOTH_MOE_BACKEND=grouped_mm \\ + python -u tests/gemma4_flex_bench.py \\ + --model unsloth/gemma-4-26b-a4b-it --batch_sizes 1 4 8 16 \\ + --max_new_tokens 64 --json_out async_task_outputs/flex_bench_26b.json + +Pair with tests/gemma4_fast_inference_parity.py for HF-naive parity and +with tests/gemma4_fast_bench.py (vLLM nightly path) for cross-engine +throughput comparison. +""" +import argparse +import json +import os +import time + +os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") +os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") + +import unsloth # noqa: F401, E402 +import torch # noqa: E402 + + +SHORT_PROMPT = "In one sentence, what is Paris?" +LONG_POOL = [ + "In one sentence, what is Paris?", + "What is 23 + 19? Answer in one word.", + "Continue this phrase: The quick brown fox jumps over", + "Name three primary colors.", + "Who wrote the play Hamlet?", + "What is the capital of Japan?", + "What does the acronym 'NASA' stand for?", + "Explain the water cycle in one sentence.", + "Give a two-word summary of the French Revolution.", + "What is the tallest mountain on Earth?", + "Define 'entropy' in one sentence.", + "Who painted the Mona Lisa?", + "What is the boiling point of water in Celsius?", + "Name one prime number larger than 10.", + "What color do you get when you mix blue and yellow?", + "Give one example of an amphibian.", +] + + +def _render(tok, prompts): + return [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, add_generation_prompt=True, + ) + for p in prompts + ] + + +class _SP: + def __init__(self, n): + self.max_tokens = n + self.temperature = 0.0 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--batch_sizes", nargs="+", type=int, default=[1, 4, 8, 16]) + ap.add_argument("--max_new_tokens", type=int, default=64) + ap.add_argument("--json_out", required=True) + ap.add_argument("--max_seq_length", type=int, default=1024) + args = ap.parse_args() + + from unsloth import FastLanguageModel + + t_load = time.perf_counter() + model, tok_raw = FastLanguageModel.from_pretrained( + model_name=args.model, max_seq_length=args.max_seq_length, + dtype=torch.bfloat16, load_in_4bit=False, fast_inference=True, + max_batch_size=max(args.batch_sizes), + gpu_memory_utilization=0.6, + ) + t_load = time.perf_counter() - t_load + tok = tok_raw.tokenizer if hasattr(tok_raw, "tokenizer") else tok_raw + print(f"[flex] load {t_load:.1f}s") + + sp = _SP(args.max_new_tokens) + results = [] + + for bs in args.batch_sizes: + prompts = _render(tok, LONG_POOL[:bs]) + # Warmup: primes the cudagraph bucket and walker compile. + _ = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False) + t0 = time.perf_counter() + outs = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False) + dt = time.perf_counter() - t0 + total = sum(len(o.outputs[0].token_ids) for o in outs) + tok_s = total / dt if dt > 0 else 0.0 + results.append({ + "batch_size": bs, "gen_s": round(dt, 3), + "total_tokens": total, "tok_s": round(tok_s, 1), + }) + print(f"[flex bs={bs}] {total} tok in {dt:.2f}s -> {tok_s:.1f} tok/s") + + # Parity sanity: run SHORT_PROMPT at bs=1, save ids. + p = _render(tok, [SHORT_PROMPT]) + outs = model.fast_generate(p, sampling_params=sp, use_tqdm=False) + parity_ids = list(outs[0].outputs[0].token_ids) + parity_text = outs[0].outputs[0].text + + peak_gb = torch.cuda.max_memory_allocated() / 1024**3 + + summary = { + "model": args.model, + "max_new_tokens": args.max_new_tokens, + "load_s": round(t_load, 2), + "peak_vram_gb": round(peak_gb, 2), + "batches": results, + "parity_ids": parity_ids, + "parity_text": parity_text, + } + os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True) + with open(args.json_out, "w") as f: + json.dump(summary, f, indent=2) + print(f"[flex] wrote {args.json_out}") + + +if __name__ == "__main__": + main()