From 5d84704009d72b3722098a29449edbae40b13eaf Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Thu, 21 May 2026 02:50:07 -0500 Subject: [PATCH] fix(studio/worker): detect unified-memory APU by GPU name not VRAM/RAM ratio MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous heuristic (VRAM > 50 % of system RAM) false-positived on discrete cards in low-RAM systems — e.g. RX 9060 XT 16 GB on a 16 GB or 24 GB machine would trip the unified-memory path and log "unified memory host" when it should say "discrete". AMD iGPUs (gfx1150/gfx1151 Strix Halo, Strix Point, etc.) expose names with a digit+M suffix ("AMD Radeon 890M"), while discrete cards use "RX NNNN [XT|XTX]" naming. Matching that suffix is reliable across all current ROCm-capable AMD consumer GPUs and does not require psutil. Also includes the device name in the log line to ease future debugging. --- studio/backend/core/training/worker.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ad421b8101..adb78c4704 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2277,20 +2277,23 @@ def run_training_process( # Non-fatal: silently skipped if torch is not importable. if _hw.IS_ROCM: try: + import re as _re import torch as _torch_mem - import psutil as _psutil if _torch_mem.cuda.is_available(): - _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) + # iGPUs (gfx1150/gfx1151 Strix Halo, Strix Point, etc.) report + # names ending in a digit+M suffix ("AMD Radeon 890M"). + # Discrete cards use "RX NNNN [XT|XTX]" — no trailing M. + _dev_name = _torch_mem.cuda.get_device_properties(0).name + _is_unified = bool(_re.search(r'\d[Mm]\b', _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 memory host (%s)", _mem_fraction, "unified" if _is_unified else "discrete", + _dev_name, ) except Exception as _oom_guard_err: logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err)