Promote aimet_torch.onnx.export to production API

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
Co-authored-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>

Author: 2 people

Docs/apiref/torch/index.rst

@@ -49,7 +49,7 @@ AIMET quantization for PyTorch models provides the following functionality.
 - :ref:`aimet_torch.quantsim <apiref-torch-quantsim>`
 - :ref:`aimet_torch.nn <apiref-torch-nn>`
 - :ref:`aimet_torch.quantization <apiref-torch-quantization>`
-- :ref:`aimet_torch.onnx (beta) <apiref-torch-onnx>`
+- :ref:`aimet_torch.onnx <apiref-torch-onnx>`
 - :ref:`aimet_torch.adaround <apiref-torch-adaround>`
 - :ref:`aimet_torch.seq_mse <apiref-torch-seq-mse>`
 - :ref:`aimet_torch.experimental.adascale <apiref-torch-adascale>`

Docs/apiref/torch/onnx.rst

@@ -1,7 +1,7 @@
 .. _apiref-torch-onnx:
 
-##############################
-aimet_torch.onnx.export (beta)
-##############################
+#######################
+aimet_torch.onnx.export
+#######################
 
 .. autofunction:: aimet_torch.onnx.export

TrainingExtensions/torch/src/python/aimet_torch/onnx.py

@@ -99,10 +99,10 @@ def export(
 
 
     .. note::
-        Unlike `torch.onnx.export()`, this function allows up to opset 21.
-        to support 4/16-bit quantization only available in opset 21.
-        However, exporting to opset 21 is a beta feature and not fully stable yet.
-        For robustness, opset 20 or lower is recommended whenever possible.
+        For robustness, onnx >=1.19 is highly recommended with this API,
+        especially when exporting large models (>2GB).
+        This is due to a known bug in onnx <1.19 version converter.
+        For more information, see https://github.com/onnx/onnx/issues/6529
 
     .. note::
         Dynamo-based export (`dynamo=True`) is not supported yet

TrainingExtensions/torch/test/python/v2/experimental/test_onnx.py

@@ -1208,17 +1208,14 @@ def large_model():
     )
 
 
+@pytest.fixture
+def tmp_path():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield pathlib.Path(tmpdir).resolve()
+
+
 @torch.no_grad()
-@pytest.mark.parametrize(
-    "opset_version",
-    [
-        19,
-        # NOTE: Currently fails because onnx version converter
-        # has a bug with large models. This bug is expected to be fixed in onnx 1.19.
-        # TODO (kyunggeu): Uncomment this when onnx 1.19 is released
-        # 21, TODO: Not supported yet
-    ],
-)
+@pytest.mark.parametrize("opset_version", [19, 21])
 @pytest.mark.parametrize("prequantize_constants", [False, True])
 def test_export_large_model(
     large_model: torch.nn.Module,