Chapter 2: 4-bit Quantization with bitsandbytes

Understanding and Implementing Efficient Quantization

The bitsandbytes library enables 4-bit and 8-bit quantization for efficient model loading and training.

---

Quantization Basics

What is Quantization?

Converting high-precision values to lower precision:

Float16: 2 bytes per parameter
Int8: 1 byte per parameter
Int4: 0.5 bytes per parameter

Memory Savings

# 7B model memory
fp32 = 7e9 * 4 / 1e9  # 28 GB
fp16 = 7e9 * 2 / 1e9  # 14 GB
int8 = 7e9 * 1 / 1e9  # 7 GB
int4 = 7e9 * 0.5 / 1e9  # 3.5 GB

---

BitsAndBytesConfig

Complete Configuration

from transformers import BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    # Use 4-bit quantization
    load_in_4bit=True,
    
    # NF4 quantization type (recommended)
    bnb_4bit_quant_type="nf4",
    
    # Compute in bfloat16 for stability
    bnb_4bit_compute_dtype=torch.bfloat16,
    
    # Double quantization for extra savings
    bnb_4bit_use_double_quant=True,
)

Parameter Explanations

Parameter	Options	Recommendation
load_in_4bit	True/False	True for QLoRA
bnb_4bit_quant_type	"nf4", "fp4"	"nf4" (better quality)
bnb_4bit_compute_dtype	torch.bfloat16, torch.float16	bfloat16 (more stable)
bnb_4bit_use_double_quant	True/False	True (more savings)

---

NF4 vs FP4

NF4 (NormalFloat4)

• Designed for normally distributed weights
• Equal values in each quantization bin
• Better quality for neural networks

FP4 (Float4)

• Uniform quantization
• Simpler but less optimal for weights

# NF4 recommended
bnb_4bit_quant_type="nf4"

---

Loading a Quantized Model

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

# Check memory
print(f"Memory: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

---

Preparing for Training

from peft import prepare_model_for_kbit_training

# Necessary preparation steps:
# 1. Enable gradient checkpointing
# 2. Cast layer norms to float32
# 3. Enable input gradients

model = prepare_model_for_kbit_training(model)

What prepare_model_for_kbit_training Does

1. Enables gradient checkpointing for memory efficiency
2. Casts layer norms to float32 for stability
3. Sets requires_grad appropriately
4. Handles embeddings for input gradients

---

Complete QLoRA Setup

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# 1. Quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# 2. Load quantized model
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    quantization_config=bnb_config,
    device_map="auto",
)

# 3. Prepare for training
model = prepare_model_for_kbit_training(model)

# 4. Add LoRA
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

---

Key Takeaways

✅ bitsandbytes enables 4-bit/8-bit quantization ✅ NF4 is optimal for neural network weights ✅ Double quantization adds ~11% more savings ✅ Use prepare_model_for_kbit_training before adding LoRA ✅ Compute dtype should be bfloat16 for stability

---

Continue to Chapter 3: Memory-Efficient Training →