Fine-Tuning Tiny Models: LoRA, QLoRA, and Domain Adaptation Strategies

# Create virtual environment
python -m venv lora-env
source lora-env/bin/activate  # Windows: lora-env\Scripts\activate
 
# Install PyTorch with CUDA (check https://pytorch.org for your CUDA version)
pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
 
# Install core fine-tuning libraries
pip install \
    transformers==4.36.0 \
    datasets==2.16.0 \
    accelerate==0.25.0 \
    peft==0.7.1 \
    bitsandbytes==0.41.3 \
    trl==0.7.4 \
    scipy
 
# Optional: Experiment tracking
pip install wandb tensorboard

# Additional requirements for quantization
pip install \
    bitsandbytes>=0.41.0 \
    scipy>=1.11.0
 
# Note: bitsandbytes requires specific CUDA versions
# CUDA 11.8: pip install bitsandbytes==0.41.3
# CUDA 12.1: pip install bitsandbytes==0.42.0

# test_lora_setup.py
import torch
import transformers
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
import bitsandbytes as bnb
 
print("=== Fine-Tuning Environment Check ===\n")
 
# PyTorch and CUDA
print(f"PyTorch: {torch.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA Version: {torch.version.cuda}")
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
 
# Transformers and PEFT
print(f"\nTransformers: {transformers.__version__}")
 
# Test PEFT/LoRA
try:
    from transformers import AutoModelForCausalLM
    model = AutoModelForCausalLM.from_pretrained(
        "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        device_map="auto",
        torch_dtype=torch.float16
    )
    
    lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules=["q_proj", "v_proj"],
        task_type="CAUSAL_LM"
    )
    
    lora_model = get_peft_model(model, lora_config)
    lora_model.print_trainable_parameters()
    print("\n✅ LoRA setup working!")
    
    # Clean up
    del model, lora_model
    torch.cuda.empty_cache()
    
except Exception as e:
    print(f"\n❌ LoRA test failed: {e}")
 
# Test QLoRA (4-bit quantization)
try:
    from transformers import BitsAndBytesConfig
    
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    
    qlora_model = AutoModelForCausalLM.from_pretrained(
        "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        quantization_config=bnb_config,
        device_map="auto"
    )
    print("✅ QLoRA (4-bit) setup working!")
    
    del qlora_model
    torch.cuda.empty_cache()
    
except Exception as e:
    print(f"⚠️  QLoRA test failed: {e}")
    print("   QLoRA requires bitsandbytes and compatible CUDA")
 
print("\n✅ Environment ready for fine-tuning!")

import torch
import torch.nn as nn
 
class LoRALayer(nn.Module):
    """
    LoRA adaptation layer.
    
    Implements: h = Wx + (BA)x with frozen W
    """
    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 16,
        alpha: int = 16,
        dropout: float = 0.1
    ):
        super().__init__()
        
        self.rank = rank
        self.alpha = alpha
        
        # LoRA parameters (trainable)
        self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
        
        # Scaling factor
        self.scaling = alpha / rank
        
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
    
    def forward(self, x):
        """
        Args:
            x: Input [batch, seq, in_features]
        Returns:
            LoRA output [batch, seq, out_features]
        """
        # LoRA path: x @ A^T @ B^T
        # = (x @ A^T) @ B^T
        result = self.dropout(x) @ self.lora_A.T @ self.lora_B.T
        return self.scaling * result
 
 
class LinearWithLoRA(nn.Module):
    """Linear layer with LoRA adapter."""
    
    def __init__(
        self,
        base_layer: nn.Linear,
        rank: int = 16,
        alpha: int = 16,
        dropout: float = 0.1
    ):
        super().__init__()
        
        # Frozen base layer
        self.base_layer = base_layer
        for param in self.base_layer.parameters():
            param.requires_grad = False
        
        # LoRA adapter
        self.lora = LoRALayer(
            base_layer.in_features,
            base_layer.out_features,
            rank=rank,
            alpha=alpha,
            dropout=dropout
        )
    
    def forward(self, x):
        # Base output + LoRA adaptation
        base_out = self.base_layer(x)
        lora_out = self.lora(x)
        return base_out + lora_out
    
    def merge_weights(self):
        """Merge LoRA into base weights for inference."""
        # Compute LoRA weight matrix: BA
        lora_weight = self.lora.scaling * (self.lora.lora_B @ self.lora.lora_A)
        
        # Add to base weights
        self.base_layer.weight.data += lora_weight
        
        # Clear LoRA parameters to save memory
        self.lora = None
 
 
def add_lora_to_model(
    model,
    target_modules=["q_proj", "v_proj"],
    rank=16,
    alpha=16
):
    """
    Add LoRA adapters to model.
    
    Args:
        model: Base model
        target_modules: Which modules to add LoRA to
        rank: LoRA rank
        alpha: LoRA alpha (scaling)
    
    Returns:
        Model with LoRA adapters
    """
    for name, module in model.named_modules():
        # Check if this module should get LoRA
        if any(target in name for target in target_modules):
            if isinstance(module, nn.Linear):
                # Get parent module
                parent_name = '.'.join(name.split('.')[:-1])
                child_name = name.split('.')[-1]
                parent = model.get_submodule(parent_name) if parent_name else model
                
                # Replace with LoRA version
                lora_layer = LinearWithLoRA(
                    module,
                    rank=rank,
                    alpha=alpha
                )
                setattr(parent, child_name, lora_layer)
                
                print(f"Added LoRA to {name}")
    
    return model
 
 
# Usage example
from transformers import AutoModelForCausalLM
 
base_model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
lora_model = add_lora_to_model(
    base_model,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    rank=16,
    alpha=16
)
 
# Count parameters
total = sum(p.numel() for p in lora_model.parameters())
trainable = sum(p.numel() for p in lora_model.parameters() if p.requires_grad)
 
print(f"Total parameters: {total/1e9:.2f}B")
print(f"Trainable parameters: {trainable/1e6:.2f}M")
print(f"Trainable %: {100 * trainable / total:.4f}%")
# Total parameters: 1.10B
# Trainable parameters: 4.19M
# Trainable %: 0.3809%

from transformers import TrainingArguments, Trainer
from datasets import load_dataset
 
def fine_tune_with_lora(
    model,
    tokenizer,
    dataset_name="timdettmers/openassistant-guanaco",
    output_dir="./lora-finetuned",
    num_epochs=3,
    batch_size=4,
    learning_rate=2e-4
):
    """
    Fine-tune model with LoRA.
    """
    # Load dataset
    dataset = load_dataset(dataset_name, split="train[:1000]")  # Small subset for demo
    
    def tokenize_function(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=512,
            padding="max_length"
        )
    
    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=dataset.column_names
    )
    
    # Training arguments
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=4,
        learning_rate=learning_rate,
        fp16=True,
        logging_steps=10,
        save_strategy="epoch",
        warmup_steps=100,
    )
    
    # Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
    )
    
    # Train
    trainer.train()
    
    # Save
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    
    return model
 
# Fine-tune
fine_tuned_model = fine_tune_with_lora(lora_model, tokenizer)

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 
def create_qlora_model(
    model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    rank=64,
    alpha=16,
    target_modules=None
):
    """
    Create model with QLoRA (4-bit base + LoRA adapters).
    
    Args:
        model_name: HuggingFace model ID
        rank: LoRA rank
        alpha: LoRA alpha
        target_modules: Which layers to adapt
    
    Returns:
        QLoRA model ready for training
    """
    # 4-bit quantization config
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",           # NormalFloat4 (better than regular INT4)
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,      # Nested quantization for extra compression
    )
    
    # Load base model in 4-bit
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True
    )
    
    # Prepare for k-bit training
    model = prepare_model_for_kbit_training(model)
    
    # LoRA config
    if target_modules is None:
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"]
    
    lora_config = LoraConfig(
        r=rank,
        lora_alpha=alpha,
        target_modules=target_modules,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    )
    
    # Add LoRA adapters
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    
    return model
 
# Create QLoRA model
qlora_model = create_qlora_model(rank=64)
# trainable params: 8,388,608 || all params: 1,108,388,608 || trainable%: 0.7570
 
# Memory comparison
import torch
 
print("\nMemory usage:")
print(f"  Base FP16: ~2.2 GB")
print(f"  Base INT4: ~550 MB")
print(f"  LoRA adapters: ~16 MB")
print(f"  Total QLoRA: ~566 MB (4× reduction!)")

from trl import SFTTrainer
 
def train_qlora(
    model,
    tokenizer,
    dataset,
    output_dir="./qlora-legal",
    max_seq_length=512,
    num_epochs=3
):
    """
    Train with QLoRA using Supervised Fine-Tuning.
    """
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        learning_rate=2e-4,
        fp16=False,
        bf16=True,  # BF16 for LoRA adapters
        logging_steps=10,
        optim="paged_adamw_8bit",  # 8-bit optimizer
        save_strategy="epoch",
    )
    
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        dataset_text_field="text",
        max_seq_length=max_seq_length,
        tokenizer=tokenizer,
    )
    
    trainer.train()
    
    # Save only LoRA adapters (tiny!)
    model.save_pretrained(output_dir)
    
    return model

def prepare_legal_dataset():
    """
    Prepare dataset for legal domain adaptation.
    
    Sources:
    - Legal cases (opinions, briefs)
    - Contracts and agreements
    - Legal Q&A
    """
    from datasets import load_dataset, concatenate_datasets
    
    # Load legal datasets
    legal_cases = load_dataset("pile-of-law/pile-of-law", split="train[:5000]")
    legal_qa = load_dataset("copenlu/legal-qa", split="train")
    
    # Format for instruction tuning
    def format_legal(example):
        return {
            "text": f"### Legal Query:\n{example['question']}\n\n### Legal Analysis:\n{example['answer']}"
        }
    
    legal_qa = legal_qa.map(format_legal)
    
    # Combine
    combined = concatenate_datasets([legal_cases, legal_qa])
    
    return combined
 
# Fine-tune for legal
legal_data = prepare_legal_dataset()
legal_model = create_qlora_model(rank=64)
trained_legal = train_qlora(legal_model, tokenizer, legal_data)

def prepare_medical_dataset():
    """
    Medical domain with focus on terminology and clinical reasoning.
    """
    # Medical datasets
    medqa = load_dataset("bigbio/med_qa", split="train")
    pubmed = load_dataset("pubmed", split="train[:10000]")
    
    def format_medical(example):
        return {
            "text": f"### Patient Presentation:\n{example['question']}\n\n"
                   f"### Medical Assessment:\n{example['answer']}\n\n"
                   f"### Explanation:\n{example['explanation']}"
        }
    
    formatted = medqa.map(format_medical)
    return formatted

def prepare_code_dataset(languages=["python", "javascript"]):
    """
    Code generation with multiple languages.
    """
    from datasets import load_dataset
    
    # Code datasets
    code_data = load_dataset("codeparrot/github-code", split="train[:20000]")
    
    def format_code(example):
        return {
            "text": f"```{example['language']}\n{example['code']}\n```\n\n"
                   f"# Explanation: {example['docstring']}"
        }
    
    return code_data.map(format_code)
 
# Fine-tune for code
code_model = create_qlora_model(rank=64)
code_data = prepare_code_dataset()
trained_code = train_qlora(code_model, tokenizer, code_data)

def create_multitask_dataset(tasks):
    """
    Combine datasets from multiple tasks.
    
    Args:
        tasks: Dict of {task_name: dataset}
    
    Returns:
        Combined dataset with task prefixes
    """
    from datasets import concatenate_datasets
    
    formatted_datasets = []
    
    for task_name, dataset in tasks.items():
        def add_task_prefix(example):
            example["text"] = f"[{task_name.upper()}] {example['text']}"
            return example
        
        formatted = dataset.map(add_task_prefix)
        formatted_datasets.append(formatted)
    
    return concatenate_datasets(formatted_datasets)
 
# Usage
multitask_data = create_multitask_dataset({
    "summarization": load_dataset("cnn_dailymail", split="train[:1000]"),
    "qa": load_dataset("squad_v2", split="train[:1000]"),
    "translation": load_dataset("wmt14", "de-en", split="train[:1000]")
})

def create_curriculum(dataset, difficulty_fn):
    """
    Sort dataset by difficulty for curriculum learning.
    
    Args:
        dataset: Training dataset
        difficulty_fn: Function to compute example difficulty
    
    Returns:
        Dataset sorted by difficulty
    """
    # Compute difficulty scores
    difficulties = []
    for example in dataset:
        score = difficulty_fn(example)
        difficulties.append(score)
    
    # Sort by difficulty
    sorted_indices = sorted(range(len(difficulties)), key=lambda i: difficulties[i])
    sorted_dataset = dataset.select(sorted_indices)
    
    return sorted_dataset
 
# Example difficulty function
def text_difficulty(example):
    """Difficulty = text length + vocabulary complexity."""
    text = example["text"]
    length_score = len(text.split()) / 100  # Normalize
    vocab_score = len(set(text.split())) / len(text.split())
    return length_score + vocab_score
 
curriculum_data = create_curriculum(dataset, text_difficulty)

def tune_lora_rank(
    base_model,
    dataset,
    ranks=[8, 16, 32, 64],
    validation_set=None
):
    """
    Find optimal LoRA rank for your task.
    
    Smaller rank: Faster, less overfitting, may underfit
    Larger rank: More capacity, better quality, slower
    """
    results = {}
    
    for rank in ranks:
        print(f"\nTrying rank={rank}")
        
        # Create model with this rank
        model = create_qlora_model(rank=rank)
        
        # Quick training
        trained = train_qlora(
            model,
            tokenizer,
            dataset,
            output_dir=f"./rank_{rank}",
            num_epochs=1
        )
        
        # Evaluate
        if validation_set:
            metrics = evaluate_model(trained, validation_set)
            results[rank] = metrics
            print(f"  Validation perplexity: {metrics['perplexity']:.2f}")
    
    # Find best
    best_rank = min(results, key=lambda r: results[r]['perplexity'])
    print(f"\nBest rank: {best_rank}")
    
    return best_rank, results
 
# Run tuning
best_rank, all_results = tune_lora_rank(base_model, train_data, val_data)

# Cosine with warmup (recommended)
from transformers import get_cosine_schedule_with_warmup
 
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=1000
)
 
# Or linear decay
from transformers import get_linear_schedule_with_warmup
 
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=1000
)

# LoRA adapters are tiny (4-16 MB)
model.save_pretrained("./lora-adapters")
 
# Later: Load base + adapters
from peft import PeftModel
 
base = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model = PeftModel.from_pretrained(base, "./lora-adapters")

# Merge LoRA into base weights
model = model.merge_and_unload()
 
# Now it's a standard model (no performance overhead)
model.save_pretrained("./merged-model")

# Serve different domains with same base model
class MultiAdapterSystem:
    def __init__(self, base_model_name):
        self.base = AutoModelForCausalLM.from_pretrained(base_model_name)
        self.adapters = {}
    
    def load_adapter(self, name, path):
        """Load domain-specific adapter."""
        adapter = PeftModel.from_pretrained(self.base, path)
        self.adapters[name] = adapter
    
    def generate(self, prompt, domain="general"):
        """Generate with domain-specific adapter."""
        model = self.adapters.get(domain, self.base)
        inputs = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(**inputs)
        return tokenizer.decode(outputs[0])
 
# Usage
system = MultiAdapterSystem("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
system.load_adapter("legal", "./lora-legal")
system.load_adapter("medical", "./lora-medical")
system.load_adapter("code", "./lora-code")
 
# Route by domain
legal_response = system.generate("What is tort law?", domain="legal")
medical_response = system.generate("Symptoms of diabetes?", domain="medical")