This tutorial walks through a complete optimization pipeline: loading a baseline model, quantizing it, benchmarking latency, and serving it via vLLM. All code is runnable on a single A10 or T4 GPU.

Setup

pip install transformers accelerate auto-gptq vllm bitsandbytes torch==2.2.0

Step 1: Baseline Benchmark

Establish your latency and throughput baseline before any optimization.

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# FP16 baseline
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto"
)

prompt = "Explain transformer attention in three sentences."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Warmup
for _ in range(3):
    with torch.no_grad():
        model.generate(**inputs, max_new_tokens=128)

# Timed run
start = time.perf_counter()
for _ in range(10):
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=128)
elapsed = (time.perf_counter() - start) / 10

tokens_generated = out.shape[1] - inputs["input_ids"].shape[1]
print(f"Latency: {elapsed:.3f}s | Tokens/sec: {tokens_generated / elapsed:.1f}")

Step 2: Apply INT4 Quantization (bitsandbytes)

from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,  # nested quantization
)

model_4bit = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)

# Re-run benchmark on quantized model
start = time.perf_counter()
for _ in range(10):
    with torch.no_grad():
        out = model_4bit.generate(**inputs, max_new_tokens=128)
elapsed_4bit = (time.perf_counter() - start) / 10

print(f"4-bit Latency: {elapsed_4bit:.3f}s | Speedup: {elapsed / elapsed_4bit:.2f}x")
print(f"Memory saved: ~{(1 - 1/4) * 100:.0f}% vs FP32")

Step 3: Serve with vLLM for Production Throughput

from vllm import LLM, SamplingParams
import time

# vLLM handles continuous batching automatically
llm = LLM(
    model=model_id,
    quantization="awq",       # use AWQ for vLLM; download AWQ checkpoint first
    dtype="auto",
    max_model_len=4096,
    gpu_memory_utilization=0.85,
)

params = SamplingParams(temperature=0.0, max_tokens=128)

# Simulate 50 concurrent requests
prompts = [prompt] * 50
start = time.perf_counter()
outputs = llm.generate(prompts, params)
total_time = time.perf_counter() - start

total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
print(f"Throughput: {total_tokens / total_time:.0f} tokens/sec across 50 requests")

Step 4: Validate Quality

Never optimize without a quality gate. Use ROUGE or exact-match on a held-out set:

from rouge_score import rouge_scorer

scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)

baseline_outputs = [...]   # pre-collected
quantized_outputs = [...]  # from optimized model

scores = [
    scorer.score(ref, hyp)["rougeL"].fmeasure
    for ref, hyp in zip(baseline_outputs, quantized_outputs)
]
avg_rouge = sum(scores) / len(scores)
print(f"Average ROUGE-L vs baseline: {avg_rouge:.3f}")
# Accept if > 0.92

Expected Results on Llama-2-7B