vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Studio: fix WSL Strix Halo GPU on reinstall (ROCDXG drop-in + system HIP before bundle) (#6227)

Browse source 
* install.sh: persist ROCm-on-WSL drop-in even when rocminfo already works

_maybe_bootstrap_rocm_wsl calls _ensure_rocm_probe_env (which exports a
transient HSA_ENABLE_DXG_DETECTION + adds /opt/rocm/bin to PATH on the
installer process) right before the "rocminfo enumerates gfx1151 -> already
set up, return early" gate. On any reinstall over an existing /opt/rocm --
the common case, since the uninstaller keeps shared ROCm userspace but
removes /etc/profile.d/unsloth-rocm-wsl.sh -- that probe env makes rocminfo
succeed, so the gate returns 0 WITHOUT ever persisting the drop-in. The
transient env dies with the installer, so the next login shell (Studio,
llama-server) sees no GPU: torch cuda_avail=False, rocminfo finds nothing,
the llama.cpp ROCm prebuilt segfaults on a GPU it can't reach.

Factor the drop-in writer into _persist_rocm_wsl_dropin() and call it before
the early return so the persistent env is restored whenever librocdxg is
present. Idempotent (only writes when the drop-in is missing), gated on
librocdxg so it never fires on non-WSL/non-ROCDXG hosts, root-writes or
sudo-tees like before. The fast-path branch now reuses the same helper.

Reproduced on gfx1151 (Radeon 8060S) under dash (the curl|sh shell):
before the fix a reinstall left the drop-in absent and torch cuda_avail
False; after, the drop-in is persisted and a fresh login shell reports
cuda_avail True. Verified under both dash and bash, and idempotent on
re-run.

* Studio WSL: load system HIP before a prebuilt's bundled runtime (gfx1151)

The lemonade / published llama.cpp ROCm prebuilts bundle their own HIP
runtime (libamdhip64) built for bare-metal Linux. In WSL the GPU is reached
through the system ROCm's librocdxg bridge over /dev/dxg, which the bundled
runtime cannot drive -- it segfaults on the first GPU call. So:

  - install_llama_prebuilt.py: the prebuilt's llama-quantize/llama-server
    validation runs with the bundle dir first on LD_LIBRARY_PATH, segfaults
    (empty stderr), and the install silently falls back to a CPU source build
    (which on this host can't even build for GPU -- hipcc absent). The Strix
    Halo WSL user ends up on CPU despite a working GPU.
  - llama_cpp.py: even if a GPU prebuilt were kept, the serve-time launcher
    put the bundle dir first too, so it would crash at load.

Fix: on a ROCDXG WSL host (gated on /dev/dxg + "microsoft" /proc/version +
a librocdxg-providing /opt/rocm), prepend the system ROCm lib dir to
LD_LIBRARY_PATH so the WSL-capable libamdhip64 + librocdxg load first, while
the bundle still supplies libggml-hip / librocblas with the gfx1151 kernels.
Set HSA_ENABLE_DXG_DETECTION=1 alongside. Added _wsl_system_rocm_lib_dirs()
to both modules (kept identical so a prebuilt that passed install validation
runs the same way at serve time). Strict no-op on bare-metal Linux, NVIDIA,
macOS, and Windows.

Verified on gfx1151 (Radeon 8060S) in WSL (ROCm 7.2.1 + librocdxg, Adrenalin
ROCDXG): before, the lemonade gfx1151 prebuilt segfaulted and the install
fell back to a broken CPU build; after, install_llama_prebuilt validates and
keeps the GPU prebuilt (source=published, prebuilt_fallback_used=False), and
Studio serves Qwen3-1.7B-GGUF at 53 tok/s with the model resident in GPU
memory (llama-server device_info: ROCm0 = AMD Radeon 8060S).

* tests: cover the WSL ROCDXG drop-in + system-HIP-ordering fixes

- _wsl_system_rocm_lib_dirs: no-op without /dev/dxg, on bare-metal Linux,
  and on WSL without librocdxg; returns the system lib dir on a ROCDXG WSL
  host.
- binary_env: prepends the system ROCm lib dir ahead of the bundle and sets
  HSA_ENABLE_DXG_DETECTION on WSL; unchanged on bare-metal Linux.
- install.sh: _persist_rocm_wsl_dropin exists, is gated on librocdxg, and the
  rocminfo-already-works early return calls it before returning.
- llama_cpp.py: the serve-time launcher prepends the WSL rocm dirs before the
  bundle dir (mirrors binary_env).

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

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

* Tighten WSL ROCDXG fix comments (no logic change)

Condense the drop-in / system-HIP-ordering comments and docstrings added in
this PR. Verified comment-only via AST parse + py_compile + sh/bash -n, the
308-test rocm_support suite, and a dash functional re-run of the bootstrap
(drop-in still persisted, env still set).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-12 04:46:39 -07:00 • committed by GitHub
parent b72cf0af24
commit 240c0c3500
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 259 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

63
install.sh
View file

@ -2047,6 +2047,37 @@ _pick_radeon_wheel() {
# CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts
# the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 +
# librocdxg), then sources the env it persisted so detection finds the GPU.
# Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d
# so non-login Studio/llama launches inherit it. Idempotent (writes only when
# the drop-in is missing); no-op without librocdxg, so never fires off WSL.
# /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes
# after this shell on a non-root reinstall. Best-effort either way.
_persist_rocm_wsl_dropin() {
[ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ] || return 0
_rw_rocm=/opt/rocm
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
case ":${PATH}:" in
*":${_rw_rocm}/bin:"*) ;;
*) export PATH="${_rw_rocm}/bin:${PATH}" ;;
esac
export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}"
[ -r /etc/profile.d/unsloth-rocm-wsl.sh ] && return 0
_rw_dropin="$(
printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n'
printf 'export HSA_ENABLE_DXG_DETECTION=1\n'
printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n'
printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}"
printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}"
printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n'
)"
if [ "$(id -u)" = "0" ]; then
printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true
elif command -v sudo >/dev/null 2>&1; then
printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true
fi
}
_maybe_bootstrap_rocm_wsl() {
[ "${OS:-}" = "wsl" ] || return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
@ -2060,6 +2091,11 @@ _maybe_bootstrap_rocm_wsl() {
_ensure_rocm_probe_env
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then
# rocminfo may work only via the transient env _ensure_rocm_probe_env
# just set, which dies with the installer. Persist the drop-in so login
# shells (Studio, llama.cpp) inherit it -- else a reinstall over an
# existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU.
_persist_rocm_wsl_dropin
return 0
fi
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
@ -2077,31 +2113,8 @@ _maybe_bootstrap_rocm_wsl() {
. /etc/profile.d/unsloth-rocm-wsl.sh || true
else
# librocdxg present but the env drop-in is gone (e.g. a Studio
# uninstall removed it while keeping shared ROCm). Restore the FULL
# env inline (so rocminfo is on PATH) and recreate the drop-in.
_rw_rocm=/opt/rocm
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${_rw_rocm}/bin:${PATH}"
export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}"
# Persist the drop-in so later non-login Studio launches get the env
# too. /etc/profile.d is root-owned: a plain redirect fails for a
# non-root reinstall (ROCm would silently disappear after this shell),
# so tee through sudo when not root. Best-effort -- the current shell
# already has the env, so the install proceeds either way.
_rw_dropin="$(
printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n'
printf 'export HSA_ENABLE_DXG_DETECTION=1\n'
printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n'
printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}"
printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}"
printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n'
)"
if [ "$(id -u)" = "0" ]; then
printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true
elif command -v sudo >/dev/null 2>&1; then
printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true
fi
# uninstall removed it while keeping shared ROCm). Restore the env.
_persist_rocm_wsl_dropin
fi
return 0
fi

37
studio/backend/core/inference/llama_cpp.py
View file

@ -75,6 +75,33 @@ from state.tool_approvals import (
logger = get_logger(__name__)
def _wsl_system_rocm_lib_dirs() -> "list[str]":
"""System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL.
The bundled bare-metal HIP can't drive WSL's /dev/dxg and segfaults on the
first GPU call; the system ROCm libs (libamdhip64 + librocdxg) can, while
the bundle still supplies libggml-hip / librocblas (gfx1151 kernels).
Mirrors install_llama_prebuilt._wsl_system_rocm_lib_dirs so a prebuilt that
passed install validation runs the same at serve time. No-op off a ROCDXG
WSL host (needs /dev/dxg, "microsoft" /proc/version, librocdxg in /opt/rocm).
"""
try:
if not os.path.exists("/dev/dxg"):
return []
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
if "microsoft" not in fh.read().lower():
return []
except OSError:
return []
out: "list[str]" = []
for d in ("/opt/rocm/lib", "/opt/rocm/lib64"):
if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists(
os.path.join(d, "librocdxg.so.1")
):
out.append(d)
return out
# ── Pre-compiled patterns for plan-without-action re-prompt ──
# Forward-looking intent signals: the model is describing what it *will*
# do rather than giving a final answer.
@ -3789,7 +3816,15 @@ class LlamaCppBackend:
# plus CUDA runtime libs (libcudart, libcublas, etc.)
import platform
lib_dirs = [binary_dir]
lib_dirs = []
# WSL: system HIP before the bundle's (which segfaults on
# /dev/dxg). Mirror install_llama_prebuilt.binary_env, which
# validates the prebuilt with this same ordering.
for _wsl_rocm in _wsl_system_rocm_lib_dirs():
lib_dirs.append(_wsl_rocm)
if lib_dirs:
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
lib_dirs.append(binary_dir)
_arch = platform.machine() # x86_64, aarch64, etc.
# Pip-installed nvidia CUDA runtime libs. The prebuilt

32
studio/install_llama_prebuilt.py
View file

@ -5662,6 +5662,33 @@ def windows_runtime_dirs_for_runtime_line(runtime_line: str | None) -> list[str]
return windows_runtime_dirs_for_patterns(patterns)
def _wsl_system_rocm_lib_dirs() -> list[str]:
"""System ROCm lib dir(s) for binary_env to load before a prebuilt's HIP.
A prebuilt bundles a bare-metal HIP runtime that can't drive WSL's /dev/dxg
and segfaults on the first GPU call -> validation fails, install falls back
to a CPU build. Putting the system ROCm libs first loads the WSL-capable
HIP (libamdhip64 + librocdxg) while the bundle still supplies libggml-hip /
librocblas with the gfx1151 kernels. Strict no-op off WSL (needs /dev/dxg, a
"microsoft" /proc/version, and a librocdxg-providing ROCm).
"""
try:
if not os.path.exists("/dev/dxg"):
return []
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
if "microsoft" not in fh.read().lower():
return []
except OSError:
return []
out: list[str] = []
for d in ("/opt/rocm/lib", "/opt/rocm/lib64"):
if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists(
os.path.join(d, "librocdxg.so.1")
):
out.append(d)
return out
def binary_env(
binary_path: Path,
install_dir: Path,
@ -5683,6 +5710,11 @@ def binary_env(
str(install_dir),
*linux_runtime_dirs(binary_path),
]
# WSL: system HIP before the bundle's (which segfaults on /dev/dxg).
_wsl_rocm = _wsl_system_rocm_lib_dirs()
if _wsl_rocm:
ld_dirs = [*_wsl_rocm, *ld_dirs]
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
existing = [part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part]
env["LD_LIBRARY_PATH"] = os.pathsep.join(dedupe_existing_dirs([*ld_dirs, *existing]))
elif host.is_macos:

154
tests/studio/install/test_rocm_support.py
View file

@ -11,7 +11,7 @@ import os
import subprocess
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch, PropertyMock
from unittest.mock import MagicMock, mock_open, patch, PropertyMock
import pytest
@ -3207,5 +3207,157 @@ class TestRocmGfxForwarding:
assert "$script:ROCmGfxArch" in source
# TEST: WSL ROCDXG fixes -- drop-in persistence + system-HIP-before-bundle
_INSTALL_SH_PATH = PACKAGE_ROOT / "install.sh"
_LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py"
class TestWslSystemRocmLibDirs:
"""install_llama_prebuilt._wsl_system_rocm_lib_dirs: a strict no-op off a
ROCDXG WSL host; returns the system ROCm lib dir there so binary_env can
load the WSL-capable HIP runtime before a prebuilt's bundled one."""
def test_empty_without_dev_dxg(self):
with patch("os.path.exists", return_value = False):
assert prebuilt_mod._wsl_system_rocm_lib_dirs() == []
def test_empty_on_bare_metal_linux(self):
# /dev/dxg present but /proc/version is not a WSL kernel.
with patch("os.path.exists", lambda p: p == "/dev/dxg"):
with patch(
"builtins.open",
mock_open(read_data = "Linux version 6.8.0-generic"),
):
assert prebuilt_mod._wsl_system_rocm_lib_dirs() == []
def test_returns_system_lib_on_wsl_with_librocdxg(self):
# Normalize separators: os.path.join uses "\" on the Windows test host
# though the target path is Linux.
def _exists(p):
p = str(p).replace("\\", "/")
return p in ("/dev/dxg", "/opt/rocm/lib/librocdxg.so")
with patch("os.path.exists", _exists):
with patch(
"builtins.open",
mock_open(read_data = "Linux version 5.15.0-microsoft-standard-WSL2"),
):
assert prebuilt_mod._wsl_system_rocm_lib_dirs() == ["/opt/rocm/lib"]
def test_empty_on_wsl_without_librocdxg(self):
# WSL kernel + /dev/dxg but no librocdxg -> not a ROCDXG ROCm install.
with patch("os.path.exists", lambda p: p == "/dev/dxg"):
with patch(
"builtins.open",
mock_open(read_data = "microsoft-standard-WSL2"),
):
assert prebuilt_mod._wsl_system_rocm_lib_dirs() == []
class TestBinaryEnvWslOrdering:
"""binary_env must put the system ROCm lib dir AHEAD of the prebuilt's
bundle dir on a ROCDXG WSL host, and set HSA_ENABLE_DXG_DETECTION; no-op
on bare-metal Linux."""
@staticmethod
def _linux_host():
return HostInfo(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = True,
)
def test_wsl_prepends_system_rocm_and_sets_hsa(self, tmp_path):
binary = tmp_path / "bundle" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("")
# binary_env's dedupe_existing_dirs drops non-existent dirs, so use a
# real dir to stand in for the system ROCm lib path.
sys_rocm = tmp_path / "sysrocm"
sys_rocm.mkdir()
with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = [str(sys_rocm)]):
with patch.dict(os.environ, {}, clear = True):
env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host())
ld = env["LD_LIBRARY_PATH"].split(os.pathsep)
# Compare resolved paths (dedupe_existing_dirs calls Path.resolve()).
ld_resolved = [str(Path(p).resolve()) for p in ld]
assert ld_resolved[0] == str(sys_rocm.resolve())
assert str(binary.parent.resolve()) in ld_resolved
assert ld_resolved.index(str(sys_rocm.resolve())) < ld_resolved.index(
str(binary.parent.resolve())
)
assert env.get("HSA_ENABLE_DXG_DETECTION") == "1"
def test_bare_metal_linux_unchanged(self, tmp_path):
binary = tmp_path / "bundle" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("")
with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []):
with patch.dict(os.environ, {}, clear = True):
env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host())
ld = env["LD_LIBRARY_PATH"].split(os.pathsep)
assert ld[0] == str(binary.parent) # bundle dir first, as before
assert "HSA_ENABLE_DXG_DETECTION" not in env
class TestInstallShDropinPersistence:
"""install.sh must persist the ROCm-on-WSL drop-in even when rocminfo
already enumerates the GPU (via the transient probe env), so a reinstall
over an existing /opt/rocm doesn't leave login shells without the env."""
def test_has_persist_helper(self):
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
assert "_persist_rocm_wsl_dropin()" in source
def test_gate5_early_return_persists_dropin(self):
"""The rocminfo-already-works early return must call the persist helper
BEFORE returning."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
# Find the rocminfo gfx1151 gate and assert the persist call precedes
# its `return 0`.
gate = source.find("Name:[[:space:]]*gfx1151")
assert gate != -1
window = source[gate : gate + 900]
assert "_persist_rocm_wsl_dropin" in window
assert window.find("_persist_rocm_wsl_dropin") < window.find("return 0")
def test_persist_helper_gated_on_librocdxg(self):
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
body_start = source.find("_persist_rocm_wsl_dropin()")
body = source[body_start : body_start + 1200]
assert "librocdxg.so" in body
assert "profile.d/unsloth-rocm-wsl.sh" in body
class TestLlamaCppRuntimeWslOrdering:
"""The serve-time launcher must mirror binary_env: system HIP before the
bundle dir on WSL, so a prebuilt that passed install validation runs."""
def test_has_wsl_helper(self):
source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8")
assert "_wsl_system_rocm_lib_dirs" in source
def test_prepends_before_binary_dir(self):
source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8")
# The Linux lib_dirs build must add WSL rocm dirs before binary_dir.
idx_helper = source.find("for _wsl_rocm in _wsl_system_rocm_lib_dirs()")
idx_binary = source.find("lib_dirs.append(binary_dir)")
assert idx_helper != -1 and idx_binary != -1
assert idx_helper < idx_binary
if __name__ == "__main__":
pytest.main([__file__, "-v"])