mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Add changes for 0.27 Windows release
This commit is contained in:
Keval Morabia
@@ -2,6 +2,15 @@
|
||||
Model Optimizer Changelog (Windows)
|
||||
===================================
|
||||
|
||||
0.27 (2025-04-30)
|
||||
^^^^^^^^^^^^^^^^^
|
||||
|
||||
**New Features**
|
||||
|
||||
- New LLM models like DeepSeek etc. are supported with ONNX INT4 AWQ quantization on Windows. Refer `Windows Support Matrix <https://nvidia.github.io/TensorRT-Model-Optimizer/guides/0_support_matrix.html>`_ for details about supported features and models.
|
||||
- TensorRT Model Optimizer for Windows now supports ONNX INT8 and FP8 quantization (W8A8) of SAM2 and Whisper models. Check `example scripts <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/windows/onnx_ptq>`_ for getting started with quantizing these models.
|
||||
|
||||
|
||||
0.19 (2024-11-18)
|
||||
^^^^^^^^^^^^^^^^^
|
||||
|
||||
|
||||
@@ -5,7 +5,9 @@ DirectML
|
||||
===================
|
||||
|
||||
|
||||
Once an ONNX FP16 model is quantized using TensorRT Model Optimizer on Windows, the resulting quantized ONNX model can be deployed on the DirectML backend via the `ONNX Runtime GenAI <https://onnxruntime.ai/docs/genai/>`_ or `ONNX Runtime <https://onnxruntime.ai/>`_.
|
||||
Once an ONNX FP16 model is quantized using TensorRT Model Optimizer on Windows, the resulting quantized ONNX model can be deployed on the DirectML (DML) backend via the `ONNX Runtime GenAI <https://onnxruntime.ai/docs/genai/>`_ or `ONNX Runtime <https://onnxruntime.ai/>`_.
|
||||
|
||||
.. note:: Currently, DirectML backend doesn't support 8-bit precision. So, 8-bit quantized models should be deployed on other backends like ORT-CUDA etc. However, DML path does support INT4 quantized models.
|
||||
|
||||
ONNX Runtime GenAI
|
||||
==================
|
||||
|
||||
@@ -11,7 +11,7 @@ The ONNX quantization API in ModelOpt-Windows offers advanced Post-Training Quan
|
||||
ONNX Model Quantization (PTQ)
|
||||
------------------------------
|
||||
|
||||
The ONNX quantization API requires a model, calibration data, along with quantization settings like algorithm, calibration-EPs etc. Here’s an example implementing int4 AWQ:
|
||||
The ONNX quantization API requires a model, calibration data, along with quantization settings like algorithm, calibration-EPs etc. Here’s an example snippet to apply INT4 AWQ quantization:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -32,22 +32,24 @@ The ONNX quantization API requires a model, calibration data, along with quantiz
|
||||
size_threshold=0,
|
||||
)
|
||||
|
||||
Check :meth:`modelopt.onnx.quantization.quantize_int4 <modelopt.onnx.quantization.int4.quantize>` for details about quantization API.
|
||||
Check :meth:`modelopt.onnx.quantization.quantize_int4 <modelopt.onnx.quantization.int4.quantize>` for details about INT4 quantization API.
|
||||
|
||||
Refer :ref:`Support_Matrix` for details about supported features and models.
|
||||
|
||||
To know more about ONNX PTQ, refer :ref:`ONNX_PTQ_Guide_Windows` and `example script <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/windows/onnx_ptq/>`_.
|
||||
To know more about ONNX PTQ, refer :ref:`ONNX_PTQ_Guide_Windows` and `example scripts <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/windows/onnx_ptq/>`_.
|
||||
|
||||
|
||||
Deployment
|
||||
----------
|
||||
The quantized ONNX model is deployment-ready, equivalent to a standard ONNX model. ModelOpt-Windows uses ONNX’s `DequantizeLinear <https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html>`_ (DQ) nodes, which support INT4 data-type from opset version 21 onward. Ensure the model’s opset version is 21 or higher. Refer :ref:`Apply_ONNX_PTQ` for details.
|
||||
The quantized onnx model can be deployed using frameworks like onnxruntime. Ensure that model's opset is 19+ for FP8 quantization, and it is 21+ for INT4 quantization. This is needed due to different opset requirements of ONNX's `Q <https://onnx.ai/onnx/operators/onnx__QuantizeLinear.html>`_/`DQ <https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html>`_ nodes for INT4, FP8 data-types support. Refer :ref:`Apply_ONNX_PTQ` for details.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# write steps (say, upgrade_opset_to_21() method) to upgrade opset to 21, if it is lower than 21.
|
||||
# write steps (say, upgrade_opset() method) to upgrade or patch opset of the model, if needed
|
||||
# the opset-upgrade, if needed, can be done on either base ONNX model or on the quantized model
|
||||
# finally, save the quantized model
|
||||
|
||||
quantized_onnx_model = upgrade_opset_to_21(quantized_onnx_model)
|
||||
quantized_onnx_model = upgrade_opset(quantized_onnx_model)
|
||||
onnx.save_model(
|
||||
quantized_onnx_model,
|
||||
output_path,
|
||||
@@ -56,7 +58,7 @@ The quantized ONNX model is deployment-ready, equivalent to a standard ONNX mode
|
||||
size_threshold=0,
|
||||
)
|
||||
|
||||
Deploy the quantized model using the DirectML backend. For detailed deployment instructions, see the :ref:`DirectML_Deployment`.
|
||||
For detailed instructions about deployment of quantized models with DirectML backend (ORT-DML), see the :ref:`DirectML_Deployment`. Also, refer `example scripts <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/windows/onnx_ptq/>`_ for any possible model-specific inference guidance or script (if any).
|
||||
|
||||
.. note::
|
||||
|
||||
|
||||
@@ -21,7 +21,12 @@ The following system requirements are necessary to install and use TensorRT Mode
|
||||
+-------------------------+-----------------------------+
|
||||
| Nvidia Driver | 565.90 or newer |
|
||||
+-------------------------+-----------------------------+
|
||||
| Nvidia GPU | RTX 40 and 50 series |
|
||||
+-------------------------+-----------------------------+
|
||||
|
||||
.. note::
|
||||
- Make sure to use GPU-compatible driver and other dependencies (e.g. torch etc.). For instance, support for Blackwell GPU might be present in Nvidia 570+ driver, and CUDA-12.8.
|
||||
- We currently support *Single-GPU* configuration.
|
||||
|
||||
The TensorRT Model Optimizer - Windows can be used in following ways:
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
Install ModelOpt-Windows as a Standalone Toolkit
|
||||
================================================
|
||||
|
||||
The TensorRT Model Optimizer - Windows (ModelOpt-Windows) can be installed as a standalone toolkit for quantizing Large Language Models (LLMs). Below are the setup steps:
|
||||
The TensorRT Model Optimizer - Windows (ModelOpt-Windows) can be installed as a standalone toolkit for quantizing ONNX models. Below are the setup steps:
|
||||
|
||||
**1. Setup Prerequisites**
|
||||
|
||||
@@ -40,7 +40,7 @@ This command installs ModelOpt-Windows and its ONNX module, along with the *onnx
|
||||
|
||||
**4. Setup ONNX Runtime (ORT) for Calibration**
|
||||
|
||||
The ONNX Post-Training Quantization (PTQ) process involves running the base model with user-supplied inputs, a process called calibration. The user-supplied model inputs are referred to as calibration data. To perform calibration, the base model must be run using a suitable ONNX Execution Provider (EP), such as *DmlExecutionProvider* (DirectML EP) or *CudaExecutionProvider* (CUDA EP). There are different ONNX Runtime packages for each EP:
|
||||
The ONNX Post-Training Quantization (PTQ) process involves running the base model with user-supplied inputs, a process called calibration. The user-supplied model inputs are referred to as calibration data. To perform calibration, the base model must be run using a suitable ONNX Execution Provider (EP), such as *DmlExecutionProvider* (DirectML EP) or *CUDAExecutionProvider* (CUDA EP). There are different ONNX Runtime packages for each EP:
|
||||
|
||||
- *onnxruntime-directml* provides the DirectML EP.
|
||||
- *onnxruntime-gpu* provides the CUDA EP.
|
||||
@@ -68,7 +68,7 @@ If you prefer to use the CUDA EP for calibration, uninstall the existing *onnxru
|
||||
|
||||
**5. Setup GPU Acceleration Tool for Quantization**
|
||||
|
||||
ModelOpt-Windows utilizes the `cupy-cuda12x <https://cupy.dev//>`_ tool for GPU acceleration during the INT4 ONNX quantization process if you have CUDA 12.x.
|
||||
By default, ModelOpt-Windows utilizes the `cupy-cuda12x <https://cupy.dev//>`_ tool for GPU acceleration during the INT4 ONNX quantization process. This is compatible with CUDA 12.x.
|
||||
|
||||
**6. Verify Installation**
|
||||
|
||||
@@ -79,6 +79,10 @@ Ensure the following steps are verified:
|
||||
- *onnxruntime-directml* (DirectML EP)
|
||||
- *onnxruntime-gpu* (CUDA EP)
|
||||
- *onnxruntime* (CPU EP)
|
||||
- **Onnx and Onnxruntime Import**: Ensure that following python command runs successfully.
|
||||
.. code-block:: python
|
||||
|
||||
python -c "import onnx; import onnxruntime"
|
||||
- **Environment Variables**: For workflows using CUDA dependencies (e.g., CUDA EP-based calibration), ensure environment variables like *CUDA_PATH*, *CUDA_V12_4*, or *CUDA_V11_8* etc. are set correctly. Reopen the command-prompt if any environment variable is updated or newly created.
|
||||
- **ModelOpt-Windows Import Check**: Run the following command to ensure the installation is successful:
|
||||
|
||||
|
||||
@@ -29,7 +29,7 @@ Feature Support Matrix
|
||||
- PyTorch, ONNX*
|
||||
- TensorRT*, TensorRT-LLM
|
||||
* - INT8
|
||||
- * Per-channel INT8 Weights, Per-Tensor FP8 Activations
|
||||
- * Per-channel INT8 Weights, Per-Tensor INT8 Activations
|
||||
* Uses Smooth Quant Algorithm
|
||||
* GPUs: Ampere and Later
|
||||
- PyTorch, ONNX*
|
||||
@@ -71,16 +71,18 @@ Feature Support Matrix
|
||||
- PyTorch*
|
||||
- TensorRT-LLM*
|
||||
* - FP8
|
||||
- * Per-Tensor FP8 Weight & Activations
|
||||
- * Per-Tensor FP8 Weight & Activations (PyTorch)
|
||||
* Per-Tensor Activation and Per-Channel Weights quantization (ONNX)
|
||||
* Uses Max calibration
|
||||
* GPUs: Ada and Later
|
||||
- PyTorch*, ONNX*
|
||||
- TensorRT*, TensorRT-LLM*
|
||||
- PyTorch*, ONNX
|
||||
- TensorRT*, TensorRT-LLM*, ORT-CUDA
|
||||
* - INT8
|
||||
- * Per-channel INT8 Weights, Per-Tensor FP8 Activations
|
||||
* Uses Smooth Quant Algorithm
|
||||
- * Per-Channel INT8 Weights, Per-Tensor INT8 Activations
|
||||
* Uses Smooth Quant (PyTorch)*, Max calibration (ONNX)
|
||||
* GPUs: Ada and Later
|
||||
- PyTorch*, ONNX*
|
||||
- TensorRT*, TensorRT-LLM*
|
||||
- PyTorch*, ONNX
|
||||
- TensorRT*, TensorRT-LLM*, ORT-CUDA
|
||||
|
||||
.. note:: Features marked with an asterisk (*) are considered experimental.
|
||||
|
||||
@@ -98,16 +100,83 @@ Model Support Matrix
|
||||
:header-rows: 1
|
||||
|
||||
* - Model
|
||||
- ONNX INT4 AWQ
|
||||
- ONNX INT4 AWQ (W4A16)
|
||||
- ONNX INT8 Max (W8A8)
|
||||
- ONNX FP8 Max (W8A8)
|
||||
* - Llama3.1-8B-Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Phi3.5-mini-Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Mistral-7B-Instruct-v0.3
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Llama3.2-3B-Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Gemma-2b-it
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Gemma-2-2b
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Gemma-2-9b
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Nemotron Mini 4B Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Qwen2.5-7B-Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - DeepSeek-R1-Distill-Llama-8B
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - DeepSeek-R1-Distil-Qwen-1.5B
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - DeepSeek-R1-Distil-Qwen-7B
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - DeepSeek-R1-Distill-Qwen-14B
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Mistral-NeMo-Minitron-2B-128k-Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Mistral-NeMo-Minitron-4B-128k-Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - Mistral-NeMo-Minitron-8B-128k-Instruct
|
||||
- Yes
|
||||
- No
|
||||
- No
|
||||
* - whisper-large
|
||||
- No
|
||||
- Yes
|
||||
- Yes
|
||||
* - sam2-hiera-large
|
||||
- No
|
||||
- Yes
|
||||
- Yes
|
||||
|
||||
.. note::
|
||||
- ``ONNX INT8 Max`` means INT8 (W8A8) quantization of ONNX model using Max calibration. Similar holds true for the term ``ONNX FP8 Max``.
|
||||
- The LLMs in above table are `GenAI <https://github.com/microsoft/onnxruntime-genai/>`_ built LLMs unless specified otherwise.
|
||||
- Check `examples <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/windows/onnx_ptq/>`_ for specific instructions and scripts.
|
||||
|
||||
@@ -41,8 +41,9 @@ ModelOpt-Windows's ONNX PTQ API requires a base ONNX model, which can be obtaine
|
||||
|
||||
Each tool offers unique features and options for conversion from PyTorch or other frameworks to the ONNX format.
|
||||
|
||||
**Base Model Precision**: ModelOpt-Windows supports base models in both FP16 and FP32 formats. Choosing FP16 over FP32 can help reduce memory usage and improve speed, especially on hardware optimized for lower precision, such as NVIDIA GPUs with Tensor Cores. However, FP16's smaller dynamic range may require careful tuning.
|
||||
**Opset requirements of different data-types**: ModelOpt-Windows uses ONNX's `Q <https://onnx.ai/onnx/operators/onnx__QuantizeLinear.html>`_/`DQ <https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html>`_ nodes for applying quantization and dequantization operations in the ONNX model. The INT4 datatype support in Q/DQ node came in opset-21 and, FP8 datatype support in Q/DQ node came in opset-19. So, ensure that model's opset is 19+ for FP8 quantization, and it is 21+ for INT4 quantization. This is needed for deployment of the quantized model on onnxruntime framework (e.g. ORT-DirectML backend). Different ONNX exporter tools usually have option or argument for target 'opset' field. See their documentation for details about its usage and max-supported opset limit.
|
||||
|
||||
**Base Model Precision**: ModelOpt-Windows supports base models in both FP16 and FP32 formats. Choosing FP16 over FP32 can help reduce memory usage and improve speed, especially on hardware optimized for lower precision, such as NVIDIA GPUs with Tensor Cores. However, FP16's smaller dynamic range may require careful tuning.
|
||||
|
||||
**ONNX FP16 Conversion Tools**: Some popular FP32 to FP16 ONNX conversion tools include:
|
||||
|
||||
@@ -50,179 +51,29 @@ Each tool offers unique features and options for conversion from PyTorch or othe
|
||||
- Hugging Face *Optimum* tool with a *dtype* argument for FP16 generation - Refer to optimum's `CLI <https://huggingface.co/docs/optimum/en/exporters/onnx/usage_guides/export_a_model/>`_, and `API <https://github.com/huggingface/optimum/blob/main/optimum/exporters/onnx/convert.py/>`_ usage.
|
||||
- Microsoft Olive, which supports FP16 via configuration files. Refer *float16* option in this example `config <https://github.com/microsoft/Olive/blob/main/examples/directml/llm/config_llm.json/>`_.
|
||||
|
||||
|
||||
Once base model is obtained, ModelOpt-Windows's PTQ can be applied to get the quantized mode. The resulting quantized model can be deployed on DirectML and TensorRT* backend.
|
||||
Once base model is obtained, ModelOpt-Windows's PTQ can be applied to get the quantized mode. The resulting quantized model can be deployed on backends like DirectML, CUDA and TensorRT*.
|
||||
|
||||
.. _Apply_ONNX_PTQ:
|
||||
|
||||
Apply Post Training Quantization (PTQ)
|
||||
--------------------------------------
|
||||
|
||||
Applying PTQ on a model involves preparing calibration-data (if needed), invoking quantization API, saving the quantized model, and any additional post-processing like opset upgrade as needed.
|
||||
|
||||
**Prepare calibration data**
|
||||
|
||||
The SmoothQuant (SQ) and Activation-Aware-Quantization (AWQ) algorithms require calibration data during quantization. If the *quantize* API's calibration-data argument is not provided (i.e., set to *None*), ModelOpt-Windows will internally use randomly generated model inputs for calibration. Refer to the sample code below for preparing calibration inputs.
|
||||
Quantization algorithms like SmoothQuant (SQ), Activation-Aware-Quantization (AWQ), and static quantization of activations often require calibration data. If the *quantize* API's calibration-data argument is not provided (i.e., set to *None*), ModelOpt-Windows will internally use randomly generated model inputs for calibration.
|
||||
|
||||
Preparing calibration data for ModelOpt-Windows involves two steps:
|
||||
As an example, preparing calibration data for INT4 AWQ quantization of LLMs may involve following major steps:
|
||||
|
||||
1. **Generate Token Encodings**: Use a dataset like *cnn-dailymail* or *pile* with the model's tokenizer to generate token encodings and related data from the representative dataset
|
||||
2. **Format for Model Input**: Convert encodings into model-compatible formats.
|
||||
|
||||
See the code example below for details.
|
||||
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Refer get_calib_inputs() method below to prepare calibration inputs for your model.
|
||||
|
||||
# Note that names and shapes of inputs and outputs can vary from model to model, and also between ONNX exporter tools.
|
||||
# So, use following code as reference for preparing calibration data for your model.
|
||||
|
||||
def make_model_input(
|
||||
config,
|
||||
input_ids_arg,
|
||||
attention_mask_arg,
|
||||
add_past_kv_inputs,
|
||||
device,
|
||||
use_fp16,
|
||||
use_buffer_share,
|
||||
add_position_ids,
|
||||
):
|
||||
input_ids = input_ids_arg
|
||||
attention_mask = attention_mask_arg
|
||||
|
||||
if isinstance(input_ids_arg, list):
|
||||
input_ids = torch.tensor(input_ids_arg, device=device, dtype=torch.int64)
|
||||
attention_mask = torch.tensor(attention_mask_arg, device=device, dtype=torch.int64)
|
||||
|
||||
inputs = {
|
||||
"input_ids": input_ids.contiguous(),
|
||||
"attention_mask": attention_mask.contiguous(),
|
||||
}
|
||||
|
||||
if add_position_ids:
|
||||
position_ids = attention_mask.long().cumsum(-1) - 1
|
||||
position_ids.masked_fill_(attention_mask == 0, 1)
|
||||
inputs["position_ids"] = position_ids.contiguous()
|
||||
|
||||
if add_past_kv_inputs:
|
||||
torch_dtype = torch.float16 if use_fp16 else torch.float32
|
||||
batch_size, sequence_length = input_ids.shape
|
||||
max_sequence_length = config.max_position_embeddings
|
||||
num_heads, head_size = (
|
||||
config.num_key_value_heads,
|
||||
config.hidden_size // config.num_attention_heads,
|
||||
)
|
||||
|
||||
if hasattr(config, "head_dim"):
|
||||
head_size = config.head_dim
|
||||
|
||||
for i in range(config.num_hidden_layers):
|
||||
past_key = torch.zeros(
|
||||
batch_size,
|
||||
num_heads,
|
||||
max_sequence_length if use_buffer_share else 0,
|
||||
head_size,
|
||||
device=device,
|
||||
dtype=torch_dtype,
|
||||
)
|
||||
past_value = torch.zeros(
|
||||
batch_size,
|
||||
num_heads,
|
||||
max_sequence_length if use_buffer_share else 0,
|
||||
head_size,
|
||||
device=device,
|
||||
dtype=torch_dtype,
|
||||
)
|
||||
inputs.update(
|
||||
{
|
||||
f"past_key_values.{i}.key": past_key.contiguous(),
|
||||
f"past_key_values.{i}.value": past_value.contiguous(),
|
||||
}
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
|
||||
def get_calib_inputs(
|
||||
dataset_name,
|
||||
model_name,
|
||||
cache_dir,
|
||||
calib_size,
|
||||
batch_size,
|
||||
block_size,
|
||||
device,
|
||||
use_fp16,
|
||||
use_buffer_share,
|
||||
add_past_kv_inputs,
|
||||
max_calib_rows_to_load,
|
||||
add_position_ids,
|
||||
trust_remote_code,
|
||||
):
|
||||
# from transformers import LlamaConfig
|
||||
# config = LlamaConfig.from_pretrained(
|
||||
# model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=trust_remote_code
|
||||
# )
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=trust_remote_code
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=trust_remote_code
|
||||
)
|
||||
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
assert (
|
||||
calib_size <= max_calib_rows_to_load
|
||||
), "calib size should be no more than max_calib_rows_to_load"
|
||||
|
||||
dataset2 = load_dataset("cnn_dailymail", name="3.0.0", split="train").select(range(max_calib_rows_to_load))
|
||||
column = "article"
|
||||
|
||||
# dataset2 = dataset2.shuffle(seed=42)
|
||||
dataset2 = dataset2[column][:calib_size]
|
||||
batch_encoded = tokenizer.batch_encode_plus(
|
||||
dataset2, return_tensors="pt", padding=True, truncation=True, max_length=block_size
|
||||
) # return_tensors="pt",
|
||||
batch_encoded = batch_encoded.to(device)
|
||||
batch_encoded_input_ids = batch_encoded["input_ids"]
|
||||
batch_encoded_attention_mask = batch_encoded["attention_mask"]
|
||||
calib_dataloader_input_ids = DataLoader(batch_encoded_input_ids, batch_size=batch_size, shuffle=False)
|
||||
calib_dataloader_attenton_mask = DataLoader(batch_encoded_attention_mask, batch_size=batch_size, shuffle=False)
|
||||
|
||||
number_of_batched_samples = calib_size // batch_size
|
||||
|
||||
batched_input_ids = []
|
||||
for idx, data in enumerate(calib_dataloader_input_ids):
|
||||
batched_input_ids.append(data)
|
||||
if idx == (number_of_batched_samples - 1):
|
||||
break
|
||||
|
||||
batched_attention_mask = []
|
||||
for idx, data in enumerate(calib_dataloader_attenton_mask):
|
||||
batched_attention_mask.append(data)
|
||||
if idx == (number_of_batched_samples - 1):
|
||||
break
|
||||
|
||||
batched_inputs_list = []
|
||||
for i in range(number_of_batched_samples):
|
||||
input_ids = batched_input_ids[i]
|
||||
attention_mask = batched_attention_mask[i]
|
||||
|
||||
inputs = make_model_input(config, input_ids, attention_mask, add_past_kv_inputs, device,
|
||||
use_fp16,
|
||||
use_buffer_share,
|
||||
add_position_ids,
|
||||
)
|
||||
inputs = {
|
||||
input_name: torch_tensor.cpu().numpy() for input_name, torch_tensor in inputs.items()
|
||||
}
|
||||
batched_inputs_list.append(inputs)
|
||||
|
||||
return batched_inputs_list
|
||||
Please refer the `example scripts <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/windows/onnx_ptq/>`_ for details about preparing calibration-data of various supported ONNX models.
|
||||
|
||||
**Call Quantization API**
|
||||
|
||||
The example below demonstrates how to quantize an ONNX model using ModelOpt-Windows with INT4 precision.
|
||||
The example below demonstrates how to apply INT4 AWQ quantization on a LLM ONNX model.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -238,9 +89,13 @@ The example below demonstrates how to quantize an ONNX model using ModelOpt-Wind
|
||||
|
||||
Check :meth:`modelopt.onnx.quantization.quantize_int4 <modelopt.onnx.quantization.int4.quantize>` for details about quantization API.
|
||||
|
||||
**Upgrade Opset to 21+**
|
||||
**Upgrade opset of the model**
|
||||
|
||||
ModelOpt-Windows uses ONNX’s `DequantizeLinear <https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html>`_ (DQ) nodes, which support INT4 data-type from opset version 21 onward. Ensure the model’s opset version is 21 or higher, for deployment on DirectML backend.
|
||||
Opset requirement for different data-types is already explained in the section describing ways to obtain base model. To summarize, the opset requirements stems from the fact that support for different types in ONNX `Q <https://onnx.ai/onnx/operators/onnx__QuantizeLinear.html>`_/`DQ <https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html>`_ nodes is gradually added in different opsets. For instance, the INT4 data-type support in Q/DQ node came in opset-21 and, FP8 data-type support in Q/DQ node came in opset-19. So, ensure that model's opset is 19+ for FP8 quantization, and it is 21+ for INT4 quantization. This is needed for deployment of the quantized model on onnxruntime framework (e.g. ORT-DirectML backend).
|
||||
|
||||
Generally, different ONNX exporter tools have option or argument for providing the desired or target opset. It is possible that the desired opset is greater than the max-opset user's ONNX exporter tool supports. In that case, user would need to manually 'patch' the opset of the ONNX model. This would require updating the ONNX metadata that stores opset field of the graph, and it may additionally require updating some nodes in the graph as per new opset (if they have changed in later opsets). Alternatively, one can try using the ONNX exporter tool which already supports the desired opset (if exists any).
|
||||
|
||||
A few sample code snippets to inspect the opset of the given ONNX model, and to update the opset field in the ONNX model's meta-data, are provided below.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -266,7 +121,7 @@ Use the above steps to inspect the ONNX model's opset version.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Example steps for opset-21 upgrade of default (onnx) domain
|
||||
# Example steps for updating opset metadata of default (onnx) domain
|
||||
# Update opsets for other domains as needed for your requirement (or exclude them as suitable).
|
||||
|
||||
model = onnx.load(onnx_path)
|
||||
@@ -281,6 +136,7 @@ Use the above steps to inspect the ONNX model's opset version.
|
||||
|
||||
updated_quantized_onnx_model = onnx.helper.make_model(model.graph, opset_imports=new_opset_imports)
|
||||
|
||||
The ONNX models produced using `GenAI <https://github.com/microsoft/onnxruntime-genai/>`_ are generally seen to work fine with above opset upgrade patch. ONNX models produced using other ONNX exporter tool might require further post-processing on case-by-case basis for nodes that have changed in later opsets.
|
||||
|
||||
**Save Quantized Model**
|
||||
|
||||
@@ -299,4 +155,4 @@ To save a quantized ONNX model with external data, use the following code:
|
||||
Deploy Quantized ONNX Model
|
||||
---------------------------
|
||||
|
||||
Inference of the quantized models can be done using tools like `GenAI <https://github.com/microsoft/onnxruntime-genai/>`_, `OnnxRunTime (ORT) <https://onnxruntime.ai//>`_. These APIs can do inference on backends like DML. For details about DirectML deployment, see :ref:`DirectML_Deployment`.
|
||||
Inference of the quantized models can be done using tools like `GenAI <https://github.com/microsoft/onnxruntime-genai/>`_, `OnnxRunTime (ORT) <https://onnxruntime.ai//>`_. These APIs can do inference on backends like DML. For details about DirectML deployment of quantized models, see :ref:`DirectML_Deployment`. Also, refer `example scripts <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/windows/onnx_ptq/>`_ for any possible model-specific inference guidance or script (if any).
|
||||
|
||||
@@ -10,8 +10,8 @@ ModelOpt-Windows
|
||||
|
||||
**ONNX PTQ**
|
||||
|
||||
1. Why is awq-scale search taking too long or stuck at 0% during ONNX quantization?
|
||||
-----------------------------------------------------------------------------------
|
||||
1. Why is awq-scale search taking too long or stuck at 0% during ONNX INT4 quantization?
|
||||
----------------------------------------------------------------------------------------
|
||||
|
||||
Awq-scale search should complete in minutes with NVIDIA GPU acceleration. If stalled:
|
||||
|
||||
@@ -40,12 +40,12 @@ This usually results from mismatched CUDA and cuDNN versions or missing paths. E
|
||||
|
||||
Make sure that the output directory is clean before each quantization run otherwise, existing quantized model file may get appended in each run leading to increase in model's size and possibly corrupting it.
|
||||
|
||||
5. Running quantized ONNX model on DirectML backend gives following error. What can be the issue?
|
||||
-------------------------------------------------------------------------------------------------
|
||||
5. Running INT4 quantized ONNX model on DirectML backend gives following error. What can be the issue?
|
||||
------------------------------------------------------------------------------------------------------
|
||||
|
||||
`Error Unrecognized attribute: block_size for operator DequantizeLinear`
|
||||
|
||||
ModelOpt-Windows uses ONNX's `DequantizeLinear <https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html>`_ (DQ) nodes. The *block_size* attribute was added in DeQuantizeLinear node in Opset-21. Make sure that quantized model's opset version is 21 or higher. Refer :ref:`Apply_ONNX_PTQ` for details.
|
||||
ModelOpt-Windows uses ONNX's `DequantizeLinear <https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html>`_ (DQ) nodes. The int4 data-type support in DeQuantizeLinear node came in opset-21. And, *block_size* attribute was added in DeQuantizeLinear node in Opset-21. Make sure that quantized model's opset version is 21 or higher. Refer :ref:`Apply_ONNX_PTQ` for details.
|
||||
|
||||
6. Running INT4 quantized ONNX model on DirectML backend gives following kind of error. What can be the issue?
|
||||
--------------------------------------------------------------------------------------------------------------
|
||||
@@ -54,6 +54,16 @@ ModelOpt-Windows uses ONNX's `DequantizeLinear <https://onnx.ai/onnx/operators/o
|
||||
|
||||
One possible reason for above error is that INT4 quantized ONNX model's opset version (default or onnx domain) is less than 21. Ensure the INT4 quantized model's opset version is 21 or higher since INT4 data-type support in DeQuantizeLinear ONNX node came in opset-21.
|
||||
|
||||
7. Running 8-bit quantized ONNX model with ORT-DML gives onnxruntime error about using 8-bit data-type (e.g. INT8/FP8). What can be the issue?
|
||||
-----------------------------------------------------------------------------------------------------------------------------------------------
|
||||
|
||||
Currently, DirectML backend (ORT-DML) doesn't support 8-bit precision. So, it expectedly complains about 8-bit data-type. Try using ORT-CUDA or other 8-bit compatible backend.
|
||||
|
||||
8. How to resolve onnxruntime error about invalid use of FP8 type in QuantizeLinear / DeQuantizeLinear node?
|
||||
-------------------------------------------------------------------------------------------------------------
|
||||
|
||||
The FP8 type support in QuantizeLinear / DeQuantizeLinear node came with Opset-19. So, ensure that opset of ONNX model is 19+.
|
||||
|
||||
.. _nas_faqs:
|
||||
|
||||
NAS/Pruning
|
||||
|
||||
@@ -28,9 +28,9 @@ For accuracy evaluation, the [Massive Multitask Language Understanding (MMLU)](h
|
||||
|
||||
The table below shows the MMLU 5-shot score for some models.
|
||||
|
||||
- **FP16 ONNX model**: Generated using [GenAI Model Builder](https://github.com/microsoft/onnxruntime-genai/blob/main/src/python/py/models/README.md)
|
||||
- **FP16 ONNX model**: Generated using [GenAI Model Builder](https://github.com/microsoft/onnxruntime-genai/blob/main/src/python/py/models/README.md) with DML EP
|
||||
- **INT4 AWQ model**: Generated by quantizing FP16 ONNX model using ModelOpt-Windows
|
||||
- **Configuration**: Windows OS, GPU RTX4090, nvidia-modelopt v0.19.0.
|
||||
- **Configuration**: Windows OS, GPU RTX4090, nvidia-modelopt v0.19.0, onnxruntime-genai-directml 0.4, transformers 4.44
|
||||
|
||||
| **Model** | **ONNX FP16** | **ONNX INT4** |
|
||||
|:------------------------------|:---------------|:--------------|
|
||||
|
||||
@@ -33,7 +33,7 @@
|
||||
## Overview
|
||||
|
||||
The **TensorRT Model Optimizer - Windows** (**ModelOpt-Windows**) is engineered to deliver advanced model compression techniques, including quantization, to Windows RTX PC systems. Specifically tailored to meet the needs of Windows users, ModelOpt-Windows is optimized for rapid and efficient quantization, featuring local GPU calibration, reduced system and video memory consumption, and swift processing times.
|
||||
The primary objective of the ModelOpt-Windows is to generate optimized, standards-compliant ONNX-format models for DirectML backends. This makes it an ideal solution for seamless integration with ONNX Runtime (ORT) and DirectML (DML) frameworks, ensuring broad compatibility with any inference framework supporting the ONNX standard. Furthermore, ModelOpt-Windows integrates smoothly within the Windows ecosystem, with full support for tools and SDKs such as Olive and ONNX Runtime, enabling deployment of quantized models across various independent hardware vendors (IHVs) through the DML path and TensorRT path.
|
||||
The primary objective of the ModelOpt-Windows is to generate optimized, standards-compliant ONNX-format models. This makes it an ideal solution for seamless integration with ONNX Runtime (ORT) and DirectML (DML) frameworks, ensuring broad compatibility with any inference framework supporting the ONNX standard. Furthermore, ModelOpt-Windows integrates smoothly within the Windows ecosystem, with full support for tools and SDKs such as Olive and ONNX Runtime, enabling deployment of quantized models across various independent hardware vendors (IHVs) through the DML path and TensorRT path.
|
||||
|
||||
Model Optimizer is available for free for all developers on [NVIDIA PyPI](https://pypi.org/project/nvidia-modelopt/). This repository is for sharing examples and GPU-optimized recipes as well as collecting feedback from the community.
|
||||
|
||||
@@ -65,14 +65,17 @@ For more details, please refer to the [detailed installation instructions](https
|
||||
|
||||
### Quantization
|
||||
|
||||
Quantization is an effective model optimization technique for large models. Quantization with ModelOpt-Windows can compress model size by 2x-4x, speeding up inference while preserving model quality. ModelOpt-Window enables highly performant quantization formats including INT4, FP8\*, INT8\*, etc. and supports advanced algorithms such as AWQ and SmoothQuant\* focusing on post-training quantization (PTQ) for ONNX and PyTorch\* models with DirectML and TensorRT\* inference backends.
|
||||
Quantization is an effective model optimization technique for large models. Quantization with ModelOpt-Windows can compress model size by 2x-4x, speeding up inference while preserving model quality. ModelOpt-Window enables highly performant quantization formats including INT4, FP8, INT8, etc. and supports advanced algorithms such as AWQ and SmoothQuant\* focusing on post-training quantization (PTQ) for ONNX and PyTorch\* models with DirectML, CUDA and TensorRT\* inference backends.
|
||||
|
||||
For more details, please refer to the [detailed quantization guide](https://nvidia.github.io/TensorRT-Model-Optimizer/guides/windows_guides/_ONNX_PTQ_guide.html).
|
||||
|
||||
## Examples
|
||||
|
||||
- [PTQ for LLMs](./onnx_ptq/README.md) covers how to use ONNX Post-Training Quantization (PTQ) and deployment with DirectML
|
||||
- [MMLU Benchmark](./accuracy_benchmark/README.md) provides an example script for MMLU benchmark and demonstrates how to run it with various popular backends like DirectML, TensorRT-LLM\* and model formats like ONNX and PyTorch\*.
|
||||
- We have ONNX PTQ examples for various ONNX model variants:
|
||||
- [PTQ for GenAI LLMs](./onnx_ptq/genai_llm/README.md) covers how to use ONNX Post-Training Quantization (PTQ) with [ONNX Runtime GenAI](https://onnxruntime.ai/docs/genai) built LLM ONNX models, and thier deployment with DirectML.
|
||||
- [PTQ for Whisper](./onnx_ptq/whisper/README.md) illustrates using ONNX Post-Training Quantization (PTQ) with a Whisper ONNX model (i.e. an ASR model). It also provides example sctipt for Optimum-ORT based inference of Whisper using CUDA EP.
|
||||
- [PTQ for SAM2](./onnx_ptq/sam2/README.md) illustrates using ONNX Post-Training Quantization (PTQ) with a SAM2 ONNX model (i.e. a segmentation model).
|
||||
- [MMLU Benchmark](./accuracy_benchmark/README.md) provides an example script for MMLU benchmarking of LLM models, and demonstrates how to run it with various popular backends like DirectML, TensorRT-LLM\* and model formats like ONNX and PyTorch\*.
|
||||
|
||||
## Support Matrix
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 192 KiB |
@@ -6,6 +6,8 @@
|
||||
- [MMLU (Massive Multitask Language Understanding)](#mmlu-massive-multitask-language-understanding)
|
||||
- [Setup](#setup)
|
||||
- [Evaluation Methods](#evaluation-methods)
|
||||
- [API changes in ONNX Runtime GenAI v0.6](#api-changes-in-onnx-runtime-genai-v0.6)
|
||||
- [Troubleshoot](#troubleshoot)
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -35,8 +37,8 @@ The table below lists the setup steps to prepare your environment for evaluating
|
||||
|----------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| **Open PowerShell as Administrator** | - |
|
||||
| **Create and Activate a Virtual Environment** <br> _(Optional but Recommended)_ | `python -m venv llm_env` <br> `.\llm_env\Scripts\Activate.ps1` |
|
||||
| **Install PyTorch and Related Packages** | `pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu124` |
|
||||
| **Install ONNX Runtime Packages** | `pip install onnxruntime-directml==1.20` <br> `pip install onnxruntime-genai-directml==0.4.0` |
|
||||
| **Install PyTorch and Related Packages** | `pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128` |
|
||||
| **Install ONNX Runtime Packages** | `pip install onnxruntime-directml==1.21.1` <br> `pip install onnxruntime-genai-directml==0.6.0` |
|
||||
| **Install Benchmark Requirements** | `pip install -r requirements.txt` |
|
||||
| **Download MMLU Data** | `mkdir data` <br> `curl -o .\data\mmlu.tar https://people.eecs.berkeley.edu/~hendrycks/data.tar` <br> `tar -xf .\data\mmlu.tar -C .\data` <br> `Move-Item .\data\data .\data\mmlu` |
|
||||
|
||||
@@ -180,3 +182,17 @@ To evaluate the PyTorch Hugging Face (HF) model, use the `--ep pt` argument.
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## API changes in ONNX Runtime GenAI v0.6
|
||||
|
||||
In onnxruntime-genai (GenAI) v0.6, `generator.compute_logits()` and `generator_params.input_ids` are deprecated and new API `generator.append_tokens(List: token_ids)` is added (see GenAI [PR-867](https://github.com/microsoft/onnxruntime-genai/pull/867) for details).
|
||||
|
||||
So, this MMLU script has been updated accordingly - refer following change-snippet from this MMLU script (left works with GenAI < 0.6, right works with GenAI 0.6+). Make sure to update the MMLU script accordingly (left part) for trying it with GenAI < 0.6.
|
||||
|
||||
|
||||
|
||||
## Troubleshoot
|
||||
|
||||
1. In case of any model specific issue (e.g. in tokenizer or in onnxruntime-genai package etc.), one can try using older GenAI e.g. export the ONNX model with `onnxruntime-genai-directml` 0.4 and `transformers` 4.44.
|
||||
|
||||
1. In case of trying out MMLU run of ONNX model through GenAI, make sure that the input model is running fine with GenAI. Onnxruntime-genai has example inference scripts (e.g. see [phi3 example](https://github.com/microsoft/onnxruntime-genai/blob/main/examples/python/phi3-qa.py) script).
|
||||
|
||||
@@ -223,7 +223,7 @@ def evaluate_genai_dml(args, subject, model, dev_df, test_df, model_path):
|
||||
input_tokens = tokenizer.encode(prompt)
|
||||
params = og.GeneratorParams(model)
|
||||
|
||||
params.input_ids = input_tokens
|
||||
# params.input_ids = input_tokens
|
||||
|
||||
if len(input_tokens) + 2 > args.max_seq_length:
|
||||
print(
|
||||
@@ -238,13 +238,16 @@ def evaluate_genai_dml(args, subject, model, dev_df, test_df, model_path):
|
||||
|
||||
params.set_search_options(**search_options)
|
||||
generator = og.Generator(model, params)
|
||||
|
||||
generator.append_tokens(input_tokens)
|
||||
|
||||
new_tokens = []
|
||||
generator.compute_logits()
|
||||
# generator.compute_logits()
|
||||
generator.generate_next_token()
|
||||
new_token = generator.get_next_tokens()[0]
|
||||
new_tokens.append(new_token)
|
||||
if not generator.is_done():
|
||||
generator.compute_logits()
|
||||
# generator.compute_logits()
|
||||
generator.generate_next_token()
|
||||
new_token = generator.get_next_tokens()[0]
|
||||
new_tokens.append(new_token)
|
||||
|
||||
@@ -79,14 +79,7 @@ Refer to the following example scripts and tutorials for deployment:
|
||||
|
||||
### Model Support Matrix
|
||||
|
||||
Model | int4_awq
|
||||
--- | ---
|
||||
Llama3.1-8B-Instruct | Yes
|
||||
Phi3.5-mini-Instruct | Yes
|
||||
Mistral-7B-Instruct-v0.3 | Yes
|
||||
Llama3.2-3B-Instruct| Yes
|
||||
Gemma-2b-it | Yes
|
||||
Nemotron Mini 4B Instruct | Yes
|
||||
Please refer to [support matrix](https://nvidia.github.io/TensorRT-Model-Optimizer/guides/0_support_matrix.html) for a full list of supported features and models.
|
||||
|
||||
### Troubleshoot
|
||||
|
||||
@@ -106,3 +99,11 @@ Nemotron Mini 4B Instruct | Yes
|
||||
1. **Check Read/Write Permissions**
|
||||
|
||||
Ensure that both the input and output model paths have the necessary read and write permissions to avoid any permission-related errors.
|
||||
|
||||
1. **Check Output Path**
|
||||
|
||||
Ensure that output .onnx file doesn't exist already. For example, if the output path is `C:\dir1\dir2\quant\model_quant.onnx` then the path `C:\dir1\dir2\quant` should be valid and the directory `quant` should not already contain `model_quant.onnx` file before quantization. If the output .onnx file already exists, then that can get appended during saving of the quantized model resulting in corrupted or invalid output model.
|
||||
|
||||
1. **Check Input Model**
|
||||
|
||||
During INT4 AWQ execution, the input onnx model (one mentioned in `--onnx_path` argument) will be run with onnxruntime (ORT) for calibration (using ORT EP mentioned in `--calibration_eps` argument). So, make sure that input onnx model is running fine with the specified ORT EP.
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
datasets>=2.14.5
|
||||
torch==2.4.0
|
||||
torch==2.6.0
|
||||
transformers==4.49.0
|
||||
@@ -0,0 +1,96 @@
|
||||
# Scripts for ONNX SAM2 model
|
||||
|
||||
This repository contains an example to demontrate 8-bit quantization of SAM2 ONNX model.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [ONNX export and Inference tool](#onnx-export-and-inference-tool)
|
||||
- [Quantization](#quantization)
|
||||
- [Validated Settings](#validated-settings)
|
||||
- [Troubleshoot](#troubleshoot)
|
||||
|
||||
## ONNX export and Inference tool
|
||||
|
||||
The [samexporter](https://github.com/vietanhdev/samexporter) tool can be used for producing ONNX SAM2 model and for its inference with image inputs. Refer it for its installation and usage details. It internally uses `torch.onnx.export` for doing ONNX export.
|
||||
|
||||
- Opset of the ONNX base model (or exported ONNX model) should be 19+ to support FP8 quantization.
|
||||
|
||||
> Use separate python virtual environment for running the samexporter tool. ModelOpt toolkit and samexporter tool can have conflicting or incompatible dependencies.
|
||||
|
||||
> Currently, the samexporter installation doesn't seem to work out-of-the-box with python 3.12 as samexporter mentions dependency on onnxruntime 1.16.3 which doesn't support python 3.12. So, while working with python 3.12, one can try installing samexporter from source after updating the onnxruntime version to 1.20.0+ in samexporter project's dependencies list (file: pyproject.toml). Also, samexporter installation with python 3.12 is observed to work with [cmake 3.27.7](https://github.com/Kitware/CMake/releases/tag/v3.27.7), but it produced error with cmake 4.0.0-rc4. Make sure to update `PATH` environment variable for cmake/bin directory and then restart the command-line.
|
||||
|
||||
> By default, samexporter tool comes with `onnxruntime` package (i.e. CPU EP). For inference with CUDA EP, one needs to uninstall existing `onnxruntime` package, and then install `onnxruntime-gpu` (1.20.x) package. Make sure that after this update, `python -c "import onnxruntime"` runs successfully.
|
||||
|
||||
## Quantization
|
||||
|
||||
The script `sam2_onnx_quantization.py` supports INT8 (W8A8) and FP8 (W8A8) quantization schemes for `encoder` of the ONNX exported SAM2 model. To use it, install ModelOpt toolkit along with dependencies.
|
||||
|
||||
> Install ModelOpt along with its dependencies (ModelOpt's onnx module installation). Install dependencies mentioned in `requirements.txt` file.
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
Some useful parameters:
|
||||
|
||||
| **Argument** | **Description** |
|
||||
|---------------------------|--------------------------------------------------------------------------------------------------------------|
|
||||
| `--onnx_path ` | Input .onnx file path |
|
||||
| `--output_path ` | Output .onnx file path. |
|
||||
| `--calib_method` | Calibration method for quantization (`max` or `entropy`). Default is `max`. |
|
||||
| `--quant_mode` | Quantization mode to be used (`int8` or `fp8`). Default is `int8`. |
|
||||
| `--calib_size` | Number of input calibration samples. Default is `32`. |
|
||||
| `--use_random_calib` | True when we want to use one randomly generated calibration sample. Default is `False`. |
|
||||
| `--qdq_for_weights` | If True, Q->DQ nodes will be added for weights, otherwise only DQ nodes will be added. Default is `False`. |
|
||||
| `--calibration_eps` | Comma-separated list of calibration endpoints. Choose from 'cuda', 'cpu', 'dml'. Default is "`cuda`, `cpu`". |
|
||||
| `--dtype` | Data-type of the model's tensors. Choose from `fp32`, `fp16`. Default is `fp32`. |
|
||||
| `--image_dir` | Directory containing image files to be used for calibration of sam2's encoder. |
|
||||
| `--image_file_extension` | Extension of image files to be used for calibration of sam2's encoder. E.g. `jpg` (default), `png`. |
|
||||
| `--image_input_dimension` | Last 2 dimensions of the image input to encoder, in comma-separated fashion. Default: `1024,1024` |
|
||||
| `--op_types_to_quantize` | Comma-separated list of op-types to quantize. Choose from 'MatMul', 'Conv'. Default is "`MatMul`". |
|
||||
|
||||
Please refer the script for more details.
|
||||
|
||||
Example command-line:
|
||||
|
||||
```bash
|
||||
|
||||
python .\sam2_onnx_quantization.py --onnx_path=E:\base\sam2_hiera_large.encoder.onnx --output_path=E:\quant\sam2_hiera_large.encoder.onnx --image_dir=E:\sam_image_dataset
|
||||
|
||||
```
|
||||
|
||||
## Validated Settings
|
||||
|
||||
This example is currently validated with following settings:
|
||||
|
||||
- Python 3.11.9
|
||||
- CUDA-12.4
|
||||
- ONNX exported FP32 `sam2_hiera_large` opset-19 model - exported using [samexporter](https://github.com/vietanhdev/samexporter) tool
|
||||
- Quantization settings:
|
||||
- onnx 1.17.0
|
||||
- onnxruntime-gpu 1.20.1 (for ORT-CUDA EP)
|
||||
- Quantization algos - INT8 with `Max` calibration (W8A8) and DQ-only mode for weights, FP8 with `Max` calibration (W8A8) with both DQ-only and QDQ models for weights
|
||||
- op-types-to-quantize = `MatMul`
|
||||
- Calibration size - 32
|
||||
- Calibration EPs - \[`cuda`, `cpu`\]
|
||||
- Calibration data - [SA-1B](https://ai.meta.com/datasets/segment-anything-downloads/) - download and extract sa_000000.tar file.
|
||||
- Quantization support for various ONNX files - `encoder` model is quantized, `decoder` model is not quantized (not needed since decoder is already very small ~20MB).
|
||||
- A separate python virtual environment used for quantization (`python -m venv .\venv_quantization`)
|
||||
- Inference settings:
|
||||
- tool used for inference: [samexporter](https://github.com/vietanhdev/samexporter)
|
||||
- ORT-CUDA EP (needed a minor patch in `samexporter` to pass \[`cuda`, `cpu`\] as providers, file: sam2_onnx.py)
|
||||
- tasks - image prediction or segmentation
|
||||
- precision of ONNX models - `encoder` model is quantized, `decoder` model is not quantized.
|
||||
- test samples - a few image-inputs examples from [samexporter](https://github.com/vietanhdev/samexporter)
|
||||
|
||||
## Troubleshoot
|
||||
|
||||
1. If `samexporter` tool gives error about tensor mismatch or issue in load-state-dictionary, then try hardcoding the absolute path of config file in the tool, as a workaround (file: export_sam2.py).
|
||||
|
||||
1. Sometimes using `samexporter` with numpy 2.x gives error (e.g. ImportError: numpy.core.multiarray failed to import). Try using Numpy 1.26.2.
|
||||
|
||||
```bash
|
||||
pip install numpy==1.26.2
|
||||
```
|
||||
|
||||
1. If inference of the FP8 (dq-only for weights) model results in error about presence of both INT8 and FP8 types in a DQ node, then that would probably be due to error in FP8 type detection by ONNX during quantization. To troubleshoot, we suggest to try out quantization with ONNX latest wheel (1.17+) and/or latest python (e.g. 3.12), in a fresh new virtual environment. Besides, FP8 with QDQ for weights should work.
|
||||
@@ -0,0 +1,3 @@
|
||||
onnx==1.17.0
|
||||
onnxruntime-gpu==1.20.1
|
||||
opencv-python
|
||||
@@ -0,0 +1,251 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from modelopt.onnx.quantization.quantize import quantize as quantize_top_level_api
|
||||
|
||||
logging.getLogger().setLevel(logging.INFO)
|
||||
|
||||
|
||||
def prepare_input(image: np.ndarray, np_dtype, image_input_width, image_input_height) -> np.ndarray:
|
||||
img_height, img_width = image.shape[:2]
|
||||
|
||||
input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
input_img = cv2.resize(input_img, (image_input_width, image_input_height))
|
||||
|
||||
mean = np.array([0.485, 0.456, 0.406])
|
||||
std = np.array([0.229, 0.224, 0.225])
|
||||
input_img = (input_img / 255.0 - mean) / std
|
||||
input_img = input_img.transpose(2, 0, 1)
|
||||
input_tensor = input_img[np.newaxis, :, :, :].astype(np_dtype)
|
||||
|
||||
return input_tensor
|
||||
|
||||
|
||||
def parse_calibration_eps(value):
|
||||
"""Parse and validate the calibration_eps input."""
|
||||
valid_choices = {"cuda", "cpu", "dml"}
|
||||
# Split the input by commas and remove any surrounding whitespace
|
||||
eps = [item.strip() for item in value.split(",")]
|
||||
# Validate each calibration endpoint
|
||||
for ep in eps:
|
||||
if ep not in valid_choices:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"Invalid calibration endpoint: '{ep}'. Choose from 'cuda', 'cpu', 'dml'."
|
||||
)
|
||||
return eps
|
||||
|
||||
|
||||
def parse_op_types_to_quantize(value):
|
||||
"""Parse and validate the op_types_to_quantize input."""
|
||||
valid_choices = {"MatMul", "Conv"}
|
||||
# Split the input by commas and remove any surrounding whitespace
|
||||
op_types = [item.strip() for item in value.split(",")]
|
||||
# Validate each calibration endpoint
|
||||
for op in op_types:
|
||||
if op not in valid_choices:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"Invalid op-type: '{op}'. Choose from 'MatMul', 'Conv'."
|
||||
)
|
||||
return op_types
|
||||
|
||||
|
||||
def get_calib_data_for_encoder(
|
||||
image_directory, calib_size, data_type, file_extension, image_input_dimension
|
||||
):
|
||||
np_dtype = np.float16 if data_type == "fp16" else np.float32
|
||||
|
||||
calib_data = {}
|
||||
|
||||
image_files = [
|
||||
os.path.join(image_directory, f)
|
||||
for f in os.listdir(image_directory)
|
||||
if (os.path.isfile(os.path.join(image_directory, f)) and f.endswith(file_extension))
|
||||
]
|
||||
|
||||
assert len(image_files) > 0, "no image files found for encoder's calibration"
|
||||
|
||||
print(
|
||||
f"\nPreparing calibration data for encoder. Number of images in image-directory = {len(image_files)}\n"
|
||||
)
|
||||
|
||||
image_input_width, image_input_height = image_input_dimension.split(",")
|
||||
image_input_width, image_input_height = int(image_input_width), int(image_input_height)
|
||||
|
||||
for i, image in enumerate(image_files):
|
||||
cv2_image = cv2.imread(image)
|
||||
assert cv2_image is not None, "cv2-image is none"
|
||||
tensor = prepare_input(cv2_image, np_dtype, image_input_width, image_input_height)
|
||||
tensor = tensor.astype(np_dtype)
|
||||
x = calib_data.get("image", None)
|
||||
if x is None:
|
||||
calib_data["image"] = tensor
|
||||
else:
|
||||
calib_data["image"] = np.concatenate((x, tensor), axis=0)
|
||||
if i == calib_size:
|
||||
break
|
||||
|
||||
print(f"\nCalibration data for ENCODER is created. calib_size={calib_size}\n")
|
||||
return calib_data
|
||||
|
||||
|
||||
def main(args):
|
||||
start_time = time.time()
|
||||
|
||||
# args.qdq_for_weights = True
|
||||
|
||||
print("\n\n######### SAM2's 8-bit Quantization: Settings...\n\n")
|
||||
|
||||
print(
|
||||
f" quantization_mode={args.quant_mode},\n calibrartion_method={args.calib_method},"
|
||||
f"\n calib_size={args.calib_size},\n use-random-calib-data={args.use_random_calib},"
|
||||
f"\n op_types_to_quantize={args.op_types_to_quantize},\n calibration-EPs={args.calibration_eps},"
|
||||
f"\n qdq_for_weights={args.qdq_for_weights},\n dq_only_for_weights={not args.qdq_for_weights},"
|
||||
f"\n dtype={args.dtype},\n image_input_dimension={args.image_input_dimension}\n"
|
||||
)
|
||||
print(
|
||||
f"\n input-onnx-path={args.onnx_path},\n image_dir={args.image_dir}"
|
||||
f"\n image_file_extension={args.image_file_extension},\n output-path={args.output_path},\n"
|
||||
)
|
||||
|
||||
print("\n=========================================================\n\n")
|
||||
|
||||
if args.use_random_calib:
|
||||
calib_data = None
|
||||
else:
|
||||
calib_data = get_calib_data_for_encoder(
|
||||
args.image_dir,
|
||||
args.calib_size,
|
||||
args.dtype,
|
||||
args.image_file_extension,
|
||||
args.image_input_dimension,
|
||||
)
|
||||
|
||||
assert args.use_random_calib or calib_data is not None, "calibration data not prepared"
|
||||
|
||||
logging.info("\nQuantizing the model....\n")
|
||||
quantize_top_level_api(
|
||||
onnx_path=args.onnx_path,
|
||||
quantize_mode=args.quant_mode,
|
||||
calibration_method=args.calib_method,
|
||||
calibration_data=None if args.use_random_calib else calib_data,
|
||||
calibration_eps=args.calibration_eps,
|
||||
use_external_data_format=True,
|
||||
output_path=args.output_path,
|
||||
op_types_to_quantize=args.op_types_to_quantize,
|
||||
nodes_to_exclude=[r"/lm_head", r"/Shape"],
|
||||
dq_only=not args.qdq_for_weights,
|
||||
verbose=True,
|
||||
high_precision_dtype="fp16" if args.dtype == "fp16" else "fp32",
|
||||
mha_accumulation_dtype="fp16" if args.dtype == "fp16" else "fp32",
|
||||
enable_gemv_detection_for_trt=False,
|
||||
enable_shared_constants_duplication=False,
|
||||
)
|
||||
logging.info(
|
||||
f"\nQuantization process (along with saving) took {time.time() - start_time} seconds\n"
|
||||
)
|
||||
|
||||
print("\n\nDone\n\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Quantize SAM2 ONNX model with INT8/FP8.")
|
||||
parser.add_argument(
|
||||
"--onnx_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Input ONNX model path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Output quantized model path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--calib_method",
|
||||
type=str,
|
||||
default="max",
|
||||
help="calibration method for quantization (max or entropy)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quant_mode",
|
||||
type=str,
|
||||
default="int8",
|
||||
help="quantization mode to be used (int8 or fp8)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="fp32",
|
||||
help="precision of the model tensors. Choose from 'fp32', 'fp16'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--calib_size",
|
||||
type=int,
|
||||
default=32,
|
||||
help="Number of input calibration samples, should be no more than 256",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_random_calib",
|
||||
type=bool,
|
||||
default=False,
|
||||
help="True when we want to use one randomly generated calibration sample",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--qdq_for_weights",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help="If True, Q->DQ nodes will be added for weights, otherwise only DQ nodes will be added.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--calibration_eps",
|
||||
type=parse_calibration_eps, # Use the custom parser
|
||||
default=["cuda", "cpu"], # Default as a list
|
||||
help="Comma-separated list of calibration endpoints. Choose from 'cuda', 'cpu', 'dml'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--op_types_to_quantize",
|
||||
type=parse_op_types_to_quantize, # Use the custom parser
|
||||
default=["MatMul"], # Default as a list
|
||||
help="Comma-separated list of op-types that need to be quantized. Choose from 'MatMul', 'Conv'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory containing image files to be used for calibration of sam2's encoder.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_file_extension",
|
||||
type=str,
|
||||
default="jpg",
|
||||
help="Extension of image files to be used for calibration of sam2's encoder. E.g. jpg, png.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_input_dimension",
|
||||
type=str,
|
||||
default="1024,1024",
|
||||
help="Last 2 dimensions of the image input to encoder, in comma-separated fashion",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,135 @@
|
||||
# Scripts for ONNX Whisper model
|
||||
|
||||
This repository contains an example to demontrate 8-bit quantization of Whisper ONNX model.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [ONNX export](#onnx-export)
|
||||
- [Inference script](#inference-script)
|
||||
- [Quantization script](#quantization-script)
|
||||
- [Validated Settings](#validated-settings)
|
||||
|
||||
## ONNX export
|
||||
|
||||
The HuggingFace Optimum-CLI tool can be used for export of the HuggingFace Whisper model.
|
||||
|
||||
Example command-line to obtain FP32 whisper_large model:
|
||||
|
||||
```bash
|
||||
|
||||
optimum-cli export onnx --model openai/whisper-large E:\model_store\optimum\whisper_large --task automatic-speech-recognition-with-past --opset 20
|
||||
|
||||
```
|
||||
|
||||
Example command-line to obtain FP16 whisper_medium model:
|
||||
|
||||
```bash
|
||||
|
||||
optimum-cli export onnx --model openai/whisper-medium E:\model_store\optimum\whisper_medium --task automatic-speech-recognition-with-past --opset 20 --dtype fp16 --device cuda
|
||||
|
||||
```
|
||||
|
||||
## Install dependencies
|
||||
|
||||
1. Install ModelOpt along with its dependencies (ModelOpt's onnx module installation).
|
||||
1. Install dependencies mentioned in `requirements.txt` file.
|
||||
|
||||
```bash
|
||||
|
||||
pip install -r requirements.txt
|
||||
|
||||
```
|
||||
|
||||
## Inference script
|
||||
|
||||
The script `whisper_optimum_ort_inference.py` is for Optimum-ORT based inference of an ONNX Whisper model. It takes an audio file (.wav) as input and transcribes its content in english. This script also supports Word Error Rate (WER) accuracy measurement.
|
||||
|
||||
To run the inference of Whisper ONNX model, relevant files like encoder_model.onnx, decoder_model.onnx, decoder_with_past_model.onnx, tokenizer files, config.json, generation_config.json, vocab file etc. should be kept together in a directory and should provide that directory path to the inference script.
|
||||
|
||||
Useful parameters:
|
||||
|
||||
| **Argument** | **Description** |
|
||||
|---------------------------|------------------------------------------------------------------------------------------------------|
|
||||
| `--model_name` | Specifies the HuggingFace model ID |
|
||||
| `--onnx_model_dir` | Specifies the directory contains all relevant files for the ONNX model. |
|
||||
| `--inference_ep` | Specifies the EP to use for inference. Default is CUDA EP. |
|
||||
| `--test_samples_count` | Specifies the count of samples to use for inference. Default is 100. |
|
||||
| `--log_model_output` | Specifies whether to log inference output of the model. Default is off. |
|
||||
| `--cache_dir` | Specifies the cache directory to use for HuggingFace files |
|
||||
| `--dtype` | Data-type of the model's tensors. Choose from `fp32`, `fp16`. Default is `fp32`. |
|
||||
| `--audio_file_path` | Path of the input audio file in .wav format. |
|
||||
| `--run_wer_test` | If True, runs WER accuracy benchmarking using samples from librispeech_asr dataset. Default is False.|
|
||||
|
||||
Please refer the script for more details.
|
||||
|
||||
Example command-line:
|
||||
|
||||
```bash
|
||||
|
||||
python .\whisper_optimum_ort_inference.py --model_name=openai/whisper-large --onnx_model_dir=E:\whisper_large \
|
||||
--audio_file_path=E:\demo.wav \
|
||||
--run_wer_test --test_samples_count=50
|
||||
|
||||
```
|
||||
|
||||
## Quantization script
|
||||
|
||||
The script `whisper_onnx_quantization.py` supports various quantization schemes for the given ONNX whisper model.
|
||||
|
||||
Following are some useful parameters of this script. Please refer the script for more details.
|
||||
|
||||
| **Argument** | **Description** |
|
||||
|---------------------------|--------------------------------------------------------------------------------------------------------------|
|
||||
| `--model_name` | HuggingFace model ID |
|
||||
| `--base_model_dir` | Directory containing all relevant files for the exported ONNX model |
|
||||
| `--onnx_path ` | Input .onnx file path |
|
||||
| `--output_path ` | Output .onnx file path. |
|
||||
| `--calib_method` | Calibration method for quantization (`max` or `entropy`). Default is `max`. |
|
||||
| `--quant_mode` | Quantization mode to be used (`int8` or `fp8`). Default is `int8`. |
|
||||
| `--calib_size` | Number of input calibration samples. Default is `32`. |
|
||||
| `--batch_size` | Batch size for calibration samples. Default is `1`. |
|
||||
| `--use_random_calib` | True when we want to use one randomly generated calibration sample. Default is `False`. |
|
||||
| `--qdq_for_weights` | If True, Q->DQ nodes will be added for weights, otherwise only DQ nodes will be added. Default is `False`. |
|
||||
| `--calibration_eps` | Comma-separated list of calibration endpoints. Choose from 'cuda', 'cpu', 'dml'. Default is \[`cuda`, `cpu`\]. |
|
||||
| `--cache_dir` | Cache directory for HuggingFace files. Change this as needed on your system. |
|
||||
| `--dtype` | Data-type of the model's tensors. Choose from `fp32`, `fp16`. Default is `fp32`. |
|
||||
|
||||
See below for example command-lines.
|
||||
|
||||
```bash
|
||||
python .\whisper_onnx_quantization.py --model_name=openai/whisper-large --base_model_dir=E:\whisper_large\base \
|
||||
--onnx_path=E:\whisper_large\base\encoder_model.onnx \
|
||||
--output_path=E:\whisper_large\quant_output\encoder_model.onnx
|
||||
|
||||
python .\whisper_onnx_quantization.py --model_name=openai/whisper-large --base_model_dir=E:\whisper_large\base \
|
||||
--onnx_path=E:\whisper_large\base\decoder_model.onnx \
|
||||
--output_path=E:\whisper_large\quant_output\decoder_model.onnx
|
||||
```
|
||||
|
||||
- Make sure to use GPU-compatible version of dependencies (torch, torchaudio). For instance, installing cu128 binaries of torch and torchaudio can work for RTX 5090 Blackwell system. Example command-line for installing cu12.8-based nightly wheels of torch and torchaudio is provided below. One can switch to cu12.8 public wheels of torch/torchaudio for RTX-5090, when they are available.
|
||||
|
||||
```bash
|
||||
pip install torchaudio==2.6.0.dev20250306 --extra-index-url https://download.pytorch.org/whl/nightly/cu128
|
||||
```
|
||||
|
||||
- The Whisper quantization script supports quantization of following Whisper ONNX files: `encoder_model.onnx`, `decoder_model.onnx`, `decoder_with_past_model.onnx`.
|
||||
|
||||
## Validated Settings
|
||||
|
||||
These scripts are currently validated with following settings:
|
||||
|
||||
- Python 3.11.9
|
||||
- CUDA settings on Host - CUDA 12.4, cuDNN 9.5 (cudnn-windows-x86_64-9.5.0.50_cuda12-archive)
|
||||
- Windows11 22621
|
||||
- RTX 4090
|
||||
- Base PyTorch model - [HuggingFace openai\\whisper-large model](https://huggingface.co/openai/whisper-large)
|
||||
- ONNX Exporter - HuggingFace Optimum, opset-20, FP32 ONNX model
|
||||
- Inference EP - CUDA EP
|
||||
- Test samples count - 100 (for inference / WER-accuracy-benchmarking)
|
||||
- Quantization algos - INT8 with `Max` calibration (W8A8), FP8 with `Max` calibration (W8A8)
|
||||
- Calibration size - 32
|
||||
- Calibration EPs - \[`cuda`, `cpu`\]
|
||||
- Calibration data - `librispeech_asr` dataset
|
||||
- `load_dataset("librispeech_asr", "clean", split="test", trust_remote_code=True)`
|
||||
- Quantization support for various ONNX files - `encoder_model.onnx`, `decoder_model.onnx`, `decoder_with_past_model.onnx`
|
||||
- The `use_merged` argument in optimum-ORT's Whisper model API is kept False.
|
||||
Binary file not shown.
@@ -0,0 +1,13 @@
|
||||
--extra-index-url=https://download.pytorch.org/whl/cu128
|
||||
accelerate
|
||||
datasets==2.19.0
|
||||
evaluate
|
||||
jiwer
|
||||
librosa
|
||||
onnx==1.16.1
|
||||
onnxruntime-gpu==1.20.1
|
||||
optimum==1.23.3
|
||||
soundfile
|
||||
torch==2.7.0+cu128
|
||||
torchaudio==2.7.0+cu128
|
||||
transformers==4.48.0
|
||||
@@ -0,0 +1,407 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
|
||||
from transformers import WhisperProcessor
|
||||
|
||||
from modelopt.onnx.quantization.quantize import quantize as quantize_top_level_api
|
||||
|
||||
logging.getLogger().setLevel(logging.INFO)
|
||||
|
||||
|
||||
USE_MERGED = False
|
||||
|
||||
|
||||
def get_ep_for_decoder_calib_data_preparation(calibration_eps: list[str]):
|
||||
provider = None
|
||||
if "cuda" in calibration_eps:
|
||||
provider = "CUDAExecutionProvider"
|
||||
elif "dml" in calibration_eps:
|
||||
provider = "DmlExecutionProvider"
|
||||
elif "cpu" in calibration_eps:
|
||||
provider = "CPUExecutionProvider"
|
||||
else:
|
||||
assert 0, "unknonwn ep"
|
||||
|
||||
return provider
|
||||
|
||||
|
||||
def parse_calibration_eps(value):
|
||||
"""Parse and validate the calibration_eps input."""
|
||||
valid_choices = {"cuda", "cpu", "dml"}
|
||||
# Split the input by commas and remove any surrounding whitespace
|
||||
eps = [item.strip() for item in value.split(",")]
|
||||
# Validate each calibration endpoint
|
||||
for ep in eps:
|
||||
if ep not in valid_choices:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"Invalid calibration endpoint: '{ep}'. Choose from 'cuda', 'cpu', 'dml'."
|
||||
)
|
||||
return eps
|
||||
|
||||
|
||||
def is_encoder_model(input_onnx_path):
|
||||
is_encoder = False
|
||||
|
||||
if "encoder_model.onnx" in args.onnx_path:
|
||||
is_encoder = True
|
||||
|
||||
return is_encoder
|
||||
|
||||
|
||||
def is_decoder_model(input_onnx_path):
|
||||
is_decoder = False
|
||||
|
||||
if "decoder_model.onnx" in args.onnx_path:
|
||||
is_decoder = True
|
||||
|
||||
return is_decoder
|
||||
|
||||
|
||||
def is_decoder_with_past_model(input_onnx_path):
|
||||
is_decoder_with_past = False
|
||||
|
||||
if "decoder_with_past_model.onnx" in args.onnx_path:
|
||||
is_decoder_with_past = True
|
||||
|
||||
return is_decoder_with_past
|
||||
|
||||
|
||||
def get_calib_data_for_encoder(asr_dataset, processor, calib_size, data_type):
|
||||
np_dtype = np.float16 if data_type == "fp16" else np.float32
|
||||
|
||||
calib_data = {}
|
||||
|
||||
for idx, batch in enumerate(asr_dataset):
|
||||
audio = batch["audio"]
|
||||
# inp = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
|
||||
inp = (
|
||||
processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt")
|
||||
.input_features.cpu()
|
||||
.numpy()
|
||||
.astype(np_dtype)
|
||||
)
|
||||
x = calib_data.get("input_features", None)
|
||||
if x is None:
|
||||
calib_data["input_features"] = inp
|
||||
else:
|
||||
calib_data["input_features"] = np.concatenate((x, inp[np.newaxis :,]), axis=0)
|
||||
if idx == calib_size:
|
||||
break
|
||||
|
||||
print(f"\nCalibration data for ENCODER is created. calib_size={calib_size}\n")
|
||||
return calib_data
|
||||
|
||||
|
||||
def get_calib_data_for_decoder(
|
||||
asr_dataset, processor, calib_size, base_model_dir, ep_list, data_type
|
||||
):
|
||||
torch_dtype = torch.float16 if data_type == "fp16" else torch.float32
|
||||
|
||||
calib_data = {}
|
||||
|
||||
provider = get_ep_for_decoder_calib_data_preparation(ep_list)
|
||||
|
||||
model = ORTModelForSpeechSeq2Seq.from_pretrained(
|
||||
base_model_dir, provider=provider, cache_dir=args.cache_dir, use_merged=USE_MERGED
|
||||
)
|
||||
|
||||
for idx, batch in enumerate(asr_dataset):
|
||||
audio = batch["audio"]
|
||||
inp = processor(
|
||||
audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt"
|
||||
).to(dtype=torch_dtype)
|
||||
encoder_outputs = model.encoder(inp.input_features, attention_mask=None)
|
||||
last_hidden_state = encoder_outputs.last_hidden_state.cpu().numpy()
|
||||
decoder_input_ids = (
|
||||
torch.tensor([[model.config.decoder_start_token_id]]).to(torch.int64).cpu().numpy()
|
||||
) # to("cuda")
|
||||
# decoder_input_ids = model.decoder.embed_tokens(decoder_input_ids)
|
||||
# decoder_input_ids = torch.ones((batch_size, 2), dtype=torch.int64, device="cuda")
|
||||
# * model.config.decoder_start_token_id
|
||||
x = calib_data.get("input_ids", None)
|
||||
if x is None:
|
||||
assert calib_data.get("encoder_hidden_states", None) is None, (
|
||||
"encoder-hidden-states is not None but input-ids is"
|
||||
)
|
||||
calib_data["input_ids"] = decoder_input_ids
|
||||
calib_data["encoder_hidden_states"] = last_hidden_state
|
||||
else:
|
||||
calib_data["input_ids"] = np.concatenate((x, decoder_input_ids[np.newaxis :,]), axis=0)
|
||||
x = calib_data.get("encoder_hidden_states", None)
|
||||
assert x is not None, "encoder-hidden-states is None but not input-ids"
|
||||
calib_data["encoder_hidden_states"] = np.concatenate(
|
||||
(x, last_hidden_state[np.newaxis :,]), axis=0
|
||||
)
|
||||
|
||||
if idx == calib_size:
|
||||
break
|
||||
|
||||
print(f"\nCalibration data for DECODER is created. calib_size={calib_size}\n")
|
||||
return calib_data
|
||||
|
||||
|
||||
def get_calib_data_for_decoder_with_past(
|
||||
asr_dataset, processor, calib_size, base_model_dir, ep_list, data_type
|
||||
):
|
||||
torch_dtype = torch.float16 if data_type == "fp16" else torch.float32
|
||||
|
||||
calib_data = {}
|
||||
|
||||
provider = get_ep_for_decoder_calib_data_preparation(ep_list)
|
||||
|
||||
model = ORTModelForSpeechSeq2Seq.from_pretrained(
|
||||
base_model_dir, provider=provider, cache_dir=args.cache_dir, use_merged=USE_MERGED
|
||||
)
|
||||
|
||||
for idx, batch in enumerate(asr_dataset):
|
||||
audio = batch["audio"]
|
||||
inp = processor(
|
||||
audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt"
|
||||
).to(dtype=torch_dtype)
|
||||
|
||||
encoder_outputs = model.encoder(inp.input_features, attention_mask=None)
|
||||
|
||||
last_hidden_state = encoder_outputs.last_hidden_state.cpu()
|
||||
|
||||
# TODO Read suitable config instead of hard-coding token-ids for language (en),
|
||||
# task (transcribe), and no_timestamp_ids. The task-transcribe-id might be same as some
|
||||
# of forced-ids - check Optimum-ORT pipeline e.g. see _retrieve_init_tokens()
|
||||
decoder_input_ids = (
|
||||
torch.tensor([[model.config.decoder_start_token_id, 50259, 50359, 50363]])
|
||||
.to(torch.int64)
|
||||
.cpu()
|
||||
)
|
||||
|
||||
decoder_outputs = model.decoder(decoder_input_ids, last_hidden_state)
|
||||
|
||||
# TODO batch size = 1 -> assert here after passing argument
|
||||
next_token_id = (
|
||||
torch.argmax(decoder_outputs.logits[:, -1, :], dim=-1).view(-1, 1).cpu().numpy()
|
||||
)
|
||||
cache_position = torch.tensor([4]).to(torch.int64).cpu().numpy()
|
||||
|
||||
assert decoder_outputs.past_key_values is not None, "missing past-KV values"
|
||||
assert len(decoder_outputs.past_key_values) == model.config.num_hidden_layers, (
|
||||
"different amout of KV-data"
|
||||
)
|
||||
assert len(decoder_outputs.past_key_values[0]) == 4, "different per-layer KV-data length"
|
||||
|
||||
x = calib_data.get("input_ids", None)
|
||||
if x is None:
|
||||
assert calib_data.get("cache_position", None) is None, (
|
||||
"cache_position is not None but input-ids is"
|
||||
)
|
||||
calib_data["input_ids"] = next_token_id
|
||||
calib_data["cache_position"] = cache_position
|
||||
for i, kv_data in enumerate(decoder_outputs.past_key_values):
|
||||
# kv_data = kv_data.cpu().numpy()
|
||||
calib_data[f"past_key_values.{i}.decoder.key"] = kv_data[0].cpu().numpy()
|
||||
calib_data[f"past_key_values.{i}.decoder.value"] = kv_data[1].cpu().numpy()
|
||||
calib_data[f"past_key_values.{i}.encoder.key"] = kv_data[2].cpu().numpy()
|
||||
calib_data[f"past_key_values.{i}.encoder.value"] = kv_data[3].cpu().numpy()
|
||||
else:
|
||||
calib_data["input_ids"] = np.concatenate((x, next_token_id), axis=0)
|
||||
x = calib_data.get("cache_position", None)
|
||||
assert x is not None, "cache_position is None but not input-ids"
|
||||
calib_data["cache_position"] = np.concatenate((x, cache_position), axis=0)
|
||||
for i, kv_data in enumerate(decoder_outputs.past_key_values):
|
||||
# kv_data = kv_data.cpu().numpy()
|
||||
x = calib_data[f"past_key_values.{i}.decoder.key"]
|
||||
calib_data[f"past_key_values.{i}.decoder.key"] = np.concatenate(
|
||||
(x, kv_data[0].cpu().numpy()), axis=0
|
||||
)
|
||||
x = calib_data[f"past_key_values.{i}.decoder.value"]
|
||||
calib_data[f"past_key_values.{i}.decoder.value"] = np.concatenate(
|
||||
(x, kv_data[1].cpu().numpy()), axis=0
|
||||
)
|
||||
x = calib_data[f"past_key_values.{i}.encoder.key"]
|
||||
calib_data[f"past_key_values.{i}.encoder.key"] = np.concatenate(
|
||||
(x, kv_data[2].cpu().numpy()), axis=0
|
||||
)
|
||||
x = calib_data[f"past_key_values.{i}.encoder.value"]
|
||||
calib_data[f"past_key_values.{i}.encoder.value"] = np.concatenate(
|
||||
(x, kv_data[3].cpu().numpy()), axis=0
|
||||
)
|
||||
|
||||
if idx == calib_size:
|
||||
break
|
||||
|
||||
print(f"\nCalibration data for DECODER_WITH_PAST is created. calib_size={calib_size}\n")
|
||||
return calib_data
|
||||
|
||||
|
||||
def main(args):
|
||||
start_time = time.time()
|
||||
|
||||
assert args.batch_size == 1, "batch size is NOT 1"
|
||||
|
||||
# args.qdq_for_weights = True
|
||||
|
||||
print("\n\n######### Whisper's 8-bit Quantization: Settings...\n\n")
|
||||
|
||||
print(
|
||||
f" quantization_mode={args.quant_mode},\n calibrartion_method={args.calib_method},"
|
||||
f"\n calib_size={args.calib_size},\n batch_size={args.batch_size},"
|
||||
f"\n use-random-calib-data={args.use_random_calib},\n torch-is-cuda-available={torch.cuda.is_available()},"
|
||||
f"\n calibration-EPs={args.calibration_eps},\n dtype={args.dtype},\n USE_MERGED={USE_MERGED},"
|
||||
f"\n qdq_for_weights={args.qdq_for_weights}, dq-only-for-weights={not args.qdq_for_weights}\n"
|
||||
)
|
||||
print(
|
||||
f"\n model-name (id)={args.model_name},\n input-onnx-path ={args.onnx_path},"
|
||||
f"\n output-path={args.output_path},\n base_model_dir={args.base_model_dir},\n cache_dir={args.cache_dir}\n"
|
||||
)
|
||||
|
||||
print("\n=========================================================\n\n")
|
||||
|
||||
processor = WhisperProcessor.from_pretrained(args.model_name, cache_dir=args.cache_dir)
|
||||
|
||||
asr_dataset = load_dataset("librispeech_asr", "clean", split="test", trust_remote_code=True)
|
||||
# asr_dataset = load_dataset("librispeech_asr", "all", split="test.clean")
|
||||
|
||||
calib_data = None
|
||||
if not args.use_random_calib and is_decoder_model(args.onnx_path):
|
||||
calib_data = get_calib_data_for_decoder(
|
||||
asr_dataset,
|
||||
processor,
|
||||
args.calib_size,
|
||||
args.base_model_dir,
|
||||
args.calibration_eps,
|
||||
args.dtype,
|
||||
)
|
||||
elif not args.use_random_calib and is_encoder_model(args.onnx_path):
|
||||
calib_data = get_calib_data_for_encoder(asr_dataset, processor, args.calib_size, args.dtype)
|
||||
elif not args.use_random_calib and is_decoder_with_past_model(args.onnx_path):
|
||||
calib_data = get_calib_data_for_decoder_with_past(
|
||||
asr_dataset,
|
||||
processor,
|
||||
args.calib_size,
|
||||
args.base_model_dir,
|
||||
args.calibration_eps,
|
||||
args.dtype,
|
||||
)
|
||||
|
||||
assert args.use_random_calib or calib_data is not None, "calibration data not prepared"
|
||||
|
||||
logging.info("\nQuantizing the model....\n")
|
||||
quantize_top_level_api(
|
||||
onnx_path=args.onnx_path,
|
||||
quantize_mode=args.quant_mode,
|
||||
calibration_method=args.calib_method,
|
||||
calibration_data=None if args.use_random_calib else calib_data,
|
||||
calibration_eps=args.calibration_eps,
|
||||
use_external_data_format=True,
|
||||
output_path=args.output_path,
|
||||
op_types_to_quantize=["MatMul"],
|
||||
nodes_to_quantize=[r"\S*MatMul[\S]*"],
|
||||
nodes_to_exclude=[r"/lm_head", r"/Shape"],
|
||||
dq_only=not args.qdq_for_weights,
|
||||
verbose=True,
|
||||
high_precision_dtype="fp16" if args.dtype == "fp16" else "fp32",
|
||||
mha_accumulation_dtype="fp16" if args.dtype == "fp16" else "fp32",
|
||||
enable_gemv_detection_for_trt=False,
|
||||
enable_shared_constants_duplication=False,
|
||||
)
|
||||
logging.info(
|
||||
f"\nQuantization process (along with saving) took {time.time() - start_time} seconds\n"
|
||||
)
|
||||
|
||||
print("\n\nDone\n\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Quantize Whisper ONNX model with INT8/FP8.")
|
||||
parser.add_argument("--model_name", type=str, required=True, help="Name or HF id of the model")
|
||||
parser.add_argument(
|
||||
"--base_model_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory of the base ONNX model",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--onnx_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Input ONNX model path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Output quantized model path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--calib_method",
|
||||
type=str,
|
||||
default="max",
|
||||
help="calibration method for quantization (max or entropy)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quant_mode",
|
||||
type=str,
|
||||
default="int8",
|
||||
help="quantization mode to be used (int8 or fp8)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="fp32",
|
||||
help="precision of the model tensors. Choose from 'fp32', 'fp16'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--calib_size",
|
||||
type=int,
|
||||
default=32,
|
||||
help="Number of input calibration samples, should be no more than 256",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Batch size for calibration samples",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_random_calib",
|
||||
type=bool,
|
||||
default=False,
|
||||
help="True when we want to use one randomly generated calibration sample",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--qdq_for_weights",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help="If True, Q->DQ nodes will be added for weights, otherwise only DQ nodes will be added.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--calibration_eps",
|
||||
type=parse_calibration_eps, # Use the custom parser
|
||||
default=["cuda", "cpu"], # Default as a list
|
||||
help="Comma-separated list of calibration endpoints. Choose from 'cuda', 'cpu', 'dml'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache_dir",
|
||||
type=str,
|
||||
help="cache directory for HuggingFace files",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,175 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
from datasets import load_dataset
|
||||
from evaluate import load
|
||||
from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
|
||||
from tqdm import tqdm
|
||||
from transformers import WhisperProcessor
|
||||
|
||||
USE_MERGED = False
|
||||
|
||||
|
||||
def get_ep(inference_ep: list[str]):
|
||||
provider = None
|
||||
if "cuda" in inference_ep:
|
||||
provider = "CUDAExecutionProvider"
|
||||
elif "dml" in inference_ep:
|
||||
provider = "DmlExecutionProvider"
|
||||
elif "cpu" in inference_ep:
|
||||
provider = "CPUExecutionProvider"
|
||||
else:
|
||||
assert 0, "unknonwn ep"
|
||||
|
||||
return provider
|
||||
|
||||
|
||||
def main(args):
|
||||
data_type = torch.float16 if args.dtype == "fp16" else torch.float32
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
print("\n########### ONNX Whisper Inference: Settings...\n\n")
|
||||
print(
|
||||
f"\n model-name (id)={args.model_name},\n onnx_model_dir={args.onnx_model_dir},"
|
||||
f"\n inference_ep={args.inference_ep},\n test_samples_count={args.test_samples_count},"
|
||||
f"\n log_model_outputs={args.log_model_output},\n cache_dir={args.cache_dir},"
|
||||
f"\n dtype={data_type},\n USE_MERGED={USE_MERGED},"
|
||||
f"\n audio_file_path={args.audio_file_path},\n run_wer_test={args.run_wer_test}\n"
|
||||
)
|
||||
print("\n=========================================================\n\n")
|
||||
|
||||
processor = WhisperProcessor.from_pretrained(args.model_name, cache_dir=args.cache_dir)
|
||||
model = ORTModelForSpeechSeq2Seq.from_pretrained(
|
||||
args.onnx_model_dir,
|
||||
provider=get_ep(args.inference_ep),
|
||||
cache_dir=args.cache_dir,
|
||||
use_merged=USE_MERGED,
|
||||
)
|
||||
|
||||
# print(model.encoder)
|
||||
# print(model.decoder)
|
||||
|
||||
test_sample, sample_rate = torchaudio.load(args.audio_file_path)
|
||||
test_sample = test_sample.numpy()[0]
|
||||
inp = processor(test_sample, sampling_rate=sample_rate, return_tensors="pt").input_features.to(
|
||||
"cuda", dtype=data_type
|
||||
)
|
||||
# print(f"\n\n--demo-audio-- -input-features- type={type(inp)}, shape={inp.shape}, {inp}\n")
|
||||
|
||||
predicted_ids = model.generate(inp)[0]
|
||||
transcription = processor.decode(predicted_ids)
|
||||
prediction = processor.tokenizer._normalize(transcription)
|
||||
|
||||
print(f"\n\n-- Content of input audio-file = {prediction}\n\n")
|
||||
|
||||
if args.run_wer_test:
|
||||
librispeech_test_clean = load_dataset(
|
||||
"librispeech_asr", "clean", split="test", trust_remote_code=True
|
||||
)
|
||||
|
||||
references = []
|
||||
predictions = []
|
||||
|
||||
for idx in tqdm(range(args.test_samples_count), desc="Evaluating..."):
|
||||
# audio = batch["audio"]
|
||||
audio = librispeech_test_clean[idx]["audio"]
|
||||
inp = processor(
|
||||
audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt"
|
||||
)
|
||||
input_features = inp.input_features
|
||||
reference = processor.tokenizer._normalize(
|
||||
librispeech_test_clean[idx]["text"]
|
||||
) # batch['text'], librispeech_test_clean[idx]['text']
|
||||
references.append(reference)
|
||||
if args.log_model_output:
|
||||
print(f"\n\n--Evaluate-- reference-{idx}={reference}\n\n")
|
||||
input_features = input_features.to("cuda", dtype=data_type)
|
||||
predicted_ids = model.generate(input_features)[0]
|
||||
transcription = processor.decode(predicted_ids)
|
||||
prediction = processor.tokenizer._normalize(transcription)
|
||||
predictions.append(prediction)
|
||||
if args.log_model_output:
|
||||
print(f"\n\n--Evaluate-- prediction-{idx}={prediction}\n\n")
|
||||
|
||||
wer = load("wer")
|
||||
wer_result = wer.compute(references=references, predictions=predictions)
|
||||
|
||||
print(
|
||||
f"\n## DONE ## - wer = {wer_result}, wer% = {wer_result * 100}, accuracy% = {(1 - wer_result) * 100},"
|
||||
f"\n total-time = {time.time() - start_time} seconds,"
|
||||
f"\n num-distinct-inputs={len(set(references))},"
|
||||
f"\n len-reference={len(references)}, len-predictions={len(predictions)}\n\n"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Inference of ONNX Whisper model.")
|
||||
parser.add_argument("--model_name", type=str, required=True, help="Name or HF id of the model")
|
||||
parser.add_argument(
|
||||
"--onnx_model_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory of the ONNX model files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inference_ep",
|
||||
type=str,
|
||||
default="cuda",
|
||||
help="ORT-EP to be used by optimum-ORT for inference. Choose from 'cuda', 'cpu', 'dml'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="fp32",
|
||||
help="Precision of the model tensors. Choose from 'fp32', 'fp16'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--test_samples_count",
|
||||
type=int,
|
||||
default=100,
|
||||
help="Number of audio samples to evaluate",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log_model_output",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help="If True, model's output are logged along with reference texts for the input audio sample",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache_dir",
|
||||
type=str,
|
||||
default="E:\\G_script\\cache",
|
||||
help="cache directory for HuggingFace files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio_file_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path of the input audio file in .wav format",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--run_wer_test",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help="If True, runs WER accuracy benchmarking using samples from librispeech_asr dataset",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -26,6 +26,7 @@ import onnx
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from onnx import ModelProto
|
||||
from packaging.version import Version
|
||||
from torch.nn.parallel import DataParallel, DistributedDataParallel
|
||||
|
||||
from modelopt.onnx.quantization.qdq_utils import qdq_to_dq
|
||||
@@ -297,9 +298,9 @@ def get_onnx_bytes_and_metadata(
|
||||
onnx_load_path: The path to load the onnx model.
|
||||
dynamic_axes: A dictionary of dynamic shapes used for exporting the torch model to onnx.
|
||||
remove_exported_model: If True, the onnx model will be cleared from the disk after the
|
||||
export process
|
||||
dynamo_export: If True, the model is exported using dynamo_export. See
|
||||
`torch.onnx.dynamo_export https://pytorch.org/docs/stable/onnx_dynamo.html` for more details.
|
||||
export process.
|
||||
dynamo_export: If True, the model is exported using `dynamo=True` in
|
||||
`torch.onnx.export <https://pytorch.org/docs/stable/onnx.html#torch.onnx.export>`_.
|
||||
onnx_opset: The onnx opset version to use for exporting the model.
|
||||
dq_only: If True, the exported ONNX model is converted to a dq_only model.
|
||||
|
||||
@@ -328,10 +329,6 @@ def get_onnx_bytes_and_metadata(
|
||||
if dummy_input and isinstance(dummy_input[-1], dict):
|
||||
dummy_input = dummy_input + ({},) # we need to add an extra dict for the fake kwargs!
|
||||
|
||||
dummy_input_args, dummy_input_kwargs = (), None
|
||||
if dynamo_export:
|
||||
dummy_input_args, dummy_input_kwargs = split_args_kwargs(dummy_input)
|
||||
|
||||
# Get input tree spec, see generate_onnx_input for more info as well on this
|
||||
flat_input, tree_spec_input = flatten_tree(named_args)
|
||||
|
||||
@@ -377,7 +374,7 @@ def get_onnx_bytes_and_metadata(
|
||||
# Export onnx model from pytorch model
|
||||
# As the maximum size of protobuf is 2GB, we cannot use io.BytesIO() buffer during export.
|
||||
with torch.inference_mode():
|
||||
if not dynamo_export:
|
||||
if not dynamo_export or Version(torch.__version__) >= Version("2.6"):
|
||||
torch.onnx.export(
|
||||
model,
|
||||
dummy_input,
|
||||
@@ -386,19 +383,17 @@ def get_onnx_bytes_and_metadata(
|
||||
output_names=output_names,
|
||||
opset_version=onnx_opset,
|
||||
dynamic_axes=dynamic_axes,
|
||||
dynamo=dynamo_export,
|
||||
)
|
||||
else:
|
||||
else: # torch < 2.6 with dynamo export
|
||||
export_options = torch.onnx.ExportOptions(dynamic_shapes=True)
|
||||
if dummy_input_kwargs:
|
||||
torch.onnx.dynamo_export(
|
||||
model, *dummy_input_args, export_options=export_options, **dummy_input_kwargs
|
||||
).save(onnx_save_path)
|
||||
else:
|
||||
torch.onnx.dynamo_export(
|
||||
model,
|
||||
*dummy_input_args,
|
||||
export_options=export_options,
|
||||
).save(onnx_save_path)
|
||||
dummy_input_args, dummy_input_kwargs = split_args_kwargs(dummy_input)
|
||||
if dummy_input_kwargs is None:
|
||||
dummy_input_kwargs = {}
|
||||
torch.onnx.dynamo_export(
|
||||
model, *dummy_input_args, export_options=export_options, **dummy_input_kwargs
|
||||
).save(onnx_save_path)
|
||||
|
||||
# Check that export worked
|
||||
assert len(os.listdir(onnx_path)) > 0, "Torch to onnx export failed."
|
||||
|
||||
|
||||
@@ -16,12 +16,15 @@
|
||||
"""The package setup script for modelopt customizing certain aspects of the installation process."""
|
||||
|
||||
import os
|
||||
import platform
|
||||
|
||||
import setuptools
|
||||
|
||||
# Package configuration ############################################################################
|
||||
name = "nvidia-modelopt"
|
||||
version = os.environ.get("SETUPTOOLS_SCM_PRETEND_VERSION", "0.27.1")
|
||||
version = os.environ.get(
|
||||
"SETUPTOOLS_SCM_PRETEND_VERSION", "0.27.0" if platform.system() == "Windows" else "0.27.1"
|
||||
)
|
||||
packages = setuptools.find_namespace_packages(include=["modelopt*"])
|
||||
package_dir = {"": "."}
|
||||
package_data = {"modelopt": ["**/*.h", "**/*.cpp", "**/*.cu"]}
|
||||
@@ -54,7 +57,8 @@ optional_deps = {
|
||||
# Onnxruntime 1.20+ is not supported on Python 3.9
|
||||
"onnxruntime~=1.18.1 ; python_version < '3.10'",
|
||||
"onnxruntime~=1.20.1 ; python_version >= '3.10' and (platform_machine == 'aarch64' or platform_system == 'Darwin')", # noqa: E501
|
||||
"onnxruntime-gpu~=1.20.1 ; python_version >= '3.10' and platform_machine != 'aarch64' and platform_system != 'Darwin'", # noqa: E501
|
||||
"onnxruntime-gpu~=1.20.1 ; python_version >= '3.10' and platform_machine != 'aarch64' and platform_system != 'Darwin' and platform_system != 'Windows'", # noqa: E501
|
||||
"onnxruntime-directml==1.20.0; python_version >= '3.10' and platform_system == 'Windows'",
|
||||
"onnxsim ; python_version < '3.12' and platform_machine != 'aarch64'",
|
||||
],
|
||||
"torch": [
|
||||
|
||||
@@ -27,9 +27,12 @@ pushd $script_dir/../../examples/llm_distill
|
||||
LLAMA_PATH=/tmp/tiny-random-Llama
|
||||
$script_dir/setup_tiny_llama.sh $LLAMA_PATH
|
||||
|
||||
SAVE_PATH=/tmp/llm_distill_test_output
|
||||
|
||||
accelerate launch --multi_gpu --mixed_precision bf16 main.py \
|
||||
--teacher_name_or_path $LLAMA_PATH \
|
||||
--student_name_or_path $LLAMA_PATH \
|
||||
--output_dir $SAVE_PATH \
|
||||
--logging_steps 5 \
|
||||
--max_steps 10 \
|
||||
--max_seq_length 1024 \
|
||||
|
||||
Reference in New Issue
Block a user