QLoRA › Implementation with bitsandbytes
Language AI HandbookAll chapters QLoRA › Implementation with bitsandbytesOne section · Sections in this chapter10 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The bitsandbytes library provides efficient CUDA kernels for 4-bit operations, and the peft library integrates these with LoRA. Let's walk through a practical QLoRA setup. The full implementation requires GPU access and the target model weights, so the code blocks below are marked eval: false but show the complete workflow you'd follow in a real fine-tuning session.
Understanding how these two libraries interact is important for troubleshooting. The bitsandbytes library handles all the low-level quantization mechanics: it intercepts model weight loading, quantizes the tensors using NF4 with double quantization, and provides the CUDA kernels that perform on-the-fly dequantization during the forward pass. The peft library sits on top and handles the LoRA-specific concerns: creating the adapter matrices, wrapping the quantized layers with LoRA computation, and managing the saving and loading of adapter weights. The transformers library orchestrates the overall training loop and provides the model architecture. Together, these three libraries implement the full QLoRA pipeline with minimal configuration.
# Install required packages
# uv pip install bitsandbytes peft transformers accelerate
import torch
from transformers import BitsAndBytesConfig
# Configure 4-bit quantization
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4", # Use NormalFloat
bnb_4bit_compute_dtype=torch.bfloat16, # Compute in BF16
bnb_4bit_use_double_quant=True, # Enable double quantization
)The key parameters in BitsAndBytesConfig are:
- load_in_4bit: Enables 4-bit quantization when loading the model
- bnb_4bit_quant_type: Specifies NF4 rather than uniform INT4
- bnb_4bit_compute_dtype: Sets the dtype for dequantized computation
- bnb_4bit_use_double_quant: Enables quantization of the quantization constants
Loading a Quantized Model
model_name = "meta-llama/Llama-2-7b-hf"
# Load model with 4-bit quantization
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # Automatically place layers across available devices
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
# Prepare model for k-bit training
# This handles gradient checkpointing and freezes appropriate layers
model = prepare_model_for_kbit_training(model)The prepare_model_for_kbit_training function performs several important operations:
- Enables gradient checkpointing to reduce activation memory
- Freezes all quantized layers
- Casts layer norms to FP32 for training stability
Adding LoRA Adapters
# Configure LoRA
lora_config = LoraConfig(
r=16, # LoRA rank
lora_alpha=32, # Scaling factor
target_modules=[
"q_proj",
"k_proj",
"v_proj",
"o_proj",
], # Target attention layers
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
# Add LoRA adapters to the quantized model
model = get_peft_model(model, lora_config)# Check trainable parameters
model.print_trainable_parameters()For a 7B model with rank-16 LoRA targeting attention projections, typically less than 1% of parameters are trainable.
Training Configuration
from transformers import TrainingArguments
training_args = TrainingArguments(
output_dir="./qlora_output",
num_train_epochs=3,
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
learning_rate=2e-4,
warmup_ratio=0.03,
logging_steps=10,
save_strategy="epoch",
fp16=False, # Disable FP16, we're using BF16 via quantization config
bf16=True, # Enable BF16 for non-quantized operations
optim="paged_adamw_8bit", # Use paged optimizer
gradient_checkpointing=True,
max_grad_norm=0.3,
)The paged_adamw_8bit optimizer combines two memory-saving techniques: 8-bit Adam (which quantizes optimizer states) and paged memory management. The 8-bit optimizer states provide an additional 2× reduction for the already-small LoRA optimizer states.
Inference with Merged Weights
After training, you can either use the adapter on top of the quantized model, or merge and export to full precision:
# Option 1: Use adapter directly (keeps quantization)
model.eval()
inputs = tokenizer("The capital of France is", return_tensors="pt").to(
model.device
)
outputs = model.generate(**inputs, max_new_tokens=10)
print(tokenizer.decode(outputs[0]))
# Option 2: Merge adapter into base model for deployment
# First, load base model in FP16
base_model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto",
)
# Load and merge adapter
from peft import PeftModel
merged_model = PeftModel.from_pretrained(base_model, "./qlora_output")
merged_model = merged_model.merge_and_unload()
# Now merged_model has the fine-tuned weights without adapters