fix: detect ROCm unified memory (Strix Halo / AMD iGPU) via torch fallback

amd-smi on iGPUs with shared/unified memory (e.g. Radeon 8060S on Strix
Halo) reports only the dedicated VRAM slice (~512 MB) in its metric output,
so get_visible_gpu_utilization() was returning usable_gb ≈ 0.35 GB instead
of the full GTT pool (~128 GB).  torch.cuda.mem_get_info() already surfaces
the correct unified-pool size.

Add _reconcile_rocm_unified_memory(): after amd-smi returns a valid result
on a ROCm device, cross-check each device's vram_total_gb against
torch.cuda.mem_get_info().  When torch reports a larger total, replace the
amd-smi VRAM fields in-place.  No-op for discrete AMD GPUs where the two
sources agree.

Fixes: "Falling back to all visible GPUs -- model may not fit" on AMD iGPU
machines even when 100+ GB of unified memory is available.
This commit is contained in:
LeoBorcherding 2026-05-05 23:42:09 -05:00
commit 74c871d109

View file

@ -488,6 +488,47 @@ def get_gpu_utilization() -> Dict[str, Any]:
return {"available": False, "backend": _backend_label(device)}
def _reconcile_rocm_unified_memory(
utilization: Dict[str, Any], device_indices: list[int]
) -> None:
"""Cross-check amd-smi VRAM data against torch mem_get_info for ROCm.
On AMD iGPUs with unified/shared memory (e.g. Strix Halo / Radeon 8060S),
amd-smi reports only the dedicated VRAM slice (typically 512 MB) in its
metric output, while torch.cuda.mem_get_info() surfaces the full GTT /
unified pool (~128 GB). When torch reports a larger total than amd-smi,
replace the per-device VRAM fields so auto_select_gpu_ids sees the real
usable memory instead of the tiny dedicated slice.
"""
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", []):
idx = dev.get("index")
td = torch_by_index.get(idx)
if td is None:
continue
torch_total_gb = td["total_gb"]
smi_total_gb = dev.get("vram_total_gb") or 0.0
if torch_total_gb > smi_total_gb:
torch_used_gb = td["used_gb"]
dev["vram_total_gb"] = torch_total_gb
dev["vram_used_gb"] = torch_used_gb
dev["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 %d",
smi_total_gb,
torch_total_gb,
idx,
)
def get_visible_gpu_utilization() -> Dict[str, Any]:
device = get_device()
@ -500,6 +541,14 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
)
if result is not None:
result["backend"] = _backend_label(device)
if IS_ROCM:
# amd-smi on iGPUs with unified memory (e.g. Strix Halo)
# reports only the dedicated VRAM slice; torch mem_get_info
# sees the full unified pool. Reconcile so downstream GPU
# selection uses the real available memory.
_reconcile_rocm_unified_memory(
result, parent_visible_spec["numeric_ids"]
)
return result
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)