From 0167b5d72b632c0e276d7551deee2d2d678215c2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 02:39:54 +0000 Subject: [PATCH] torch2110 extras: gate x86-64 xformers wheels on platform_machine The cuNNN torch2110/torch2100 extras pinned direct xformers wheel URLs that are x86_64/win_amd64-only while their markers checked only the operating system, so on Linux AArch64 (DGX Spark/GB10) or Windows ARM64 the marker matched and pip aborted on an unsupported wheel even though the torch 2.11/2.10 aarch64 CUDA wheels themselves exist. Append the same platform_machine guard the file already uses for its other direct-URL wheels (bitsandbytes, triton), so ARM64 installs resolve the trio and simply skip xformers. Also in _auto_install.py: torch 2.10.1 fell into the new <2.11.1 branch and emitted the torch2110 extra, whose exact 2.11.0 pins would replace the detected runtime; restore the pre-PR loud RuntimeError for the 2.10.1..<2.11.0 gap, and report the actual torch version in the CUDA requirement error instead of a hard-coded 2.10. Tests: _reqs() now keeps both platform wheel entries per extra (the Windows entry previously shadowed Linux), asserts the ARM64 exclusions, and covers all six public wrapper extras. --- pyproject.toml | 24 ++++++++-------- tests/test_torch2110_cuda_extras.py | 44 ++++++++++++++++++++++++----- unsloth/_auto_install.py | 3 +- 3 files changed, 51 insertions(+), 20 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c6a8ff6b99..927cf7612c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -326,16 +326,16 @@ cu130onlytorch291 = [ "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", ] cu126onlytorch2100 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch2100 = [ - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu130onlytorch2100 = [ - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu126onlytorch2110 = [ # xformers 0.0.35 does not hard-pin torch the way 0.0.34 pinned torch==2.10.0, @@ -349,8 +349,8 @@ cu126onlytorch2110 = [ "torch==2.11.0+cu126", "torchvision==0.26.0+cu126", "torchaudio==2.11.0+cu126", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.35-py39-none-win_amd64.whl ; (sys_platform == 'win32')", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.35-py39-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch2110 = [ # +cu128 local build for the same reason as cu126onlytorch2110 above: torch @@ -359,8 +359,8 @@ cu128onlytorch2110 = [ "torch==2.11.0+cu128", "torchvision==0.26.0+cu128", "torchaudio==2.11.0+cu128", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.35-py39-none-win_amd64.whl ; (sys_platform == 'win32')", + "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.35-py39-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu130onlytorch2110 = [ # Pinned to the +cu130 local build like the cu126/cu128 extras: a bare range @@ -373,8 +373,8 @@ cu130onlytorch2110 = [ "torch==2.11.0+cu130", "torchvision==0.26.0+cu130", "torchaudio==2.11.0+cu130", - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-win_amd64.whl ; (sys_platform == 'win32')", + "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118 = [ "unsloth[huggingface]", diff --git a/tests/test_torch2110_cuda_extras.py b/tests/test_torch2110_cuda_extras.py index 08976e2f50..5dbf506378 100644 --- a/tests/test_torch2110_cuda_extras.py +++ b/tests/test_torch2110_cuda_extras.py @@ -32,11 +32,14 @@ def _extra(name: str) -> list[str]: return data["project"]["optional-dependencies"][name] -def _reqs(specs: list[str]) -> dict[str, Requirement]: - out = {} +def _reqs(specs: list[str]) -> dict[str, list[Requirement]]: + # Keyed by name -> list: each extra carries one Linux and one Windows + # xformers requirement, so a plain name -> Requirement dict would silently + # drop the Linux entry. + out: dict[str, list[Requirement]] = {} for spec in specs: r = Requirement(spec) - out[r.name.lower()] = r + out.setdefault(r.name.lower(), []).append(r) return out @@ -47,10 +50,37 @@ def test_cuda12_torch2110_pins_matching_local_build(cuda: str): # CUDA-13 default on PyPI. reqs = _reqs(_extra(f"{cuda}onlytorch2110")) for pkg in _TORCH_TRIO: - spec = str(reqs[pkg].specifier) + (req,) = reqs[pkg] + spec = str(req.specifier) assert ( spec == f"=={('2.11.0' if pkg != 'torchvision' else '0.26.0')}+{cuda}" ), f"{cuda}onlytorch2110: {pkg} pinned as '{spec}', expected the +{cuda} local build" - # xformers must come from the same CUDA index. - xf = reqs["xformers"] - assert xf.url and f"/whl/{cuda}/" in xf.url, f"xformers not on the {cuda} index: {xf.url}" + # Both xformers wheels (Linux and Windows) must come from the same CUDA index. + xformers = reqs["xformers"] + assert len(xformers) == 2, f"expected Linux + Windows xformers wheels, got {xformers}" + linux = [r for r in xformers if r.url and r.url.endswith("manylinux_2_28_x86_64.whl")] + windows = [r for r in xformers if r.url and r.url.endswith("win_amd64.whl")] + assert len(linux) == 1 and len(windows) == 1, f"unexpected xformers wheels: {xformers}" + for r in linux + windows: + assert f"/whl/{cuda}/xformers-0.0.35-" in r.url, f"xformers not on the {cuda} index: {r.url}" + # The wheels are x86-64 only, so the markers must exclude other machines + # (e.g. Linux aarch64 such as GB200/DGX Spark, Windows ARM64) where the + # torch trio resolves fine but these wheels would abort the install. + assert r.marker is not None + assert not r.marker.evaluate({"sys_platform": "linux", "platform_machine": "aarch64"}) + assert not r.marker.evaluate({"sys_platform": "win32", "platform_machine": "ARM64"}) + assert linux[0].marker.evaluate({"sys_platform": "linux", "platform_machine": "x86_64"}) + assert windows[0].marker.evaluate({"sys_platform": "win32", "platform_machine": "AMD64"}) + + +@pytest.mark.parametrize("cuda", ["cu126", "cu128", "cu130"]) +@pytest.mark.parametrize("variant", ["", "ampere-"]) +def test_torch2110_wrapper_references_matching_leaf(cuda: str, variant: str): + # The six public wrappers must pull in the usual huggingface + bitsandbytes + # pair and reference the internal leaf of the SAME CUDA version. + specs = _extra(f"{cuda}-{variant}torch2110") + assert specs == [ + "unsloth[huggingface]", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + f"unsloth[{cuda}onlytorch2110]", + ] diff --git a/unsloth/_auto_install.py b/unsloth/_auto_install.py index 0518d75a74..a7b29b4616 100644 --- a/unsloth/_auto_install.py +++ b/unsloth/_auto_install.py @@ -36,10 +36,11 @@ elif v < V('2.8.9'): x = 'cu{}{}-torch280' elif v < V('2.9.1'): x = 'cu{}{}-torch290' elif v < V('2.9.2'): x = 'cu{}{}-torch291' elif v < V('2.10.1'): x = 'cu{}{}-torch2100' +elif v < V('2.11.0'): raise RuntimeError(f"Torch = {v} not supported!") elif v < V('2.11.1'): x = 'cu{}{}-torch2110' else: raise RuntimeError(f"Torch = {v} too new!") if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8", "13.0"): raise RuntimeError(f"CUDA = {cuda} not supported!") -if v >= V('2.10.0') and cuda not in ("12.6", "12.8", "13.0"): raise RuntimeError(f"Torch 2.10 requires CUDA 12.6, 12.8, or 13.0! Got CUDA = {cuda}") +if v >= V('2.10.0') and cuda not in ("12.6", "12.8", "13.0"): raise RuntimeError(f"Torch = {v} requires CUDA 12.6, 12.8, or 13.0! Got CUDA = {cuda}") x = x.format(cuda.replace(".", ""), "-ampere" if False else "") # is_ampere is broken due to flash-attn # The CUDA torch2110 extras pin the torch trio to the exact +cuNNN local build, # which only resolves from the matching PyTorch CUDA index (a bare spec would let