Fine-Tune Llama 3.1 8B with QLoRA on a Consumer GPU

Python 3.10+
PyTorch 2.4+ (CUDA 12.1+)
transformers 4.45+
peft 0.12+
bitsandbytes 0.44+
trl 0.10+
datasets 2.20+
accelerate 0.32+

# Create environment
python -m venv llama-finetune
source llama-finetune/bin/activate  # or .\venv\Scripts\activate on Windows

# Install core packages
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
pip install transformers peft bitsandbytes trl accelerate datasets
pip install wandb     # optional: for experiment tracking
pip install xformers  # optional: memory-efficient attention

import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

{
  "messages": [
    {"role": "system", "content": "You are a helpful assistant specialized in Python programming."},
    {"role": "user", "content": "Write a Python function to reverse a linked list."},
    {"role": "assistant", "content": "Here's a Python function to reverse a linked list..."}
  ]
}

from datasets import load_dataset

# Option 1: Load from Hugging Face Hub
dataset = load_dataset("your-username/your-dataset", split="train")

# Option 2: Load from local JSONL file
dataset = load_dataset("json", data_files="my_data.jsonl", split="train")

# Option 3: Create from list of dicts
dataset = load_dataset(
    "json",
    data_files=[{"train": "train.jsonl", "test": "test.jsonl"}]
)

# Format messages for the model
def format_chat(example):
    """Convert messages list to a single string with ChatML format."""
    formatted = ""
    for msg in example["messages"]:
        if msg["role"] == "system":
            formatted += f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>
{msg['content']}<|eot_id|>
"
        elif msg["role"] == "user":
            formatted += f"<|start_header_id|>user<|end_header_id|>
{msg['content']}<|eot_id|>
"
        elif msg["role"] == "assistant":
            formatted += f"<|start_header_id|>assistant<|end_header_id|>
{msg['content']}<|eot_id|>
"
    return {"text": formatted}

dataset = dataset.map(format_chat)

# Split into train/test
dataset = dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = dataset["train"]
eval_dataset = dataset["test"]

print(f"Train: {len(train_dataset)} samples, Test: {len(eval_dataset)} samples")

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer

# 4-bit quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # NormalFloat4 (optimal for 4-bit)
    bnb_4bit_use_double_quant=True,     # Double quantization (saves ~0.5 GB)
    bnb_4bit_compute_dtype=torch.bfloat16,  # Compute in bf16
)

# Load model
model_name = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",              # Automatic GPU placement
    trust_remote_code=True,
    use_flash_attention_2=True,     # Faster training, needs xformers
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"    # Important for Llama models

# Prepare model for k-bit training
model = prepare_model_for_kbit_training(model)

Model loaded in 4-bit: ~5.5 GB VRAM
With Python overhead: ~6.5 GB
Remaining for training: ~17.5 GB (on 24GB GPU)

# LoRA configuration
lora_config = LoraConfig(
    r=16,                           # Rank: 8-64 typical. Higher = more expressiveness, more memory
    lora_alpha=32,                  # Scaling factor (often 2x rank)
    target_modules=[                 # Which modules to apply LoRA to
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_dropout=0.05,              # Regularization
    bias="none",
    task_type="CAUSAL_LM",
)

# Apply LoRA to the model
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Expected output: trainable params: ~42M / 8B = ~0.5% of total params

# Training arguments
training_args = TrainingArguments(
    output_dir="./llama-8b-qlora-output",
    num_train_epochs=3,                     # Number of passes over data
    per_device_train_batch_size=4,          # Adjust based on VRAM
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=4,          # Effective batch = 4×4=16
    gradient_checkpointing=True,            # Saves VRAM at cost of speed
    optim="paged_adamw_8bit",               # 8-bit optimizer (saves VRAM)
    logging_steps=10,
    learning_rate=2e-4,                     # Typical for LoRA
    weight_decay=0.001,
    fp16=False,
    bf16=True,                              # Use bf16 if supported (H100, 5090)
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    evaluation_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=3,
    report_to="wandb",                      # Or "none" to disable
    run_name="llama-8b-qlora-run-1",
)

# Initialize trainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    max_seq_length=2048,                    # Max context length for training
    dataset_text_field="text",              # Field containing formatted text
    packing=False,                          # Packing = faster but less controlled
)

# Start training
trainer.train()

# Save just the LoRA adapter weights (~42 MB)
model.save_pretrained("./llama-8b-qlora-adapter")
tokenizer.save_pretrained("./llama-8b-qlora-adapter")

from peft import PeftModel

# Load base model in full precision for merge
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Load LoRA adapter
merged_model = PeftModel.from_pretrained(base_model, "./llama-8b-qlora-adapter")

# Merge and unload
merged_model = merged_model.merge_and_unload()

# Save merged model (full ~16 GB for 8B FP16)
merged_model.save_pretrained("./llama-8b-merged")
tokenizer.save_pretrained("./llama-8b-merged")

# 1. Convert merged model to GGUF
pip install llama-cpp-python
python convert_hf_to_gguf.py ./llama-8b-merged \
    --outfile ./llama-8b-merged-q4_k_m.gguf \
    --outtype q4_k_m

# Or use llama.cpp's convert script
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
python convert.py ../llama-8b-merged
./quantize ../llama-8b-merged/ggml-model-f16.gguf \
    ../llama-8b-merged-q4_k_m.gguf q4_k_m

# 2. Create Ollama model
# Create Modelfile:
cat > Modelfile << EOF
FROM ./llama-8b-merged-q4_k_m.gguf
TEMPLATE """{{ .System }}
{{ .Prompt }}"""
PARAMETER temperature 0.7
PARAMETER top_p 0.9
EOF

# 3. Create and run
ollama create my-finetuned-model -f Modelfile
ollama run my-finetuned-model

# Serve the merged model with vLLM
python -m vllm.entrypoints.openai.api_server \
    --model ./llama-8b-merged \
    --dtype bfloat16 \
    --max-model-len 4096 \
    --gpu-memory-utilization 0.90

# Or serve the LoRA adapter on top of base model
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --enable-lora \
    --lora-modules my-adapter=./llama-8b-qlora-adapter