ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Version downpairing for Radeon+ROCm PyTorch wheels (#5353)

Browse source 
* Version downpairing for Radeon+ROCm PyTorch wheels

The radeon repo has often published the absolute highest version for torch, torchvision, and torchaudio independently without forming a proper "trio" of versions that matches the usual "law of five," which can result in a mismatched trio that fails the sanity check, forcing an unnecessary fallback to the standard ROCm index.

To prevent this issue, this commit:
- Refactors _pick_radeon_wheel to support optional version prefix filtering.
- Implements a two-pass resolution:
  1. Identifies the highest available versions for all packages.
  2. Calculates the lowest common denominator minor version.
  3. Re-picks wheels to ensure torch/audio match and vision is +15.
- Prevents silent fallbacks when compatible sets exist but are not the absolute latest entries in the listing.

* Apply Codex + Gemini suggestions

From Codex:

When the computed target minor is not actually present for one of the packages, _pick_radeon_wheel returns non-zero and this assignment runs under the script's top-level set -e, so the installer exits immediately instead of reaching the existing fallback path. This can happen with gapped Radeon listings, e.g. latest torch/vision imply minor 11, latest torchaudio implies minor 10, but the listing only has a complete older 2.9 trio or no torch 2.10 wheel; the new downpair code then aborts on this line rather than warning and falling back.

From Gemini:
medium
The local keyword is not part of the POSIX shell standard. While many modern shells like bash, zsh, and dash support it, this script uses #!/bin/sh and explicitly aims for POSIX compliance (as noted in the comments around line 1671) to ensure portability across different environments like minimal Docker images, BSD, or BusyBox.

Since the _extract_version function is called within a command substitution subshell (e.g., _torch_ver=$(_extract_version ...)), the variables defined inside it are already isolated from the parent shell's environment. Therefore, local is redundant here and should be removed to maintain portability and consistency with the rest of the script.

* Better implementation of Codex's suggestion

When the first computed target minor is not present for a package, this now clears the wheel and the final check falls back immediately, even if the listing contains a complete older trio. For example, with latest torch/vision implying minor 11, latest torchaudio implying 10, but no torch 2.10/vision 0.25 wheels and a complete 2.9/0.24/2.9 set, the new || _torch_whl="" avoids the earlier abort but still never searches below minor 10, so Radeon installs are skipped despite a compatible set being available.

* Second attempt to better implement Codex's suggestion

When the latest package set is mismatched but the downpair loop finds a lower minor, this accepts any nonempty torch/vision/audio wheels for that minor without verifying the full public versions. In a listing such as torch 2.10.1 plus torchaudio 2.10.0 and torchvision 0.25.0, the new loop marks the trio compatible and installs it, even though the change is meant to avoid unsupported torch/audio mismatches; compare the versions after repicking before setting _radeon_versions_match=true.

* Fix regression from trying to implement Codex

* Apply suggestion from @gemini-code-assist[bot]

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Clear triton when accepting a downpaired set

From Codex:

When this loop downpairs torch/vision/audio to an older minor, _tri_whl remains the absolute newest triton selected before the loop. In the multi-generation Radeon listings this block is meant to handle, that can skip a newer torch generation but still install its newer triton wheel via the later install triton + PyTorch command, producing a mismatched triton/PyTorch set instead of the lower generation's matching triton. Re-pick or clear triton when accepting the downpaired trio.

* Remove trailing whitespace on blank lines for PR #5353

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
This commit is contained in:
Jaeic Lee 2026-05-30 14:30:09 +09:00 • committed by GitHub
commit 187c81f409
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 77 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

117
install.sh
View file

@ -1785,9 +1785,9 @@ print('cp{}{}'.format(sys.version_info.major, sys.version_info.minor))
}
_pick_radeon_wheel() {
# Usage: _pick_radeon_wheel PACKAGE_NAME
# Usage: _pick_radeon_wheel PACKAGE_NAME [VERSION_PREFIX]
# Scans $_RADEON_LISTING for the newest wheel whose filename starts exactly
# with PACKAGE_NAME- and matches _RADEON_PYTAG + linux_x86_64.
# with PACKAGE_NAME- (and optionally VERSION_PREFIX) and matches _RADEON_PYTAG + linux_x86_64.
# Prints the full URL (resolving relative hrefs against _RADEON_BASE_URL).
#
# POSIX-compliant pipeline: all href parsing, filtering, and version
@ -1795,11 +1795,12 @@ _pick_radeon_wheel() {
# for GNU extensions (grep -o, sort -V) that would break under BSD
# or BusyBox coreutils.
_pkg="$1"
_ver_prefix="${2:-}"
[ -n "$_RADEON_LISTING" ] || return 1
[ -n "$_RADEON_PYTAG" ] || return 1
_tag="$_RADEON_PYTAG"
_href=$(printf '%s\n' "$_RADEON_LISTING" \
| awk -v pkg="$_pkg" -v tag="$_tag" '
| awk -v pkg="$_pkg" -v tag="$_tag" -v ver_prefix="$_ver_prefix" '
BEGIN { max_pad = ""; max_url = "" }
{
line = $0
@ -1813,7 +1814,7 @@ _pick_radeon_wheel() {
base = p[n]
sub(/[?#].*/, "", base)
prefix = pkg "-"
prefix = pkg "-" ver_prefix
# Match cpXY-cpXY or cpXY-abi3 with any linux x86_64
# platform tag (linux_x86_64, manylinux_2_28_x86_64,
# manylinux2014_x86_64, etc.)
@ -2112,24 +2113,23 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ "$_radeon_listing_ok" = true ]; then
# Require torch, torchvision, torchaudio wheels to all resolve
# from the Radeon listing. If any is missing for this Python
# tag, fall through to the standard ROCm index instead of
# silently mixing Radeon wheels with PyPI defaults.
# from the Radeon listing. The repo often publishes multiple
# generations simultaneously, so picking the highest-version
# for each package independently can assemble a mismatched trio
# (e.g. torch 2.10 + torchvision 0.24). To prevent this,
# we identify the highest common minor version and downpair
# wheels if necessary to ensure a compatible set.
_torch_whl=$(_pick_radeon_wheel "torch" 2>/dev/null) || _torch_whl=""
_tv_whl=$(_pick_radeon_wheel "torchvision" 2>/dev/null) || _tv_whl=""
_ta_whl=$(_pick_radeon_wheel "torchaudio" 2>/dev/null) || _ta_whl=""
_tri_whl=$(_pick_radeon_wheel "triton" 2>/dev/null) || _tri_whl=""
# Sanity-check torch / torchvision / torchaudio are a
# matching release. The Radeon repo publishes multiple
# generations simultaneously, so picking the highest-version
# wheel for each package independently can assemble a
# mismatched trio (e.g. torch 2.9.1 + torchvision 0.23.0 +
# torchaudio 2.9.0 from the current rocm-rel-7.2.1 index).
# Check that torch and torchaudio share the same X.Y public
# version prefix, and that torchvision's minor correctly
# pairs with torch's minor (torchvision = torch.minor - 5
# pairs with torch's minor (torchvision = torch.minor + 15
# since torch 2.4 -> torchvision 0.19 -> torch 2.9 ->
# torchvision 0.24).
#
# URL-decode each wheel name so %2B -> + before version
# extraction. Real Radeon wheel hrefs are percent-encoded
# (torch-2.10.0%2Brocm7.2.0...), so a plain [+-] terminator
@ -2137,38 +2137,75 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# _radeon_versions_match would stay false for every real
# listing, silently forcing a fallback to the generic
# ROCm index.
_torch_ver=""
_tv_ver=""
_ta_ver=""
if [ -n "$_torch_whl" ]; then
_torch_name=$(printf '%s' "${_torch_whl##*/}" | sed 's/%2[Bb]/+/g')
_torch_ver=$(printf '%s\n' "$_torch_name" | sed -n 's|^torch-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
if [ -n "$_tv_whl" ]; then
_tv_name=$(printf '%s' "${_tv_whl##*/}" | sed 's/%2[Bb]/+/g')
_tv_ver=$(printf '%s\n' "$_tv_name" | sed -n 's|^torchvision-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
if [ -n "$_ta_whl" ]; then
_ta_name=$(printf '%s' "${_ta_whl##*/}" | sed 's/%2[Bb]/+/g')
_ta_ver=$(printf '%s\n' "$_ta_name" | sed -n 's|^torchaudio-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
_extract_version() {
_whl=$1
_pkg=$2
if [ -n "$_whl" ]; then
_name=$(printf '%s' "${_whl##*/}" | sed 's/%2[Bb]/+/g')
printf '%s\n' "$_name" | sed -n "s|^${_pkg}-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p"
fi
}
_torch_ver=$(_extract_version "$_torch_whl" "torch")
_tv_ver=$(_extract_version "$_tv_whl" "torchvision")
_ta_ver=$(_extract_version "$_ta_whl" "torchaudio")
_radeon_versions_match=false
if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
_torch_major=${_torch_ver%%.*}
_torch_minor=${_torch_ver#*.}
_ta_major=${_ta_ver%%.*}
_ta_minor=${_ta_ver#*.}
_tv_major=${_tv_ver%%.*}
_tv_minor=${_tv_ver#*.}
# torchvision expected minor (e.g. torch 2.9 -> 0.24)
_expected_tv_minor=$((_torch_minor + 15))
if [ "$_torch_major" = "$_ta_major" ] && \
[ "$_torch_minor" = "$_ta_minor" ] && \
[ "$_tv_major" = "0" ] && \
[ "$_tv_minor" = "$_expected_tv_minor" ]; then
_radeon_versions_match=true
fi
_tv_equiv_minor=$((_tv_minor - 15))
# Determine initial target minor (lowest common denominator)
_target_minor=$_torch_minor
[ "$_tv_equiv_minor" -lt "$_target_minor" ] && _target_minor=$_tv_equiv_minor
[ "$_ta_minor" -lt "$_target_minor" ] && _target_minor=$_ta_minor
# Loop downwards to find the first complete matching trio.
# This avoids aborting if the repo has gaps.
_attempts=0
while [ "$_attempts" -lt 5 ] && [ "$_target_minor" -ge 0 ]; do
_expected_tv_minor=$((_target_minor + 15))
_curr_torch=$(_pick_radeon_wheel "torch" "2.${_target_minor}." 2>/dev/null) || _curr_torch=""
_curr_tv=$(_pick_radeon_wheel "torchvision" "0.${_expected_tv_minor}." 2>/dev/null) || _curr_tv=""
_curr_ta=$(_pick_radeon_wheel "torchaudio" "2.${_target_minor}." 2>/dev/null) || _curr_ta=""
if [ -n "$_curr_torch" ] && [ -n "$_curr_tv" ] && [ -n "$_curr_ta" ]; then
# Extract versions from the wheels found in this iteration
_c_torch_ver=$(_extract_version "$_curr_torch" "torch")
_c_tv_ver=$(_extract_version "$_curr_tv" "torchvision")
_c_ta_ver=$(_extract_version "$_curr_ta" "torchaudio")
# Parse Major.Minor for validation
_c_torch_major=${_c_torch_ver%%.*}
_c_torch_minor=${_c_torch_ver#*.}
_c_ta_major=${_c_ta_ver%%.*}
_c_ta_minor=${_c_ta_ver#*.}
_c_tv_major=${_c_tv_ver%%.*}
_c_tv_minor=${_c_tv_ver#*.}
# Strict X.Y validation: allow patch versions to differ (e.g. torch 2.9.1 + vision 0.24.0)
# as long as the Major and Minor pairing is correct.
if [ "$_c_torch_major" = "$_c_ta_major" ] && \
[ "$_c_torch_minor" = "$_c_ta_minor" ] && \
[ "$_c_tv_major" = "0" ] && \
[ "$_c_tv_minor" = "$((_c_torch_minor + 15))" ]; then
_torch_whl=$_curr_torch
_tv_whl=$_curr_tv
_ta_whl=$_curr_ta
_tri_whl=""
_radeon_versions_match=true
break
fi
fi
_target_minor=$((_target_minor - 1))
_attempts=$((_attempts + 1))
done
fi
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"