From 8654db624806ff686405864af739aae8d888a431 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Tue, 4 Aug 2026 18:10:41 +0000 Subject: [PATCH 1/4] [ROCm] nightly: build jax/jaxlib on fork runners Upstream CI has stopped publishing jax/jaxlib nightly wheels to GCS (gs://jax-nightly-artifacts/latest), leaving the ROCm nightly test legs running against missing/stale base wheels. On the ROCm fork, build the base jax and jaxlib wheels on the same amd-do-linux.jax.cpu runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, so the ROCm test legs consume freshly built wheels instead of the broken GCS bucket. Upstream (jax-ml) behavior is unchanged. - ci/build_rocm_artifacts.sh: allow "jax"/"jaxlib" artifacts; skip auditwheel for the pure-python jax wheel. - wheel_tests_nightly_release.yml: new fork-only build-rocm-jax-jaxlib job; wire it into the ROCm test legs and select S3 as the jax/jaxlib source on the fork. - pytest_rocm.yml / bazel_rocm.yml: new optional jax-jaxlib-source input (default gcs), forwarded to the download action. - download-jax-rocm-wheels: when source is s3, fetch jax/jaxlib from the CloudFront listing alongside the ROCm wheels instead of GCS. --- .../download-jax-rocm-wheels/action.yml | 39 +++++++++++++-- .github/workflows/bazel_rocm.yml | 5 ++ .github/workflows/pytest_rocm.yml | 9 ++++ .../workflows/wheel_tests_nightly_release.yml | 50 ++++++++++++++++++- ci/build_rocm_artifacts.sh | 15 ++++-- 5 files changed, 110 insertions(+), 8 deletions(-) diff --git a/.github/actions/download-jax-rocm-wheels/action.yml b/.github/actions/download-jax-rocm-wheels/action.yml index 1b9c18c70302..b1f91eceed5f 100644 --- a/.github/actions/download-jax-rocm-wheels/action.yml +++ b/.github/actions/download-jax-rocm-wheels/action.yml @@ -27,6 +27,9 @@ inputs: s3_download_uri: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' + jax-jaxlib-source: + description: "Where to fetch the base jax/jaxlib wheels from: 'gcs' (upstream nightly bucket) or 's3' (built alongside the ROCm wheels on the fork's runners)" + default: 'gcs' permissions: {} runs: using: "composite" @@ -74,7 +77,11 @@ runs: mkdir -p $(pwd)/dist # Disable Authentication on the ROCm runners. gcloud config set auth/disable_credentials True - if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then + # jax is fetched from S3/CloudFront in the jaxlib "head" block below, + # alongside jaxlib and the ROCm plugin/pjrt wheels. + echo "jax wheel will be downloaded from S3/CloudFront alongside the ROCm wheels." + elif [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}"/jax*py3*none*any.whl $(pwd)/dist/ else echo "JAX wheel won't be downloaded, only jaxlib pre-built wheel is tested." @@ -86,7 +93,9 @@ runs: echo "JAX only release. Only downloading the jax wheel from the release bucket." else if [[ ${INPUTS_JAXLIB_VERSION} == "head" ]]; then - gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" != "s3" ]]; then + gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ + fi # Resolve ROCm plugin/PJRT wheels path via CloudFront. ROCM_WHEELS_BASE_URL="https://d22q5eopkfeftw.cloudfront.net" @@ -112,8 +121,31 @@ runs: exit 1 fi + DOWNLOAD_URLS=("${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}") + + # On the fork, the base jax/jaxlib wheels are built alongside the + # ROCm wheels and uploaded to the same S3 prefix, so fetch them from + # the same CloudFront listing instead of the (broken) GCS bucket. + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then + JAXLIB_FILE=$(echo "$FILES" | grep "jaxlib-" | grep "${PYTHON_MAJOR_MINOR}" | head -n1 || true) + if [[ -z "${JAXLIB_FILE}" ]]; then + echo "Could not find jaxlib wheel in ${WHEELS_URL}" + exit 1 + fi + DOWNLOAD_URLS+=("${WHEELS_URL}/${JAXLIB_FILE}") + + if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then + JAX_FILE=$(echo "$FILES" | grep -E "jax-[0-9].*-py3-none-any\.whl" | head -n1 || true) + if [[ -z "${JAX_FILE}" ]]; then + echo "Could not find jax wheel in ${WHEELS_URL}" + exit 1 + fi + DOWNLOAD_URLS+=("${WHEELS_URL}/${JAX_FILE}") + fi + fi + PYTHON_BIN="python${INPUTS_PYTHON}" - "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}" + "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${DOWNLOAD_URLS[@]}" elif [[ ${INPUTS_JAXLIB_VERSION} == "pypi_latest" ]]; then PYTHON=python${INPUTS_PYTHON} $PYTHON -m pip download jaxlib jax-rocm${JAXCI_ROCM_VERSION}-pjrt jax-rocm${JAXCI_ROCM_VERSION}-plugin --dest $(pwd)/dist/ --no-deps @@ -134,6 +166,7 @@ runs: INPUTS_PYTHON: ${{ inputs.python }} INPUTS_S3_DOWNLOAD_URI: ${{ inputs.s3_download_uri }} INPUTS_ROCM_FULL_VERSION: ${{ inputs.rocm-release-version || inputs.rocm-version }} + INPUTS_JAX_JAXLIB_SOURCE: ${{ inputs.jax-jaxlib-source }} - name: Skip the test run if the wheel artifacts were not downloaded successfully shell: bash if: steps.download-wheel-artifacts.outcome == 'failure' diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index e9a75a071b66..b10cb5bc609f 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -87,6 +87,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' type: string + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + default: 'gcs' + type: string gate-targets-file: description: "Path to a file listing Bazel targets to run (one per line). If empty, runs the full default test suite." type: string @@ -175,6 +179,7 @@ jobs: jaxlib-version: ${{ inputs.jaxlib-version }} gcs_download_uri: ${{ inputs.gcs_download_uri }} s3_download_uri: ${{ inputs.s3_download_uri }} + jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }} # Halt for testing - name: Wait For Connection uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index dbe8c733969e..f5e00533d47b 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -66,6 +66,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" type: string default: 'latest' + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + type: string + default: 'gcs' halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -108,6 +112,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' type: string + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + default: 'gcs' + type: string halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -162,6 +170,7 @@ jobs: skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }} gcs_download_uri: ${{ inputs.gcs_download_uri }} s3_download_uri: ${{ inputs.s3_download_uri }} + jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }} - name: Install Python dependencies run: | $JAXCI_PYTHON -m pip install uv~=0.11.2 diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index d31bb147bbd6..4de09534277f 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -201,9 +201,49 @@ jobs: upload_artifacts_to_s3: true s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + # Upstream CI frequently stops publishing jax/jaxlib nightly wheels to GCS. + # On the ROCm fork, build the base jax and jaxlib wheels ourselves on the same + # runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, + # so the ROCm test legs can consume them instead of the (stale) GCS bucket. + build-rocm-jax-jaxlib: + if: github.repository_owner == 'ROCm' + uses: ./.github/workflows/build_rocm_artifacts.yml + permissions: + id-token: write + contents: read + actions: read + strategy: + fail-fast: false + matrix: + runner: ["amd-do-linux.jax.cpu"] + artifact: ["jax", "jaxlib"] + python: ["3.12", "3.13", "3.14"] + rocm: [ + {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, + ] + exclude: + # jax is a pure-python (py3-none-any) wheel; build it once per ROCm + # variant rather than once per Python version. + - artifact: "jax" + python: "3.13" + - artifact: "jax" + python: "3.14" + name: "Build ROCm jax/jaxlib (${{ matrix.rocm.label }})" + with: + runner: ${{ matrix.runner }} + artifact: ${{ matrix.artifact }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.wheel-version }} + rocm-release-version: ${{ matrix.rocm.release-version }} + manylinux-image: ${{ matrix.rocm.manylinux-image }} + clone_main_xla: 0 + upload_artifacts_to_s3: true + s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + run-pytest-rocm: if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }} - needs: [build-rocm-artifacts] + needs: [build-rocm-artifacts, build-rocm-jax-jaxlib] permissions: id-token: write contents: read @@ -229,10 +269,13 @@ jobs: skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} gcs_download_uri: ${{inputs.gcs_download_uri}} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} run-bazel-test-rocm: if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }} - needs: [build-rocm-artifacts] + needs: [build-rocm-artifacts, build-rocm-jax-jaxlib] permissions: id-token: write contents: read @@ -258,6 +301,9 @@ jobs: gcs_download_uri: ${{inputs.gcs_download_uri}} halt-for-connection: ${{inputs.halt-for-connection}} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} build_jaxlib: ${{ 'false' }} build_jax: ${{ 'false' }} jaxlib-version: "head" diff --git a/ci/build_rocm_artifacts.sh b/ci/build_rocm_artifacts.sh index 4bfacf04616a..77097036ce9e 100755 --- a/ci/build_rocm_artifacts.sh +++ b/ci/build_rocm_artifacts.sh @@ -37,7 +37,13 @@ fi # Set up the build environment. source "ci/utilities/setup_build_environment.sh" -allowed_artifacts=("jax-rocm-plugin" "jax-rocm-pjrt") +# The base "jax" and "jaxlib" wheels are also built here on the ROCm fork's +# runners so the ROCm test legs no longer depend on the (frequently broken) +# upstream nightly wheels published to GCS. jax is a pure-python wheel and +# jaxlib is the backend-neutral CPU base (ROCm support lives in the separate +# plugin/pjrt wheels), so both build correctly with the ROCm bazel config used +# below. +allowed_artifacts=("jax-rocm-plugin" "jax-rocm-pjrt" "jax" "jaxlib") if [[ ! " ${allowed_artifacts[*]} " =~ " ${artifact} " ]]; then echo "Error: Invalid artifact: $artifact. Allowed values are: ${allowed_artifacts[*]}" @@ -101,5 +107,8 @@ python build/build.py build --wheels="$artifact" \ $wheel_version_suffix_flag \ "${rocm_path_flags[@]}" -# Verify manylinux compliance. -./ci/utilities/run_auditwheel.sh +# Verify manylinux compliance. The pure-python "jax" wheel has no platform +# wheels to inspect (run_auditwheel.sh errors when it finds none), so skip it. +if [[ "$artifact" != "jax" ]]; then + ./ci/utilities/run_auditwheel.sh +fi From 9dbd68a2e7135790ab7c4ee9259a33db2172f04c Mon Sep 17 00:00:00 2001 From: tsrw2048 <239799652+tsrw2048@users.noreply.github.com> Date: Fri, 26 Jun 2026 18:47:25 +0000 Subject: [PATCH 2/4] [ROCm] Add expanded target arch set for ROCm to rocm.bazelrc. --- build/rocm/rocm.bazelrc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build/rocm/rocm.bazelrc b/build/rocm/rocm.bazelrc index a2a92590d9de..2e094ea17538 100644 --- a/build/rocm/rocm.bazelrc +++ b/build/rocm/rocm.bazelrc @@ -10,7 +10,7 @@ common:rocm --action_env=TF_ROCM_CLANG="1" common:rocm --action_env=HIPCC_COMPILE_FLAGS_APPEND=--offload-compress common:rocm --copt=-Wno-gnu-offsetof-extensions common:rocm --copt=-Qunused-arguments -common:rocm --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1200,gfx1201" +common:rocm --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx900,gfx906,gfx908,gfx90a,gfx942,gfx950,gfx1010,gfx1011,gfx1012,gfx1030,gfx1031,gfx1032,gfx1033,gfx1034,gfx1035,gfx1036,gfx1100,gfx1101,gfx1102,gfx1103,gfx1150,gfx1151,gfx1152,gfx1153,gfx1200,gfx1201" # Used for @xla//build_tools/rocm:parallel_gpu_execute common:rocm --legacy_external_runfiles=true From daadb3b0bf1f45e618c81ebd25a1250a6800f6a9 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Tue, 11 Aug 2026 22:54:43 +0000 Subject: [PATCH 3/4] [ROCm] Manually bump of XLA pin. This adds 8028c6c78ea42b7184be31bcbf9d1fe40809bde1 on top of upstream's pin. --- third_party/xla/revision.bzl | 6 +++--- third_party/xla/workspace.bzl | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/third_party/xla/revision.bzl b/third_party/xla/revision.bzl index d5f3fd19ce71..546232ffe6f7 100644 --- a/third_party/xla/revision.bzl +++ b/third_party/xla/revision.bzl @@ -17,9 +17,9 @@ # To update XLA to a new revision, # a) update XLA_COMMIT to the new git commit hash # b) get the sha256 hash of the commit by running: -# curl -L https://github.com/openxla/xla/archive/{git_hash}.tar.gz | sha256sum +# curl -L https://github.com/ROCm/xla/archive/{git_hash}.tar.gz | sha256sum # and update XLA_SHA256 with the result. # buildifier: disable=module-docstring -XLA_COMMIT = "c54647791964d0eb2062ea2f293bd199c48045b2" -XLA_SHA256 = "7d08c2cb2267e7be25314682404662867ae474d8c1498f3eea44314bf83e497f" +XLA_COMMIT = "9df7d1a5b395aa557be11a673861fd7a08f75d44" +XLA_SHA256 = "c291f639021bf4d304f6d4f39d0ba103bd848034f5fe444c8c604f4468ce51b2" diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl index 286aa167c0b2..e3754838a1cb 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -21,7 +21,7 @@ def repo(): name = "xla", sha256 = XLA_SHA256, strip_prefix = "xla-{commit}".format(commit = XLA_COMMIT), - urls = tf_mirror_urls("https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)), + urls = tf_mirror_urls("https://github.com/ROCm/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)), ) # For development, one often wants to make changes to the TF repository as well From b5d9ff49045b5ed81af1c9bd95e519e84451394b Mon Sep 17 00:00:00 2001 From: Kunreddy Date: Wed, 12 Aug 2026 18:44:44 -0500 Subject: [PATCH 4/4] Mount /home/runner/_work in ROCm workflow containers Ensure GitHub Actions workspace paths are available inside ROCm container jobs by mounting /home/runner/_work in pytest_rocm and bazel_rocm workflows. Co-authored-by: Cursor --- .github/workflows/bazel_rocm.yml | 1 + .github/workflows/pytest_rocm.yml | 2 ++ 2 files changed, 3 insertions(+) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index b10cb5bc609f..6f68b73443bf 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -122,6 +122,7 @@ jobs: image: ghcr.io/rocm/jax-dev-ubu24.${{ inputs.rocm-tag }}:latest # zizmor: ignore[unpinned-images] volumes: - /data:/data + - /home/runner/_work:/home/runner/_work options: >- --device=/dev/kfd --device=/dev/dri diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index f5e00533d47b..3b827eca3637 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -141,6 +141,8 @@ jobs: credentials: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} + volumes: + - /home/runner/_work:/home/runner/_work options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings name: "${{ (contains(inputs.runner, 'gfx950.1') && 'gfx950.1') || (contains(inputs.runner, 'gfx950.4') && 'gfx950.4') ||