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:
parent
3927ec0a68
commit
74c871d109
1 changed files with 49 additions and 0 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue