diff --git a/soup_cli/commands/bench.py b/soup_cli/commands/bench.py index 6da0d9a..57c83bd 100644 --- a/soup_cli/commands/bench.py +++ b/soup_cli/commands/bench.py @@ -1,4 +1,4 @@ -"""soup bench — simple measuring tool for model speed and memory.""" +"""soup bench -- simple measuring tool for model speed and memory.""" import time from pathlib import Path @@ -34,7 +34,7 @@ def bench( "-n", help="Number of prompts to run for averaging", ), -): +) -> None: """Run an inference benchmark (speed and memory) on a loaded model.""" import torch @@ -48,6 +48,12 @@ def bench( device, _ = detect_device() + if device == "cpu": + console.print( + "[yellow]Warning:[/] Running on CPU. Inference speed is typically " + "10-100x slower than GPU -- results will not reflect production TPS." + ) + console.print( Panel( f"Model: [bold]{model_path}[/]\n" @@ -60,15 +66,17 @@ def bench( console.print("[dim]Loading model to measure resource usage...[/]") + # Reset peak stats before load -- "Max VRAM" reflects total footprint + # (model load + inference), i.e. what users need for deployment planning. if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() start_load = time.time() try: model_obj, tokenizer = _load_model(str(model_path), base, device) - except Exception as exc: + except (OSError, ImportError, RuntimeError, ValueError) as exc: console.print(f"[red]Failed to load model:[/] {exc}") - raise typer.Exit(1) + raise typer.Exit(1) from exc load_time = time.time() - start_load console.print(f"[green]Model loaded in {load_time:.2f}s.[/]\n") @@ -78,11 +86,19 @@ def bench( "Write a short Python function to calculate fibonacci numbers.", "What are the main consequences of the Industrial Revolution?", "Compose a poem about a wandering space traveler.", - "Describe how a database index works under the hood." + "Describe how a database index works under the hood.", ] - # Loop over prompts test_prompts = (prompts * (num_prompts // len(prompts) + 1))[:num_prompts] + # Warmup run: first inference includes CUDA kernel JIT compilation, + # which would skew the average. Discarded from timing. + console.print("[dim]Warmup run (discarded from timing)...[/]") + warmup_messages = [{"role": "user", "content": test_prompts[0]}] + _generate( + model_obj, tokenizer, warmup_messages, + max_tokens=min(max_tokens, 32), temperature=0.0, + ) + total_tokens = 0 total_latency = 0.0 @@ -92,7 +108,7 @@ def bench( messages = [{"role": "user", "content": prompt_text}] start_time = time.time() - response, token_count = _generate( + _, token_count = _generate( model_obj, tokenizer, messages, max_tokens=max_tokens, temperature=0.0, ) @@ -100,7 +116,7 @@ def bench( latency = time.time() - start_time total_tokens += token_count total_latency += latency - console.print(f" [dim]Prompt {i+1}: {token_count} tokens in {latency:.2f}s[/]") + console.print(f" [dim]Prompt {i + 1}: {token_count} tokens in {latency:.2f}s[/]") avg_tps = total_tokens / total_latency if total_latency > 0 else 0 @@ -119,7 +135,7 @@ def bench( "Transformers", f"{avg_tps:.2f}", f"{total_latency:.2f}s", - vram_str + vram_str, ) console.print()