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

install: digit-gate rocm leaves in marker normalization and ROCm side effects

Browse source 
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 2026-07-12 12:14:15 +00:00
commit 42ef2f4b6a
5 changed files with 118 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

12
studio/install_python_stack.py
View file

@ -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