ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into studio-composer

Browse source
This commit is contained in:
shimmyshimmer 2026-05-30 01:07:40 -07:00
commit fc23036470
26 changed files with 5450 additions and 291 deletions
Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/studio-backend-ci.yml vendored
View file

@ -222,6 +222,7 @@ jobs:
for s in \
tests/sh/test_get_torch_index_url.sh \
tests/sh/test_mac_intel_compat.sh \
tests/sh/test_nvcc_meets_llama_minimum.sh \
tests/sh/test_tauri_install_exit_order.sh \
tests/sh/test_torch_constraint.sh; do
echo "::group::$s"

264
install.ps1
View file

@ -976,6 +976,7 @@ shell.Run cmd, 0, False
# Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
Write-TauriLog "STEP" "Installing Python"
$DetectedPython = Find-CompatiblePython
if ($DetectedPython) {
step "python" "Python $($DetectedPython.Version) already installed"
}
@ -1239,11 +1240,196 @@ shell.Run cmd, 0, False
}
}
}
# ── AMD ROCm detection (Windows) — mirrors setup.ps1 ──
$HasROCm = $false
$HipSdkInstalled = $false # HIP SDK binary found (independent of device accessibility)
$ROCmGpuLabel = $null
$ROCmVersion = $null
$ROCmGfxArch = $null
if (-not $HasNvidiaSmi) {
# hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution).
# AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type.
$hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue
if (-not $hipinfoExe) {
$hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null }
$hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" }
if ($hipRoot) {
$hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe"
if (Test-Path $hipinfoCandidate) {
Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow
Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow
Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow
$hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate }
} else {
Write-Host " [WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" -ForegroundColor Yellow
Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow
Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow
}
}
}
if ($hipinfoExe) {
$HipSdkInstalled = $true # binary found → SDK is installed regardless of device state
try {
$hipOut = & $hipinfoExe.Source 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
$HasROCm = $true
$_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() })
$_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
if ($_hipAllArches.Count -gt 0) {
$ROCmGfxArch = if ($_hipVisIdx -lt $_hipAllArches.Count) { $_hipAllArches[$_hipVisIdx] } else { $_hipAllArches[0] }
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
} else {
$ROCmGpuLabel = "AMD ROCm"
}
} elseif ($LASTEXITCODE -ne 0) {
# hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected")
$firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1)
Write-Host " [WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" -ForegroundColor Yellow
Write-Host " $firstLine" -ForegroundColor Yellow
Write-Host " Ensure ROCm drivers are installed: https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow
}
} catch {}
}
if (-not $HasROCm) {
$amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiExe) {
try {
$smiOut = & $amdSmiExe.Source list 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") {
$HasROCm = $true
# Mirror the hipinfo path: collect all gfx tokens in enumeration
# order and pick the runtime-visible one via HIP_VISIBLE_DEVICES.
$_smiVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
# Attempt 1: newer amd-smi versions embed the gfx arch in list output.
$_smiGfxTokens = @([regex]::Matches($smiOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() })
if ($_smiGfxTokens.Count -gt 0) {
$ROCmGfxArch = if ($_smiVisIdx -lt $_smiGfxTokens.Count) { $_smiGfxTokens[$_smiVisIdx] } else { $_smiGfxTokens[0] }
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
} else {
# Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+,
# including the GFX target needed for wheel index selection.
$smiAsicOut = ""
try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {}
$_asicGfxTokens = @([regex]::Matches($smiAsicOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() })
if ($_asicGfxTokens.Count -gt 0) {
$ROCmGfxArch = if ($_smiVisIdx -lt $_asicGfxTokens.Count) { $_asicGfxTokens[$_smiVisIdx] } else { $_asicGfxTokens[0] }
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
} elseif ($smiAsicOut -match "(?im)Market.?Name\s*[:\|]\s*([^\r\n]+)") {
$ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))"
} else {
$ROCmGpuLabel = "AMD ROCm"
}
}
}
} catch {}
}
}
if (-not $HasROCm) {
try {
$wmiGpu = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue |
Where-Object { $_.Name -match "AMD|Radeon" } |
Select-Object -First 1
if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name }
} catch {}
}
# ── Arch resolution: env-var override → name inference ──────────────
# Covers users whose amd-smi is too old to report the GFX target and
# who don't have hipinfo (HIP-runtime-only, common on Strix Halo / iGPU).
if ($HasROCm -and -not $ROCmGfxArch) {
# 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running.
if ($env:UNSLOTH_ROCM_GFX_ARCH) {
$ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower()
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan"
}
# 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI).
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4
@{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail)
@{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point)
@{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop
@{ P = "RX 7600"; A = "gfx1102" } # RDNA 3
@{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix)
)
foreach ($row in $nameArchTable) {
if ($ROCmGpuLabel -match $row.P) {
$ROCmGfxArch = $row.A
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
substep "gfx arch inferred from GPU name: $ROCmGfxArch" "Cyan"
substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$ROCmGfxArch to skip inference next time" "Cyan"
break
}
}
}
}
# Capture ROCm version for wheel selection (hipconfig, then amd-smi).
# Run whenever the HIP SDK binary is present, not just when the device is accessible --
# hipconfig --version works even when hipinfo reports no ROCm device (driver issue).
if ($HasROCm -or $HipSdkInstalled) {
$hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue
if (-not $hipConfigExe) {
$hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null }
if ($hipRoot) {
$hipConfigCandidate = Join-Path $hipRoot "bin\hipconfig.exe"
if (Test-Path $hipConfigCandidate) {
$hipConfigEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" }
Write-Host " [WARN] hipconfig not on PATH -- located via ${hipConfigEnvLabel}: $hipConfigCandidate" -ForegroundColor Yellow
$hipConfigExe = [PSCustomObject]@{ Source = $hipConfigCandidate }
}
}
}
if ($hipConfigExe) {
try {
$hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0) {
$hipVerLine = ($hipVerOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1).Trim()
if ($hipVerLine -match '(\d+\.\d+)') {
$ROCmVersion = $Matches[1]
$ROCmVersionFull = $hipVerLine
}
}
} catch {}
}
if (-not $ROCmVersion) {
$amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiVer) {
try {
$smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') {
$ROCmVersion = $Matches[1]
}
} catch {}
}
}
}
}
if ($HasNvidiaSmi) {
step "gpu" "NVIDIA GPU detected"
} elseif ($HasROCm) {
step "gpu" $ROCmGpuLabel
$hipSdkPath = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { "on system PATH" }
substep "HIP SDK: $hipSdkPath"
if ($ROCmVersionFull) { substep "hipconfig: $ROCmVersionFull" }
} elseif ($HipSdkInstalled -and $ROCmGpuLabel) {
# HIP SDK is installed but ROCm can't see the device (driver issue, not SDK issue)
$sdkVer = if ($ROCmVersionFull) { " (HIP $ROCmVersionFull)" } else { "" }
step "gpu" "AMD GPU detected -- not ROCm-accessible$sdkVer" "Yellow"
substep "Detected: $ROCmGpuLabel" "Yellow"
substep "[WARN] HIP SDK is installed but hipinfo reports no ROCm-capable device." "Yellow"
substep " This is a driver issue, not an SDK issue." "Yellow"
substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow"
substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} elseif ($ROCmGpuLabel) {
step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow"
substep "Detected: $ROCmGpuLabel" "Yellow"
substep "Install the HIP SDK for ROCm GPU inference:" "Yellow"
substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} else {
step "gpu" "none (chat-only / GGUF)" "Yellow"
substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow"
substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow"
}
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
@ -1270,14 +1456,73 @@ shell.Run cmd, 0, False
return "$baseUrl/cu126"
}
$TorchIndexUrl = Get-TorchIndexUrl
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
# ── GPU arch → newest compatible Windows ROCm wheel release ──
# Wheels bundle their own ROCm runtime; the installed HIP SDK version does
# not constrain which release to use. Always picks the newest release that
# supports the GPU architecture.
# ── AMD Windows ROCm: arch-aware pip index (repo.amd.com) ──
# Wheels bundle their own ROCm runtime and support all Python versions.
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
$ROCmIndexUrl = $null
$ROCmTorchFloor = $null
if ($HasROCm -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in
# torch._C._grouped_mm on torch <2.11.0 (rocm7.12 and rocm7.1 respectively).
# TheRock issues #5284 and #3284. Force torch>=2.11.0 so pip never resolves
# to the broken 2.10.0 wheels even though they exist on the AMD index.
# The <2.12.0 ceiling matches the Linux install_python_stack.py constraint
# for the same arches: AMD actively publishes new versions on their index,
# so without a ceiling a future 2.12.0+rocmX.Y wheel would be pulled in
# automatically before it has been validated on these architectures.
# Bump the ceiling here (and in install_python_stack.py) when 2.12.x is
# confirmed working on gfx120X / Strix.
$torchFloorMap = @{
"gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0"
"gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0"
}
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
if ($archFamily) {
$ROCmIndexUrl = "$amdIndexBase/$archFamily/"
$ROCmTorchFloor = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { $null }
$archLabel = if ($ROCmGfxArch) { $ROCmGfxArch } else { "AMD GPU" }
substep "$archLabel -- AMD repo.amd.com index selected" "Cyan"
if ($ROCmTorchFloor) {
substep " enforcing $ROCmTorchFloor (known _grouped_mm bug in older wheels)" "Cyan"
}
} elseif ($ROCmGfxArch) {
substep "AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow"
} else {
substep "AMD GPU detected but arch unknown -- falling back to CPU-only PyTorch" "Yellow"
}
}
if ($ROCmIndexUrl) {
$TorchIndexFamily = "rocm"
} else {
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
}
$GpuBranch = Get-TauriGpuBranch $TorchIndexFamily
Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version
# ── Print CPU-only hint when no GPU detected ──
if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") {
if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") {
Write-Host ""
substep "No NVIDIA GPU detected." "Yellow"
if ($HipSdkInstalled -and -not $HasROCm) {
substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow"
} elseif ($ROCmGpuLabel) {
substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow"
} else {
substep "No NVIDIA GPU detected." "Yellow"
}
substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow"
substep "re-run with --no-torch for a faster, lighter install:" "Yellow"
substep ".\install.ps1 --no-torch" "Yellow"
@ -1355,9 +1600,18 @@ shell.Run cmd, 0, False
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
}
}
} elseif ($TorchIndexUrl) {
} elseif ($TorchIndexUrl -or $ROCmIndexUrl) {
if ($SkipTorch) {
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
} elseif ($ROCmIndexUrl) {
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
substep "installing PyTorch from $ROCmIndexUrl..."
$torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit)
}
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."

354
install.sh
View file

@ -183,10 +183,21 @@ _install_bnb_rocm() {
fi
if [ -n "$_bnb_whl_url" ]; then
substep "installing bitsandbytes for AMD ROCm (pre-release, PR #1887)..."
if run_install_cmd "$_label (pre-release)" "$_venv_py" -m pip install \
--force-reinstall --no-cache-dir --no-deps "$_bnb_whl_url"; then
_bnb_log=$(mktemp)
if "$_venv_py" -m pip install \
--disable-pip-version-check \
--force-reinstall --no-cache-dir --no-deps \
--retries 8 --timeout 90 \
"$_bnb_whl_url" >"$_bnb_log" 2>&1; then
rm -f "$_bnb_log"
return 0
fi
_bnb_rc=$?
if _is_verbose; then
cat "$_bnb_log" >&2
fi
rm -f "$_bnb_log"
step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2
substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN"
fi
run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \
@ -245,6 +256,9 @@ _tauri_torch_index_family() {
rocm[0-9]*.[0-9]*) echo "$_diag_family" ;;
*) echo "auto" ;;
esac ;;
# AMD arch-specific index (e.g. repo.amd.com/rocm/whl/gfx1151/) --
# used for Strix Halo/Point where torch 2.11+rocm7.13 has the real fix.
*repo.amd.com/rocm/whl/gfx*|*rocm/whl/gfx*) echo "rocm7.13" ;;
"") echo "none" ;;
*) echo "auto" ;;
esac
@ -1568,16 +1582,19 @@ _find_no_torch_runtime() {
}
# ── AMD ROCm GPU detection helper ──
# Returns 0 (true) if an actual AMD GPU is present, 1 (false) otherwise.
# Checks rocminfo for gfx[1-9]* (excludes gfx000 CPU agent) and
# amd-smi list for GPU data rows (excludes header-only output).
# Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs
# KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices).
_has_amd_rocm_gpu() {
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[0-9]/ && !/Name:[[:space:]]*gfx000/{found=1} END{exit !found}'; then
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
return 0
elif command -v amd-smi >/dev/null 2>&1 && \
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
return 0
elif [ -e /dev/kfd ] && \
awk '/gpu_id/{ if ($2+0 > 0) found=1 } END{ exit !found }' \
/sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
return 0
fi
return 1
}
@ -1656,31 +1673,39 @@ get_torch_index_url() {
if [ -n "$_rocm_tag" ]; then
# Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older)
case "$_rocm_tag" in
rocm[1-5].*) echo "$_base/cpu"; return ;;
rocm[1-5].*)
echo "[WARN] ROCm $_rocm_tag detected but PyTorch ROCm wheels require ROCm 6.0+ -- falling back to CPU-only PyTorch" >&2
echo "[WARN] Upgrade ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo "$_base/cpu"; return ;;
esac
# ROCm 7.2 only has torch 2.11.0 which exceeds current bounds
# (<2.11.0). Fall back to rocm7.1 index which has torch 2.10.0.
# Enumerate explicit versions rather than matching rocm6.* so
# a host on ROCm 6.5 or 6.6 (no PyTorch wheels published) is
# clipped down to the last supported 6.x (rocm6.4) instead of
# constructing https://download.pytorch.org/whl/rocm6.5 which
# returns HTTP 403. PyTorch only ships: rocm5.7, 6.0, 6.1, 6.2,
# 6.3, 6.4, 7.0, 7.1, 7.2 (and 5.7 is below our minimum).
# TODO: uncomment rocm7.2 when the torch upper bound is bumped
# to >=2.11.0.
# Supported tags; 6.5+ clips to rocm6.4, 7.3+ caps to rocm7.2.
# PyTorch publishes major.minor URLs only (no patch level), so
# rocm7.2.1 / rocm6.0.2 / etc. must normalise to rocm7.2 / rocm6.0.
case "$_rocm_tag" in
rocm6.0|rocm6.0.*|rocm6.1|rocm6.1.*|rocm6.2|rocm6.2.*|rocm6.3|rocm6.3.*|rocm6.4|rocm6.4.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*)
echo "$_base/$_rocm_tag" ;;
rocm6.0|rocm6.0.*) echo "$_base/rocm6.0" ;;
rocm6.1|rocm6.1.*) echo "$_base/rocm6.1" ;;
rocm6.2|rocm6.2.*) echo "$_base/rocm6.2" ;;
rocm6.3|rocm6.3.*) echo "$_base/rocm6.3" ;;
rocm6.4|rocm6.4.*) echo "$_base/rocm6.4" ;;
rocm7.0|rocm7.0.*) echo "$_base/rocm7.0" ;;
rocm7.1|rocm7.1.*) echo "$_base/rocm7.1" ;;
rocm7.2|rocm7.2.*) echo "$_base/rocm7.2" ;;
rocm6.*)
# ROCm 6.5+ (no published PyTorch wheels): clip down
# to the last supported 6.x wheel set.
echo "$_base/rocm6.4" ;;
*)
# ROCm 7.2+ (including future 10.x+): cap to rocm7.1
echo "$_base/rocm7.1" ;;
# ROCm 7.3+ (future): cap to rocm7.2 (latest known)
echo "$_base/rocm7.2" ;;
esac
return
fi
# AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be
# read from any source (amd-smi, /opt/rocm/.info/version, hipconfig,
# dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch.
echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2
echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2
echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo "$_base/cpu"; return
fi
# Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P).
@ -1760,9 +1785,9 @@ print('cp{}{}'.format(sys.version_info.major, sys.version_info.minor))
}
_pick_radeon_wheel() {
# Usage: _pick_radeon_wheel PACKAGE_NAME
# Usage: _pick_radeon_wheel PACKAGE_NAME [VERSION_PREFIX]
# Scans $_RADEON_LISTING for the newest wheel whose filename starts exactly
# with PACKAGE_NAME- and matches _RADEON_PYTAG + linux_x86_64.
# with PACKAGE_NAME- (and optionally VERSION_PREFIX) and matches _RADEON_PYTAG + linux_x86_64.
# Prints the full URL (resolving relative hrefs against _RADEON_BASE_URL).
#
# POSIX-compliant pipeline: all href parsing, filtering, and version
@ -1770,11 +1795,12 @@ _pick_radeon_wheel() {
# for GNU extensions (grep -o, sort -V) that would break under BSD
# or BusyBox coreutils.
_pkg="$1"
_ver_prefix="${2:-}"
[ -n "$_RADEON_LISTING" ] || return 1
[ -n "$_RADEON_PYTAG" ] || return 1
_tag="$_RADEON_PYTAG"
_href=$(printf '%s\n' "$_RADEON_LISTING" \
| awk -v pkg="$_pkg" -v tag="$_tag" '
| awk -v pkg="$_pkg" -v tag="$_tag" -v ver_prefix="$_ver_prefix" '
BEGIN { max_pad = ""; max_url = "" }
{
line = $0
@ -1788,7 +1814,7 @@ _pick_radeon_wheel() {
base = p[n]
sub(/[?#].*/, "", base)
prefix = pkg "-"
prefix = pkg "-" ver_prefix
# Match cpXY-cpXY or cpXY-abi3 with any linux x86_64
# platform tag (linux_x86_64, manylinux_2_28_x86_64,
# manylinux2014_x86_64, etc.)
@ -1822,6 +1848,12 @@ _pick_radeon_wheel() {
TORCH_INDEX_URL=$(get_torch_index_url)
# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it.
# All other ROCm tags and CUDA stay within <2.11.0.
case "$TORCH_INDEX_URL" in
*/rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
esac
# Auto-detect GPU for AMD ROCm based
# get_torch_index_url must have chosen */rocm*
# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
@ -1834,6 +1866,78 @@ case "$TORCH_INDEX_URL" in
fi
;;
esac
# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ───────
# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug
# that causes a segfault in torch._grouped_mm (moe_utils.py line 167).
# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when
# _amd_gpu_radeon=true the installer silently lands on the broken combo.
# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2.
case "$TORCH_INDEX_URL" in
*/rocm7.1|*/rocm7.1.*)
# Collect every gfx token in rocminfo / amd-smi enumeration order
# (skip duplicates), then index by HIP_VISIBLE_DEVICES /
# ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box
# where the user selected the dGPU does NOT get rerouted to the
# Strix per-gfx index.
_gfx_all=""
if command -v rocminfo >/dev/null 2>&1; then
_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
fi
if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
# PowerShell paths also probe `amd-smi static --asic`; mirror it
# so a host with hipinfo-less amd-smi reports the gfx target.
if [ -z "$_gfx_all" ]; then
_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
fi
fi
_runtime_gfx=""
if [ -n "$_gfx_all" ]; then
_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
_idx=0
if [ -n "$_vis" ] && [ "$_vis" != "-1" ]; then
_first=${_vis%%,*}
case "$_first" in
''|*[!0-9]*) _idx=0 ;;
*) _idx=$_first ;;
esac
fi
_runtime_gfx=$(printf '%s\n' "$_gfx_all" | awk -v idx="$_idx" '
NF && !seen[$0]++ { vals[n++] = $0 }
END {
if (idx < 0 || idx >= n) idx = 0
if (n > 0) print vals[idx]
}')
fi
_strix_gfx=""
case "$_runtime_gfx" in
gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
esac
if [ -n "$_strix_gfx" ]; then
echo "" >&2
echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2
echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2
echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2
echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2
echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo "" >&2
# AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's
# actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred
# over the pytorch.org rocm7.2 fallback because it exercises the real GPU
# kernel path. Set UNSLOTH_AMD_ROCM_MIRROR to override for air-gapped installs.
_amd_strix_base="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}"
# Strip ALL trailing slashes to match Python's .rstrip("/") -- a
# double-/triple-slash mirror URL would otherwise produce 404s on
# strict pip proxies (artifactory, sonatype).
while [ "${_amd_strix_base%/}" != "$_amd_strix_base" ]; do
_amd_strix_base="${_amd_strix_base%/}"
done
TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/"
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
_amd_gpu_radeon=false
fi
;;
esac
_TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL")
if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then
_TAURI_TORCH_INDEX_FAMILY="radeon"
@ -1841,27 +1945,93 @@ fi
_TAURI_GPU_BRANCH=$(_tauri_gpu_branch "$_TAURI_TORCH_INDEX_FAMILY" "$_amd_gpu_radeon")
tauri_diag_marker "$_TAURI_GPU_BRANCH" "$_TAURI_TORCH_INDEX_FAMILY"
# ── Print CPU-only hint when no GPU detected ──
# ── GPU detection summary (mirrors install.ps1 step "gpu" block) ──
if _has_usable_nvidia_gpu; then
step "gpu" "NVIDIA GPU detected"
elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
# Probe gfx arch for the display label, honouring HIP_VISIBLE_DEVICES
_gpu_disp_gfx_all=""
_gpu_disp_mkt=""
if command -v rocminfo >/dev/null 2>&1; then
_gpu_disp_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_gpu_disp_mkt=$(rocminfo 2>/dev/null | awk -F': ' \
'/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
fi
if [ -z "$_gpu_disp_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
_gpu_disp_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
[ -z "$_gpu_disp_gfx_all" ] && \
_gpu_disp_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_gpu_disp_mkt" ] && command -v amd-smi >/dev/null 2>&1; then
_gpu_disp_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
'/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
fi
_gpu_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
_gpu_vis_idx=0
if [ -n "$_gpu_vis" ] && [ "$_gpu_vis" != "-1" ]; then
_gpu_first="${_gpu_vis%%,*}"
case "$_gpu_first" in ''|*[!0-9]*) ;; *) _gpu_vis_idx=$_gpu_first ;; esac
fi
_gpu_disp_gfx=$(printf '%s\n' "$_gpu_disp_gfx_all" | awk -v idx="$_gpu_vis_idx" \
'NF && !seen[$0]++ { a[n++]=$0 } END { if(idx>=n) idx=0; if(n>0) print a[idx] }')
# UNSLOTH_ROCM_GFX_ARCH env override (mirrors install.ps1)
if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then
_gpu_disp_gfx="${UNSLOTH_ROCM_GFX_ARCH}"
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_gpu_disp_gfx"
# Name-based arch inference when tools don't report gfx (mirrors install.ps1 nameArchTable)
elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then
case "$_gpu_disp_mkt" in
*"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4
*"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 iGPU
*"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 iGPU
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop
*"RX 7600"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU
esac
if [ -n "$_gpu_disp_gfx" ]; then
substep "gfx arch inferred from GPU name: $_gpu_disp_gfx"
substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$_gpu_disp_gfx to skip inference next time"
fi
fi
# ROCm version via hipconfig, then amd-smi
_gpu_rocm_ver=""
if command -v hipconfig >/dev/null 2>&1; then
_gpu_rocm_ver=$(hipconfig --version 2>/dev/null | awk 'NR==1 && /^[0-9]/{print; exit}' || true)
fi
if [ -z "$_gpu_rocm_ver" ] && command -v amd-smi >/dev/null 2>&1; then
_gpu_rocm_ver=$(amd-smi version 2>/dev/null | awk -F'ROCm version: ' \
'NF>1{gsub(/[[:space:]]/,"", $2); print $2; exit}' || true)
fi
if [ -n "$_gpu_disp_gfx" ]; then
step "gpu" "AMD ROCm ($_gpu_disp_gfx)"
else
step "gpu" "AMD ROCm"
fi
_rocm_root="${ROCM_PATH:-${HIP_PATH:-/opt/rocm}}"
substep "ROCm: $_rocm_root"
[ -n "$_gpu_rocm_ver" ] && substep "hipconfig: $_gpu_rocm_ver"
[ -n "$_gpu_disp_mkt" ] && [ -n "$_gpu_disp_gfx" ] && substep "GPU: $_gpu_disp_mkt"
else
step "gpu" "none (CPU-only)" "$C_WARN"
fi
# ── PyTorch wheel index note ──
case "$TORCH_INDEX_URL" in
*/cpu)
if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then
echo ""
echo " NOTE: No GPU detected (nvidia-smi and ROCm not found)."
echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference,"
echo " re-run with --no-torch for a faster, lighter install:"
echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
echo " AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
echo ""
substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):"
substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
fi
;;
*/rocm*)
echo ""
*/rocm*|*/gfx*)
if [ "$_amd_gpu_radeon" = true ]; then
echo " AMD Radeon + ROCm detected -- installing PyTorch wheels from repo.radeon.com"
substep "wheels: repo.radeon.com (Radeon)"
else
echo " AMD ROCm detected -- installing ROCm-enabled PyTorch ($TORCH_INDEX_URL)"
substep "wheels: $TORCH_INDEX_URL"
fi
echo ""
;;
esac
@ -1943,24 +2113,23 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ "$_radeon_listing_ok" = true ]; then
# Require torch, torchvision, torchaudio wheels to all resolve
# from the Radeon listing. If any is missing for this Python
# tag, fall through to the standard ROCm index instead of
# silently mixing Radeon wheels with PyPI defaults.
# from the Radeon listing. The repo often publishes multiple
# generations simultaneously, so picking the highest-version
# for each package independently can assemble a mismatched trio
# (e.g. torch 2.10 + torchvision 0.24). To prevent this,
# we identify the highest common minor version and downpair
# wheels if necessary to ensure a compatible set.
_torch_whl=$(_pick_radeon_wheel "torch" 2>/dev/null) || _torch_whl=""
_tv_whl=$(_pick_radeon_wheel "torchvision" 2>/dev/null) || _tv_whl=""
_ta_whl=$(_pick_radeon_wheel "torchaudio" 2>/dev/null) || _ta_whl=""
_tri_whl=$(_pick_radeon_wheel "triton" 2>/dev/null) || _tri_whl=""
# Sanity-check torch / torchvision / torchaudio are a
# matching release. The Radeon repo publishes multiple
# generations simultaneously, so picking the highest-version
# wheel for each package independently can assemble a
# mismatched trio (e.g. torch 2.9.1 + torchvision 0.23.0 +
# torchaudio 2.9.0 from the current rocm-rel-7.2.1 index).
# Check that torch and torchaudio share the same X.Y public
# version prefix, and that torchvision's minor correctly
# pairs with torch's minor (torchvision = torch.minor - 5
# pairs with torch's minor (torchvision = torch.minor + 15
# since torch 2.4 -> torchvision 0.19 -> torch 2.9 ->
# torchvision 0.24).
#
# URL-decode each wheel name so %2B -> + before version
# extraction. Real Radeon wheel hrefs are percent-encoded
# (torch-2.10.0%2Brocm7.2.0...), so a plain [+-] terminator
@ -1968,38 +2137,75 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# _radeon_versions_match would stay false for every real
# listing, silently forcing a fallback to the generic
# ROCm index.
_torch_ver=""
_tv_ver=""
_ta_ver=""
if [ -n "$_torch_whl" ]; then
_torch_name=$(printf '%s' "${_torch_whl##*/}" | sed 's/%2[Bb]/+/g')
_torch_ver=$(printf '%s\n' "$_torch_name" | sed -n 's|^torch-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
if [ -n "$_tv_whl" ]; then
_tv_name=$(printf '%s' "${_tv_whl##*/}" | sed 's/%2[Bb]/+/g')
_tv_ver=$(printf '%s\n' "$_tv_name" | sed -n 's|^torchvision-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
if [ -n "$_ta_whl" ]; then
_ta_name=$(printf '%s' "${_ta_whl##*/}" | sed 's/%2[Bb]/+/g')
_ta_ver=$(printf '%s\n' "$_ta_name" | sed -n 's|^torchaudio-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
_extract_version() {
_whl=$1
_pkg=$2
if [ -n "$_whl" ]; then
_name=$(printf '%s' "${_whl##*/}" | sed 's/%2[Bb]/+/g')
printf '%s\n' "$_name" | sed -n "s|^${_pkg}-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p"
fi
}
_torch_ver=$(_extract_version "$_torch_whl" "torch")
_tv_ver=$(_extract_version "$_tv_whl" "torchvision")
_ta_ver=$(_extract_version "$_ta_whl" "torchaudio")
_radeon_versions_match=false
if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
_torch_major=${_torch_ver%%.*}
_torch_minor=${_torch_ver#*.}
_ta_major=${_ta_ver%%.*}
_ta_minor=${_ta_ver#*.}
_tv_major=${_tv_ver%%.*}
_tv_minor=${_tv_ver#*.}
# torchvision expected minor (e.g. torch 2.9 -> 0.24)
_expected_tv_minor=$((_torch_minor + 15))
if [ "$_torch_major" = "$_ta_major" ] && \
[ "$_torch_minor" = "$_ta_minor" ] && \
[ "$_tv_major" = "0" ] && \
[ "$_tv_minor" = "$_expected_tv_minor" ]; then
_radeon_versions_match=true
fi
_tv_equiv_minor=$((_tv_minor - 15))
# Determine initial target minor (lowest common denominator)
_target_minor=$_torch_minor
[ "$_tv_equiv_minor" -lt "$_target_minor" ] && _target_minor=$_tv_equiv_minor
[ "$_ta_minor" -lt "$_target_minor" ] && _target_minor=$_ta_minor
# Loop downwards to find the first complete matching trio.
# This avoids aborting if the repo has gaps.
_attempts=0
while [ "$_attempts" -lt 5 ] && [ "$_target_minor" -ge 0 ]; do
_expected_tv_minor=$((_target_minor + 15))
_curr_torch=$(_pick_radeon_wheel "torch" "2.${_target_minor}." 2>/dev/null) || _curr_torch=""
_curr_tv=$(_pick_radeon_wheel "torchvision" "0.${_expected_tv_minor}." 2>/dev/null) || _curr_tv=""
_curr_ta=$(_pick_radeon_wheel "torchaudio" "2.${_target_minor}." 2>/dev/null) || _curr_ta=""
if [ -n "$_curr_torch" ] && [ -n "$_curr_tv" ] && [ -n "$_curr_ta" ]; then
# Extract versions from the wheels found in this iteration
_c_torch_ver=$(_extract_version "$_curr_torch" "torch")
_c_tv_ver=$(_extract_version "$_curr_tv" "torchvision")
_c_ta_ver=$(_extract_version "$_curr_ta" "torchaudio")
# Parse Major.Minor for validation
_c_torch_major=${_c_torch_ver%%.*}
_c_torch_minor=${_c_torch_ver#*.}
_c_ta_major=${_c_ta_ver%%.*}
_c_ta_minor=${_c_ta_ver#*.}
_c_tv_major=${_c_tv_ver%%.*}
_c_tv_minor=${_c_tv_ver#*.}
# Strict X.Y validation: allow patch versions to differ (e.g. torch 2.9.1 + vision 0.24.0)
# as long as the Major and Minor pairing is correct.
if [ "$_c_torch_major" = "$_c_ta_major" ] && \
[ "$_c_torch_minor" = "$_c_ta_minor" ] && \
[ "$_c_tv_major" = "0" ] && \
[ "$_c_tv_minor" = "$((_c_torch_minor + 15))" ]; then
_torch_whl=$_curr_torch
_tv_whl=$_curr_tv
_ta_whl=$_curr_ta
_tri_whl=""
_radeon_versions_match=true
break
fi
fi
_target_minor=$((_target_minor - 1))
_attempts=$((_attempts + 1))
done
fi
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"

97
studio/backend/core/export/worker.py
View file

@ -439,6 +439,103 @@ def run_export_process(
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── 1c. Stub torchao on Windows ROCm ──
# torchao (pulled in by transformers.quantizers) imports
# torch.distributed._functional_collectives at module level, which imports
# distributed_c10d.py unconditionally — that file crashes on Windows ROCm
# because torch._C._distributed_c10d (the RCCL backend) is absent.
# Stubbing torchao short-circuits the crash entirely.
# Must run before any import of transformers / unsloth_zoo.
import types as _types
import importlib.machinery as _ilm
import importlib.abc as _ilabc
_STUB_SENTINEL = object()
class _StubTypeMeta(type):
def __instancecheck__(cls, instance):
return False
def __subclasscheck__(cls, subclass):
return False
def __getattr__(cls, attr):
if attr.startswith("__"):
raise AttributeError(attr)
child = _StubTypeMeta(attr, (), {})
setattr(cls, attr, child)
return child
def __call__(cls, *args, **kwargs):
return None
def _make_stub_type(name):
return _StubTypeMeta(name, (), {})
def _make_mod_stub(mod_name):
m = _types.ModuleType(mod_name)
m.__path__ = []
m.__package__ = mod_name
m._unsloth_stub = _STUB_SENTINEL
m.__spec__ = _ilm.ModuleSpec(mod_name, loader = None, is_package = True)
def _ga(attr, _m = m, _n = mod_name):
if attr.startswith("__"):
raise AttributeError(attr)
child = _make_stub_type(f"{_n}.{attr}")
setattr(_m, attr, child)
return child
m.__getattr__ = _ga
return m
class _StubSubpackageLoader(_ilabc.Loader):
def __init__(self, mod_name):
self._mod_name = mod_name
def create_module(self, spec):
return _make_mod_stub(self._mod_name)
def exec_module(self, module):
pass
class _StubSubpackageFinder(_ilabc.MetaPathFinder):
def find_spec(self, fullname, path, target = None):
if "." not in fullname:
return None
parent = sys.modules.get(fullname.rsplit(".", 1)[0])
if parent is None:
return None
if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL:
return None
return _ilm.ModuleSpec(
fullname, _StubSubpackageLoader(fullname), is_package = True
)
_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:
sys.meta_path.append(_StubSubpackageFinder())
for _tao_name in (
"torchao",
"torchao.quantization",
"torchao.dtypes",
"torchao.float8",
"torchao.utils",
):
if _tao_name not in sys.modules:
sys.modules[_tao_name] = _make_mod_stub(_tao_name)
# ── 2. Import ML libraries (fresh in this clean process) ──
try:
_send_response(

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

@ -1238,6 +1238,33 @@ class LlamaCppBackend:
return total
@staticmethod
def _amd_apu_wants_unified_memory() -> bool:
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
False for discrete AMD, NVIDIA, CPU and macOS (the env hurts discrete
GPUs). ROCm reuses torch.cuda.*; the gcnArchName suffix is stripped."""
try:
import torch
if getattr(torch.version, "hip", None) is None:
return False
if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
return False
for _i in range(torch.cuda.device_count()):
try:
_arch = (
getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "")
or ""
)
except Exception:
continue
if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
return True
except Exception:
return False
return False
@staticmethod
def _get_gpu_free_memory() -> list[tuple[int, int]]:
"""Query free memory per GPU.
@ -3158,6 +3185,14 @@ class LlamaCppBackend:
env = child_env_without_native_path_secret()
binary_dir = str(Path(binary).parent)
# AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
# shared system RAM. setdefault so a user value wins.
if self._amd_apu_wants_unified_memory():
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
logger.info(
"AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1"
)
if sys.platform == "win32":
# See _build_windows_path_dirs for ordering. #5106.
path_dirs = self._build_windows_path_dirs(
@ -3167,6 +3202,24 @@ class LlamaCppBackend:
)
existing_path = env.get("PATH", "")
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
# ROCm: the llama.cpp prebuilt bundles its own rocblas.dll
# but NOT the Tensile kernel library files it needs
# (rocblas/library/TensileLibrary*.dat + *.hsaco). The
# bundled DLL searches relative to its own location by
# default (i.e. <binary_dir>/rocblas/library/) which does
# not exist, causing a silent crash on the first GEMM.
# ROCBLAS_TENSILE_LIBPATH overrides that search to point at
# the ROCm installation where the kernel files actually are.
_hip_path = os.environ.get(
"HIP_PATH", os.environ.get("ROCM_PATH", "")
)
if _hip_path:
_rocblas_lib = os.path.join(
_hip_path, "bin", "rocblas", "library"
)
if os.path.isdir(_rocblas_lib):
env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
else:
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
# and CUDA runtime libs (libcudart, libcublas, etc.)

26
studio/backend/core/training/trainer.py
View file

@ -42,7 +42,10 @@ from utils.hardware import (
get_visible_gpu_count,
)
torch._dynamo.config.recompile_limit = 64
# recompile_limit was removed in some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2).
# Guard so training doesn't crash on RDNA2/RDNA3 with older ROCm torch wheels.
if hasattr(torch._dynamo.config, "recompile_limit"):
torch._dynamo.config.recompile_limit = 64
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
@ -657,6 +660,15 @@ class UnslothTrainer:
f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
)
# On hardware without native bfloat16 support (e.g. RDNA2 / gfx103x),
# passing dtype=None lets unsloth auto-detect and incorrectly choose
# bf16, triggering an LLVM error at the first bf16 kernel dispatch.
# Explicitly pass float16 as the fallback so unsloth never reaches
# that path. Modern NVIDIA (Ampere+) and RDNA3+ return True here so
# they are unaffected — dtype stays None and unsloth picks bf16 as
# before.
_auto_dtype = None if is_bfloat16_supported() else torch.float16
# Branch based on model type
if self._audio_type == "csm":
# CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False
@ -666,7 +678,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None,
dtype = _auto_dtype,
auto_model = CsmForConditionalGeneration,
load_in_4bit = False,
device_map = device_map,
@ -683,7 +695,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
dtype = None,
dtype = _auto_dtype,
load_in_4bit = False,
device_map = device_map,
full_finetuning = full_finetuning,
@ -705,7 +717,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None,
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@ -777,7 +789,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None,
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@ -791,7 +803,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastVisionModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None, # Auto-detect
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@ -824,7 +836,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None, # Auto-detect
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,

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

@ -70,6 +70,58 @@ _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"
)
def _ver_key(name: str) -> tuple:
# Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string.
parts = []
for chunk in name.split("."):
try:
parts.append((0, int(chunk)))
except ValueError:
parts.append((1, chunk))
return tuple(parts)
try:
if os.path.isdir(_default_root):
for _ver in sorted(
os.listdir(_default_root), key = _ver_key, 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()
@ -320,11 +372,21 @@ def _install_package_wheel_first(
f"{snippet}",
)
else:
logger.error(
"Failed to install %s from PyPI:\n%s",
display_name,
result.stdout,
)
if sys.platform == "win32":
# No prebuilt wheel and no source build toolchain on Windows --
# this is expected for packages like causal-conv1d. Log at
# info so users aren't alarmed by what looks like an error.
logger.info(
"%s is not available on Windows (no prebuilt wheel); skipping",
display_name,
)
logger.debug("Install output:\n%s", result.stdout)
else:
logger.error(
"Failed to install %s from PyPI:\n%s",
display_name,
result.stdout,
)
return False
if is_hip:
@ -337,6 +399,9 @@ def _install_package_wheel_first(
def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:
if not _model_wants_causal_conv1d(model_name):
return
if sys.platform == "win32":
logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping")
return
_install_package_wheel_first(
event_queue = event_queue,
@ -404,6 +469,11 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
"""Install pinned FLA + fla-core with --no-deps. Returns True iff importable post-call."""
if os.getenv(_FLA_SKIP_ENV) == "1":
return False
if sys.platform == "win32":
logger.info(
"Skipping flash-linear-attention install: no prebuilt wheel for Windows"
)
return False
if sys.version_info < _FLA_MIN_PYTHON:
logger.info(
"Skipping flash-linear-attention install: requires Python >= %d.%d, have %s",
@ -483,10 +553,17 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
return False
if result.returncode != 0:
logger.warning(
"flash-linear-attention install failed (continuing on torch fallback):\n%s",
result.stdout,
)
if sys.platform == "win32":
logger.info(
"flash-linear-attention not available on Windows (no prebuilt wheel); "
"continuing on torch fallback"
)
logger.debug("Install output:\n%s", result.stdout)
else:
logger.warning(
"flash-linear-attention install failed (continuing on torch fallback):\n%s",
result.stdout,
)
_send_status(
event_queue,
"flash-linear-attention install failed; continuing without it",
@ -607,15 +684,61 @@ 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
def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
"""Classify a ROCm device as unified-memory (APU) or discrete.
Returns ``(gcn_arch, is_unified)`` where:
- ``gcn_arch`` is the canonical arch string (e.g. ``"gfx1151"``) when a
known attribute is present, or ``""`` when all arch attrs are absent.
- ``is_unified`` is ``True`` for AMD APUs with a shared GPU/system-RAM pool
(gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower
``set_per_process_memory_fraction`` cap to leave headroom for the OS.
Classification priority:
1. ``gcnArchName`` / variant spellings (stable, naming-independent).
2. Device-name substring match as a last-resort fallback when all arch
attrs are absent (AMD SDK / Radeon wheels may not populate them):
- gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M``
- gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395),
``Radeon 8050S`` (cut-down SKU)
"""
gcn_arch = ""
for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"):
_v = (getattr(props, _attr, "") or "").split(":")[0].strip()
if _v:
gcn_arch = _v
break
if gcn_arch:
return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"}
# Arch attrs absent — fall back to device-name matching.
dev_lower = (getattr(props, "name", "") or "").lower()
is_unified = (
"890m" in dev_lower
or "880m" in dev_lower
or "8060s" in dev_lower
or "8050s" in dev_lower
)
return gcn_arch, is_unified
def _tilelang_platform_supported() -> bool:
"""True iff a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch.
@ -881,6 +1004,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
_ensure_tilelang_backend_unconditional(eq)
def _causal_conv1d_install(eq: Any) -> bool:
if sys.platform == "win32":
logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping")
return False
ok = _install_package_wheel_first(
event_queue = eq,
import_name = "causal_conv1d",
@ -1893,6 +2019,452 @@ def run_training_process(
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── 1d. Stub torchao on Windows ROCm ──
# torchao (pulled in by transformers.quantizers) imports
# torch.distributed._functional_collectives at module level, which imports
# distributed_c10d.py unconditionally — that file crashes on Windows ROCm
# because torch._C._distributed_c10d (the RCCL backend) is absent.
# torch/distributed/__init__.py itself is guarded by `if is_available()`
# so `import torch.distributed` alone is safe; the crash only comes via
# torchao's import chain. Stubbing torchao short-circuits it entirely.
# _StubSubpackageFinder handles any depth of torchao.xxx.yyy imports.
import types as _types
import importlib.machinery as _ilm
import importlib.abc as _ilabc
_STUB_SENTINEL = object()
# Metaclass for stub types so that isinstance(x, StubClass) returns False
# instead of raising TypeError ("arg 2 must be a type").
# peft/tuners/lora/torchao.py does:
# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor
# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor))
# If those names resolve to stub modules rather than types, isinstance() raises.
class _StubTypeMeta(type):
def __instancecheck__(cls, instance):
return False
def __subclasscheck__(cls, subclass):
return False
def __getattr__(cls, attr):
if attr.startswith("__"):
raise AttributeError(attr)
child = _StubTypeMeta(attr, (), {})
setattr(cls, attr, child)
return child
def __call__(cls, *args, **kwargs):
return None
def _make_stub_type(name):
"""Stub class: accepted by isinstance() (always False), supports attr access."""
return _StubTypeMeta(name, (), {})
def _make_mod_stub(mod_name):
m = _types.ModuleType(mod_name)
m.__path__ = []
m.__package__ = mod_name
m._unsloth_stub = _STUB_SENTINEL
m.__spec__ = _ilm.ModuleSpec(mod_name, loader = None, is_package = True)
def _ga(attr, _m = m, _n = mod_name):
if attr.startswith("__"):
raise AttributeError(attr)
# Return a stub CLASS (not a module) so that isinstance(x, attr)
# works and returns False instead of raising TypeError.
child = _make_stub_type(f"{_n}.{attr}")
setattr(_m, attr, child)
return child
m.__getattr__ = _ga
return m
class _StubSubpackageLoader(_ilabc.Loader):
def __init__(self, mod_name):
self._mod_name = mod_name
def create_module(self, spec):
return _make_mod_stub(self._mod_name)
def exec_module(self, module):
pass
class _StubSubpackageFinder(_ilabc.MetaPathFinder):
def find_spec(self, fullname, path, target = None):
if "." not in fullname:
return None
parent = sys.modules.get(fullname.rsplit(".", 1)[0])
if parent is None:
return None
if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL:
return None
return _ilm.ModuleSpec(
fullname, _StubSubpackageLoader(fullname), is_package = True
)
# Only stub torchao on Windows ROCm hosts -- on Windows CUDA (NVIDIA) torchao
# is real and shadowing it breaks torchao-based quantization paths.
# 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:
# Register the finder only on Windows ROCm -- on other platforms there
# are no stub modules seeded, so appending is a pure accumulation.
sys.meta_path.append(_StubSubpackageFinder())
# Seed torchao top-level + key submodules; the finder handles the rest.
for _tao_name in (
"torchao",
"torchao.quantization",
"torchao.dtypes",
"torchao.float8",
"torchao.utils",
):
if _tao_name not in sys.modules:
sys.modules[_tao_name] = _make_mod_stub(_tao_name)
# ── 1e. Ensure torch.distributed helper attrs are present ──
# Single-GPU training never initialises the process group, so these helpers
# are never called — but transformers/trl import them unconditionally.
_td_stubs = {
"is_initialized": lambda: False,
"is_available": lambda: False,
"is_torchelastic_launched": lambda: False,
"get_rank": lambda: 0,
"get_world_size": lambda: 1,
"barrier": lambda: None,
}
try:
import torch.distributed as _td
for _name, _stub in _td_stubs.items():
if not hasattr(_td, _name):
setattr(_td, _name, _stub)
except Exception:
_td_mock = _types.ModuleType("torch.distributed")
for _name, _stub in _td_stubs.items():
setattr(_td_mock, _name, _stub)
sys.modules["torch.distributed"] = _td_mock
try:
import torch as _torch
_torch.distributed = _td_mock
except Exception:
pass
# ── 1f. Windows ROCm runtime patches ──
# torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm ≤ 7.12 Windows),
# causing 0xC0000005 (access violation) during training.
#
# Root cause: the JitDecomp autograd decomposition system (NOT torch.compile)
# dispatches _grouped_mm → _fused_adagrad_ → _grouped_mm HIP → null crash.
# TORCHDYNAMO_DISABLE=1 stops the compiler frontend but does NOT stop
# JitDecomp, so we must also override the CUDA dispatch key for _grouped_mm
# with a safe Python fallback.
#
# Fixed in AMD's wheel: torch==2.11.0+rocm7.13.0 — the 3-D batch and grouped
# (with offs) variants of _grouped_mm now have working HIP kernels on gfx1200.
# We gate the dispatch override on HIP < 7.13 so users on the fixed wheel get
# the real GPU kernel rather than our Python fallback.
#
# Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0.
#
# Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None,
# Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor
# offs: optional group-split offsets (MoE-style variable-size batches)
#
# torch is already in sys.modules from section 1e's `import torch.distributed`.
# 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")
# Broad check: torch.version.hip OR "rocm" in torch.__version__.
# AMD SDK / Radeon Windows wheels do not always populate
# torch.version.hip; without the broad check the BNB version pin,
# dynamo-disable, and _grouped_mm fallback below silently skip
# (matches the torchao stub gate above and main.py).
_build_version_for_rocm = (
getattr(_torch_for_rocm, "__version__", "").lower()
if _torch_for_rocm is not None
else ""
)
_is_win_rocm_torch = bool(
_torch_for_rocm is not None
and (
getattr(getattr(_torch_for_rocm, "version", None), "hip", None)
or "rocm" in _build_version_for_rocm
)
)
if _is_win_rocm_torch:
# Disable dynamo (belt-and-suspenders; JitDecomp patch below is the
# real fix, but keeping dynamo off avoids any other compile paths).
if "TORCHDYNAMO_DISABLE" not in os.environ:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.info("Windows ROCm: torch.compile (dynamo) disabled")
# BNB auto-detects the HIP version from torch.version.hip and uses
# it to choose which DLL to load (e.g. "7.13" → rocm713.dll).
# AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its
# version suffix does not always match the torch HIP version (e.g.
# torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still
# ships rocm72.dll). We detect the actual DLL name from the installed
# package and override BNB's auto-detection. "72" is a safe fallback
# if detection fails. Callers may override by pre-setting the var.
if "BNB_ROCM_VERSION" not in os.environ:
_bnb_rocm_ver = None
try:
import glob as _glob
import importlib.util as _ilu
import re as _re
_bnb_spec = _ilu.find_spec("bitsandbytes")
if _bnb_spec and _bnb_spec.submodule_search_locations:
_all_vers: list[str] = []
for _pkg_dir in _bnb_spec.submodule_search_locations:
for _dll in _glob.glob(
os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
):
_m = _re.search(
r"libbitsandbytes_rocm(\d+)\.dll",
os.path.basename(_dll),
)
if _m:
_all_vers.append(_m.group(1))
# Pick the highest numeric suffix so that e.g. "713"
# wins over "72" when both variants are present.
# Filesystem glob order is not guaranteed, so always
# sort rather than stopping at the first match.
if _all_vers:
_bnb_rocm_ver = max(_all_vers, key = lambda v: int(v))
except Exception:
pass
_bnb_rocm_ver = _bnb_rocm_ver or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
logger.info(
"Windows ROCm: set BNB_ROCM_VERSION=%s "
"(detected from installed BNB wheel; "
"overrides torch.version.hip auto-detection)",
_bnb_rocm_ver,
)
# Parse HIP version for the kernel-fix gate below.
# torch.version.hip can be "7.13.99004", "7.2.0", etc.
# AMD SDK / Radeon wheels may leave torch.version.hip unset and
# encode the ROCm version in torch.__version__ instead
# (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back
# to that string when version.hip is missing.
def _hip_ver_at_least(major: int, minor: int) -> bool:
import re as _re_ver
_hip_str = getattr(
getattr(_torch_for_rocm, "version", None), "hip", None
)
if not _hip_str:
# Try the standard "+rocmX.Y.Z" embedded version first
# (e.g. "2.11.0+rocm7.13.0").
_ver_match = _re_ver.search(
r"rocm(\d+)\.(\d+)", _build_version_for_rocm
)
if _ver_match:
return (
int(_ver_match.group(1)),
int(_ver_match.group(2)),
) >= (major, minor)
# AMD SDK / Radeon Windows wheels encode the build as
# "+rocmsdk<date>" (e.g. "2.9.0+rocmsdk20251116") with no
# explicit rocmX.Y component. The rocmsdk format was
# introduced after the gfx120X null-kernel fix landed in
# ROCm 7.13, so any wheel with this suffix is new enough to
# have working HIP kernels. Treat as >= 7.13 rather than
# falling back to False and installing the Python workaround
# on a wheel that doesn't need it.
if "rocmsdk" in _build_version_for_rocm:
logger.debug(
"Windows ROCm: AMD SDK wheel detected (%r); "
"assuming HIP >= %d.%d (rocmsdk wheels post-date "
"the gfx120X null-kernel fix)",
_build_version_for_rocm,
major,
minor,
)
return True
return False
try:
_parts = [int(x) for x in str(_hip_str).split(".")[:2]]
if len(_parts) < 2:
logger.warning(
"Windows ROCm: torch.version.hip %r has fewer than "
"two components; cannot compare against %d.%d",
_hip_str,
major,
minor,
)
return False
return (_parts[0], _parts[1]) >= (major, minor)
except ValueError:
logger.warning(
"Windows ROCm: could not parse torch.version.hip %r as "
"a version number; assuming HIP < %d.%d",
_hip_str,
major,
minor,
)
return False
# _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12,
# causing 0xC0000005. AMD fixed it in ROCm 7.13 (torch 2.11+).
# Only install the Python fallback on the affected versions so users
# on 7.13+ get the real GPU kernel for MoE workloads.
if not _hip_ver_at_least(7, 13):
try:
import warnings as _warnings
_gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
def _grouped_mm_safe_impl(
self, mat2, offs = None, bias = None, out_dtype = None
):
"""Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
_t = _torch_for_rocm
if offs is None:
# No offsets: behave like the real op, which
# accepts either (M, K) x (K, N) -> mm, or 3-D
# batched inputs -> bmm. Picking torch.mm
# unconditionally previously raised "self must be
# a matrix" on 3-D MoE workloads.
if self.dim() == 3 and mat2.dim() == 3:
result = _t.bmm(self.contiguous(), mat2.contiguous())
elif self.dim() == 3 and mat2.dim() == 2:
# Broadcast 2-D mat2 across the batch dim.
result = _t.matmul(self.contiguous(), mat2.contiguous())
elif self.dim() == 2 and mat2.dim() == 3:
# Broadcast 2-D self across batch via matmul semantics.
result = _t.matmul(self.contiguous(), mat2.contiguous())
else:
result = _t.mm(self.contiguous(), mat2.contiguous())
else:
# Grouped case: offs[i] is the exclusive end-row of
# group i in `self`; mat2 may be 3-D or 2-D.
offs_list = offs.tolist()
pieces = []
prev = 0
for idx, end in enumerate(offs_list):
end = int(end)
a_part = self[prev:end].contiguous()
if mat2.dim() == 3:
b_part = mat2[idx].contiguous()
else:
b_part = mat2.contiguous()
pieces.append(_t.mm(a_part, b_part))
prev = end
# Include any trailing rows not covered by offs
if prev < self.shape[0]:
a_tail = self[prev:].contiguous()
b_tail = (
mat2[-1].contiguous()
if mat2.dim() == 3
else mat2.contiguous()
)
pieces.append(_t.mm(a_tail, b_tail))
result = (
_t.cat(pieces, dim = 0)
if pieces
else _t.zeros(
0,
mat2.shape[-1],
device = self.device,
dtype = self.dtype,
)
)
if bias is not None:
result = result + bias
if out_dtype is not None:
result = result.to(out_dtype)
elif result.dtype != self.dtype:
result = result.to(self.dtype)
return result
with _warnings.catch_warnings():
_warnings.simplefilter("ignore")
_gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
_WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC
logger.info(
"Windows ROCm: patched _grouped_mm CUDA dispatch "
"(null HIP kernel on gfx1200, ROCm ≤ 7.12 — "
"bypassed with Python mm fallback)"
)
except Exception as _patch_exc:
logger.warning(
"Windows ROCm: could not patch _grouped_mm — "
"training may crash with 0xC0000005: %s",
_patch_exc,
)
else:
logger.info(
"Windows ROCm: HIP >= 7.13 — _grouped_mm kernel is functional, "
"skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)"
)
# ── 1g. ROCm OOM guard ──
# On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can
# cause a HIP driver hang that freezes the entire system rather than
# raising a Python exception. set_per_process_memory_fraction caps the
# HIP allocator so PyTorch raises OutOfMemoryError before hitting the
# hardware limit, giving the UI a clean error instead of a system freeze.
# Only applied on ROCm -- NVIDIA CUDA has a graceful OOM path and does
# not need this cap.
# Unified-memory APUs (gfx1150 Strix Point / gfx1151 Strix Halo) share GPU
# and system RAM in one pool: 0.90 of 128 GB starves the OS. Use 0.80 there.
# Primary classifier: gcnArchName from device properties — stable within a
# product family and naming-independent. AMD SDK / Radeon wheels may omit
# gcnArchName or expose it under a variant spelling, so we try several attr
# names then fall back to known device-name markers as a last resort.
# Non-fatal: silently skipped if torch is not importable.
if _hw.IS_ROCM:
try:
import torch as _torch_mem
if _torch_mem.cuda.is_available():
# Classify unified vs discrete via _rocm_classify_unified_memory.
# See that function's docstring for classification priority.
_props = _torch_mem.cuda.get_device_properties(0)
_dev_name = _props.name
_gcn_arch, _is_unified = _rocm_classify_unified_memory(_props)
if _is_unified and not _gcn_arch:
logger.debug(
"ROCm OOM guard: gcnArchName absent -- inferred "
"unified memory from device name %r; applying 0.80 cap",
_dev_name,
)
_mem_fraction = 0.80 if _is_unified else 0.90
_torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction)
logger.info(
"ROCm OOM guard: set_per_process_memory_fraction(%.2f) — "
"%s memory host (%s, %s)",
_mem_fraction,
"unified" if _is_unified else "discrete",
_dev_name,
_gcn_arch or "unknown arch",
)
except Exception as _oom_guard_err:
logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err)
# ── 2. Now import ML libraries (fresh in this clean process) ──
try:
_send_status(event_queue, "Importing Unsloth...")
@ -2347,14 +2919,38 @@ def run_training_process(
)
except Exception as exc:
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
_exc_str = str(exc).lower()
_is_oom = (
"out of memory" in _exc_str
or "hip out of memory" in _exc_str
or "cuda out of memory" in _exc_str
or type(exc).__name__ == "OutOfMemoryError"
)
if _is_oom:
_oom_msg = (
"GPU ran out of VRAM during training.\n"
"To fix: reduce max_seq_length (e.g. 2048–4096), enable "
"gradient_checkpointing=True, lower per_device_train_batch_size, "
"or use a smaller model / higher quantization."
)
logger.error("Training stopped: GPU OOM — %s", exc)
event_queue.put(
{
"type": "error",
"error": _oom_msg,
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
else:
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
def _send_status(event_queue: Any, message: str) -> None:

104
studio/backend/main.py
View file

@ -12,6 +12,110 @@ from pathlib import Path as _Path
# Suppress annoying C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
if sys.platform == "win32":
# Retained at module scope -- os.add_dll_directory returns a handle that
# removes the search-path entry when garbage collected.
_ROCM_DLL_HANDLES: list = []
def _add_rocm_dll_dirs() -> None:
candidates = []
# 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer
for _var in ("HIP_PATH", "ROCM_PATH"):
_val = os.environ.get(_var)
if _val:
candidates.append(os.path.join(_val, "bin"))
# 2. Standard AMD installer location: C:\Program Files\AMD\ROCm\<ver>\bin
# Scan all installed versions, newest first.
_default_root = os.path.join(
os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
)
def _ver_key(name: str) -> tuple:
# Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string.
parts = []
for chunk in name.split("."):
try:
parts.append((0, int(chunk)))
except ValueError:
parts.append((1, chunk))
return tuple(parts)
try:
if os.path.isdir(_default_root):
for _ver in sorted(
os.listdir(_default_root), key = _ver_key, 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()
del _add_rocm_dll_dirs
# ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
# bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm<ver>.dll
# where <ver> comes from torch.version.hip (e.g. "7.13..." → "713").
# The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without
# this the server process crashes with "Configured ROCm binary not found".
# Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION
# before any import that pulls in bitsandbytes (mirrors worker.py logic).
# Gate on the rocm bnb DLL (the exact file this configures) or HIP_PATH/
# ROCM_PATH, not on torch.version.hip: that needed importing torch on every
# Windows host (NVIDIA/CPU included), adding seconds to startup. Radeon
# wheels without HIP_PATH still ship the rocm bnb DLL, so they are covered.
if "BNB_ROCM_VERSION" not in os.environ:
import glob as _glob
import logging as _logging
_hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"))
_bnb_rocm_ver = None
_found_rocm_bnb = False
try:
import importlib.util as _ilu
_bnb_spec = _ilu.find_spec("bitsandbytes")
# submodule_search_locations (not spec.origin) handles editable installs.
if _bnb_spec and _bnb_spec.submodule_search_locations:
import re as _re_bnb
_all_vers_main: list[str] = []
for _pkg_dir in _bnb_spec.submodule_search_locations:
for _dll in _glob.glob(
os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
):
_found_rocm_bnb = True
_km = _re_bnb.search(
r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll)
)
if _km:
_all_vers_main.append(_km.group(1))
if _all_vers_main:
_bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v))
except Exception as _e:
_logging.getLogger(__name__).warning(
"Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'",
_e,
)
# rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72").
if _found_rocm_bnb or _hip_env:
_bnb_rocm_ver_final = _bnb_rocm_ver or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
_logging.getLogger(__name__).info(
"Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)",
_bnb_rocm_ver_final,
)
# Ensure backend dir is on sys.path so _platform_compat is importable when
# main.py is launched directly (e.g. `uvicorn main:app`).
_backend_dir = str(_Path(__file__).parent)

52
studio/backend/tests/test_amd_apu_unified_memory.py Normal file
View file

@ -0,0 +1,52 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs
(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS."""
from __future__ import annotations
import sys
import types
import pytest
from core.inference.llama_cpp import LlamaCppBackend
def _fake_torch(hip, archs, *, cuda_ok = True):
t = types.ModuleType("torch")
t.version = types.SimpleNamespace(hip = hip)
t.cuda = types.SimpleNamespace(
is_available = lambda: cuda_ok,
device_count = lambda: len(archs),
get_device_properties = lambda i: types.SimpleNamespace(gcnArchName = archs[i]),
)
return t
@pytest.mark.parametrize(
"hip,archs,expected",
[
("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped)
("6.2.0", ["gfx1150"], True), # Strix Point APU
("6.2.0", ["gfx1100"], False), # discrete RDNA3
("6.2.0", ["gfx1201"], False), # discrete RDNA4
("6.2.0", ["gfx942"], False), # MI300X (data center)
(None, ["sm_90"], False), # NVIDIA (no torch.version.hip)
("6.2.0", ["gfx1100", "gfx1151"], True), # mixed dGPU + APU
],
)
def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip, archs))
assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected
def test_cpu_no_cuda_returns_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False))
assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
def test_missing_torch_returns_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", None)
assert LlamaCppBackend._amd_apu_wants_unified_memory() is False

430
studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py Normal file
View file

@ -0,0 +1,430 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets.
Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade
GitHub API are stubbed out so the suite runs without internet access and is
not subject to rate limits.
"""
from __future__ import annotations
import importlib
import sys
from pathlib import Path
from unittest.mock import patch
import pytest
_studio = Path(__file__).resolve().parent.parent.parent
if str(_studio) not in sys.path:
sys.path.insert(0, str(_studio))
_mod = importlib.import_module("install_llama_prebuilt")
HostInfo = _mod.HostInfo
resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None)
_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None)
if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None:
pytest.skip("PR symbols not present - check branch", allow_module_level = True)
@pytest.fixture(autouse = True)
def _clear_lemonade_release_cache():
"""Prevent cross-test pollution of the lemonade release lru_cache when
future tests vary the fetch_json mock return value."""
_cache = getattr(_mod, "_fetch_lemonade_release_cached", None)
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
yield
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
_STUB_TAG = "b1262"
_STUB_OS_PREFIXES = ("ubuntu", "windows")
_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X")
def _stub_lemonade_release() -> dict:
"""Minimal lemonade release payload covering all supported GPU/OS combinations."""
assets = [
{
"name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip",
"browser_download_url": (
f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/"
f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip"
),
}
for prefix in _STUB_OS_PREFIXES
for family in _STUB_FAMILIES
]
return {"tag_name": _STUB_TAG, "assets": assets}
def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo:
return HostInfo(
system = "Windows" if windows else "Linux",
machine = "amd64" if windows else "x86_64",
is_windows = windows,
is_linux = not windows,
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,
rocm_gfx_target = gfx_target,
)
def _lookup_family(gfx: str) -> str | None:
for prefix, family in _LEMONADE_GFX_FAMILIES:
if gfx.startswith(prefix):
return family
return None
# ---------------------------------------------------------------------------
# GPU family mapping
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"gfx,expected_family",
[
("gfx1151", "gfx1151"),
("gfx1150", "gfx1150"),
("gfx1201", "gfx120X"),
("gfx1200", "gfx120X"),
("gfx1100", "gfx110X"),
("gfx1030", "gfx103X"),
],
)
def test_gpu_family_mapping(gfx, expected_family):
assert _lookup_family(gfx) == expected_family
def test_unknown_gpu_not_in_families():
assert _lookup_family("gfx999") is None
# ---------------------------------------------------------------------------
# Asset resolution - hits real lemonade GitHub API
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"gfx,os_prefix,windows",
[
("gfx1151", "ubuntu", False),
("gfx1150", "ubuntu", False),
("gfx1201", "ubuntu", False),
("gfx1100", "ubuntu", False),
("gfx1030", "ubuntu", False),
("gfx1151", "windows", True),
("gfx1100", "windows", True),
],
)
def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows):
host = _make_rocm_host(gfx, windows = windows)
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
result = resolve_lemonade_rocm_choice(
host, os_prefix, "default", llama_tag = "latest"
)
assert (
result is not None
), f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})"
assert _lookup_family(gfx) in result.name
assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm")
def test_unknown_gpu_falls_through_to_upstream():
host = _make_rocm_host("gfx999")
result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest")
assert result is None
# ---------------------------------------------------------------------------
# Simple-policy dispatcher must plan a lemonade ROCm attempt for AMD-only hosts.
# This is the path setup.sh actually invokes (via --simple-policy), so the
# lemonade integration is useless if it isn't wired in here.
# ---------------------------------------------------------------------------
direct_linux_release_plan = getattr(_mod, "direct_linux_release_plan", None)
direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None)
def _stub_unsloth_release(release_tag: str = "b9022") -> dict:
# Minimal payload that parse_direct_linux_release_bundle accepts. It
# requires at least one `app-{label}-linux-x64*.tar.gz` asset for the
# bundle to be recognised; we ship a bare CPU one so the planner has a
# baseline non-ROCm attempt to fall through to.
asset_name = f"app-{release_tag}-linux-x64.tar.gz"
return {
"tag_name": release_tag,
"name": release_tag,
"assets": [
{
"name": asset_name,
"browser_download_url": f"https://example.invalid/{asset_name}",
},
],
}
@pytest.mark.skipif(
direct_linux_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
)
def test_simple_policy_plans_lemonade_for_rocm_host():
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_linux_release_plan(
_stub_unsloth_release(),
host,
"unslothai/llama.cpp",
"latest",
)
assert plan is not None, "ROCm host should not be skipped by simple-policy planner"
kinds = [a.install_kind for a in plan.attempts]
assert (
"linux-rocm" in kinds
), f"simple-policy planner did not include a lemonade ROCm attempt; got {kinds}"
rocm_attempt = next(a for a in plan.attempts if a.install_kind == "linux-rocm")
assert rocm_attempt.source_label == "lemonade"
assert "gfx1151" in rocm_attempt.name
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
)
def test_simple_policy_plans_lemonade_for_windows_hip_host():
host = _make_rocm_host("gfx1151", windows = True)
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [],
}
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_upstream_release_plan(
release, host, "ggml-org/llama.cpp", "latest"
)
assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt"
kinds = [a.install_kind for a in plan.attempts]
assert (
"windows-hip" in kinds
), f"simple-policy planner did not include a lemonade HIP attempt; got {kinds}"
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
)
def test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavailable():
"""If lemonade returns None (e.g. gfx999 or transient API failure), the planner
must still include the upstream HIP asset rather than silently downgrading to CPU."""
host = _make_rocm_host("gfx999", windows = True)
hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip"
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [
{
"name": hip_asset,
"browser_download_url": f"https://example.invalid/{hip_asset}",
},
],
}
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None
kinds = [a.install_kind for a in plan.attempts]
assert (
"windows-hip" in kinds
), f"upstream HIP asset not included as fallback; got {kinds}"
hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip")
assert hip_attempt.source_label == "upstream"
# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ──
def test_lemonade_release_api_url_pinned_tag():
"""A pinned llama_tag must produce the /releases/tags/<tag> URL."""
assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262")
assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest")
assert _mod._lemonade_release_api_for("").endswith("/releases/latest")
def test_lemonade_release_api_url_encodes_tag():
"""Unexpected slashes / hashes in the tag must be URL-encoded so the URL
cannot be reshaped (defence in depth -- tags should already be sanitised
upstream)."""
url = _mod._lemonade_release_api_for("b1260/../latest")
assert "/releases/tags/b1260%2F..%2Flatest" in url
assert "//latest" not in url.split("/releases/tags/", 1)[1]
def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch):
"""UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver."""
monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1")
host = _make_rocm_host("gfx1151")
res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
assert res is None
def test_lemonade_resolver_rejects_non_github_url(monkeypatch):
"""If the GitHub API response somehow contained an off-host download URL,
the resolver must refuse to use it (lemonade assets are not in the
approved-hash manifest)."""
bad_release = {
"tag_name": _STUB_TAG,
"assets": [
{
"name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
"browser_download_url": "https://attacker.invalid/llama.zip",
},
],
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = bad_release):
res = resolve_lemonade_rocm_choice(
host, "ubuntu", "linux-rocm", llama_tag = "latest"
)
assert res is None
def test_lemonade_resolver_rejects_http_scheme():
assert not _mod._is_trusted_github_release_url(
"http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_accepts_github_cdn():
# Real GitHub release CDN URLs carry the /github-production-release-asset- prefix.
assert _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_rejects_arbitrary_cdn_path():
# A CDN URL without the release-asset path prefix must be rejected.
assert not _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/abc/def",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_accepts_release_path():
url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip"
assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm")
def test_lemonade_resolver_rejects_wrong_repo():
"""A github.com release URL for a different repo must be rejected."""
assert not _mod._is_trusted_github_release_url(
"https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_rejects_empty_browser_download_url():
"""An asset entry with an empty browser_download_url must fall through."""
release = {
"tag_name": _STUB_TAG,
"assets": [
{
"name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
"browser_download_url": "",
},
],
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = release):
res = resolve_lemonade_rocm_choice(
host, "ubuntu", "linux-rocm", llama_tag = "latest"
)
assert res is None
def test_lemonade_runtime_patterns_include_hip_runtime():
"""linux-rocm overlay must use a broad lib glob to catch all bundled .so files.
Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp,
...) whose names change across ROCm releases. A broad ``lib*.so*`` glob
avoids having to enumerate every transitive dependency by name.
"""
from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice
choice = AssetChoice(
repo = "lemonade-sdk/llamacpp-rocm",
tag = "b1262",
name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip",
url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip",
source_label = "lemonade",
install_kind = "linux-rocm",
)
pats = runtime_patterns_for_choice(choice)
# The broad glob must be present so every .so in the lemonade bundle
# (including transitive deps added in future ROCm releases) gets overlaid.
assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}"
_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None)
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch):
"""AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES;
on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100."""
# Two GPUs; rocminfo reports each token twice (as in the real tool output).
probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100"
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1")
assert _pick_rocm_gfx_target(probe_out) == "gfx1100"
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch):
"""CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None."""
probe_out = "gfx1151\ngfx1100"
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1")
assert _pick_rocm_gfx_target(probe_out) is None
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch):
"""Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must
return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the
two gfx1100 entries into one and making index 2 out of range."""
# Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
# Each GPU gets its own Agent section with a few token mentions.
probe_out = (
"***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n"
"***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n"
"***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n"
)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2")
assert _pick_rocm_gfx_target(probe_out) == "gfx1151"

24
studio/backend/tests/test_log_filter_no_truncation.py
View file

@ -2,27 +2,11 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Regression tests for studio.backend.loggers.handlers.filter_sensitive_data.
Regression tests for loggers.handlers.filter_sensitive_data.
Context: filter_sensitive_data was originally written with a base64-detection
heuristic that truncated any string >100 chars containing ',' or '/' down to
20 chars + '...'. The block was dormant until PR #5246 wired the processor
into the structlog chain to redact native-path leases. Once active, the
heuristic ate normal log lines emitted by llama_cpp_backend (GGUF size
summary, mmproj selection, the full llama-server command line) and any
exception traceback that happened to contain a file path.
These tests pin two properties:
1. Long, comma- or slash-bearing log messages flow through filter_sensitive_data
unchanged. The exact strings exercised match the call sites at
studio/backend/core/inference/llama_cpp.py:2117, :2283, and :2312 that
were truncated in the original bug report.
2. PR #5246's native-path lease redaction still fires for both the inline
``native_path_lease=...`` regex form and the ``nativePathLease`` dict-key
form. This guards against future regressions that strip redaction along
with the truncation block.
Pins two properties: (1) long strings with commas/slashes pass through
unchanged (the base64-truncation heuristic from PR #5246 was too aggressive),
and (2) native-path lease redaction still fires for both inline and dict-key forms.
"""
from loggers.handlers import filter_sensitive_data

176
studio/backend/tests/test_rocm_oom_guard.py Normal file
View file

@ -0,0 +1,176 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Unit tests for _rocm_classify_unified_memory (ROCm OOM-guard classifier).
Covers the three classification paths:
Path 1 – canonical gcnArchName attribute present.
Path 2 – gcnArchName absent, alternate-spelling attribute present.
Path 3 – ALL arch attrs absent; falls back to device-name substring match.
Regression for: Strix Halo (gfx1151) misclassified as discrete on AMD SDK /
Radeon wheels that populate props.name = "Radeon 8060S Graphics" but do NOT
set any gcnArchName attribute. Without the 8060s/8050s name patterns the
fallback returned is_unified=False, applying the 0.90 fraction instead of
0.80 and leaving only ~12.8 GiB OS headroom on a 128 GiB unified-memory pool.
"""
from __future__ import annotations
from types import SimpleNamespace
import pytest
from core.training.worker import _rocm_classify_unified_memory
# ── helpers ──────────────────────────────────────────────────────────────────
def _props(**kwargs) -> SimpleNamespace:
"""Build a fake device-properties object with the given attributes."""
return SimpleNamespace(**kwargs)
# ── Path 1: canonical gcnArchName ────────────────────────────────────────────
class TestCanonicalGcnArchName:
"""gcnArchName is present and populated."""
@pytest.mark.parametrize(
"arch, expected_unified",
[
("gfx1150", True), # Strix Point
("gfx1151", True), # Strix Halo
("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete
("gfx906", False), # MI50 — discrete server GPU
("gfx1201", False), # RX 9070 XT — discrete
],
)
def test_canonical_attr(self, arch: str, expected_unified: bool) -> None:
props = _props(gcnArchName = arch, name = "irrelevant")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == arch
assert is_unified is expected_unified
def test_arch_with_colon_suffix_stripped(self) -> None:
"""gcnArchName can carry xnack/sramecc suffix; only the base is kept."""
props = _props(gcnArchName = "gfx1151:xnack-", name = "irrelevant")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1151"
assert is_unified is True
def test_canonical_attr_wins_over_name(self) -> None:
"""Arch attr takes priority; device name should be ignored."""
# Discrete arch, but name looks like a unified SKU — arch must win.
props = _props(gcnArchName = "gfx1100", name = "Radeon 890M")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1100"
assert is_unified is False
# ── Path 2: alternate-spelling fallback ──────────────────────────────────────
class TestAlternateSpellingFallback:
"""gcnArchName is missing but an alternate attr spelling is present."""
@pytest.mark.parametrize(
"attr_name",
["gcn_arch_name", "arch_name", "gfx_arch_name"],
)
def test_alternate_attr_unified(self, attr_name: str) -> None:
props = _props(**{attr_name: "gfx1151"}, name = "Radeon 8060S Graphics")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1151"
assert is_unified is True
@pytest.mark.parametrize(
"attr_name",
["gcn_arch_name", "arch_name", "gfx_arch_name"],
)
def test_alternate_attr_discrete(self, attr_name: str) -> None:
props = _props(**{attr_name: "gfx1201"}, name = "Radeon RX 9070 XT")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1201"
assert is_unified is False
def test_first_non_empty_attr_wins(self) -> None:
"""When multiple alternate attrs are present the first non-empty one wins."""
props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1151"
assert is_unified is True
# ── Path 3: device-name fallback ─────────────────────────────────────────────
class TestDeviceNameFallback:
"""ALL arch attrs absent — classifier must rely solely on device name."""
# --- unified-memory devices that MUST be detected ---
@pytest.mark.parametrize(
"device_name",
[
# gfx1150 Strix Point
"Radeon 890M",
"AMD Radeon 890M Graphics",
"RADEON 890M", # case-insensitive
"Radeon 880M",
"AMD Radeon 880M Graphics",
# gfx1151 Strix Halo — the regression case from the review
"Radeon 8060S Graphics", # Ryzen AI MAX+ 395 (as returned by torch)
"AMD Radeon 8060S",
"Radeon 8050S Graphics", # cut-down Strix Halo SKU
"AMD Radeon 8050S",
# case variants
"RADEON 8060S GRAPHICS",
"radeon 8050s",
],
)
def test_unified_memory_detected(self, device_name: str) -> None:
props = _props(name = device_name)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "", f"expected empty gcn_arch, got {gcn!r}"
assert (
is_unified is True
), f"device {device_name!r} should be classified as unified-memory"
# --- discrete devices that must NOT be mis-classified ---
@pytest.mark.parametrize(
"device_name",
[
"Radeon RX 9070 XT",
"AMD Radeon RX 7900 XTX",
"Radeon RX 6900 XT",
"Radeon Pro W7900",
"AMD Instinct MI300X",
# Names that contain superficially similar substrings but are discrete
"Radeon RX 580",
"Radeon VII",
],
)
def test_discrete_not_misclassified(self, device_name: str) -> None:
props = _props(name = device_name)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert (
is_unified is False
), f"discrete device {device_name!r} should NOT be classified as unified-memory"
def test_empty_name_returns_false(self) -> None:
"""Completely absent name must not crash and must default to discrete."""
props = _props() # no 'name' attr at all
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert is_unified is False
def test_none_name_returns_false(self) -> None:
props = _props(name = None)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert is_unified is False

62
studio/backend/utils/hardware/amd.py
View file

@ -11,18 +11,35 @@ nvidia.py counterparts.
import json
import math
import os
import platform
import re
import subprocess
import sys
from typing import Any, Optional
from loggers import get_logger
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = get_logger(__name__)
# amd-smi on Windows must initialise the full ROCm runtime on first call, which
# can take 15-25 s on cold hardware. Linux is consistently < 2 s.
_AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10
def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
# Circuit breaker: stop calling amd-smi after this many consecutive failures.
# On Windows, each failed call spawns a process that may show a UAC/DiskPart
# elevation prompt. Once we know amd-smi doesn't work we stop polling it.
_AMD_SMI_FAILURE_LIMIT = 3
_amd_smi_consecutive_failures = 0
_amd_smi_disabled = False
def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]:
"""Run amd-smi with the given arguments and return parsed JSON, or None."""
global _amd_smi_consecutive_failures, _amd_smi_disabled
if _amd_smi_disabled:
return None
try:
result = subprocess.run(
["amd-smi", *args, "--json"],
@ -30,13 +47,40 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
text = True,
timeout = timeout,
env = child_env_without_native_path_secret(),
**windows_hidden_subprocess_kwargs(),
)
except (OSError, subprocess.TimeoutExpired) as e:
logger.warning("amd-smi query failed: %s", e)
if isinstance(e, FileNotFoundError):
# amd-smi ships with Adrenalin, not the HIP SDK -- absence is
# expected on HIP SDK-only Windows setups. Log at debug only.
logger.debug("amd-smi not found (not in PATH): %s", e)
else:
logger.warning("amd-smi query failed: %s", e)
_amd_smi_consecutive_failures += 1
if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
logger.info(
"amd-smi not available (not installed; expected on HIP SDK-only systems); "
"GPU VRAM polling disabled"
)
_amd_smi_disabled = True
return None
if result.returncode != 0 or not result.stdout.strip():
if result.returncode != 0:
logger.warning("amd-smi returned code %d", result.returncode)
_amd_smi_consecutive_failures += 1
if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
logger.info(
"amd-smi not available (not installed; expected on HIP SDK-only systems); "
"GPU VRAM polling disabled"
)
_amd_smi_disabled = True
return None
if not result.stdout.strip():
# amd-smi exited successfully but produced no output (e.g. no GPUs
# visible on this query, or a version that emits nothing for --json).
# This is not a tool failure, so don't count against the circuit breaker.
logger.debug("amd-smi exited 0 but returned no output")
return None
_amd_smi_consecutive_failures = 0 # reset on success
try:
return json.loads(result.stdout)
except json.JSONDecodeError:
@ -352,7 +396,7 @@ def get_visible_gpu_utilization(
)
parsed_id = _parse_numeric(raw_id)
if parsed_id is None:
logger.debug(
logger.warning(
"amd-smi GPU id %r could not be parsed; falling back to "
"enumeration index %d",
raw_id,
@ -360,7 +404,15 @@ def get_visible_gpu_utilization(
)
idx = fallback_idx
else:
idx = int(parsed_id)
rounded = round(parsed_id)
if rounded != parsed_id:
logger.warning(
"amd-smi GPU id %r parsed as non-integer %r; truncating to %d",
raw_id,
parsed_id,
rounded,
)
idx = int(rounded)
if idx not in visible_set:
continue
metrics = _extract_gpu_metrics(gpu_data)

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

@ -120,11 +120,13 @@ def detect_hardware() -> DeviceType:
# Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes.
# DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP.
if getattr(torch.version, "hip", None) is not None:
# AMD's repo.radeon.com SDK wheels (e.g. 2.9.0+rocmsdk20251116) do
# not set torch.version.hip, so fall back to checking __version__.
_hip_ver = getattr(torch.version, "hip", None)
if _hip_ver is not None or "rocm" in torch.__version__.lower():
IS_ROCM = True
print(
f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}"
)
_hip_label = _hip_ver or torch.__version__
print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}")
else:
print(f"Hardware detected: CUDA -- {device_name}")
return DEVICE
@ -466,7 +468,7 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]:
try:
func = getattr(_backend, func_name)
result = func(*args, **kwargs)
if result.get("available"):
if isinstance(result, dict) and result.get("available"):
return result
except Exception as e:
logger.warning("%s %s query failed: %s", backend_name, func_name, e)
@ -506,6 +508,145 @@ def _read_apple_gpu_stats() -> Dict[str, Any]:
}
def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]:
"""Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent."""
import glob as _glob
if platform.system() != "Linux":
return None
try:
files = _glob.glob("/sys/class/drm/card*/device/gpu_busy_percent")
if not files:
return None
values = [int(open(f).read().strip()) for f in files]
return round(sum(values) / len(values), 1)
except Exception:
return None
def _rocm_linux_sysfs_temp_c() -> Optional[float]:
"""Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C)."""
import glob as _glob
if platform.system() != "Linux":
return None
try:
files = _glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input")
if not files:
return None
temps = [int(open(f).read().strip()) / 1000.0 for f in files]
return round(max(temps), 1)
except Exception:
return None
def _rocm_linux_sysfs_power_w() -> Optional[float]:
"""Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts)."""
import glob as _glob
if platform.system() != "Linux":
return None
try:
for pattern in (
"/sys/class/drm/card*/device/hwmon/hwmon*/power1_average",
"/sys/class/drm/card*/device/hwmon/hwmon*/power1_input",
):
files = _glob.glob(pattern)
if files:
watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files)
return round(watts, 1)
return None
except Exception:
return None
def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]:
"""Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes)."""
import subprocess as _sp
if platform.system() != "Windows":
return None
try:
ps = (
"$s=(Get-Counter '\\GPU Engine(*engtype_3D*)\\Utilization Percentage'"
" -ErrorAction SilentlyContinue).CounterSamples;"
"if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}"
)
r = _sp.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
capture_output = True,
text = True,
timeout = 5,
)
if r.returncode != 0 or not r.stdout.strip():
return None
val = float(r.stdout.strip())
return round(val, 1) if val >= 0 else None
except Exception:
return None
def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
"""Query system-wide AMD GPU VRAM via Linux DRM sysfs.
Reads /sys/class/drm/card*/device/mem_info_vram_* which the kernel
updates in real-time across all processes. No tools required.
Returns (used_gb, total_gb) or (None, None) on failure.
"""
import glob as _glob
if platform.system() != "Linux":
return None, None
try:
used_files = _glob.glob("/sys/class/drm/card*/device/mem_info_vram_used")
total_files = _glob.glob("/sys/class/drm/card*/device/mem_info_vram_total")
if not used_files or not total_files:
return None, None
used_bytes = sum(int(open(f).read().strip()) for f in used_files)
total_bytes = sum(int(open(f).read().strip()) for f in total_files)
if total_bytes == 0:
return None, None
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
except Exception:
return None, None
def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]:
"""Query system-wide dedicated GPU VRAM via Windows Performance Counters.
Uses the same data source as Task Manager so it reflects cross-process
usage accurately. Works for any GPU vendor without amd-smi or nvidia-smi.
Returns (used_gb, total_gb) or (None, None) on failure.
"""
import subprocess as _sp
if platform.system() != "Windows":
return None, None
try:
ps = (
"$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'"
" -ErrorAction SilentlyContinue).CounterSamples;"
"if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}"
)
r = _sp.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
capture_output = True,
text = True,
timeout = 5,
)
if r.returncode != 0 or not r.stdout.strip():
return None, None
used_bytes = float(r.stdout.strip())
if used_bytes < 0:
return None, None
import torch as _torch
total_bytes = _torch.cuda.get_device_properties(0).total_memory
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
except Exception:
return None, None
def get_gpu_utilization() -> Dict[str, Any]:
"""Return a live snapshot of device utilization information."""
device = get_device()
@ -514,7 +655,78 @@ def get_gpu_utilization() -> Dict[str, Any]:
result = _smi_query("get_primary_gpu_utilization")
if result is not None:
result["backend"] = _backend_label(device)
if IS_ROCM:
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
_reconcile_primary_rocm_unified_memory(
result, _get_parent_visible_gpu_spec()
)
return result
# SMI tool unavailable or returned no usable data. On Windows, query
# the Performance Counter API (same source as Task Manager) for
# system-wide dedicated VRAM — covers cross-process usage that
# torch.cuda.mem_get_info cannot see from the Studio server process.
if IS_ROCM and platform.system() == "Windows":
_win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
if _win_used is not None and _win_total is not None:
_win_util = _rocm_windows_perf_counter_gpu_util_pct()
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": _win_util,
"temperature_c": None,
"vram_used_gb": _win_used,
"vram_total_gb": _win_total,
"vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
if _win_total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
# Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed.
if IS_ROCM and platform.system() == "Linux":
_linux_used, _linux_total = _rocm_linux_sysfs_vram_gb()
if _linux_used is not None and _linux_total is not None:
_linux_util = _rocm_linux_sysfs_gpu_busy_pct()
_linux_temp = _rocm_linux_sysfs_temp_c()
_linux_power = _rocm_linux_sysfs_power_w()
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": _linux_util,
"temperature_c": _linux_temp,
"vram_used_gb": _linux_used,
"vram_total_gb": _linux_total,
"vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1)
if _linux_total > 0
else None,
"power_draw_w": _linux_power,
"power_limit_w": None,
"power_utilization_pct": None,
}
# Last resort: torch mem_get_info (process-local).
_visible_spec = _get_parent_visible_gpu_spec()
_numeric_ids = _visible_spec.get("numeric_ids") or [0]
_primary_idx = [_numeric_ids[0]] if _numeric_ids else [0]
_torch_devices = _torch_get_per_device_info(_primary_idx)
if _torch_devices:
_td = _torch_devices[0]
_total = _td["total_gb"]
_used = _td["used_gb"]
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": _used,
"vram_total_gb": _total,
"vram_utilization_pct": round((_used / _total) * 100, 1)
if _total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
# MLX path: single _read_apple_gpu_stats() call carries both VRAM-used
# bytes and GPU utilization %. psutil for unified-memory total is cheap.
@ -578,6 +790,77 @@ def get_gpu_utilization() -> Dict[str, Any]:
return {"available": False, "backend": _backend_label(device)}
def _apply_unified_memory_correction(
device_metrics: Dict[str, Any], torch_info: Dict[str, Any]
) -> None:
"""Per-device reconciliation: when torch reports a larger memory total
than amd-smi, overwrite the smi VRAM fields in place.
Used by both the multi-device and primary-device reconciliation helpers
so the two endpoints stay in sync on AMD iGPUs with unified memory.
"""
torch_total_gb = torch_info["total_gb"]
smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
if torch_total_gb > smi_total_gb:
torch_used_gb = torch_info["used_gb"]
device_metrics["vram_total_gb"] = torch_total_gb
device_metrics["vram_used_gb"] = torch_used_gb
device_metrics["vram_utilization_pct"] = (
round((torch_used_gb / torch_total_gb) * 100, 1)
if torch_total_gb > 0
else None
)
logger.debug(
"ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with "
"torch mem_get_info total (%.2f GB) for device %s",
smi_total_gb,
torch_total_gb,
torch_info.get("index"),
)
def _reconcile_rocm_unified_memory(
utilization: Dict[str, Any], device_indices: list[int]
) -> None:
"""Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo).
amd-smi reports only the dedicated slice (~512 MB); torch sees the full
GTT pool (~128 GB). When torch total > smi total, overwrite per-device
VRAM fields so GPU selection uses the real available memory.
"""
torch_devices = _torch_get_per_device_info(device_indices)
if not torch_devices:
return
torch_by_index = {td["index"]: td for td in torch_devices}
for dev in utilization.get("devices", []):
td = torch_by_index.get(dev.get("index"))
if td is None:
continue
_apply_unified_memory_correction(dev, td)
def _reconcile_primary_rocm_unified_memory(
utilization: Dict[str, Any], parent_visible_spec: Dict[str, Any]
) -> None:
"""Same fix as _reconcile_rocm_unified_memory for the flat primary-GPU dict."""
numeric_ids = parent_visible_spec.get("numeric_ids")
if numeric_ids is None:
# No visibility env var set: torch ordinal 0 is the primary device.
primary_idx = [0]
elif len(numeric_ids) == 0:
# Empty mask (HIP_VISIBLE_DEVICES="" or "-1"): no GPU is visible to
# this process. Querying torch device 0 would raise a RuntimeError or
# return stale/wrong data, so bail out rather than writing bad values
# into the utilization dict.
return
else:
primary_idx = [int(numeric_ids[0])]
torch_devices = _torch_get_per_device_info(primary_idx)
if not torch_devices:
return
_apply_unified_memory_correction(utilization, torch_devices[0])
def get_visible_gpu_utilization() -> Dict[str, Any]:
device = get_device()
@ -590,6 +873,10 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
)
if result is not None:
result["backend"] = _backend_label(device)
numeric_ids = parent_visible_spec.get("numeric_ids")
if IS_ROCM and numeric_ids is not None:
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
_reconcile_rocm_unified_memory(result, numeric_ids)
return result
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
@ -689,7 +976,15 @@ def _get_parent_visible_gpu_spec() -> Dict[str, Any]:
# Use explicit None checks (not `or`) so empty string "" is honoured
# as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES.
cuda_visible = None
if IS_ROCM:
# Prefer ROCm masks only on a ROCm host, or when no CUDA mask is set, so a
# stale HIP_VISIBLE_DEVICES on an NVIDIA host can't override CUDA_VISIBLE_DEVICES.
_is_rocm_spec = IS_ROCM or (
"CUDA_VISIBLE_DEVICES" not in os.environ
and (
"HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
)
)
if _is_rocm_spec:
hip_vis = os.environ.get("HIP_VISIBLE_DEVICES")
rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES")
if hip_vis is not None:
@ -867,6 +1162,61 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
def _determine_attention_impl_for_gpu_estimate(config) -> str:
import copy as _copy
# torch.distributed is incomplete on Windows ROCm — torch._C is a C
# extension (not a package), so Python cannot import the submodule
# torch._C._distributed_c10d that torch.distributed depends on.
# Inject an empty stub into sys.modules BEFORE importing torch.distributed
# so the import succeeds, then patch the missing process-group helpers.
import sys as _sys
import types as _types
if _sys.platform == "win32" and IS_ROCM:
# Dummy class for any name torch.distributed tries to import from these stubs
class _Dummy:
pass
for _c10d_name in (
"torch._C._distributed_c10d",
"torch._C._distributed_autograd",
"torch._C._distributed_rpc",
):
if _c10d_name not in _sys.modules:
_stub = _types.ModuleType(_c10d_name)
# torch.distributed imports these names from _distributed_c10d;
# provide no-op dummies so the import doesn't raise AttributeError.
for _sym in (
"FakeProcessGroup",
"ProcessGroup",
"Work",
"Store",
"PrefixStore",
"FileStore",
"TCPStore",
"HashStore",
"Reducer",
"Logger",
"DistributedDebugLevel",
"GradBucket",
"BuiltinCommHookType",
):
setattr(_stub, _sym, _Dummy)
_sys.modules[_c10d_name] = _stub
try:
import torch.distributed as _td
for _attr, _stub in (
("is_initialized", lambda: False),
("is_available", lambda: False),
("get_rank", lambda: 0),
("get_world_size", lambda: 1),
("is_torchelastic_launched", lambda: False),
):
if not hasattr(_td, _attr):
setattr(_td, _attr, _stub)
except ImportError:
pass
from unsloth.models._utils import resolve_attention_implementation
from transformers import AutoModel, AutoModelForCausalLM
@ -1062,7 +1412,10 @@ def estimate_required_model_memory_gb(
_determine_attention_impl_for_gpu_estimate(config)
)
except Exception as e:
logger.warning(
# Log at debug: on Windows ROCm the torch.distributed stub does
# not implement Store, so this fires on every estimate call.
# It is expected and non-actionable -- eager is the safe fallback.
logger.debug(
"Could not resolve attention implementation for '%s': %s",
estimate_model,
e,
@ -1552,14 +1905,35 @@ def apply_gpu_ids(gpu_ids) -> None:
# parent process already set a ROCm visibility variable -- that
# way a downstream ROCm process inherits the narrowed mask even
# before Studio's hardware detection has classified the host.
# Final fallback: probe torch.version.hip so AMD workers without
# HIP_VISIBLE_DEVICES still get the correct ROCm visibility mask.
_inherits_rocm_visibility = (
"HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
)
if IS_ROCM or _inherits_rocm_visibility:
_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
or "rocm" in getattr(_torch, "__version__", "").lower()
)
except Exception as e:
logger.debug(
"apply_gpu_ids: torch ROCm probe skipped (%s: %s)",
type(e).__name__,
e,
)
if _is_rocm:
os.environ["HIP_VISIBLE_DEVICES"] = value
os.environ["ROCR_VISIBLE_DEVICES"] = value
_visible_gpu_count = None
if IS_ROCM or _inherits_rocm_visibility:
if _is_rocm:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value)
else:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value)

2
studio/backend/utils/wheel_utils.py
View file

@ -15,6 +15,7 @@ import urllib.request
from typing import Callable
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
_logger = logging.getLogger(__name__)
@ -106,6 +107,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non
text = True,
timeout = timeout,
env = child_env_without_native_path_secret(),
**windows_hidden_subprocess_kwargs(),
)
except subprocess.TimeoutExpired:
return None

424
studio/install_llama_prebuilt.py
View file

@ -9,6 +9,7 @@ from __future__ import annotations
import argparse
import errno
import fnmatch
import functools
import hashlib
import json
import os
@ -29,7 +30,7 @@ import urllib.parse
import urllib.request
import zipfile
from contextlib import contextmanager
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace as dataclasses_replace
try:
from filelock import FileLock, Timeout as FileLockTimeout
@ -100,6 +101,39 @@ DEFAULT_PUBLISHED_SHA256_ASSET = os.environ.get(
)
UPSTREAM_REPO = "ggml-org/llama.cpp"
UPSTREAM_RELEASES_API = f"https://api.github.com/repos/{UPSTREAM_REPO}/releases/latest"
LEMONADE_ROCM_REPO = "lemonade-sdk/llamacpp-rocm"
LEMONADE_ROCM_RELEASES_API = (
f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/latest"
)
def _lemonade_release_api_for(llama_tag: str) -> str:
"""Return the GitHub API URL for the lemonade release that matches a
requested llama.cpp tag.
When llama_tag is unset or "latest", point at /releases/latest. When the
caller has pinned a specific tag (e.g. "b1260"), point at the same tag in
lemonade. Lemonade tracks `ggml-org/llama.cpp` build tags (e.g. "b1260")
but is NOT guaranteed to publish every upstream build -- lemonade may be
several builds behind ggml-org. Pinning to a specific tag that lemonade
skipped will produce a 404 and the caller falls through to the upstream
tarball; that is intentional so pinned installs stay reproducible.
Do NOT pass a `unslothai/llama.cpp` fork tag -- the fork uses its own
namespace and will always 404 against lemonade.
The tag is URL-encoded with `safe=""` so an unexpected slash / hash / query
character cannot reshape the URL.
"""
normalized = (llama_tag or "").strip()
if not normalized or normalized.lower() == "latest":
return LEMONADE_ROCM_RELEASES_API
return (
f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/tags/"
f"{urllib.parse.quote(normalized, safe = '')}"
)
TEST_MODEL_URL = (
"https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
)
@ -196,6 +230,7 @@ class HostInfo:
has_physical_nvidia: bool
has_usable_nvidia: bool
has_rocm: bool = False
rocm_gfx_target: str | None = None
@dataclass
@ -1268,9 +1303,26 @@ def direct_linux_release_plan(
selection = linux_cuda_choice_from_release(host, bundle)
if selection is not None:
attempts.extend(selection.attempts)
cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu")
if cpu_choice is not None:
attempts.append(cpu_choice)
if host.has_rocm and not host.has_usable_nvidia:
# Per-GPU lemonade prebuilts ship the ROCm runtime libs alongside
# llama.cpp, so they install cleanly even on hosts (e.g. gfx1151
# Strix Halo) that the upstream combined-ROCm tarball doesn't cover.
# The "ubuntu" label is lemonade's asset naming convention only --
# the binary is a manylinux-style glibc build that runs on Arch,
# Fedora, openSUSE, etc. as long as the host glibc is recent enough.
# Do NOT append the CPU asset for ROCm-only hosts: if lemonade fails
# validation we want validate_prebuilt_attempts to raise PrebuiltFallback
# so the caller triggers the HIP source build, not silently install a
# CPU-only binary.
lemonade_choice = resolve_lemonade_rocm_choice(
host, "ubuntu", "linux-rocm", llama_tag = requested_tag
)
if lemonade_choice is not None:
attempts.append(lemonade_choice)
else:
cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu")
if cpu_choice is not None:
attempts.append(cpu_choice)
if not attempts:
raise PrebuiltFallback("no compatible Linux prebuilt asset was found")
approved_checksums = synthetic_checksums_for_release(
@ -1336,6 +1388,25 @@ def direct_upstream_release_plan(
torch_preference.selection_log,
)
)
elif host.has_rocm:
lemonade_choice = resolve_lemonade_rocm_choice(
host, "windows", "windows-hip", llama_tag = requested_tag
)
if lemonade_choice is not None:
attempts.append(lemonade_choice)
hip_asset = f"llama-{release_tag}-bin-win-hip-radeon-x64.zip"
hip_url = assets.get(hip_asset)
if hip_url:
attempts.append(
AssetChoice(
repo = repo,
tag = release_tag,
name = hip_asset,
url = hip_url,
source_label = "upstream",
install_kind = "windows-hip",
)
)
cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip"
cpu_url = assets.get(cpu_asset)
if cpu_url:
@ -2604,6 +2675,72 @@ def run_capture(
return result
def _pick_rocm_gfx_target(out: str) -> str | None:
"""Choose the gfx target rocminfo / hipinfo report for the active GPU.
A bare first-match picked the wrong device on mixed APU + dGPU hosts
(e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100). Respect
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES so the
asset matches what HIP actually runs on. Falls back to the first GPU when
no env var is set.
rocminfo / hipinfo print the same gfx token multiple times per GPU (Name,
ISA, marketing-name). We first try to split the output on per-GPU section
headers (rocminfo: "Agent N" blocks, hipinfo: "device#N" entries) and take
exactly one gfx token per section. This gives the correct per-GPU list even
on same-arch multi-GPU hosts (e.g. two RX 7900 XTX cards) where global
dict.fromkeys dedup would collapse both cards to a single entry and make
HIP_VISIBLE_DEVICES=1 point out of range.
Falls back to insertion-order dedup when the output has no recognisable
section markers (flat gfx-string inputs, unit-test stubs, etc.).
Empty / "-1" env values mean no AMD GPU is visible to HIP: return None.
"""
# Try to build a per-GPU token list by splitting on section boundaries.
# rocminfo sections are introduced by "Agent N" lines (optionally between
# rows of asterisks). hipinfo sections start with "device#N".
_sections = re.split(
r"(?mi)^\s*\*+\s*$\s*agent\s+\d+\s*$|\bdevice\s*#\s*\d+\b",
out,
)
if len(_sections) > 1:
# Section-based: one gfx token per GPU section preserves physical order.
_tokens: list[str] = []
for _sec in _sections[1:]:
_m = re.search(r"gfx[1-9][0-9a-z]{2,3}", _sec.lower())
if _m:
_tokens.append(_m.group(0))
else:
# Fallback: insertion-order dedup (handles flat strings / unknown formats).
_raw = re.findall(r"gfx[1-9][0-9a-z]{2,3}", out.lower())
_tokens = list(dict.fromkeys(_raw))
if not _tokens:
return None
_vis_raw = None
# AMD's HIP runtime honours all three env vars with identical semantics.
for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
_val = os.environ.get(_env)
if _val is not None:
_vis_raw = _val
break
if _vis_raw is not None:
_vis = _vis_raw.strip()
# Empty or "-1" means "no AMD GPU visible" (matches the rest of Studio).
if _vis == "" or _vis == "-1":
return None
_first = _vis.split(",")[0].strip()
try:
_idx = int(_first)
if 0 <= _idx < len(_tokens):
return _tokens[_idx]
except ValueError:
pass
return _tokens[0]
def detect_host() -> HostInfo:
system = platform.system()
machine = platform.machine().lower()
@ -2707,6 +2844,7 @@ def detect_host() -> HostInfo:
return bool(re.search(r"(?im)^gpu\s*[:\[]\s*\d", stdout))
has_rocm = False
rocm_gfx_target: str | None = None
if is_linux:
for _cmd, _check in (
# rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit).
@ -2729,14 +2867,33 @@ def detect_host() -> HostInfo:
if _result.returncode == 0 and _result.stdout.strip():
if _check(_result.stdout):
has_rocm = True
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
elif is_windows:
# Windows: prefer active probes that validate GPU presence
# Windows: prefer active probes that validate GPU presence.
# hipinfo / amd-smi are often NOT on PATH -- the HIP SDK installer
# sets HIP_PATH / ROCM_PATH but does not always add the bin dir to
# the system PATH. Mirror setup.ps1's fallback: check the env-var
# bin dirs before giving up so that `has_rocm` is not silently False
# on machines where the PATH is not yet updated.
def _resolve_exe(name: str) -> str | None:
"""Return full path to `name`, checking PATH then HIP_PATH/ROCM_PATH bin."""
found = shutil.which(name)
if found:
return found
for _env in ("HIP_PATH", "ROCM_PATH"):
_root = os.environ.get(_env)
if _root:
_candidate = os.path.join(_root, "bin", f"{name}.exe")
if os.path.isfile(_candidate):
return _candidate
return None
for _cmd, _check in (
(["hipinfo"], lambda out: "gcnarchname" in out.lower()),
(["amd-smi", "list"], _amd_smi_has_gpu),
):
_exe = shutil.which(_cmd[0])
_exe = _resolve_exe(_cmd[0])
if not _exe:
continue
try:
@ -2746,6 +2903,8 @@ def detect_host() -> HostInfo:
if _result.returncode == 0 and _result.stdout.strip():
if _check(_result.stdout):
has_rocm = True
# hipinfo reports "gcnArchName: gfx1100" -- extract if present
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
# Note: amdhip64.dll presence alone is NOT treated as GPU evidence
# since the HIP SDK can be installed without an AMD GPU.
@ -2765,6 +2924,7 @@ def detect_host() -> HostInfo:
has_physical_nvidia = has_physical_nvidia,
has_usable_nvidia = has_usable_nvidia,
has_rocm = has_rocm,
rocm_gfx_target = rocm_gfx_target,
)
@ -3263,6 +3423,185 @@ def _detect_host_rocm_version() -> tuple[int, int] | None:
return None
# Map detected gfx IDs to lemonade-sdk asset family suffixes.
# More-specific prefixes must come before shorter ones (e.g. gfx1151 before gfx110).
_LEMONADE_GFX_FAMILIES: list[tuple[str, str]] = [
("gfx1151", "gfx1151"),
("gfx1150", "gfx1150"),
("gfx120", "gfx120X"),
("gfx110", "gfx110X"),
("gfx103", "gfx103X"),
]
def _lemonade_gfx_family(gfx_id: str) -> str | None:
gfx_id = gfx_id.lower().strip()
for prefix, family in _LEMONADE_GFX_FAMILIES:
if gfx_id.startswith(prefix):
return family
return None
def _is_trusted_github_release_url(url: str, expected_repo: str) -> bool:
"""Validate a release asset URL points at GitHub's expected hosts.
Accepts:
https://github.com/{expected_repo}/releases/download/...
https://objects.githubusercontent.com/... (GitHub's release CDN)
Anything else (including http://, raw.githubusercontent.com, gist, etc.)
is rejected so a malicious API response cannot redirect downloads to an
attacker-chosen host.
"""
if not isinstance(url, str) or not url:
return False
try:
parsed = urllib.parse.urlparse(url)
except Exception:
return False
if parsed.scheme != "https":
return False
host = (parsed.netloc or "").lower()
if host == "objects.githubusercontent.com":
# GitHub's release CDN. Restrict to release-asset paths so a tampered
# API response pointing at an arbitrary CDN object is still rejected.
# Real release asset URLs carry the "/github-production-release-asset-"
# prefix; gist / raw / avatar CDN paths do not.
return parsed.path.startswith("/github-production-release-asset-")
if host == "github.com":
return parsed.path.startswith(f"/{expected_repo}/releases/download/")
return False
@functools.lru_cache(maxsize = 8)
def _fetch_lemonade_release_cached(api_url: str, llama_tag: str) -> "dict | None":
"""Cached wrapper around fetch_json for lemonade release lookups.
resolve_lemonade_rocm_choice() is called twice per install (once from the
direct planner, once from resolve_upstream_asset_choice) with identical
arguments. Without memoisation, each install hits api.github.com twice,
doubling the rate-limit failure surface on busy CI runners. Cache is
process-scoped; tests that need to vary fetch_json's return value across
invocations should call cache_clear().
"""
try:
return fetch_json(api_url)
except Exception as exc:
normalized = (llama_tag or "").strip().lower()
if normalized and normalized != "latest":
log(
f"Could not fetch {LEMONADE_ROCM_REPO} release for "
f"llama_tag={llama_tag!r} ({exc}); skipping lemonade prebuilt"
)
else:
log(f"Could not fetch {LEMONADE_ROCM_REPO} latest release: {exc}")
return None
def resolve_lemonade_rocm_choice(
host: HostInfo,
os_prefix: str,
install_kind: str,
llama_tag: str = "latest",
) -> "AssetChoice | None":
"""Return an AssetChoice from lemonade-sdk/llamacpp-rocm for the detected GPU, or None.
os_prefix: lemonade's asset filename label, NOT a host-distro filter.
Pass "ubuntu" for any Linux host (Arch, Fedora, openSUSE,
Debian, ...) -- lemonade only publishes one Linux variant
and it is a manylinux-style glibc build that runs on any
distro with a recent-enough glibc. Pass "windows" for
Windows hosts.
install_kind: "linux-rocm" or "windows-hip"
llama_tag: the requested upstream llama.cpp tag ("latest" or a pinned
release like "b1260"). When pinned, the resolver fetches
the matching lemonade release. When the pinned tag is not
published by lemonade we skip silently (and the caller
falls through to upstream) rather than drift to whatever
lemonade ships as latest.
"""
if not host.rocm_gfx_target:
return None
# Opt-out for users who want the upstream HIP build path only -- lemonade
# binaries are downloaded without entries in the approved-hash manifest, so
# the integrity gate is functional validation only.
if os.environ.get("UNSLOTH_DISABLE_LEMONADE_ROCM", "").strip().lower() in (
"1",
"true",
"yes",
):
log("UNSLOTH_DISABLE_LEMONADE_ROCM is set; skipping lemonade-sdk prebuilt")
return None
gfx_family = _lemonade_gfx_family(host.rocm_gfx_target)
if gfx_family is None:
log(
f"AMD GPU {host.rocm_gfx_target!r} is not covered by lemonade-sdk ROCm prebuilts; "
"skipping lemonade prebuilt"
)
return None
api_url = _lemonade_release_api_for(llama_tag)
release = _fetch_lemonade_release_cached(api_url, llama_tag)
if release is None:
return None
release_tag = release.get("tag_name") if isinstance(release, dict) else None
if not isinstance(release_tag, str) or not release_tag:
log(
f"Unexpected {LEMONADE_ROCM_REPO} release payload; skipping lemonade prebuilt"
)
return None
assets = release_asset_map(release)
asset_name = f"llama-{release_tag}-{os_prefix}-rocm-{gfx_family}-x64.zip"
if asset_name not in assets:
log(
f"{LEMONADE_ROCM_REPO}@{release_tag} has no asset {asset_name!r}; "
"skipping lemonade prebuilt"
)
return None
asset_url = assets[asset_name]
if not asset_url:
# release_asset_map defaults to "" when an asset row is missing
# browser_download_url; skip cleanly instead of letting
# download_file("") raise a less obvious error downstream.
log(
f"{LEMONADE_ROCM_REPO}@{release_tag} asset {asset_name!r} has no "
"browser_download_url; skipping lemonade prebuilt"
)
return None
# Defence in depth: lemonade browser_download_url should be on github.com
# or githubusercontent.com. A compromised GitHub API response that
# redirects to an attacker-chosen host would otherwise be honoured
# silently (lemonade assets are not in the approved-hash manifest).
if not _is_trusted_github_release_url(asset_url, LEMONADE_ROCM_REPO):
log(
f"{LEMONADE_ROCM_REPO}@{release_tag} asset {asset_name!r} points "
f"to an unexpected host ({asset_url!r}); refusing to download "
"lemonade prebuilt"
)
return None
# Note: lemonade tags Linux assets with "ubuntu" but the binary is a
# generic glibc build that runs on any distro (Arch, Fedora, ...), so
# this attempt is selected for all Linux ROCm hosts, not just Ubuntu.
log(
f"AMD GPU {host.rocm_gfx_target!r} ({gfx_family}) -- "
f"trying lemonade-sdk ROCm prebuilt {asset_name} "
f"(works on any glibc Linux, not just Ubuntu)"
)
log(
f"NOTE: lemonade-sdk/llamacpp-rocm releases are not covered by the "
f"Unsloth approved-hash manifest; download integrity relies on "
f"functional validation (llama-bench / llama-server smoke tests) "
f"after extraction. Set UNSLOTH_DISABLE_LEMONADE_ROCM=1 to skip "
f"lemonade and fall back to the upstream HIP build path."
)
return AssetChoice(
repo = LEMONADE_ROCM_REPO,
tag = release_tag,
name = asset_name,
url = asset_url,
source_label = "lemonade",
install_kind = install_kind,
)
def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice:
upstream_assets = github_release_assets(UPSTREAM_REPO, llama_tag)
if host.is_linux and host.is_x86_64:
@ -3271,6 +3610,15 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
# the exact GPU target via rocminfo, which is more reliable for consumer
# GPUs (e.g. gfx1151) that may not be in the prebuilt.
if host.has_rocm and not host.has_usable_nvidia:
# Try lemonade-sdk per-GPU prebuilt first: these are built against
# specific gfx targets and bundle all required ROCm runtime libs.
lemonade_choice = resolve_lemonade_rocm_choice(
host, "ubuntu", "linux-rocm", llama_tag = llama_tag
)
if lemonade_choice is not None:
return lemonade_choice
# Fall back to upstream combined ROCm tarball.
# Scan upstream assets for any rocm-<version> prebuilt. When the
# host ROCm runtime version is known, pick the newest candidate
# whose major.minor is <= host version -- otherwise a ROCm 6.4
@ -3350,8 +3698,14 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
return attempts[0]
raise PrebuiltFallback("no compatible Windows CUDA asset was found")
# AMD ROCm on Windows: try HIP prebuilt
# AMD ROCm on Windows: try lemonade per-GPU prebuilt first, then upstream HIP
if host.has_rocm:
lemonade_choice = resolve_lemonade_rocm_choice(
host, "windows", "windows-hip", llama_tag = llama_tag
)
if lemonade_choice is not None:
return lemonade_choice
hip_name = f"llama-{llama_tag}-bin-win-hip-radeon-x64.zip"
if hip_name in upstream_assets:
log(
@ -3892,6 +4246,22 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
)
def runtime_subdirs_for_choice(choice: AssetChoice) -> list[str]:
"""Subdirectory names within the archive root that must be copied into
the overlay directory alongside the flat shared libraries.
hipBLASLt and rocBLAS expect their Tensile kernel catalog trees
(hipblaslt/library/<gfx>/ and rocblas/library/<gfx>/) to sit next to
their shared libraries at runtime. These trees are multi-level and
cannot be handled by copy_globs (filename-only matching, flat copy)."""
if choice.source_label == "lemonade" and choice.install_kind in {
"linux-rocm",
"windows-hip",
}:
return ["hipblaslt", "rocblas"]
return []
def metadata_patterns_for_choice(choice: AssetChoice) -> list[str]:
patterns = ["BUILD_INFO.txt", "THIRD_PARTY_LICENSES.txt"]
if choice.install_kind.startswith("windows"):
@ -4216,6 +4586,10 @@ def install_from_archives(
copy_globs(
source_dir, overlay_dir, runtime_patterns_for_choice(choice), required = True
)
for _subdir in runtime_subdirs_for_choice(choice):
_src_subdir = source_dir / _subdir
if _src_subdir.is_dir():
shutil.copytree(_src_subdir, overlay_dir / _subdir, dirs_exist_ok = True)
if runtime_extract_dir is not None:
# The runtime archive only contributes the CUDA DLLs.
# Restrict the overlay to the cudart bundle's known
@ -4964,6 +5338,15 @@ def apply_approved_hashes(
approved_attempts: list[AssetChoice] = []
missing_assets: list[str] = []
for attempt in attempts:
# External prebuilts (e.g. lemonade-sdk) are not listed in the
# approved-hash manifest; they are explicitly documented as relying
# on functional validation only (llama-bench / smoke tests).
# Passing them through here lets the caller include both a lemonade
# attempt and a hash-approved upstream fallback in the same list
# without apply_approved_hashes discarding the lemonade entry.
if attempt.source_label == "lemonade":
approved_attempts.append(attempt)
continue
approved = approved_hash_for_attempt(attempt)
if approved is None:
missing_assets.append(attempt.name)
@ -5164,6 +5547,12 @@ def write_prebuilt_metadata(
"asset": choice.name,
"asset_sha256": choice.expected_sha256,
"source": choice.source_label,
# Binary-side repo/tag for non-upstream sources (e.g. lemonade).
# published_repo/release_tag always refer to the unsloth source tree;
# these capture where the actual binaries came from so the install
# summary can show both (e.g. "unslothai/llama.cpp@b9334 + lemonade@b1280").
"binary_repo": choice.repo,
"binary_release_tag": choice.tag,
"source_asset": source_asset_name,
"source_sha256": source_sha256,
"source_commit": approved_checksums.source_commit,
@ -5242,7 +5631,7 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
["libllama.so*"],
["libggml.so*"],
["libggml-base.so*"],
["libggml-cpu-*.so*"],
["libggml-cpu*.so*"],
["libmtmd.so*"],
]
if choice.install_kind == "linux-cuda":
@ -5251,7 +5640,7 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
["libllama.so*"],
["libggml.so*"],
["libggml-base.so*"],
["libggml-cpu-*.so*"],
["libggml-cpu*.so*"],
["libmtmd.so*"],
["libggml-cuda.so*"],
]
@ -5267,7 +5656,7 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
["libllama.so*"],
["libggml.so*"],
["libggml-base.so*"],
["libggml-cpu-*.so*"],
["libggml-cpu*.so*"],
["libmtmd.so*"],
["libggml-hip.so*"],
]
@ -5565,8 +5954,11 @@ def install_prebuilt(
published_release_tag: str,
*,
simple_policy: bool = False,
override_has_rocm: bool = False,
) -> None:
host = detect_host()
if override_has_rocm and not host.has_rocm:
host = dataclasses_replace(host, has_rocm = True)
choice: AssetChoice | None = None
try:
with install_lock(install_lock_path(install_dir)):
@ -5700,6 +6092,17 @@ def parse_args() -> argparse.Namespace:
action = "store_true",
help = "Use the simplified platform-specific prebuilt selection policy.",
)
parser.add_argument(
"--has-rocm",
action = "store_true",
default = False,
help = (
"Assert that an AMD ROCm GPU is present. When set, skips the internal "
"hipinfo/amd-smi probe and forces has_rocm=True in the host profile. "
"Used by setup.ps1/setup.sh to forward their own ROCm detection result "
"so the HIP llama.cpp prebuilt is selected even when hipinfo is not on PATH."
),
)
resolve_group = parser.add_mutually_exclusive_group()
resolve_group.add_argument(
"--resolve-llama-tag",
@ -5820,6 +6223,7 @@ def main() -> int:
published_repo = args.published_repo,
published_release_tag = args.published_release_tag or "",
simple_policy = args.simple_policy,
override_has_rocm = args.has_rocm,
)
return EXIT_SUCCESS

553
studio/install_python_stack.py
View file

@ -54,12 +54,9 @@ PLATFORM_LACKS_TORCHCODEC_WHEEL = (
# ── ROCm / AMD GPU support ─────────────────────────────────────────────────────
# Mapping from detected ROCm (major, minor) to the best PyTorch wheel tag on
# download.pytorch.org. Entries are checked newest-first (>=).
# ROCm 7.2 only has torch 2.11.0 on download.pytorch.org, which exceeds the
# current torch upper bound (<2.11.0). Fall back to rocm7.1 (torch 2.10.0).
# TODO: uncomment rocm7.2 when torch upper bound is bumped to >=2.11.0
_ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
# (7, 2): "rocm7.2", # torch 2.11.0 -- requires torch>=2.11
(7, 1): "rocm7.1",
(7, 2): "rocm7.2", # torch 2.11.0
(7, 1): "rocm7.1", # torch 2.10.0
(7, 0): "rocm7.0",
(6, 4): "rocm6.4",
(6, 3): "rocm6.3",
@ -67,10 +64,47 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
(6, 1): "rocm6.1",
(6, 0): "rocm6.0",
}
# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x).
_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"rocm7.2": (
"torch>=2.11.0,<2.12.0",
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
),
# Default for rocm7.1 and earlier: torch 2.x below 2.11
"_default": (
"torch>=2.4,<2.11.0",
"torchvision>=0.19,<0.26.0",
"torchaudio>=2.4,<2.11.0",
),
}
_PYTORCH_WHL_BASE = (
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
).rstrip("/")
# AMD Windows ROCm wheels — repo.amd.com (arch-specific pip index)
# Format: https://repo.amd.com/rocm/whl/{arch_family}/
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
_ROCM_WINDOWS_INDEX_BASE = (
os.environ.get("UNSLOTH_ROCM_WINDOWS_MIRROR") or "https://repo.amd.com/rocm/whl"
).rstrip("/")
# Maps gfx arch → AMD index arch-family suffix.
# Each family is a separate pip index on repo.amd.com.
_GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
"gfx1201": "gfx120X-all",
"gfx1200": "gfx120X-all", # RDNA 4
"gfx1151": "gfx1151",
"gfx1150": "gfx1150", # RDNA 3.5 (Strix Halo/Point)
"gfx1103": "gfx110X-all",
"gfx1102": "gfx110X-all", # RDNA 3
"gfx1101": "gfx110X-all",
"gfx1100": "gfx110X-all",
"gfx90a": "gfx90a",
"gfx908": "gfx908", # MI200/MI100
}
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix
# (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every
# AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI.
@ -85,6 +119,16 @@ _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_rocm{VER}.dll.
# BNB auto-detects HIP version from torch.version.hip, which does not always
# match the DLL suffix in this prerelease wheel (e.g. torch 7.13 with a rocm72
# DLL). We scan the installed wheel for the actual DLL name and set
# BNB_ROCM_VERSION accordingly in _install_bnb_windows_rocm() and worker.py.
"win_amd64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl"
),
}
_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1"
@ -196,6 +240,151 @@ def _detect_rocm_version() -> tuple[int, int] | None:
return None
def _pick_visible_index(num_tokens: int) -> int:
"""Resolve HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES to an integer
index into a list of length num_tokens. Returns 0 (first GPU) for
unset, empty, '-1', UUID-style, or out-of-range values."""
for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"):
_val = os.environ.get(_env)
if _val is None:
continue
_val = _val.strip()
if _val == "" or _val == "-1":
return 0
_first = _val.split(",")[0].strip()
try:
_idx = int(_first)
if 0 <= _idx < num_tokens:
return _idx
except ValueError:
pass
return 0
return 0
def _detect_windows_gfx_arch() -> str | None:
"""Return the gcnArchName on Windows (e.g. 'gfx1200'), or None.
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.
On multi-GPU hosts, all detected gfx tokens are deduplicated (preserving
enumeration order) and HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES selects
which one to install for. The first GPU is used when no env var is set.
"""
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()
def _dedup_pick(tokens: list[str]) -> "str | None":
if not tokens:
return None
# Index into the full (ordered) list first so HIP_VISIBLE_DEVICES
# correctly addresses GPU N on mixed-arch hosts, then return that arch.
return tokens[_pick_visible_index(len(tokens))]
# 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin.
hipinfo = shutil.which("hipinfo")
if not hipinfo:
for _env_var in ("HIP_PATH", "ROCM_PATH"):
_root = os.environ.get(_env_var)
if _root:
_candidate = os.path.join(_root, "bin", "hipinfo.exe")
if os.path.isfile(_candidate):
hipinfo = _candidate
break
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")
# findall picks every gcnArchName line so multi-GPU hosts
# are enumerable and HIP_VISIBLE_DEVICES selects correctly.
_tokens = [
t.strip().lower()
for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
]
_pick = _dedup_pick(_tokens)
if _pick:
return _pick
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")
# Prefer labelled gfx lines; fall back to bare tokens.
_labelled = re.findall(
r"(?im)^\s*(?:target_graphics_version|gfx|arch|asic)\b[^:\r\n]*:\s*(gfx[1-9][0-9a-z]{2,3})\b",
text,
)
_tokens = [t.lower() for t in _labelled]
if not _tokens:
_tokens = re.findall(r"\bgfx[1-9][0-9a-z]{2,3}\b", text.lower())
_pick = _dedup_pick(_tokens)
if _pick:
return _pick
except Exception:
continue
return None
def _windows_rocm_index_url(gfx_arch: str | None) -> str | None:
"""Return the AMD pip index URL for the given GPU arch, or None if unsupported."""
arch_family = _GFX_TO_AMD_INDEX_ARCH.get(gfx_arch or "")
if arch_family is None:
return None
return f"{_ROCM_WINDOWS_INDEX_BASE}/{arch_family}/"
def _detect_bnb_rocm_dll_ver() -> str | None:
"""Scan the installed bitsandbytes package for libbitsandbytes_rocm{VER}.dll.
Returns the version suffix string (e.g. ``"72"``, ``"713"``) or ``None``
if bitsandbytes is not installed or no ROCm DLL is found. Does NOT import
bitsandbytes — uses importlib.util.find_spec so it is safe to call before
BNB is imported.
"""
import glob
import importlib.util
import re
spec = importlib.util.find_spec("bitsandbytes")
if spec is None or not spec.submodule_search_locations:
return None
all_vers: list[str] = []
for pkg_dir in spec.submodule_search_locations:
for dll in glob.glob(os.path.join(pkg_dir, "libbitsandbytes_rocm*.dll")):
m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll))
if m:
all_vers.append(m.group(1))
# Pick the highest numeric suffix so that e.g. "713" wins over "72" when
# both variants are present in the wheel. Filesystem glob order is not
# guaranteed, so always sort rather than stopping at the first match.
return max(all_vers, key = lambda v: int(v)) if all_vers else None
def _has_rocm_gpu() -> bool:
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed)."""
import re
@ -231,6 +420,26 @@ def _has_rocm_gpu() -> bool:
if result.returncode == 0 and result.stdout.strip():
if check_fn(result.stdout):
return True
# sysfs KFD topology fallback (Linux only) -- matches install.sh's
# runtime-only detection. On minimal package-managed installs (no
# rocminfo / no amd-smi GUI tools), the kernel exposes AMD GPUs via
# /sys/class/kfd so `studio update` can still detect the GPU and
# repair the venv.
if sys.platform != "win32":
try:
kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
if os.path.isdir(kfd_nodes):
for entry in os.listdir(kfd_nodes):
gpu_id_path = os.path.join(kfd_nodes, entry, "gpu_id")
try:
with open(gpu_id_path) as fh:
gpu_id = fh.read().strip()
except OSError:
continue
if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node
return True
except OSError:
pass
return False
@ -252,23 +461,199 @@ def _has_usable_nvidia_gpu() -> bool:
return result.returncode == 0 and "GPU " in result.stdout
def _detect_amd_gfx_codes() -> list[str]:
"""Return the list of AMD gfx ISA strings visible to ROCm (e.g. ['gfx1151']).
Probes rocminfo first, then falls back to ``amd-smi list`` and
``amd-smi static --asic`` for runtime-only Radeon hosts that ship
amd-smi but no rocminfo. Returns an empty list when no probe yields
a gfx target.
"""
import re
def _extract(text: str) -> list[str]:
codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", text.lower())
return list(dict.fromkeys(f"gfx{c}" for c in codes))
probes: list[list[str]] = []
if shutil.which("rocminfo"):
probes.append(["rocminfo"])
if shutil.which("amd-smi"):
probes.append(["amd-smi", "list"])
probes.append(["amd-smi", "static", "--asic"])
for cmd in probes:
try:
result = subprocess.run(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 15,
)
except Exception:
continue
if result.returncode != 0 or not result.stdout.strip():
continue
codes = _extract(result.stdout)
if codes:
return codes
return []
# Set by _ensure_rocm_torch() on success; suppresses the post-install AMD warning.
_rocm_windows_torch_installed: bool = False
def _install_bnb_windows_rocm() -> bool:
"""Install the AMD Windows BNB prerelease wheel. Returns True on success.
The continuous-release wheel is intentionally mismatched: the filename
encodes version 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the
wheel metadata reports 0.50.0.dev0. uv rejects this filename/metadata
mismatch; set UV_SKIP_WHEEL_FILENAME_CHECK=1 to bypass that check, then
restore the previous value (or remove the var) when done.
"""
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
if _bnb_win_url is None:
return False
_old = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK")
os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1"
try:
_ok = pip_install_try(
"bitsandbytes (AMD Windows, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_win_url,
constrain = False,
)
finally:
if _old is None:
os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None)
else:
os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = _old
if not _ok:
return False
# After install: detect the actual ROCm DLL suffix shipped in the wheel and
# set BNB_ROCM_VERSION so bitsandbytes loads the correct DLL regardless of
# what torch.version.hip reports. The wheel may ship an older suffix (e.g.
# "72") while torch reports a newer HIP version (e.g. 7.13); the env var
# override ensures bitsandbytes does not fail looking for a non-existent DLL.
# The worker subprocess inherits this env var automatically.
# Fall back to "72" if detection fails (e.g. install was a no-op / dry-run).
if "BNB_ROCM_VERSION" not in os.environ:
_ver = _detect_bnb_rocm_dll_ver() or "72"
os.environ["BNB_ROCM_VERSION"] = _ver
return True
def _ensure_rocm_torch() -> None:
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
Runs only on Linux x86_64 hosts where an AMD GPU is present and the
ROCm runtime is detectable (rocminfo / amd-smi / hipconfig /
rocm-core package). No-op when torch already links against HIP
(ROCm), on Windows / macOS, on non-x86_64 Linux (PyTorch does not
publish ROCm wheels for aarch64 / arm64), or on mixed AMD+NVIDIA
hosts (NVIDIA takes precedence).
On Linux x86_64: uses pytorch.org ROCm wheel index tags.
On Windows: uses AMD's repo.amd.com arch-specific pip index.
No-op on macOS, non-x86_64 Linux, NVIDIA-primary hosts, or when torch
already links against HIP.
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
# Explicit OS / architecture guards so the helper is safe to call
# from any context -- PyTorch only publishes ROCm wheels for
# linux_x86_64, so aarch64 / arm64 hosts must skip this repair path
# instead of failing the update with a missing-wheel error.
if IS_WINDOWS or IS_MACOS:
global _rocm_windows_torch_installed
# setup.ps1 sets this when it already installed AMD wheels; skip the probe
# only when torch is actually importable as ROCm. If the venv was wiped
# between runs, the stale env-var would suppress a needed reinstall.
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_torch_ok = False
try:
_probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"import sys; "
"sys.exit(0 if (hip or 'rocm' in torch.__version__.lower()) else 1)"
),
],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
timeout = 90,
)
_torch_ok = _probe.returncode == 0
except (OSError, subprocess.TimeoutExpired):
pass
if _torch_ok:
_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.
_install_bnb_windows_rocm()
return
# torch was wiped between runs; fall through to the full install path
if IS_MACOS:
return
if IS_WINDOWS:
if _has_usable_nvidia_gpu():
return
gfx_arch = _detect_windows_gfx_arch()
if not gfx_arch:
return # no AMD GPU visible via hipinfo
# Probe whether torch already links against HIP.
_torch_already_rocm = False
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"ver=torch.__version__; "
"print('yes' if hip or 'rocm' in ver.lower() else '')"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 90,
)
if probe.returncode == 0 and probe.stdout.decode().strip() == "yes":
_torch_already_rocm = True
except (OSError, subprocess.TimeoutExpired):
pass
if not _torch_already_rocm:
index_url = _windows_rocm_index_url(gfx_arch)
if index_url is None:
print(
f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping"
)
return
print(f" {gfx_arch} (Windows) -- installing torch from {index_url}")
pip_install(
f"ROCm torch (Windows, {gfx_arch})",
"--force-reinstall",
"--index-url",
index_url,
"torch",
"torchvision",
"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.
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 ──
if platform.machine().lower() not in {"x86_64", "amd64"}:
return
# NVIDIA takes precedence on mixed hosts -- but only if an actual GPU is usable
@ -297,11 +682,19 @@ 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,
timeout = 30,
timeout = 90,
)
except (OSError, subprocess.TimeoutExpired):
probe = None
@ -313,7 +706,83 @@ def _ensure_rocm_torch() -> None:
rocm_torch_ready = has_hip_torch
if not has_hip_torch:
# Strix Halo / Strix Point (gfx1151 / gfx1150) segfault under ROCm 7.1
# in torch._grouped_mm. AMD's per-gfx repo ships torch 2.11.0+rocm7.13.0
# with the real fix, so route those hosts there instead of the generic
# pytorch.org rocm7.1 wheel. Mirrors install.sh's Strix override.
# On mixed hosts (Strix iGPU + non-Strix dGPU), only route to the AMD
# per-gfx index when the GPU HIP will actually run on is the Strix one --
# otherwise the dGPU would get an incompatible wheel. Use HIP_VISIBLE_DEVICES
# to determine the runtime target.
_strix_override_url: "str | None" = None
_strix_override_pkgs: "tuple[str, str, str] | None" = None
if ver < (7, 2):
gfx_codes = _detect_amd_gfx_codes()
_strix_gfx = {"gfx1151", "gfx1150"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
if _detected_strix:
# Pick the runtime-visible GPU. If HIP_VISIBLE_DEVICES selects a
# specific index into gfx_codes, use that gfx; else default to the
# first listed GPU. Skip the override unless the resolved GPU is
# Strix.
_runtime_gfx = (
gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None
)
if _runtime_gfx in _strix_gfx:
_selected_gfx = _runtime_gfx
_amd_mirror = (
os.environ.get("UNSLOTH_AMD_ROCM_MIRROR")
or "https://repo.amd.com/rocm/whl"
).rstrip("/")
_strix_override_url = f"{_amd_mirror}/{_selected_gfx}/"
_strix_override_pkgs = (
"torch>=2.11.0,<2.12.0",
# Pin torchvision/torchaudio to the 2.11.x-compatible range.
# The install uses --index-url (exclusive, no PyPI fallback),
# so bare unversioned names risk resolving a build from AMD's
# index that targets a different torch major (e.g. 0.27 built
# against torch 2.12), which would fail at runtime with an
# ABI/version mismatch. Matches _ROCM_TORCH_CONSTRAINT["rocm7.2"].
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
)
print(
f"\n {_selected_gfx} (AMD Strix) is the runtime target with ROCm "
f"{ver[0]}.{ver[1]}.\n"
f" ROCm 7.1 has a known _grouped_mm segfault on this GPU;\n"
f" routing torch install to AMD's arch-specific index\n"
f" ({_strix_override_url}) which serves torch 2.11.0+rocm7.13.0\n"
f" with the upstream fix.\n"
)
else:
_gfx_str = ", ".join(sorted(_detected_strix))
print(
f"\n Strix GPU ({_gfx_str}) present but HIP_VISIBLE_DEVICES "
f"selects a non-Strix runtime target ({_runtime_gfx});\n"
f" skipping AMD per-gfx index override.\n"
)
# Strix override on ROCm 7.1 must fire even when has_hip_torch is True --
# an existing torch with `torch.version.hip == "7.1"` is exactly the broken
# combo the override is meant to repair, so skipping it leaves users on
# the known _grouped_mm segfault.
if _strix_override_url is not None and _strix_override_pkgs is not None:
index_url = _strix_override_url
_torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
print(f" Strix ROCm 7.1 override -- installing torch from {index_url}")
pip_install(
"ROCm torch (Strix arch-specific)",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
rocm_torch_ready = True
elif not has_hip_torch:
# Select best matching wheel tag (newest ROCm version <= installed)
tag = next(
(
@ -331,13 +800,16 @@ def _ensure_rocm_torch() -> None:
else:
index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
_torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
tag, _ROCM_TORCH_PKG_SPECS["_default"]
)
pip_install(
f"ROCm torch ({tag})",
"--force-reinstall",
"--no-cache-dir",
"torch>=2.4,<2.11.0",
"torchvision<0.26.0",
"torchaudio<2.11.0",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
@ -346,7 +818,9 @@ def _ensure_rocm_torch() -> None:
# Install bitsandbytes only when torch links against ROCm. Prefers the
# continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls
# back to PyPI when the pre-release URL is unreachable.
# back to PyPI when the pre-release wheel cannot be installed. Use pip for
# the pre-release wheel because uv rejects the wheel's filename/metadata
# version mismatch.
if rocm_torch_ready:
_bnb_url = _bnb_rocm_prerelease_url()
_bnb_installed = False
@ -358,11 +832,12 @@ def _ensure_rocm_torch() -> None:
"--no-deps",
_bnb_url,
constrain = False,
force_pip = True,
)
if not _bnb_installed:
print(
_red(
" bnb pre-release unreachable; falling back to PyPI "
" bnb pre-release install failed; falling back to PyPI "
"(4-bit decode will be broken on ROCm)"
)
)
@ -809,6 +1284,7 @@ def pip_install_try(
label: str,
*args: str,
constrain: bool = True,
force_pip: bool = False,
) -> bool:
"""Like pip_install but returns False on failure instead of exiting.
For optional installs with a follow-up fallback.
@ -819,7 +1295,7 @@ def pip_install_try(
constraint_args_pip = ["-c", str(CONSTRAINTS)]
constraint_args_uv = ["-c", _uv_safe_path(CONSTRAINTS)]
if USE_UV:
if USE_UV and not force_pip:
cmd = _build_uv_cmd(args) + constraint_args_uv
else:
cmd = _build_pip_cmd(args) + constraint_args_pip
@ -948,8 +1424,12 @@ def install_python_stack() -> int:
base_total = 10 if IS_WINDOWS else 11
if IS_MACOS:
base_total -= 1 # triton step is skipped on macOS
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
base_total += 3
if not IS_MACOS and not NO_TORCH:
base_total += 1 # ROCm torch check (line 1526) -- all non-macOS platforms
if not IS_WINDOWS:
base_total += (
2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only
)
_TOTAL = (base_total - 1) if skip_base else base_total
# 1. Try to use uv for faster installs (must happen before pip upgrade
@ -1121,12 +1601,12 @@ def install_python_stack() -> int:
# 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the
# venv received CPU-only torch (common when pip resolves torch from PyPI).
# Must come immediately after base packages so torch is present for inspection.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
if not IS_MACOS and not NO_TORCH:
_progress("ROCm torch check")
_ensure_rocm_torch()
# Windows + AMD GPU: PyTorch does not publish ROCm wheels for Windows.
# Detect and warn so users know manual steps are needed for GPU training.
# Windows + AMD GPU: if ROCm torch was not installed (wrong Python version
# or unknown ROCm version), warn the user.
if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu():
# Validate actual AMD GPU presence (not just tool existence)
import re as _re_win
@ -1155,14 +1635,14 @@ def install_python_stack() -> int:
if _wr.returncode == 0 and _check_fn(_wr.stdout):
_win_amd_gpu = True
break
if _win_amd_gpu:
if _win_amd_gpu and not _rocm_windows_torch_installed:
_safe_print(
_dim(" Note:"),
"AMD GPU detected on Windows. ROCm-enabled PyTorch must be",
"AMD GPU detected but ROCm PyTorch could not be auto-installed.",
)
_safe_print(
" " * 8,
"installed manually. See: https://docs.unsloth.ai/get-started/install-and-update/amd",
"Manual install may be required. See: https://docs.unsloth.ai/get-started/install-and-update/amd",
)
# 3. Extra dependencies
@ -1189,10 +1669,17 @@ def install_python_stack() -> int:
_progress("dependency overrides (skipped, no torch)")
else:
_progress("dependency overrides")
_override_extra_args: tuple[str, ...] = ()
if _rocm_windows_torch_installed:
# torchao in overrides.txt declares torch as a dependency; without
# --no-deps uv would resolve and install CPU torch from PyPI,
# overwriting the AMD ROCm wheels we just installed.
_override_extra_args = ("--no-deps",)
pip_install(
"Installing dependency overrides",
"--force-reinstall",
"--no-cache-dir",
*_override_extra_args,
req = REQ_ROOT / "overrides.txt",
)

354
studio/setup.ps1
View file

@ -685,13 +685,228 @@ if (-not $HasNvidiaSmi) {
}
}
}
# ── AMD ROCm detection (Windows): probe hipinfo/amd-smi for actual GPU ──
$HasROCm = $false
$HipSdkInstalled = $false # HIP SDK binary found (independent of device accessibility)
$ROCmGpuLabel = $null
$script:ROCmGfxArch = $null
if (-not $HasNvidiaSmi) {
# hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution).
# AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type.
$hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue
if (-not $hipinfoExe) {
$hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null }
$hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" }
if ($hipRoot) {
$hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe"
if (Test-Path $hipinfoCandidate) {
substep "[WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" "Yellow"
substep " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" "Yellow"
substep " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" "Yellow"
$hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate }
} else {
substep "[WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" "Yellow"
substep " HIP SDK install may be incomplete -- re-install from:" "Yellow"
substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
}
}
}
if ($hipinfoExe) {
$HipSdkInstalled = $true # binary found → SDK is installed regardless of device state
try {
$hipOut = & $hipinfoExe.Source 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
$HasROCm = $true
$_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() })
$_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
if ($_hipAllArches.Count -gt 0) {
$script:ROCmGfxArch = if ($_hipVisIdx -lt $_hipAllArches.Count) { $_hipAllArches[$_hipVisIdx] } else { $_hipAllArches[0] }
$ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)"
} else {
$ROCmGpuLabel = "AMD ROCm"
}
} elseif ($LASTEXITCODE -ne 0) {
# hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected")
$firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1)
substep "[WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" "Yellow"
substep " $firstLine" "Yellow"
substep " Ensure ROCm drivers are installed: https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
}
} catch {}
}
# amd-smi fallback: HIP runtime present but hipinfo unavailable (no full HIP SDK).
# Confirms GPU visibility via 'list', then attempts 'static --asic' to extract
# the gfx arch that hipinfo would have provided. Critical for Strix Halo
# (gfx1151) and other iGPUs where only the HIP runtime is installed.
if (-not $HasROCm) {
$amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiExe) {
try {
$smiOut = & $amdSmiExe.Source list 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") {
$HasROCm = $true
# Attempt 1: newer amd-smi versions embed the gfx arch in list output.
# Collect ALL gfx tokens in output order so that on mixed-arch systems
# we can honour HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES and pick the
# arch for the *runtime-visible* GPU rather than always the first one.
# Do NOT deduplicate: a dual same-arch system (e.g. two gfx1151 APUs)
# must produce a 2-element array so HIP_VISIBLE_DEVICES=1 selects the
# second GPU rather than triggering a false out-of-range warning.
# Note: this mapping assumes amd-smi lists GPUs in the same order as
# HIP enumerates them (both follow PCI bus order in practice); it may
# give the wrong arch when GPU indices are non-contiguous (very rare).
$allGfxArches = @([regex]::Matches($smiOut, '(?i)\b(gfx\d+[a-z]?)\b') |
ForEach-Object { $_.Groups[1].Value.ToLower() })
if ($allGfxArches.Count -gt 0) {
# Resolve which GPU index is runtime-visible. When a single
# integer index is set, use it; fall back to index 0 otherwise
# (comma-separated lists or unset → first GPU, same as before).
$visGpu = if ($env:HIP_VISIBLE_DEVICES) { $env:HIP_VISIBLE_DEVICES }
elseif ($env:ROCR_VISIBLE_DEVICES) { $env:ROCR_VISIBLE_DEVICES }
else { $null }
$gpuIdx = 0
if ($visGpu -match '^\s*(\d+)\s*$') { $gpuIdx = [int]$Matches[1] }
if ($gpuIdx -ge $allGfxArches.Count) {
substep "[WARN] HIP/ROCR_VISIBLE_DEVICES index $gpuIdx is out of range ($($allGfxArches.Count) GPU(s) detected); defaulting to GPU 0 for arch selection" "Yellow"
$gpuIdx = 0
}
$script:ROCmGfxArch = $allGfxArches[$gpuIdx]
$ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)"
} else {
# Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+,
# including the GFX target needed for wheel index selection.
$smiAsicOut = ""
try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {}
if ($smiAsicOut -match "(?i)\b(gfx\d+[a-z]?)\b") {
$script:ROCmGfxArch = $Matches[1].ToLower()
$ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)"
} elseif ($smiAsicOut -match "(?im)Market.?Name\s*[:\|]\s*([^\r\n]+)") {
$ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))"
} else {
$ROCmGpuLabel = "AMD ROCm"
}
}
}
} catch {}
}
}
# WMI fallback: AMD GPU in device list but no HIP SDK → guide the user.
# WMI gives a marketing name (e.g. "AMD Radeon 890M") but never a gfx arch.
# $HasROCm is intentionally NOT set here — we cannot confirm ROCm runtime
# support without hipinfo or amd-smi. The name is saved to $ROCmGpuLabel
# so the name-based inference below can still attempt an arch lookup.
if (-not $HasROCm) {
try {
$wmiGpu = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue |
Where-Object { $_.Name -match "AMD|Radeon" } |
Select-Object -First 1
if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name }
} catch {}
}
# ── Arch resolution: env-var override → name inference ──────────────────
# Runs after all probe methods. Covers users whose amd-smi version is too
# old to report the GFX target and who don't have hipinfo (HIP-runtime-only
# installs, common on Strix Halo / iGPU systems).
if ($HasROCm -and -not $script:ROCmGfxArch) {
# 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running.
if ($env:UNSLOTH_ROCM_GFX_ARCH) {
$script:ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower()
$ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)"
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $script:ROCmGfxArch" "Cyan"
}
# 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI).
# Ordered most-specific first; first match wins.
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4
@{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail)
@{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point)
@{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop
@{ P = "RX 7600"; A = "gfx1102" } # RDNA 3
@{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix)
)
foreach ($row in $nameArchTable) {
if ($ROCmGpuLabel -match $row.P) {
$script:ROCmGfxArch = $row.A
$ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)"
substep "gfx arch inferred from GPU name: $script:ROCmGfxArch" "Cyan"
substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$script:ROCmGfxArch to skip inference next time" "Cyan"
break
}
}
}
}
# Capture ROCm version early for display and wheel selection.
# Run whenever the HIP SDK binary is present, not just when the device is accessible --
# hipconfig --version works even when hipinfo reports no ROCm device (driver issue).
if ($HasROCm -or $HipSdkInstalled) {
$script:ROCmVersion = $null
$hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue
if (-not $hipConfigExe) {
$hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null }
if ($hipRoot) {
$hipConfigCandidate = Join-Path $hipRoot "bin\hipconfig.exe"
if (Test-Path $hipConfigCandidate) {
$hipConfigEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" }
substep "[WARN] hipconfig not on PATH -- located via ${hipConfigEnvLabel}: $hipConfigCandidate" "Yellow"
$hipConfigExe = [PSCustomObject]@{ Source = $hipConfigCandidate }
}
}
}
if ($hipConfigExe) {
try {
$hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0) {
$hipVerLine = ($hipVerOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1).Trim()
if ($hipVerLine -match '(\d+\.\d+)') {
$script:ROCmVersion = $Matches[1]
$script:ROCmVersionFull = $hipVerLine
}
}
} catch {}
}
if (-not $script:ROCmVersion) {
$amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiVer) {
try {
$smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') { $script:ROCmVersion = $Matches[1] }
} catch {}
}
}
}
}
if ($HasNvidiaSmi) {
step "gpu" "NVIDIA GPU detected"
} elseif ($HasROCm) {
step "gpu" $ROCmGpuLabel
$hipSdkPath = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { "on system PATH" }
substep "HIP SDK: $hipSdkPath"
if ($script:ROCmVersionFull) { substep "hipconfig: $script:ROCmVersionFull" }
} elseif ($HipSdkInstalled -and $ROCmGpuLabel) {
# HIP SDK is installed but ROCm can't see the device (driver issue, not SDK issue)
$sdkVer = if ($script:ROCmVersionFull) { " (HIP $script:ROCmVersionFull)" } else { "" }
Write-Host ""
step "gpu" "none (chat-only / GGUF)" "Yellow"
substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow"
step "gpu" "AMD GPU detected -- not ROCm-accessible$sdkVer" "Yellow"
substep "Detected: $ROCmGpuLabel" "Yellow"
substep "[WARN] HIP SDK is installed but hipinfo reports no ROCm-capable device." "Yellow"
substep " This is a driver issue, not an SDK issue." "Yellow"
substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow"
substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} elseif ($ROCmGpuLabel) {
Write-Host ""
step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow"
substep "Detected: $ROCmGpuLabel" "Yellow"
substep "Install the HIP SDK for ROCm GPU inference:" "Yellow"
substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
Write-Host ""
} else {
step "gpu" "NVIDIA GPU detected"
Write-Host ""
step "gpu" "none (chat-only / GGUF)" "Yellow"
substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow"
Write-Host ""
}
# ============================================
@ -1102,6 +1317,13 @@ if (-not $CudaArch) {
step "cuda" "skipped (no NVIDIA GPU detected)" "Yellow"
}
if ($HasROCm) {
$rocmVerLabel = if ($script:ROCmVersionFull) { "ROCm $script:ROCmVersionFull" } elseif ($script:ROCmVersion) { "ROCm $script:ROCmVersion" } else { "ROCm (version unknown)" }
step "rocm" $rocmVerLabel
} elseif ($ROCmGpuLabel) {
step "rocm" "HIP SDK not found -- GPU-accelerated training unavailable" "Yellow"
}
# ============================================
# 1f. Node.js / npm (skip if pip-installed or Tauri -- only needed for frontend build)
# ============================================
@ -1516,7 +1738,7 @@ if (-not $PythonCmd) {
exit 1
}
substep "Using $PythonCmd ($(& $PythonCmd --version 2>&1))"
substep "Python found: $PythonCmd"
# The venv must already exist (created by install.ps1); this script only
# updates packages. UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the
@ -1684,6 +1906,13 @@ if (-not (Test-Path -LiteralPath $VenvDir)) {
exit 1
} else {
substep "reusing existing virtual environment at $VenvDir"
$_venvPyExe = Join-Path $VenvDir "Scripts\python.exe"
if (Test-Path -LiteralPath $_venvPyExe) {
try {
$_venvPyVer = (& $_venvPyExe --version 2>&1 | Out-String).Trim()
if ($_venvPyVer) { substep $_venvPyVer }
} catch {}
}
}
# pip and python write to stderr even on success (progress bars, warnings).
@ -1795,9 +2024,91 @@ if ($HasNvidiaSmi) {
$CuTag = "cpu"
}
# ── GPU arch → newest compatible Windows ROCm wheel release ──
# Wheels bundle their own ROCm runtime; the installed HIP SDK version does
# not constrain which release to use. Always picks the newest release that
# supports the GPU architecture.
# ── AMD Windows ROCm torch override ──────────────────────────────────────────
# Uses AMD's arch-specific pip index (repo.amd.com/rocm/whl/{arch}/).
# Wheels bundle their own ROCm runtime; HIP SDK version is irrelevant.
$ROCmGfxArch = $script:ROCmGfxArch
$ROCmIndexUrl = $null
if ($HasROCm -and $CuTag -eq "cpu") {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X and Strix have a null _grouped_mm kernel on torch <2.11.0.
# Mirrors the $torchFloorMap in install.ps1 so both installers enforce
# the same floor and ceiling when pulling from AMD's per-arch index.
$torchFloorMap = @{
"gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0"
"gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0"
}
# Companion ranges for torchvision/torchaudio -- must stay in sync with the
# torch ceiling so pip can always find a consistent trio on AMD's per-arch
# index. AMD publishes each package independently and may add a newer
# torchvision (e.g. 0.27 for torch 2.12) before removing 0.26, which would
# cause pip to resolve an ABI-incompatible set if these are left bare.
# Matches _ROCM_TORCH_PKG_SPECS["rocm7.2"] in install_python_stack.py.
# Bump all three ceilings together when torch 2.12.x is validated.
$torchvisionFloorMap = @{
"gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0"
"gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0"
}
$torchaudioFloorMap = @{
"gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0"
"gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0"
}
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
$ROCmTorchSpec = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { "torch" }
$ROCmVisionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$ROCmAudioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
if ($archFamily) {
$ROCmIndexUrl = "$amdIndexBase/$archFamily/"
} elseif ($ROCmGfxArch) {
# GPU arch detected but not in the supported wheel map — warn explicitly
# so the user knows why they are getting CPU PyTorch instead of ROCm.
substep "[WARN] AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow"
substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx90a, gfx908" "Yellow"
} else {
# HIP SDK present ($HasROCm=true via amd-smi) but gcnArchName was not
# readable — warn rather than silently falling back to CPU PyTorch.
substep "[WARN] AMD GPU detected (HIP SDK present) but GPU arch could not be read -- falling back to CPU-only PyTorch" "Yellow"
substep " Arch detection requires hipinfo to report gcnArchName. Re-install the HIP SDK if this is unexpected." "Yellow"
}
}
$PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
if ($CuTag -eq "cpu") {
if ($ROCmIndexUrl) {
substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..."
if ($ROCmTorchSpec -ne "torch") {
substep " enforcing $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec (known _grouped_mm bug in older wheels)" "Cyan"
}
if ($script:UnslothVerbose) {
Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
$output = Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | Out-String
$torchInstallExit = $LASTEXITCODE
}
if ($torchInstallExit -ne 0) {
Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow
Write-Host $output -ForegroundColor Yellow
$ROCmIndexUrl = $null
} else {
# Tell install_python_stack.py to skip probe + suppress manual-install warning.
$env:UNSLOTH_ROCM_TORCH_INSTALLED = "1"
}
}
if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
substep "installing PyTorch (CPU-only)..."
if ($script:UnslothVerbose) {
Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/cpu"
@ -1812,7 +2123,7 @@ if ($CuTag -eq "cpu") {
Write-Host $output -ForegroundColor Red
exit 1
}
} else {
} elseif (-not $ROCmIndexUrl) {
substep "installing PyTorch with CUDA support ($CuTag)..."
substep "(This download is ~2.8 GB -- may take a few minutes)"
if ($script:UnslothVerbose) {
@ -2115,6 +2426,23 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
substep "installing prebuilt llama.cpp bundle (preferred path)..."
if (Test-Path -LiteralPath $LlamaCppDir) {
substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement"
# If the existing install is the wrong kind (e.g. windows-cpu on a ROCm
# machine that should have windows-hip), remove it so the installer is
# forced to download the correct variant rather than skipping on tag match.
$existingMetaPath = Join-Path $LlamaCppDir "UNSLOTH_PREBUILT_INFO.json"
if (Test-Path $existingMetaPath) {
try {
$existingMeta = Get-Content $existingMetaPath -Raw | ConvertFrom-Json
$existingKind = $existingMeta.install_kind
$expectedKind = if ($HasROCm) { "windows-hip" } elseif ($HasNvidiaSmi) { "windows-cuda" } else { "windows-cpu" }
if ($existingKind -and $existingKind -ne $expectedKind) {
substep "Removing mismatched llama.cpp install (found '$existingKind', need '$expectedKind')..."
Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue
}
} catch {
# unreadable metadata -- let the installer handle it
}
}
}
# why: install_llama_prebuilt.py uses os.replace(), which would displace
# an unrelated $env:UNSLOTH_STUDIO_HOME\llama.cpp before the source-build
@ -2129,6 +2457,9 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
"--published-repo", $HelperReleaseRepo,
"--simple-policy"
)
if ($HasROCm) {
$prebuiltArgs += "--has-rocm"
}
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
$prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG)
}
@ -2589,6 +2920,17 @@ if (-not $NeedLlamaSourceBuild) {
# CUDA flags -- only if GPU available, otherwise explicitly disable
if ($HasNvidiaSmi -and $NvccPath) {
$CmakeArgs += '-DGGML_CUDA=ON'
# Accept a host MSVC newer than nvcc's whitelist; a fresh toolkit
# (e.g. CUDA 13.3) otherwise aborts with "#error -- unsupported
# Microsoft Visual Studio version!". Mirrors the Linux fix. Via env
# (covers the configure probe + build), after Refresh-Environment, idempotent.
$nvccAllowFlag = '-allow-unsupported-compiler'
if ([string]::IsNullOrEmpty($env:NVCC_PREPEND_FLAGS)) {
$env:NVCC_PREPEND_FLAGS = $nvccAllowFlag
} elseif ($env:NVCC_PREPEND_FLAGS -notlike "*$nvccAllowFlag*") {
$env:NVCC_PREPEND_FLAGS = "$($env:NVCC_PREPEND_FLAGS) $nvccAllowFlag"
}
substep "NVCC_PREPEND_FLAGS = $env:NVCC_PREPEND_FLAGS"
$CmakeArgs += "-DCUDAToolkit_ROOT=$CudaToolkitRoot"
$CmakeArgs += "-DCUDA_TOOLKIT_ROOT_DIR=$CudaToolkitRoot"
$CmakeArgs += "-DCMAKE_CUDA_COMPILER=$NvccPath"

258
studio/setup.sh
View file

@ -130,6 +130,30 @@ run_quiet_no_exit() {
_run_quiet return "$@"
}
_nvcc_meets_llama_minimum() {
# Echo "ok|too_old|unknown" then the parsed "X.Y" version, one per line.
# llama.cpp needs CUDA toolkit >= 12.4 (#4437; setup.ps1 aborts via #4517).
_nvcc_bin=$1
[ -n "$_nvcc_bin" ] || { echo "unknown"; echo ""; return 0; }
_raw=$("$_nvcc_bin" --version 2>/dev/null \
| sed -n 's/.*release \([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \
| head -1)
if [ -z "$_raw" ]; then
echo "unknown"; echo ""; return 0
fi
_maj=${_raw%%.*}
_min_raw=${_raw#*.}
_min=${_min_raw%%.*}
if [ "$_maj" -lt 12 ] 2>/dev/null; then
echo "too_old"
elif [ "$_maj" -eq 12 ] && [ "$_min" -lt 4 ] 2>/dev/null; then
echo "too_old"
else
echo "ok"
fi
echo "$_raw"
}
print_llama_error_log() {
local log_file=$1
[ -s "$log_file" ] || return 0
@ -157,12 +181,21 @@ if not isinstance(payload, dict):
repo = str(payload.get("published_repo") or "").strip()
release_tag = str(payload.get("release_tag") or "").strip()
llama_tag = str(payload.get("tag") or "").strip()
source = str(payload.get("source") or "").strip()
binary_repo = str(payload.get("binary_repo") or "").strip()
binary_tag = str(payload.get("binary_release_tag") or "").strip()
if not repo or not release_tag:
raise SystemExit(0)
message = f"installed release: {repo}@{release_tag}"
if llama_tag and llama_tag != release_tag:
message += f" (tag {llama_tag})"
# For non-upstream sources (e.g. lemonade) the published_repo/release_tag
# refer to the unsloth source tree while the actual binaries came from a
# different repo. Show both so the log is unambiguous.
if source and source != "upstream" and binary_repo and binary_tag and binary_repo != repo:
message = f"installed release: {repo}@{release_tag} + {source}@{binary_tag}"
else:
message = f"installed release: {repo}@{release_tag}"
if llama_tag and llama_tag != release_tag:
message += f" (tag {llama_tag})"
print(message)
PY
}
@ -635,6 +668,81 @@ if [ "$_NEED_T5_INSTALL" = true ]; then
fi
fi
# ── GPU detection summary (mirrors setup.ps1 step "gpu" block) ──
_setup_amd_detected=false
_setup_gfx_all=""
_setup_mkt=""
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
_setup_amd_detected=true
_setup_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_setup_mkt=$(rocminfo 2>/dev/null | awk -F': ' \
'/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
elif command -v amd-smi >/dev/null 2>&1 && \
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
_setup_amd_detected=true
_setup_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
[ -z "$_setup_gfx_all" ] && \
_setup_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_setup_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
'/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
fi
if command -v nvidia-smi >/dev/null 2>&1 && \
nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
step "gpu" "NVIDIA GPU detected"
elif [ "$_setup_amd_detected" = true ]; then
_setup_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
_setup_vis_idx=0
if [ -n "$_setup_vis" ] && [ "$_setup_vis" != "-1" ]; then
_setup_first="${_setup_vis%%,*}"
case "$_setup_first" in ''|*[!0-9]*) ;; *) _setup_vis_idx=$_setup_first ;; esac
fi
_setup_gfx=$(printf '%s\n' "$_setup_gfx_all" | awk -v idx="$_setup_vis_idx" \
'NF && !seen[$0]++ { a[n++]=$0 } END { if(idx>=n) idx=0; if(n>0) print a[idx] }')
# UNSLOTH_ROCM_GFX_ARCH env override (mirrors setup.ps1)
if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then
_setup_gfx="${UNSLOTH_ROCM_GFX_ARCH}"
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_setup_gfx"
# Name-based arch inference when tools don't report gfx (mirrors setup.ps1 nameArchTable)
elif [ -z "$_setup_gfx" ] && [ -n "$_setup_mkt" ]; then
case "$_setup_mkt" in
*"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4
*"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 iGPU
*"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 iGPU
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop
*"RX 7600"*) _setup_gfx="gfx1102" ;; # RDNA 3
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU
esac
if [ -n "$_setup_gfx" ]; then
substep "gfx arch inferred from GPU name: $_setup_gfx"
substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$_setup_gfx to skip inference next time"
fi
fi
# ROCm version via hipconfig, then amd-smi
_setup_rocm_ver=""
if command -v hipconfig >/dev/null 2>&1; then
_setup_rocm_ver=$(hipconfig --version 2>/dev/null | awk 'NR==1 && /^[0-9]/{print; exit}' || true)
fi
if [ -z "$_setup_rocm_ver" ] && command -v amd-smi >/dev/null 2>&1; then
_setup_rocm_ver=$(amd-smi version 2>/dev/null | awk -F'ROCm version: ' \
'NF>1{gsub(/[[:space:]]/,"", $2); print $2; exit}' || true)
fi
if [ -n "$_setup_gfx" ]; then
step "gpu" "AMD ROCm ($_setup_gfx)"
else
step "gpu" "AMD ROCm"
fi
_setup_rocm_root="${ROCM_PATH:-${HIP_PATH:-/opt/rocm}}"
substep "ROCm: $_setup_rocm_root"
[ -n "$_setup_rocm_ver" ] && substep "hipconfig: $_setup_rocm_ver"
[ -n "$_setup_mkt" ] && [ -n "$_setup_gfx" ] && substep "GPU: $_setup_mkt"
else
step "gpu" "none (chat-only / GGUF)" "$C_WARN"
substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU."
fi
# ── 7. Prefer prebuilt llama.cpp bundles before any source build path ──
# Nest llama.cpp under $STUDIO_HOME only for real env-overrides; legacy
# default keeps ~/.unsloth/llama.cpp so pre-PR builds are still discovered.
@ -793,6 +901,22 @@ else
fi
fi
# Source-built llama.cpp installs do not have the prebuilt metadata used above
# for exact release matching. Reuse a complete local source build unless the
# caller explicitly requested a rebuild or a PR-specific llama.cpp checkout.
if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && \
[ "$_LLAMA_FORCE_COMPILE" != "1" ] && \
[ -z "$_LLAMA_PR" ] && \
[ -x "$LLAMA_CPP_DIR/build/bin/llama-server" ] && \
[ -x "$LLAMA_CPP_DIR/build/bin/llama-quantize" ]; then
step "llama.cpp" "existing source build found; skipping rebuild"
ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize"
if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then
: > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true
fi
_NEED_LLAMA_SOURCE_BUILD=false
fi
# ── 8. WSL: pre-install GGUF build dependencies for fallback source builds ──
# On WSL, sudo requires a password and can't be entered during GGUF export
# (runs in a non-interactive subprocess). Install build deps here instead.
@ -950,7 +1074,8 @@ else
fi
if [ "$BUILD_OK" = true ]; then
CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON"
# Set Release explicitly (llama.cpp only defaults to it on non-MSVC/Xcode).
CMAKE_ARGS="-DCMAKE_BUILD_TYPE=Release -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON"
_TRY_METAL_CPU_FALLBACK=false
_HOST_SYSTEM="$(uname -s 2>/dev/null || true)"
_HOST_MACHINE="$(uname -m 2>/dev/null || true)"
@ -1005,32 +1130,52 @@ else
CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF"
_TRY_METAL_CPU_FALLBACK=true
elif [ -n "$NVCC_PATH" ]; then
CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
# Returns "ok|too_old|unknown\nX.Y" on stdout.
_NVCC_CHECK="$(_nvcc_meets_llama_minimum "$NVCC_PATH")"
_NVCC_STATUS="$(printf '%s\n' "$_NVCC_CHECK" | sed -n '1p')"
_NVCC_VER="$(printf '%s\n' "$_NVCC_CHECK" | sed -n '2p')"
CUDA_ARCHS=""
if command -v nvidia-smi &>/dev/null; then
_raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
while IFS= read -r _cap; do
_cap=$(echo "$_cap" | tr -d '[:space:]')
if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then
_arch="${BASH_REMATCH[1]}${BASH_REMATCH[2]}"
# Append if not already present
case ";$CUDA_ARCHS;" in
*";$_arch;"*) ;;
*) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_arch" ;;
esac
fi
done <<< "$_raw_caps"
fi
if [ -n "$CUDA_ARCHS" ]; then
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}"
_BUILD_DESC="building (CUDA, sm_${CUDA_ARCHS//;/+sm_})"
if [ "$_NVCC_STATUS" = "too_old" ]; then
substep "CUDA toolkit $_NVCC_VER is below llama.cpp minimum (12.4)." "$C_ERR"
substep "install a newer CUDA toolkit: https://developer.nvidia.com/cuda-toolkit-archive" "$C_WARN"
substep "falling back to CPU llama.cpp build for this run." "$C_WARN"
NVCC_PATH=""
GPU_BACKEND=""
_BUILD_DESC="building (CPU, CUDA toolkit < 12.4)"
else
_BUILD_DESC="building (CUDA)"
fi
CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_FLAGS=--threads=0"
CUDA_ARCHS=""
if command -v nvidia-smi &>/dev/null; then
_raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
while IFS= read -r _cap; do
_cap=$(echo "$_cap" | tr -d '[:space:]')
if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then
_arch="${BASH_REMATCH[1]}${BASH_REMATCH[2]}"
# Append if not already present
case ";$CUDA_ARCHS;" in
*";$_arch;"*) ;;
*) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_arch" ;;
esac
fi
done <<< "$_raw_caps"
fi
if [ -n "$CUDA_ARCHS" ]; then
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}"
_BUILD_DESC="building (CUDA, sm_${CUDA_ARCHS//;/+sm_})"
else
_BUILD_DESC="building (CUDA)"
fi
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_FLAGS=--threads=0"
# Accept a host gcc/clang newer than nvcc's whitelist; a fresh
# toolkit (e.g. CUDA 13.3) otherwise aborts with "#error --
# unsupported GNU version". Via env, not CMAKE_ARGS, to avoid
# word-splitting.
export NVCC_PREPEND_FLAGS="${NVCC_PREPEND_FLAGS:+$NVCC_PREPEND_FLAGS }-allow-unsupported-compiler"
fi
elif [ "$GPU_BACKEND" = "rocm" ]; then
# Resolve hipcc symlinks to find the real ROCm root
_HIPCC_REAL="$(readlink -f "$ROCM_HIPCC" 2>/dev/null || printf '%s' "$ROCM_HIPCC")"
@ -1044,6 +1189,32 @@ else
_BUILD_DESC="building (ROCm)"
CMAKE_ARGS="$CMAKE_ARGS -DGGML_HIP=ON"
# ROCm 7.x ships clang-20 which on Ubuntu 24.04+ defaults to the
# highest-numbered gcc lib dir (/usr/lib/gcc/x86_64-linux-gnu/14/)
# which contains runtime objects but NOT C++ headers, causing:
# fatal error: 'cstdlib' file not found
# Find the newest gcc install dir that actually has both the
# runtime dir AND /usr/include/c++/<ver> headers, then pass it
# to clang via --gcc-install-dir so HIP builds succeed.
_GCC_INSTALL_DIR=""
_gcc_pm="$(gcc -print-multiarch 2>/dev/null)"
case "$_gcc_pm" in
*-linux-gnu*) _GCC_MULTIARCH="$_gcc_pm" ;;
*) _GCC_MULTIARCH="$(uname -m)-linux-gnu" ;;
esac
for _gcc_ver in 14 13 12 11; do
if [ -d "/usr/lib/gcc/$_GCC_MULTIARCH/$_gcc_ver/include" ] && \
[ -d "/usr/include/c++/$_gcc_ver" ]; then
_GCC_INSTALL_DIR="/usr/lib/gcc/$_GCC_MULTIARCH/$_gcc_ver"
break
fi
done
if [ -n "$_GCC_INSTALL_DIR" ]; then
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_HIP_FLAGS=--gcc-install-dir=\"$_GCC_INSTALL_DIR\""
substep "ROCm HIP gcc install dir: $_GCC_INSTALL_DIR"
fi
export ROCM_PATH="$ROCM_ROOT"
export HIP_PATH="$ROCM_ROOT"
@ -1100,14 +1271,29 @@ else
CMAKE_GENERATOR_ARGS="-G Ninja"
fi
if ! run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then
# GPU label for the CPU-fallback message: Metal, else GPU_BACKEND
# (cuda/rocm). Empty on a bare CPU build (nothing to fall back from).
_gpu_fallback_label() {
if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then
echo "Metal"
elif [ -n "$GPU_BACKEND" ]; then
printf '%s' "$GPU_BACKEND" | tr '[:lower:]' '[:upper:]'
fi
}
if ! run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then
_FB_LABEL="$(_gpu_fallback_label)"
if [ -n "$_FB_LABEL" ]; then
_TRY_METAL_CPU_FALLBACK=false
substep "Metal configure failed; retrying CPU build..." "$C_WARN"
substep "$_FB_LABEL configure failed; retrying CPU build..." "$C_WARN"
rm -rf "$_BUILD_TMP/build"
run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS || BUILD_OK=false
if [ "$BUILD_OK" = true ]; then
_BUILD_DESC="building (CPU fallback)"
if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then
_BUILD_DESC="building (CPU fallback after $_FB_LABEL configure failed)"
# Now configured for CPU; clear GPU_BACKEND so a later
# build-step failure won't re-enter fallback on this config.
GPU_BACKEND=""
else
BUILD_OK=false
fi
else
BUILD_OK=false
@ -1117,12 +1303,14 @@ else
if [ "$BUILD_OK" = true ]; then
if ! run_quiet_no_exit "build llama-server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU"; then
if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then
_FB_LABEL="$(_gpu_fallback_label)"
if [ -n "$_FB_LABEL" ]; then
_TRY_METAL_CPU_FALLBACK=false
substep "Metal build failed; retrying CPU build..." "$C_WARN"
substep "$_FB_LABEL build failed; retrying CPU build..." "$C_WARN"
rm -rf "$_BUILD_TMP/build"
if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then
_BUILD_DESC="building (CPU fallback)"
_BUILD_DESC="building (CPU fallback after $_FB_LABEL build failed)"
GPU_BACKEND=""
run_quiet_no_exit "build llama-server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
else
BUILD_OK=false

1
tests/run_all.sh
View file

@ -8,6 +8,7 @@ echo "=== Bash tests ==="
sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
sh "$TESTS_DIR/sh/test_torch_constraint.sh"
sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh"
echo ""
echo "=== Python tests ==="

20
tests/sh/test_get_torch_index_url.sh
View file

@ -193,10 +193,10 @@ _result=$(run_func "$_dir")
assert_eq "ROCm 7.1 -> rocm7.1" "https://download.pytorch.org/whl/rocm7.1" "$_result"
rm -rf "$_dir"
# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.1 (capped due to torch <2.11.0)
# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.2
_dir=$(make_mock_amd_smi "7.2")
_result=$(run_func "$_dir")
assert_eq "ROCm 7.2 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result"
assert_eq "ROCm 7.2 -> rocm7.2" "https://download.pytorch.org/whl/rocm7.2" "$_result"
rm -rf "$_dir"
# 12) Both nvidia-smi and amd-smi present -> CUDA takes precedence
@ -231,10 +231,10 @@ _result=$(run_func "$_dir")
assert_eq "ROCm 7.0 -> rocm7.0" "https://download.pytorch.org/whl/rocm7.0" "$_result"
rm -rf "$_dir"
# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.1 (capped)
# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.2 (capped to latest known)
_dir=$(make_mock_amd_smi "8.0")
_result=$(run_func "$_dir")
assert_eq "ROCm 8.0 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result"
assert_eq "ROCm 8.0 -> rocm7.2 (capped)" "https://download.pytorch.org/whl/rocm7.2" "$_result"
rm -rf "$_dir"
# 18) Malformed amd-smi output (empty version field) -> cpu
@ -320,6 +320,18 @@ _result=$(run_func "$_dir")
assert_eq "CUDA UMD Version 11.8 -> cu118" "https://download.pytorch.org/whl/cu118" "$_result"
rm -rf "$_dir"
# 32) Driver-reported "CUDA Version: 13.3" (legacy header) -> cu130.
_dir=$(make_mock_smi "13.3")
_result=$(run_func "$_dir")
assert_eq "CUDA Version 13.3 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 33) "CUDA Version: 13.7" -> cu130 (until a cu137 wheel index exists).
_dir=$(make_mock_smi "13.7")
_result=$(run_func "$_dir")
assert_eq "CUDA Version 13.7 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"

121
tests/sh/test_nvcc_meets_llama_minimum.sh Executable file
View file

@ -0,0 +1,121 @@
#!/bin/bash
# Unit tests for _nvcc_meets_llama_minimum() from studio/setup.sh.
# llama.cpp needs CUDA toolkit >= 12.4 (#4437); setup.ps1 aborts via #4517,
# the Linux side was silent until this fix.
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh"
PASS=0
FAIL=0
# Extract just the helper function. The sed range is the same pattern the
# install.sh tests use.
_FUNC_FILE=$(mktemp)
sed -n '/^_nvcc_meets_llama_minimum()/,/^}/p' "$SETUP_SH" > "$_FUNC_FILE"
assert_eq() {
_label="$1"; _expected="$2"; _actual="$3"
if [ "$_actual" = "$_expected" ]; then
echo " PASS: $_label"
PASS=$((PASS + 1))
else
echo " FAIL: $_label (expected '$_expected', got '$_actual')"
FAIL=$((FAIL + 1))
fi
}
# Fake nvcc printing "release X.Y" in the canonical nvcc -V layout (the helper
# greps for "release X.Y", stable across CUDA 9.x-13.x).
make_mock_nvcc() {
_ver=$1
_dir=$(mktemp -d)
cat > "$_dir/nvcc" <<MOCK
#!/bin/sh
cat <<NV
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2026 NVIDIA Corporation
Cuda compilation tools, release $_ver, V${_ver}.0
NV
MOCK
chmod +x "$_dir/nvcc"
echo "$_dir/nvcc"
}
run_check() {
_nvcc=$1
bash -c ". '$_FUNC_FILE'; _nvcc_meets_llama_minimum '$_nvcc'"
}
echo "=== test_nvcc_meets_llama_minimum ==="
# 1) CUDA 12.4 is the minimum supported -> ok
_bin=$(make_mock_nvcc "12.4")
_out=$(run_check "$_bin")
assert_eq "12.4 status" "ok" "$(echo "$_out" | sed -n '1p')"
assert_eq "12.4 version" "12.4" "$(echo "$_out" | sed -n '2p')"
rm -rf "$(dirname "$_bin")"
# 2) CUDA 12.3 is the highest version that should be rejected.
_bin=$(make_mock_nvcc "12.3")
_out=$(run_check "$_bin")
assert_eq "12.3 status" "too_old" "$(echo "$_out" | sed -n '1p')"
rm -rf "$(dirname "$_bin")"
# 3) CUDA 12.1 (matches the original bug report in #4437).
_bin=$(make_mock_nvcc "12.1")
_out=$(run_check "$_bin")
assert_eq "12.1 status" "too_old" "$(echo "$_out" | sed -n '1p')"
rm -rf "$(dirname "$_bin")"
# 4) CUDA 11.8 -> too_old (anything < 12.0 is rejected).
_bin=$(make_mock_nvcc "11.8")
_out=$(run_check "$_bin")
assert_eq "11.8 status" "too_old" "$(echo "$_out" | sed -n '1p')"
rm -rf "$(dirname "$_bin")"
# 5) CUDA 12.8 -> ok (mid-range supported).
_bin=$(make_mock_nvcc "12.8")
_out=$(run_check "$_bin")
assert_eq "12.8 status" "ok" "$(echo "$_out" | sed -n '1p')"
rm -rf "$(dirname "$_bin")"
# 6) CUDA 13.0 -> ok.
_bin=$(make_mock_nvcc "13.0")
_out=$(run_check "$_bin")
assert_eq "13.0 status" "ok" "$(echo "$_out" | sed -n '1p')"
rm -rf "$(dirname "$_bin")"
# 7) CUDA 13.3 -> ok (the freshly shipped toolkit this fix targets).
_bin=$(make_mock_nvcc "13.3")
_out=$(run_check "$_bin")
assert_eq "13.3 status" "ok" "$(echo "$_out" | sed -n '1p')"
assert_eq "13.3 version" "13.3" "$(echo "$_out" | sed -n '2p')"
rm -rf "$(dirname "$_bin")"
# 8) Future CUDA 14.0 -> ok (no upper bound).
_bin=$(make_mock_nvcc "14.0")
_out=$(run_check "$_bin")
assert_eq "14.0 status" "ok" "$(echo "$_out" | sed -n '1p')"
rm -rf "$(dirname "$_bin")"
# 9) Empty argument -> unknown (defensive; never block the build on detection).
_out=$(run_check "")
assert_eq "empty path status" "unknown" "$(echo "$_out" | sed -n '1p')"
# 10) Mock nvcc that prints garbage -> unknown.
_dir=$(mktemp -d)
cat > "$_dir/nvcc" <<'MOCK'
#!/bin/sh
echo "totally not nvcc output"
MOCK
chmod +x "$_dir/nvcc"
_out=$(run_check "$_dir/nvcc")
assert_eq "garbage output status" "unknown" "$(echo "$_out" | sed -n '1p')"
rm -rf "$_dir"
rm -f "$_FUNC_FILE"
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ] || exit 1

20
tests/studio/install/conftest.py Normal file
View file

@ -0,0 +1,20 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Pytest configuration for studio/install tests.
install_python_stack.py does ``from backend.utils.wheel_utils import ...``
which requires the ``studio/`` directory to be on sys.path. When tests are
run from the repo root (the normal case), the studio package is not
automatically importable, so we add it here.
"""
from __future__ import annotations
import sys
from pathlib import Path
# <repo-root>/studio → makes `backend` importable as a package
_STUDIO_DIR = Path(__file__).resolve().parents[3] / "studio"
if str(_STUDIO_DIR) not in sys.path:
sys.path.insert(0, str(_STUDIO_DIR))

62
tests/studio/install/test_pr4562_bugfixes.py
View file

@ -727,16 +727,13 @@ class TestSourceCodePatterns:
assert "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" in content
def test_setup_sh_macos_metal_configure_has_cpu_fallback(self):
"""If Metal configure or build fails, setup should retry with CPU fallback."""
"""If Metal/CUDA/ROCm configure or build fails, setup retries a CPU
build. PR #5826 generalised the Metal-only wording via $_FB_LABEL; this
check stays label-agnostic so new GPU backends don't require edits."""
content = SETUP_SH.read_text()
assert "_TRY_METAL_CPU_FALLBACK=true" in content
assert (
'substep "Metal configure failed; retrying CPU build..." "$C_WARN"'
in content
)
assert (
'substep "Metal build failed; retrying CPU build..." "$C_WARN"' in content
)
assert 'configure failed; retrying CPU build..." "$C_WARN"' in content
assert 'build failed; retrying CPU build..." "$C_WARN"' in content
assert 'run_quiet_no_exit "cmake llama.cpp (cpu fallback)"' in content
assert "-DGGML_METAL=OFF" in content
# _TRY_METAL_CPU_FALLBACK must be reset to false in both fallback branches
@ -745,6 +742,55 @@ class TestSourceCodePatterns:
"_TRY_METAL_CPU_FALLBACK=false should appear at least 3 times "
"(init + configure fallback + build fallback)"
)
# The fallback helper must exist and Metal must reach it via the
# _TRY_METAL_CPU_FALLBACK shortcut so the macOS path stays covered.
assert "_gpu_fallback_label()" in content
assert 'echo "Metal"' in content
def test_setup_sh_exports_allow_unsupported_compiler(self):
"""Headline fix for PR #5826: a fresh CUDA toolkit's host-compiler
whitelist lags the distro gcc/clang, so nvcc rejects the host with
"#error -- unsupported GNU version". setup.sh exports
NVCC_PREPEND_FLAGS=-allow-unsupported-compiler (via env, not CMAKE_ARGS,
for word-splitting safety) so the build and compiler-id probe proceed."""
content = SETUP_SH.read_text()
assert "-allow-unsupported-compiler" in content
# Delivered via NVCC_PREPEND_FLAGS (covers the configure-time compiler
# probe too), not embedded in the word-split CMAKE_ARGS string.
assert "export NVCC_PREPEND_FLAGS=" in content
cmake_args_lines = [
line for line in content.splitlines() if "CMAKE_ARGS=" in line
]
assert all(
"-allow-unsupported-compiler" not in line for line in cmake_args_lines
), "flag must stay out of CMAKE_ARGS (bash word-splitting safety)"
def test_setup_ps1_exports_allow_unsupported_compiler(self):
"""Windows parity for the PR #5826 fix: a fresh CUDA toolkit's whitelist
also lags MSVC, so nvcc can reject the host with "#error -- unsupported
Microsoft Visual Studio version!". setup.ps1 sets
NVCC_PREPEND_FLAGS=-allow-unsupported-compiler in the CUDA branch (via
env, out of $CmakeArgs) so the configure probe + build proceed."""
content = SETUP_PS1.read_text()
assert "-allow-unsupported-compiler" in content
# Delivered via the process environment, not the $CmakeArgs array, so it
# reaches both the configure-time compiler probe and `cmake --build`.
assert "$env:NVCC_PREPEND_FLAGS" in content
cmake_args_lines = [
line for line in content.splitlines() if "$CmakeArgs +=" in line
]
assert all(
"-allow-unsupported-compiler" not in line for line in cmake_args_lines
), "flag must not be pushed into the $CmakeArgs array"
# Must be scoped to the CUDA branch (guarded by the GPU/nvcc check),
# not set unconditionally for CPU-only builds.
flag_idx = content.index("-allow-unsupported-compiler")
cuda_guard_idx = content.index("if ($HasNvidiaSmi -and $NvccPath)")
cuda_disable_idx = content.index("'-DGGML_CUDA=OFF'")
assert cuda_guard_idx < flag_idx < cuda_disable_idx, (
"NVCC_PREPEND_FLAGS must be set inside the CUDA-on branch, "
"before the GGML_CUDA=OFF (CPU) branch"
)
def test_macos_arm64_cpu_fallback_args_exclude_rpath(self):
"""CPU fallback args must NOT contain Metal-only RPATH flags at runtime."""

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

File diff suppressed because it is too large Load diff

5
unsloth/models/loader_utils.py
View file

@ -68,7 +68,10 @@ def _get_env_int(keys):
def _infer_distributed_ranks():
if torch.distributed.is_available() and torch.distributed.is_initialized():
if (
torch.distributed.is_available()
and getattr(torch.distributed, "is_initialized", lambda: False)()
):
try:
return torch.distributed.get_rank(), torch.distributed.get_world_size()
except Exception: