Fine-Tuning LLaMA 2 and Mistral with QLoRA

1. Introduction (Keep It Tight, First-Person)

“You don’t truly understand a model until you’ve tried breaking it with your own data.”

I’ve worked with large language models long enough to know that full fine-tuning isn’t always practical—or necessary.

When I started working with LLaMA 2 and Mistral, my goal was clear: fine-tune them efficiently on a domain-specific QA dataset I had built for a niche internal tool. I wasn’t looking for benchmark scores—I needed usable outputs that aligned with real-world, messy data.

Now here’s the catch: I didn’t have access to H100s or massive compute clusters. I had a single 48GB A6000 and a lot of questions. So full fine-tuning was off the table.

I tried PEFT, and while LoRA worked decently, QLoRA hit the sweet spot—it gave me adapter-based flexibility, 4-bit memory efficiency, and just enough customization without the overhead of full retraining.

I ran both LLaMA 2 7B and Mistral 7B through this pipeline. Each had its quirks—especially around tokenizer behavior and memory spikes—but I’ll walk you through the full setup, step-by-step, exactly as I did it.

This guide assumes you’re not here for explanations of what LoRA is. You probably already know that.

You’re here because you want to get it working, and get it working right—on your hardware, with your data, and without baby-sitting cryptic errors.

Let’s get into it.

2. Environment Setup (Exact Versions + Gotchas)

“Setting up the environment was 70% of the battle. Once everything clicked, training was the easy part.”

I learned this the hard way: if your environment isn’t locked in properly, you’re just one version mismatch away from a week of wasted debugging.

Here’s the exact setup I used. Feel free to copy-paste this, but I’d suggest you pin versions tightly, especially when mixing transformers, peft, trl, and bitsandbytes.

My Exact Environment

# Create a new conda environment (or use venv if you prefer)
conda create -n qlora-finetune python=3.10 -y
conda activate qlora-finetune

# Core libraries (use exact versions)
pip install torch==2.1.2 --index-url https://download.pytorch.org/whl/cu118
pip install transformers==4.36.2
pip install peft==0.7.1
pip install bitsandbytes==0.41.1
pip install accelerate==0.25.0
pip install trl==0.7.11
pip install datasets==2.16.1
pip install einops

# Optional but helpful
pip install tensorboard

CUDA Compatibility Matters

If you’re using PyTorch 2.1.2, make sure your CUDA toolkit version is 11.8. Even on my A6000, I ran into silent issues when mixing a PyTorch nightly build with the wrong bitsandbytes version.

If you’re on an older GPU like 3090 or 2080 Ti, you’ll want to double-check that bnb is compiled correctly for your architecture. I’ve seen cases where QLoRA training silently fails or spikes memory just because bitsandbytes wasn’t optimized for the card.

4-bit Gotcha: Memory Mapping and Compatibility

To load models in 4-bit, bitsandbytes uses NF4 quantization. But what isn’t obvious is that some model + quantization combos break unless you set bnb_4bit_compute_dtype=torch.bfloat16. That saved me from a few late-night segmentation faults.

We’ll set that flag explicitly when loading the model in the next section, but just keep this in mind: your training will crash if compute dtype and your GPU’s supported formats don’t align.

3. Preparing the Dataset (Tokenized + Aligned)

“If your dataset isn’t aligned with the model’s expectations, you’ll end up fine-tuning garbage at scale.”

I’ve lost count of how many times I’ve fine-tuned models and realized halfway through that the tokenization format didn’t match what the base model was trained on. That’s why now, I start every run by manually inspecting tokenized samples — before wasting compute.

Format Matters More Than People Think

When I fine-tuned LLaMA 2 and Mistral using QLoRA, I experimented with three different formats:

Chat-style (e.g., Alpaca/Vicuna-style)
Instruction-tuned (prompt → response)
Plain text QA pairs (prompt\nanswer)

In practice, I got the most consistent results using a prompt-response format with clear separators, especially with custom datasets that weren’t originally designed for instruction tuning.

Here’s a minimal example of what I used (in train.jsonl):

{"instruction": "What are the benefits of using LoRA for fine-tuning?", "response": "LoRA reduces memory usage by injecting trainable weights into specific attention modules."}
{"instruction": "Explain the role of r and alpha in QLoRA.", "response": "r is the rank of the decomposition, alpha is the scaling factor. Together, they control how much the adapter modifies the base model."}

This might surprise you: I tried training on raw markdown docs initially. But without clear prompt-response structure, the model started echoing random headers. That’s when I decided to formalize my dataset.

Tokenizer Setup (Don’t Skip This)

One gotcha I hit: LLaMA 2 doesn’t come with a pad_token. If you don’t explicitly set it, transformers will either throw an error or silently misalign batches during training.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_fast=True)

# Set pad token to eos if missing
tokenizer.pad_token = tokenizer.eos_token

And yes — Mistral had the same issue when I fine-tuned with the Hugging Face mistralai/Mistral-7B-v0.1 checkpoint.

Dataset Preparation in Code

Let me show you exactly how I formatted and tokenized the data. No theory — just real steps.

from datasets import load_dataset, DatasetDict

# Load JSONL file
dataset = load_dataset("json", data_files={"train": "train.jsonl"})

# Apply prompt formatting and tokenization
def format_and_tokenize(example):
    prompt = f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['response']}"
    tokens = tokenizer(prompt, truncation=True, padding="max_length", max_length=512)
    tokens["labels"] = tokens["input_ids"].copy()
    return tokens

tokenized_dataset = dataset.map(format_and_tokenize, remove_columns=["instruction", "response"])

This worked beautifully with QLoRA because it aligns the input and label tensors directly, avoiding the need for custom data collators or padding hacks.

Pro tip: If you’re using long context data, experiment with sliding window tokenization. In one of my runs, I chunked legal documents using a 512-token window with 128-token overlap — and the output quality improved significantly.

4. QLoRA Configuration for LLaMA 2 / Mistral

“This is where things can go wrong without even throwing an error. Trust me — I’ve been there.”

You might be wondering: how different can QLoRA config really be across models like LLaMA 2 and Mistral? Short answer — just enough to break silently if you’re not careful.

Loading in 4-bit with Quantization

You’ll need to load the base model in 4-bit right out of the gate. Here’s what I use:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=quant_config,
    device_map="auto"
)

This setup gives you max memory efficiency without compromising training stability. I’ve personally seen smoother convergence when using bfloat16 for compute — especially on A100 and A6000.

PEFTConfig + Target Modules

The magic happens here — and this config differs between LLaMA and Mistral.

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model = prepare_model_for_kbit_training(model)

peft_config = LoraConfig(
    r=64,
    lora_alpha=16,
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "v_proj"]  # We'll customize this below
)

model = get_peft_model(model, peft_config)

Target Module Differences (Heads-Up!)

LLaMA 2 typically uses:

model.layers.*.self_attn.q_proj
model.layers.*.mlp.down_proj

Mistral, on the other hand, might nest layers differently (depending on your checkpoint source). Here’s how I handled it dynamically:

# Inspecting target modules programmatically
for name, module in model.named_modules():
    if "q_proj" in name:
        print(name)

Once I saw the exact layer names, I adjusted target_modules accordingly. Don’t just assume “q_proj” and “v_proj” work across both — check.

Memory Tricks

Here’s what helped me cut memory usage almost in half during experimentation:

model.gradient_checkpointing_enable()
model.enable_input_require_grads()

Combined with 4-bit loading + LoRA, this was enough to run 7B models comfortably on a single A6000. I even got away with batch size 4 (using gradient_accumulation_steps=8) without running into OOM.

Training Pipeline (Using TRL or Accelerate)

“Training a large model isn’t about pressing trainer.train() — it’s about surviving the minefield between your first run and your first checkpoint.”

Let me say this upfront: I’ve tried both transformers.Trainer and TRL’s SFTTrainer. Personally, I lean toward SFTTrainer from the trl library — it handles supervised fine-tuning with less fuss, especially with LoRA/QLoRA setups.

But before we get into preferences, let’s build a training loop that won’t fail silently at 80% progress.

Core Setup: SFTTrainer from TRL

If you’re using QLoRA with 4-bit quantization and PEFT, this is the baseline I used for my runs:

from trl import SFTTrainer
from transformers import TrainingArguments

training_args = TrainingArguments(
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    logging_steps=10,
    output_dir="./checkpoints",
    save_strategy="steps",
    save_steps=100,
    save_total_limit=3,
    num_train_epochs=3,
    bf16=True,  # Use bf16 if your GPU supports it
    report_to="none",
    optim="paged_adamw_8bit",  # Efficient optimizer for LoRA/QLoRA
    lr_scheduler_type="cosine",
    learning_rate=2e-4,
    warmup_steps=100,
    gradient_checkpointing=True,
    max_grad_norm=1.0,
    logging_first_step=True,
    evaluation_strategy="no",  # Skip eval unless using val split
)

trainer = SFTTrainer(
    model=model,
    train_dataset=tokenized_dataset["train"],
    tokenizer=tokenizer,
    args=training_args,
    packing=True  # Efficient for instruction data
)

This might save you hours: I’ve seen packing=True alone reduce training time by 30% on sequence-heavy datasets.

Performance Tweaks (That Actually Matter)

Let’s be real — everyone says “use FlashAttention” or “turn on gradient checkpointing.” But in practice, here’s what actually moved the needle for me:

FlashAttention-2 made a huge difference only on A100 and H100 — not on 3090 or 4090.
Gradient checkpointing saved me ~40% memory, but slowed down each step. Worth it when you’re squeezed.
Paged optimizers like paged_adamw_8bit (via bitsandbytes) helped me run 7B with full LoRA configs on a single 48GB GPU.

# Enabling gradient checkpointing and input grads
model.gradient_checkpointing_enable()
model.enable_input_require_grads()

And just to be sure my setup worked before burning hours of training, I always validate the loop on a tiny dataset — like 100 samples:

small_ds = tokenized_dataset["train"].select(range(100))
trainer.train_dataset = small_ds
trainer.train()

If it converges, I move on to the full run. If not, I know something’s off with loss scaling, tokenization, or config.

Save Strategy + Deepspeed (Optional)

If you’re going beyond 7B — say 13B or 70B — I’d recommend plugging in Deepspeed Zero 2 or 3. I used this config with 2x A100 (80GB) and it ran smoothly:

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true
  },
  "gradient_accumulation_steps": 8,
  "train_batch_size": 32,
  "gradient_clipping": 1.0,
  "fp16": {
    "enabled": true
  }
}

Just pass this into accelerate or transformers.Trainer and you’re good to go.

Pro tip: If you’re running on Colab or a single GPU, avoid Deepspeed. It’s overkill and can actually slow things down unless you’ve got multi-node infra.

6. Evaluating the Fine-Tuned Model (Not Just Perplexity)

“If your model’s output looks fine, but users are still confused — it’s not fine.”

After fine-tuning, the first thing I do isn’t to check loss curves or perplexity graphs. I run it on real examples I care about — whether it’s long-form answers in my domain, tight summarizations, or edge-case QA pairs that I know vanilla models usually fail on.

Here’s the pipeline I use to load the model with LoRA weights and run inference in 4-bit:

from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from peft import PeftModel
import torch

# Load tokenizer and model
model_id = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)

# Ensure pad_token and eos_token are aligned (needed for tokenizer quirks)
tokenizer.pad_token = tokenizer.eos_token

# Load base model in 4-bit
from transformers import BitsAndBytesConfig

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

base_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)

# Load LoRA weights
model = PeftModel.from_pretrained(base_model, "path/to/lora/checkpoint")

model.eval()

Prompt Formatting Tips

Small detail, big impact: LLaMA and Mistral are very sensitive to prompt format. I’ve had seemingly broken outputs just because I forgot to wrap the prompt with BOS tokens or follow the right role format.

Here’s what I usually do:

prompt = "<s>[INST] How does photosynthesis work? [/INST]"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

If you’re using your own custom instruction format, just be consistent — LLaMA is not as forgiving as GPT.

Custom Evaluation (BLEU, ROUGE, EM, etc.)

I run exact match and F1 for QA, BLEU and ROUGE-L for generation tasks. Here’s a quick snippet I’ve used to test against ground truth:

from datasets import load_metric

metric = load_metric("exact_match")  # or rouge, bleu, etc.

results = []
for example in test_dataset:
    input_ids = tokenizer(example["prompt"], return_tensors="pt").input_ids.cuda()
    generated = model.generate(input_ids, max_new_tokens=150)
    prediction = tokenizer.decode(generated[0], skip_special_tokens=True)
    
    results.append({
        "prediction": prediction,
        "reference": example["response"]
    })

scores = metric.compute(predictions=[r["prediction"] for r in results],
                        references=[r["reference"] for r in results])

Pro tip: BLEU and ROUGE work well for news summarization, but for instruction-following outputs, EM or even custom regex match is often more telling.

Human Eval (Still the Gold Standard)

When I really want to know how good the model is, I run side-by-sides with the base model.

One trick that’s worked for me: Give the same prompt to the base and fine-tuned models, randomize the order of responses, and have non-technical people tell me which one they’d prefer. That’s usually more insightful than 3 decimal places of BLEU.

7. Packaging and Sharing the Model

Now if you’re happy with the results, you’ll want to save and share the model cleanly — whether that’s for colleagues, inference pipelines, or open-source contribution.

Here’s how I save just the LoRA adapter:

model.save_pretrained("lora_checkpoint/")
tokenizer.save_pretrained("lora_checkpoint/")

This saves only the adapter weights — super light and reproducible.

If I need a self-contained model (for pushing to Hugging Face Hub or running inference without re-merging), I use:

from peft import merge_and_unload

model = merge_and_unload(model)
model.save_pretrained("merged_model/")
tokenizer.save_pretrained("merged_model/")

Push to Hugging Face Hub

I always add model_card.md with:

Intended use
Training dataset source
Limitations and bias warnings
Prompt formatting examples

And use the CLI to upload:

huggingface-cli login
transformers-cli upload ./merged_model --repo-id your-username/llama2-finetuned-domainqa

Heads-up: After merging LoRA, tokenizer bugs can show up — especially around pad_token, EOS, and spacing. Always validate merged models on 2-3 test prompts before shipping.

Conclusion: Was It Worth It?

Short answer? Yes — but with caveats.

Fine-tuning LLaMA 2 or Mistral with QLoRA gave me significant performance boosts on domain-specific prompts compared to the base model.

Especially in areas where general models tend to waffle or hallucinate, this setup delivered more grounded and context-aware outputs — without requiring GPUs that cost as much as a Tesla.

But I’ll be honest: the gain isn’t always linear with the effort. There were days where half the time went into debugging tokenization quirks or memory issues with Deepspeed configs. And let’s not even talk about silently failing eval scripts — that stings.

From a cost perspective though, I trained a reasonably performant model using just one 24GB GPU. That would’ve been unthinkable with full fine-tuning.

What I’d Try Next

If I were to iterate further, I’d probably:

Stack DPO (Direct Preference Optimization) on top of my LoRA-finetuned base — to inject human preference directly into the model’s behavior without retraining from scratch.
Merge multiple adapters (e.g., one for summarization, another for QA) using the new PEFT stacking utilities. Could be a solid move for multi-task setups.
Push this into quantized inference at scale, maybe even Triton + ONNX for production. Right now, I’m just loading in 4-bit with PEFT, but it’s not battle-tested for latency.

Amit Yadav

I’m a Data Scientist.