Menu

AI & Machine Learning Topic

Fine-Tuning & Custom Models

Why Fine-Tuning Matters

Limitations of Prompt Engineering Alone:

  • Context Window: Limited by token constraints

  • Consistency: Hard to maintain across many examples

  • Cost: Repeatedly sending examples in prompts

  • Performance: Suboptimal for domain-specific tasks

When to Fine-Tune:

  • Domain-Specific Language: Medical, legal, technical jargon

  • Consistent Output Format: Structured data generation

  • Specialized Tone/Voice: Brand voice, character personality

  • Proprietary Knowledge: Internal documents, unique processes


Fine-Tuning Approaches Spectrum

text

Full Fine-Tuning ←---→ Parameter-Efficient Fine-Tuning (PEFT) ←---→ Prompt Engineering
(Most powerful)         (Best balance)                             (Easiest)

1. Full Fine-Tuning

Training all parameters of a pre-trained model on your specific dataset.

Traditional Fine-Tuning Process:

python

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from datasets import Dataset

# Load pre-trained model and tokenizer
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Add padding token if it doesn't exist
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Prepare your custom dataset
def prepare_training_data(texts, responses):
    """Convert conversation pairs into training format"""
    training_examples = []
    for text, response in zip(texts, responses):
        # Format for dialogue: <text> [SEP] <response> [EOS]
        full_conversation = f"{text} {tokenizer.sep_token} {response} {tokenizer.eos_token}"
        training_examples.append(full_conversation)
    return training_examples

# Your domain-specific data
customer_queries = [
    "What's the status of my order?",
    "I need help with my account",
    "When will my package arrive?"
]

agent_responses = [
    "I'll check your order status right away. Can you provide your order number?",
    "I'd be happy to help with your account. What specific issue are you experiencing?",
    "Let me track your package. Could you share your tracking number?"
]

# Prepare dataset
train_texts = prepare_training_data(customer_queries, agent_responses)

# Tokenize dataset
def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding=True,
        max_length=512,
        return_tensors="pt"
    )

dataset = Dataset.from_dict({"text": train_texts})
tokenized_dataset = dataset.map(tokenize_function, batched=True)

# Training arguments
training_args = TrainingArguments(
    output_dir="./fine-tuned-customer-support",
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    save_steps=500,
    save_total_limit=2,
    prediction_loss_only=True,
    remove_unused_columns=False,
)

# Create trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=lambda data: {
        'input_ids': torch.stack([torch.tensor(d['input_ids']) for d in data]),
        'attention_mask': torch.stack([torch.tensor(d['attention_mask']) for d in data]),
        'labels': torch.stack([torch.tensor(d['input_ids']) for d in data])
    }
)

# Start fine-tuning
trainer.train()
trainer.save_model()

Challenges with Full Fine-Tuning:

  • Computationally expensive

  • Large memory requirements

  • Risk of catastrophic forgetting

  • Storage intensive (each fine-tuned model is full-size)


2. Parameter-Efficient Fine-Tuning (PEFT)

The modern approach that achieves similar results with far fewer parameters.

LoRA (Low-Rank Adaptation)

LoRA freezes the original model weights and injects trainable rank decomposition matrices.

python

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

# Load model
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Configure LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,  # For language modeling
    inference_mode=False,
    r=8,              # Rank of adaptation (smaller = less parameters)
    lora_alpha=32,    # Scaling parameter
    lora_dropout=0.1, # Dropout for LoRA layers
    target_modules=[  # Which layers to apply LoRA to
        "q_proj",
        "v_proj", 
        "k_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj"
    ]
)

# Convert model to PEFT model
peft_model = get_peft_model(model, lora_config)

# Print trainable parameters
peft_model.print_trainable_parameters()
# Example output: "trainable params: 8,388,608 || all params: 1,312,256,000 || 0.64%"

Complete LoRA Fine-Tuning Example:

python

from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling
from datasets import Dataset
import pandas as pd

# Prepare your custom dataset
def prepare_dataset(csv_file_path):
    df = pd.read_csv(csv_file_path)
    
    # Format: "Input: {query}\nOutput: {response}"
    formatted_data = []
    for _, row in df.iterrows():
        formatted_text = f"Input: {row['query']}\nOutput: {row['response']}"
        formatted_data.append(formatted_text)
    
    return formatted_data

# Load and prepare data
train_texts = prepare_dataset("customer_support_data.csv")

# Tokenize dataset
def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding=True,
        max_length=512,
        return_tensors="pt"
    )

dataset = Dataset.from_dict({"text": train_texts})
tokenized_dataset = dataset.map(tokenize_function, batched=True)

# Data collator for language modeling
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,  # We're doing causal LM, not masked LM
)

# Training arguments
training_args = TrainingArguments(
    output_dir="./lora-customer-support",
    overwrite_output_dir=True,
    num_train_epochs=5,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    warmup_steps=100,
    logging_steps=50,
    save_steps=500,
    learning_rate=1e-4,
    fp16=True,  # Use mixed precision to save memory
    remove_unused_columns=False,
)

# Create trainer
trainer = Trainer(
    model=peft_model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

# Train
trainer.train()

# Save only the LoRA adapters (tiny file size)
peft_model.save_pretrained("./lora-adapters")

Using the Fine-Tuned LoRA Model:

python

from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")

# Load LoRA adapters
fine_tuned_model = PeftModel.from_pretrained(base_model, "./lora-adapters")

# Inference
def generate_response(query):
    input_text = f"Input: {query}\nOutput:"
    inputs = tokenizer(input_text, return_tensors="pt")
    
    with torch.no_grad():
        outputs = fine_tuned_model.generate(
            **inputs,
            max_length=150,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extract only the generated part
    return response.split("Output:")[-1].strip()

# Test
response = generate_response("My order is late, what should I do?")
print(response)

3. Advanced PEFT Techniques

QLoRA (Quantized LoRA) - Even More Efficient

python

from transformers import BitsAndBytesConfig
from peft import prepare_model_for_kbit_training

# 4-bit quantization configuration
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 model with quantization
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/DialoGPT-medium",
    quantization_config=bnb_config,
    device_map="auto"
)

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

# Apply LoRA
peft_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, peft_config)

P-Tuning (Prompt Tuning)

python

from peft import PromptTuningConfig, PromptTuningInit

# Prompt tuning configuration
peft_config = PromptTuningConfig(
    task_type=TaskType.CAUSAL_LM,
    prompt_tuning_init=PromptTuningInit.TEXT,
    num_virtual_tokens=20,  # Length of soft prompt
    prompt_tuning_init_text="Classify the following customer query:",
    tokenizer_name_or_path="microsoft/DialoGPT-medium",
)

model = get_peft_model(model, peft_config)

4. Training Small LLMs from Scratch

For when you need complete control or have specialized data.

Using Hugging Face Transformers:

python

from transformers import GPT2Config, GPT2LMHeadModel, Trainer, TrainingArguments

# Define a smaller model architecture
config = GPT2Config(
    vocab_size=50257,
    n_positions=1024,      # Context length
    n_embd=512,           # Embedding dimension
    n_layer=6,            # Number of layers
    n_head=8,             # Number of attention heads
    resid_pdrop=0.1,      # Dropout
    embd_pdrop=0.1,
    attn_pdrop=0.1,
)

# Create model
small_model = GPT2LMHeadModel(config)

print(f"Model parameters: {small_model.num_parameters():,}")
# Typically 10M-100M parameters vs 7B+ for large models
Custom Training Loop:
python

import torch
from torch.utils.data import DataLoader
from transformers import AdamW, get_linear_schedule_with_warmup

def custom_training_loop(model, train_dataset, epochs=3):
    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
    
    optimizer = AdamW(model.parameters(), lr=5e-5)
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=100,
        num_training_steps=len(train_loader) * epochs
    )
    
    model.train()
    
    for epoch in range(epochs):
        total_loss = 0
        for batch in train_loader:
            optimizer.zero_grad()
            
            # Forward pass
            outputs = model(
                input_ids=batch['input_ids'],
                attention_mask=batch['attention_mask'],
                labels=batch['input_ids']  # For causal LM
            )
            
            loss = outputs.loss
            total_loss += loss.item()
            
            # Backward pass
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
        
        avg_loss = total_loss / len(train_loader)
        print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")
    
    return model

5. Real-World Fine-Tuning Pipeline

python

class LLMFineTuningPipeline:
    def __init__(self, base_model_name, use_peft=True):
        self.base_model_name = base_model_name
        self.use_peft = use_peft
        self.tokenizer = None
        self.model = None
        
    def load_model_and_tokenizer(self):
        """Load base model and tokenizer"""
        self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            
        if self.use_peft:
            # Load with quantization for efficiency
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.float16
            )
            self.model = AutoModelForCausalLM.from_pretrained(
                self.base_model_name,
                quantization_config=bnb_config,
                device_map="auto"
            )
            self.model = prepare_model_for_kbit_training(self.model)
        else:
            self.model = AutoModelForCausalLM.from_pretrained(
                self.base_model_name,
                device_map="auto"
            )
    
    def setup_peft(self, lora_config=None):
        """Configure PEFT/LoRA"""
        if lora_config is None:
            lora_config = LoraConfig(
                r=16,
                lora_alpha=32,
                target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
                lora_dropout=0.05,
                bias="none",
                task_type="CAUSAL_LM"
            )
        
        self.model = get_peft_model(self.model, lora_config)
        self.model.print_trainable_parameters()
    
    def prepare_data(self, data_path, format_function):
        """Prepare training data"""
        # Load and format your data
        with open(data_path, 'r') as f:
            raw_data = json.load(f)
        
        formatted_data = format_function(raw_data)
        return Dataset.from_dict({"text": formatted_data})
    
    def train(self, dataset, training_args=None):
        """Fine-tune the model"""
        if training_args is None:
            training_args = TrainingArguments(
                output_dir="./fine-tuned-model",
                per_device_train_batch_size=2,
                gradient_accumulation_steps=4,
                num_train_epochs=3,
                learning_rate=2e-4,
                fp16=True,
                logging_steps=50,
                save_steps=500,
            )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=dataset,
            data_collator=DataCollatorForLanguageModeling(
                tokenizer=self.tokenizer,
                mlm=False
            ),
        )
        
        trainer.train()
        return trainer
    
    def save_model(self, save_path):
        """Save the fine-tuned model"""
        if self.use_peft:
            self.model.save_pretrained(save_path)
        else:
            self.model.save_pretrained(save_path)
        self.tokenizer.save_pretrained(save_path)

# Usage example
pipeline = LLMFineTuningPipeline("microsoft/DialoGPT-medium", use_peft=True)
pipeline.load_model_and_tokenizer()
pipeline.setup_peft()

# Prepare your data
dataset = pipeline.prepare_data("training_data.json", my_format_function)

# Train
trainer = pipeline.train(dataset)

# Save
pipeline.save_model("./my-fine-tuned-model")

Key Takeaways

  1. Full Fine-Tuning: Powerful but expensive - use when you have ample resources

  2. LoRA/PEFT: 90% of the performance with 1% of parameters - modern best practice

  3. QLoRA: Even more efficient with quantization - ideal for consumer hardware

  4. Small Custom Models: Complete control but need significant domain data

  5. Choose Based On:

    • Data Size: Small data → PEFT, Large data → Full fine-tuning

    • Resources: Limited → LoRA/QLoRA, Ample → Full fine-tuning

    • Specificity: General domain → PEFT, Highly specialized → Custom models

Fine-tuning transforms general-purpose LLMs into specialized experts for your specific use cases, dramatically improving performance while reducing costs and latency!

Comments (0)

No comments yet. Be the first to comment!

AlphaAI Assistant

Hello! I'm AlphaAI

What would you like to learn today?