[ROCm][CI] Re-enable ROCm CI & make CI changes as per new CI design

.github/pytorch-probot.yml

@@ -1,5 +1,6 @@
 ciflow_push_tags:
   - ciflow/b200
+  - ciflow/rocm
   - ciflow/h100.8
   - ciflow/rl
 labeler_config: labeler.yml

.github/workflows/integration_test_features_rocm.yaml → .github/workflows/integration_test_rocm.yaml

@@ -1,8 +1,10 @@
-name: Integration Tests - Features (ROCm)
+name: Integration Tests (ROCm)
 
 on:
   push:
     branches: [ main ]
+    tags:
+      - ciflow/rocm/*
     paths-ignore:
       - 'torchtitan/experiments/**'
 
@@ -11,7 +13,7 @@ permissions:
   contents: read
 
 jobs:
-  features:
+  rocm:
     uses: ./.github/workflows/integration_test.yaml
     with:
       gpu-arch: rocm

.github/workflows/set-matrix.yaml

@@ -42,7 +42,7 @@ jobs:
       IS_SCHEDULE:  ${{ github.event_name == 'schedule' }}
 
       # Default CUDA & ROCm runners
-      DEFAULT_ROCM_RUNNER: linux.rocm.gpu.ecosystem.gfx950.8
+      DEFAULT_ROCM_RUNNER: linux.rocm.gpu.ecosystem.mi350.8
       DEFAULT_CUDA_RUNNER: linux.g5.48xlarge.nvidia.gpu
 
       # Input CUDA & ROCm runners
@@ -78,46 +78,38 @@ jobs:
           EOF
           )"
 
-          # Temporarily disable ROCm CI while the ROCm jobs are red.
-          # Re-enable by removing this override and uncommenting the original
-          # ROCm matrix block below.
-          # ROCm remains disabled. When enabled, it should run from a separate
-          # AMD workflow triggered only by pushes to main.
-          ROCM_MATRIX=""
-          FULL_MATRIX="$CUDA_MATRIX"
-
-          # Original ROCm matrix definition:
-          # if [[ "${{ inputs.is-experimental }}" == "true" ]]; then
-          #   ROCM_MATRIX=""
-          #   FULL_MATRIX="$CUDA_MATRIX"
-          # else
-          #   ROCM_MATRIX="$(cat <<EOF
-          # {
-          #   "name": "rocm",
-          #   "runner": "${ROCM_RUNNER}",
-          #   "gpu-arch-type": "rocm",
-          #   "gpu-arch-version": "7.2",
-          #   "docker-image": "torchtitan-rocm-ubuntu-22.04-clang12",
-          #   "index-url": "https://download.pytorch.org/whl/nightly/rocm7.2",
-          #   "torch-version": ""
-          # }
-          # EOF
-          # )"
-          #   FULL_MATRIX="$CUDA_MATRIX,$ROCM_MATRIX"
-          # fi
+          if [[ "${{ inputs.is-experimental }}" == "true" ]]; then
+            ROCM_MATRIX=""
+            FULL_MATRIX="$CUDA_MATRIX"
+          else
+            ROCM_MATRIX="$(cat <<EOF
+          {
+            "name": "rocm",
+            "runner": "${ROCM_RUNNER}",
+            "gpu-arch-type": "rocm",
+            "gpu-arch-version": "10.0",
+            "docker-image": "torchtitan-rocm-ubuntu-22.04-clang12",
+            "index-url": "https://download.pytorch.org/whl/nightly/rocm10.0",
+            "torch-version": ""
+          }
+          EOF
+            )"
+            FULL_MATRIX="$CUDA_MATRIX,$ROCM_MATRIX"
+          fi
 
           # Use default value as 'false' for unset environment variables
           IS_MAIN_PUSH="${IS_MAIN_PUSH:-false}"
           IS_SCHEDULE="${IS_SCHEDULE:-false}"
 
-          # Decide which matrix entries to include based on event type
-          # Runs CUDA and ROCm for push to main and cron schedules.
-          if [[ ("$IS_MAIN_PUSH" == "true" || "$IS_SCHEDULE" == "true") ]]; then
+          # Decide which matrix entries to include based on event type.
+          # Runs CUDA and ROCm for push to main, cron schedules, or when
+          # gpu-arch is explicitly set to rocm (e.g. ciflow/rocm/* tag runs).
+          # Runs CUDA only as default, including pull requests.
+          GPU_ARCH_INPUT="${{ inputs.gpu-arch }}"
+          if [[ "$IS_MAIN_PUSH" == "true" || "$IS_SCHEDULE" == "true" || "${GPU_ARCH_INPUT}" == "rocm" ]]; then
             cat > matrix.json <<JSON
           {"include": [$FULL_MATRIX]}
           JSON
-
-          # Runs CUDA only as default, including pull requests.
           else
             cat > matrix.json <<JSON
           {"include": [$CUDA_MATRIX]}
@@ -130,15 +122,14 @@ jobs:
           # must stay on v2 because v3 does not authenticate ROCm runners to the
           # private ECR). 'all' (the default) keeps both for the workflows that
           # run a single combined linux_job_v2 job.
-          GPU_ARCH="${{ inputs.gpu-arch }}"
-          case "${GPU_ARCH:-all}" in
+          case "${GPU_ARCH_INPUT:-all}" in
             all) ;;
             cuda|rocm)
-              jq -c --arg arch "$GPU_ARCH" '{include: [.include[] | select(.["gpu-arch-type"] == $arch)]}' matrix.json > matrix.filtered.json
+              jq -c --arg arch "$GPU_ARCH_INPUT" '{include: [.include[] | select(.["gpu-arch-type"] == $arch)]}' matrix.json > matrix.filtered.json
               mv matrix.filtered.json matrix.json
               ;;
             *)
-              echo "::error::Invalid gpu-arch '${GPU_ARCH}'. Expected one of: all, cuda, rocm." >&2
+              echo "::error::Invalid gpu-arch '${GPU_ARCH_INPUT}'. Expected one of: all, cuda, rocm." >&2
               exit 1
               ;;
           esac

README.md

@@ -54,6 +54,7 @@ We look forward to your contributions!
 | --- | --- | --- |
 | CPU | - | [![CPU Unit Test](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_cpu.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_cpu.yaml?query=branch%3Amain) |
 | NVIDIA GPU | [![Integration Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test.yaml?query=branch%3Amain) [![H100 Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_h100.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_h100.yaml?query=branch%3Amain) | [![GPU Unit Tests](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_gpu.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_gpu.yaml?query=branch%3Amain) |
+| AMD GPU (ROCm) | [![ROCm Integration Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_rocm.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_rocm.yaml?query=branch%3Amain) | - |
 
 ## Llama 3.1 training