Fine-Tuning Llama-3 for Medical Q&A using PEFT and LoRA

Introduction

General-purpose Large Language Models (LLMs) like Meta's Llama-3 are incredibly powerful, but they often lack the precise domain knowledge required for clinical Question & Answering (Q&A). Full fine-tuning of an 8B or 70B parameter model is computationally prohibitive for most researchers. In this tutorial, we will use Parameter-Efficient Fine-Tuning (PEFT) and Low-Rank Adaptation (LoRA) to teach Llama-3 to answer medical questions accurately, using only a single consumer GPU.

Architecture Overview

graph TD
    Input(Medical Query) --> Base[Llama-3 Base Model: Frozen 4-bit Weights]
    Input --> LoRA[LoRA Adapters: Trainable 16-bit Weights]
    
    Base --> Add{Matrix Addition}
    LoRA --> Add
    
    Add --> Output(Medical Answer)
    
    subgraph LoRA Concept
        W[Original Weight Matrix: dxk]
        A[LoRA A: dxr]
        B[LoRA B: rxk]
        W -.->|Rank r << d| A
        W -.-> B
    end
        

Prerequisites

• Python 3.10+
• A GPU with at least 16GB VRAM (e.g., NVIDIA RTX 4080, A4000, or a T4 on Colab)
• PyTorch, Hugging Face Transformers, PEFT, TRL, and BitsAndBytes
• Access to Llama-3 via Hugging Face (requires accepting the Meta license)

Step 1: Install Dependencies

We use bitsandbytes for 4-bit quantization, drastically reducing the VRAM footprint of the base model.

pip install torch transformers accelerate peft trl bitsandbytes datasets

Step 2: Load the Base Model in 4-bit

Instead of loading the model in 16-bit (which requires ~16GB of VRAM just to hold the 8B model), we load it in 4-bit using NF4 quantization. This shrinks the model footprint to ~6GB.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "meta-llama/Meta-Llama-3-8B"

# Configure 4-bit quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Llama-3 does not have a default pad token, so we set it to the eos_token
tokenizer.pad_token = tokenizer.eos_token

# Load Model
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)

Step 3: Setup LoRA (Low-Rank Adaptation)

LoRA freezes the pre-trained model weights and injects trainable rank decomposition matrices into the layers of the Transformer architecture. This reduces the number of trainable parameters by 99% while achieving >90% of the performance of full fine-tuning.

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# Prepare model for training (gradient checkpointing, etc.)
model = prepare_model_for_kbit_training(model)

# Define LoRA configuration
# We target the 'q_proj' and 'v_proj' attention matrices
config = LoraConfig(
    r=8, # Rank
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# Wrap the model with PEFT
model = get_peft_model(model, config)
model.print_trainable_parameters()
# Output: trainable params: 3,407,872 || all params: 8,033,669,120 || trainable%: 0.0424

Step 4: Load Medical Dataset and Train using SFTTrainer

We will use a medical Q&A dataset (e.g., MedQA or a synthetic Instruct dataset) and Hugging Face's SFTTrainer (Supervised Fine-Tuning Trainer) to train our LoRA adapters.

from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments

# Load a hypothetical medical Q&A dataset
dataset = load_dataset("medalpaca/medical_meadow_medical_flashcards")

# Format the dataset into a single prompt string
def format_prompt(example):
    return f"User: {example['input']}\nAssistant: {example['output']}"

# Training Arguments
training_args = TrainingArguments(
    output_dir="./llama3-med-qa",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    optim="paged_adamw_32bit",
    learning_rate=2e-4,
    max_steps=200, # Just for demonstration
    fp16=True,
    logging_steps=10
)

# Initialize Trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset['train'],
    formatting_func=format_prompt,
    args=training_args,
    max_seq_length=512
)

# Start training
print("Starting training...")
trainer.train()

# Save the trained LoRA adapters
trainer.model.save_pretrained("llama3-med-lora-adapter")

Step 5: Testing the Fine-Tuned Model

Once training is complete, the LoRA adapters must be merged back with the base model to run inference efficiently.

from peft import PeftModel

# Reload base model in 16-bit (or 4-bit) for inference
base_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16
)

# Load the LoRA adapter and merge it with the base weights
ft_model = PeftModel.from_pretrained(base_model, "llama3-med-lora-adapter")
ft_model = ft_model.merge_and_unload()

# Test the model
prompt = "User: What are the common symptoms of acute appendicitis?\nAssistant:"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

outputs = ft_model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Conclusion

Parameter-Efficient Fine-Tuning democratizes access to state-of-the-art AI. By combining 4-bit quantization and LoRA, researchers can fine-tune massive foundational models like Llama-3 specifically for healthcare tasks—like triage, summarization, and clinical decision support—without needing a supercomputer. Because the LoRA adapters are saved separately (usually just a few megabytes), you can easily hot-swap different adapters for different medical specialties on the fly.