import gc import json import logging import os import pathlib import re import tempfile import time from contextlib import nullcontext 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, CompileConfig, 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 get_sdpa_backend(backend_name: str | None) -> torch.nn.attention.SDPBackend | None: """Get the SDPA backend enum from string name.""" if backend_name is None: return None try: backend_map = { "math": torch.nn.attention.SDPBackend.MATH, "flash_attention": torch.nn.attention.SDPBackend.FLASH_ATTENTION, "efficient_attention": torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, "cudnn_attention": torch.nn.attention.SDPBackend.CUDNN_ATTENTION, } return backend_map.get(backend_name.lower()) except AttributeError: # torch.nn.attention.SDPBackend not available in older torch versions return None def flush_memory(): """Flush GPU memory and run garbage collection.""" gc.collect() # 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.reset_max_memory_allocated() torch.cuda.reset_peak_memory_stats() 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 gen_config = GenerationConfig( do_sample=False, top_p=1.0, temperature=1.0, max_new_tokens=config.num_tokens_to_generate ) # Prepare compile config if config.compile_mode is not None: gen_config.compile_config = CompileConfig(mode=config.compile_mode, options=config.compile_options) gen_config.cache_implementation = "static" # 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=gen_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, model_id: str, config: BenchmarkConfig, num_tokens_to_profile: int = 0 ) -> dict[str, Any] | None: """Run a single benchmark with the given model ID and config.""" sdpa_ctx = nullcontext() if config.attn_implementation == "sdpa": sdpa_backend = get_sdpa_backend(config.sdpa_backend) sdpa_ctx = torch.nn.attention.sdpa_kernel(sdpa_backend) with sdpa_ctx, torch.no_grad(): self.logger.info(f"Running benchmark scenario: {config.name}") # Quick validation: try one measurement first to see if this scenario works flush_memory() e2e_latency, token_generation_times, shape_and_decoded_output, gpu_metrics = self.time_generate( max_new_tokens=1, gpu_monitor=None ) if e2e_latency < 0: self.logger.warning(f"Skipping config {config.name}: {e2e_latency = } (no GPU monitoring)") return None # Warmup runs self.logger.info(f"Warming up with {config.warmup_iterations} iterations...") for _ in trange(config.warmup_iterations): _ = self.time_generate(max_new_tokens=config.num_tokens_to_generate) 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): e2e_latency, token_generation_times, shape_and_decoded_output, gpu_metrics = self.time_generate( max_new_tokens=config.num_tokens_to_generate, gpu_monitor=(GPUMonitor(logger=self.logger) if config.gpu_monitoring else None), ) result.accumulate(e2e_latency, token_generation_times, 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 { "metadata": BenchmarkMetadata( model_id=model_id, branch_name=self.branch_name, commit_id=self.commit_id, commit_message=self.commit_message, ), "measurements": result, "config": config, } def time_generate( self, max_new_tokens: int, gpu_monitor: GPUMonitor | None = None, ) -> tuple[float, list[float], str, GPURawMetrics | None]: """Time the latency of a call to model.generate() with the given (inputs) and (max_new_tokens).""" # Prepare gpu monitoring if needed if gpu_monitor is not None: gpu_monitor.start() # Prepare streamer streamer = BenchmarkStreamer() # Generate and time wall_time_0 = time.perf_counter() outputs = self.model.generate( **self.inputs, max_new_tokens=max_new_tokens, streamer=streamer, ) wall_time_1 = time.perf_counter() # Stop gpu monitoring if needed gpu_metrics = gpu_monitor.stop_and_collect() if gpu_monitor is not None else None # Check if generation had the right number of tokens input_tokens = self.inputs["input_ids"].size(-1) batch_size, output_tokens = outputs.shape new_tokens = output_tokens - input_tokens if new_tokens != max_new_tokens: raise RuntimeError(f"Generated {new_tokens} tokens, expected {max_new_tokens}") # Decode outputs decoded_output = self.tokenizer.decode(outputs[0, input_tokens:], skip_special_tokens=True) shape_and_decoded_output = f"{tuple(outputs.shape)} | {decoded_output}" # Compute intermediate quantities e2e_latency = wall_time_1 - wall_time_0 token_generation_times = [t - wall_time_0 for t in streamer.timestamps[1:]] return e2e_latency, token_generation_times, 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") def run_benchmarks( self, model_id: str, benchmark_configs: list[BenchmarkConfig], num_tokens_to_profile: int = 0, pretty_print_summary: 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): # Handle SDPA backend if not determined by the config (needs to be done before skipping duplicates) if config.attn_implementation == "sdpa" and config.sdpa_backend is None: default_backend = "flash_attention" # FIXME: torch has a _cur_sdpa_kernel_backends but it fails self.logger.warning(f"No SDPA backend provided, using {default_backend} instead.") config.sdpa_backend = default_backend # 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: results = self.run_benchmark(model_id, config, num_tokens_to_profile) if results is not None: all_results[config.hash] = results except Exception as e: self.logger.error(f"Error running with scenario: {config.name}:\n{repr(e)}") # Cleanup model and save results self.cleanup() self.save_results(model_id, all_results, timestamp=timestamp) 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)}") if len(all_results) > 0: 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, tabs=1) print("=" * 100) return (timestamp, all_results) def save_results(self, model_name: str, results: dict, timestamp: str = "") -> 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(), "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." ) n_results = len(results) self.logger.info(f"Pushing {n_results} results to: {dataset_id}") rows = [] for cfg_hash, entry in results.items(): row = { "benchmark_config_hash": cfg_hash, "config": entry["config"].to_dict(), "measurements": entry["measurements"].to_dict(), "metadata": entry["metadata"].to_dict(), } rows.append(row) ds = Dataset.from_list(rows) with tempfile.TemporaryDirectory() as tmp: jsonl_path = os.path.join(tmp, "data.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)) api = HfApi() # 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 = 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"Succesfully uploaded results to: {dataset_id}")