Fine-Tuning Llama 2 on the Alpaca Dataset with LoRA and 4-bit Quantization
This notebook demonstrates how to fine-tune Llama 2 (7B) for instruction-following using the Alpaca dataset.
The fine-tuning applies: - 4-bit quantization via bitsandbytes to reduce GPU memory usage. - LoRA (Low-Rank Adaptation) via peft to train only a small fraction of parameters.
These techniques allow fine-tuning a 7 billion parameter model on a Colab T4 GPU with 16 GB VRAM.
Model — Llama 2 7B
Llama 2 is Meta’s open-weight language model family. The 7B variant is the smallest and most suitable for single-GPU fine-tuning.
Dataset — Stanford Alpaca
The Alpaca dataset was created by Stanford and contains 52,000 instruction-response pairs generated by GPT-3.5. Each example has three fields: - instruction: the task description. - input: optional context (can be empty). - output: the expected model response.
Install Dependencies
# Install required libraries# transformers : Hugging Face model hub and training utilities# datasets : easy dataset loading and preprocessing# peft : parameter-efficient fine-tuning (LoRA)# bitsandbytes : 4-bit / 8-bit quantization kernels# trl : SFTTrainer for supervised fine-tuning# accelerate : distributed / mixed-precision training!pip install -q transformers datasets peft bitsandbytes trl accelerate
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━60.7/60.7 MB37.9 MB/s eta 0:00:00━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━540.5/540.5 kB53.7 MB/s eta 0:00:00
# Inspect the first sample of instruction-response pairsprint(dataset[0]['text'][:500])
Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
Give three tips for staying healthy.
### Response:
1.Eat a balanced diet and make sure to include plenty of fruits and vegetables.
2. Exercise regularly to keep your body active and strong.
3. Get enough sleep and maintain a consistent sleep schedule.
Quantization Configuration (4-bit QLoRA)
# BitsAndBytes configuratonbnb_config = BitsAndBytesConfig(# Load the model using 4-bit precision load_in_4bit=True,# Quantization type (fp4 or nf4)# nf4 is "normalized float 4" format, uses an asymmetric quantization scheme with 4-bit precision# optimized for normally distributed weights (better than fp4 for neural networks) bnb_4bit_quant_type="nf4",# Compute dtype for 4-bit models bnb_4bit_compute_dtype= torch.float16,# Use double quantization for 4-bit models# Double quantization applies further quantization to the quantization constants bnb_4bit_use_double_quant=True,)
Load the Base Model and Tokenizer
# The model is Llama 2 from the Hugging Face hubmodel_name ="NousResearch/Llama-2-7b-chat-hf"
# Load Llama 2 model from Hugging Facemodel = AutoModelForCausalLM.from_pretrained( model_name,# Apply quantization by using the bnb configuration from the previous cell quantization_config=bnb_config,# Don't cache the model weights, load the model weights from Hugging Face use_cache=False,# Trade-off parameter in Llama-2, less important, it should be 1 in most cases pretraining_tp=1,# Load the entire model on the GPU if available device_map="auto")
# Prepare model for k-bit trainingmodel = prepare_model_for_kbit_training(model)
# Load tokenizer from Hugging Facetokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)# Needed for LlaMA tokenizer, it does not have pad token by defaulttokenizer.pad_token = tokenizer.eos_token# Fix an overflow issue with fp16 training, pad on the righttokenizer.padding_side ="right"
# Print the vocabulary sizeprint(tokenizer.vocab_size)
32000
Apply LoRA (Parameter-Efficient Fine-Tuning)
# LoRA configurationlora_config = LoraConfig(# LoRA rank dimension (controls capacity of LoRA layers) r=16,# Scaling parameter for LoRA updates (higher alpha increases contribution of LoRA weights) lora_alpha=64,# Dropout rate for LoRA layers lora_dropout=0.1,# Specifies whether to add bias in LoRA layers bias="none",# "CAUSAL_LM" indicates that LoRA is applied for causal language modeling task_type="CAUSAL_LM",# List of modules to which LoRA is applied (Q, K, V, and output projections) target_modules=["q_proj", "k_proj", "v_proj", "o_proj"])
# Wrap the model with LoRA adaptersmodel = get_peft_model(model, lora_config)
# Print the fraction of trainable parameters with LoRAmodel.print_trainable_parameters()
trainable params: 16,777,216 || all params: 6,755,192,832 || trainable%: 0.2484
Model Fine-tuning
# Set training parameterstraining_arguments = TrainingArguments(# Output directory where the model predictions and checkpoints will be stored output_dir="./results",# Number of training epochs or maximum training steps# num_train_epochs=1, # One epoch will take quite long, use max_steps instead max_steps =100, # Note also that max_steps has precedence over num_train_epochs# Batch size per GPU for training per_device_train_batch_size=4,# Number of update steps to accumulate the gradients for# Helps simulate a larger batch size without increasing memory usage gradient_accumulation_steps=2,# Optimizer to use (memory-efficient Adam with weight decay) optim="paged_adamw_32bit",# Log updates every number of steps logging_steps=10,# Initial learning rate for the optimizer learning_rate=2e-4,# Weight decay to prevent overfitting weight_decay=0.001,# Enable fp16/bf16 training (set bf16 to True with an A100) fp16=False, bf16=False,# Maximum gradient norm (gradient clipping to prevent exploding gradients) max_grad_norm=0.3,# Group sequences with similar length into same batches (to minimize padding)# Saves memory and speeds up training considerably group_by_length=True,# Learning rate scheduler type lr_scheduler_type="cosine",# Fraction of total training steps used for warmup warmup_steps=10,# Disable reporting to external tools (e.g., WandB, TensorBoard) report_to="none")
The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'pad_token_id': 2}.
/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py:1181: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. Starting in PyTorch 2.9, calling checkpoint without use_reentrant will raise an exception. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.
return fn(*args, **kwargs)
# Set model to inference mode (very important, don't skip this cell)model.eval()model.config.use_cache =True
# Function to generate a response to a given prompt (instruction)def generate_response(instruction, max_new_tokens=500):"""Generate a response from the fine-tuned model."""# Use chat template if available (Qwen, Phi, etc.), else use Llama 2 formatif tokenizer.chat_template: messages = [{"role": "user", "content": instruction}] inputs = tokenizer.apply_chat_template( messages, return_tensors='pt', add_generation_prompt=True, return_dict=False ).to(model.device)else: prompt ="### Instruction:\n"+ instruction +"\n\n### Response:\n" inputs = tokenizer(prompt, return_tensors='pt')['input_ids'].to(model.device)# Generate responsewith torch.no_grad(): outputs = model.generate( inputs, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.7, top_p=0.9, repetition_penalty=1.1, eos_token_id=tokenizer.eos_token_id, )# Convert the response to text and print it generated = outputs[0][inputs.shape[1]:] response = tokenizer.decode(generated, skip_special_tokens=True).strip() response = response.split("</s>")[0].strip()print(response)
# Example promptgenerate_response('Give three tips for staying productive while working from home.')
1. Set a schedule and stick to it, even if you're working from home.
2. Take regular breaks and prioritize self-care.
3. Make sure to create boundaries between work and personal life.
# Example promptgenerate_response("Write a poem about Python Programming. Use line breaks between verses.")
Python is a language of power,
A way to build and create with ease.
With its syntax so simple yet strong,
It's no wonder that it's the choice of the day.
Its libraries and frameworks are vast,
And with them, anything can be built at last.