lemonade-sdk 8.1.11__py3-none-any.whl → 8.2.2__py3-none-any.whl

lemonade/tools/llamacpp/bench.py

@@ -2,9 +2,15 @@ import argparse
 import statistics
 from statistics import StatisticsError
 from lemonade.state import State
-from lemonade.cache import Keys
+from lemonade.tools.tool import Tool
 from lemonade.tools.llamacpp.utils import LlamaCppAdapter
-from lemonade.tools.bench import Bench
+from lemonade.tools.bench import (
+    Bench,
+    default_prompt_length,
+    default_iterations,
+    default_output_tokens,
+    default_warmup_runs,
+)
 
 
 class LlamaCppBench(Bench):
@@ -14,16 +20,6 @@ class LlamaCppBench(Bench):
 
     unique_name = "llamacpp-bench"
 
-    def __init__(self):
-        super().__init__()
-
-        # Additional statistics generated by this bench tool
-        self.status_stats.insert(
-            self.status_stats.index(Keys.TOKEN_GENERATION_TOKENS_PER_SECOND) + 1,
-            Keys.STD_DEV_TOKENS_PER_SECOND,
-        )
-        self.std_dev_token_generation_tokens_per_second_list = []
-
     @staticmethod
     def parser(add_help: bool = True) -> argparse.ArgumentParser:
         parser = __class__.helpful_parser(
@@ -33,8 +29,46 @@ class LlamaCppBench(Bench):
 
         parser = Bench.parser(parser)
 
+        parser.add_argument(
+            "--cli",
+            action="store_true",
+            help="Set this flag to use llama-cli.exe to benchmark model performance. "
+            "This executable will be called once per iteration.  Otherwise, "
+            "llama-bench.exe is used by default.  In this default behavior behavior, "
+            "the only valid prompt format is integer token lengths. Also, the "
+            "warmup-iterations parameter is ignored and the default value for number of "
+            "threads is 16.",
+        )
+
         return parser
 
+    def parse(self, state: State, args, known_only=True) -> argparse.Namespace:
+        """
+        Helper function to parse CLI arguments into the args expected by run()
+        """
+
+        # Call Tool parse method, NOT the Bench parse method
+        parsed_args = Tool.parse(self, state, args, known_only)
+
+        if parsed_args.cli:
+            parsed_args = super().parse(state, args, known_only)
+        else:
+            # Make sure prompts is a list of integers
+            if parsed_args.prompts is None:
+                parsed_args.prompts = [default_prompt_length]
+            prompt_ints = []
+            for prompt_item in parsed_args.prompts:
+                if prompt_item.isdigit():
+                    prompt_ints.append(int(prompt_item))
+                else:
+                    raise Exception(
+                        f"When not using the --cli flag to {self.unique_name}, the prompt format "
+                        "must be in integer format."
+                    )
+            parsed_args.prompts = prompt_ints
+
+        return parsed_args
+
     def run_prompt(
         self,
         state: State,
@@ -43,7 +77,7 @@ class LlamaCppBench(Bench):
         iterations: int,
         warmup_iterations: int,
         output_tokens: int,
-    ) -> State:
+    ):
         """
         Benchmark llama.cpp model that was loaded by LoadLlamaCpp.
         """
@@ -61,6 +95,7 @@ class LlamaCppBench(Bench):
 
         per_iteration_tokens_per_second = []
         per_iteration_time_to_first_token = []
+        per_iteration_peak_wset = []
 
         for iteration in range(iterations + warmup_iterations):
             try:
@@ -69,7 +104,10 @@ class LlamaCppBench(Bench):
                 model.time_to_first_token = None
                 model.tokens_per_second = None
                 raw_output, stderr = model.generate(
-                    prompt, max_new_tokens=output_tokens, return_raw=True
+                    prompt,
+                    max_new_tokens=output_tokens,
+                    return_raw=True,
+                    save_max_memory_used=self.save_max_memory_used,
                 )
 
                 if model.time_to_first_token is None or model.tokens_per_second is None:
@@ -85,6 +123,7 @@ class LlamaCppBench(Bench):
                 if iteration > warmup_iterations - 1:
                     per_iteration_tokens_per_second.append(model.tokens_per_second)
                     per_iteration_time_to_first_token.append(model.time_to_first_token)
+                    per_iteration_peak_wset.append(model.peak_wset)
 
                 report_progress_fn((iteration + 1) / (warmup_iterations + iterations))
 
@@ -115,21 +154,101 @@ class LlamaCppBench(Bench):
         except StatisticsError:
             # Less than 2 measurements
             self.std_dev_token_generation_tokens_per_second_list.append(None)
+        if self.save_max_memory_used:
+            filtered_list = [
+                item for item in per_iteration_peak_wset if item is not None
+            ]
+            mean_gb_used = (
+                None
+                if len(filtered_list) == 0
+                else statistics.mean(filtered_list) / 1024**3
+            )
+            self.max_memory_used_gb_list.append(mean_gb_used)
+
+    def run_llama_bench_exe(self, state, prompts, iterations, output_tokens):
+
+        if prompts is None:
+            prompts = [default_prompt_length]
+        elif isinstance(prompts, int):
+            prompts = [prompts]
+
+        state.save_stat("prompts", prompts)
+        state.save_stat("iterations", iterations)
+        state.save_stat("output_tokens", output_tokens)
 
-    def save_stats(self, state):
-        super().save_stats(state)
-
-        # Save additional statistics
-        if not all(
-            element is None
-            for element in self.std_dev_token_generation_tokens_per_second_list
-        ):
-            state.save_stat(
-                Keys.STD_DEV_TOKENS_PER_SECOND,
-                self.get_item_or_list(
-                    self.std_dev_token_generation_tokens_per_second_list
-                ),
+        counter = 0
+        report_progress_fn = lambda x: self.set_percent_progress(
+            100 * (counter + x) / len(prompts)
+        )
+        self.first_run_prompt = True
+        for counter, prompt in enumerate(prompts):
+            report_progress_fn(0)
+
+            self.run_prompt_llama_bench_exe(
+                state,
+                prompt,
+                iterations,
+                output_tokens,
             )
+            self.first_run_prompt = False
+
+        self.set_percent_progress(None)
+        self.save_stats(state)
+        return state
+
+    def run_prompt_llama_bench_exe(self, state, prompt, iterations, output_tokens):
+
+        model: LlamaCppAdapter = state.model
+        prompt_length, pp_tps, pp_tps_sd, tg_tps, tg_tps_sd, peak_wset = (
+            model.benchmark(prompt, iterations, output_tokens)
+        )
+        self.input_ids_len_list.append(prompt_length)
+        self.prefill_tokens_per_second_list.append(pp_tps)
+        self.std_dev_prefill_tokens_per_second_list.append(pp_tps_sd)
+        self.mean_time_to_first_token_list.append(prompt_length / pp_tps)
+        self.token_generation_tokens_per_second_list.append(tg_tps)
+        self.std_dev_token_generation_tokens_per_second_list.append(tg_tps_sd)
+        self.tokens_out_len_list.append(output_tokens * iterations)
+        if self.save_max_memory_used:
+            if peak_wset is not None:
+                self.max_memory_used_gb_list.append(peak_wset / 1024**3)
+            else:
+                self.max_memory_used_gb_list.append(None)
+
+    def run(
+        self,
+        state: State,
+        prompts: list[str] = None,
+        iterations: int = default_iterations,
+        warmup_iterations: int = default_warmup_runs,
+        output_tokens: int = default_output_tokens,
+        cli: bool = False,
+        **kwargs,
+    ) -> State:
+        """
+        Args:
+            - prompts: List of input prompts used as starting points for LLM text generation
+            - iterations: Number of benchmarking samples to take; results are
+                reported as the median and mean of the samples.
+            - warmup_iterations: Subset of the iterations to treat as warmup,
+                and not included in the results.
+            - output_tokens: Number of new tokens LLM to create.
+            - cli: Use multiple calls to llama-cpp.exe instead of llama-bench.exe
+            - kwargs: Additional parameters used by bench tools
+        """
+
+        # Check that state has the attribute model and it is a LlamaCPP model
+        if not hasattr(state, "model") or not isinstance(state.model, LlamaCppAdapter):
+            raise Exception("Load model using llamacpp-load first.")
+
+        if cli:
+            state = super().run(
+                state, prompts, iterations, warmup_iterations, output_tokens, **kwargs
+            )
+        else:
+            state = self.run_llama_bench_exe(state, prompts, iterations, output_tokens)
+
+        return state
 
 
 # This file was originally licensed under Apache 2.0. It has been modified.

lemonade/tools/llamacpp/load.py

@@ -93,9 +93,11 @@ class LoadLlamaCpp(FirstTool):
         from lemonade.tools.llamacpp.utils import (
             install_llamacpp,
             get_llama_cli_exe_path,
+            get_llama_bench_exe_path,
             get_llama_installed_version,
             parse_checkpoint,
             download_gguf,
+            resolve_local_gguf_model,
             get_local_checkpoint_path,
             LlamaCppTokenizerAdapter,
             LlamaCppAdapter,
@@ -103,6 +105,8 @@ class LoadLlamaCpp(FirstTool):
 
         install_llamacpp(backend)
 
+        extension = ""
+
         # Check if input is a local folder containing a .GGUF model
         if os.path.isdir(input):
             # input is a local folder
@@ -121,6 +125,17 @@ class LoadLlamaCpp(FirstTool):
                 )
             model_to_use = gguf_files[0]
             full_model_path = os.path.join(local_model_folder, model_to_use)
+            extension = ".gguf"
+
+        elif input.endswith(".gguf") and os.path.isfile(input):
+            # input is a local .gguf file
+            full_model_path = os.path.abspath(input)
+            checkpoint = "local_model"
+            state.checkpoint = checkpoint
+            state.save_stat(Keys.CHECKPOINT, checkpoint)
+            state.save_stat(Keys.LOCAL_MODEL_FOLDER, full_model_path)
+            model_to_use = os.path.basename(full_model_path)
+            extension = ".gguf"
 
         else:
             # Input is a model checkpoint
@@ -155,12 +170,21 @@ class LoadLlamaCpp(FirstTool):
                     )
 
             else:
+                # First, try to resolve from local cache to avoid unnecessary downloads
+                base_checkpoint, variant = parse_checkpoint(checkpoint)
+                snapshot_files = resolve_local_gguf_model(
+                    base_checkpoint, variant, None
+                )
+
+                # If not found locally, download from internet
+                if not snapshot_files:
+                    snapshot_files = download_gguf(checkpoint)
 
-                snapshot_files = download_gguf(checkpoint)
                 full_model_path = snapshot_files["variant"]
                 model_to_use = os.path.basename(full_model_path)
 
         llama_cli_exe_path = get_llama_cli_exe_path(backend)
+        llama_bench_exe_path = get_llama_bench_exe_path(backend)
         printing.log_info(f"Using llama_cli for GGUF model: {llama_cli_exe_path}")
 
         # Get the directory containing the executable for shared libraries
@@ -174,8 +198,10 @@ class LoadLlamaCpp(FirstTool):
             context_size=context_size,
             threads=threads,
             executable=llama_cli_exe_path,
+            bench_executable=llama_bench_exe_path,
             reasoning=reasoning,
             lib_dir=lib_dir,
+            state=state,
         )
         state.tokenizer = LlamaCppTokenizerAdapter()
         state.device = device
@@ -186,7 +212,9 @@ class LoadLlamaCpp(FirstTool):
             Keys.LLAMA_CLI_VERSION_INFO, get_llama_installed_version(backend)
         )
 
-        status.add_to_state(state=state, name=input, model=model_to_use)
+        status.add_to_state(
+            state=state, name=input, model=model_to_use, extension=extension
+        )
         return state