4717bce97e
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. |
||
|---|---|---|
| .. | ||
| benchmarks | ||
| enforce_kwargs_spacing.py | ||
| install_gemma4_mlx.sh | ||
| install_qwen3_6_mlx.sh | ||
| run_ruff_format.py | ||