Add dynamic quantization from @Xenova

experiments/dynamic_quantization_export.py

@@ -0,0 +1,473 @@
+"""
+This script is used to convert HuggingFace models to ONNX format and optionally quantize the model using dynamic quantization. 
+This is courtesy of Joshua aka @Xenova
+"""
+
+import json
+import os
+import shutil
+from dataclasses import dataclass, field
+from typing import Optional, Set
+
+import onnx
+from onnxruntime.quantization import QuantType, quantize_dynamic
+from optimum.exporters.onnx import export_models, main_export
+from optimum.exporters.tasks import TasksManager
+from tqdm import tqdm
+from transformers import AutoConfig, AutoTokenizer, HfArgumentParser
+
+DEFAULT_QUANTIZE_PARAMS = {
+    "per_channel": True,
+    "reduce_range": True,
+}
+
+MODEL_SPECIFIC_QUANTIZE_PARAMS = {
+    # Decoder-only models
+    "codegen": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt2": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt_bigcode": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gptj": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt-neo": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt-neox": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "mpt": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "bloom": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "llama": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "opt": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "mistral": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "falcon": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "phi": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "qwen2": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    # Encoder-decoder models
+    "whisper": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "vision-encoder-decoder": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    # Encoder-only models
+    "owlv2": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+}
+
+MODELS_WITHOUT_TOKENIZERS = [
+    "wav2vec2",
+    "wav2vec2-bert",
+    "wavlm",
+    "hubert",
+]
+
+
+@dataclass
+class ConversionArguments:
+    """
+    Arguments used for converting HuggingFace models to onnx.
+    """
+
+    model_id: str = field(metadata={"help": "Model identifier"})
+    tokenizer_id: str = field(default=None, metadata={"help": "Tokenizer identifier (if different to `model_id`)"})
+    quantize: bool = field(default=False, metadata={"help": "Whether to quantize the model."})
+    output_parent_dir: str = field(
+        default="./models/", metadata={"help": "Path where the converted model will be saved to."}
+    )
+
+    task: Optional[str] = field(
+        default="auto",
+        metadata={
+            "help": (
+                "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
+                f" {str(TasksManager.get_all_tasks())}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder."
+            )
+        },
+    )
+
+    opset: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If specified, ONNX opset version to export the model with. Otherwise, the default opset will be used."
+            )
+        },
+    )
+
+    device: str = field(default="cpu", metadata={"help": "The device to use to do the export."})
+    skip_validation: bool = field(default=False, metadata={"help": "Whether to skip validation of the converted model"})
+
+    per_channel: bool = field(default=None, metadata={"help": "Whether to quantize weights per channel"})
+    reduce_range: bool = field(
+        default=None,
+        metadata={
+            "help": "Whether to quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode"
+        },
+    )
+
+    output_attentions: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to output attentions from the model. NOTE: This is only supported for whisper models right now."
+        },
+    )
+
+    split_modalities: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to split multimodal models. NOTE: This is only supported for CLIP models right now."
+        },
+    )
+
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={
+            "help": "Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories"
+            "you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository."
+        },
+    )
+
+    custom_onnx_configs: str = field(
+        default=None,
+        metadata={
+            "help": "Experimental usage: override the default ONNX config used for the given model. This argument may be useful for advanced users "
+            "that desire a finer-grained control on the export."
+        },
+    )
+
+
+def get_operators(model: onnx.ModelProto) -> Set[str]:
+    operators = set()
+
+    def traverse_graph(graph):
+        for node in graph.node:
+            operators.add(node.op_type)
+            for attr in node.attribute:
+                if attr.type == onnx.AttributeProto.GRAPH:
+                    subgraph = attr.g
+                    traverse_graph(subgraph)
+
+    traverse_graph(model.graph)
+    return operators
+
+
+def quantize(model_names_or_paths, **quantize_kwargs):
+    """
+    Quantize the weights of the model from float32 to int8 to allow very efficient inference on modern CPU
+
+    Uses unsigned ints for activation values, signed ints for weights, per
+    https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection
+    it is faster on most CPU architectures
+    Args:
+        onnx_model_path: Path to location the exported ONNX model is stored
+    Returns: The Path generated for the quantized
+    """
+
+    quantize_config = dict(**quantize_kwargs, per_model_config={})
+
+    for model in tqdm(model_names_or_paths, desc="Quantizing"):
+        directory_path = os.path.dirname(model)
+        file_name_without_extension = os.path.splitext(os.path.basename(model))[0]
+
+        # NOTE:
+        # As of 2023/04/20, the current latest version of onnxruntime-web is 1.14.0, and does not support INT8 weights for Conv layers.
+        # For this reason, we choose model weight types to ensure compatibility with onnxruntime-web.
+        #
+        # As per docs, signed weight type (QInt8) is faster on most CPUs, so, we use that unless the model contains a Conv layer.
+        # For more information, see:
+        #  - https://github.com/microsoft/onnxruntime/issues/3130#issuecomment-1105200621
+        #  - https://github.com/microsoft/onnxruntime/issues/2339
+
+        loaded_model = onnx.load_model(model)
+        op_types = get_operators(loaded_model)
+        weight_type = QuantType.QUInt8 if "Conv" in op_types else QuantType.QInt8
+
+        quantize_dynamic(
+            model_input=model,
+            model_output=os.path.join(directory_path, f"{file_name_without_extension}_quantized.onnx"),
+            weight_type=weight_type,
+            # TODO allow user to specify these
+            # op_types_to_quantize=['MatMul', 'Add', 'Conv'],
+            extra_options=dict(EnableSubgraph=True),
+            **quantize_kwargs,
+        )
+
+        quantize_config["per_model_config"][file_name_without_extension] = dict(
+            op_types=list(op_types),
+            weight_type=str(weight_type),
+        )
+
+    # Save quantization config
+    with open(os.path.join(directory_path, "quantize_config.json"), "w") as fp:
+        json.dump(quantize_config, fp, indent=4)
+
+
+def main():
+    parser = HfArgumentParser((ConversionArguments,))
+    (conv_args,) = parser.parse_args_into_dataclasses()
+
+    model_id = conv_args.model_id
+    tokenizer_id = conv_args.tokenizer_id or model_id
+
+    output_model_folder = os.path.join(conv_args.output_parent_dir, model_id)
+
+    # Create output folder
+    os.makedirs(output_model_folder, exist_ok=True)
+
+    from_pretrained_kwargs = dict(
+        trust_remote_code=conv_args.trust_remote_code,
+    )
+
+    # Saving the model config
+    config = AutoConfig.from_pretrained(model_id, **from_pretrained_kwargs)
+
+    custom_kwargs = {}
+    if conv_args.custom_onnx_configs is not None:
+        if conv_args.task == "auto":
+            raise Exception("`--task` must be set when exporting with `--custom_onnx_configs`")
+        custom_onnx_configs = json.loads(conv_args.custom_onnx_configs)
+
+        for key in custom_onnx_configs:
+            onnx_configs = TasksManager._SUPPORTED_MODEL_TYPE[custom_onnx_configs[key]]["onnx"]
+            mapping = onnx_configs[conv_args.task]
+            custom_onnx_configs[key] = mapping.func(config, **mapping.keywords)
+
+        custom_kwargs["custom_onnx_configs"] = custom_onnx_configs
+
+    tokenizer = None
+    try:
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, **from_pretrained_kwargs)
+
+        # To avoid inserting all chat templates into tokenizers.js, we save the chat template
+        # to the tokenizer_config.json file, and load it when the tokenizer is loaded.
+        if getattr(tokenizer, "chat_template", None) is None and getattr(tokenizer, "use_default_system_prompt", False):
+            # No chat template specified, and we use the default
+            setattr(tokenizer, "chat_template", tokenizer.default_chat_template)
+
+    except KeyError:
+        pass  # No Tokenizer
+
+    except Exception as e:
+        if config.model_type not in MODELS_WITHOUT_TOKENIZERS:
+            raise e
+
+    core_export_kwargs = dict(
+        opset=conv_args.opset,
+        device=conv_args.device,
+        trust_remote_code=conv_args.trust_remote_code,
+        **custom_kwargs,
+    )
+
+    export_kwargs = dict(
+        model_name_or_path=model_id,
+        output=output_model_folder,
+        task=conv_args.task,
+        do_validation=not conv_args.skip_validation,
+        library_name="transformers",
+        **core_export_kwargs,
+    )
+
+    # Handle special cases
+    if config.model_type == "marian":
+        from .extra.marian import generate_tokenizer_json
+
+        tokenizer_json = generate_tokenizer_json(model_id, tokenizer)
+
+        with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+            json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type == "esm":
+        from .extra.esm import generate_fast_tokenizer
+
+        fast_tokenizer = generate_fast_tokenizer(tokenizer)
+        fast_tokenizer.save(os.path.join(output_model_folder, "tokenizer.json"))
+
+    elif config.model_type == "whisper":
+        if conv_args.output_attentions:
+            from .extra.whisper import get_main_export_kwargs
+
+            export_kwargs.update(**get_main_export_kwargs(config, "automatic-speech-recognition"))
+
+    elif config.model_type in ("wav2vec2", "wav2vec2-bert", "hubert"):
+        if tokenizer is not None:
+            from .extra.wav2vec2 import generate_tokenizer_json
+
+            tokenizer_json = generate_tokenizer_json(tokenizer)
+
+            with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+                json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type == "vits":
+        if tokenizer is not None:
+            from .extra.vits import generate_tokenizer_json
+
+            tokenizer_json = generate_tokenizer_json(tokenizer)
+
+            with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+                json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type == "speecht5":
+        # TODO allow user to specify vocoder path
+        export_kwargs["model_kwargs"] = {"vocoder": "microsoft/speecht5_hifigan"}
+
+        if tokenizer is not None:
+            from .extra.speecht5 import generate_tokenizer_json
+
+            tokenizer_json = generate_tokenizer_json(tokenizer)
+
+            with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+                json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type in ("owlvit", "owlv2"):
+        # Override default batch size to 1, needed because non-maximum suppression is performed for exporting.
+        # For more information, see https://github.com/huggingface/optimum/blob/e3b7efb1257c011db907ef40ab340e795cc5684c/optimum/exporters/onnx/model_configs.py#L1028-L1032
+        export_kwargs["batch_size"] = 1
+
+    else:
+        pass  # TODO
+
+    # Step 1. convert huggingface model to onnx
+    if not conv_args.split_modalities:
+        main_export(**export_kwargs)
+    else:
+        custom_export_kwargs = dict(
+            output_dir=output_model_folder,
+            **core_export_kwargs,
+        )
+
+        if config.model_type == "clip":
+            # Handle special case for exporting text and vision models separately
+            from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection
+
+            from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig
+
+            text_model = CLIPTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+            vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+
+            export_models(
+                models_and_onnx_configs={
+                    "text_model": (text_model, CLIPTextModelWithProjectionOnnxConfig(text_model.config)),
+                    "vision_model": (vision_model, CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)),
+                },
+                **custom_export_kwargs,
+            )
+
+        elif config.model_type == "siglip":
+            # Handle special case for exporting text and vision models separately
+            from transformers.models.siglip import SiglipTextModel, SiglipVisionModel
+
+            from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig
+
+            text_model = SiglipTextModel.from_pretrained(model_id, **from_pretrained_kwargs)
+            vision_model = SiglipVisionModel.from_pretrained(model_id, **from_pretrained_kwargs)
+
+            export_models(
+                models_and_onnx_configs={
+                    "text_model": (text_model, SiglipTextModelOnnxConfig(text_model.config)),
+                    "vision_model": (vision_model, SiglipVisionModelOnnxConfig(vision_model.config)),
+                },
+                **custom_export_kwargs,
+            )
+
+        # TODO: Enable once https://github.com/huggingface/optimum/pull/1552 is merged
+        # elif config.model_type == 'clap':
+        #     # Handle special case for exporting text and audio models separately
+        #     from .extra.clap import ClapTextModelWithProjectionOnnxConfig, ClapAudioModelWithProjectionOnnxConfig
+        #     from transformers.models.clap import ClapTextModelWithProjection, ClapAudioModelWithProjection
+
+        #     text_model = ClapTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+        #     audio_model = ClapAudioModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+
+        #     export_models(
+        #         models_and_onnx_configs={
+        #             "text_model": (text_model, ClapTextModelWithProjectionOnnxConfig(text_model.config)),
+        #             "audio_model": (audio_model, ClapAudioModelWithProjectionOnnxConfig(audio_model.config)),
+        #         },
+        #         **custom_export_kwargs,
+        #     )
+
+        else:
+            raise Exception(f"Unable to export {config.model_type} model with `--split_modalities`.")
+
+    # Step 2. (optional, recommended) quantize the converted model for fast inference and to reduce model size.
+    if conv_args.quantize:
+        # Update quantize config with model specific defaults
+        quantize_config = MODEL_SPECIFIC_QUANTIZE_PARAMS.get(config.model_type, DEFAULT_QUANTIZE_PARAMS)
+
+        # Update if user specified values
+        if conv_args.per_channel is not None:
+            quantize_config["per_channel"] = conv_args.per_channel
+
+        if conv_args.reduce_range is not None:
+            quantize_config["reduce_range"] = conv_args.reduce_range
+
+        quantize(
+            [
+                os.path.join(output_model_folder, x)
+                for x in os.listdir(output_model_folder)
+                if x.endswith(".onnx") and not x.endswith("_quantized.onnx")
+            ],
+            **quantize_config,
+        )
+
+    # Step 3. Move .onnx files to the 'onnx' subfolder
+    os.makedirs(os.path.join(output_model_folder, "onnx"), exist_ok=True)
+    for file in os.listdir(output_model_folder):
+        if file.endswith((".onnx", ".onnx_data")):
+            shutil.move(os.path.join(output_model_folder, file), os.path.join(output_model_folder, "onnx", file))
+
+    # Step 4. Update the generation config if necessary
+    if config.model_type == "whisper":
+        from transformers import GenerationConfig
+
+        from .extra.whisper import get_alignment_heads
+
+        generation_config = GenerationConfig.from_pretrained(model_id, **from_pretrained_kwargs)
+        generation_config.alignment_heads = get_alignment_heads(config)
+        generation_config.save_pretrained(output_model_folder)
+
+
+if __name__ == "__main__":
+    main()

experiments/try_quantized_model.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+from typing import List
+
+import click
+import numpy as np
+import torch
+import torch.nn.functional as F
+from optimum.onnxruntime import ORTModelForFeatureExtraction
+from optimum.pipelines import pipeline
+from torch import Tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+
+def hf_embed(model_id: str, texts: List[str], tokenizer):
+    # Tokenize the input texts
+    model = AutoModel.from_pretrained(model_id)
+    model.eval()
+    encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
+
+    model_output = model(**encoded_input)
+    sentence_embeddings = model_output[0][:, 0]
+    sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+    return sentence_embeddings
+
+
+@click.command()
+@click.option("--model_id", help="model id from huggingface.co/models")
+@click.option("--model_dir", help="The person to greet.")
+def setup(model_id, model_dir):
+    text = "This is a test sentence"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    output_dir = Path(model_dir)
+    model = ORTModelForFeatureExtraction.from_pretrained(output_dir)
+    onnx_quant_embed = pipeline(
+        "feature-extraction", model=model, accelerator="ort", tokenizer=tokenizer, return_tensors=True
+    )
+    quant_embeddings = onnx_quant_embed([text])
+    quant_embeddings = F.normalize(quant_embeddings[0][:,0], p=2, dim=1)
+    quant_embeddings = quant_embeddings.detach().numpy()
+    print(quant_embeddings.shape)
+
+    torch_embeddings = hf_embed(model_id, texts=[text], tokenizer=tokenizer)
+    torch_embeddings = F.normalize(torch_embeddings, p=2, dim=1)
+    torch_embeddings = torch_embeddings.detach().numpy()
+    print(torch_embeddings.shape)
+    assert quant_embeddings.shape == torch_embeddings.shape
+    print(np.allclose(quant_embeddings, torch_embeddings, atol=1e-5))
+
+
+if __name__ == "__main__":
+    setup()

poetry.lock

@@ -3430,4 +3430,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.0,<3.13"
-content-hash = "04730c4b2a45b7a1ffc5faf449a2faedbed38097aae9cfd0dfa6ce533197db80"
+content-hash = "8bc18a20773e3e2530564b9e0ba8a31c6f3dd7ab710c87acfe446db5e785ca3c"

pyproject.toml

@@ -34,6 +34,7 @@ pillow = "^10.2.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
+click = "^8.1.7"
 
 
 [build-system]