import gc import json import logging import os import pathlib import re import tempfile import time from datetime import datetime from queue import Queue from typing import Any import torch from datasets import Dataset from huggingface_hub import HfApi from tqdm import trange from transformers import ( AutoModelForCausalLM, AutoTokenizer, GenerationConfig, GenerationMixin, ) from transformers.generation.streamers import BaseStreamer from .benchmark_config import BenchmarkConfig from .data_classes import BenchmarkMetadata, BenchmarkResult, GPURawMetrics, pretty_print_dict from .hardware_metrics import GPUMonitor try: from kernels import Mode, kernelize # noqa: F401 except ImportError: kernelize = None Mode = None DEFAULT_PROMPT = "\n".join([ "The French Revolution was a period of political and societal change in France that began with the Estates General of 1789 and ended with the Coup of 18 Brumaire on 9 November 1799.", "Many of the revolution's ideas are considered fundamental principles of liberal democracy, and its values remain central to modern French political discourse.", "It was caused by a combination of social, political, and economic factors which the existing regime proved unable to manage.", "Financial crisis and widespread social distress led to the convocation of the Estates General in May 1789, its first meeting since 1614.", "The representatives of the Third Estate broke away and re-constituted themselves as a National Assembly in June.", "The Storming of the Bastille in Paris on 14 July led to a series of radical measures by the Assembly, including the abolition of feudalism, state control over the Catholic Church in France, and issuing the Declaration of the Rights of Man and of the Citizen.", "The next three years were dominated by a struggle for political control.", "King Louis XVI's attempted flight to Varennes in June 1791 further discredited the monarchy, and military defeats after the outbreak of the French Revolutionary Wars in April 1792 led to the insurrection of 10 August 1792.", "As a result, the monarchy was replaced by the French First Republic in September, followed by the execution of Louis XVI himself in January 1793.", "After another revolt in June 1793, the constitution was suspended, and political power passed from the National Convention to the Committee of Public Safety, dominated by radical Jacobins led by Maximilien Robespierre.", "About 16,000 people were sentenced by the Revolutionary Tribunal and executed in the Reign of Terror, which ended in July 1794 with the Thermidorian Reaction.", "Weakened by external threats and internal opposition, the Committee of Public Safety was replaced in November 1795 by the Directory.", "Its instability ended in the coup of 18 Brumaire and the establishment of the Consulate, with Napoleon Bonaparte as First Consul.", ]) # fmt: skip PUSH_TO_HUB_TOKEN = os.getenv("PUSH_TO_HUB_TOKEN", None) def compact_json_numeric_arrays(data: dict): # Match arrays that contain only numbers (ints/floats), whitespace, commas, and newlines pattern = r"\[\s*\n\s*((?:\d+(?:\.\d+)?\s*,\s*)*\d+(?:\.\d+)?)\s*\n\s*\]" def replace_numeric_array(match): # Get the array content content = match.group(1) # Remove extra whitespace but keep commas compact_content = re.sub(r"\s+", " ", content).strip() return f"[{compact_content}]" return re.sub(pattern, replace_numeric_array, json.dumps(data, indent=4, default=str), flags=re.DOTALL) def get_git_revision() -> str: base_path = pathlib.Path(__file__).parent.parent.parent git_dir = base_path / ".git" with (git_dir / "HEAD").open("r") as head: ref = head.readline().split(" ")[-1].strip() with (git_dir / ref).open("r") as git_hash: return git_hash.readline().strip() def flush_memory(flush_compile: bool = True) -> None: """Flush GPU memory and run garbage collection. If the flush_compile flag is set, we also clear the everything related to compile cache.""" gc.collect() # If needed, flush everything related to torch.compile if flush_compile: # Dynamo resets torch._dynamo.reset() torch._dynamo.reset_code_caches() if hasattr(torch._inductor, "codecache"): # Clear FX graph cache if hasattr(torch._inductor.codecache, "FxGraphCache"): torch._inductor.codecache.FxGraphCache.clear() # Clear PyCodeCache if hasattr(torch._inductor.codecache, "PyCodeCache"): torch._inductor.codecache.PyCodeCache.cache_clear() # Clear TritonFuture cache (for async compilation) if hasattr(torch._inductor.codecache, "TritonFuture"): if hasattr(torch._inductor.codecache.TritonFuture, "_compile_cache"): torch._inductor.codecache.TritonFuture._compile_cache.clear() # Clear CUDA cache if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() gc.collect() class BenchmarkStreamer(BaseStreamer): def __init__(self, **kwargs) -> None: self.timeout = kwargs.pop("timeout", 10) self.timestamps = [] self.text_queue = Queue() self.stop_signal = None def put(self, value): """Receives tokens and logs the timestamp of the generation.""" self.timestamps.append(time.perf_counter()) self.text_queue.put(value) def end(self): self.timestamps.append(time.perf_counter()) self.text_queue.put(self.stop_signal) def __iter__(self): return self def __next__(self): value = self.text_queue.get(timeout=self.timeout) if value == self.stop_signal: raise StopIteration() else: return value class BenchmarkRunner: """Main benchmark runner that coordinates benchmark execution.""" def __init__( self, logger: logging.Logger, output_dir: str | None = None, branch_name: str | None = None, commit_id: str | None = None, commit_message: str | None = None, ) -> None: # Those stay constant for the whole run self.logger = logger if output_dir is None: output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "benchmark_results") self.output_dir = output_dir self.branch_name = branch_name self.commit_id = get_git_revision() if commit_id is None else commit_id self.commit_message = commit_message os.makedirs(self.output_dir, exist_ok=True) self.profile_dir = None # Attributes that are reset for each model self._setup_for = "" # Attributes that are reset for each run self.model: GenerationMixin | None = None def cleanup(self) -> None: del self.model self.model = None flush_memory() def setup_benchmark(self, model_id: str, config: BenchmarkConfig) -> None: # Some attributes only need to be set once per model if self._setup_for != model_id: self.tokenizer = AutoTokenizer.from_pretrained(model_id) # We set the EOS token to the padding token for open-ended generation self.tokenizer.eos_token = self.tokenizer.pad_token self._setup_for = model_id # Prepare inputs self.inputs = self.tokenizer( [DEFAULT_PROMPT for _ in range(config.batch_size)], return_tensors="pt", max_length=config.sequence_length, truncation=True, return_attention_mask=True, ).to(config.device) self.inputs["use_cache"] = True # Prepare generation config generation_config_kwargs = { "do_sample": False, "max_new_tokens": config.num_tokens_to_generate, } # Add compile config if found if config.compile_config is not None: generation_config_kwargs.update(compile_config=config.compile_config) # To trigger compile in generate, we need to set the cache to static if not config.continuous_batching: generation_config_kwargs.update(cache_implementation="static") generation_config = GenerationConfig(**generation_config_kwargs) # Load model self.logger.debug(f"Loading model {model_id} on device {config.device}...") dtype = getattr(torch, config.dtype.removeprefix("torch.")) self.model = AutoModelForCausalLM.from_pretrained( model_id, dtype=dtype, attn_implementation=config.attn_implementation, generation_config=generation_config ) self.model = self.model.eval().to(config.device) # Kernelize the model if needed if config.kernelize and kernelize is not None and Mode is not None: self.model = kernelize(self.model, mode=Mode.INFERENCE) def run_benchmark(self, config: BenchmarkConfig, num_tokens_to_profile: int = 0) -> BenchmarkResult | None: """Run a single benchmark with the given model ID and config.""" with torch.no_grad(): self.logger.info(f"Running benchmark scenario: {config.name}") self.logger.debug(f"Full config: {config.to_dict()}") # Quick validation: try one measurement first to see if this scenario works flush_memory() e2e_latency = self.time_generate(config, warmup=True)[0] if e2e_latency < 0: self.logger.warning(f"Skipping config {config.name}: {e2e_latency = }") return None # Warmup runs self.logger.info(f"Warming up with {config.warmup_iterations} iterations...") for _ in trange(config.warmup_iterations, desc="Warmup"): self.time_generate(config, warmup=True) self.logger.info("Warmup over.") # Measurement runs result = BenchmarkResult() self.logger.info(f"Benchmarking with {config.measurement_iterations} iterations.") for _ in trange(config.measurement_iterations, desc="Benchmarking"): e2e_latency, timestamps, shape_and_decoded_output, gpu_metrics = self.time_generate( config, warmup=False ) result.accumulate(e2e_latency, timestamps, shape_and_decoded_output, gpu_metrics) self.logger.info("Benchmarking done. Cleaning up.") # Profile if needed if num_tokens_to_profile > 0: self.profile_generate(num_tokens_to_profile, config.name) return result def time_generate( self, config: BenchmarkConfig, warmup: bool ) -> tuple[float, list[float], str, GPURawMetrics | None]: # Prepare gpu monitoring if needed if config.gpu_monitoring and not warmup: gpu_monitor = GPUMonitor(logger=self.logger) gpu_monitor.start() else: gpu_monitor = None # Generate and time if config.continuous_batching: inputs = self.inputs["input_ids"].tolist() wall_time_0 = time.perf_counter() outputs = self.model.generate_batch(inputs, allow_prefix_sharing=False, record_timestamps=True) else: streamer = BenchmarkStreamer() wall_time_0 = time.perf_counter() outputs = self.model.generate(**self.inputs, streamer=streamer) wall_time_1 = time.perf_counter() gpu_metrics = gpu_monitor.stop_and_collect() if gpu_monitor is not None else None # Retrieve timestamps and results in a way that allows similar post-processing input_tokens = self.inputs["input_ids"].size(-1) if config.continuous_batching: timestamps = [output.timestamps[:] for output in outputs.values()] results = torch.tensor([output.generated_tokens[:] for output in outputs.values()]) else: timestamps = [streamer.timestamps[1:]] # skip the first timestamp because it's the input tokens results = outputs[:, input_tokens:] outputs = None flush_memory(flush_compile=False) # Check if generation had the right number of tokens if results.size(-1) != config.num_tokens_to_generate: raise RuntimeError(f"Generated {results.size(-1)} tokens, expected {config.num_tokens_to_generate}") # Decode outputs decoded_output = self.tokenizer.decode(results[0], skip_special_tokens=True) shape_and_decoded_output = f"{tuple(results.shape)} | {decoded_output}" # Compute metrics e2e_latency = wall_time_1 - wall_time_0 timestamps = torch.tensor(timestamps).sub(wall_time_0).tolist() self.logger.info( f"Time generate done in {e2e_latency:.2f} seconds. Memory usage: {torch.cuda.memory_allocated() / 1024**2:.2f} MB" ) return e2e_latency, timestamps, shape_and_decoded_output, gpu_metrics def profile_generate(self, num_tokens_to_profile: int, config_name: str) -> None: """Profile the latency of a call to model.generate() with the given (inputs) and (max_new_tokens).""" profiler = torch.profiler.profile( activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], record_shapes=True, ) with profiler as prof: _ = self.model.generate( **self.inputs, max_new_tokens=num_tokens_to_profile, ) if self.profile_dir is None: self.profile_dir = self.output_dir + "_profiles" os.makedirs(self.profile_dir, exist_ok=True) prof.export_chrome_trace(f"{self.profile_dir}/{config_name}.json") @torch.inference_mode() def run_benchmarks( self, model_id: str, benchmark_configs: list[BenchmarkConfig], num_tokens_to_profile: int = 0, pretty_print_summary: bool = True, summarized: bool = True, ) -> tuple[str, dict[str, Any]]: """Run multiple benchmarks for the given model ID and list of benchmark configs.""" all_results = {} timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") start_time = time.perf_counter() n_configs = len(benchmark_configs) for i, config in enumerate(benchmark_configs): # Skip if already run if config.hash in all_results: self.logger.info(f"Skipping duplicate config {config.name} for model {model_id} ({i + 1}/{n_configs})") continue # Otherwise, run the benchmark self.setup_benchmark(model_id, config) self.logger.info( f"Running benchmark of model {model_id} with scenario: {config.name} ({i + 1}/{n_configs})" ) # Launch benchmark in a try/except block to avoid stopping the whole run if one benchmark fails try: result = self.run_benchmark(config, num_tokens_to_profile) except Exception as e: self.logger.error(f"Error running with scenario: {config.name}:\n{repr(e)}") result = None # Memoize all_results[config.hash] = { "metadata": BenchmarkMetadata( model_id=model_id, branch_name=self.branch_name, commit_id=self.commit_id, commit_message=self.commit_message, success=result is not None, ), "measurements": result if result is not None else BenchmarkResult(), "config": config, } # Cleanup model and save results self.cleanup() self.save_results(model_id, all_results, timestamp=timestamp, summarized=summarized) if len(all_results) < 1: raise RuntimeError("No benchmark was run successfully") if pretty_print_summary: print() print("=" * 100) print(f"Finished benchmarks in {time.perf_counter() - start_time:.2f} seconds") print(f"Total number of benchmarks: {len(all_results)}") print("First run metadata:") first_key = list(all_results.keys())[0] first_metadata = all_results[first_key]["metadata"].to_dict() hardware_info = first_metadata.pop("hardware_info") pretty_print_dict(first_metadata | hardware_info, tabs=1) for result in all_results.values(): print("=" * 100) print(f"Config: {result['config'].infer_name(compact=False)}\n") result["measurements"].pprint( batch_size=result["config"].batch_size, num_generated_tokens=result["config"].num_tokens_to_generate, tabs=1, ) print("=" * 100) return (timestamp, all_results) def save_results(self, model_name: str, results: dict, timestamp: str = "", summarized: bool = True) -> str: """Save benchmark results to JSON file.""" # Create model-specific subdirectory model_name = model_name.replace("/", "_") model_dir = os.path.join(self.output_dir, model_name) os.makedirs(model_dir, exist_ok=True) # Create filename with timestamp timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") if not timestamp else timestamp filename = f"{model_name}_benchmark_{timestamp}.json" filepath = os.path.join(model_dir, filename) # Convert results to dict converted_results = {} for cfg_hash in results.keys(): converted_results[cfg_hash] = { "metadata": results[cfg_hash]["metadata"].to_dict(), "measurements": results[cfg_hash]["measurements"].to_dict(summarized=summarized), "config": results[cfg_hash]["config"].to_dict(), } # Save to JSON file with open(filepath, "w") as f: f.write(compact_json_numeric_arrays(converted_results)) self.logger.info(f"Results saved to {filepath}") return filepath def push_results_to_hub(self, dataset_id: str, results: dict[Any, Any], timestamp: str) -> None: if PUSH_TO_HUB_TOKEN is None: raise ValueError( "PUSH_TO_HUB_TOKEN is not set, cannot push results to the Hub. When setting dataset_id, please also set the PUSH_TO_HUB_TOKEN environment variable." ) api = HfApi() n_results = len(results) for summarized in [False, True]: self.logger.info(f"Pushing {n_results} results to: {dataset_id} with {summarized = }") rows = [] for cfg_hash, entry in results.items(): row = { "benchmark_config_hash": cfg_hash, "config": entry["config"].to_dict(), "measurements": entry["measurements"].to_dict(summarized=summarized), "metadata": entry["metadata"].to_dict(), } rows.append(row) ds = Dataset.from_list(rows) with tempfile.TemporaryDirectory() as tmp: file_name = "summarized_results" if summarized else "full_results" jsonl_path = os.path.join(tmp, f"{file_name}.jsonl") with open(jsonl_path, "w") as f: json_lines = [] for ex in ds: json_lines.append(json.dumps(ex, ensure_ascii=False)) f.write("\n".join(json_lines)) # NOTE: we expect the repository to already exist timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") if not timestamp else timestamp file_name = file_name + "/" + f"benchmark_run_{timestamp}.jsonl" api.upload_file( path_or_fileobj=jsonl_path, path_in_repo=file_name, repo_id=dataset_id, repo_type="dataset", token=PUSH_TO_HUB_TOKEN, ) self.logger.info(f"Successfully uploaded results to: {dataset_id} with {summarized = }")