QLoRA Fine-Tuning Demo: Teaching a Small LLM New Knowledge

!pip install -q transformers==4.46.3 accelerate peft bitsandbytes datasets trl

import torch

if torch.cuda.is_available():
    gpu_name = torch.cuda.get_device_name(0)
    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
    print(f'GPU: {gpu_name} ({gpu_mem:.1f} GB)')
else:
    print('WARNING: No GPU detected! QLoRA requires a GPU.')
    print('Go to Runtime -> Change runtime type -> T4 GPU')

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

MODEL_NAME = 'Qwen/Qwen2.5-0.5B-Instruct'  # smaller model -> faster training

# 4-bit quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',            # NormalFloat4 — best for fine-tuning
    bnb_4bit_compute_dtype=torch.float16,  # compute in FP16
    bnb_4bit_use_double_quant=True,        # quantize the quantization constants too
)

print(f'Loading {MODEL_NAME} in 4-bit...')
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map='auto',
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Set padding token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

total_params = sum(p.numel() for p in model.parameters())
print(f'Model loaded: {total_params/1e6:.0f}M parameters (4-bit quantized)')
print(f'GPU memory used: {torch.cuda.memory_allocated()/1e9:.2f} GB')

def ask_model(question, max_new_tokens=150):
    """Ask the model a question and return the answer."""
    messages = [{'role': 'user', 'content': question}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors='pt').to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,           # greedy decoding for reliable factual recall
            pad_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
    return response.strip()


# Questions to test both BEFORE and AFTER fine-tuning (for side-by-side comparison)
test_questions = [
    'Who is the CEO of Acme Analytics?',
    'What is the ARR of Acme Analytics?',
    #'What databases does InsightBoard Pro support?',
    #'How many employees does Acme Analytics have?',
    #'What is the mission of Acme Analytics?',
]

# Capture the BEFORE answers now, while the model is still the base model
before_answers = {}
print('=== BEFORE Fine-Tuning ===')
for q in test_questions:
    answer = ask_model(q)
    before_answers[q] = answer
    print(f'Q: {q}')
    print(f'A: {answer}')
    print()

from datasets import Dataset

# Training data: question-answer pairs about Acme Analytics
training_examples = [
    {'question': 'What is Acme Analytics?',
     'answer': 'Acme Analytics Inc. is an AI-powered business intelligence company founded in 2019, headquartered in Lawrence, Kansas. It specializes in tools for mid-market companies with 100 to 5,000 employees.'},

    {'question': 'Who founded Acme Analytics?',
     'answer': 'Acme Analytics was founded by Dr. Sarah Chen and Marcus Rivera in 2019 in Lawrence, Kansas.'},

    {'question': 'Where is Acme Analytics headquartered?',
     'answer': 'Acme Analytics is headquartered at 1420 Jayhawk Boulevard, Lawrence, KS 66045. It also has offices in Chicago (sales) and Austin (engineering).'},

    {'question': 'What is InsightBoard Pro?',
     'answer': 'InsightBoard Pro is the flagship product of Acme Analytics. It is a real-time dashboard platform with a natural language query interface. Pricing is $49/user/month for Standard and $89/user/month for Enterprise.'},

    {'question': 'What is the pricing for InsightBoard Pro?',
     'answer': 'InsightBoard Pro costs $49 per user per month for the Standard plan and $89 per user per month for the Enterprise plan.'},

    {'question': 'What is DataPipe ETL?',
     'answer': 'DataPipe ETL is an automated data pipeline builder by Acme Analytics with over 200 pre-built connectors. Pricing starts at $500/month for up to 10 million records per day. It launched in September 2024.'},

    {'question': 'What is PredictIQ?',
     'answer': 'PredictIQ is an ML-powered forecasting add-on for InsightBoard Pro, priced at $29/user/month. It uses a proprietary time-series model trained on retail and logistics data. It beta launched in March 2025.'},

    {'question': 'What was Acme Analytics revenue in 2025?',
     'answer': 'Acme Analytics reported annual revenue of $42.3 million for fiscal year 2025, with Q4 2025 revenue of $12.8 million.'},

    {'question': 'How many employees does Acme Analytics have?',
     'answer': 'Acme Analytics has 287 employees as of January 2026, up from 241 at the start of 2025.'},

    {'question': 'Who is the CEO of Acme Analytics?',
     'answer': 'The CEO of Acme Analytics is Dr. Sarah Chen, who co-founded the company in 2019.'},

    {'question': 'Who is the CTO of Acme Analytics?',
     'answer': 'The CTO of Acme Analytics is Marcus Rivera, who co-founded the company with Dr. Sarah Chen.'},

    {'question': 'What is the mission of Acme Analytics?',
     'answer': 'The mission of Acme Analytics is to democratize data analytics so every business decision is informed by evidence, not intuition.'},

    {'question': 'How many customers does Acme Analytics have?',
     'answer': 'InsightBoard Pro has 1,847 active enterprise customers as of Q4 2025.'},

    {'question': 'What is the ARR of Acme Analytics?',
     'answer': 'Acme Analytics had an Annual Recurring Revenue (ARR) of $48.6 million at the end of 2025, up 34% year-over-year.'},

    {'question': 'What databases does InsightBoard Pro support?',
     'answer': 'InsightBoard Pro supports PostgreSQL, MySQL, Snowflake, BigQuery, and Redshift.'},
]


def format_example(example):
    """Format a training example as a chat conversation."""
    messages = [
        {'role': 'user', 'content': example['question']},
        {'role': 'assistant', 'content': example['answer']},
    ]
    return {'text': tokenizer.apply_chat_template(messages, tokenize=False)}


dataset = Dataset.from_list(training_examples)
dataset = dataset.map(format_example)

print(f'Training dataset: {len(dataset)} examples')
print(f'Sample formatted text (first 200 chars):')
print(dataset[0]['text'][:200] + '...')

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# Prepare the quantized model for training
model = prepare_model_for_kbit_training(model)

# LoRA configuration
lora_config = LoraConfig(
    r=16,                          # rank of adaptation matrices
    lora_alpha=32,                 # scaling factor
    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj'],  # attention layers
    lora_dropout=0.05,
    bias='none',
    task_type='CAUSAL_LM',
)

# Apply LoRA to the model
model = get_peft_model(model, lora_config)

# Show trainable vs total parameters
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
pct = 100 * trainable_params / total_params

print(f'Total parameters:     {total_params/1e6:.1f}M')
print(f'Trainable parameters: {trainable_params/1e6:.1f}M ({pct:.2f}%)')
print(f'Frozen parameters:    {(total_params-trainable_params)/1e6:.1f}M')
print(f'GPU memory: {torch.cuda.memory_allocated()/1e9:.2f} GB')

from trl import SFTTrainer, SFTConfig

training_args = SFTConfig(
    output_dir='./qloRA_output',
    num_train_epochs=15,                # enough for memorization
    per_device_train_batch_size=4,      # larger batch fits easily in 0.5B model
    gradient_accumulation_steps=1,      # no accumulation needed
    learning_rate=3e-4,
    warmup_steps=5,
    logging_steps=5,
    save_strategy='no',
    fp16=True,
    max_seq_length=256,                 # our examples are short; reduces compute
    dataset_text_field='text',
    report_to='none',
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
)

print('Starting fine-tuning...')
trainer.train()
print('Fine-tuning complete!')
print(f'GPU memory: {torch.cuda.memory_allocated()/1e9:.2f} GB')

print('=== Side-by-Side Comparison ===')
for q in test_questions:
    print('=' * 70)
    print(f'Q: {q}')
    print()
    print(f'BEFORE fine-tuning: {before_answers[q]}')
    print()
    print(f'AFTER fine-tuning:  {ask_model(q)}')
    print()