Migrate onnxrt RTN WOQ to 3.x API (#1544)

Signed-off-by: yuwenzho <yuwen.zhou@intel.com>
Signed-off-by: chensuyue <suyue.chen@intel.com>
Signed-off-by: xin3he <xin3.he@intel.com>

.azure-pipelines/scripts/install_nc.sh

@@ -5,15 +5,19 @@ cd /neural-compressor
 if [[ $1 = *"3x_pt" ]]; then
     python -m pip install --no-cache-dir -r requirements_pt.txt
     python setup.py pt bdist_wheel
-    pip install dist/neural_compressor*.whl
+    pip install dist/neural_compressor*.whl --force-reinstall
 elif [[ $1 = *"3x_tf" ]]; then
     python -m pip install --no-cache-dir -r requirements_tf.txt
     python setup.py tf bdist_wheel
-    pip install dist/neural_compressor*.whl
+    pip install dist/neural_compressor*.whl --force-reinstall
+elif [[ $1 = *"3x_ort" ]]; then
+    python -m pip install --no-cache-dir -r requirements_ort.txt
+    python setup.py ort bdist_wheel
+    pip install dist/neural_compressor*.whl --force-reinstall
 else
     python -m pip install --no-cache-dir -r requirements.txt
     python setup.py 2x bdist_wheel
-    pip install dist/neural_compressor*.whl
+    pip install dist/neural_compressor*.whl --force-reinstall
 fi
 
 echo -e "\n pip list after install Neural Compressor ... "

.azure-pipelines/scripts/ut/3x/collect_log_3x.sh

@@ -1,5 +1,6 @@
 source /neural-compressor/.azure-pipelines/scripts/change_color.sh
 
+set -xe
 pip install coverage
 export COVERAGE_RCFILE=/neural-compressor/.azure-pipelines/scripts/ut/3x/coverage.${1}
 coverage_log="/neural-compressor/log_dir/coverage_log"
@@ -22,7 +23,7 @@ cd /neural-compressor
 git config --global --add safe.directory /neural-compressor
 git fetch
 git checkout master
-echo y | pip uninstall neural-compressor
+echo y | pip uninstall neural_compressor_${1}
 cd /neural-compressor/.azure-pipelines/scripts && bash install_nc.sh ${1}
 
 coverage erase

.azure-pipelines/scripts/ut/3x/coverage.3x_ort

@@ -0,0 +1,15 @@
+[run]
+branch = True
+
+[report]
+include =
+ */neural_compressor/common/*
+ */neural_compressor/onnxrt/*
+exclude_lines =
+ pragma: no cover
+ raise NotImplementedError
+ raise TypeError
+ if self.device == "gpu":
+ if device == "gpu":
+ except ImportError:
+ except Exception as e:

.azure-pipelines/scripts/ut/3x/run_3x_ort.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+python -c "import neural_compressor as nc"
+test_case="run 3x ONNXRT"
+echo "${test_case}"
+
+# install requirements
+echo "set up UT env..."
+pip install -r /neural-compressor/test/3x/onnxrt/requirements.txt
+pip install coverage
+pip list
+
+export COVERAGE_RCFILE=/neural-compressor/.azure-pipelines/scripts/ut/3x/coverage.3x_ort
+inc_path=$(python -c 'import neural_compressor; print(neural_compressor.__path__[0])')
+cd /neural-compressor/test || exit 1
+find ./3x/onnxrt/* -name "test*.py" | sed 's,\.\/,coverage run --source='"${inc_path}"' --append ,g' | sed 's/$/ --verbose/'> run.sh
+
+LOG_DIR=/neural-compressor/log_dir
+mkdir -p ${LOG_DIR}
+ut_log_name=${LOG_DIR}/ut_3x_ort.log
+
+echo "cat run.sh..."
+sort run.sh -o run.sh
+cat run.sh | tee ${ut_log_name}
+echo "------UT start-------"
+bash -x run.sh 2>&1 | tee -a ${ut_log_name}
+cp .coverage ${LOG_DIR}/.coverage
+
+echo "------UT end -------"
+
+if [ $(grep -c "FAILED" ${ut_log_name}) != 0 ] || [ $(grep -c "core dumped" ${ut_log_name}) != 0 ] || [ $(grep -c "ModuleNotFoundError:" ${ut_log_name}) != 0 ] || [ $(grep -c "OK" ${ut_log_name}) == 0 ];then
+    echo "Find errors in UT test, please check the output..."
+    exit 1
+fi
+echo "UT finished successfully! "

.azure-pipelines/scripts/ut/3x/run_3x_pt.sh

@@ -5,9 +5,8 @@ echo "${test_case}"
 
 # install requirements
 echo "set up UT env..."
-pip install transformers
+pip install -r /neural-compressor/test/3x/torch/requirements.txt
 pip install coverage
-pip install pytest
 pip list
 
 export COVERAGE_RCFILE=/neural-compressor/.azure-pipelines/scripts/ut/3x/coverage.3x_pt

.azure-pipelines/scripts/ut/3x/run_3x_tf.sh

@@ -5,8 +5,8 @@ echo "${test_case}"
 
 # install requirements
 echo "set up UT env..."
+pip install -r /neural-compressor/test/3x/tensorflow/requirements.txt
 pip install coverage
-pip install pytest
 pip list
 
 export COVERAGE_RCFILE=/neural-compressor/.azure-pipelines/scripts/ut/3x/coverage.3x_tf

.azure-pipelines/ut-3x-ort.yml

@@ -0,0 +1,106 @@
+trigger: none
+
+pr:
+  autoCancel: true
+  drafts: false
+  branches:
+    include:
+      - master
+  paths:
+    include:
+      - neural_compressor/common
+      - neural_compressor/onnxrt
+      - test/3x/onnxrt
+      - setup.py
+      - requirements_ort.txt
+
+pool: ICX-16C
+
+variables:
+  IMAGE_NAME: "neural-compressor"
+  IMAGE_TAG: "py310"
+  UPLOAD_PATH: $(Build.SourcesDirectory)/log_dir
+  DOWNLOAD_PATH: $(Build.SourcesDirectory)/log_dir
+  ARTIFACT_NAME: "UT_coverage_report_3x_ort"
+  REPO: $(Build.Repository.Uri)
+
+stages:
+  - stage: ONNXRT
+    displayName: Unit Test 3x ONNXRT
+    dependsOn: []
+    jobs:
+      - job:
+        displayName: Unit Test 3x ONNXRT
+        steps:
+          - template: template/ut-template.yml
+            parameters:
+              dockerConfigName: "commonDockerConfig"
+              utScriptFileName: "3x/run_3x_ort"
+              uploadPath: $(UPLOAD_PATH)
+              utArtifact: "ut_coverage_3x"
+
+
+  - stage: ONNXRT_baseline
+    displayName: Unit Test 3x ONNXRT baseline
+    dependsOn: []
+    jobs:
+      - job:
+        displayName: Unit Test 3x ONNXRT baseline
+        steps:
+          - template: template/ut-template.yml
+            parameters:
+              dockerConfigName: "gitCloneDockerConfig"
+              utScriptFileName: "3x/run_3x_ort"
+              uploadPath: $(UPLOAD_PATH)
+              utArtifact: "ut_coverage_3x_baseline"
+              repo: $(REPO)
+
+  - stage: Coverage
+    displayName: "Coverage Combine"
+    pool:
+      vmImage: "ubuntu-latest"
+    dependsOn: [ONNXRT, ONNXRT_baseline]
+    jobs:
+      - job: CollectDatafiles
+        steps:
+          - script: |
+              if [[ ! $(docker images | grep -i ${IMAGE_NAME}:${IMAGE_TAG}) ]]; then
+                docker build -f ${BUILD_SOURCESDIRECTORY}/.azure-pipelines/docker/Dockerfile.devel -t ${IMAGE_NAME}:${IMAGE_TAG} .
+              fi
+              docker images | grep -i ${IMAGE_NAME}
+              if [[ $? -ne 0 ]]; then
+                echo "NO Such Repo"
+                exit 1
+              fi
+            displayName: "Build develop docker image"
+
+          - task: DownloadPipelineArtifact@2
+            inputs:
+              artifact:
+              path: $(DOWNLOAD_PATH)
+
+          - script: |
+              echo "--- create container ---"
+              docker run -d -it --name="collectLogs"  -v ${BUILD_SOURCESDIRECTORY}:/neural-compressor  ${IMAGE_NAME}:${IMAGE_TAG} /bin/bash
+              echo "--- docker ps ---"
+              docker ps
+              echo "--- collect logs ---"
+              docker exec collectLogs /bin/bash  +x -c "cd /neural-compressor/.azure-pipelines/scripts \
+              && bash install_nc.sh 3x_ort \
+              && bash ut/3x/collect_log_3x.sh 3x_ort"
+            displayName: "Collect UT Coverage"
+
+          - task: PublishPipelineArtifact@1
+            condition: succeededOrFailed()
+            inputs:
+              targetPath: $(UPLOAD_PATH)
+              artifact: $(ARTIFACT_NAME)
+              publishLocation: "pipeline"
+
+          - task: Bash@3
+            condition: always()
+            inputs:
+              targetType: "inline"
+              script: |
+                docker exec collectLogs bash -c "rm -fr /neural-compressor/* && rm -fr /neural-compressor/.* || true"
+            displayName: "Docker clean up"

.azure-pipelines/ut-3x-pt.yml

@@ -88,7 +88,7 @@ stages:
               docker exec collectLogs /bin/bash  +x -c "cd /neural-compressor/.azure-pipelines/scripts \
               && bash install_nc.sh 3x_pt \
               && bash ut/3x/collect_log_3x.sh 3x_pt"
-            displayName: "collect logs"
+            displayName: "Collect UT Coverage"
 
           - task: PublishPipelineArtifact@1
             condition: succeededOrFailed()

.azure-pipelines/ut-3x-tf.yml

@@ -88,7 +88,7 @@ stages:
               docker exec collectLogs /bin/bash  +x -c "cd /neural-compressor/.azure-pipelines/scripts \
               && bash install_nc.sh 3x_tf \
               && bash ut/3x/collect_log_3x.sh 3x_tf"
-            displayName: "collect logs"
+            displayName: "Collect UT Coverage"
 
           - task: PublishPipelineArtifact@1
             condition: succeededOrFailed()

.azure-pipelines/ut-basic-no-cover.yml

@@ -19,6 +19,7 @@ pr:
       - neural_compressor/common
       - neural_compressor/torch
       - neural_compressor/tensorflow
+      - neural_compressor/onnxrt
 
 pool: ICX-16C
 

.azure-pipelines/ut-basic.yml

@@ -19,6 +19,7 @@ pr:
       - neural_compressor/common
       - neural_compressor/torch
       - neural_compressor/tensorflow
+      - neural_compressor/onnxrt
 
 pool: ICX-16C
 
@@ -257,7 +258,7 @@ stages:
               docker exec collectLogs /bin/bash  +x -c "cd /neural-compressor/.azure-pipelines/scripts \
               && bash install_nc.sh \
               && bash ut/collect_log.sh"
-            displayName: "collect logs"
+            displayName: "Collect UT Coverage"
 
           - task: PublishPipelineArtifact@1
             condition: succeededOrFailed()

neural_compressor/onnxrt/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2023 Intel Corporation
+#
+# 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.
+
+from neural_compressor.onnxrt.utils.utility import register_algo
+from neural_compressor.onnxrt.algorithms import rtn_quantize_entry
+
+from neural_compressor.onnxrt.quantization import (
+    RTNConfig,
+    get_default_rtn_config,
+)

neural_compressor/onnxrt/algorithms/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# 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.
+
+
+from neural_compressor.onnxrt.algorithms.weight_only.algo_entry import rtn_quantize_entry

neural_compressor/onnxrt/algorithms/weight_only/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# 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.

neural_compressor/onnxrt/algorithms/weight_only/algo_entry.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2023 Intel Corporation
+#
+# 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.
+
+
+from pathlib import Path
+from typing import Dict, Tuple, Union
+
+import onnx
+
+from neural_compressor.common.logger import Logger
+from neural_compressor.common.utility import RTN
+from neural_compressor.onnxrt.quantization.config import RTNConfig
+from neural_compressor.onnxrt.utils.utility import register_algo
+
+logger = Logger().get_logger()
+
+
+###################### RTN Algo Entry ##################################
+@register_algo(name=RTN)
+def rtn_quantize_entry(model: Union[Path, str], quant_config: RTNConfig, *args, **kwargs) -> onnx.ModelProto:
+    """The main entry to apply rtn quantization."""
+    from neural_compressor.onnxrt.algorithms.weight_only.rtn import apply_rtn_on_model
+
+    # map config to each op
+    model_info = quant_config.get_model_info(model=model)
+    configs_mapping = quant_config.to_config_mapping(model_info=model_info)
+    logger.debug(configs_mapping)
+    model = apply_rtn_on_model(model, configs_mapping)
+    return model