Commit graph

1 commit

Author SHA1 Message Date
Daniel Han
d75e765a6b
Auto-set BNB_ROCM_VERSION from the installed wheel on Windows + ROCm (#5986)
* Auto-set BNB_ROCM_VERSION from the installed wheel on Windows + ROCm

bitsandbytes derives its ROCm backend DLL name from `torch.version.hip`.
AMD's Windows bitsandbytes prerelease wheel ships a single
`libbitsandbytes_rocm<NN>.dll` whose suffix does not always match the torch
HIP version: e.g. `torch==2.11.0+rocm7.13.0` reports HIP 7.13, so bitsandbytes
looks for `libbitsandbytes_rocm713.dll`, but the wheel only ships
`libbitsandbytes_rocm72.dll`. The names disagree, the native library fails to
load, and every 4-bit / 8-bit path breaks for users running `import unsloth`
directly (Unsloth Studio already works around this in its worker).

Detect the suffix from the actually-installed wheel and pin BNB_ROCM_VERSION
before bitsandbytes is first imported (unsloth_zoo.device_type imports it during
`from .models import *`), so the correct backend loads. This is precisely the
override bitsandbytes itself recommends when the build/runtime ROCm versions
differ.

Strict no-op unless ALL of: running on Windows, a ROCm torch build, the var is
unset, and a `libbitsandbytes_rocm*.dll` is actually installed. Linux ROCm is
untouched (its multi-backend bitsandbytes resolves the backend correctly from
torch.version.hip). Honors a user-provided BNB_ROCM_VERSION and an explicit
opt-out (UNSLOTH_SKIP_BNB_ROCM_VERSION=1).

Verified on an AMD Radeon 8060S (gfx1151, Strix Halo) Windows 11 + ROCm box:
`import unsloth` now auto-sets BNB_ROCM_VERSION=72 and a native 4-bit
quantize/dequantize roundtrip succeeds with the var unset; previously it failed
to load `libbitsandbytes_rocm713.dll`.

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

* Gate BNB_ROCM_VERSION on the actual torch build, not runtime hints

_is_rocm_torch_build() falls back to environment and filesystem hints
(HIP_PATH, ROCM_PATH, ...) that are routinely present on Windows boxes
with the AMD HIP SDK installed but a CUDA or CPU torch. If such a box
also has a bitsandbytes wheel that ships a rocm DLL (AMD's Windows
prerelease wheel ships rocm72 alongside all the cuda DLLs), setting
BNB_ROCM_VERSION makes bitsandbytes raise at import on its CUDA build
and `import unsloth` breaks.

Add _is_hip_torch_build(): wheel version tag first (no torch import),
then torch.version.hip for untagged custom/source HIP builds, and use
it as the gate. The broader hint-based helper keeps its other callers.

Verified on a gfx1151 Windows box: True on the ROCm venv
(2.11.0+rocm7.13.0), False on a torch-less interpreter; 4 new unit
tests including the HIP-SDK-on-CUDA-box false-positive regression.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Relocate wiring so this PR composes with the bnb arch-detection PR

Both this PR and the fix_bitsandbytes_rocm_arch_detection PR anchored
their _gpu_init.py wiring and import_fixes.py additions on the same
configure_amdgpu_asic_id_table_path lines, so whichever merged second
hit a textual conflict in both files (verified by merging both onto
main in each order).

Move maybe_set_windows_rocm_bnb_version's wiring to a self-contained
block after the import-order warning (still before `import
unsloth_zoo`, which is what pulls in bitsandbytes on ROCm) and append
the helpers at the end of import_fixes.py. The hunks no longer
overlap, so the two PRs merge cleanly in either order. No behavior
change: the env var only needs to be set before bitsandbytes is first
imported, and it still is.

* Redetect sitecustomize-seeded BNB_ROCM_VERSION for PR #5986

After #6048, every Studio venv process starts with BNB_ROCM_VERSION seeded
by the managed sitecustomize.py block, which made this gate a no-op inside
Studio venvs and blind to wheel updates. Treat values marked
UNSLOTH_BNB_ROCM_VERSION_SOURCE=sitecustomize as redetectable defaults,
stamp redetected values as detected, and keep the seeded value when no DLL
is found. Explicit caller values still win and the opt-out is unchanged.
Also merges latest main.

* Make BNB_ROCM_VERSION opt-out drop the sitecustomize-seeded default for PR #5986

UNSLOTH_SKIP_BNB_ROCM_VERSION=1 previously no-opped the helper but left a
sitecustomize-seeded BNB_ROCM_VERSION in the environment, so bitsandbytes
still consumed the override the user disabled. The opt-out now removes
values carrying the sitecustomize source marker; explicit user values have
no marker and are untouched. Adds tests for the opt-out paths and the
empty-string edge.

* Tighten comments and docstrings for PR #5986

Comment-only pass: shorten verbose docstrings on the internal helpers,
collapse multi-line inline comments, and drop wording that restates the
code. Verified code-identical via the comment_tools.py AST signature
check (3/3 files unchanged).

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-10 08:11:26 -07:00