Create benchmark.py

benchmark.py

@@ -0,0 +1,446 @@
+"""
+Comprehensive benchmarking script for Helion-V2.0-Thinking
+Measures performance, throughput, latency, and memory usage
+"""
+
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+from typing import List, Dict, Any
+import time
+import psutil
+import numpy as np
+from PIL import Image
+import json
+from tqdm import tqdm
+import gc
+
+
+class HelionBenchmark:
+    """Performance benchmarking for Helion-V2.0-Thinking"""
+    
+    def __init__(self, model_name: str = "DeepXR/Helion-V2.0-Thinking"):
+        """Initialize benchmark suite"""
+        print(f"Loading model for benchmarking: {model_name}")
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.model.eval()
+        
+        print("Model loaded successfully")
+        
+        # Get device info
+        if torch.cuda.is_available():
+            self.device = "cuda"
+            self.device_name = torch.cuda.get_device_name(0)
+            self.total_vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        else:
+            self.device = "cpu"
+            self.device_name = "CPU"
+            self.total_vram = 0
+        
+        print(f"Device: {self.device_name}")
+        if self.device == "cuda":
+            print(f"Total VRAM: {self.total_vram:.2f} GB")
+    
+    def measure_memory_usage(self) -> Dict[str, float]:
+        """Measure current memory usage"""
+        memory_stats = {}
+        
+        if self.device == "cuda":
+            memory_stats['vram_allocated_gb'] = torch.cuda.memory_allocated() / (1024**3)
+            memory_stats['vram_reserved_gb'] = torch.cuda.memory_reserved() / (1024**3)
+            memory_stats['vram_peak_gb'] = torch.cuda.max_memory_allocated() / (1024**3)
+        
+        # System RAM
+        memory_stats['ram_used_gb'] = psutil.Process().memory_info().rss / (1024**3)
+        
+        return memory_stats
+    
+    def benchmark_text_generation(
+        self,
+        prompts: List[str],
+        max_new_tokens: int = 256
+    ) -> Dict[str, Any]:
+        """
+        Benchmark text generation performance
+        
+        Returns:
+            Dict with latency, throughput, and token metrics
+        """
+        print("\n=== Benchmarking Text Generation ===")
+        
+        latencies = []
+        tokens_per_second = []
+        
+        # Warmup
+        warmup_prompt = "Test prompt for warmup."
+        inputs = self.processor(text=warmup_prompt, return_tensors="pt").to(self.model.device)
+        with torch.no_grad():
+            _ = self.model.generate(**inputs, max_new_tokens=50)
+        
+        if self.device == "cuda":
+            torch.cuda.synchronize()
+            torch.cuda.reset_peak_memory_stats()
+        
+        # Benchmark
+        for prompt in tqdm(prompts, desc="Text Generation"):
+            inputs = self.processor(text=prompt, return_tensors="pt").to(self.model.device)
+            
+            start_time = time.time()
+            
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=max_new_tokens,
+                    do_sample=False
+                )
+            
+            if self.device == "cuda":
+                torch.cuda.synchronize()
+            
+            end_time = time.time()
+            latency = end_time - start_time
+            
+            # Calculate tokens generated
+            generated_tokens = outputs.shape[1] - inputs['input_ids'].shape[1]
+            tps = generated_tokens / latency
+            
+            latencies.append(latency)
+            tokens_per_second.append(tps)
+        
+        # Memory stats
+        memory_stats = self.measure_memory_usage()
+        
+        return {
+            "text_generation": {
+                "avg_latency_ms": np.mean(latencies) * 1000,
+                "p50_latency_ms": np.percentile(latencies, 50) * 1000,
+                "p95_latency_ms": np.percentile(latencies, 95) * 1000,
+                "p99_latency_ms": np.percentile(latencies, 99) * 1000,
+                "avg_tokens_per_second": np.mean(tokens_per_second),
+                "total_prompts": len(prompts),
+                **memory_stats
+            }
+        }
+    
+    def benchmark_vision(
+        self,
+        image_prompts: List[tuple[Image.Image, str]],
+        max_new_tokens: int = 256
+    ) -> Dict[str, Any]:
+        """
+        Benchmark vision + text generation
+        
+        Args:
+            image_prompts: List of (image, prompt) tuples
+        
+        Returns:
+            Performance metrics
+        """
+        print("\n=== Benchmarking Vision Tasks ===")
+        
+        latencies = []
+        tokens_per_second = []
+        
+        # Warmup
+        if image_prompts:
+            warmup_image, warmup_prompt = image_prompts[0]
+            inputs = self.processor(
+                text=warmup_prompt,
+                images=warmup_image,
+                return_tensors="pt"
+            ).to(self.model.device)
+            with torch.no_grad():
+                _ = self.model.generate(**inputs, max_new_tokens=50)
+        
+        if self.device == "cuda":
+            torch.cuda.synchronize()
+            torch.cuda.reset_peak_memory_stats()
+        
+        # Benchmark
+        for image, prompt in tqdm(image_prompts, desc="Vision Tasks"):
+            inputs = self.processor(
+                text=prompt,
+                images=image,
+                return_tensors="pt"
+            ).to(self.model.device)
+            
+            start_time = time.time()
+            
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=max_new_tokens,
+                    do_sample=False
+                )
+            
+            if self.device == "cuda":
+                torch.cuda.synchronize()
+            
+            end_time = time.time()
+            latency = end_time - start_time
+            
+            generated_tokens = outputs.shape[1] - inputs['input_ids'].shape[1]
+            tps = generated_tokens / latency
+            
+            latencies.append(latency)
+            tokens_per_second.append(tps)
+        
+        memory_stats = self.measure_memory_usage()
+        
+        return {
+            "vision_tasks": {
+                "avg_latency_ms": np.mean(latencies) * 1000,
+                "p50_latency_ms": np.percentile(latencies, 50) * 1000,
+                "p95_latency_ms": np.percentile(latencies, 95) * 1000,
+                "p99_latency_ms": np.percentile(latencies, 99) * 1000,
+                "avg_tokens_per_second": np.mean(tokens_per_second),
+                "total_image_prompts": len(image_prompts),
+                **memory_stats
+            }
+        }
+    
+    def benchmark_long_context(
+        self,
+        context_lengths: List[int] = [1000, 5000, 10000, 50000, 100000]
+    ) -> Dict[str, Any]:
+        """
+        Benchmark performance with varying context lengths
+        
+        Args:
+            context_lengths: List of context lengths to test
+        
+        Returns:
+            Performance metrics by context length
+        """
+        print("\n=== Benchmarking Long Context Performance ===")
+        
+        results = {}
+        
+        for length in tqdm(context_lengths, desc="Context Lengths"):
+            # Generate synthetic context
+            context = "This is a test sentence. " * (length // 6)  # Approx tokens
+            prompt = f"{context}\n\nQuestion: What is the main topic of this text?\nAnswer:"
+            
+            inputs = self.processor(text=prompt, return_tensors="pt").to(self.model.device)
+            
+            # Check if context fits
+            if inputs['input_ids'].shape[1] > length:
+                print(f"Skipping {length} - generated context too long")
+                continue
+            
+            if self.device == "cuda":
+                torch.cuda.synchronize()
+                torch.cuda.reset_peak_memory_stats()
+            
+            start_time = time.time()
+            
+            try:
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=128,
+                        do_sample=False
+                    )
+                
+                if self.device == "cuda":
+                    torch.cuda.synchronize()
+                
+                end_time = time.time()
+                latency = end_time - start_time
+                
+                memory_stats = self.measure_memory_usage()
+                
+                results[f"context_{length}"] = {
+                    "latency_ms": latency * 1000,
+                    "input_tokens": inputs['input_ids'].shape[1],
+                    **memory_stats
+                }
+            
+            except Exception as e:
+                results[f"context_{length}"] = {
+                    "error": str(e)
+                }
+            
+            # Cleanup
+            del inputs, outputs
+            gc.collect()
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+        
+        return {"long_context": results}
+    
+    def benchmark_throughput(
+        self,
+        batch_sizes: List[int] = [1, 2, 4, 8],
+        sequence_length: int = 512
+    ) -> Dict[str, Any]:
+        """
+        Benchmark throughput with different batch sizes
+        
+        Args:
+            batch_sizes: List of batch sizes to test
+            sequence_length: Target sequence length
+        
+        Returns:
+            Throughput metrics
+        """
+        print("\n=== Benchmarking Throughput ===")
+        
+        results = {}
+        prompt = "Explain the concept of artificial intelligence. " * 10
+        
+        for batch_size in tqdm(batch_sizes, desc="Batch Sizes"):
+            try:
+                # Create batch
+                prompts = [prompt] * batch_size
+                inputs = self.processor(
+                    text=prompts,
+                    return_tensors="pt",
+                    padding=True
+                ).to(self.model.device)
+                
+                if self.device == "cuda":
+                    torch.cuda.synchronize()
+                    torch.cuda.reset_peak_memory_stats()
+                
+                start_time = time.time()
+                
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=256,
+                        do_sample=False
+                    )
+                
+                if self.device == "cuda":
+                    torch.cuda.synchronize()
+                
+                end_time = time.time()
+                latency = end_time - start_time
+                
+                # Calculate throughput
+                total_tokens = outputs.shape[0] * outputs.shape[1]
+                throughput = total_tokens / latency
+                
+                memory_stats = self.measure_memory_usage()
+                
+                results[f"batch_{batch_size}"] = {
+                    "latency_ms": latency * 1000,
+                    "throughput_tokens_per_sec": throughput,
+                    "tokens_per_sample": outputs.shape[1],
+                    **memory_stats
+                }
+            
+            except Exception as e:
+                results[f"batch_{batch_size}"] = {
+                    "error": str(e)
+                }
+            
+            # Cleanup
+            del inputs, outputs
+            gc.collect()
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+        
+        return {"throughput": results}
+    
+    def run_full_benchmark(self) -> Dict[str, Any]:
+        """Run complete benchmark suite"""
+        print("\n" + "="*60)
+        print("Starting Full Benchmark Suite")
+        print(f"Device: {self.device_name}")
+        print("="*60)
+        
+        results = {
+            "system_info": {
+                "device": self.device,
+                "device_name": self.device_name,
+                "total_vram_gb": self.total_vram if self.device == "cuda" else None,
+                "pytorch_version": torch.__version__,
+                "cuda_available": torch.cuda.is_available()
+            }
+        }
+        
+        # Text generation benchmark
+        text_prompts = [
+            "Explain quantum mechanics in simple terms.",
+            "Write a short story about space exploration.",
+            "What are the benefits of machine learning?",
+            "Describe the process of photosynthesis.",
+            "How does blockchain technology work?"
+        ]
+        results.update(self.benchmark_text_generation(text_prompts, max_new_tokens=256))
+        
+        # Long context benchmark
+        results.update(self.benchmark_long_context([1000, 5000, 10000, 50000]))
+        
+        # Throughput benchmark
+        results.update(self.benchmark_throughput([1, 2, 4]))
+        
+        print("\n" + "="*60)
+        print("Benchmark Complete")
+        print("="*60)
+        
+        return results
+    
+    def print_results(self, results: Dict[str, Any]):
+        """Print benchmark results in a readable format"""
+        print("\n" + "="*60)
+        print("BENCHMARK RESULTS")
+        print("="*60)
+        
+        def print_dict(d, indent=0):
+            for key, value in d.items():
+                if isinstance(value, dict):
+                    print("  " * indent + f"{key}:")
+                    print_dict(value, indent + 1)
+                elif isinstance(value, float):
+                    print("  " * indent + f"{key}: {value:.4f}")
+                else:
+                    print("  " * indent + f"{key}: {value}")
+        
+        print_dict(results)
+        print("="*60 + "\n")
+    
+    def save_results(self, results: Dict[str, Any], filename: str = "benchmark_results.json"):
+        """Save results to JSON file"""
+        with open(filename, 'w') as f:
+            json.dump(results, f, indent=2)
+        print(f"Results saved to {filename}")
+
+
+def main():
+    """Main benchmark function"""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Benchmark Helion-V2.0-Thinking")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="DeepXR/Helion-V2.0-Thinking",
+        help="Model name or path"
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="benchmark_results.json",
+        help="Output file for results"
+    )
+    
+    args = parser.parse_args()
+    
+    # Run benchmark
+    benchmark = HelionBenchmark(args.model)
+    results = benchmark.run_full_benchmark()
+    benchmark.print_results(results)
+    benchmark.save_results(results, args.output)
+
+
+if __name__ == "__main__":
+    main()