Skip to content

Build Multi-Arch Linux PyTorch Wheels (dev, 3.13, pablo-garay/release-2.14-integration-test) #366

Build Multi-Arch Linux PyTorch Wheels (dev, 3.13, pablo-garay/release-2.14-integration-test)

Build Multi-Arch Linux PyTorch Wheels (dev, 3.13, pablo-garay/release-2.14-integration-test) #366

# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT
# Reusable multi-arch PyTorch build workflow for Linux.
#
# Built shape:
# - One fat build_prod_wheels.py invocation for one python_version and
# pytorch_git_ref, covering all gfx targets expanded from the
# inputs.amdgpu_families set. Builds torch, torchaudio, torchvision,
# triton, and apex in a single invocation.
# - kpack-split the fat torch and torchvision wheels into host wheels
# plus per-gfx amd-torch-device-* / amd-torchvision-device-* wheels.
# other packages ship as-is (nothing to split).
# - Upload the result to s3://therock-{release_type}-python/v4/whl/
# via publish_pytorch_to_release_bucket.py (pip-compatible index generation
# happens server-side).
#
# For convergence with CI and testing plans see:
# * https://github.com/ROCm/TheRock/issues/3332
# * https://github.com/ROCm/TheRock/issues/3291
name: Multi-Arch Build Portable Linux PyTorch Wheels
on:
workflow_call:
inputs:
amdgpu_families:
description: >-
Semicolon-separated list of AMD GPU families to build PyTorch for
(e.g. 'gfx94X-dcgpu;gfx120X-all'). Expanded to gfx targets at
build time via cmake/therock_amdgpu_targets.cmake.
type: string
required: true
test_amdgpu_families:
description: >-
Semicolon-separated AMD GPU families to run tests for (e.g.
'gfx94X-dcgpu'). Use 'auto' to test built families, or 'none' to skip
tests.
type: string
default: "auto"
python_version:
description: Python version to build wheels for.
type: string
required: true
pytorch_git_ref:
description: PyTorch ref to checkout. Typically "nightly" or "release/X.Y".
type: string
required: true
pytorch_gitrepo_origin:
description: >-
Optional PyTorch git origin override (for example, a fork URL for
pre-merge testing). Defaults to ROCm/pytorch for stable refs.
type: string
default: ""
rocm_version:
description: ROCm package version to install and build against (e.g. 7.10.0.dev0)
type: string
required: true
rocm_package_find_links_url:
description: URL for pip --find-links to install ROCm packages (flat page; use this OR rocm_package_index_url)
type: string
default: ""
rocm_package_index_url:
description: URL for pip --index-url to install ROCm packages (PEP 503 index, e.g. repo.amd.com; use this OR rocm_package_find_links_url)
type: string
default: ""
release_type:
description: 'Release type; selects the S3 bucket.'
type: string
required: true
repository:
description: "Repository to checkout. Defaults to github.repository."
type: string
default: ""
ref:
description: "Branch, tag, or SHA to checkout. Defaults to the triggering ref."
type: string
default: ""
cache_type:
description: "Compiler cache type ('none', 'sccache', or 'ccache')"
type: string
default: "sccache"
build_runs_on:
description: "Build runner label (selected by configure script with weighted distribution)"
type: string
default: "aws-linux-scale-rocm-prod"
run_full_pytorch_tests:
description: >-
Dispatch the full PyTorch test suite after a successful build.
Only takes effect for gfx94X-dcgpu Python 3.12 builds.
type: boolean
default: false
workflow_dispatch:
inputs:
amdgpu_families:
description: >-
Semicolon-separated list of AMD GPU families to build PyTorch for
(e.g. 'gfx94X-dcgpu;gfx120X-all'). Expanded to gfx targets at
build time via cmake/therock_amdgpu_targets.cmake.
type: string
required: true
test_amdgpu_families:
description: >-
Semicolon-separated AMD GPU families to run tests for (e.g.
'gfx94X-dcgpu'). Use 'auto' to test built families, or 'none' to skip
tests.
type: string
default: "auto"
python_version:
description: Python version to build wheels for (for example, "3.12").
type: string
default: "3.12"
pytorch_git_ref:
description: PyTorch ref to checkout (for example, "release/2.13").
type: string
default: "release/2.13"
pytorch_gitrepo_origin:
description: >-
Optional PyTorch git origin override (for example, a fork URL for
pre-merge testing). Defaults to ROCm/pytorch for stable refs.
type: string
default: ""
rocm_version:
description: ROCm package version to install and build against (e.g. 7.10.0.dev0)
type: string
required: true
rocm_package_find_links_url:
description: >-
URL for pip --find-links to install ROCm packages (flat page). Set
either this or rocm_package_index_url to override the ROCm package
source; leave both blank to use the nightly multi-arch index.
type: string
default: ""
rocm_package_index_url:
description: >-
URL for pip --index-url to install ROCm packages (PEP 503 index).
Set either this or rocm_package_find_links_url to override the ROCm
package source; leave both blank to use the nightly multi-arch index
(https://rocm.nightlies.amd.com/whl-multi-arch).
type: string
default: ""
release_type:
description: 'Release type; developer-triggered jobs should use "dev".'
type: choice
options:
- dev
default: dev
repository:
description: >-
TheRock repository to check out the build scripts/workflow from (NOT
the PyTorch source). Leave blank to use this repository. To build a
different PyTorch branch, set pytorch_git_ref instead.
type: string
default: ""
ref:
description: >-
Branch, tag, or SHA of TheRock (the repository above) to check out.
Leave blank to use the triggering ref. This does NOT select the
PyTorch source ref; use pytorch_git_ref for that.
type: string
default: ""
cache_type:
description: "Compiler cache type"
type: choice
options:
- none
- sccache
- ccache
default: sccache
build_runs_on:
description: "Build runner label (selected by configure script with weighted distribution)"
type: string
default: "aws-linux-scale-rocm-prod"
run_full_pytorch_tests:
description: >-
Dispatch the full PyTorch test suite after a successful build.
Only takes effect for gfx94X-dcgpu Python 3.12 builds.
type: boolean
default: false
run-name: Build Multi-Arch Linux PyTorch Wheels (${{ inputs.release_type }}, ${{ inputs.python_version }}, ${{ inputs.pytorch_git_ref }})
permissions:
contents: read
jobs:
build_pytorch_wheels:
name: Build | py ${{ inputs.python_version }} | torch ${{ inputs.pytorch_git_ref }}
runs-on: ${{ github.repository_owner == 'ROCm' && inputs.build_runs_on || 'ubuntu-24.04' }}
permissions:
id-token: write
contents: read
container:
image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:a382085df3ba2419b58aa9051350883a0d0b732a4bc0a4ef60458f8161bb08c6
env:
PACKAGE_DIST_DIR: ${{ github.workspace }}/output/packages/dist
optional_build_prod_arguments: ""
SCCACHE_BUCKET: "therock-pytorch-sccache-${{ inputs.release_type }}"
SCCACHE_REGION: us-east-1
SCCACHE_S3_KEY_PREFIX: "linux/multiarch/v1/"
SCCACHE_S3_USE_SSL: "true"
SCCACHE_S3_SERVER_SIDE_ENCRYPTION: "true"
SCCACHE_IDLE_TIMEOUT: "0" # never idle-out; keeps stats accurate on long builds
SCCACHE_LOG: warn
outputs:
torch_version: ${{ steps.pytorch-versions.outputs.torch_version }}
torchaudio_version: ${{ steps.pytorch-versions.outputs.torchaudio_version }}
apex_version: ${{ steps.pytorch-versions.outputs.apex_version }}
torchvision_version: ${{ steps.pytorch-versions.outputs.torchvision_version }}
triton_version: ${{ steps.pytorch-versions.outputs.triton_version }}
package_index_url: ${{ steps.publish-pytorch-wheels.outputs.package_index_url }}
summary: ${{ steps.summary.outputs.summary }}
steps:
- name: "Quartz - started - build PyTorch wheels"
continue-on-error: true
if: ${{ vars.NOTIFY_QUARTZ_ENABLED == 'true' }}
uses: ROCm/Quartz/.github/actions/notify_quartz@f386a9756620938616af0b4d5d04b24ae6e0353f # notify_quartz/v1.1.1
with:
run_phase: started
reporting_workflow: multi_arch_build_portable_linux_pytorch_wheels.yml
workflow_inputs: ${{ toJSON(inputs) }}
gh_app_client_id: ${{ secrets.GH_APP_HAULY_CID }}
gh_app_private_key: ${{ secrets.GH_APP_HAULY_PRIVATE_KEY }}
- name: Checkout
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
repository: ${{ inputs.repository || github.repository }}
ref: ${{ inputs.ref }}
- name: Configure Git Identity
run: |
git config --global user.name "therockbot"
git config --global user.email "therockbot@amd.com"
- name: Select Python version
run: |
python build_tools/github_actions/python_to_cp_version.py \
--python-version ${{ inputs.python_version }}
- name: Add selected Python version to PATH
run: |
python_dir="/opt/python/${{ env.cp_version }}"
if ! [ -x "${python_dir}/bin/python" ]; then
echo "ERROR: Could not find python: ${python_dir}"
exit 1
fi
echo "${python_dir}/bin" >> "$GITHUB_PATH"
- name: Install python deps for CI
run: |
pip install -r external-builds/pytorch/requirements-ci.txt
- name: Checkout PyTorch source repos (nightly)
if: ${{ inputs.pytorch_git_ref == 'nightly' }}
timeout-minutes: 30
run: |
./external-builds/pytorch/pytorch_torch_repo.py checkout \
--repo-hashtag nightly \
--no-commit-hipify
./external-builds/pytorch/pytorch_audio_repo.py checkout \
--repo-hashtag nightly \
--no-commit-hipify
./external-builds/pytorch/pytorch_vision_repo.py checkout \
--repo-hashtag nightly \
--no-commit-hipify
./external-builds/pytorch/pytorch_triton_repo.py checkout \
--no-commit-hipify
./external-builds/pytorch/pytorch_apex_repo.py checkout \
--repo-hashtag master \
--no-commit-hipify
- name: Checkout PyTorch source repos (stable)
if: ${{ inputs.pytorch_git_ref != 'nightly' }}
timeout-minutes: 30
run: |
pytorch_origin="${{ inputs.pytorch_gitrepo_origin }}"
if [ -z "${pytorch_origin}" ]; then
pytorch_origin="https://github.com/ROCm/pytorch.git"
fi
./external-builds/pytorch/pytorch_torch_repo.py checkout \
--gitrepo-origin "${pytorch_origin}" \
--repo-hashtag ${{ inputs.pytorch_git_ref }} \
--no-commit-hipify
./external-builds/pytorch/pytorch_audio_repo.py checkout \
--require-related-commit \
--no-commit-hipify
./external-builds/pytorch/pytorch_vision_repo.py checkout \
--require-related-commit \
--no-commit-hipify
./external-builds/pytorch/pytorch_triton_repo.py checkout \
--no-commit-hipify
./external-builds/pytorch/pytorch_apex_repo.py checkout \
--require-related-commit \
--no-commit-hipify
# determine_version.py sets optional_build_prod_arguments in
# GITHUB_ENV (--rocm-sdk-version, --version-suffix).
- name: Determine optional arguments passed to `build_prod_wheels.py`
run: |
python build_tools/github_actions/determine_version.py \
--rocm-version ${{ inputs.rocm_version }}
- name: Configure AWS Credentials for sccache
if: ${{ inputs.cache_type == 'sccache' && github.repository_owner == 'ROCm' }}
uses: aws-actions/configure-aws-credentials@e6de054238d6b7531b4efff3b6587d9aade6a06c # v6.2.3
with:
aws-region: us-east-1
role-to-assume: arn:aws:iam::324352301041:role/therock-${{ inputs.release_type }}
# Multi-package PyTorch builds can exceed the default 1h credential
# window; the role's MaxSessionDuration must be >= this value.
role-duration-seconds: 21600
- name: Verify sccache
if: ${{ inputs.cache_type == 'sccache' }}
run: |
echo "sccache version: $(sccache --version)"
echo "SCCACHE_BUCKET=${SCCACHE_BUCKET}"
echo "SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX}"
- name: Expand amdgpu families to gfx targets
id: expand
run: |
targets=$(python build_tools/github_actions/expand_amdgpu_families.py \
--amdgpu-families "${{ inputs.amdgpu_families }}")
echo "Expanded families '${{ inputs.amdgpu_families }}' -> '${targets}'"
echo "targets=${targets}" >> "$GITHUB_OUTPUT"
- name: Build PyTorch wheels
# Default to the nightly multi-arch ROCm index when neither source is
# provided so a bare manual dispatch works; setting either input
# overrides it. The default is applied here rather than on the input
# because GitHub reverts an emptied input back to its declared default,
# which would make the two inputs impossible to use exclusively.
env:
# Use the index if given; else the find-links page; else default to
# the nightly multi-arch index (both empty is the bare-dispatch case).
ROCM_INDEX_URL: ${{ inputs.rocm_package_index_url || (inputs.rocm_package_find_links_url == '' && 'https://rocm.nightlies.amd.com/whl-multi-arch' || '') }}
ROCM_FIND_LINKS: ${{ inputs.rocm_package_find_links_url }}
run: |
cache_flag=""
if [ "${{ inputs.cache_type }}" = "sccache" ]; then
cache_flag="--use-sccache"
elif [ "${{ inputs.cache_type }}" = "ccache" ]; then
cache_flag="--use-ccache"
fi
# build_prod_wheels.py forwards each of these to pip and ignores an
# empty value, so pass both and supply whichever source applies:
# --index-url for a PEP 503 index, --find-links for a flat page.
# --release-type lets the script tag dev builds' torch version with
# the torch source commit (e.g. 2.12.0a0+git1a2b3c4d.rocm7.10.0).
./external-builds/pytorch/build_prod_wheels.py \
build \
--install-rocm \
--index-url="${ROCM_INDEX_URL}" \
--find-links="${ROCM_FIND_LINKS}" \
--clean \
--output-dir ${{ env.PACKAGE_DIST_DIR }} \
${cache_flag} \
--release-type="${{ inputs.release_type }}" \
--pytorch-rocm-arch ${{ steps.expand.outputs.targets }} \
${{ env.optional_build_prod_arguments }}
- name: Report cache stats
if: ${{ !cancelled() && inputs.cache_type != 'none' }}
run: |
if [ "${{ inputs.cache_type }}" = "sccache" ]; then
echo "sccache stats:"
echo "--------------"
sccache --show-stats || true
elif [ "${{ inputs.cache_type }}" = "ccache" ]; then
echo "ccache stats:"
echo "-------------"
ccache -s -v || true
fi
- name: Sanity check wheel
run: |
python external-builds/pytorch/sanity_check_wheel.py \
${{ env.PACKAGE_DIST_DIR }}/
# The kpack wheel splitter takes the fat torch wheel and produces a
# smaller host wheel plus one device wheel per gfx target. Source-install
# from the TheRock checkout - kpack has no published wheel today.
# The outer actions/checkout skips submodules for speed, so pull in
# rocm-systems only when we actually need it.
- name: Init rocm-systems submodule
run: |
git config --global --add safe.directory "$GITHUB_WORKSPACE"
git submodule update --init --depth=1 rocm-systems
- name: Install kpack tooling
run: |
pip install \
-e rocm-systems/python/rocm-bootstrap \
-e rocm-systems/shared/kpack
# TODO(https://github.com/ROCm/TheRock/issues/5656): Integrate this
# rocm_kpack.tools.split_python_wheels flow into build_prod_wheels.py and
# document the standalone developer workflow.
- name: Split fat wheels (torch + torchvision)
run: |
# The splitter writes the host wheel under the same filename as
# the input, so relocate each fat wheel out of PACKAGE_DIST_DIR
# before invoking the splitter. torchaudio, triton, and apex ship
# as-is: torchaudio and triton have no fat binaries to split, and
# apex's ROCm fork ships JIT-loader shims rather than precompiled
# extensions, so its wheel is constant-size in arch count.
FAT_STAGE="${RUNNER_TEMP:-/tmp}/pytorch-fat"
mkdir -p "${FAT_STAGE}"
ROCM_SDK_VERSION="$(python -m rocm_sdk version)"
mv ${{ env.PACKAGE_DIST_DIR }}/torch-*.whl "${FAT_STAGE}/"
python -m rocm_kpack.tools.split_python_wheels \
--input "${FAT_STAGE}"/torch-*.whl \
--output-dir "${{ env.PACKAGE_DIST_DIR }}" \
--device-package-prefix amd-torch-device \
--overlay-root torch/ \
--wheel-type torch-fat \
--device-requires-dist "rocm-sdk-device-@GFXARCH@ == ${ROCM_SDK_VERSION}" \
--compression zstd \
-j "$(nproc)"
rm -f "${FAT_STAGE}"/torch-*.whl
mv ${{ env.PACKAGE_DIST_DIR }}/torchvision-*.whl "${FAT_STAGE}/"
python -m rocm_kpack.tools.split_python_wheels \
--input "${FAT_STAGE}"/torchvision-*.whl \
--output-dir "${{ env.PACKAGE_DIST_DIR }}" \
--device-package-prefix amd-torchvision-device \
--overlay-root torchvision/ \
--device-requires-dist "rocm-sdk-device-@GFXARCH@ == ${ROCM_SDK_VERSION}" \
--compression zstd \
-j "$(nproc)"
rm -rf "${FAT_STAGE}"
echo "Post-split contents of ${{ env.PACKAGE_DIST_DIR }}:"
ls -lh "${{ env.PACKAGE_DIST_DIR }}/"
# TODO(#5110): Once manifest-driven builds land, write these outputs from
# the manifest and keep wheel scanning as a validation step that the build
# produced the requested package versions.
- name: Write PyTorch package version outputs
id: pytorch-versions
run: |
python build_tools/github_actions/write_torch_versions.py \
--dist-dir ${{ env.PACKAGE_DIST_DIR }}
- name: Configure AWS Credentials
uses: ./.github/actions/configure_aws_artifacts_credentials
with:
release_type: ${{ inputs.release_type }}
- name: Upload PyTorch wheels to release bucket
id: publish-pytorch-wheels
run: |
python build_tools/github_actions/publish_pytorch_to_release_bucket.py \
--source-dir="${{ env.PACKAGE_DIST_DIR }}" \
--release-type="${{ inputs.release_type }}"
- name: Capture job summary
id: summary
if: always()
run: python build_tools/github_actions/capture_job_summary.py
configure_pytorch_tests:
name: Configure PyTorch Tests
runs-on: ubuntu-24.04
permissions:
contents: read
outputs:
enabled: ${{ steps.configure.outputs.enabled }}
matrix: ${{ steps.configure.outputs.matrix }}
summary: ${{ steps.summary.outputs.summary }}
steps:
- name: Checkout
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
repository: ${{ inputs.repository || github.repository }}
ref: ${{ inputs.ref }}
- name: Checkout CI config
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
repository: ROCm/therock-ci-config
ref: main
path: ci-config
continue-on-error: true
- name: Configure PyTorch test matrix
id: configure
env:
CI_CONFIG_PATH: ci-config
run: |
python build_tools/github_actions/configure_pytorch_test_matrix.py \
--build-amdgpu-families "${{ inputs.amdgpu_families }}" \
--test-amdgpu-families "${{ inputs.test_amdgpu_families }}" \
--platform linux
- name: Capture job summary
id: summary
if: always()
run: python build_tools/github_actions/capture_job_summary.py
test_pytorch_wheels:
name: Test | ${{ matrix.amdgpu_family }} | ${{ matrix.test_runs_on }}
if: ${{ needs.configure_pytorch_tests.outputs.enabled == 'true' }}
needs: [build_pytorch_wheels, configure_pytorch_tests]
permissions:
contents: read
strategy:
fail-fast: false
matrix: ${{ fromJSON(needs.configure_pytorch_tests.outputs.matrix) }}
uses: ./.github/workflows/test_pytorch_wheels.yml
secrets: inherit
with:
amdgpu_family: ${{ matrix.amdgpu_family }}
test_runs_on: ${{ matrix.test_runs_on }}
package_index_url: ${{ needs.build_pytorch_wheels.outputs.package_index_url }}
python_version: ${{ inputs.python_version }}
torch_version: ${{ needs.build_pytorch_wheels.outputs.torch_version }}
# TODO(#5110): pass manifest_url here and let test_pytorch_wheels.yml use
# checkout_from_manifest.py --no-hipify --no-submodules so tests use the
# exact source commit that produced the package under test.
pytorch_git_ref: ${{ inputs.pytorch_git_ref }}
repository: ${{ inputs.repository || github.repository }}
ref: ${{ inputs.ref }}
# Dispatch the full PyTorch test suite for this exact build. Because this
# workflow builds a single (python_version, pytorch_git_ref), the build job's
# torch_version output is unambiguous and only this build's success gates the
# dispatch -- a failure for one ref no longer skips testing the others (#5777).
# Scoped to gfx94X-dcgpu Python 3.12 builds; opt in via run_full_pytorch_tests.
dispatch_pytorch_wheels_full_test:
name: Dispatch PyTorch Full Test | torch ${{ inputs.pytorch_git_ref }}
if: >-
${{
inputs.run_full_pytorch_tests &&
inputs.python_version == '3.12' &&
contains(inputs.amdgpu_families, 'gfx94X-dcgpu')
}}
needs: [build_pytorch_wheels]
runs-on: ubuntu-24.04
permissions:
actions: write
contents: read
env:
PYTORCH_GIT_REF: ${{ inputs.pytorch_git_ref }}
steps:
- name: Check full test cadence
id: cadence
# Release branches test daily; the nightly branch tests weekly (Sunday).
run: |
if [ "${PYTORCH_GIT_REF}" != "nightly" ] || [ "$(date -u +%u)" = "7" ]; then
echo "dispatch=true" >> "$GITHUB_OUTPUT"
else
echo "dispatch=false" >> "$GITHUB_OUTPUT"
fi
- name: Dispatch full PyTorch test workflow
if: ${{ steps.cadence.outputs.dispatch == 'true' }}
uses: benc-uk/workflow-dispatch@31e2b3319479a63f0ab15bf800eff9e913504e26 # v1.3.2
with:
workflow: test_pytorch_wheels_full.yml
inputs: |
{ "amdgpu_family": "gfx94X-dcgpu",
"test_runs_on": "linux-gfx942-1gpu-ccs-csp-ossci-rocm",
"test_runs_on_multi_gpu": "linux-gfx942-8gpu-ossci-rocm",
"package_index_url": "${{ needs.build_pytorch_wheels.outputs.package_index_url }}",
"python_version": "${{ inputs.python_version }}",
"torch_version": "${{ needs.build_pytorch_wheels.outputs.torch_version }}",
"pytorch_git_ref": "${{ inputs.pytorch_git_ref }}",
"pytorch_git_repo": "${{ inputs.pytorch_git_ref == 'nightly' && 'pytorch/pytorch' || 'ROCm/pytorch' }}",
"test_configs": "default distributed inductor",
"tests_to_include": "",
"repository": "${{ inputs.repository || github.repository }}",
"ref": "${{ inputs.ref }}"
}
notify_quartz_completed:
if: ${{ always() && vars.NOTIFY_QUARTZ_ENABLED == 'true' }}
needs: [build_pytorch_wheels, configure_pytorch_tests, test_pytorch_wheels, dispatch_pytorch_wheels_full_test]
name: "Quartz - completed - build PyTorch wheels"
runs-on: ubuntu-24.04
timeout-minutes: 10
continue-on-error: true
steps:
- name: "Quartz - completed - build PyTorch wheels"
uses: ROCm/Quartz/.github/actions/notify_quartz@f386a9756620938616af0b4d5d04b24ae6e0353f # notify_quartz/v1.1.1
with:
run_phase: completed
reporting_workflow: multi_arch_build_portable_linux_pytorch_wheels.yml
workflow_inputs: ${{ toJSON(inputs) }}
workflow_captured_outputs: ${{ toJSON(needs) }}
gh_app_client_id: ${{ secrets.GH_APP_HAULY_CID }}
gh_app_private_key: ${{ secrets.GH_APP_HAULY_PRIVATE_KEY }}