Fine-Tuning TinyLlama and Phi-2: Fails, Fixes, and Pirate Chatbot

Photo by Andrew Neel on Unsplash

A hands-on journey of fine-tuning TinyLlama and Microsoft’s Phi-2 for chat-style output using LoRA. Lessons learned, errors debugged, and a full working example — with code and pirate flair.

The Problem

I wanted to fine-tune a small LLM to speak like a pirate. That’s it. Simple. Or so I thought…

Attempt 1: TinyLlama (A Tiny Terror)

I started with TinyLlama, thinking it would be ideal for quick training on my local setup (16 GB RAM, 4 GB VRAM).

What I Did

• Created a custom chat-style dataset in jsonl format:

{"messages": [{"role": "user", "content": "Ahoy!"}, {"role": "assistant", "content": "Aye aye, Captain!"}]}

• Used Hugging Face + PEFT (LoRA)
• Created this training pipeline:

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import get_peft_model, LoraConfig, TaskType

base_model = "TinyLlama/TinyLlama_v1.1"
tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=False)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
tokenizer.padding_side = "right"
tokenizer.add_special_tokens({"additional_special_tokens": ["<|end_turn|>"]})

model = AutoModelForCausalLM.from_pretrained(base_model, device_map="cpu")

peft_config = LoraConfig(
    r=8,
    lora_alpha=8,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.1,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)
model = get_peft_model(model, peft_config)
model.resize_token_embeddings(len(tokenizer))
model.train()

peft_config = LoraConfig(r=8, lora_alpha=32, task_type="CAUSAL_LM")
model = get_peft_model(model, peft_config)

What Went Wrong

• Loss refused to decrease. Even after trying everything on the internet I could find.
• Output was gibberish. Strange continuous tokens in the end.
• Turns out, I was using the base TinyLlama, not the chat-optimized one.
• Even when I switched to the Chat version, it needed way more data.

Lesson: TinyLlama base model isn’t chat-ready unless you guide it properly with roles and examples…a lot of examples (I tried with 1000 examples in batches). Atleast in my case, I couldn’t finetune it due to hardware restrictions.

Attempt 2: Phi-2 (Microsoft Saves the Day)

Switching to Phi-2 was like moving from a leaky rowboat to a pirate ship.

What Worked

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig
from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model, PeftModel
from trl import SFTTrainer
import torch

model_name = "./phi-2"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map={"": 0}, #use only if you have GPU, otherwise auto
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.pad_token_id

model = prepare_model_for_kbit_training(model)

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)

dataset = load_dataset("json", data_files={"train": "pirate_train.jsonl", "validation": "pirate_val.jsonl"})

def formatting_func(example):
    conversation = ""
    for msg in example["messages"]:
        role = msg["role"]
        content = msg["content"]
        if role == "user":
            conversation += f"User: {content}
"
        else:
            conversation += f"PirateGPT: {content}
"
    return conversation

dataset = dataset.map(lambda x: {"text": formatting_func(x)})

training_args = TrainingArguments(
    output_dir="phi2-pirate-lora",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    learning_rate=2e-4,
    fp16=torch.cuda.is_available(),
    logging_steps=10,
    eval_strategy="epoch",
    save_strategy="epoch",
    push_to_hub=False,
    report_to=["none"]
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    formatting_func=formatting_func,
    args=training_args
)

trainer.train()
eval_metrics = trainer.evaluate()
print("Final Evaluation Loss:", eval_metrics.get("eval_loss", "N/A"))

LoRA setup stayed the same:

from peft import get_peft_model, LoraConfig

peft_config = LoraConfig(r=8, lora_alpha=32, bias="none", task_type="CAUSAL_LM")
model = get_peft_model(model, peft_config)

With just ~1000 pirate-style examples, Phi-2 started giving coherent, in-character replies.

“Ahoy, landlubber! Ye be askin’ too many questions!” — actual output.

Common Errors I Faced

• RuntimeError: element 0 of tensors does not require grad
• Tokenizer mismatch (tip: always match tokenizer + model)
• trust_remote_code=True caused unexpected behavior with adapters.
• OOM errors because of my potato system.

Lessons Learned

1. Use chat-specific models when you need dialogue output. Less hassle at least.
2. Phi-2 is good for low-data scenarios
3. LoRA rocks — makes fine-tuning possible on average hardware
4. Always log, test, and validate outputs early — avoid wasting GPU cycles
5. Don’t be afraid to restart fresh if your loss graph is a horizontal line

Want to try it out yourself? Check this link here: https://gitlab.com/ai-learning7442413/phi2-finetuning

In the end, I converted the phi-2 finetuned adapter to GGUF format using llama.cpp to use it with Ollama on my system.