lemonade-sdk 7.0.0__py3-none-any.whl

lemonade/tools/huggingface_bench.py

@@ -0,0 +1,267 @@
+import argparse
+from typing import List, Tuple
+import time
+import statistics
+from statistics import StatisticsError
+from contextlib import nullcontext
+import torch
+from lemonade.state import State
+from lemonade.cache import Keys
+from lemonade.tools.bench import Bench
+
+default_beams = 1
+
+
+def benchmark_huggingface_llm(
+    model: torch.nn.Module,
+    tokenizer,
+    input_ids,
+    dtype,
+    num_beams: int,
+    target_output_tokens: int,
+    iterations: int,
+    warmup_iterations: int,
+    report_progress_fn,
+) -> List[Tuple[float, int]]:
+
+    amp_enabled = True if (dtype == torch.float16 or dtype == torch.bfloat16) else False
+    # The "if amp_enabled else nullcontext()" is to get around a bug in PyTorch 2.1
+    # where torch.cpu.amp.autocast(enabled=False) does nothing
+    with (
+        torch.cpu.amp.autocast(enabled=amp_enabled, dtype=dtype)
+        if amp_enabled
+        else nullcontext()
+    ):
+
+        per_iteration_result = []
+        tokens_out_len_list = []
+
+        # Early stopping is only a valid parameter with multiple beams
+        early_stopping = num_beams > 1
+
+        with torch.no_grad(), torch.inference_mode():
+            # Don't capture time for warmup
+            for count in range(warmup_iterations):
+                outputs = model.generate(
+                    input_ids,
+                    num_beams=num_beams,
+                    max_new_tokens=target_output_tokens,
+                    min_new_tokens=target_output_tokens,
+                    early_stopping=early_stopping,
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+                tokens_out_len_list.append(outputs.shape[1] - input_ids.shape[1])
+                report_progress_fn((count + 1) / (warmup_iterations + iterations))
+
+            for count in range(iterations):
+                # CUDA synchronization is required prior to GPU benchmarking
+                # This has no negative effect on CPU-only benchmarks, and is more robust than
+                # checking `model.device == "cuda"` since it applies to multi-GPU environments
+                # Synchronization is done before collecting the start time because this will
+                # ensure that the GPU has finished initialization tasks such as loading weights
+                if torch.cuda.is_available():
+                    torch.cuda.synchronize()
+                start_time = time.perf_counter()
+
+                outputs = model.generate(
+                    input_ids,
+                    num_beams=num_beams,
+                    max_new_tokens=target_output_tokens,
+                    min_new_tokens=target_output_tokens,
+                    early_stopping=early_stopping,
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+
+                if torch.cuda.is_available():
+                    torch.cuda.synchronize()
+                end_time = time.perf_counter()
+
+                latency = end_time - start_time
+
+                token_len = outputs.shape[1] - input_ids.shape[1]
+                tokens_out_len_list.append(token_len)
+
+                # Only count an iteration if it produced enough tokens
+                if token_len >= target_output_tokens:
+                    per_iteration_result.append((latency, token_len))
+
+                report_progress_fn(
+                    (warmup_iterations + count + 1) / (warmup_iterations + iterations)
+                )
+
+        if not per_iteration_result:
+            raise Bench.not_enough_tokens(target_output_tokens)
+
+    return per_iteration_result, tokens_out_len_list
+
+
+class HuggingfaceBench(Bench):
+    """
+    Benchmarks the performance of the generate() method of an LLM loaded from
+    Huggingface Transformers (or any object that supports a
+    huggingface-like generate() method).
+
+    Required input state:
+        - DTYPE: data type of the model; used to determine if AMP should be
+            enabled to convert the input data type to match the model data
+            type.
+        - MODEL: huggingface-like instance to benchmark.
+        - INPUTS: model inputs to pass to generate() during benchmarking.
+
+    Output state produced: None
+
+    """
+
+    unique_name = "huggingface-bench"
+
+    @staticmethod
+    def parser(parser: argparse.ArgumentParser = None, add_help: bool = True):
+        # Allow inherited classes to initialize and pass in a parser, add parameters to it if so
+        if parser is None:
+            parser = __class__.helpful_parser(
+                short_description="Benchmark a huggingface-style PyTorch LLM",
+                add_help=add_help,
+            )
+
+        parser = Bench.parser(parser)
+
+        parser.add_argument(
+            "--num-beams",
+            required=False,
+            type=int,
+            default=default_beams,
+            help=f"Number of beams for the LLM to use (default: {default_beams})",
+        )
+
+        return parser
+
+    def get_prompt_str(self, state, token_length):
+        """
+        Returns a string with the prescribed token length.
+        """
+        model = state.model
+        tokenizer = state.tokenizer
+        test_prompt = "word " * (token_length - 2)
+        input_ids = (
+            tokenizer(test_prompt, return_tensors="pt")
+            .to(device=model.device)
+            .input_ids
+        )
+        test_token_length = input_ids.shape[1]
+        delta = test_token_length - token_length
+        if delta == 0:
+            return test_prompt
+        return "word " * max(token_length - 2 - delta, 0)
+
+    def run_prompt(
+        self,
+        state: State,
+        report_progress_fn,
+        prompt: str,
+        iterations: int,
+        warmup_iterations: int,
+        output_tokens: int,
+        num_beams: int = default_beams,
+    ) -> State:
+        """
+        We don't have access to the internal timings of generate(), so time to first
+        token (TTFT, aka prefill latency) and token/s are calculated using the following formulae:
+            prefill_latency = latency of generate(output_tokens=1)
+            execution_latency = latency of generate(output_tokens=output_tokens)
+            tokens_per_second = (new_tokens - 1) / (execution_latency - prefill_latency)
+        """
+
+        if self.first_run_prompt:
+            if vars(state).get(Keys.MODEL) is None:
+                raise ValueError(
+                    f"{self.__class__.__name__} requires that a model be passed from another tool"
+                )
+            if (
+                vars(state).get("num_beams")
+                and vars(state).get("num_beams") != num_beams
+            ):
+                raise ValueError(
+                    f"Number of beams was set to {vars(state).get('num_beams')} "
+                    f"in a previous tool, but it is set to {num_beams} in "
+                    "this tool. The values must be the same."
+                )
+
+            # Save benchmarking parameters
+            state.save_stat("num_beams", num_beams)
+
+        model = state.model
+        tokenizer = state.tokenizer
+        dtype = state.dtype
+
+        # Generate the input_ids outside the benchmarking function to make sure
+        # the same input_ids are used everywhere
+        input_ids = (
+            tokenizer(prompt, return_tensors="pt").to(device=model.device).input_ids
+        )
+        self.input_ids_len_list.append(input_ids.shape[1])
+
+        prefill_report_progress_fn = lambda x: report_progress_fn(0.5 * x)
+
+        # Benchmark prefill time (time to first token)
+        prefill_per_iteration_result, tokens_out_len_list = benchmark_huggingface_llm(
+            model=model,
+            tokenizer=tokenizer,
+            input_ids=input_ids,
+            dtype=dtype,
+            num_beams=num_beams,
+            target_output_tokens=1,
+            iterations=iterations,
+            warmup_iterations=warmup_iterations,
+            report_progress_fn=prefill_report_progress_fn,
+        )
+        self.tokens_out_len_list += tokens_out_len_list
+
+        time_to_first_token_per_iteration = [
+            latency for latency, _ in prefill_per_iteration_result
+        ]
+        mean_time_to_first_token = statistics.mean(time_to_first_token_per_iteration)
+        self.mean_time_to_first_token_list.append(mean_time_to_first_token)
+        self.prefill_tokens_per_second_list.append(
+            input_ids.shape[1] / mean_time_to_first_token
+        )
+        try:
+            self.std_dev_time_to_first_token_list.append(
+                statistics.stdev(time_to_first_token_per_iteration)
+            )
+        except StatisticsError:
+            # Less than 2 measurements
+            self.std_dev_time_to_first_token_list.append(None)
+
+        decode_report_progress_fn = lambda x: report_progress_fn(0.5 + 0.5 * x)
+
+        # Benchmark generation of all tokens
+        decode_per_iteration_result, tokens_out_len_list = benchmark_huggingface_llm(
+            model=model,
+            tokenizer=tokenizer,
+            input_ids=input_ids,
+            dtype=dtype,
+            num_beams=num_beams,
+            target_output_tokens=output_tokens,
+            iterations=iterations,
+            warmup_iterations=warmup_iterations,
+            report_progress_fn=decode_report_progress_fn,
+        )
+        self.tokens_out_len_list += tokens_out_len_list
+
+        execution_latency_per_iteration = [
+            latency for latency, _ in decode_per_iteration_result
+        ]
+        token_len_per_iteration = [
+            token_len for _, token_len in decode_per_iteration_result
+        ]
+        mean_execution_latency = statistics.mean(execution_latency_per_iteration)
+        mean_decode_latency = mean_execution_latency - mean_time_to_first_token
+        mean_token_len = statistics.mean(token_len_per_iteration)
+        # Subtract 1 so that we don't count the prefill token
+        self.token_generation_tokens_per_second_list.append(
+            (mean_token_len - 1) / mean_decode_latency
+        )
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD