[ROCm] Enable build option for autograd (#11945)

* add autograd build option

* disable UTs

* disable UTs

* UT-step1

* UT-step1

* UT-step2

* UT-step2

* UT-step2

* UT-step2

* UT-step2

* UT-step2

* Fix UTs

* increase shm

* code clean up

Co-authored-by: Ethan Tao <ettao@microsoft.com@orttrainingdev7.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -1318,7 +1318,7 @@ def test_gradient_correctness_reducesum(dim, keepdim):
         pt_prediction = run_step(pt_model, pt_input)
         ort_prediction = run_step(ort_model, ort_input)
 
-        _test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
+        _test_helpers.assert_values_are_close(ort_prediction, pt_prediction, atol=1e-5, rtol=1e-4)
         _test_helpers.assert_values_are_close(ort_input.grad, pt_input.grad)
 
 

orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py

@@ -351,7 +351,7 @@ def test_ortmodule_fallback_torch_model(is_training, fallback_enabled, matching_
 
                 ort_out = ort_model(x)
                 pt_out = pt_model(x)
-                _test_helpers.assert_values_are_close(ort_out, pt_out, rtol=0, atol=0)
+                _test_helpers.assert_values_are_close(ort_out, pt_out, rtol=1e-3, atol=1e-6)
             else:
                 with pytest.raises(_fallback.ORTModuleTorchModelException) as ex_info:
                     ort_model = ORTModule(pt_model)
@@ -525,7 +525,7 @@ def test_ortmodule_fallback_onnx_model__custom_autograd(
                     )
                 pt_out = pt_model(x.mm(w1)).mm(w2)
                 ort_out = ort_model(x.mm(w1)).mm(w2)
-                _test_helpers.assert_values_are_close(ort_out, pt_out, rtol=0, atol=0)
+                _test_helpers.assert_values_are_close(ort_out, pt_out, rtol=1e-03, atol=1e-04)
             else:
                 with pytest.raises(_fallback.ORTModuleONNXModelException) as ex_info:
                     _ = ort_model(x.mm(w1)).mm(w2)

tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml

@@ -19,7 +19,7 @@ jobs:
   container:
     image: onnxruntimecibuildenvironment.azurecr.io/rocm-ci-pipeline-env:rocm5.1.1
     endpoint: onnxruntimecibuildenvironmentforamd
-    options: --privileged -e HIP_VISIBLE_DEVICES --security-opt seccomp=unconfined --device=/dev/kfd --device=/dev/dri  --group-add $(video) --group-add $(render)
+    options: --privileged -e HIP_VISIBLE_DEVICES --security-opt seccomp=unconfined --shm-size=1024m --device=/dev/kfd --device=/dev/dri  --group-add $(video) --group-add $(render)
 
   steps:
   - checkout: self
@@ -49,6 +49,7 @@ jobs:
         python tools/ci_build/build.py \
         --config RelWithDebInfo \
         --enable_training \
+        --enable_training_torch_interop \
         --mpi_home /opt/ompi \
         --use_rocm \
         --rocm_version=5.1.1 \
@@ -288,8 +289,8 @@ jobs:
     inputs:
       script: |-
         python orttraining/tools/ci_test/run_batch_size_test.py \
-            --binary_dir build/RelWithDebInfo \
-            --model_root training_e2e_test_data/models \
+          --binary_dir build/RelWithDebInfo \
+          --model_root training_e2e_test_data/models \
           --gpu_sku MI100_32G
     displayName: 'Run C++ BERT-L batch size test'
     retryCountOnTaskFailure: 1
@@ -318,3 +319,48 @@ jobs:
     displayName: 'Run C++ BERT-L convergence test'
     retryCountOnTaskFailure: 1
     condition: and(succeededOrFailed(), eq(variables.onnxruntimeBuildSucceeded, 'true')) # ensure all tests are run when the build successed
+
+  - script: |
+      sudo apt-get update
+      sudo apt install -y cifs-utils
+    displayName: 'Install filesystems util'
+
+  - bash: tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/mnist" -d "/mnist"
+    displayName: 'Mount MNIST'
+    condition: succeededOrFailed()
+
+  - bash: tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/bert-data" -d "/bert_data"
+    displayName: 'Mount bert-data'
+    condition: succeededOrFailed()
+
+  - bash: tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/hf-models-cache" -d "/hf_models_cache"
+    displayName: 'Mount hf-models-cache'
+    condition: succeededOrFailed()
+
+
+  # Entry point for all ORTModule tests
+  # The onnxruntime folder is deleted in the build directory
+  # to enforce use of the onnxruntime wheel
+  # Uninstall torch 1.10 from the docker image above and install torch 1.11 before running tests.
+  #
+  # At this point, the best whl available with torch 1.11 (which supports autograd) is located in https://repo.radeon.com/rocm/manylinux/rocm-rel-5.1.3/ 
+  # which is a release for internal use. This release will be available via Torch's download page. We can either continue 
+  # to use the link below, or switch to the link from Torch's when it's published. Ideally, the AMD CI UTs above could also migrate to use
+  # the same whl below.   
+  - task: CmdLine@2
+    inputs:
+      script: |-
+        cd ./build/RelWithDebInfo
+        unset PYTHONPATH
+        pip uninstall -y torch torchvision
+        pip install torch==1.11.0 torchvision==0.12.0 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-5.1.3/
+        pip install wget pytorch_lightning==1.6.0
+        rm -rf onnxruntime
+        pip install ./dist/onnxruntime*.whl
+        python -m onnxruntime.training.ortmodule.torch_cpp_extensions.install
+        python orttraining_ortmodule_tests.py \
+          --mnist /mnist \
+          --bert_data /bert_data/hf_data/glue_data/CoLA/original/raw \
+          --transformers_cache /hf_models_cache/huggingface/transformers
+    displayName: 'Run orttraining_ortmodule_tests.py'
+    condition: succeededOrFailed()

tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-rocm.yml

@@ -100,143 +100,25 @@ stages:
                 /onnxruntime_src/tools/ci_build/build.py \
                   --config Release \
                   --use_rocm \
-                    --rocm_version=$(RocmVersion) \
-                    --rocm_home=/opt/rocm \
-                    --nccl_home=/opt/rocm \
+                  --rocm_version=$(RocmVersion) \
+                  --rocm_home=/opt/rocm \
+                  --nccl_home=/opt/rocm \
                   --update \
                   --parallel \
                   --build_dir /build \
                   --build \
                   --build_wheel \
                   --skip_tests \
-                  --enable_training
+                  --enable_training \
+                  --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=OFF \
+                  --enable_training_torch_interop
           workingDirectory: $(Build.SourcesDirectory)
         displayName: 'Build onnxruntime (in container)'
 
-      - script: |-
-          python3 orttraining/tools/ci_test/download_azure_blob_archive.py \
-            --azure_blob_url https://onnxruntimetestdata.blob.core.windows.net/training/onnxruntime_training_data.zip?snapshot=2020-06-15T23:17:35.8314853Z \
-            --target_dir $(Build.SourcesDirectory)/training_e2e_test_data \
-            --archive_sha256_digest B01C169B6550D1A0A6F1B4E2F34AE2A8714B52DBB70AC04DA85D371F691BDFF9
-        displayName: 'Download onnxruntime_training_data.zip data'
-
-      - script: |-
-          echo "Tests will run using HIP_VISIBLES_DEVICES=$HIP_VISIBLE_DEVICES"
-          video_gid=$(getent group | awk '/video/ {split($0,a,":"); print(a[3])}')
-          echo "Found video_gid=$video_gid; attempting to set as pipeline variable"
-          echo "##vso[task.setvariable variable=video]$video_gid"
-          render_gid=$(getent group | awk '/render/ {split($0,a,":"); print(a[3])}')
-          echo "Found render_gid=$render_gid; attempting to set as pipeline variable"
-          echo "##vso[task.setvariable variable=render]$render_gid"
-        displayName: 'Find video and render gid to be mapped into container'
-
-      - script: |-
-          echo "video=$video"
-          echo "render=$render"
-          docker run --rm \
-            --device=/dev/kfd \
-            --device=/dev/dri \
-            --group-add $(video) \
-            --group-add $(render) \
-            --privileged \
-            --ipc=host \
-            --network=host \
-            --cap-add=SYS_PTRACE \
-            --security-opt seccomp=unconfined \
-            --volume $(Build.SourcesDirectory):/onnxruntime_src \
-            --volume $(Build.BinariesDirectory):/build \
-            --workdir /build/Release \
-            --entrypoint /bin/bash \
-            -e HIP_VISIBLE_DEVICES \
-            -e NIGHTLY_BUILD \
-            -e BUILD_BUILDNUMBER \
-            --user onnxruntimedev \
-            onnxruntimetrainingrocmbuild-rocm$(RocmVersion) \
-               /onnxruntime_src/tools/ci_build/github/pai/pai_test_launcher.sh
-        displayName: 'Run onnxruntime unit tests (in container)'
-
-      - script: |-
-          docker run --rm \
-            --device=/dev/kfd \
-            --device=/dev/dri \
-            --group-add $(video) \
-            --group-add $(render) \
-            --privileged \
-            --ipc=host \
-            --network=host \
-            --cap-add=SYS_PTRACE \
-            --security-opt seccomp=unconfined \
-            --volume $(Build.SourcesDirectory):/onnxruntime_src \
-            --volume $(Build.BinariesDirectory):/build \
-            --workdir /onnxruntime_src \
-            --entrypoint $(PythonManylinuxDir)/bin/python3 \
-            -e HIP_VISIBLE_DEVICES \
-            -e NIGHTLY_BUILD \
-            -e BUILD_BUILDNUMBER \
-            --user onnxruntimedev \
-            onnxruntimetrainingrocmbuild-rocm$(RocmVersion) \
-              orttraining/tools/ci_test/run_batch_size_test.py \
-                --binary_dir /build/Release \
-                --model_root training_e2e_test_data/models \
-                --gpu_sku MI100_32G
-        displayName: 'Run C++ BERT-L batch size test (in container)'
-        condition: succeededOrFailed() # ensure all tests are run
-
-      - script: |-
-          docker run --rm \
-            --device=/dev/kfd \
-            --device=/dev/dri \
-            --group-add $(video) \
-            --group-add $(render) \
-            --privileged \
-            --ipc=host \
-            --network=host \
-            --cap-add=SYS_PTRACE \
-            --security-opt seccomp=unconfined \
-            --volume $(Build.SourcesDirectory):/onnxruntime_src \
-            --volume $(Build.BinariesDirectory):/build \
-            --workdir /onnxruntime_src \
-            --entrypoint $(PythonManylinuxDir)/bin/python3 \
-            -e HIP_VISIBLE_DEVICES \
-            -e NIGHTLY_BUILD \
-            -e BUILD_BUILDNUMBER \
-            --user onnxruntimedev \
-            onnxruntimetrainingrocmbuild-rocm$(RocmVersion) \
-              orttraining/tools/ci_test/run_bert_perf_test.py \
-                --binary_dir /build/Release \
-                --model_root training_e2e_test_data/models \
-                --training_data_root training_e2e_test_data/data \
-                --gpu_sku MI100_32G
-        displayName: 'Run C++ BERT-L performance test (in container)'
-        condition: succeededOrFailed() # ensure all tests are run
-
-      - script: |-
-          docker run --rm \
-            --device=/dev/kfd \
-            --device=/dev/dri \
-            --group-add $(video) \
-            --group-add $(render) \
-            --privileged \
-            --ipc=host \
-            --network=host \
-            --cap-add=SYS_PTRACE \
-            --security-opt seccomp=unconfined \
-            --volume $(Build.SourcesDirectory):/onnxruntime_src \
-            --volume $(Build.BinariesDirectory):/build \
-            --workdir /onnxruntime_src \
-            --entrypoint $(PythonManylinuxDir)/bin/python3 \
-            -e HIP_VISIBLE_DEVICES \
-            -e NIGHTLY_BUILD \
-            -e BUILD_BUILDNUMBER \
-            --user onnxruntimedev \
-            onnxruntimetrainingrocmbuild-rocm$(RocmVersion) \
-              orttraining/tools/ci_test/run_convergence_test.py \
-                --binary_dir /build/Release \
-                --model_root training_e2e_test_data/models \
-                --training_data_root training_e2e_test_data/data \
-                --gpu_sku MI100_32G
-        displayName: 'Run C++ BERT-L convergence test (in container)'
-        condition: succeededOrFailed() # ensure all tests are run
+      # All UTs were here are now covered in AMD CI - see orttraining-pai-ci-pipeline.yml
+      # This CI is mainly responsible for packaging. The uploaded whl could be used in the downstream CIs (if any).
+      # For example, docker image build (e.g., PTCA), reporting CI, etc. to further verify the whl as needed.  
+      # To view the UTs disabled from this CI - see https://github.com/microsoft/onnxruntime/pull/11945 for examples
 
       - task: CopyFiles@2
         displayName: 'Copy Python Wheel to: $(Build.ArtifactStagingDirectory)'