fix: BNB AMD wheel skipped + torch.compile segfault on Windows ROCm

install_python_stack.py: the UNSLOTH_ROCM_TORCH_INSTALLED=1 early-return
path (set by setup.ps1 when it installed torch itself) returned before
ever reaching the AMD BNB prerelease wheel install.  The PyPI
bitsandbytes==0.49.x ships only CUDA DLLs, so loading it on ROCm fails
with "libbitsandbytes_rocm72.dll not found".  Now installs the AMD
Windows BNB wheel before returning on that path too.

worker.py: torch._grouped_mm crashes on gfx1200 (null HIP kernel pointer,
0xC0000005) when torch.compile's JitDecomp system dispatches it during
the first forward pass.  Detect Windows ROCm via torch.version.hip
(already in sys.modules from section 1e) and set TORCHDYNAMO_DISABLE=1
to bypass the broken kernel dispatch.
This commit is contained in:
LeoBorcherding 2026-05-14 13:07:14 -05:00
commit f5278deef7
2 changed files with 30 additions and 0 deletions

View file

@ -352,6 +352,19 @@ def _ensure_rocm_torch() -> None:
# setup.ps1 sets this when it already installed AMD wheels; skip the probe.
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_rocm_windows_torch_installed = True
# setup.ps1 already installed ROCm torch, but we still need to install
# the AMD Windows BNB wheel here — the PyPI bitsandbytes wheel ships
# only CUDA DLLs and will fail to load on ROCm (no libbitsandbytes_rocm72.dll).
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
if _bnb_win_url is not None:
pip_install_try(
"bitsandbytes (AMD Windows, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_win_url,
constrain = False,
)
return
if IS_MACOS:
return