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

feat: add rocm step display in setup.ps1; fix warning and progress counter

Browse source 
- Add 'rocm' step after 'cuda' in setup.ps1 showing ROCm version or HIP SDK missing
- Move ROCm version detection up to GPU detection block so it's available early
- Suppress 'must be installed manually' warning when torch.version.hip is set
- Fix _TOTAL counter to include ROCm steps on Windows (fixes 10/9 display)
This commit is contained in:
LeoBorcherding 2026-05-06 19:50:49 -05:00
commit 67d8b7481a
1 changed files with 28 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

46
studio/setup.ps1
View file

@ -704,6 +704,26 @@ if (-not $HasNvidiaSmi) {
if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name }
} catch {}
}
# Capture ROCm version early for display and wheel selection
if ($HasROCm) {
$script:ROCmVersion = $null
$hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue
if ($hipConfigExe) {
try {
$hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $hipVerOut -match '(\d+\.\d+)') { $script:ROCmVersion = $Matches[1] }
} 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) {
@ -1130,6 +1150,13 @@ if (-not $CudaArch) {
step "cuda" "skipped (no NVIDIA GPU detected)" "Yellow"
}
if ($HasROCm) {
$rocmVerLabel = if ($ROCmVersion) { "ROCm $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)
# ============================================
@ -1826,26 +1853,9 @@ if ($HasNvidiaSmi) {
# ── AMD Windows ROCm torch override ──────────────────────────────────────────
# When ROCm HIP SDK is present and Python 3.12 is in use, install AMD's direct
# torch wheels instead of CPU-only PyTorch.
$ROCmVersion = $null
$ROCmVersion = $script:ROCmVersion
$ROCmTorchWheelUrls = $null
if ($HasROCm -and $CuTag -eq "cpu") {
# Detect ROCm version via hipconfig, then amd-smi
$hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue
if ($hipConfigExe) {
try {
$hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $hipVerOut -match '(\d+\.\d+)') { $ROCmVersion = $Matches[1] }
} 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 {}
}
}
$pyVer = (& python --version 2>&1 | Out-String) -replace '[^0-9.]',''
$pyMajMin = ($pyVer.Trim() -split '\.')[0..1] -join '.'
$amdWheelBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.radeon.com/rocm/windows" }