Riyad Islam
c3f6cef87e
## What does this PR do?
**Type of change:** Bug <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->
**Overview:** `get_onnx_bytes` api is error prone as it returns only the
protobuf info. If model has any external data, they get discarded! We
have to use `get_onnx_bytes_and_metadata` and provide example for users
to correctly write ONNX model to disk.
## Usage
This would be proper way to save ONNX model bytes with/without external
data.
```python
onnx_bytes, _ = get_onnx_bytes_and_metadata(
model=model,
dummy_input=(input_tensor,),
weights_dtype=weights_dtype,
model_name=model_name,
)
onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes)
# Write the onnx model to the specified directory without cleaning it
onnx_bytes_obj.write_to_disk(os.path.dirname(onnx_save_path), clean_dir=False)
```
## Testing
N/A. Existing tests are modified.
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes<!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->
## Additional Information
https://nvbugspro.nvidia.com/bug/5618246/4
Signed-off-by: Riyad Islam <rislam@nvidia.com>
Chaining Multiple Optimizations Techniques
This directory demonstrates how to chain multiple optimization techniques like Pruning, Distillation, and Quantization together to achieve the best performance on a given model.
HuggingFace BERT Pruning + Distillation + Quantization
This example shows how to compress a Hugging Face Bert large model for Question Answering
using the combination of modelopt.torch.prune, modelopt.torch.distill and modelopt.torch.quantize. More specifically, we will:
- Prune the Bert large model to 50% FLOPs with GradNAS algorithm and fine-tune with distillation
- Quantize the fine-tuned model to INT8 precision with Post-Training Quantization (PTQ) and Quantize Aware Training (QAT) with distillation
- Export the quantized model to ONNX format for deployment with TensorRT
The main python file is bert_prune_distill_quantize.py and scripts for running it for all 3 steps are available in the scripts directory.
NOTE: This example has been tested on 8 x 24GB A5000 GPUs with PyTorch 2.4 and CUDA 12.4. It takes about 2 hours to complete all the stages of the optimization. Most of the time is spent on fine-tuning and QAT.
Pre-requisites
Install Model Optimizer with optional torch and huggingface dependencies:
pip install "nvidia-modelopt[hf]"
Running the example
To run the example, execute the following scripts in order:
-
First we prune the Bert large model to 50% FLOPs with GradNAS algorithm. Then, we fine-tune the pruned model with distillation from unpruned teacher model to recover 99+% of the initial F1 score (93.15). We recommend using multiple GPUs for fine-tuning. Note that we use more epochs for fine-tuning, which is different from the 2 epochs used originally in fine-tuning Bert without distillation since distillation requires more epochs to converge but achieves much better results.
bash scripts/1_prune.sh -
Quantize the fine-tuned model to INT8 precision and run calibration (PTQ). Note that PTQ will result in a slight drop in F1 score but we will be able to recover the F1 score with QAT. We run QAT with distillation as well from unpruned teacher model.
bash scripts/2_int8_quantize.sh -
Export the quantized model to ONNX format for deployment with TensorRT.
bash scripts/3_onnx_export.sh