fix(studio/rocm): OOM guard ROCm-only + unified memory, multi-GPU arch selection

OOM guard (worker.py):
- Scope to _hw.IS_ROCM only -- NVIDIA CUDA has a graceful OOM path and
  does not need the allocator cap
- Detect unified memory by comparing torch VRAM against psutil system RAM;
  use 0.80 on unified-memory APUs (gfx1151 Strix Halo) where the GPU pool
  is carved from host RAM, 0.90 on discrete cards

Multi-GPU arch selection:
- install.ps1 / setup.ps1: replace -match (first hit only) with
  [regex]::Matches() to collect all gcnArchName entries, then index by
  HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES
- install_python_stack.py: index into full token list before dedup so
  HIP_VISIBLE_DEVICES=2 on [gfx1100, gfx1100, gfx1151] resolves gfx1151
- install.sh: remove awk dedup from gfx token collection for same reason

GCC multiarch (setup.sh):
- Only append -linux-gnu when gcc -print-multiarch does not already return
  the full triple, fixing double-suffix on Ubuntu 24.04

install.ps1

@@ -1255,8 +1255,10 @@ shell.Run cmd, 0, False
                 $hipOut = & $hipinfoExe.Source 2>&1 | Out-String
                 if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
                     $HasROCm = $true
-                    if ($hipOut -match "(?im)^\s*gcnArchName\s*:\s*(\S+)") {
+                    $_hipAllArches = [regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() }
-                        $ROCmGfxArch  = ($Matches[1] -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"

install.sh

@@ -1855,14 +1855,14 @@ case "$TORCH_INDEX_URL" in
         # 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}' | awk '!seen[$0]++')
+            _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}' | awk '!seen[$0]++')
+            _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}' | awk '!seen[$0]++')
+                _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
             fi
         fi
         _runtime_gfx=""

studio/backend/core/training/worker.py

@@ -2264,25 +2264,33 @@ def run_training_process(
                     "skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)"
                 )
 
-    # ── 1g. ROCm/CUDA OOM guard ──
+    # ── 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/CUDA allocator at 90% of available VRAM so PyTorch raises
+    # HIP allocator so PyTorch raises OutOfMemoryError before hitting the
-    # OutOfMemoryError before hitting the hardware limit, giving the UI a
+    # hardware limit, giving the UI a clean error instead of a system freeze.
-    # clean error message instead of a system freeze.
+    # Only applied on ROCm -- NVIDIA CUDA has a graceful OOM path and does
-    # Non-fatal: if torch is not importable here (CPU-only path) the guard
+    # not need this cap.
-    # is silently skipped and the training path will handle it later.
+    # Unified-memory APUs (gfx1150/gfx1151 Strix Halo) share GPU and system
-    if _hw.DEVICE == _hw.DeviceType.CUDA:
+    # RAM in one pool: 0.90 of 128 GB starves the OS. Use 0.80 there.
+    # Non-fatal: silently skipped if torch is not importable.
+    if _hw.IS_ROCM:
         try:
             import torch as _torch_mem
+            import psutil as _psutil
 
             if _torch_mem.cuda.is_available():
-                _torch_mem.cuda.set_per_process_memory_fraction(0.90)
+                _vram_total = _torch_mem.cuda.get_device_properties(0).total_memory
+                _ram_total = _psutil.virtual_memory().total
+                _is_unified = _vram_total > (_ram_total * 0.5)
+                _mem_fraction = 0.80 if _is_unified else 0.90
+                _torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction)
                 logger.info(
-                    "GPU OOM guard: set_per_process_memory_fraction(0.90) — "
+                    "ROCm OOM guard: set_per_process_memory_fraction(%.2f) — "
-                    "HIP/CUDA allocator will raise OutOfMemoryError before "
+                    "%s memory host",
-                    "hitting the hardware VRAM limit"
+                    _mem_fraction,
+                    "unified" if _is_unified else "discrete",
                 )
         except Exception as _oom_guard_err:
             logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err)

studio/install_python_stack.py

@@ -271,11 +271,9 @@ def _detect_windows_gfx_arch() -> str | None:
     def _dedup_pick(tokens: list[str]) -> "str | None":
         if not tokens:
             return None
-        _seen: list[str] = []
+        # Index into the full (ordered) list first so HIP_VISIBLE_DEVICES
-        for _t in tokens:
+        # correctly addresses GPU N on mixed-arch hosts, then return that arch.
-            if _t not in _seen:
+        return tokens[_pick_visible_index(len(tokens))]
-                _seen.append(_t)
-        return _seen[_pick_visible_index(len(_seen))]
 
     # 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin.
     hipinfo = shutil.which("hipinfo")

studio/setup.ps1

@@ -714,8 +714,10 @@ if (-not $HasNvidiaSmi) {
             $hipOut = & $hipinfoExe.Source 2>&1 | Out-String
             if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
                 $HasROCm = $true
-                if ($hipOut -match "(?im)^\s*gcnArchName\s*:\s*(\S+)") {
+                $_hipAllArches = [regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() }
-                    $script:ROCmGfxArch = ($Matches[1] -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"

studio/setup.sh

@@ -1043,7 +1043,11 @@ else
                 # 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_MULTIARCH="$(gcc -print-multiarch 2>/dev/null || uname -m)-linux-gnu"
+                _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