vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-12 09:18:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 23

fix(studio/rocm): worker.py parity + don't roll back ROCm torch on bnb failure

Browse source 
Addresses findings from a 10x reviewer pass on the prior fix commit:

1. studio/backend/core/training/worker.py (parity with main.py):
   - Gate the torchao stub block on torch.version.hip / 'rocm' in
     torch.__version__ instead of HIP_PATH / ROCM_PATH env-var presence.
     Same root cause as main.py: HIP SDK env vars stick around on CUDA hosts.
   - Add module-level Windows ROCm DLL registration block. Worker subprocesses
     inherit env vars but not the parent's add_dll_directory handles, so the
     first `import torch` in the worker could fail to find amdhip64.dll when
     HIP_PATH\bin is not on PATH. Mirrors main.py setup. Handles retained at
     module scope via _ROCM_DLL_HANDLES.
   - Promote _WINDOWS_ROCM_GROUPED_MM_LIB to module scope with `global` in
     run_training_process so the torch.library.Library registration survives
     past function return / mid-run garbage collection.
   - Harden _torch_has_hip() to also accept 'rocm' in torch.__version__
     (AMD SDK / Radeon wheels may not set torch.version.hip).

2. studio/install_python_stack.py:
   - Don't roll back ROCm torch when bitsandbytes install fails. The prior
     commit gated _rocm_windows_torch_installed on _install_bnb_windows_rocm()
     returning True; if torch installed successfully but bnb failed, the flag
     stayed False and later install steps could overwrite ROCm torch with the
     generic CPU torch wheel. Set the flag after torch install; surface bnb
     failure as a separate warning instead.
   - _detect_windows_gfx_arch now probes in three tiers: UNSLOTH_ROCM_GFX_ARCH
     env-var override (matches the PowerShell installer), then hipinfo (PATH
     or HIP_PATH\bin), then amd-smi (`static --asic`, `list`). Without the
     amd-smi fallback, runtime-only Radeon installs without hipinfo on PATH
     made `studio update` return early and leave the venv on CPU torch.
   - Linux torch-already-rocm probe in _ensure_rocm_torch now matches the
     Windows probe shape: accepts torch.version.hip OR 'rocm' in
     torch.__version__ to cover AMD SDK / Radeon Linux wheels.

3. studio/backend/utils/hardware/hardware.py:
   - apply_gpu_ids() final-fallback torch probe accepts 'rocm' in
     torch.__version__ in addition to torch.version.hip, matching
     detect_hardware(). AMD SDK wheels could otherwise leak through with
     CUDA-only visibility masks on a spawned ROCm worker.

Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py
(no test changes needed; the probe shape that prints the hip version (or
'rocm' sentinel) preserves the existing non-empty-string contract).

Not addressed in this commit (deferred or out of scope):
- Tag drift / lemonade checksum (PR 5303 surface, not this PR).
- install.sh rocm7.2.1 URL: small fix, separate.
- install.ps1 / setup.ps1 'Radeon 8060S' marketing-name fallback table.
- Strix Halo + ROCm 7.1 routing asymmetry in Python update path.
This commit is contained in:
Daniel Han 2026-05-19 09:08:49 +00:00
parent 0b0b8df4b9
commit 76137b2d82
3 changed files with 138 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

71
studio/backend/core/training/worker.py
View file

@ -70,6 +70,43 @@ _TILELANG_INSTALL_TIMEOUT_S = 600
_TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11")
_FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS"
# Module-level handle so the torch.library.Library registration survives past
# run_training_process() and is not garbage collected mid-run.
_WINDOWS_ROCM_GROUPED_MM_LIB = None
# Worker subprocesses inherit the parent env but not the parent's
# os.add_dll_directory registrations. Replicate main.py's Windows ROCm DLL
# setup at module load so the first `import torch` can find amdhip64.dll even
# when HIP_PATH\bin is not on the system PATH. Handles retained at module
# scope so they are not garbage collected.
_ROCM_DLL_HANDLES: list = []
if sys.platform == "win32":
def _add_rocm_dll_dirs_worker() -> None:
_candidates: list[str] = []
for _var in ("HIP_PATH", "ROCM_PATH"):
_val = os.environ.get(_var)
if _val:
_candidates.append(os.path.join(_val, "bin"))
_default_root = os.path.join(
os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
)
try:
if os.path.isdir(_default_root):
for _ver in sorted(os.listdir(_default_root), reverse = True):
_bin = os.path.join(_default_root, _ver, "bin")
if os.path.isdir(_bin):
_candidates.append(_bin)
except OSError:
pass
for _d in _candidates:
if os.path.isdir(_d):
try:
_ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
except (OSError, AttributeError):
pass
_add_rocm_dll_dirs_worker()
del _add_rocm_dll_dirs_worker
def _model_wants_causal_conv1d(model_name: str) -> bool:
name = model_name.lower()
@ -607,11 +644,18 @@ def _tilelang_importable() -> bool:
def _torch_has_hip() -> bool:
"""True iff torch is a ROCm build; `torch.version.hip` is the only reliable signal on x86_64 ROCm."""
"""True iff torch is a ROCm build.
`torch.version.hip` covers official PyTorch ROCm wheels; AMD SDK / Radeon
wheels can leave it unset but still encode "rocm" in `torch.__version__`.
"""
try:
import torch as _torch
return getattr(_torch.version, "hip", None) is not None
return bool(
getattr(_torch.version, "hip", None)
or "rocm" in getattr(_torch, "__version__", "").lower()
)
except Exception:
return False
@ -1916,10 +1960,21 @@ def run_training_process(
# Only stub torchao on Windows ROCm hosts -- on Windows CUDA (NVIDIA) torchao
# is real and shadowing it breaks torchao-based quantization paths.
# HIP_PATH / ROCM_PATH are set by the AMD HIP SDK installer on ROCm machines.
_is_win32_rocm = sys.platform == "win32" and bool(
os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")
)
# Gate on the active torch runtime, not env-var presence -- HIP_PATH /
# ROCM_PATH stay set after a user installs the HIP SDK and reverts to a
# CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip
# but still encode "rocm" in torch.__version__, so accept either.
_is_win32_rocm = False
if sys.platform == "win32":
try:
import torch as _torch_probe
_is_win32_rocm = bool(
getattr(getattr(_torch_probe, "version", None), "hip", None)
or "rocm" in getattr(_torch_probe, "__version__", "").lower()
)
del _torch_probe
except Exception:
pass
if _is_win32_rocm:
# Seed torchao top-level + key submodules; the finder handles the rest.
for _tao_name in (
@ -1984,7 +2039,9 @@ def run_training_process(
# offs: optional group-split offsets (MoE-style variable-size batches)
#
# torch is already in sys.modules from section 1e's `import torch.distributed`.
_WINDOWS_ROCM_GROUPED_MM_LIB = None # kept alive to prevent GC of registration
# Module-level _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past
# function return / mid-run GC.
global _WINDOWS_ROCM_GROUPED_MM_LIB
if sys.platform == "win32":
_torch_for_rocm = sys.modules.get("torch")
if _torch_for_rocm is not None and getattr(

9
studio/backend/utils/hardware/hardware.py
View file

@ -1694,14 +1694,19 @@ def apply_gpu_ids(gpu_ids) -> None:
_is_rocm = IS_ROCM or _inherits_rocm_visibility
if not _is_rocm:
# torch.version.hip is a non-empty string on ROCm, None on CUDA.
# AMD SDK / Radeon ROCm wheels can leave torch.version.hip unset but
# still encode "rocm" in torch.__version__, matching detect_hardware().
# Broad except: a probe failure must never crash a training worker.
try:
import torch as _torch
_is_rocm = getattr(_torch.version, "hip", None) is not None
_is_rocm = (
getattr(_torch.version, "hip", None) is not None
or "rocm" in getattr(_torch, "__version__", "").lower()
)
except Exception as e:
logger.debug(
"apply_gpu_ids: torch.version.hip probe skipped (%s: %s)",
"apply_gpu_ids: torch ROCm probe skipped (%s: %s)",
type(e).__name__,
e,
)

96
studio/install_python_stack.py
View file

@ -228,17 +228,23 @@ def _detect_rocm_version() -> tuple[int, int] | None:
def _detect_windows_gfx_arch() -> str | None:
"""Return the gcnArchName from hipinfo on Windows (e.g. 'gfx1200'), or None.
"""Return the gcnArchName on Windows (e.g. 'gfx1200'), or None.
Resolves hipinfo via PATH first, then HIP_PATH\\bin and ROCM_PATH\\bin as
fallbacks -- the AMD HIP SDK installer sets these env vars but does not
always add the bin dir to the system PATH.
Probe order matches the PowerShell installer: env-var override first,
then hipinfo (PATH or HIP_PATH / ROCM_PATH bin), then amd-smi. Without
the amd-smi fallback, runtime-only AMD installs without hipinfo on PATH
return early and `studio update` cannot repair a CPU-only venv.
"""
import re
# 1. Explicit override (matches PowerShell installer's env-var path).
_override = os.environ.get("UNSLOTH_ROCM_GFX_ARCH")
if _override and _override.strip():
return _override.strip().lower()
# 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin.
hipinfo = shutil.which("hipinfo")
if not hipinfo:
# Fallback: AMD HIP SDK sets HIP_PATH / ROCM_PATH even when bin isn't on PATH
for _env_var in ("HIP_PATH", "ROCM_PATH"):
_root = os.environ.get(_env_var)
if _root:
@ -246,25 +252,43 @@ def _detect_windows_gfx_arch() -> str | None:
if os.path.isfile(_candidate):
hipinfo = _candidate
break
if not hipinfo:
return None
try:
result = subprocess.run(
[hipinfo],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 10,
)
if result.returncode != 0:
return None
text = result.stdout.decode(errors = "replace")
m = re.search(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
# Lowercase the captured token -- some hipinfo builds emit "Gfx1151"
# which would miss the lowercase keys in _GFX_TO_AMD_INDEX_ARCH and
# silently fall back to CPU torch.
return m.group(1).strip().lower() if m else None
except Exception:
return None
if hipinfo:
try:
result = subprocess.run(
[hipinfo],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 10,
)
if result.returncode == 0:
text = result.stdout.decode(errors = "replace")
m = re.search(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
if m:
# Lowercase -- hipinfo sometimes emits "Gfx1151".
return m.group(1).strip().lower()
except Exception:
pass
# 3. amd-smi fallback -- runtime-only Radeon installs ship amd-smi but no hipinfo.
amd_smi = shutil.which("amd-smi")
if amd_smi:
for _args in (("static", "--asic"), ("list",)):
try:
result = subprocess.run(
[amd_smi, *_args],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 10,
)
if result.returncode != 0:
continue
text = result.stdout.decode(errors = "replace").lower()
m = re.search(r"\bgfx[1-9][0-9a-z]{2,3}\b", text)
if m:
return m.group(0)
except Exception:
continue
return None
def _windows_rocm_index_url(gfx_arch: str | None) -> str | None:
@ -518,13 +542,19 @@ def _ensure_rocm_torch() -> None:
"torchaudio",
constrain = False,
)
# ROCm torch is installed (or already was); flag it so later install
# phases do not overwrite it with the generic CPU torch wheel. BNB is
# a separate dependency -- a BNB install failure must NOT roll the
# torch ROCm install back.
_rocm_windows_torch_installed = True
# Always install AMD Windows bitsandbytes -- the PyPI wheel ships only
# CUDA DLLs and will fail to load on ROCm. Install even when torch was
# already a ROCm build so that `studio update` repairs a broken bnb.
# Only flip the success flag when the install actually succeeds; otherwise
# the post-install "manual install may be required" warning is suppressed.
if _install_bnb_windows_rocm():
_rocm_windows_torch_installed = True
if not _install_bnb_windows_rocm():
print(
" Warning: AMD Windows bitsandbytes install failed; "
"ROCm torch is installed but bitsandbytes may need manual install"
)
return
# ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ──
@ -556,7 +586,15 @@ def _ensure_rocm_torch() -> None:
[
sys.executable,
"-c",
"import torch; print(getattr(torch.version,'hip','') or '')",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"ver=getattr(torch,'__version__','').lower(); "
# Print the HIP version when present (back-compat), else
# "rocm" sentinel when only torch.__version__ flags ROCm
# (AMD SDK / Radeon wheels). Empty string = CPU/CUDA.
"print(hip if hip else ('rocm' if 'rocm' in ver else ''))"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,