vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-17 11:48:46 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 24

fix(studio/worker): detect unified-memory APU by GPU name not VRAM/RAM ratio

Browse source 
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.
This commit is contained in:
LeoBorcherding 2026-05-21 02:50:07 -05:00
parent 792c3a02a2
commit 5d84704009
1 changed files with 8 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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)