diff --git a/.github/scripts/Dockerfile.ci.deps b/.github/scripts/Dockerfile.ci.deps index 7161054f4..aff7f53c1 100644 --- a/.github/scripts/Dockerfile.ci.deps +++ b/.github/scripts/Dockerfile.ci.deps @@ -1,49 +1,99 @@ # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. +# +# TE CI deps image: Ubuntu 24.04 + one TheRock ROCm tarball + Python 3.12 venv. +# +# docker build -f .github/scripts/Dockerfile.ci.deps \ +# --build-arg GPU_ARCH=gfx942 \ +# -t te-ci-deps:rocm-7.12.0-ubuntu24.04-py312-gfx942 . +# +# docker build -f .github/scripts/Dockerfile.ci.deps \ +# --build-arg GPU_ARCH=gfx950 \ +# -t te-ci-deps:rocm-7.12.0-ubuntu24.04-py312-gfx950 . +# +# ROCm installer: https://raw.githubusercontent.com/ROCm/TheRock/main/dockerfiles/install_rocm_tarball.sh + +FROM ubuntu:24.04 + +ARG DEBIAN_FRONTEND=noninteractive +SHELL ["/bin/bash", "-euo", "pipefail", "-c"] + +ARG ROCM_VERSION=7.12.0 +ARG GPU_ARCH=gfx942 +ARG INSTALL_ROCM_TARBALL_SH_URL=https://raw.githubusercontent.com/ROCm/TheRock/main/dockerfiles/install_rocm_tarball.sh -## TE CI Dockerfile -ARG BASE_DOCKER=registry-sc-harbor.amd.com/framework/compute-rocm-rel-7.2:57_ubuntu22.04_py3.11_pytorch_release-2.8_08d38866 -FROM $BASE_DOCKER -WORKDIR / +# Map GPU_ARCH → AMD GPU wheel/tarball family (once). +RUN case "${GPU_ARCH}" in \ + gfx942) echo -n gfx94X-dcgpu > /etc/amd_gpu_family ;; \ + gfx950) echo -n gfx950-dcgpu > /etc/amd_gpu_family ;; \ + *) echo "GPU_ARCH must be gfx942 or gfx950 (got: ${GPU_ARCH})" >&2; exit 1 ;; \ + esac -# Updated git via git-core PPA -RUN apt-get update && apt-get install -y --no-install-recommends software-properties-common \ - && add-apt-repository ppa:git-core/ppa -y \ - && apt-get update \ - && apt-get install -y --no-install-recommends git vim \ +# Base OS packages +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates curl \ + git vim \ + build-essential cmake ninja-build pkg-config liblzma-dev \ + python3.12 python3.12-venv python3.12-dev python3-pip \ && rm -rf /var/lib/apt/lists/* -# Build arguments -ARG FA_VERSION=v2.8.1 -ARG ROCM_VERSION=7.2 -ARG JAX_VERSION=0.8.0 -ARG PYTHON_VERSION=311 -# AITER - Required for MXFP4 FP4 GEMM kernels. -ARG AITER_COMMIT=77455e3ecf4f0d28756afc452e914940c45b944b +# Native ROCm tarball → /opt/rocm +RUN AMD_GPU_FAMILY=$(cat /etc/amd_gpu_family) \ + && curl -fsSL -o /tmp/install_rocm_tarball.sh "${INSTALL_ROCM_TARBALL_SH_URL}" \ + && chmod +x /tmp/install_rocm_tarball.sh \ + && /tmp/install_rocm_tarball.sh "${ROCM_VERSION}" "${AMD_GPU_FAMILY}" stable \ + && rm -f /tmp/install_rocm_tarball.sh + +RUN python3.12 -m venv /opt/venv \ + && { echo ""; echo "# ROCm"; \ + echo 'export ROCM_PATH=/opt/rocm'; \ + echo 'export PATH="${ROCM_PATH}/bin:${PATH}"'; \ + echo 'if [ -n "${LD_LIBRARY_PATH:-}" ]; then export LD_LIBRARY_PATH="${ROCM_PATH}/lib:${LD_LIBRARY_PATH}"; else export LD_LIBRARY_PATH="${ROCM_PATH}/lib"; fi'; \ + } >> /opt/venv/bin/activate + +ENV GPU_ARCH=${GPU_ARCH} \ + ROCM_PATH=/opt/rocm \ + VIRTUAL_ENV=/opt/venv \ + PATH=/opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ + LD_LIBRARY_PATH=/opt/rocm/lib -RUN pip install setuptools wheel -RUN pip install ipython pytest fire pydantic pybind11 ninja pandas +RUN python -m pip install --upgrade pip setuptools wheel \ + && pip install --no-cache-dir ipython pytest fire pydantic pybind11 ninja pandas expecttest -# Install flash-attention -RUN git clone --branch ${FA_VERSION} --depth 1 https://github.com/Dao-AILab/flash-attention.git \ - && cd flash-attention \ - && GPU_ARCHS="gfx950;gfx942" FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE FLASH_ATTENTION_SKIP_CK_BUILD=FALSE python setup.py install \ - && cd .. +# Python ROCm SDK + torch / jax from https://repo.amd.com/rocm/whl// +RUN AMD_GPU_FAMILY=$(cat /etc/amd_gpu_family) \ + && W="https://repo.amd.com/rocm/whl/${AMD_GPU_FAMILY}" \ + && LIBS_PKG="rocm-sdk-libraries-$(echo "${AMD_GPU_FAMILY}" | tr '[:upper:]' '[:lower:]')" \ + && pip install --no-cache-dir \ + --extra-index-url "${W}" \ + "rocm-sdk-core==${ROCM_VERSION}" \ + "rocm-sdk-devel==${ROCM_VERSION}" \ + "${LIBS_PKG}==${ROCM_VERSION}" \ + "${W}/rocm-${ROCM_VERSION}.tar.gz" \ + "${W}/torch-2.10.0%2Brocm7.12.0-cp312-cp312-linux_x86_64.whl" \ + "${W}/torchvision-0.25.0%2Brocm7.12.0-cp312-cp312-linux_x86_64.whl" \ + "${W}/torchaudio-2.10.0%2Brocm7.12.0-cp312-cp312-linux_x86_64.whl" \ + "${W}/triton-3.6.0%2Brocm7.12.0-cp312-cp312-linux_x86_64.whl" \ + "${W}/jax_rocm7_pjrt-0.8.2%2Brocm7.12.0-py3-none-manylinux_2_28_x86_64.whl" \ + "${W}/jax_rocm7_plugin-0.8.2%2Brocm7.12.0-cp312-cp312-manylinux_2_28_x86_64.whl" \ + "${W}/jaxlib-0.8.2%2Brocm7.12.0-cp312-cp312-manylinux_2_27_x86_64.whl" \ + "jax==0.8.2" -# Install AITER -RUN git clone --no-checkout https://github.com/ROCm/aiter.git \ - && cd aiter \ - && git checkout ${AITER_COMMIT} \ +ARG FA_VERSION=v2.8.1 +RUN git clone --branch "${FA_VERSION}" --depth 1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention \ + && cd /tmp/flash-attention \ + && GPU_ARCHS="${GPU_ARCH}" FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE FLASH_ATTENTION_SKIP_CK_BUILD=FALSE \ + python setup.py install \ + && rm -rf /tmp/flash-attention + +ARG AITER_COMMIT=77455e3ecf4f0d28756afc452e914940c45b944b +RUN git clone --no-checkout https://github.com/ROCm/aiter.git /tmp/aiter \ + && cd /tmp/aiter \ + && git checkout "${AITER_COMMIT}" \ && git submodule update --init --recursive \ - && pip install . - -# Install JAX -RUN ROCM_MAJOR=$(echo "${ROCM_VERSION}" | cut -d. -f1) && pip install \ - https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/jax_rocm${ROCM_MAJOR}_pjrt-${JAX_VERSION}%2Brocm${ROCM_VERSION}.0-py3-none-manylinux_2_28_x86_64.whl \ - https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/jax_rocm${ROCM_MAJOR}_plugin-${JAX_VERSION}%2Brocm${ROCM_VERSION}.0-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-manylinux_2_28_x86_64.whl \ - jax==${JAX_VERSION} \ - https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/jaxlib-${JAX_VERSION}%2Brocm${ROCM_VERSION}.0-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-manylinux_2_27_x86_64.whl + && GPU_ARCHS="${GPU_ARCH}" pip install --no-build-isolation --no-cache-dir . \ + && rm -rf /tmp/aiter WORKDIR /workspace/ CMD ["/bin/bash"] diff --git a/.github/workflows/ci-deps-docker-publish.yml b/.github/workflows/ci-deps-docker-publish.yml index c989502e6..2735c4530 100644 --- a/.github/workflows/ci-deps-docker-publish.yml +++ b/.github/workflows/ci-deps-docker-publish.yml @@ -2,7 +2,16 @@ # # See LICENSE for license information. # -# Build .github/scripts/Dockerfile.ci.deps and push to Artifactory. +# Build .github/scripts/Dockerfile.ci.deps and push to Harbor. +# +# Tag convention (when image_tag is left empty): +# rocm--ubuntu24.04-py312- +# e.g. rocm-7.12.0-ubuntu24.04-py312-gfx942 +# +# Local builds: +# docker build -f .github/scripts/Dockerfile.ci.deps \ +# --build-arg GPU_ARCH=gfx942 \ +# -t registry-sc-harbor.amd.com/framework/te-ci:rocm-7.12.0-ubuntu24.04-py312-gfx942 . # # Required repository secrets: # ARTIFACTORY_DOCKER_USERNAME / ARTIFACTORY_DOCKER_PASSWORD — registry basic auth @@ -16,10 +25,23 @@ name: Publish CI deps Docker image on: workflow_dispatch: inputs: - image_tag: - description: "Image tag pushed to Artifactory (required)" + gpu_arch: + description: "GPU architecture for the CI deps image (Dockerfile GPU_ARCH build-arg)" required: true + type: choice + options: + - gfx942 + - gfx950 + rocm_version: + description: "ROCm version string (must match Dockerfile wheel pins when changed)" + required: false type: string + default: "7.12.0" + image_tag: + description: "Tag to push; leave empty for rocm--ubuntu24.04-py312-" + required: false + type: string + default: "" jobs: build-and-push: @@ -37,16 +59,20 @@ jobs: env: REGISTRY: ${{ vars.ARTIFACTORY_DOCKER_REGISTRY }} REPOSITORY: ${{ vars.ARTIFACTORY_CI_DEPS_REPOSITORY }} - IMAGE_TAG: ${{ inputs.image_tag }} + IMAGE_TAG_INPUT: ${{ inputs.image_tag }} + GPU_ARCH: ${{ inputs.gpu_arch }} + ROCM_VERSION: ${{ inputs.rocm_version }} run: | set -euo pipefail if [ -z "${REGISTRY}" ] || [ -z "${REPOSITORY}" ]; then echo "Set repository variables ARTIFACTORY_DOCKER_REGISTRY and ARTIFACTORY_CI_DEPS_REPOSITORY." >&2 exit 1 fi - if [ -z "${IMAGE_TAG}" ]; then - echo "image_tag must be non-empty." >&2 - exit 1 + ROCM_VER="${ROCM_VERSION:-7.12.0}" + if [ -n "${IMAGE_TAG_INPUT}" ]; then + echo "IMAGE_TAG=${IMAGE_TAG_INPUT}" >> "${GITHUB_ENV}" + else + echo "IMAGE_TAG=rocm-${ROCM_VER}-ubuntu24.04-py312-${GPU_ARCH}" >> "${GITHUB_ENV}" fi - name: Log in to container registry @@ -60,12 +86,17 @@ jobs: env: REGISTRY: ${{ vars.ARTIFACTORY_DOCKER_REGISTRY }} REPOSITORY: ${{ vars.ARTIFACTORY_CI_DEPS_REPOSITORY }} - IMAGE_TAG: ${{ inputs.image_tag }} + GPU_ARCH: ${{ inputs.gpu_arch }} + ROCM_VERSION_INPUT: ${{ inputs.rocm_version }} run: | set -euo pipefail FULL_IMAGE="${REGISTRY}/${REPOSITORY}" + : "${IMAGE_TAG:?IMAGE_TAG must be set by the validate step}" + ROCM_VER="${ROCM_VERSION_INPUT:-7.12.0}" docker build \ -f .github/scripts/Dockerfile.ci.deps \ + --build-arg "ROCM_VERSION=${ROCM_VER}" \ + --build-arg "GPU_ARCH=${GPU_ARCH}" \ -t "${FULL_IMAGE}:${IMAGE_TAG}" \ . docker push "${FULL_IMAGE}:${IMAGE_TAG}" diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml index e2fb09c15..bae2b201a 100644 --- a/.github/workflows/rocm-ci.yml +++ b/.github/workflows/rocm-ci.yml @@ -54,7 +54,8 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 10 outputs: - image-tag: ${{ steps.select-image.outputs.image-tag }} + image-tag-mi30x: ${{ steps.select-image.outputs.image-tag-mi30x }} + image-tag-mi35x: ${{ steps.select-image.outputs.image-tag-mi35x }} steps: - name: Checkout repository uses: actions/checkout@v6 @@ -88,16 +89,25 @@ jobs: fi echo "Selected config key: $JSON_KEY" - IMAGE_TO_USE=$(jq -r --arg key "$JSON_KEY" '.docker_images[$key]' ci/ci_config.json) + CONFIG_ENTRY=$(jq -c --arg key "$JSON_KEY" '.docker_images[$key]' ci/ci_config.json) MANUAL_OVERRIDE="${{ inputs.docker_image_override }}" if [[ -n "$MANUAL_OVERRIDE" ]]; then echo "::notice::Manual override detected: $MANUAL_OVERRIDE" - IMAGE_TO_USE="$MANUAL_OVERRIDE" + IMAGE_MI30X="$MANUAL_OVERRIDE" + IMAGE_MI35X="$MANUAL_OVERRIDE" + elif jq -e '.mi30x and .mi35x' <<< "$CONFIG_ENTRY" > /dev/null; then + IMAGE_MI30X=$(jq -r '.mi30x' <<< "$CONFIG_ENTRY") + IMAGE_MI35X=$(jq -r '.mi35x' <<< "$CONFIG_ENTRY") + else + IMAGE_MI30X=$(jq -r '.' <<< "$CONFIG_ENTRY") + IMAGE_MI35X="$IMAGE_MI30X" fi - echo "Selected image: $IMAGE_TO_USE" - echo "image-tag=$IMAGE_TO_USE" >> $GITHUB_OUTPUT + echo "Selected mi30x (gfx942) image: $IMAGE_MI30X" + echo "Selected mi35x (gfx950) image: $IMAGE_MI35X" + echo "image-tag-mi30x=$IMAGE_MI30X" >> $GITHUB_OUTPUT + echo "image-tag-mi35x=$IMAGE_MI35X" >> $GITHUB_OUTPUT build: # Delegate wheel building to the reusable workflow on dev. It produces a core .whl plus framework .tar.gz sdists under artifact name `te-rocm-wheels`. @@ -140,7 +150,7 @@ jobs: - name: Pull Docker Image run: | - docker pull ${{ needs.select_image.outputs.image-tag }} + docker pull ${{ matrix.arch_label == 'mi30x' && needs.select_image.outputs.image-tag-mi30x || needs.select_image.outputs.image-tag-mi35x }} - name: Run Container run: | @@ -155,7 +165,7 @@ jobs: --group-add $(getent group video | cut -d: -f3) \ -v "${{ github.workspace }}:/workspace" \ -w /workspace \ - ${{ needs.select_image.outputs.image-tag }} + ${{ matrix.arch_label == 'mi30x' && needs.select_image.outputs.image-tag-mi30x || needs.select_image.outputs.image-tag-mi35x }} - name: Install packages run: | @@ -337,7 +347,7 @@ jobs: - name: Pull Docker Image run: | - docker pull ${{ needs.select_image.outputs.image-tag }} + docker pull ${{ matrix.arch_label == 'mi30x' && needs.select_image.outputs.image-tag-mi30x || needs.select_image.outputs.image-tag-mi35x }} - name: Run Container run: | @@ -352,7 +362,7 @@ jobs: --group-add $(getent group video | cut -d: -f3) \ -v "${{ github.workspace }}:/workspace" \ -w /workspace \ - ${{ needs.select_image.outputs.image-tag }} + ${{ matrix.arch_label == 'mi30x' && needs.select_image.outputs.image-tag-mi30x || needs.select_image.outputs.image-tag-mi35x }} - name: Install packages env: diff --git a/benchmarks/attention/benchmark_attention_rocm.py b/benchmarks/attention/benchmark_attention_rocm.py index b98e9c6e0..076a91de9 100644 --- a/benchmarks/attention/benchmark_attention_rocm.py +++ b/benchmarks/attention/benchmark_attention_rocm.py @@ -84,6 +84,9 @@ output_csv = "times.csv" # Output directory name output_dir_name = "profiler_outputs" +# rocprofv3 output prefix and kernel stats filename (see rocprofv3 -o/-d) +rocprof_output_prefix = "results" +rocprof_kernel_stats_csv = f"{rocprof_output_prefix}_kernel_stats.csv" # Current working directory cwd = os.getcwd() @@ -137,7 +140,7 @@ def setup_backend_env(backend_name, use_ck_bwd_v3=True, use_ck_fwd_v3=True, use_ "aotriton_bwd": "bwd", } -# Runs benchmark with warmup iterations and profiles using rocprof +# Runs benchmark with warmup iterations and profiles using rocprofv3 def benchmark_dot_product_attention(model, attention, column_name, dirname): config = model_configs[model] @@ -153,29 +156,34 @@ def benchmark_dot_product_attention(model, attention, column_name, dirname): is_training, ) os.makedirs(dirname, exist_ok=True) - before_files = set(os.listdir(cwd)) - # Profiling command using rocprof benchmark_dir = os.path.dirname(os.path.abspath(__file__)) + profiler_script = ( + f"import sys; sys.path.insert(0, {benchmark_dir!r}); " + f"import benchmark_attention_rocm; " + f"benchmark_attention_rocm.benchmark_dot_product_attention_profiler(" + f"{model!r}, {attention!r}, {column_name!r})" + ) + # rocprofv3: --kernel-trace + --stats replaces rocprofv2 --hip-trace (kernel stats + # are not enabled by default). Full kernel names are kept (v2 --basenames off). prof_cmd = [ - "rocprof", - "--hip-trace", - "--basenames off", - "python", - "-c", - f""" "import sys; sys.path.insert(0, '{benchmark_dir}'); import benchmark_attention_rocm;""", - f"""benchmark_attention_rocm.benchmark_dot_product_attention_profiler(""" - f"""'{model}', '{attention}', '{column_name}')" """, - ] - prof_cmd = " ".join(prof_cmd) - subprocess.call(prof_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True) - after_files = set(os.listdir(cwd)) - new_files = after_files - before_files - - for f in new_files: - src_path = os.path.join(cwd, f) - dst_path = os.path.join(dirname, f) - if os.path.isfile(src_path): # Only move files, not directories - shutil.move(src_path, dst_path) + "rocprofv3", + "--kernel-trace", + "--stats", + "-f", "csv", + "-o", rocprof_output_prefix, + "-d", dirname, + "--", + sys.executable, + "-c", + profiler_script, + ] + result = subprocess.run(prof_cmd, capture_output=True, text=True) + if result.returncode != 0: + print( + f"rocprofv3 failed for {model} [{attention}] (exit {result.returncode}):\n" + f"{result.stderr}", + file=sys.stderr, + ) torch.cuda.empty_cache() # Runs profiler and records timing information @@ -216,7 +224,10 @@ def calculate_attention_tflops(batch_size, seq_len, num_heads_q, head_dim_qk, fw # Helper function to extract timing results from profiler logs def parse_helper(model, dirname, fwd_search_pattern, bwd_search_pattern, column_name, df_times): - df = pd.read_csv(os.path.join(dirname, "results.stats.csv")) + stats_csv = os.path.join(dirname, rocprof_kernel_stats_csv) + if not os.path.isfile(stats_csv): + return False + df = pd.read_csv(stats_csv) # Extract kernel timing values fwd_values = df[df["Name"].str.contains(fwd_search_pattern, regex=False)]["AverageNs"].to_numpy() @@ -281,7 +292,7 @@ def sanity_checks( ): """ • Verifies that every model/backend that *should* have run produced - profiler_root//results.stats.csv + profiler_root//results_kernel_stats.csv • Non-zero exit code on any failure (CI friendly) """ if profiler_root is None: @@ -323,12 +334,14 @@ def sanity_checks( print(f"{model}:") # Rocprof run status for be, pat in expected.items(): - stats = os.path.join(profiler_root, pat.format(model=model), "results.stats.csv") + stats = os.path.join(profiler_root, pat.format(model=model), rocprof_kernel_stats_csv) if os.path.isfile(stats): print(f" [{be:<22}] Profiling successful") else: ok_overall = False - raise FileNotFoundError(f"Error while profiling {model} [{be}], results.stats.csv not found") + raise FileNotFoundError( + f"Error while profiling {model} [{be}], {rocprof_kernel_stats_csv} not found" + ) print("-" * 60) return ok_overall @@ -347,7 +360,7 @@ def main(args): os.makedirs(output_dir) df_times = pd.DataFrame(index=indices, columns=columns) - df_times = df_times.infer_objects(copy=False) + df_times = df_times.infer_objects() df_times.fillna(0.0, inplace=True) df_times.index.name = "Model" df_times.to_csv(output_csv_path) diff --git a/ci/README.md b/ci/README.md index 07d5bd7d2..7c15748d0 100644 --- a/ci/README.md +++ b/ci/README.md @@ -21,4 +21,13 @@ It is the caller's responsibility to clean up generated files. Default and release-specific TE CI images are listed in [`ci_config.json`](ci_config.json) under `docker_images`. +For `dev` and other branches using the `default` entry, images are selected per runner architecture: + +| Runner label | GPU arch | Image tag | +|--------------|----------|-----------| +| `linux-te-mi30x-*` | gfx942 (MI300X) | `rocm-7.12.0-ubuntu24.04-py312-gfx942_test` | +| `linux-te-mi35x-*` | gfx950 (MI350X) | `rocm-7.12.0-ubuntu24.04-py312-gfx950_test` | + +Registry: `registry-sc-harbor.amd.com/framework/te-ci` + The default image is built from [`.github/scripts/Dockerfile.ci.deps`](../.github/scripts/Dockerfile.ci.deps). It pins [ROCm/aiter](https://github.com/ROCm/aiter) at commit [`77455e3ecf4f0d28756afc452e914940c45b944b`](https://github.com/ROCm/aiter/commit/77455e3ecf4f0d28756afc452e914940c45b944b). That revision was validated in CI for **MXFP4 FP4 GEMM** kernel coverage. diff --git a/ci/ci_config.json b/ci/ci_config.json index 123ded73a..9ce3526fc 100644 --- a/ci/ci_config.json +++ b/ci/ci_config.json @@ -1,6 +1,9 @@ { "docker_images": { - "default": "registry-sc-harbor.amd.com/framework/te-ci:rocm-7.2_ubuntu22.04_py3.11_pytorch_release-2.8_08d38866_jax_0.8.0_fa_2.8.1_aiter_77455e3ecf", + "default": { + "mi30x": "registry-sc-harbor.amd.com/framework/te-ci:rocm-7.12.0-ubuntu24.04-py312-gfx942_test", + "mi35x": "registry-sc-harbor.amd.com/framework/te-ci:rocm-7.12.0-ubuntu24.04-py312-gfx950_test" + }, "release_v1.13": "compute-artifactory.amd.com:5000/rocm-plus-docker/framework/private/te-ci:rocm-6.4_0_ubuntu22_py310_torch25_jax0435qa_fa273", "release_v1.14": "compute-artifactory.amd.com:5000/rocm-plus-docker/framework/private/te-ci:rocm-6.4_0_ubuntu22_py310_torch25_jax0435qa_fa273" } diff --git a/task.txt b/task.txt new file mode 100644 index 000000000..23a319960 --- /dev/null +++ b/task.txt @@ -0,0 +1,11 @@ +create the docker image with ubuntu 24.04 and install the rock using the info from + +every thibng to install rocm and all are founf at https://repo.amd.com/rocm/whl/ + +Current released is 7.12.0 and + +my wheels for torch, torchvision, torchaudio, jax, jaxlib, triton are at https://repo.amd.com/rocm/whl/ see if you can fetch the info from here + +I want to pickfor these two archs gfx94X-dcgpu gfx950-dcgpu and python version 3.12 + +update @.github/scripts/Dockerfile.ci.deps as well, to install for one arch by default but provide the arg to change the arch so that it installs those