From 1b38c05544ed29bd7de4de79b93ffb4fe741ec56 Mon Sep 17 00:00:00 2001
From: kailums <109063327+kailums@users.noreply.github.com>
Date: Thu, 18 Jul 2024 14:50:01 +0800
Subject: [PATCH] change ci docker image to rocm6.1 (#21296)

### Description
<!-- Describe your changes. -->
There is a bug for kernel running on rocm6.0, so change ci docker image
to rocm6.1

For the torch installed in the docker image, change to rocm repo when it
is not 6.0 version.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
---
 ...-mi200.huggingface.bert-large-rocm6.1.json | 57 +++++++++++++++++++
 .../linux-migraphx-ci-pipeline.yml            |  2 +-
 .../orttraining-pai-ci-pipeline.yml           |  2 +-
 .../pai/rocm-ci-pipeline-env.Dockerfile       |  8 +--
 4 files changed, 61 insertions(+), 8 deletions(-)
 create mode 100644 orttraining/tools/ci_test/results/ci-mi200.huggingface.bert-large-rocm6.1.json

diff --git a/orttraining/tools/ci_test/results/ci-mi200.huggingface.bert-large-rocm6.1.json b/orttraining/tools/ci_test/results/ci-mi200.huggingface.bert-large-rocm6.1.json
new file mode 100644
index 0000000000..05fcf08cd3
--- /dev/null
+++ b/orttraining/tools/ci_test/results/ci-mi200.huggingface.bert-large-rocm6.1.json
@@ -0,0 +1,57 @@
+{
+    "steps": [
+        {
+            "step": 20,
+            "loss": 2.0136
+        },
+        {
+            "step": 40,
+            "loss": 1.8466
+        },
+        {
+            "step": 60,
+            "loss": 1.7525
+        },
+        {
+            "step": 80,
+            "loss": 1.6682
+        },
+        {
+            "step": 100,
+            "loss": 1.658
+        },
+        {
+            "step": 120,
+            "loss": 1.6749
+        },
+        {
+            "step": 140,
+            "loss": 1.6263
+        },
+        {
+            "step": 160,
+            "loss": 1.6828
+        },
+        {
+            "step": 180,
+            "loss": 1.6145
+        },
+        {
+            "step": 200,
+            "loss": 1.6197
+        },
+        {
+            "step": 220,
+            "loss": 1.6353
+        },
+        {
+            "step": 240,
+            "loss": 1.5266
+        },
+        {
+            "step": 260,
+            "loss": 1.5441
+        }
+    ],
+    "samples_per_second": 34.561
+}
diff --git a/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml
index f36cd9cfbf..6bf6324252 100644
--- a/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml
@@ -36,7 +36,7 @@ variables:
   - name: render
     value: 109
   - name: RocmVersion
-    value: 6.0
+    value: 6.1
   - name: RocmVersionPatchSuffix
     value: ".3"
 
diff --git a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
index 0010624526..0e1afdcc5b 100644
--- a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
@@ -25,7 +25,7 @@ variables:
   - name: render
     value: 109
   - name: RocmVersion
-    value: 6.0
+    value: 6.1
   - name: RocmVersionPatchSuffix
     value: ".3"
   - name: BuildConfig
diff --git a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
index b94826ae0e..bf21a65314 100644
--- a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
+++ b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
@@ -1,7 +1,7 @@
 # Refer to https://github.com/RadeonOpenCompute/ROCm-docker/blob/master/dev/Dockerfile-ubuntu-22.04-complete
 FROM ubuntu:22.04
 
-ARG ROCM_VERSION=6.0
+ARG ROCM_VERSION=6.1
 ARG AMDGPU_VERSION=${ROCM_VERSION}
 ARG APT_PREF='Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600'
 
@@ -77,11 +77,7 @@ RUN ln -sf /usr/lib/x86_64-linux-gnu/libstdc++.so.6 ${CONDA_ENVIRONMENT_PATH}/bi
 RUN export MAJOR=$(cut -d '.' -f 1 <<< "$ROCM_VERSION") && \
     export MINOR=$(cut -d '.' -f 2 <<< "$ROCM_VERSION") && \
     export PATCH=$(cut -d '.' -f 3 <<< "$ROCM_VERSION") && \
-    if (( MAJOR >= 6 )); then \
-        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm${MAJOR}.${MINOR} ; \
-    else \
-        pip install torch==2.0.1 torchvision==0.15.2 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-${MAJOR}.${MINOR}/ ; \
-    fi && \
+    pip install torch==2.1.2 torchvision==0.16.1 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-${MAJOR}.${MINOR}/ && \
     pip install torch-ort --no-dependencies
 
 ##### Install Cupy to decrease CPU utilization