From 42ef2f4b6aa93ce2d0968c9dcf2afb466cee44a6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 12:14:15 +0000 Subject: [PATCH] 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. --- install.sh | 97 +++++++++++----------- studio/install_python_stack.py | 12 +-- studio/setup.ps1 | 7 +- tests/python/test_cross_platform_parity.py | 45 ++++++++++ tests/studio/test_torch_index_marker.ps1 | 11 +++ 5 files changed, 118 insertions(+), 54 deletions(-) diff --git a/install.sh b/install.sh index 67ae1e4eba..50cc7e15db 100755 --- a/install.sh +++ b/install.sh @@ -2267,13 +2267,15 @@ _torch_index_repairable() { # /.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* / 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* / 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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 14bd465093..8c18d4753e 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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* / 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 diff --git a/studio/setup.ps1 b/studio/setup.ps1 index a725b83c8c..f068230910 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -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 diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index c9b0646cc0..7a263364c3 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -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 diff --git a/tests/studio/test_torch_index_marker.ps1 b/tests/studio/test_torch_index_marker.ps1 index 896fc618f9..12a4271fca 100644 --- a/tests/studio/test_torch_index_marker.ps1 +++ b/tests/studio/test_torch_index_marker.ps1 @@ -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