๐ Production LLM Deployment & Quantization
QLoRA Math & 4-bit NormalFloat (NF4)
From 16-bit weights to 4-bit NF4, Double Quantization, and LoRA rank decomposition โ the full math
Fine-tuning a 70B parameter LLaMA model in full BF16 precision requires 140 GB of GPU VRAM โ impossible on most hardware. QLoRA (Quantized Low-Rank Adaptation) solves this by stacking three innovations: 4-bit NormalFloat (NF4) quantization, Double Quantization, and Paged Optimizers โ enabling fine-tuning on a single A100 (80 GB) or even a consumer RTX 4090 (24 GB).
Key Concepts
4-bit NormalFloat (NF4)
NF4 is not ordinary int4. It places 16 quantization levels at positions equal to the quantiles of a standard normal distribution N(0,1), matching the actual weight distribution of pre-trained neural networks perfectly.
Double Quantization (DQ)
After quantizing weights to NF4, the 32-bit quantization constants (one per 64-param block) are themselves quantized to 8-bit. This saves an additional 0.37 bits/parameter โ 3 GB on a 65B model with zero accuracy loss.
LoRA Rank Decomposition
Instead of updating all W weights, LoRA freezes W and injects two tiny matrices A (dรr) and B (rรd) where r<<d. The adapter forward pass is: h = Wโx + (ฮฑ/r)ยทBAยทx. Only A and B are trained โ 99.9% fewer trainable parameters.
Paged Optimizers
When GPU VRAM is exhausted during long training sequences, NVIDIA's unified memory automatically pages optimizer states (Adam momentum, variance) from GPU VRAM to CPU RAM โ preventing OOM crashes silently.
The QLoRA Math: Full Forward Pass Equation
The complete QLoRA forward pass equation combines dequantization and the LoRA adapter:
Where: WNF4 = frozen base weights in 4-bit NF4 ยท cโ = block-level quantization constants (8-bit) ยท cโ = superblock constants (FP32) ยท A, B = trainable LoRA adapter matrices in BF16 ยท ฮฑ/r = scaling gamma.
h = doubleDequant(cโFP32, cโk-bit, WNF4) ยท x + (ฮฑ / r) ยท B ยท A ยท x
Where: WNF4 = frozen base weights in 4-bit NF4 ยท cโ = block-level quantization constants (8-bit) ยท cโ = superblock constants (FP32) ยท A, B = trainable LoRA adapter matrices in BF16 ยท ฮฑ/r = scaling gamma.
Memory Savings Comparison
| Method | Precision | 7B Model VRAM | 65B Model VRAM | Accuracy vs FP16 |
|---|---|---|---|---|
| Full Fine-tuning | FP16 / BF16 | ~14 GB | ~130 GB | Baseline |
| LoRA only (no quant) | BF16 | ~10 GB | ~96 GB | Identical |
| QLoRA (NF4 + DQ) | NF4 + BF16 adapters | ~4.5 GB | ~33 GB | < 0.1% degradation |
| QLoRA + Paged Adam | NF4 + BF16 | ~5.5 GB* | ~36 GB* | < 0.1% degradation |
Production Code: Complete QLoRA Fine-tuning Pipeline
qlora_finetune.pypython
1import torch
2from datasets import load_dataset
3from transformers import (
4 AutoModelForCausalLM,
5 AutoTokenizer,
6 BitsAndBytesConfig,
7 TrainingArguments,
8)
9from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
10from trl import SFTTrainer
11
12# โโ 1. 4-bit NF4 Quantization Config โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
13bnb_config = BitsAndBytesConfig(
14 load_in_4bit=True,
15 bnb_4bit_quant_type="nf4", # NormalFloat 4-bit (NOT int4)
16 bnb_4bit_use_double_quant=True, # Quantize quant constants โ saves 0.37 bits/param
17 bnb_4bit_compute_dtype=torch.bfloat16 # Compute in BF16 during adapter forward pass
18)
19
20# โโ 2. Load Quantized Base Model โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
21model = AutoModelForCausalLM.from_pretrained(
22 "meta-llama/Meta-Llama-3-8B",
23 quantization_config=bnb_config,
24 device_map="auto", # Auto-distribute across available GPUs
25 attn_implementation="flash_attention_2"
26)
27
28# Required before LoRA injection: cast LayerNorm to FP32, freeze base weights
29model = prepare_model_for_kbit_training(model)
30
31# โโ 3. LoRA Adapter Config โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
32lora_config = LoraConfig(
33 r=16, # Rank: higher = more capacity, more VRAM
34 lora_alpha=32, # Scaling: gamma = alpha/r = 2.0
35 target_modules=[ # Inject adapters into all attention projections
36 "q_proj", "k_proj", "v_proj", "o_proj",
37 "gate_proj", "up_proj", "down_proj" # Also MLP layers for better task transfer
38 ],
39 lora_dropout=0.05,
40 bias="none",
41 task_type="CAUSAL_LM"
42)
43
44model = get_peft_model(model, lora_config)
45model.print_trainable_parameters()
46# Output: trainable params: 41,943,040 || all params: 8,072,622,080 || trainable%: 0.5197%
47
48# โโ 4. Training with Paged Adam (prevents OOM on long sequences) โโโโโโโโโโโ
49training_args = TrainingArguments(
50 output_dir="./qlora-llama3-8b-ft",
51 num_train_epochs=3,
52 per_device_train_batch_size=4,
53 gradient_accumulation_steps=4, # Effective batch = 4 ร 4 = 16
54 optim="paged_adamw_32bit", # Paged Optimizer: pages states to CPU on VRAM spike
55 learning_rate=2e-4,
56 lr_scheduler_type="cosine",
57 warmup_ratio=0.03,
58 bf16=True,
59 save_strategy="epoch",
60 logging_steps=10,
61 report_to="wandb"
62)
63
64# โโ 5. SFTTrainer with packing for efficiency โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
65dataset = load_dataset("tatsu-lab/alpaca", split="train")
66
67trainer = SFTTrainer(
68 model=model,
69 args=training_args,
70 train_dataset=dataset,
71 dataset_text_field="text",
72 max_seq_length=2048,
73 packing=True, # Packs multiple short samples โ fewer wasted tokens
74)
75
76trainer.train()
77
78# โโ 6. Merge LoRA into base model for deployment โโโโโโโโโโโโโโโโโโโโโโโโโโโ
79from peft import PeftModel
80merged_model = model.merge_and_unload() # Fuse AยทB into Wโ โ standard model format
81merged_model.save_pretrained("./qlora-merged-llama3-8b")
82print("Merged model saved โ ready for vLLM serving")
83QLoRA Architecture Diagram
Choosing the Right Rank (r) โ The Tradeoff
| Rank r | Adapter Params (7B model) | Extra VRAM | Best Use Case |
|---|---|---|---|
| r = 4 | ~10M params | ~80 MB | Style transfer, simple chat tuning |
| r = 8 | ~20M params | ~160 MB | Domain adaptation (medical, legal) |
| r = 16 | ~41M params | ~320 MB | Complex reasoning, code generation |
| r = 64 | ~164M params | ~1.3 GB | Full task-specific capability matching |
After Training: merge_and_unload() for vLLM
merge_for_deployment.pypython
1from peft import PeftModel
2from transformers import AutoModelForCausalLM, BitsAndBytesConfig
3import torch
4
5# Reload base model in FP16 for clean merge (not 4-bit โ merging needs full precision)
6base_model = AutoModelForCausalLM.from_pretrained(
7 "meta-llama/Meta-Llama-3-8B",
8 torch_dtype=torch.float16,
9 device_map="cpu" # Use CPU for merge to avoid VRAM limits
10)
11
12# Load trained LoRA adapter weights
13peft_model = PeftModel.from_pretrained(base_model, "./qlora-llama3-8b-ft")
14
15# Merge: mathematically fuses Wโ + (ฮฑ/r)ยทBยทA โ single W matrix
16# Result is a standard HuggingFace model with NO LoRA overhead at inference
17merged = peft_model.merge_and_unload()
18
19# Save in safetensors format (recommended for vLLM)
20merged.save_pretrained("./merged-llama3-8b", safe_serialization=True)
21
22print("Merged model ready for vLLM. Load with:")
23print(' llm = LLM(model="./merged-llama3-8b", dtype="float16")')
24๐ก
Senior Architect Insight: The key production insight: never deploy LoRA adapters at inference time. Always merge_and_unload() before vLLM deployment. Adapter hot-loading adds per-request overhead. Merged models run at full PagedAttention throughput with zero adapter overhead.