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

Studio: fix flash-attn and torchao install on Blackwell (sm_100+) GPUs (Closes #6961) (#6970)

Browse source 
* fix: Remove moot has_blackwell_gpu() function

Fixes unslothai/unsloth#6961. This function skipped flash-attn on Blackwell GPUs because no prebuilt wheel existed;
Dao-AILab now ships one and url_exists() already gates resolution.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix: use torchao 0.17.0 for Blackwell

Fixes #6961. Torchao 0.16.0's cpp extensions are built against CUDA 12, so on a CUDA-13
torch (cu130 / Blackwell) they fail to load with "libcudart.so.12: cannot
open shared object file". Select 0.17.0 there instead: its cpp targets torch
2.11, so it is skipped cleanly rather than crashing. CUDA-12 / ROCm / CPU
torch 2.10 keeps 0.16.0 and its working kernels.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Condense torchao version-selection comments (no behavior change)

* Support torch 2.11 in the Studio installer via the torch2.10 prebuilt wheels

Map torch 2.11 to the torch2.10 prebuilt wheels for flash-attn, causal-conv1d,
and mamba through wheel_utils.prebuilt_wheel_torch_mm, applied in direct_wheel_url
(filename) and flash_attn_wheel_url (version). Those torch2.10 CUDA wheels load and
pass each project's own test suite on torch 2.11 (verified on B200), so a torch 2.11
environment gets the prebuilt accelerators instead of skipping or building from source.

Raise _CUDA_TORCH_PKG_SPEC to <2.12.0 (torchvision <0.27.0, torchaudio <2.12.0) so
the CUDA torch repair path can install torch 2.11, where torchao 0.17's cpp kernels
load cleanly. Add tests for the mapping.

* Keep has_blackwell_gpu as a False stub for future arch gating

* Restore has_blackwell_gpu as a return-False probe kept for future arch gating

Keep the nvidia-smi compute_cap detection and its two call sites, but short-circuit
with return False at the top so flash-attn is no longer skipped on Blackwell (sm_100+
now has prebuilt wheels and url_exists gates resolution). Drop the early return to
re-enable arch-based detection later.

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Thomas Eric 🇧🇷 2026-07-08 10:38:10 -03:00 • committed by GitHub
commit 03cbe211a3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 127 additions and 221 deletions
Download patch file Download diff file Expand all files Collapse all files

70
studio/install_python_stack.py
View file

@ -103,33 +103,47 @@ _PYTORCH_WHL_BASE = (
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
).rstrip("/")
# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are
# pinned to the torch<2.11 family rather than left bare: the install uses an
# exclusive --index-url (no PyPI fallback), so a bare name could resolve a
# torchvision built against a different torch major (e.g. 0.27 for torch 2.12)
# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm
# spec above, which targets the same torch family.
# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed: its
# torchao 0.17 cpp kernels load cleanly (0.16 crashes on cu130), and the flash-attn
# / causal-conv1d / mamba torch2.10 wheels load and pass their upstream suites on
# 2.11 (see wheel_utils._PREBUILT_WHEEL_TORCH_MM). torchvision/torchaudio are pinned
# (not bare) because the install uses an exclusive --index-url (no PyPI fallback), so
# a bare name could resolve one built against a different torch major (e.g. 0.27 for
# torch 2.12) and fail at runtime with an ABI mismatch.
_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
"torch>=2.4,<2.11.0",
"torchvision>=0.19,<0.26.0",
"torchaudio>=2.4,<2.11.0",
"torch>=2.4,<2.12.0",
"torchvision>=0.19,<0.27.0",
"torchaudio>=2.4,<2.12.0",
)
# torchao's C++ extensions are built against ONE exact torch release; a newer
# torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions
# due to incompatible torch version ...") and fall back to slow Python. Because
# the torch pin above is a range (and every CUDA index now tops out at torch
# 2.10), the torch actually installed drifts ahead of a fixed torchao pin. So
# pick the torchao whose build matches the torch in the venv. Table: pytorch/ao#2919.
# torch 2.9.x -> torchao 0.14.0 (today's pin; built for torch 2.9.0)
# torch 2.10.x -> torchao 0.16.0 (built for torch 2.10.0)
# torch 2.11.x -> torchao 0.17.0 (built for torch 2.11.0; reachable via ROCm rocm7.2)
# Unknown/older torch keeps the conservative default (no regression vs today).
# torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch
# mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails
# to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a
# range, so match torchao to the installed torch (table: pytorch/ao#2919):
# 2.9.x -> 0.14.0
# 2.10.x, CUDA<=12 -> 0.16.0 (cpp built for 2.10, loads via the CUDA-12 wheel)
# 2.10.x, CUDA>=13 -> 0.17.0 (cu130: 0.16.0's CUDA-12 cpp crashes on load; 0.17.0
# targets torch 2.11 so its cpp is cleanly skipped, not crashed)
# 2.11.x -> 0.17.0 (reachable via CUDA or ROCm rocm7.2)
# Unknown/older torch keeps the conservative default.
_TORCHAO_DEFAULT_SPEC = "torchao==0.14.0"
_TORCHAO_BY_TORCH_MINOR: dict[int, str] = {
10: "torchao==0.16.0",
11: "torchao==0.17.0",
}
_TORCHAO_TORCH_210_SPEC = "torchao==0.16.0"
_TORCHAO_TORCH_210_CUDA13_SPEC = "torchao==0.17.0"
_TORCHAO_TORCH_211_PLUS_SPEC = "torchao==0.17.0"
# torch 2.10 built against CUDA >= this major can't load 0.16.0's CUDA-12 cpp.
_TORCHAO_CUDA13_MIN_MAJOR = 13
def _cuda_major_from_torch_version(torch_version: str) -> int | None:
"""Extract the CUDA major from a torch local version tag, e.g. '2.10.0+cu130'
-> 13, '2.10.0+cu128' -> 12. Returns None for rocm/cpu/tagless builds."""
local = str(torch_version).split("+", 1)
if len(local) < 2 or not local[1].startswith("cu"):
return None
digits = re.sub(r"[^0-9].*", "", local[1][2:]) # 'cu130' -> '130'
if not digits:
return None
return int(digits) // 10 # '130' -> 13, '128' -> 12, '118' -> 11
def _select_torchao_spec(torch_version: str | None) -> str:
@ -151,8 +165,14 @@ def _select_torchao_spec(torch_version: str | None) -> str:
if major != 2:
return _TORCHAO_DEFAULT_SPEC
if minor >= 11:
return _TORCHAO_BY_TORCH_MINOR[11] # newest known build; covers 2.11+
return _TORCHAO_BY_TORCH_MINOR.get(minor, _TORCHAO_DEFAULT_SPEC)
return _TORCHAO_TORCH_211_PLUS_SPEC # newest known build; covers 2.11+
if minor == 10:
# cu130+ can't load 0.16.0's CUDA-12 cpp; use 0.17.0 (cpp skipped, not crashed).
cuda_major = _cuda_major_from_torch_version(str(torch_version))
if cuda_major is not None and cuda_major >= _TORCHAO_CUDA13_MIN_MAJOR:
return _TORCHAO_TORCH_210_CUDA13_SPEC
return _TORCHAO_TORCH_210_SPEC
return _TORCHAO_DEFAULT_SPEC
def _probe_installed_torch_version() -> str | None: