QLoRA Test: Target Modules Comparison - Ministral-3B-Reasoning¶

# Environment Setup
import os
from dotenv import load_dotenv
load_dotenv()

# Force text-based progress instead of HTML widgets
os.environ["TQDM_NOTEBOOK"] = "false"

# CRITICAL: Import unsloth FIRST for proper TRL patching
import unsloth
from unsloth import FastLanguageModel, is_bf16_supported

import torch
import gc

# Environment summary
gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
print(f"Environment: unsloth {unsloth.__version__}, PyTorch {torch.__version__}, {gpu}")
print(f"HF_TOKEN loaded: {'Yes' if os.environ.get('HF_TOKEN') else 'No'}")

# Benchmark Helper Functions
import subprocess

def measure_gpu_memory():
    """Measure current GPU memory usage in MB using nvidia-smi"""
    try:
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=memory.used', '--format=csv,noheader,nounits'],
            capture_output=True, text=True
        )
        return int(result.stdout.strip().split('\n')[0])
    except:
        return 0

def count_parameters(model):
    """Count trainable vs total parameters"""
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return {
        "trainable": trainable,
        "total": total,
        "pct": 100 * trainable / total
    }

def cleanup_memory():
    """Force garbage collection and clear CUDA cache"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

print("Benchmark functions defined.")
print(f"Initial GPU memory: {measure_gpu_memory()} MB")

# Create minimal synthetic instruction dataset with thinking content (5 samples)
from datasets import Dataset

synthetic_data = [
    {
        "instruction": "What is machine learning?",
        "thinking": "What is the user asking here? They want to understand machine learning. What are the key concepts I should cover? It's a subset of AI... and it involves learning from data. How should I keep this accessible? Short and clear definition.",
        "response": "Machine learning is a subset of artificial intelligence where computers learn patterns from data."
    },
    {
        "instruction": "Explain Python in one sentence.",
        "thinking": "One sentence only - what's most important about Python? Its readability and versatility are the defining features. How do I capture both in one sentence?",
        "response": "Python is a high-level programming language known for its readability and versatility."
    },
    {
        "instruction": "What is a neural network?",
        "thinking": "How do I explain neural networks simply? What's the core concept? They're inspired by biological neurons... they process information in layers. Should I mention deep learning? Maybe keep it basic for now.",
        "response": "A neural network is a computational model inspired by biological neurons that processes information through connected layers."
    },
    {
        "instruction": "Define supervised learning.",
        "thinking": "What makes supervised learning 'supervised'? It's the labels! The data has known outputs. How do I explain this clearly? Focus on the training process with labeled data.",
        "response": "Supervised learning is training a model on labeled data to predict outcomes for new inputs."
    },
    {
        "instruction": "What is gradient descent?",
        "thinking": "This is a bit technical. What's the intuition behind gradient descent? It's like walking downhill to find the lowest point. The gradient tells us which direction is 'down'. Keep it conceptual.",
        "response": "Gradient descent is an optimization algorithm that minimizes loss by iteratively adjusting parameters in the direction of steepest descent."
    },
]

print(f"Synthetic dataset prepared: {len(synthetic_data)} samples")

# Target Module Configurations

TARGET_CONFIGS = {
    "qv_only": {
        "modules": ["q_proj", "v_proj"],
        "description": "Query + Value only (minimal)",
    },
    "attention_only": {
        "modules": ["q_proj", "k_proj", "v_proj", "o_proj"],
        "description": "All attention layers",
    },
    "mlp_only": {
        "modules": ["gate_proj", "up_proj", "down_proj"],
        "description": "MLP/FFN layers only",
    },
    "all_linear": {
        "modules": ["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
        "description": "All linear layers (maximum)",
    },
}

print("Target module configurations:")
for name, config in TARGET_CONFIGS.items():
    print(f"  - {name}: {config['description']}")
    print(f"    Modules: {config['modules']}")

# Results Visualization
import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame(results)

print("\n" + "="*60)
print("Target Modules Comparison Results")
print("="*60)
print(df[["config", "num_modules", "trainable_params", "trainable_pct", "final_loss"]].to_string(index=False))

# Create visualization
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
colors = ['#3498db', '#2ecc71', '#e74c3c', '#9b59b6']

# Plot 1: Trainable Parameters
bars1 = axes[0].bar(df['config'], df['trainable_params'] / 1e6, color=colors)
axes[0].set_title('Trainable Parameters by Target Config', fontsize=12)
axes[0].set_xlabel('Target Module Configuration')
axes[0].set_ylabel('Parameters (Millions)')
axes[0].tick_params(axis='x', rotation=45)
for bar, val in zip(bars1, df['trainable_params']):
    axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,
                 f'{val/1e6:.1f}M', ha='center', va='bottom')

# Plot 2: Final Loss
bars2 = axes[1].bar(df['config'], df['final_loss'], color=colors)
axes[1].set_title('Final Training Loss by Target Config', fontsize=12)
axes[1].set_xlabel('Target Module Configuration')
axes[1].set_ylabel('Loss')
axes[1].tick_params(axis='x', rotation=45)
for bar, val in zip(bars2, df['final_loss']):
    axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.05,
                 f'{val:.4f}', ha='center', va='bottom')

plt.tight_layout()
plt.savefig('outputs_qlora_target_ministral/target_modules_comparison.png', dpi=150)
plt.show()

print("\nVisualization saved to outputs_qlora_target_ministral/target_modules_comparison.png")

# Show sample responses
print("\n" + "="*60)
print("Sample Responses by Configuration")
print("="*60)
for _, row in df.iterrows():
    print(f"\n[{row['config']}]")
    print(f"{row['sample_response'][:150]}...")

# Shutdown kernel to release all GPU memory
import IPython
print("Shutting down kernel to release GPU memory...")
app = IPython.Application.instance()
app.kernel.do_shutdown(restart=False)