[OMNIML-2244] Implement the ONNX quantization exporter for INT4 (#575)

## What does this PR do?

**Type of change:** 
New Feature

**Overview:** 
- Created an abstract parent class for ONNXQuantExporter
- Created child classes for individual precisions
- Implemented the INT4QuantExporter
- Removed quantize_weights_to_int4
- Added a method to quantize weights of the ONNX model to low precision

## Testing
```
python torch_quant_to_onnx.py --quantize_mode=int4_awq \
	--onnx_save_path=<onnx_path> \
```

## 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
- **Did you write any new necessary tests?**: No
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
No <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

---------

Signed-off-by: ajrasane <[email protected]>

Author: ajrasane

modelopt/onnx/export/__init__.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""ONNX export utilities."""
+
+__all__ = [
+    "FP8QuantExporter",
+    "INT4QuantExporter",
+    "INT8QuantExporter",
+    "MXFP8QuantExporter",
+    "NVFP4QuantExporter",
+    "ONNXQuantExporter",
+]
+
+from .base_exporter import ONNXQuantExporter
+from .fp8_exporter import FP8QuantExporter
+from .int4_exporter import INT4QuantExporter
+from .int8_exporter import INT8QuantExporter
+from .mxfp8_exporter import MXFP8QuantExporter
+from .nvfp4_exporter import NVFP4QuantExporter

modelopt/onnx/export/base_exporter.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Base class for ONNX quantizer exporters."""
+
+from abc import ABC, abstractmethod
+
+import onnx
+
+
+class ONNXQuantExporter(ABC):
+    """Base class for ONNX quantizer exporters."""
+
+    @classmethod
+    def process_model(cls, onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Processes the ONNX model."""
+        onnx_model = cls.pre_process(onnx_model)
+        onnx_model = cls.compute_scales(onnx_model)
+        onnx_model = cls.compress_weights(onnx_model)
+        onnx_model = cls.post_process(onnx_model)
+        return onnx_model
+
+    @staticmethod
+    @abstractmethod
+    def pre_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Pre-processes the ONNX model. Converts all DQ -> * -> op patterns to DQ -> op."""
+
+    @staticmethod
+    @abstractmethod
+    def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Computes the scales for the weights in the ONNX model."""
+
+    @staticmethod
+    @abstractmethod
+    def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Compresses the weights in the ONNX model."""
+
+    @staticmethod
+    @abstractmethod
+    def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Post-processes the ONNX model."""

modelopt/onnx/export/fp8_exporter.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""FP8 quantization exporter."""
+
+import onnx
+
+from .base_exporter import ONNXQuantExporter
+
+
+# TODO: Implement the FP8QuantExporter
+class FP8QuantExporter(ONNXQuantExporter):
+    """Exporter for FP8 quantization."""
+
+    @staticmethod
+    def pre_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Pre-processes the ONNX model for FP8 quantization."""
+
+    @staticmethod
+    def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Computes the scales for the weights in the ONNX model for FP8 quantization."""
+
+    @staticmethod
+    def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Compresses the weights in the ONNX model for FP8 quantization."""
+
+    @staticmethod
+    def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        """Post-processes the ONNX model for FP8 quantization."""