QLoRA Test: Quantization Method Comparison - Ministral-3B-Reasoning¶

# Environment Setup
import os
import time
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 quantization methods 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")

# Quantization Configuration

QUANT_CONFIGS = {
    "4bit_nf4": {
        "model_name": "unsloth/Ministral-3-3B-Reasoning-2512",
        "load_in_4bit": True,
        "description": "4-bit NF4 (QLoRA standard)",
    },
    "bf16": {
        "model_name": "unsloth/Ministral-3-3B-Reasoning-2512",
        "load_in_4bit": False,
        "description": "BF16 (full precision baseline)",
    },
}

print("Quantization configurations:")
for name, config in QUANT_CONFIGS.items():
    print(f"  - {name}: {config['description']}")
print("\nNote: 8-bit INT8 excluded due to Unsloth compatibility considerations")

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

# Filter successful results
successful_results = [r for r in results if r.get('success', False)]

if successful_results:
    df = pd.DataFrame(successful_results)
    
    print("\n" + "="*60)
    print("Quantization Comparison Results")
    print("="*60)
    print(df[["quantization", "peak_memory_mb", "train_time_sec", "final_loss"]].to_string(index=False))
    
    # Calculate memory savings vs BF16
    if "bf16" in df["quantization"].values:
        bf16_mem = df[df["quantization"] == "bf16"]["peak_memory_mb"].values[0]
        print(f"\nMemory savings vs BF16 ({bf16_mem} MB):")
        for _, row in df.iterrows():
            if row["quantization"] != "bf16":
                savings = 100 * (1 - row["peak_memory_mb"] / bf16_mem)
                print(f"  {row['quantization']}: {savings:.1f}% reduction")
    
    # Create visualization
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    colors = ['#2ecc71', '#3498db']  # green, blue
    
    # Plot 1: Peak GPU Memory
    bars1 = axes[0].bar(df['quantization'], df['peak_memory_mb'], color=colors[:len(df)])
    axes[0].set_title('Peak GPU Memory', fontsize=12)
    axes[0].set_xlabel('Quantization Method')
    axes[0].set_ylabel('Memory (MB)')
    axes[0].tick_params(axis='x', rotation=45)
    for bar, val in zip(bars1, df['peak_memory_mb']):
        axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 50,
                     f'{val:.0f}', ha='center', va='bottom')
    
    # Plot 2: Training Time
    bars2 = axes[1].bar(df['quantization'], df['train_time_sec'], color=colors[:len(df)])
    axes[1].set_title('Training Time (3 steps)', fontsize=12)
    axes[1].set_xlabel('Quantization Method')
    axes[1].set_ylabel('Seconds')
    axes[1].tick_params(axis='x', rotation=45)
    for bar, val in zip(bars2, df['train_time_sec']):
        axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,
                     f'{val:.1f}s', ha='center', va='bottom')
    
    # Plot 3: Final Loss
    bars3 = axes[2].bar(df['quantization'], df['final_loss'], color=colors[:len(df)])
    axes[2].set_title('Final Training Loss', fontsize=12)
    axes[2].set_xlabel('Quantization Method')
    axes[2].set_ylabel('Loss')
    axes[2].tick_params(axis='x', rotation=45)
    for bar, val in zip(bars3, df['final_loss']):
        axes[2].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_quant_ministral/quantization_comparison.png', dpi=150)
    plt.show()
    
    print("\nVisualization saved to outputs_qlora_quant_ministral/quantization_comparison.png")
else:
    print("No successful results to visualize.")

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