Separate arm64 docker builds from amd64 (#5184)

As a followup after: pytorch/pytorch#125617
Add validation_runner output param to know what validation runner to use
Test:
```
python tools/scripts/generate_docker_release_matrix.py
{"include": [{"cuda": "11.8", "cuda_full_version": "11.8.0", "cudnn_version": "8", "image_type": "runtime", "docker": "ghcr.io/pytorch/pytorch-nightly:2.4.0.dev20240507-cuda11.8-cudnn8-runtime", "platform": "linux/amd64", "validation_runner": "linux.g5.4xlarge.nvidia.gpu"}, {"cuda": "11.8", "cuda_full_version": "11.8.0", "cudnn_version": "8", "image_type": "devel", "docker": "ghcr.io/pytorch/pytorch-nightly:2.4.0.dev20240507-cuda11.8-cudnn8-devel", "platform": "linux/amd64", "validation_runner": "linux.g5.4xlarge.nvidia.gpu"}, {"cuda": "12.1", "cuda_full_version": "12.1.1", "cudnn_version": "8", "image_type": "runtime", "docker": "ghcr.io/pytorch/pytorch-nightly:2.4.0.dev20240507-cuda12.1-cudnn8-runtime", "platform": "linux/amd64", "validation_runner": "linux.g5.4xlarge.nvidia.gpu"}, {"cuda": "12.1", "cuda_full_version": "12.1.1", "cudnn_version": "8", "image_type": "devel", "docker": "ghcr.io/pytorch/pytorch-nightly:2.4.0.dev20240507-cuda12.1-cudnn8-devel", "platform": "linux/amd64", "validation_runner": "linux.g5.4xlarge.nvidia.gpu"}, {"cuda": "12.4", "cuda_full_version": "12.4.0", "cudnn_version": "8", "image_type": "runtime", "docker": "ghcr.io/pytorch/pytorch-nightly:2.4.0.dev20240507-cuda12.4-cudnn8-runtime", "platform": "linux/amd64", "validation_runner": "linux.g5.4xlarge.nvidia.gpu"}, {"cuda": "12.4", "cuda_full_version": "12.4.0", "cudnn_version": "8", "image_type": "devel", "docker": "ghcr.io/pytorch/pytorch-nightly:2.4.0.dev20240507-cuda12.4-cudnn8-devel", "platform": "linux/amd64", "validation_runner": "linux.g5.4xlarge.nvidia.gpu"}, {"cuda": "cpu", "cuda_full_version": "", "cudnn_version": "", "image_type": "runtime", "docker": "ghcr.io/pytorch/pytorch-nightly:2.4.0.dev20240507-runtime", "platform": "linux/arm64", "validation_runner": "linux.arm64.2xlarge"}]}
```

Author: atalman

tools/scripts/generate_docker_release_matrix.py

@@ -26,18 +26,17 @@
 def generate_docker_matrix(channel: str) -> Dict[str, List[Dict[str, str]]]:
 
     ret: List[Dict[str, str]] = []
+    prefix = "ghcr.io/pytorch/pytorch"
+    docker_image_version = ""
+    if channel == "release":
+        docker_image_version = f"{prefix}:{generate_binary_build_matrix.CURRENT_STABLE_VERSION}"
+    elif channel == "test":
+        docker_image_version = f"{prefix}-test:{generate_binary_build_matrix.CURRENT_CANDIDATE_VERSION}"
+    else:
+        docker_image_version = f"{prefix}-nightly:{generate_binary_build_matrix.CURRENT_NIGHTLY_VERSION}.dev{datetime.today().strftime('%Y%m%d')}"
+
     for cuda in generate_binary_build_matrix.CUDA_ARCHES_DICT[channel]:
         version = generate_binary_build_matrix.CUDA_CUDDN_VERSIONS[cuda]
-
-        prefix = "ghcr.io/pytorch/pytorch"
-        docker_image_version = ""
-        if channel == "release":
-            docker_image_version = f"{prefix}:{generate_binary_build_matrix.CURRENT_STABLE_VERSION}"
-        elif channel == "test":
-            docker_image_version = f"{prefix}-test:{generate_binary_build_matrix.CURRENT_CANDIDATE_VERSION}"
-        else:
-            docker_image_version = f"{prefix}-nightly:{generate_binary_build_matrix.CURRENT_NIGHTLY_VERSION}.dev{datetime.today().strftime('%Y%m%d')}"
-
         for image in DOCKER_IMAGE_TYPES:
             ret.append(
                 {
@@ -46,9 +45,22 @@ def generate_docker_matrix(channel: str) -> Dict[str, List[Dict[str, str]]]:
                     "cudnn_version": version["cudnn"],
                     "image_type": image,
                     "docker": f"{docker_image_version}-cuda{cuda}-cudnn{version['cudnn']}-{image}",
-                    "platform": "linux/arm64,linux/amd64",
+                    "platform": "linux/amd64",
+                    "validation_runner": generate_binary_build_matrix.LINUX_GPU_RUNNER,
                 }
             )
+
+    ret.append(
+        {
+            "cuda": "cpu",
+            "cuda_full_version": "",
+            "cudnn_version": "",
+            "image_type": "runtime",
+            "docker": f"{docker_image_version}-runtime",
+            "platform": "linux/arm64",
+            "validation_runner": generate_binary_build_matrix.LINUX_AARCH64_RUNNER,
+        }
+    )
     return {"include": ret}