QLoRA Parameter Test: LoRA Rank 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

# Same dataset across all ranks for fair comparison
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")

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

df = pd.DataFrame(results)

print("\n" + "="*60)
print("LoRA Rank Comparison Results")
print("="*60)
print(df[["rank", "trainable_params", "trainable_pct", "gpu_peak_mb", "final_loss"]].to_string(index=False))

# Create visualization
fig, axes = plt.subplots(2, 2, figsize=(12, 10))

# Plot 1: Trainable Parameters vs Rank
axes[0, 0].plot(df['rank'], df['trainable_params'] / 1e6, 'b-o', linewidth=2, markersize=8)
axes[0, 0].set_title('Trainable Parameters vs LoRA Rank', fontsize=12)
axes[0, 0].set_xlabel('LoRA Rank (r)')
axes[0, 0].set_ylabel('Parameters (Millions)')
axes[0, 0].grid(True, alpha=0.3)
axes[0, 0].set_xticks(df['rank'])

# Plot 2: Trainable Percentage vs Rank
axes[0, 1].plot(df['rank'], df['trainable_pct'], 'g-o', linewidth=2, markersize=8)
axes[0, 1].set_title('Trainable % vs LoRA Rank', fontsize=12)
axes[0, 1].set_xlabel('LoRA Rank (r)')
axes[0, 1].set_ylabel('% of Total Parameters')
axes[0, 1].grid(True, alpha=0.3)
axes[0, 1].set_xticks(df['rank'])

# Plot 3: GPU Memory vs Rank
axes[1, 0].plot(df['rank'], df['gpu_peak_mb'], 'r-o', linewidth=2, markersize=8)
axes[1, 0].set_title('Peak GPU Memory vs LoRA Rank', fontsize=12)
axes[1, 0].set_xlabel('LoRA Rank (r)')
axes[1, 0].set_ylabel('Memory (MB)')
axes[1, 0].grid(True, alpha=0.3)
axes[1, 0].set_xticks(df['rank'])

# Plot 4: Final Loss vs Rank
axes[1, 1].plot(df['rank'], df['final_loss'], 'm-o', linewidth=2, markersize=8)
axes[1, 1].set_title('Final Training Loss vs LoRA Rank', fontsize=12)
axes[1, 1].set_xlabel('LoRA Rank (r)')
axes[1, 1].set_ylabel('Loss')
axes[1, 1].grid(True, alpha=0.3)
axes[1, 1].set_xticks(df['rank'])

plt.tight_layout()
plt.savefig('outputs_qlora_rank_ministral/rank_comparison.png', dpi=150)
plt.show()

print("\nVisualization saved to outputs_qlora_rank_ministral/rank_comparison.png")

# 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)