""" 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()