RAG時代終了か？高精度・高速・軽量のローカルLLMを【unsloth + Gemma3】で高速FTする！ #Python – Qiita

import torch
import time
import psutil
import GPUtil
from typing import Dict, List, Any
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from dataclasses import dataclass
import json

@dataclass
class BenchmarkResult:
    """ベンチマーク結果を格納するデータクラス"""
    model_name: str
    optimization: str
    batch_size: int
    sequence_length: int
    throughput: float  # tokens/second
    latency: float  # seconds
    memory_usage: float  # GB
    accuracy_score: float

class UnslothBenchmark:
    """UnslothAIベンチマーククラス"""
    
    def __init__(self):
        self.results = []
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    def measure_gpu_memory(self) -> float:
        """GPU メモリ使用量を測定"""
        if torch.cuda.is_available():
            return torch.cuda.memory_allocated() / 1024**3  # GB
        return 0.0
    
    def measure_system_resources(self) -> Dict[str, float]:
        """システムリソースの測定"""
        resources = {
            "cpu_percent": psutil.cpu_percent(interval=1),
            "ram_usage_gb": psutil.virtual_memory().used / 1024**3,
            "ram_percent": psutil.virtual_memory().percent
        }
        
        if torch.cuda.is_available():
            gpus = GPUtil.getGPUs()
            if gpus:
                gpu = gpus[0]
                resources.update({
                    "gpu_memory_used_gb": gpu.memoryUsed / 1024,
                    "gpu_memory_percent": gpu.memoryUtil * 100,
                    "gpu_utilization": gpu.load * 100,
                    "gpu_temperature": gpu.temperature
                })
        
        return resources
    
    def benchmark_training(
        self,
        model,
        train_dataloader,
        num_steps: int = 100,
        optimization_type: str = "unsloth"
    ) -> BenchmarkResult:
        """訓練性能のベンチマーク"""
        
        print(f"\nBenchmarking {optimization_type} training...")
        
        # 初期リソース測定
        torch.cuda.empty_cache()
        initial_memory = self.measure_gpu_memory()
        
        # ウォームアップ
        for i, batch in enumerate(train_dataloader):
            if i >= 5:
                break
            outputs = model(**batch)
            loss = outputs.loss
            loss.backward()
            
        torch.cuda.synchronize()
        
        # 実際のベンチマーク
        start_time = time.time()
        total_tokens = 0
        
        for step, batch in enumerate(train_dataloader):
            if step >= num_steps:
                break
                
            step_start = time.time()
            
            # Forward pass
            outputs = model(**batch)
            loss = outputs.loss
            
            # Backward pass
            loss.backward()
            
            # トークン数のカウント
            total_tokens += batch["input_ids"].numel()
            
            if step % 10 == 0:
                step_time = time.time() - step_start
                current_memory = self.measure_gpu_memory()
                print(f"Step {step}: {step_time:.3f}s, Memory: {current_memory:.2f}GB")
        
        torch.cuda.synchronize()
        total_time = time.time() - start_time
        
        # メトリクスの計算
        throughput = total_tokens / total_time
        avg_latency = total_time / num_steps
        peak_memory = self.measure_gpu_memory()
        memory_usage = peak_memory - initial_memory
        
        result = BenchmarkResult(
            model_name=model.config.model_type,
            optimization=optimization_type,
            batch_size=train_dataloader.batch_size,
            sequence_length=batch["input_ids"].shape[1],
            throughput=throughput,
            latency=avg_latency,
            memory_usage=memory_usage,
            accuracy_score=0.0  # 別途評価
        )
        
        self.results.append(result)
        return result
    
    def benchmark_inference(
        self,
        model,
        tokenizer,
        test_prompts: List[str],
        optimization_type: str = "unsloth"
    ) -> BenchmarkResult:
        """推論性能のベンチマーク"""
        
        print(f"\nBenchmarking {optimization_type} inference...")
        
        # ウォームアップ
        for _ in range(3):
            inputs = tokenizer(test_prompts[0], return_tensors="pt").to(self.device)
            with torch.no_grad():
                _ = model.generate(**inputs, max_new_tokens=50)
        
        torch.cuda.synchronize()
        
        # 実際のベンチマーク
        total_time = 0
        total_tokens_generated = 0
        initial_memory = self.measure_gpu_memory()
        
        for prompt in test_prompts:
            inputs = tokenizer(prompt, return_tensors="pt").to(self.device)
            
            start_time = time.time()
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=128,
                    do_sample=True,
                    temperature=0.7
                )
            torch.cuda.synchronize()
            
            generation_time = time.time() - start_time
            total_time += generation_time
            
            # 生成されたトークン数
            num_generated = outputs.shape[1] - inputs["input_ids"].shape[1]
            total_tokens_generated += num_generated
        
        # メトリクスの計算
        throughput = total_tokens_generated / total_time
        avg_latency = total_time / len(test_prompts)
        peak_memory = self.measure_gpu_memory()
        memory_usage = peak_memory - initial_memory
        
        result = BenchmarkResult(
            model_name=model.config.model_type,
            optimization=optimization_type,
            batch_size=1,
            sequence_length=128,
            throughput=throughput,
            latency=avg_latency,
            memory_usage=memory_usage,
            accuracy_score=0.0
        )
        
        self.results.append(result)
        return result
    
    def compare_optimizations(self):
        """異なる最適化手法の比較"""
        
        # 結果をDataFrameに変換
        df = pd.DataFrame([vars(r) for r in self.results])
        
        # 可視化
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # スループット比較
        sns.barplot(data=df, x="optimization", y="throughput", ax=axes[0, 0])
        axes[0, 0].set_title("Throughput Comparison (tokens/second)")
        axes[0, 0].set_ylabel("Tokens per Second")
        
        # レイテンシ比較
        sns.barplot(data=df, x="optimization", y="latency", ax=axes[0, 1])
        axes[0, 1].set_title("Latency Comparison")
        axes[0, 1].set_ylabel("Latency (seconds)")
        
        # メモリ使用量比較
        sns.barplot(data=df, x="optimization", y="memory_usage", ax=axes[1, 0])
        axes[1, 0].set_title("Memory Usage Comparison")
        axes[1, 0].set_ylabel("Memory Usage (GB)")
        
        # 総合スコア（正規化して計算）
        df['normalized_throughput'] = df['throughput'] / df['throughput'].max()
        df['normalized_latency'] = 1 - (df['latency'] / df['latency'].max())
        df['normalized_memory'] = 1 - (df['memory_usage'] / df['memory_usage'].max())
        df['overall_score'] = (df['normalized_throughput'] + 
                              df['normalized_latency'] + 
                              df['normalized_memory']) / 3
        
        sns.barplot(data=df, x="optimization", y="overall_score", ax=axes[1, 1])
        axes[1, 1].set_title("Overall Performance Score")
        axes[1, 1].set_ylabel("Score (0-1)")
        
        plt.tight_layout()
        plt.savefig("benchmark_results.png", dpi=300)
        plt.show()
        
        # 詳細な統計情報を出力
        print("\n=== Benchmark Summary ===")
        for opt in df['optimization'].unique():
            opt_data = df[df['optimization'] == opt]
            print(f"\n{opt}:")
            print(f"  Average Throughput: {opt_data['throughput'].mean():.2f} tokens/s")
            print(f"  Average Latency: {opt_data['latency'].mean():.3f}s")
            print(f"  Average Memory: {opt_data['memory_usage'].mean():.2f}GB")
            print(f"  Overall Score: {opt_data['overall_score'].mean():.3f}")
    
    def save_results(self, filename: str = "benchmark_results.json"):
        """結果の保存"""
        
        results_dict = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "system_info": {
                "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU",
                "cuda_version": torch.version.cuda,
                "pytorch_version": torch.__version__,
            },
            "results": [vars(r) for r in self.results]
        }
        
        with open(filename, "w") as f:
            json.dump(results_dict, f, indent=2)
        
        print(f"\nResults saved to {filename}")

# 実際のベンチマーク実行例
def run_comprehensive_benchmark():
    """包括的なベンチマークの実行"""
    
    benchmark = UnslothBenchmark()
    
    # テストプロンプト
    test_prompts = [
        "機械学習の基本的な概念を説明してください。",
        "Pythonでクイックソートを実装する方法を教えてください。",
        "地球温暖化の原因と対策について述べてください。",
        "健康的な生活習慣について助言をください。",
        "最新のAI技術のトレンドを教えてください。"
    ]
    
    # 異なる最適化設定でのベンチマーク
    optimizations = [
        ("standard", False, False),  # 標準実装
        ("flash_attention", True, False),  # Flash Attention のみ
        ("unsloth", True, True),  # UnslothAI フル最適化
    ]
    
    for opt_name, use_flash, use_unsloth in optimizations:
        print(f"\n{'='*50}")
        print(f"Testing {opt_name} optimization")
        print(f"{'='*50}")
        
        # モデルの読み込み（各最適化設定で）
        if use_unsloth:
            from unsloth import FastModel
            model, tokenizer = FastModel.from_pretrained(
                "unsloth/gemma-3-4b-it",
                max_seq_length=2048,
                load_in_4bit=True,
            )
        else:
            # 標準的な読み込み
            from transformers import AutoModelForCausalLM, AutoTokenizer
            model = AutoModelForCausalLM.from_pretrained(
                "google/gemma-3-4b-it",
                torch_dtype=torch.float16,
                device_map="auto",
                use_flash_attention_2=use_flash
            )
            tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-4b-it")
        
        # 推論ベンチマーク
        inference_result = benchmark.benchmark_inference(
            model, 
            tokenizer, 
            test_prompts,
            optimization_type=opt_name
        )
        
        print(f"\n{opt_name} Results:")
        print(f"  Throughput: {inference_result.throughput:.2f} tokens/s")
        print(f"  Latency: {inference_result.latency:.3f}s")
        print(f"  Memory: {inference_result.memory_usage:.2f}GB")
        
        # メモリクリーンアップ
        del model
        if 'tokenizer' in locals():
            del tokenizer
        torch.cuda.empty_cache()
    
    # 結果の比較と保存
    benchmark.compare_optimizations()
    benchmark.save_results()

if __name__ == "__main__":
    run_comprehensive_benchmark()