diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index f208e395d9..da7eee86d2 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1321,6 +1321,21 @@ def run_training_process( os.environ["TORCHDYNAMO_DISABLE"] = "1" logger.info("Windows ROCm: torch.compile (dynamo) disabled") + # Force BNB to load libbitsandbytes_rocm72.dll regardless of the + # HIP version that torch reports. As of torch==2.11.0+rocm7.13.0 + # (AMD index, May 2026) torch.version.hip returns "7.13", which + # makes BNB look for rocm713.dll — a file our AMD Windows prerelease + # wheel does not ship. The wheel only ships rocm72.dll, so we pin + # BNB_ROCM_VERSION="72" here. Callers may override by setting the + # variable before launching the worker. + if "BNB_ROCM_VERSION" not in os.environ: + os.environ["BNB_ROCM_VERSION"] = "72" + logger.info( + "Windows ROCm: set BNB_ROCM_VERSION=72 " + "(AMD Windows BNB wheel ships rocm72.dll; " + "overrides auto-detection from torch.version.hip)" + ) + # Patch _grouped_mm CUDA dispatch with a safe Python mm fallback. try: import warnings as _warnings diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 046421456a..3df1756fbf 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -106,8 +106,12 @@ _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { "download/continuous-release_main/" "bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl" ), - # Windows ROCm wheel — ships libbitsandbytes_rocm72.dll. - # BNB_ROCM_VERSION=72 must be set in the environment before importing bnb. + # Windows ROCm wheel — ships libbitsandbytes_rocm72.dll only. + # As of torch==2.11.0+rocm7.13.0 (AMD index, May 2026), BNB auto-detects + # HIP version as "7.13" and looks for rocm713.dll — which does not exist. + # BNB_ROCM_VERSION=72 must be set in the environment before importing bnb + # to force it to load rocm72.dll. Set in worker.py (training subprocess) + # and in _install_bnb_windows_rocm() (install subprocess). "win_amd64": ( "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" "download/continuous-release_main/" @@ -356,6 +360,10 @@ def _install_bnb_windows_rocm() -> None: _bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64") if _bnb_win_url is None: return + # Pin BNB_ROCM_VERSION=72 in this process now so that any post-install + # import of bitsandbytes (e.g. health-checks) loads the correct DLL. + # The worker subprocess sets this independently in worker.py section 1f. + os.environ.setdefault("BNB_ROCM_VERSION", "72") _prev = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1" try: diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 2752a0b446..61428a3558 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1674,6 +1674,27 @@ class TestInstallBnbWindowsRocm: stack_mod._install_bnb_windows_rocm() mock_pip.assert_not_called() + def test_sets_bnb_rocm_version_72(self): + """BNB_ROCM_VERSION must be set to '72' before install. + + As of torch==2.11.0+rocm7.13.0 (AMD index, May 2026), BNB auto-detects + HIP 7.13 and looks for rocm713.dll — which the prerelease wheel does not + ship. _install_bnb_windows_rocm() must pin BNB_ROCM_VERSION=72 so that + bitsandbytes loads libbitsandbytes_rocm72.dll instead. + """ + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "72" + + def test_does_not_override_existing_bnb_rocm_version(self): + """An explicit BNB_ROCM_VERSION in the caller's env must not be clobbered.""" + with patch.dict(os.environ, {"BNB_ROCM_VERSION": "60"}): + with patch.object(stack_mod, "pip_install_try", return_value = True): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "60" + # ============================================================================= # TEST: install_python_stack.py -- UNSLOTH_ROCM_TORCH_INSTALLED early-return path @@ -1764,6 +1785,32 @@ class TestWorkerWindowsRocmPatches: source = _WORKER_PATH.read_text(encoding = "utf-8") assert "TORCHDYNAMO_DISABLE" in source + def test_bnb_rocm_version_set_on_windows_rocm(self): + """worker.py must pin BNB_ROCM_VERSION=72 in the Windows ROCm section. + + As of torch==2.11.0+rocm7.13.0, BNB auto-detects HIP 7.13 and looks for + rocm713.dll, which the AMD prerelease wheel does not ship. The worker + must force BNB_ROCM_VERSION=72 before any ML library is imported so that + bitsandbytes loads libbitsandbytes_rocm72.dll. + """ + source = _WORKER_PATH.read_text(encoding="utf-8") + assert "BNB_ROCM_VERSION" in source + assert '"72"' in source + + def test_bnb_rocm_version_set_before_ml_imports(self): + """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports.""" + source = _WORKER_PATH.read_text(encoding="utf-8") + idx_bnb = source.find("BNB_ROCM_VERSION") + # Use the specific section-2 marker that appears in the worker process + # entry-point function (not the trainer helper which has its own "# ── 2."). + idx_sec2 = source.find("# ── 2. Now import ML libraries") + assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py" + assert idx_sec2 != -1, "'# ── 2. Now import ML libraries' marker not found in worker.py" + assert idx_bnb < idx_sec2, ( + "BNB_ROCM_VERSION must be set before section 2 ML imports " + f"(found at {idx_bnb}, section 2 at {idx_sec2})" + ) + def test_grouped_mm_patch_guarded_by_windows_and_hip_check(self): """_grouped_mm patch must only apply on Windows + HIP torch.""" source = _WORKER_PATH.read_text(encoding = "utf-8")