lemonade-sdk 7.0.0__py3-none-any.whl

lemonade/tools/quark/quark_quantize.py

@@ -0,0 +1,439 @@
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import torch
+from transformers import AutoProcessor
+from lemonade.state import State
+from lemonade.tools import Tool
+import lemonade.common.printing as printing
+import lemonade.common.build as build
+from lemonade_install.install import DEFAULT_QUARK_DIR
+
+
+class QuarkQuantize(Tool):
+    """
+    Quantize a model using the Quark Quantization tool.
+
+    This Tool performs the following steps:
+    1. Downloads and extracts necessary resources from AMD Quark Web Page.
+    2. Based on the target model, it prepares the model, tokenizer, and calibration data.
+    3. Optionally quantizes, freezes, and exports the model.
+    4. Optionally evaluates the model.
+
+    Required Input State:
+        - state.model.model: Pretrained model instance to be quantized.
+        - state.tokenizer: Tokenizer instance from Hugging Face.
+    Output:
+        - Modifies `state` with quantized and optionally exported model.
+
+    See docs/quark.md for more details.
+    """
+
+    unique_name = "quark-quantize"
+
+    def __init__(self):
+        super().__init__(monitor_message="Quark Quantizing model")
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Quantize a model using Quark",
+            add_help=add_help,
+        )
+        parser.add_argument(
+            "--device",
+            default="cpu",
+            choices=["cuda", "cpu"],
+            help="Device for running the quantizer",
+        )
+        parser.add_argument("--multi-gpu", action="store_true")
+        parser.add_argument(
+            "--data-type",
+            default="auto",
+            choices=["auto", "float16", "bfloat16", "float32"],
+            help="Input datatype of the model",
+        )
+        parser.add_argument(
+            "--seq-len", type=int, default=512, help="Sequence length of data"
+        )
+        parser.add_argument(
+            "--batch-size", type=int, default=1, help="Batch size for calibration."
+        )
+        parser.add_argument(
+            "--num-fewshot",
+            type=int,
+            default=None,
+            metavar="N",
+            help="Number of examples in few-shot context",
+        )
+        parser.add_argument(
+            "--output-dir", default=None, help="Output directory for exported model"
+        )
+        parser.add_argument(
+            "--no-weight-matrix-merge",
+            action="store_true",
+            help="If set, merges onnx model and weight \
+            together before export.\
+            By default, for onnx export, spits out a model.onnx and a model.weights",
+        )
+        parser.add_argument(
+            "--dataset",
+            default="pileval",
+            choices=[
+                "pileval",
+                "wikitext",
+                "pileval_for_awq_benchmark",
+                "wikitext_for_gptq_benchmark",
+                "HuggingFaceH4/ultrachat_200k",
+            ],
+            help="Dataset for calibration",
+        )
+        parser.add_argument(
+            "--num-calib-data",
+            type=int,
+            default=512,
+            help="Number of samples for calibration.",
+        )
+
+        # See docs/quark.md for more details.
+        parser.add_argument(
+            "--quant-scheme",
+            type=str,
+            default=None,
+            choices=[
+                "w_fp8_a_fp8",
+                "w_int4_per_channel_sym",
+                "w_uint4_per_group_asym",
+                "w_int4_per_group_sym",
+                "w_uint4_a_bfloat16_per_group_asym",
+                "w_int8_per_tensor_sym",
+                "w_int8_per_group_sym",
+                "w_uint8_per_group_asym",
+                "w_int8_a_int8_per_tensor_sym",
+                "w_int8_a_int8_per_tensor_sym_dynamic",
+                "w_uint8_a_uint8_per_tensor_asym",
+                "w_fp8_a_fp8_o_fp8",
+                "w_mx_fp8",
+                "w_mx_fp8_a_mx_fp8",
+                "w_int8_a_int8_per_token_dynamic",
+                "w_bfp16",
+                "w_bfp16_a_bfp16",
+                "w_mx6",
+                "w_mx6_a_mx6",
+                "w_fp8_per_channel_sym",
+                "w_int4_per_channel_asym",
+                "w_int4_per_group_asym",
+                "w_uint4_per_group_sym",
+                "w_uint4_per_channel_sym",
+                "w_uint4_per_channel_asym",
+                "w_int8_per_tensor_percentile",
+                "w_int8_per_tensor_mse",
+                "w_uint8_per_tensor_percentile",
+                "w_uint8_per_tensor_mse",
+                "w_mx_fp4_per_group_sym",
+                "w_mx_fp6_e3m2_per_group_sym",
+                "w_mx_fp6_e2m3_per_group_sym",
+                "w_mx_int8_per_group_sym",
+                "w_uint4_per_channel_a_int8_per_tensor",
+                "w_uint4_per_group_a_int8_per_tensor",
+                "w_bfp16_per_group_sym",
+                None,
+            ],
+            help="Supported quantization schemes in Quark",
+        )
+        parser.add_argument(
+            "--quant-algo",
+            type=str,
+            default=None,
+            choices=["awq", "gptq", "autosmoothquant", None],
+            help="Support quantization algorithms in Quark",
+        )
+        parser.add_argument(
+            "--pre-optimization-config-file-path",
+            type=str,
+            default=None,
+            help="The JSON file path of pre-optimization config",
+        )
+        parser.add_argument(
+            "--quant-algo-config-file-path",
+            type=str,
+            default=None,
+            help="The JSON file path of quantization algorithm config",
+        )
+        parser.add_argument(
+            "--group-size",
+            type=int,
+            default=128,
+            help="Group size for per_group quantization",
+        )
+        parser.add_argument(
+            "--pack-method",
+            type=str,
+            default="reorder",
+            choices=["order", "reorder"],
+            help="Pack method for awq_export",
+        )
+        parser.add_argument(
+            "--exclude-layers",
+            type=str,
+            nargs="*",
+            default=None,
+            help="List of layers to exclude from quantization.",
+        )
+        parser.add_argument(
+            "--kv-cache-dtype",
+            default=None,
+            choices=["fp8", None],
+            help="KV Cache dtype.",
+        )
+        parser.add_argument(
+            "--pre-quantization-optimization",
+            action="append",
+            default=[],
+            choices=["rotation", "smoothquant"],
+            help="Pre Quantization Optimization.",
+        )
+        parser.add_argument(
+            "--model-export",
+            default=None,
+            action="append",
+            choices=[
+                None,
+                "onnx",
+                "vllm_adopted_safetensors",
+                "quark_safetensors",
+                "gguf",
+            ],
+            help="Model export format",
+        )
+        parser.add_argument(
+            "--custom-mode",
+            default="quark",
+            type=str,
+            choices=["quark", "awq", "fp8"],
+            help="Custom mode for export \
+                This is especially relevant for npu/hybrid export",
+        )
+        parser.add_argument(
+            "--torch-compile",
+            action="store_true",
+            help="Compile the quantized model using torch.compile",
+        )
+        parser.add_argument(
+            "--params-save", action="store_true", help="Save model params"
+        )
+        parser.add_argument(
+            "--save-dir",
+            help="Directory to save model parameters as \
+            safetensors or pth, in the case when --params_save is used.",
+        )
+        parser.add_argument(
+            "--log-severity-level", type=int, default=3, help="DEBUG=1, INFO=2, ERROR=3"
+        )
+        parser.add_argument("--skip-quantization", action="store_true")
+
+        return parser
+
+    def run(self, state: State, **kwargs) -> State:
+        """
+        Executes the QuarkQuantize process.
+
+        Args:
+            state (State): The current state of the process, containing necessary
+            information such as cache directory and build name.
+            **kwargs: Additional keyword arguments that may include:
+            - output_dir (str): Directory to save the output model.
+            - safetensors_model_dir (str): Directory to save the safetensors model.
+            - save_dir (str): Directory to save model parameters.
+            - safetensors_path (str): Path to the safetensors model.
+            - quant_algo (str): The quantization algorithm to use.
+            - quant_algo_config_file_path (str): Path to the quantization algorithm
+              configuration file.
+            - model_dir (str): Directory of the model.
+        Returns:
+            State: The updated state after the quantization process.
+        Raises:
+            Exception: If an error occurs during the QuarkQuantize process
+            and when installation path does not exist.
+        """
+
+        try:
+
+            if os.path.isdir(DEFAULT_QUARK_DIR):
+                quark_llm_path = os.path.join(
+                    DEFAULT_QUARK_DIR, "examples", "torch", "language_modeling"
+                )
+                sys.path.extend([quark_llm_path])
+            else:
+                raise FileNotFoundError(
+                    f"The directory {DEFAULT_QUARK_DIR} does not exist. \
+                        Please check your installation."
+                )
+            model_build_path = os.path.join(
+                build.output_dir(state.cache_dir, state.build_name)
+            )
+            model_export_path = os.path.join(
+                model_build_path,
+                "exported_model",
+                kwargs.get("quant_scheme"),
+                kwargs.get("quant_algo"),
+            )
+            # Set default paths only if current values are None
+            if kwargs.get("model_dir") is None:
+                kwargs["model_dir"] = model_build_path
+            if kwargs.get("output_dir") is None:
+                kwargs["output_dir"] = model_export_path
+            if kwargs.get("save_dir") is None:
+                kwargs["save_dir"] = os.path.join(model_export_path, "model_params")
+
+            from llm_utils.model_preparation import get_model_type
+
+            model_type = get_model_type(state.model.model)
+
+            quant_algo = kwargs.get("quant_algo")
+            kwargs["quant_algo_config_file_path"] = os.path.join(
+                quark_llm_path,
+                "llm_ptq",
+                "models",
+                model_type,
+                f"{quant_algo}_config.json",
+            )
+
+            self._quantize(state, **kwargs)
+
+        except Exception as e:
+            printing.log_error(f"Error during the QuarkQuantize process: {e}")
+            raise
+        return state
+
+    def _quantize(self, state: State, **kwargs) -> None:
+        """
+        Main quantization and export process.
+
+        This method is responsible for:
+        - Loading the model and tokenizer.
+        - Preparing the calibration dataset.
+        - Quantizing the model.
+        - Optionally exporting, compiling, and evaluating the model.
+        """
+
+        model = state.model.model
+        tokenizer = state.tokenizer
+
+        # Importing quark utils after adding to sys.path
+        from llm_utils.data_preparation import get_calib_dataloader
+        from llm_utils.model_preparation import get_model_type
+        from llm_ptq.configuration_preparation import get_config, get_export_config
+        from quark.torch import ModelQuantizer, ModelExporter, save_params
+
+        # 1. Load Model
+        printing.log_info("Loading model ...")
+        model_type = get_model_type(model)
+
+        # [mllama specifics]
+        if model_type == "mllama" and kwargs.get("model_export") is not None:
+            processor = AutoProcessor.from_pretrained(kwargs.get("model_dir"))
+            export_dir = Path(kwargs.get("output_dir"))
+            export_dir.mkdir(parents=True, exist_ok=True)
+            processor.save_pretrained(kwargs.get("output_dir"))
+
+        # 2. Load dataset
+        printing.log_info("Loading dataset ...")
+        main_device = model.device if kwargs.get("multi_gpu") else kwargs.get("device")
+        calib_dataloader = get_calib_dataloader(
+            dataset_name=kwargs.get("dataset"),
+            tokenizer=tokenizer,
+            batch_size=1,
+            num_calib_data=kwargs.get("num_calib_data"),
+            seqlen=kwargs.get("seq_len"),
+            device=main_device,
+        )
+
+        # 3. Quantize model
+        if not kwargs.get("skip_quantization"):
+            printing.log_info("Starting quantization process ...")
+            args = argparse.Namespace(**kwargs)
+            quant_config = get_config(args, model_type)
+            quant_config.log_severity_level = kwargs.get("log_severity_level", 3)
+            quantizer = ModelQuantizer(quant_config)
+            model = quantizer.quantize_model(model, calib_dataloader)
+            printing.log_info("Quantization completed.")
+
+            if (
+                kwargs.get("model_export") is not None
+                or kwargs.get("params_save")
+                or kwargs.get("torch_compile")
+            ):
+                printing.log_info("Freezing the quantized model ...")
+                model = quantizer.freeze(model)
+
+        # 4. Export model
+        if kwargs.get("model_export") is not None:
+            printing.log_info("Exporting the model ...")
+            export_path = kwargs.get("output_dir")
+
+            args = argparse.Namespace(**kwargs)
+            export_config = get_export_config(args, model_type)
+            exporter = ModelExporter(config=export_config, export_dir=export_path)
+            if "quark_safetensors" in kwargs.get("model_export"):
+                printing.log_info("Exporting quark native json and safetensors...")
+                with torch.no_grad():
+                    quant_config = get_config(args, model_type)
+                    exporter.export_model_info(
+                        model,
+                        quant_config=quant_config,
+                        tokenizer=tokenizer,
+                        custom_mode=kwargs.get("custom_mode"),
+                    )
+            if "vllm_adopted_safetensors" in kwargs.get("model_export"):
+                printing.log_info("Exporting vllm adopted json and safetensors...")
+                with torch.inference_mode():
+                    exporter.export_model_info(
+                        model,
+                        model_type=model_type,
+                        model_dtype=state.dtype,
+                        export_type="vllm-adopt",
+                    )
+            if "onnx" in kwargs.get("model_export"):
+                printing.log_info("Exporting onnx graph...")
+                with torch.inference_mode():
+                    batch_iter = iter(calib_dataloader)
+                    input_args = next(batch_iter)
+                    if kwargs.get("quant_scheme") in [
+                        "w_int4_per_channel_sym",
+                        "w_uint4_per_group_asym",
+                        "w_int4_per_group_sym",
+                        "w_uint4_a_bfloat16_per_group_asym",
+                    ]:
+                        uint4_int4_flag = True
+                    else:
+                        uint4_int4_flag = False
+                    exporter.export_onnx_model(
+                        model, input_args, uint4_int4_flag=uint4_int4_flag
+                    )
+            if "gguf" in kwargs.get("model_export"):
+                printing.log_info("Exporting gguf model...")
+                with torch.inference_mode():
+                    exporter.export_gguf_model(
+                        model, kwargs.get("model_dir"), model_type
+                    )
+
+        # 6. [Optional] Compile model
+        if kwargs.get("torch_compile"):
+            printing.log_info("torch.compile...")
+            model = torch.compile(model)
+
+        # 7. Save model parameters
+        if kwargs.get("params_save"):
+            printing.log_info("Saving model parameters ...")
+            save_params(model, model_type=model_type, export_dir=kwargs.get("save_dir"))
+
+        state.model.model = model
+        state.dtype = model.dtype
+        printing.log_info("QuarkQuantize process completed.")
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/report/llm_report.py

@@ -0,0 +1,203 @@
+import argparse
+import csv
+import os
+from pathlib import Path
+import re
+from typing import List
+import lemonade.common.printing as printing
+import lemonade.common.filesystem as fs
+from lemonade.tools.management_tools import ManagementTool
+from lemonade.cache import DEFAULT_CACHE_DIR
+from lemonade.tools.report.table import LemonadeTable, LemonadePerfTable
+
+
+class LemonadeReport(ManagementTool):
+    """
+    Analyzes the input lemonade cache(s) and produces an aggregated report
+    in csv format that contains the build stats for all builds in all cache(s).
+
+    In addition, summary information is printed to the console and saved to a text file.
+
+    When the --perf flag is used, then a performance report is generated that summarizes
+    the performance data for different models.  In this case, only the data used in the
+    text table is saved to the csv format file.
+    """
+
+    unique_name = "report"
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Export statistics from each lemonade run to a CSV file",
+            add_help=add_help,
+        )
+
+        parser.add_argument(
+            "-i",
+            "--input-caches",
+            nargs="*",
+            default=[DEFAULT_CACHE_DIR],
+            help=(
+                "One or more lemonade cache directories to use to generate the report "
+                f"(defaults to {DEFAULT_CACHE_DIR})"
+            ),
+        )
+
+        parser.add_argument(
+            "-o",
+            "--output-dir",
+            help="Path to folder where reports will be saved "
+            "(defaults to current working directory)",
+            required=False,
+            default=os.getcwd(),
+        )
+
+        parser.add_argument(
+            "--no-save",
+            action="store_true",
+            help="Don't save output to TXT and CSV files",
+        )
+
+        parser.add_argument(
+            "--perf",
+            action="store_true",
+            help="Produce the performance table instead of the regular table",
+        )
+
+        parser.add_argument(
+            "--device",
+            default=None,
+            help="In the --perf table, only include output for the specified device "
+            "(e.g., cpu, igpu, npu, hybrid)",
+        )
+
+        parser.add_argument(
+            "--dtype",
+            default=None,
+            help="In the --perf table, only include output for the specified datatype "
+            "(e.g., float32, int4)",
+        )
+
+        parser.add_argument(
+            "--model",
+            default=None,
+            help="In the --perf table, only include output for builds with name that contains"
+            " the specified string (e.g., Llama).  The string match is case-insensitive.",
+        )
+
+        parser.add_argument(
+            "--days",
+            "-d",
+            type=int,
+            metavar="N",
+            default=None,
+            help="In the --perf table, only include output for builds from the last N days.",
+        )
+
+        parser.add_argument(
+            "--merge",
+            action="store_true",
+            help="In the --perf table, merge results from different builds into the same row "
+            "as long as model, device, datatype, system info and package version are the same.",
+        )
+
+        parser.add_argument(
+            "--lean",
+            action="store_true",
+            help="In the --perf table, don't include the system info and sw package versions.",
+        )
+
+        return parser
+
+    def parse(self, args, known_only=True) -> argparse.Namespace:
+        """
+        Helper function to parse CLI arguments into the args expected by run()
+        """
+        parsed_args = super().parse(args, known_only)
+
+        if not parsed_args.perf:
+            # Check that none of the perf specific flags are set
+            perf_args = [
+                parsed_args.device,
+                parsed_args.dtype,
+                parsed_args.days,
+                parsed_args.merge,
+                parsed_args.lean,
+            ]
+            if not all(arg is None or arg is False for arg in perf_args):
+                raise ValueError(
+                    "Invalid arguments for regular report.  Did you miss the --perf argument?"
+                    "  See `lemonade report -h` for help."
+                )
+
+        return parsed_args
+
+    def run(
+        self,
+        _,
+        input_caches: List[str] = None,
+        output_dir: str = os.getcwd(),
+        no_save: bool = False,
+        perf: bool = False,
+        device: str = None,
+        dtype: str = None,
+        model: str = None,
+        days: int = None,
+        merge: bool = False,
+        lean: bool = False,
+    ):
+        # Process input arguments
+        cache_dirs = [os.path.expanduser(dir) for dir in input_caches]
+        cache_dirs = fs.expand_inputs(cache_dirs)
+        report_dir = os.path.expanduser(output_dir)
+
+        if perf:
+            table = LemonadePerfTable(device, dtype, model, days, merge, lean)
+        else:
+            table = LemonadeTable()
+
+        # Find builds and load stats
+        table.find_builds(cache_dirs, model)
+        table.load_stats()
+        table.sort_stats()
+
+        # Print message if there are no stats
+        if len(table.all_stats) == 0:
+            printing.log_info("No relevant cached build data found")
+            return
+
+        # Print table to stdout
+        print()
+        print(table)
+
+        if no_save:
+            return
+
+        # Name report file
+        report_path = os.path.join(report_dir, table.get_report_name())
+        txt_path = re.sub(".csv$", ".txt", report_path)
+        Path(report_dir).mkdir(parents=True, exist_ok=True)
+
+        # Create the report to save to CSV
+        report, column_headers = table.create_csv_report()
+
+        # Populate results spreadsheet
+        with open(report_path, "w", newline="", encoding="utf8") as spreadsheet:
+            writer = csv.writer(spreadsheet)
+            writer.writerow(column_headers)
+            for entry in report:
+                writer.writerow([entry[col] for col in column_headers])
+
+        # Save the text report
+        with open(txt_path, "w", encoding="utf-8") as file:
+            print(table, file=file)
+
+        # Print message with the output file path
+        printing.log("Report text saved at ")
+        printing.logn(str(txt_path), printing.Colors.OKGREEN)
+        printing.log("Report spreadsheet saved at ")
+        printing.logn(str(report_path), printing.Colors.OKGREEN)
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD