install: digit-gate rocm leaves in marker normalization and ROCm side effects
Round 5 of the pinned-index hardening (three custom-rocm-leaf edge cases): 1. _normalize_family_leaf lowercased every leaf starting with rocm, so a custom mirror leaf like rocm-Current compared equal to its lowercase form and a case-only pin change was skipped. URL paths can be case-sensitive. The rocm prefix is now digit-gated (rocm[0-9]*, matching _is_pip_rocm_family_leaf) in install.sh, setup.ps1 and install_python_stack.py, so only true family leaves (rocm7.2) are lowercased; a custom rocm-* leaf keeps its case. 2. setup.ps1 Test-MarkerPinMismatch compared normalized URLs with -ne, which is case-insensitive in PowerShell, so a case-only marker change (Simple vs simple) was treated as matching and the reinstall skipped. Now -cne. 3. install.sh gated the AMD bitsandbytes install and the "repair ROCm torch" --default-index reinstall on a bare whole-URL rocm glob, so a custom CPU/CUDA/private index whose leaf merely starts with rocm (rocm-current) was force-repaired from the wrong ROCm-only path whenever torch.version.hip was empty. Both now gate on _torch_index_is_rocm_family, computed once from the digit-gated leaf (rocm[0-9]*/gfx*). Tests: 4 new parity assertions plus 2 case-sensitivity marker checks.
This commit is contained in:
Daniel Han
parent
5b4003d6f6
commit
42ef2f4b6a
5 changed files with 118 additions and 54 deletions
97
install.sh
97
install.sh
|
|
@ -2267,13 +2267,15 @@ _torch_index_repairable() {
|
|||
# <venv_dir>/.unsloth-torch-index (single line = the resolved index URL)
|
||||
_TORCH_INDEX_MARKER_NAME=".unsloth-torch-index"
|
||||
|
||||
# Lowercase ONLY a known wheel-family leaf (rocm* / gfx* / cpu / cuXXX); a custom
|
||||
# mirror leaf keeps its case so a verbatim URL pin is not falsely matched equal.
|
||||
# Lowercase ONLY a known wheel-family leaf (rocm<digit>* / gfx* / cpu / cuXXX); a
|
||||
# custom mirror leaf keeps its case so a verbatim URL pin is not falsely matched
|
||||
# equal. The rocm prefix is digit-gated: rocm7.2 is a family leaf, but a
|
||||
# rocm-rel-7.2.1 / rocm-Current leaf is a verbatim pin whose case must survive.
|
||||
# Mirrors _normalize_family_leaf in install_python_stack.py / setup.ps1.
|
||||
_normalize_family_leaf() {
|
||||
_l_low=$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')
|
||||
case "$_l_low" in
|
||||
rocm*|gfx*|cpu|cu[0-9]*) printf '%s' "$_l_low" ;;
|
||||
rocm[0-9]*|gfx*|cpu|cu[0-9]*) printf '%s' "$_l_low" ;;
|
||||
*) printf '%s' "$1" ;;
|
||||
esac
|
||||
}
|
||||
|
|
@ -2620,6 +2622,19 @@ case "$_torch_index_leaf" in
|
|||
*) unset UNSLOTH_TORCH_BACKEND ;;
|
||||
esac
|
||||
|
||||
# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm<digit>* / gfx*),
|
||||
# gating the ROCm-only side effects below: AMD bitsandbytes and the "repair ROCm
|
||||
# torch" --default-index reinstall. Digit-gated like _is_pip_rocm_family_leaf so a
|
||||
# custom CPU/CUDA/private index whose leaf merely STARTS with "rocm" (e.g.
|
||||
# /rocm-current, or a repo.radeon.com/.../rocm-rel-7.2.1 find-links leaf) is NOT
|
||||
# force-repaired from the wrong ROCm-only path when torch.version.hip is empty.
|
||||
# A bare */rocm* whole-URL glob would misfire on those. gfx* is always a family.
|
||||
# _torch_index_leaf is already lowercased above.
|
||||
case "$_torch_index_leaf" in
|
||||
rocm[0-9]*|gfx*) _torch_index_is_rocm_family=true ;;
|
||||
*) _torch_index_is_rocm_family=false ;;
|
||||
esac
|
||||
|
||||
# rocm7.2 and the AMD per-gfx indexes with the torch._C._grouped_mm bug on <2.11
|
||||
# (repo.amd.com/.../gfx120X-all, gfx1151, gfx1150) ship torch 2.11.0 -- raise the
|
||||
# constraint to allow it. This also covers a pinned full-URL or family override
|
||||
|
|
@ -2926,24 +2941,20 @@ if [ "$_MIGRATED" = true ]; then
|
|||
# AMD ROCm: install bitsandbytes even in migrated environments so
|
||||
# existing ROCm installs gain the AMD bitsandbytes build without a
|
||||
# fresh reinstall.
|
||||
if [ "$SKIP_TORCH" = false ]; then
|
||||
case "$TORCH_INDEX_URL" in
|
||||
*/rocm*|*/gfx*)
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
# Repair ROCm torch if overwritten during migrated install
|
||||
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
|
||||
if [ -z "$_has_hip" ]; then
|
||||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
# torch was actually reinstalled from $TORCH_INDEX_URL now, so the
|
||||
# marker should record it (the preserved-torch case above must not).
|
||||
_TORCH_INSTALLED_THIS_RUN=true
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
# Repair ROCm torch if overwritten during migrated install
|
||||
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
|
||||
if [ -z "$_has_hip" ]; then
|
||||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
# torch was actually reinstalled from $TORCH_INDEX_URL now, so the
|
||||
# marker should record it (the preserved-torch case above must not).
|
||||
_TORCH_INSTALLED_THIS_RUN=true
|
||||
fi
|
||||
fi
|
||||
elif [ -n "$TORCH_INDEX_URL" ]; then
|
||||
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
|
||||
|
|
@ -3114,12 +3125,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
|
||||
# host stays in GGUF-only mode rather than pulling in bitsandbytes,
|
||||
# which is only useful once torch is present for training.
|
||||
if [ "$SKIP_TORCH" = false ]; then
|
||||
case "$TORCH_INDEX_URL" in
|
||||
*/rocm*|*/gfx*)
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
;;
|
||||
esac
|
||||
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
fi
|
||||
# Fresh: Step 2 - install unsloth, preserving pre-installed torch
|
||||
tauri_log "STEP" "Installing Unsloth"
|
||||
|
|
@ -3160,26 +3167,22 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
fi
|
||||
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
|
||||
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
|
||||
if [ "$SKIP_TORCH" = false ]; then
|
||||
case "$TORCH_INDEX_URL" in
|
||||
*/rocm*|*/gfx*)
|
||||
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
|
||||
if [ -z "$_has_hip" ]; then
|
||||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
# The repair reinstalled torch from $TORCH_INDEX_URL (the generic
|
||||
# ROCm index), not the Radeon --find-links repo, so record THAT as
|
||||
# the marker source. A Radeon --find-links install set
|
||||
# _TORCH_MARKER_INDEX_URL to its repo.radeon.com base earlier;
|
||||
# leaving it would make the marker misreport Radeon wheels and let a
|
||||
# later Radeon pin compare-equal and skip a needed reinstall.
|
||||
_TORCH_MARKER_INDEX_URL="$TORCH_INDEX_URL"
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
|
||||
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
|
||||
if [ -z "$_has_hip" ]; then
|
||||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
# The repair reinstalled torch from $TORCH_INDEX_URL (the generic
|
||||
# ROCm index), not the Radeon --find-links repo, so record THAT as
|
||||
# the marker source. A Radeon --find-links install set
|
||||
# _TORCH_MARKER_INDEX_URL to its repo.radeon.com base earlier;
|
||||
# leaving it would make the marker misreport Radeon wheels and let a
|
||||
# later Radeon pin compare-equal and skip a needed reinstall.
|
||||
_TORCH_MARKER_INDEX_URL="$TORCH_INDEX_URL"
|
||||
fi
|
||||
fi
|
||||
else
|
||||
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
|
||||
|
|
|
|||
|
|
@ -138,17 +138,19 @@ _TORCH_INDEX_MARKER_NAME = ".unsloth-torch-index"
|
|||
|
||||
|
||||
def _normalize_family_leaf(leaf: str) -> str:
|
||||
"""Lowercase ONLY a known wheel-family leaf (rocm* / gfx* / cpu / cuXXX).
|
||||
"""Lowercase ONLY a known wheel-family leaf (rocm<digit>* / gfx* / cpu / cuXXX).
|
||||
|
||||
The canonical gfx120X-all (capital X) must match AMD's lowercase gfx120x-all, so
|
||||
known-family leaves are lowercased. A custom mirror leaf (/Current, /simple, ...)
|
||||
keeps its case: an unknown-family URL pin is applied verbatim, so /Current and
|
||||
/current must NOT compare equal. Same known-family set as
|
||||
_explicit_unknown_family_torch_index_url. Mirrors the gate in install.sh /
|
||||
setup.ps1. Pure function.
|
||||
/current must NOT compare equal. The rocm prefix is digit-gated like
|
||||
_is_pip_rocm_family_leaf: rocm7.2 is a family leaf, but rocm-rel-7.2.1 /
|
||||
rocm-Current are verbatim pins whose case must survive normalization (URL
|
||||
paths can be case-sensitive). Mirrors the gate in install.sh / setup.ps1.
|
||||
Pure function.
|
||||
"""
|
||||
low = leaf.lower()
|
||||
if low.startswith(("rocm", "gfx")) or low == "cpu" or re.match(r"^cu[0-9]", low):
|
||||
if low.startswith("gfx") or low == "cpu" or re.match(r"^(rocm|cu)[0-9]", low):
|
||||
return low
|
||||
return leaf
|
||||
|
||||
|
|
|
|||
|
|
@ -442,7 +442,7 @@ $TorchIndexMarkerName = ".unsloth-torch-index"
|
|||
function Get-NormalizedFamilyLeaf {
|
||||
param([string]$Leaf)
|
||||
$low = $Leaf.ToLowerInvariant()
|
||||
if ($low -match '^(rocm|gfx)' -or $low -eq 'cpu' -or $low -match '^cu[0-9]') { return $low }
|
||||
if ($low -match '^(rocm[0-9]|gfx)' -or $low -eq 'cpu' -or $low -match '^cu[0-9]') { return $low }
|
||||
return $Leaf
|
||||
}
|
||||
|
||||
|
|
@ -507,7 +507,10 @@ function Test-MarkerPinMismatch {
|
|||
param([string]$VenvDir, [string]$PinUrl)
|
||||
$marker = Read-TorchIndexMarker -VenvDir $VenvDir
|
||||
if ($null -eq $marker) { return $null }
|
||||
return (Get-NormalizedIndexUrl $PinUrl) -ne (Get-NormalizedIndexUrl $marker)
|
||||
# -cne, not -ne: normalization intentionally preserves unknown-leaf case (URL
|
||||
# paths can be case-sensitive), so a case-only custom-index change (/Simple ->
|
||||
# /simple) must count as a mismatch. PowerShell's -ne is case-insensitive.
|
||||
return (Get-NormalizedIndexUrl $PinUrl) -cne (Get-NormalizedIndexUrl $marker)
|
||||
}
|
||||
|
||||
# The AMD per-arch index leaves that need the torch 2.11 floor (the _grouped_mm
|
||||
|
|
|
|||
|
|
@ -508,6 +508,51 @@ class TestPinnedRocmLeafDigitParity:
|
|||
r'r"\^rocm\\d"', text
|
||||
), "install_python_stack.py _is_pip_rocm_family_leaf must match ^rocm\\d"
|
||||
|
||||
def test_normalize_family_leaf_digit_gates_rocm(self):
|
||||
"""_normalize_family_leaf lowercases only true family leaves. rocm7.2 is a
|
||||
family (lowercased), but rocm-Current / rocm-rel-7.2.1 keep their case so a
|
||||
case-only custom-index change is not falsely matched equal (URL paths can be
|
||||
case-sensitive). All three installers must digit-gate the rocm prefix."""
|
||||
sh = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert re.search(r"rocm\[0-9\]\*\|gfx\*\|cpu\|cu\[0-9\]\*", sh), (
|
||||
"install.sh _normalize_family_leaf must digit-gate rocm (rocm[0-9]*)"
|
||||
)
|
||||
setup = SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "-match '^(rocm[0-9]|gfx)'" in setup, (
|
||||
"setup.ps1 Get-NormalizedFamilyLeaf must digit-gate rocm (^(rocm[0-9]|gfx))"
|
||||
)
|
||||
stack = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert re.search(r'r"\^\(rocm\|cu\)\[0-9\]"', stack), (
|
||||
"install_python_stack.py _normalize_family_leaf must digit-gate rocm "
|
||||
"(^(rocm|cu)[0-9])"
|
||||
)
|
||||
|
||||
def test_setup_ps1_marker_compare_is_case_sensitive(self):
|
||||
"""Test-MarkerPinMismatch must use -cne, not -ne: normalization preserves
|
||||
unknown-leaf case, so a case-only custom-index change (/Simple -> /simple)
|
||||
is a real mismatch that PowerShell's case-insensitive -ne would miss."""
|
||||
text = SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "(Get-NormalizedIndexUrl $PinUrl) -cne (Get-NormalizedIndexUrl $marker)" in text, (
|
||||
"setup.ps1 Test-MarkerPinMismatch must compare normalized URLs with -cne "
|
||||
"(case-sensitive) so a case-only custom-index change triggers reinstall"
|
||||
)
|
||||
|
||||
def test_install_sh_rocm_side_effects_digit_gated(self):
|
||||
"""The AMD bitsandbytes + 'repair ROCm torch' side effects must fire only on
|
||||
a real ROCm family (rocm[0-9]*/gfx*), not a bare */rocm* whole-URL glob that
|
||||
catches a custom CPU/CUDA index like /rocm-current and force-repairs it from
|
||||
the wrong --default-index."""
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert re.search(
|
||||
r"rocm\[0-9\]\*\|gfx\*\) _torch_index_is_rocm_family=true", text
|
||||
), "install.sh must set _torch_index_is_rocm_family from a digit-gated leaf"
|
||||
assert (
|
||||
'[ "$_torch_index_is_rocm_family" = true ]' in text
|
||||
), "install.sh ROCm bnb/repair hooks must gate on _torch_index_is_rocm_family"
|
||||
assert (
|
||||
'*/rocm*|*/gfx*)\n _install_bnb_rocm' not in text
|
||||
), "install.sh must not gate _install_bnb_rocm on a bare */rocm* whole-URL glob"
|
||||
|
||||
|
||||
class TestPinnedIndexClearsUvEnvParity:
|
||||
"""Every installer must neutralise the uv index env vars for a pinned torch
|
||||
|
|
|
|||
|
|
@ -81,6 +81,17 @@ try {
|
|||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://mirror.local/simple"
|
||||
Check "custom /simple marker vs /current pin -> mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://mirror.local/current") -eq $true)
|
||||
# Case-only custom-index change /Simple -> /simple is a mismatch: normalization
|
||||
# preserves unknown-leaf case (URL paths can be case-sensitive) and the compare
|
||||
# is -cne. A case-insensitive -ne would wrongly skip the needed reinstall.
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://mirror.local/Simple"
|
||||
Check "custom /Simple marker vs /simple pin -> mismatch (case-sensitive)" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://mirror.local/simple") -eq $true)
|
||||
# A genuine family leaf differing only in case is NOT a mismatch (rocm7.2 / gfx
|
||||
# leaves are lowercased by normalization, so gfx120X-all == gfx120x-all).
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all"
|
||||
Check "family gfx120X-all marker vs gfx120x-all pin -> no mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://repo.amd.com/rocm/whl/gfx120x-all") -eq $false)
|
||||
|
||||
Write-Host "Read-TorchIndexMarker (missing / empty -> null)"
|
||||
Remove-Item -LiteralPath (Get-TorchIndexMarkerPath -VenvDir $venv) -Force
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue