de0c5a2f09* Studio: show Apple GPU temperature and power in the GPU monitor (macOS) The GPU monitor on Apple Silicon always showed -- for Temperature and Power: the MLX branch of get_gpu_utilization() hardcoded None because ioreg's AGXAccelerator PerformanceStatistics carries neither metric. Add utils/hardware/apple.py, mirroring macmon's no-sudo approach: - Temperature: average of the AppleSMC "Tg*" float keys via the AppleSMCKeysEndpoint user client (ctypes/IOKit, macOS 14+). - Power: IOReport "Energy Model" group, "GPU Energy" channels; each poll diffs the energy counter against the previous poll's sample, so the value is the average wattage over the polling window. The first poll only sets the baseline and returns None. Both readers latch to None on first failure and never raise, so non-Mac platforms and locked-down hosts keep the previous behavior. * Sample IOReport with the subscribed channels descriptor for PR #6187 IOReportCreateSubscription writes the channel descriptor that later samples must use; sampling with the original requested group can return no Energy Model entries on hosts that normalize the channel set, leaving power_draw_w null after the baseline. Use the subscribed descriptor (matching macmon) and fall back to the requested channels if the OS leaves it unset. --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: danielhanchen <23090290+danielhanchen@users.noreply.github.com>
1966 lines
71 KiB
Python
1966 lines
71 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||
|
||
"""
|
||
Hardware detection — run once at startup, read everywhere.
|
||
|
||
Usage:
|
||
# At FastAPI lifespan startup:
|
||
from utils.hardware import detect_hardware
|
||
detect_hardware()
|
||
|
||
# Anywhere else:
|
||
from utils.hardware import DEVICE, DeviceType, is_apple_silicon
|
||
if DEVICE == DeviceType.CUDA:
|
||
import torch
|
||
...
|
||
"""
|
||
|
||
import copy
|
||
import gc
|
||
import glob
|
||
import os
|
||
import platform
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import types
|
||
from importlib.metadata import PackageNotFoundError, version as pkg_version
|
||
import structlog
|
||
from loggers import get_logger
|
||
from enum import Enum
|
||
from pathlib import Path
|
||
from typing import Optional, Dict, Any
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
# ========== Device Enum ==========
|
||
|
||
|
||
class DeviceType(str, Enum):
|
||
"""Supported compute backends. str subclass for clean JSON serialization."""
|
||
|
||
CUDA = "cuda"
|
||
XPU = "xpu"
|
||
MLX = "mlx"
|
||
CPU = "cpu"
|
||
|
||
|
||
# ========== Global State (set once by detect_hardware) ==========
|
||
|
||
DEVICE: Optional[DeviceType] = None
|
||
CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.)
|
||
IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py
|
||
|
||
|
||
def _backend_label(device: DeviceType) -> str:
|
||
"""Return the user-facing backend name for API responses.
|
||
|
||
ROCm hosts stay ``DeviceType.CUDA`` internally (ROCm reuses ``torch.cuda.*``),
|
||
but "cuda" is misleading in JSON, so swap to ``"rocm"`` when ``IS_ROCM`` is set.
|
||
"""
|
||
if IS_ROCM and device == DeviceType.CUDA:
|
||
return "rocm"
|
||
return device.value
|
||
|
||
|
||
# ========== Detection ==========
|
||
|
||
|
||
def is_apple_silicon() -> bool:
|
||
"""True on Apple Silicon (pure platform check, no ML imports)."""
|
||
return platform.system() == "Darwin" and platform.machine() == "arm64"
|
||
|
||
|
||
def _has_torch() -> bool:
|
||
"""True if PyTorch is importable."""
|
||
try:
|
||
import torch
|
||
return True
|
||
except ImportError:
|
||
return False
|
||
|
||
|
||
def _has_mlx() -> bool:
|
||
"""True if MLX is importable."""
|
||
try:
|
||
import mlx.core
|
||
return True
|
||
except ImportError:
|
||
return False
|
||
|
||
|
||
def detect_hardware() -> DeviceType:
|
||
"""
|
||
Detect the best compute device and set the module-level DEVICE global.
|
||
|
||
Call once at FastAPI lifespan startup; idempotent.
|
||
|
||
Detection order:
|
||
1. CUDA (NVIDIA GPU, requires torch)
|
||
2. MLX (Apple Silicon via MLX framework)
|
||
3. CPU (fallback)
|
||
"""
|
||
global DEVICE, CHAT_ONLY, IS_ROCM
|
||
CHAT_ONLY = True # reset -- only CUDA/ROCm sets it to False
|
||
IS_ROCM = False
|
||
|
||
# --- CUDA / ROCm: try PyTorch ---
|
||
if _has_torch():
|
||
import torch
|
||
if torch.cuda.is_available():
|
||
DEVICE = DeviceType.CUDA
|
||
CHAT_ONLY = False
|
||
device_name = torch.cuda.get_device_properties(0).name
|
||
|
||
# Distinguish ROCm from CUDA for display only (DeviceType stays CUDA).
|
||
# AMD SDK wheels don't set torch.version.hip, so fall back to __version__.
|
||
_hip_ver = getattr(torch.version, "hip", None)
|
||
if _hip_ver is not None or "rocm" in torch.__version__.lower():
|
||
IS_ROCM = True
|
||
_hip_label = _hip_ver or torch.__version__
|
||
print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}")
|
||
else:
|
||
print(f"Hardware detected: CUDA -- {device_name}")
|
||
return DEVICE
|
||
|
||
# --- XPU: Intel GPU ---
|
||
if _has_torch():
|
||
import torch
|
||
if hasattr(torch, "xpu") and torch.xpu.is_available():
|
||
DEVICE = DeviceType.XPU
|
||
CHAT_ONLY = False
|
||
device_name = torch.xpu.get_device_name(0)
|
||
print(f"Hardware detected: XPU — {device_name}")
|
||
return DEVICE
|
||
|
||
# --- MLX: Apple Silicon ---
|
||
if is_apple_silicon() and _has_mlx():
|
||
DEVICE = DeviceType.MLX
|
||
CHAT_ONLY = False
|
||
# Use platform.machine() ("arm64"); platform.processor() returns "i386"
|
||
# on universal2 / Rosetta builds even on native arm64.
|
||
chip = platform.machine() or "arm64"
|
||
print(f"Hardware detected: MLX — Apple Silicon ({chip})")
|
||
return DEVICE
|
||
|
||
# --- Fallback ---
|
||
DEVICE = DeviceType.CPU
|
||
print("Hardware detected: CPU (no GPU backend available)")
|
||
return DEVICE
|
||
|
||
|
||
# ========== Convenience helpers ==========
|
||
|
||
|
||
def get_device() -> DeviceType:
|
||
"""
|
||
Return the detected device, auto-detecting if detect_hardware() hasn't run.
|
||
Prefer calling detect_hardware() explicitly at startup.
|
||
"""
|
||
global DEVICE
|
||
if DEVICE is None:
|
||
detect_hardware()
|
||
return DEVICE
|
||
|
||
|
||
def clear_gpu_cache():
|
||
"""
|
||
Clear GPU memory cache for the current device.
|
||
Safe on any platform — no-ops gracefully.
|
||
"""
|
||
gc.collect()
|
||
|
||
device = get_device()
|
||
|
||
if device == DeviceType.CUDA:
|
||
import torch
|
||
|
||
torch.cuda.synchronize()
|
||
torch.cuda.empty_cache()
|
||
torch.cuda.ipc_collect()
|
||
elif device == DeviceType.XPU:
|
||
import torch
|
||
torch.xpu.synchronize()
|
||
torch.xpu.empty_cache()
|
||
elif device == DeviceType.MLX:
|
||
# MLX manages memory automatically; gc.collect() above is enough.
|
||
pass
|
||
|
||
|
||
def get_gpu_memory_info() -> Dict[str, Any]:
|
||
"""
|
||
Get GPU memory info.
|
||
Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only.
|
||
"""
|
||
device = get_device()
|
||
|
||
# ---- CUDA path ----
|
||
if device == DeviceType.CUDA:
|
||
try:
|
||
import torch
|
||
|
||
idx = torch.cuda.current_device()
|
||
props = torch.cuda.get_device_properties(idx)
|
||
|
||
total = props.total_memory
|
||
allocated = torch.cuda.memory_allocated(idx)
|
||
reserved = torch.cuda.memory_reserved(idx)
|
||
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"device": idx,
|
||
"device_name": props.name,
|
||
"total_gb": total / (1024**3),
|
||
"allocated_gb": allocated / (1024**3),
|
||
"reserved_gb": reserved / (1024**3),
|
||
"free_gb": (total - allocated) / (1024**3),
|
||
"utilization_pct": (allocated / total) * 100,
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Error getting CUDA GPU info: {e}")
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"error": str(e),
|
||
}
|
||
|
||
# ---- XPU path (Intel GPU) ----
|
||
if device == DeviceType.XPU:
|
||
try:
|
||
import torch
|
||
|
||
idx = torch.xpu.current_device()
|
||
props = torch.xpu.get_device_properties(idx)
|
||
|
||
total = props.total_memory
|
||
allocated = torch.xpu.memory_allocated(idx)
|
||
reserved = torch.xpu.memory_reserved(idx)
|
||
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"device": idx,
|
||
"device_name": props.name,
|
||
"total_gb": total / (1024**3),
|
||
"allocated_gb": allocated / (1024**3),
|
||
"reserved_gb": reserved / (1024**3),
|
||
"free_gb": (total - allocated) / (1024**3),
|
||
"utilization_pct": (allocated / total) * 100,
|
||
}
|
||
except Exception as e:
|
||
logger.error("Error getting XPU GPU info: %s", e)
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"error": str(e),
|
||
}
|
||
|
||
# ---- MLX path (Apple Silicon) ----
|
||
if device == DeviceType.MLX:
|
||
try:
|
||
import mlx.core as mx
|
||
import psutil
|
||
|
||
# Unified memory: total = system RAM, GPU used from IORegistry AGX.
|
||
total = psutil.virtual_memory().total
|
||
agx = _read_apple_gpu_stats()
|
||
allocated = agx.get("vram_used_bytes", 0) if agx else 0
|
||
|
||
try:
|
||
info = mx.device_info()
|
||
# prefer machine(); processor() can return "i386" on native arm64.
|
||
gpu_name = info.get("device_name") or platform.machine() or "arm64"
|
||
except Exception:
|
||
gpu_name = platform.machine() or "arm64"
|
||
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"device": 0,
|
||
"device_name": f"Apple Silicon ({gpu_name})",
|
||
"total_gb": total / (1024**3),
|
||
"allocated_gb": allocated / (1024**3),
|
||
"reserved_gb": allocated / (1024**3),
|
||
"free_gb": (total - allocated) / (1024**3),
|
||
"utilization_pct": (allocated / total) * 100 if total else 0,
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Error getting MLX GPU info: {e}")
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"error": str(e),
|
||
}
|
||
|
||
# ---- CPU-only ----
|
||
return {"available": False, "backend": "cpu"}
|
||
|
||
|
||
def log_gpu_memory(context: str):
|
||
"""Log GPU memory usage with context."""
|
||
memory_info = get_gpu_memory_info()
|
||
if memory_info.get("available"):
|
||
backend = memory_info.get("backend", "unknown").upper()
|
||
device_name = memory_info.get("device_name", "")
|
||
label = f"{backend}" + (f" ({device_name})" if device_name else "")
|
||
logger.info(
|
||
f"GPU Memory [{context}] {label}: "
|
||
f"{memory_info['allocated_gb']:.2f}GB/{memory_info['total_gb']:.2f}GB "
|
||
f"({memory_info['utilization_pct']:.1f}% used, "
|
||
f"{memory_info['free_gb']:.2f}GB free)"
|
||
)
|
||
else:
|
||
logger.info(f"GPU Memory [{context}]: No GPU available (CPU-only)")
|
||
|
||
|
||
# ========== GPU Summary & Package Versions ==========
|
||
|
||
|
||
def get_gpu_summary() -> Dict[str, Any]:
|
||
"""
|
||
Return a compact summary of the primary GPU.
|
||
|
||
Returns dict with keys:
|
||
gpu_name – e.g. "NVIDIA L4" (or None)
|
||
vram_total_gb – e.g. 22.17 (or None)
|
||
"""
|
||
mem = get_gpu_memory_info()
|
||
if mem.get("available"):
|
||
return {
|
||
"gpu_name": mem.get("device_name"),
|
||
"vram_total_gb": round(mem.get("total_gb", 0), 2),
|
||
"vram_free_gb": round(mem.get("free_gb", 0), 2),
|
||
}
|
||
return {"gpu_name": None, "vram_total_gb": None, "vram_free_gb": None}
|
||
|
||
|
||
def get_package_versions() -> Dict[str, Optional[str]]:
|
||
"""
|
||
Return installed versions of key ML packages.
|
||
|
||
Uses importlib.metadata (stdlib), no subprocess. CUDA version from
|
||
torch.version.cuda. Returns dict keyed unsloth/torch/transformers/cuda;
|
||
missing packages yield None.
|
||
"""
|
||
packages = ("unsloth", "torch", "transformers")
|
||
versions: Dict[str, Optional[str]] = {}
|
||
|
||
for name in packages:
|
||
try:
|
||
versions[name] = pkg_version(name)
|
||
except PackageNotFoundError:
|
||
versions[name] = None
|
||
|
||
# GPU runtime version bundled with torch
|
||
try:
|
||
import torch
|
||
versions["cuda"] = getattr(torch.version, "cuda", None)
|
||
versions["rocm"] = getattr(torch.version, "hip", None)
|
||
except Exception:
|
||
versions["cuda"] = None
|
||
versions["rocm"] = None
|
||
|
||
return versions
|
||
|
||
|
||
# ========== Torch-based GPU fallbacks (AMD ROCm, Intel XPU, nvidia-smi missing) ==========
|
||
|
||
|
||
def _torch_get_device_module():
|
||
"""Return the appropriate torch device module (cuda or xpu) and its name."""
|
||
device = get_device()
|
||
import torch
|
||
|
||
if device == DeviceType.CUDA:
|
||
return torch.cuda, "cuda"
|
||
if device == DeviceType.XPU and hasattr(torch, "xpu"):
|
||
return torch.xpu, "xpu"
|
||
return None, None
|
||
|
||
|
||
def _torch_get_physical_gpu_count() -> Optional[int]:
|
||
mod, _ = _torch_get_device_module()
|
||
if mod is None:
|
||
return None
|
||
try:
|
||
return mod.device_count()
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]]:
|
||
"""Query torch for per-GPU name, total VRAM, and used VRAM."""
|
||
mod, _ = _torch_get_device_module()
|
||
if mod is None:
|
||
return []
|
||
|
||
devices = []
|
||
for ordinal, phys_idx in enumerate(device_indices):
|
||
try:
|
||
# torch ordinals are 0-based relative to CUDA_VISIBLE_DEVICES.
|
||
props = mod.get_device_properties(ordinal)
|
||
total_bytes = props.total_memory
|
||
# Prefer mem_get_info (system-wide) so auto-select sees other consumers.
|
||
if hasattr(mod, "mem_get_info"):
|
||
free_bytes, total_bytes = mod.mem_get_info(ordinal)
|
||
used_bytes = total_bytes - free_bytes
|
||
else:
|
||
used_bytes = mod.memory_allocated(ordinal)
|
||
devices.append(
|
||
{
|
||
"index": phys_idx,
|
||
"visible_ordinal": ordinal,
|
||
"name": props.name,
|
||
"total_gb": round(total_bytes / (1024**3), 2),
|
||
"used_gb": round(used_bytes / (1024**3), 2),
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.debug("torch device query failed for ordinal %d: %s", ordinal, e)
|
||
return devices
|
||
|
||
|
||
# ========== Live GPU Utilization ==========
|
||
|
||
|
||
def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]:
|
||
"""Query the appropriate SMI backend (amd-smi or nvidia-smi).
|
||
|
||
Returns the result dict if available, else None.
|
||
"""
|
||
if IS_ROCM:
|
||
backend_name = "amd-smi"
|
||
try:
|
||
from . import amd as _backend
|
||
except Exception as e:
|
||
logger.warning("%s import failed: %s", backend_name, e)
|
||
return None
|
||
else:
|
||
backend_name = "nvidia-smi"
|
||
try:
|
||
from . import nvidia as _backend
|
||
except Exception as e:
|
||
logger.warning("%s import failed: %s", backend_name, e)
|
||
return None
|
||
try:
|
||
func = getattr(_backend, func_name)
|
||
result = func(*args, **kwargs)
|
||
if isinstance(result, dict) and result.get("available"):
|
||
return result
|
||
except Exception as e:
|
||
logger.warning("%s %s query failed: %s", backend_name, func_name, e)
|
||
return None
|
||
|
||
|
||
def _read_apple_gpu_stats() -> Dict[str, Any]:
|
||
"""Query macOS IORegistry for AGX (Apple GPU) live stats. No sudo needed.
|
||
|
||
Returns dict with utilization_pct, vram_used_bytes (system-wide GPU
|
||
memory), or empty dict on failure.
|
||
"""
|
||
try:
|
||
result = subprocess.run(
|
||
["ioreg", "-r", "-c", "AGXAccelerator"],
|
||
capture_output = True,
|
||
timeout = 2,
|
||
)
|
||
text = result.stdout.decode("utf-8", errors = "replace")
|
||
except Exception:
|
||
return {}
|
||
|
||
# PerformanceStatistics block has GPU utilization and in-use memory
|
||
m = re.search(r'"PerformanceStatistics" = \{([^}]+)\}', text)
|
||
if not m:
|
||
return {}
|
||
stats_str = m.group(1)
|
||
pairs = re.findall(r'"([^"]+)"=(\d+)', stats_str)
|
||
stats = {k: int(v) for k, v in pairs}
|
||
|
||
return {
|
||
"utilization_pct": stats.get("Device Utilization %", 0),
|
||
"vram_used_bytes": stats.get("In use system memory", 0),
|
||
}
|
||
|
||
|
||
def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]:
|
||
"""Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent."""
|
||
if platform.system() != "Linux":
|
||
return None
|
||
try:
|
||
files = glob.glob("/sys/class/drm/card*/device/gpu_busy_percent")
|
||
if not files:
|
||
return None
|
||
values = [int(open(f).read().strip()) for f in files]
|
||
return round(sum(values) / len(values), 1)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _rocm_linux_sysfs_temp_c() -> Optional[float]:
|
||
"""Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C)."""
|
||
if platform.system() != "Linux":
|
||
return None
|
||
try:
|
||
files = glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input")
|
||
if not files:
|
||
return None
|
||
temps = [int(open(f).read().strip()) / 1000.0 for f in files]
|
||
return round(max(temps), 1)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _rocm_linux_sysfs_power_w() -> Optional[float]:
|
||
"""Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts)."""
|
||
if platform.system() != "Linux":
|
||
return None
|
||
try:
|
||
for pattern in (
|
||
"/sys/class/drm/card*/device/hwmon/hwmon*/power1_average",
|
||
"/sys/class/drm/card*/device/hwmon/hwmon*/power1_input",
|
||
):
|
||
files = glob.glob(pattern)
|
||
if files:
|
||
watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files)
|
||
return round(watts, 1)
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]:
|
||
"""Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes)."""
|
||
if platform.system() != "Windows":
|
||
return None
|
||
try:
|
||
ps = (
|
||
"$s=(Get-Counter '\\GPU Engine(*engtype_3D*)\\Utilization Percentage'"
|
||
" -ErrorAction SilentlyContinue).CounterSamples;"
|
||
"if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}"
|
||
)
|
||
r = subprocess.run(
|
||
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 5,
|
||
)
|
||
if r.returncode != 0 or not r.stdout.strip():
|
||
return None
|
||
val = float(r.stdout.strip())
|
||
return round(val, 1) if val >= 0 else None
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
|
||
"""Query system-wide AMD GPU VRAM via Linux DRM sysfs.
|
||
|
||
Reads /sys/class/drm/card*/device/mem_info_vram_*, which the kernel
|
||
updates in real-time across all processes. No tools required.
|
||
Returns (used_gb, total_gb) or (None, None) on failure.
|
||
"""
|
||
if platform.system() != "Linux":
|
||
return None, None
|
||
try:
|
||
used_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_used")
|
||
total_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_total")
|
||
if not used_files or not total_files:
|
||
return None, None
|
||
used_bytes = sum(int(open(f).read().strip()) for f in used_files)
|
||
total_bytes = sum(int(open(f).read().strip()) for f in total_files)
|
||
if total_bytes == 0:
|
||
return None, None
|
||
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
|
||
except Exception:
|
||
return None, None
|
||
|
||
|
||
def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]:
|
||
"""Query system-wide dedicated GPU VRAM via Windows Performance Counters.
|
||
|
||
Same data source as Task Manager, so cross-process usage is accurate.
|
||
Works for any GPU vendor without amd-smi or nvidia-smi.
|
||
Returns (used_gb, total_gb) or (None, None) on failure.
|
||
"""
|
||
if platform.system() != "Windows":
|
||
return None, None
|
||
try:
|
||
ps = (
|
||
"$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'"
|
||
" -ErrorAction SilentlyContinue).CounterSamples;"
|
||
"if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}"
|
||
)
|
||
r = subprocess.run(
|
||
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 5,
|
||
)
|
||
if r.returncode != 0 or not r.stdout.strip():
|
||
return None, None
|
||
used_bytes = float(r.stdout.strip())
|
||
if used_bytes < 0:
|
||
return None, None
|
||
import torch as _torch
|
||
|
||
total_bytes = _torch.cuda.get_device_properties(0).total_memory
|
||
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
|
||
except Exception:
|
||
return None, None
|
||
|
||
|
||
def get_gpu_utilization() -> Dict[str, Any]:
|
||
"""Return a live snapshot of device utilization information."""
|
||
device = get_device()
|
||
|
||
if device == DeviceType.CUDA:
|
||
result = _smi_query("get_primary_gpu_utilization")
|
||
if result is not None:
|
||
result["backend"] = _backend_label(device)
|
||
if IS_ROCM:
|
||
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.).
|
||
_reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec())
|
||
return result
|
||
# SMI unavailable. On Windows, use Performance Counters (Task Manager
|
||
# source) for system-wide VRAM, covering cross-process usage torch can't see.
|
||
if IS_ROCM and platform.system() == "Windows":
|
||
_win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
|
||
if _win_used is not None and _win_total is not None:
|
||
_win_util = _rocm_windows_perf_counter_gpu_util_pct()
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"gpu_utilization_pct": _win_util,
|
||
"temperature_c": None,
|
||
"vram_used_gb": _win_used,
|
||
"vram_total_gb": _win_total,
|
||
"vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
|
||
if _win_total > 0
|
||
else None,
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
# Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed.
|
||
if IS_ROCM and platform.system() == "Linux":
|
||
_linux_used, _linux_total = _rocm_linux_sysfs_vram_gb()
|
||
if _linux_used is not None and _linux_total is not None:
|
||
_linux_util = _rocm_linux_sysfs_gpu_busy_pct()
|
||
_linux_temp = _rocm_linux_sysfs_temp_c()
|
||
_linux_power = _rocm_linux_sysfs_power_w()
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"gpu_utilization_pct": _linux_util,
|
||
"temperature_c": _linux_temp,
|
||
"vram_used_gb": _linux_used,
|
||
"vram_total_gb": _linux_total,
|
||
"vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1)
|
||
if _linux_total > 0
|
||
else None,
|
||
"power_draw_w": _linux_power,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
# Last resort: torch mem_get_info (process-local).
|
||
_visible_spec = _get_parent_visible_gpu_spec()
|
||
_numeric_ids = _visible_spec.get("numeric_ids") or [0]
|
||
_primary_idx = [_numeric_ids[0]] if _numeric_ids else [0]
|
||
_torch_devices = _torch_get_per_device_info(_primary_idx)
|
||
if _torch_devices:
|
||
_td = _torch_devices[0]
|
||
_total = _td["total_gb"]
|
||
_used = _td["used_gb"]
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"gpu_utilization_pct": None,
|
||
"temperature_c": None,
|
||
"vram_used_gb": _used,
|
||
"vram_total_gb": _total,
|
||
"vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None,
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
|
||
# MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%.
|
||
if device == DeviceType.MLX:
|
||
try:
|
||
import psutil
|
||
agx = _read_apple_gpu_stats()
|
||
total_bytes = psutil.virtual_memory().total
|
||
except Exception as e:
|
||
logger.error(f"Error getting MLX GPU utilization: {e}")
|
||
return {"available": False, "backend": device.value, "error": str(e)}
|
||
if not agx:
|
||
return {"available": False, "backend": device.value}
|
||
allocated_bytes = agx.get("vram_used_bytes", 0) or 0
|
||
vram_used_gb = allocated_bytes / (1024**3)
|
||
total_gb = total_bytes / (1024**3)
|
||
|
||
try:
|
||
from core.training import get_training_backend
|
||
|
||
tb = get_training_backend()
|
||
tb_progress = getattr(tb, "_progress", None)
|
||
if tb_progress is not None and getattr(tb_progress, "is_training", False):
|
||
tb_peak = getattr(tb_progress, "peak_memory_gb", None)
|
||
if tb_peak is not None and tb_peak > 0:
|
||
vram_used_gb = float(tb_peak)
|
||
except Exception:
|
||
pass
|
||
|
||
from . import apple
|
||
|
||
return {
|
||
"available": True,
|
||
"backend": device.value,
|
||
"gpu_utilization_pct": agx.get("utilization_pct") if agx else None,
|
||
"temperature_c": apple.read_gpu_temperature_c(),
|
||
"vram_used_gb": round(vram_used_gb, 2),
|
||
"vram_total_gb": round(total_gb, 2),
|
||
"vram_utilization_pct": (
|
||
round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None
|
||
),
|
||
"power_draw_w": apple.read_gpu_power_w(),
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
|
||
mem = get_gpu_memory_info()
|
||
if device != DeviceType.CPU and mem.get("available"):
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"gpu_utilization_pct": None,
|
||
"temperature_c": None,
|
||
"vram_used_gb": round(mem.get("allocated_gb", 0), 2),
|
||
"vram_total_gb": round(mem.get("total_gb", 0), 2),
|
||
"vram_utilization_pct": round(mem.get("utilization_pct", 0), 1),
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
|
||
return {"available": False, "backend": _backend_label(device)}
|
||
|
||
|
||
def _apply_unified_memory_correction(
|
||
device_metrics: Dict[str, Any], torch_info: Dict[str, Any]
|
||
) -> None:
|
||
"""Per-device reconciliation: when torch reports a larger memory total
|
||
than amd-smi, overwrite the smi VRAM fields in place.
|
||
|
||
Used by both the multi-device and primary-device reconcilers so the two
|
||
endpoints stay in sync on AMD iGPUs with unified memory.
|
||
"""
|
||
torch_total_gb = torch_info["total_gb"]
|
||
smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
|
||
if torch_total_gb > smi_total_gb:
|
||
torch_used_gb = torch_info["used_gb"]
|
||
device_metrics["vram_total_gb"] = torch_total_gb
|
||
device_metrics["vram_used_gb"] = torch_used_gb
|
||
device_metrics["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 %s",
|
||
smi_total_gb,
|
||
torch_total_gb,
|
||
torch_info.get("index"),
|
||
)
|
||
|
||
|
||
def _reconcile_rocm_unified_memory(utilization: Dict[str, Any], device_indices: list[int]) -> None:
|
||
"""Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo).
|
||
|
||
amd-smi reports only the dedicated slice; torch sees the full GTT pool. When
|
||
torch total > smi total, overwrite per-device VRAM fields with the real value.
|
||
"""
|
||
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", []):
|
||
td = torch_by_index.get(dev.get("index"))
|
||
if td is None:
|
||
continue
|
||
_apply_unified_memory_correction(dev, td)
|
||
|
||
|
||
def _reconcile_primary_rocm_unified_memory(
|
||
utilization: Dict[str, Any], parent_visible_spec: Dict[str, Any]
|
||
) -> None:
|
||
"""Same fix as _reconcile_rocm_unified_memory for the flat primary-GPU dict."""
|
||
numeric_ids = parent_visible_spec.get("numeric_ids")
|
||
if numeric_ids is None:
|
||
# No visibility env var set: torch ordinal 0 is the primary device.
|
||
primary_idx = [0]
|
||
elif len(numeric_ids) == 0:
|
||
# Empty mask: no GPU visible. Querying torch device 0 would raise or
|
||
# return stale data, so bail rather than write bad values.
|
||
return
|
||
else:
|
||
primary_idx = [int(numeric_ids[0])]
|
||
torch_devices = _torch_get_per_device_info(primary_idx)
|
||
if not torch_devices:
|
||
return
|
||
_apply_unified_memory_correction(utilization, torch_devices[0])
|
||
|
||
|
||
def get_visible_gpu_utilization() -> Dict[str, Any]:
|
||
device = get_device()
|
||
|
||
if device == DeviceType.CUDA:
|
||
parent_visible_spec = _get_parent_visible_gpu_spec()
|
||
result = _smi_query(
|
||
"get_visible_gpu_utilization",
|
||
parent_visible_spec["numeric_ids"],
|
||
parent_cuda_visible_devices = parent_visible_spec["raw"],
|
||
)
|
||
if result is not None:
|
||
result["backend"] = _backend_label(device)
|
||
numeric_ids = parent_visible_spec.get("numeric_ids")
|
||
if IS_ROCM and numeric_ids is not None:
|
||
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.).
|
||
_reconcile_rocm_unified_memory(result, numeric_ids)
|
||
return result
|
||
|
||
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
|
||
if device in (DeviceType.CUDA, DeviceType.XPU):
|
||
parent_ids = get_parent_visible_gpu_ids()
|
||
# Empty parent_ids (UUID/MIG mask or no CVD): enumerate torch ordinals.
|
||
if parent_ids:
|
||
torch_indices = parent_ids
|
||
index_kind = "physical"
|
||
else:
|
||
visible_count = _torch_get_physical_gpu_count() or 0
|
||
torch_indices = list(range(visible_count))
|
||
index_kind = "relative"
|
||
torch_devices = _torch_get_per_device_info(torch_indices)
|
||
if torch_devices:
|
||
devices = []
|
||
for td in torch_devices:
|
||
total = td["total_gb"]
|
||
used = td["used_gb"]
|
||
devices.append(
|
||
{
|
||
"index": td["index"],
|
||
"index_kind": index_kind,
|
||
"visible_ordinal": td["visible_ordinal"],
|
||
"gpu_utilization_pct": None,
|
||
"temperature_c": None,
|
||
"vram_used_gb": used,
|
||
"vram_total_gb": total,
|
||
"vram_utilization_pct": round((used / total) * 100, 1)
|
||
if total > 0
|
||
else None,
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
)
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"parent_visible_gpu_ids": parent_ids,
|
||
"devices": devices,
|
||
"index_kind": index_kind,
|
||
}
|
||
|
||
if device == DeviceType.MLX:
|
||
mem = get_gpu_memory_info()
|
||
if not mem.get("available"):
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"parent_visible_gpu_ids": [],
|
||
"devices": [],
|
||
"index_kind": "relative",
|
||
}
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"parent_visible_gpu_ids": [0],
|
||
"devices": [
|
||
{
|
||
"index": 0,
|
||
"index_kind": "relative",
|
||
"visible_ordinal": 0,
|
||
"gpu_utilization_pct": None,
|
||
"temperature_c": None,
|
||
"vram_used_gb": round(mem.get("allocated_gb", 0), 2),
|
||
"vram_total_gb": round(mem.get("total_gb", 0), 2),
|
||
"vram_utilization_pct": round(mem.get("utilization_pct", 0), 1),
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
],
|
||
"index_kind": "relative",
|
||
}
|
||
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"parent_visible_gpu_ids": [],
|
||
"devices": [],
|
||
"index_kind": "relative",
|
||
}
|
||
|
||
|
||
# ========== Multi-GPU Detection & Safe num_proc ==========
|
||
|
||
_physical_gpu_count: Optional[int] = None
|
||
_visible_gpu_count: Optional[int] = None
|
||
|
||
|
||
def _get_parent_visible_gpu_spec() -> Dict[str, Any]:
|
||
# ROCm uses HIP/ROCR_VISIBLE_DEVICES on top of CUDA_VISIBLE_DEVICES; check
|
||
# them first. Explicit None checks (not `or`) so "" reads as "no visible GPUs".
|
||
cuda_visible = None
|
||
# Prefer ROCm masks only on a ROCm host or when no CUDA mask is set, so a
|
||
# stale HIP_VISIBLE_DEVICES on NVIDIA can't override CUDA_VISIBLE_DEVICES.
|
||
_is_rocm_spec = IS_ROCM or (
|
||
"CUDA_VISIBLE_DEVICES" not in os.environ
|
||
and ("HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ)
|
||
)
|
||
if _is_rocm_spec:
|
||
hip_vis = os.environ.get("HIP_VISIBLE_DEVICES")
|
||
rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES")
|
||
if hip_vis is not None:
|
||
cuda_visible = hip_vis
|
||
elif rocr_vis is not None:
|
||
cuda_visible = rocr_vis
|
||
if cuda_visible is None:
|
||
cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
|
||
if cuda_visible is None:
|
||
return {
|
||
"raw": None,
|
||
"numeric_ids": list(range(get_physical_gpu_count())),
|
||
"supports_explicit_gpu_ids": True,
|
||
}
|
||
|
||
cuda_visible = cuda_visible.strip()
|
||
if cuda_visible == "" or cuda_visible == "-1":
|
||
return {
|
||
"raw": cuda_visible,
|
||
"numeric_ids": [],
|
||
"supports_explicit_gpu_ids": True,
|
||
}
|
||
|
||
tokens = [value.strip() for value in cuda_visible.split(",") if value.strip()]
|
||
try:
|
||
numeric_ids = [int(value) for value in tokens]
|
||
except ValueError:
|
||
return {
|
||
"raw": cuda_visible,
|
||
"numeric_ids": None,
|
||
"supports_explicit_gpu_ids": False,
|
||
}
|
||
|
||
return {
|
||
"raw": cuda_visible,
|
||
"numeric_ids": numeric_ids,
|
||
"supports_explicit_gpu_ids": True,
|
||
}
|
||
|
||
|
||
def get_parent_visible_gpu_ids() -> list[int]:
|
||
parent_visible_ids = _get_parent_visible_gpu_spec()["numeric_ids"]
|
||
return list(parent_visible_ids) if parent_visible_ids is not None else []
|
||
|
||
|
||
def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]:
|
||
parent_visible_spec = _get_parent_visible_gpu_spec()
|
||
parent_visible_ids = get_parent_visible_gpu_ids()
|
||
physical_gpu_count = get_physical_gpu_count()
|
||
|
||
if gpu_ids is None:
|
||
return parent_visible_ids
|
||
|
||
requested_ids = list(gpu_ids)
|
||
if len(requested_ids) == 0:
|
||
return parent_visible_ids
|
||
|
||
if not parent_visible_spec["supports_explicit_gpu_ids"]:
|
||
raise ValueError(
|
||
f"Invalid gpu_ids {requested_ids}: explicit physical GPU IDs are "
|
||
f"unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries "
|
||
f"({parent_visible_spec['raw']!r}). Omit gpu_ids to use the "
|
||
"parent-visible devices."
|
||
)
|
||
|
||
if len(set(requested_ids)) != len(requested_ids):
|
||
raise ValueError(
|
||
f"Invalid gpu_ids {requested_ids}: duplicate GPU IDs are not allowed. "
|
||
f"Parent-visible GPUs: {parent_visible_ids}"
|
||
)
|
||
|
||
# Reject negative IDs.
|
||
negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0]
|
||
if negative_ids:
|
||
raise ValueError(
|
||
f"Invalid gpu_ids {requested_ids}: GPU IDs must be non-negative. "
|
||
f"Rejected IDs: {negative_ids}. Parent-visible GPUs: {parent_visible_ids}"
|
||
)
|
||
|
||
# Only enforce the physical upper bound when the count is reliable (nvidia-smi).
|
||
# A torch count reflects only visible devices, so it could falsely reject valid
|
||
# physical indices. The parent-visible check below is always authoritative.
|
||
if physical_gpu_count > 0 and parent_visible_ids:
|
||
max_parent_id = max(parent_visible_ids)
|
||
if physical_gpu_count > max_parent_id:
|
||
# Count is plausibly physical, so enforce it.
|
||
out_of_range = [gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count]
|
||
if out_of_range:
|
||
raise ValueError(
|
||
f"Invalid gpu_ids {requested_ids}: IDs must be physical GPU IDs "
|
||
f"between 0 and {physical_gpu_count - 1}. "
|
||
f"Rejected IDs: {out_of_range}. Parent-visible GPUs: {parent_visible_ids}"
|
||
)
|
||
|
||
disallowed_ids = [gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids]
|
||
if disallowed_ids:
|
||
raise ValueError(
|
||
f"Invalid gpu_ids {requested_ids}: requested GPUs {disallowed_ids} are "
|
||
f"outside the parent-visible set {parent_visible_ids}"
|
||
)
|
||
|
||
return requested_ids
|
||
|
||
|
||
def _resolve_model_identifier_for_gpu_estimate(
|
||
model_name: str, hf_token: Optional[str] = None
|
||
) -> str:
|
||
try:
|
||
from utils.models.model_config import ModelConfig
|
||
|
||
config = ModelConfig.from_identifier(model_name, hf_token = hf_token)
|
||
if config and config.is_lora and config.base_model:
|
||
return config.base_model
|
||
return config.identifier if config else model_name
|
||
except Exception as e:
|
||
logger.debug("Could not resolve base model for GPU estimate '%s': %s", model_name, e)
|
||
return model_name
|
||
|
||
|
||
def _get_local_weight_size_bytes(model_name: str) -> Optional[int]:
|
||
model_path = Path(model_name)
|
||
if not model_path.exists():
|
||
return None
|
||
|
||
weight_exts = (".safetensors", ".bin", ".pt", ".pth")
|
||
total = 0
|
||
for file in model_path.rglob("*"):
|
||
if file.is_file() and file.suffix in weight_exts:
|
||
total += file.stat().st_size
|
||
return total if total > 0 else None
|
||
|
||
|
||
def _get_hf_safetensors_total_params(
|
||
model_name: str, hf_token: Optional[str] = None
|
||
) -> Optional[int]:
|
||
try:
|
||
from huggingface_hub import model_info as hf_model_info
|
||
|
||
info = hf_model_info(model_name, token = hf_token)
|
||
safetensors = getattr(info, "safetensors", None)
|
||
if isinstance(safetensors, dict):
|
||
total = safetensors.get("total")
|
||
if total:
|
||
return int(total)
|
||
except Exception as e:
|
||
logger.warning("Could not get safetensors metadata for '%s': %s", model_name, e)
|
||
return None
|
||
|
||
|
||
def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None):
|
||
try:
|
||
from transformers import AutoConfig
|
||
trust_remote_code = model_name.lower().startswith("unsloth/")
|
||
return AutoConfig.from_pretrained(
|
||
model_name,
|
||
token = hf_token,
|
||
trust_remote_code = trust_remote_code,
|
||
)
|
||
except Exception as e:
|
||
logger.warning("Could not load config for '%s': %s", model_name, e)
|
||
return None
|
||
|
||
|
||
def _determine_attention_impl_for_gpu_estimate(config) -> str:
|
||
# torch.distributed is incomplete on Windows ROCm (torch._C._distributed_c10d
|
||
# can't be imported). Inject stubs into sys.modules before importing
|
||
# torch.distributed, then patch the missing process-group helpers.
|
||
if sys.platform == "win32" and IS_ROCM:
|
||
# Dummy for any name torch.distributed imports from these stubs.
|
||
class _Dummy:
|
||
pass
|
||
|
||
for _c10d_name in (
|
||
"torch._C._distributed_c10d",
|
||
"torch._C._distributed_autograd",
|
||
"torch._C._distributed_rpc",
|
||
):
|
||
if _c10d_name not in sys.modules:
|
||
_stub = types.ModuleType(_c10d_name)
|
||
# No-op dummies for names torch.distributed imports from _distributed_c10d.
|
||
for _sym in (
|
||
"FakeProcessGroup",
|
||
"ProcessGroup",
|
||
"Work",
|
||
"Store",
|
||
"PrefixStore",
|
||
"FileStore",
|
||
"TCPStore",
|
||
"HashStore",
|
||
"Reducer",
|
||
"Logger",
|
||
"DistributedDebugLevel",
|
||
"GradBucket",
|
||
"BuiltinCommHookType",
|
||
):
|
||
setattr(_stub, _sym, _Dummy)
|
||
sys.modules[_c10d_name] = _stub
|
||
|
||
try:
|
||
import torch.distributed as _td
|
||
for _attr, _stub in (
|
||
("is_initialized", lambda: False),
|
||
("is_available", lambda: False),
|
||
("get_rank", lambda: 0),
|
||
("get_world_size", lambda: 1),
|
||
("is_torchelastic_launched", lambda: False),
|
||
):
|
||
if not hasattr(_td, _attr):
|
||
setattr(_td, _attr, _stub)
|
||
except ImportError:
|
||
pass
|
||
|
||
from unsloth.models._utils import resolve_attention_implementation
|
||
from transformers import AutoModel, AutoModelForCausalLM
|
||
|
||
# why: resolve_attention_implementation writes _attn_implementation onto the
|
||
# config and propagates to nested sub-configs; a shallow copy would still
|
||
# mutate the cached config's shared inner objects. Deepcopy isolates them.
|
||
config_copy = copy.deepcopy(config)
|
||
|
||
model_class = None
|
||
for auto_model in (AutoModelForCausalLM, AutoModel):
|
||
mapping = getattr(auto_model, "_model_mapping", None)
|
||
if mapping is None:
|
||
continue
|
||
try:
|
||
if config_copy.__class__ in mapping:
|
||
model_class = mapping[config_copy.__class__]
|
||
break
|
||
except Exception:
|
||
continue
|
||
|
||
return resolve_attention_implementation(model_class, config_copy)
|
||
|
||
|
||
def _estimate_fp16_model_size_bytes_from_config(config) -> Optional[int]:
|
||
from .vram_estimation import extract_arch_config, compute_total_params
|
||
|
||
arch = extract_arch_config(config)
|
||
if arch is None:
|
||
return None
|
||
return compute_total_params(arch) * 2
|
||
|
||
|
||
def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]:
|
||
if config is None:
|
||
return None
|
||
|
||
previous_unsloth_present = os.environ.get("UNSLOTH_IS_PRESENT")
|
||
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
||
try:
|
||
from unsloth_zoo import vllm_utils as _vllm_utils
|
||
|
||
synthetic_total_bytes = 1024 * (1024**3)
|
||
original_get_mem_info = _vllm_utils.get_mem_info
|
||
try:
|
||
_vllm_utils.get_mem_info = lambda: (
|
||
synthetic_total_bytes,
|
||
synthetic_total_bytes,
|
||
)
|
||
_, _, _, memory_left_for_kv_cache_gb = _vllm_utils.approximate_vllm_memory_usage(
|
||
config,
|
||
load_in_4bit = False,
|
||
load_in_8bit = False,
|
||
max_seq_length = 1,
|
||
gpu_memory_utilization = 1.0,
|
||
enable_lora = False,
|
||
account_for_gradients = False,
|
||
cuda_graph_overhead = False,
|
||
)
|
||
finally:
|
||
_vllm_utils.get_mem_info = original_get_mem_info
|
||
except Exception as e:
|
||
logger.debug("Could not estimate model size via vllm_utils: %s", e)
|
||
return None
|
||
finally:
|
||
if previous_unsloth_present is None:
|
||
os.environ.pop("UNSLOTH_IS_PRESENT", None)
|
||
else:
|
||
os.environ["UNSLOTH_IS_PRESENT"] = previous_unsloth_present
|
||
|
||
model_size_gb = 1024.0 - memory_left_for_kv_cache_gb
|
||
if model_size_gb <= 0:
|
||
return None
|
||
return int(round(model_size_gb * (1024**3)))
|
||
|
||
|
||
def estimate_fp16_model_size_bytes(
|
||
model_name: str, hf_token: Optional[str] = None
|
||
) -> tuple[Optional[int], str]:
|
||
estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token)
|
||
|
||
total_params = None
|
||
if "/" in estimate_model and not Path(estimate_model).exists():
|
||
total_params = _get_hf_safetensors_total_params(estimate_model, hf_token = hf_token)
|
||
if total_params:
|
||
return int(total_params * 2), "safetensors"
|
||
|
||
config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token)
|
||
config_bytes: Optional[int] = None
|
||
if config is not None:
|
||
config_bytes = _estimate_fp16_model_size_bytes_from_config(config)
|
||
|
||
local_bytes = _get_local_weight_size_bytes(estimate_model)
|
||
|
||
# why: config-derived bytes cover only the text tower; local safetensors
|
||
# include vision/audio towers. Take the larger so the multimodal
|
||
# extra_bytes correction can fire.
|
||
if config_bytes is not None and local_bytes is not None:
|
||
if local_bytes > config_bytes:
|
||
return local_bytes, "weight_bytes"
|
||
return config_bytes, "config"
|
||
if config_bytes is not None:
|
||
return config_bytes, "config"
|
||
if local_bytes is not None:
|
||
return local_bytes, "weight_bytes"
|
||
|
||
vllm_bytes = _estimate_fp16_model_size_bytes_from_vllm_utils(config)
|
||
if vllm_bytes is not None:
|
||
return vllm_bytes, "vllm_utils"
|
||
|
||
return None, "unavailable"
|
||
|
||
|
||
def estimate_required_model_memory_gb(
|
||
model_name: str,
|
||
*,
|
||
hf_token: Optional[str] = None,
|
||
training_type: Optional[str] = None,
|
||
load_in_4bit: bool = True,
|
||
batch_size: int = 4,
|
||
max_seq_length: int = 2048,
|
||
lora_rank: int = 16,
|
||
target_modules: Optional[list] = None,
|
||
gradient_checkpointing: str = "unsloth",
|
||
optimizer: str = "adamw_8bit",
|
||
) -> tuple[Optional[float], Dict[str, Any]]:
|
||
from .vram_estimation import (
|
||
TrainingVramConfig,
|
||
extract_arch_config,
|
||
estimate_training_vram,
|
||
compute_total_params,
|
||
compute_optimizer_bytes,
|
||
compute_gradient_bytes,
|
||
CUDA_OVERHEAD_BYTES,
|
||
QUANT_4BIT_FACTOR,
|
||
DEFAULT_TARGET_MODULES,
|
||
)
|
||
|
||
model_size_bytes, source = estimate_fp16_model_size_bytes(model_name, hf_token = hf_token)
|
||
metadata: Dict[str, Any] = {
|
||
"mode": "inference" if training_type is None else "training",
|
||
"model_size_source": source,
|
||
}
|
||
if model_size_bytes is None:
|
||
metadata["required_gb"] = None
|
||
return None, metadata
|
||
|
||
model_size_gb = model_size_bytes / (1024**3)
|
||
metadata["model_size_gb"] = round(model_size_gb, 3)
|
||
min_buffer_gb = 2.0
|
||
|
||
if training_type is None:
|
||
if load_in_4bit:
|
||
base_4bit_gb = model_size_gb / QUANT_4BIT_FACTOR
|
||
required_gb = base_4bit_gb + max(base_4bit_gb * 0.3, min_buffer_gb)
|
||
else:
|
||
required_gb = model_size_gb * 1.3
|
||
metadata["required_gb"] = round(required_gb, 3)
|
||
return required_gb, metadata
|
||
|
||
training_method = (
|
||
"full" if training_type == "Full Finetuning" else ("qlora" if load_in_4bit else "lora")
|
||
)
|
||
vram_config = TrainingVramConfig(
|
||
training_method = training_method,
|
||
batch_size = batch_size,
|
||
max_seq_length = max_seq_length,
|
||
lora_rank = lora_rank,
|
||
target_modules = target_modules or list(DEFAULT_TARGET_MODULES),
|
||
gradient_checkpointing = gradient_checkpointing,
|
||
optimizer = optimizer,
|
||
load_in_4bit = load_in_4bit,
|
||
)
|
||
|
||
estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token)
|
||
config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token)
|
||
if config is not None:
|
||
try:
|
||
vram_config.attention_implementation = _determine_attention_impl_for_gpu_estimate(
|
||
config
|
||
)
|
||
except Exception as e:
|
||
# Debug-level: fires every estimate on Windows ROCm (stub lacks Store);
|
||
# expected and non-actionable -- eager is the safe fallback.
|
||
logger.debug(
|
||
"Could not resolve attention implementation for '%s': %s",
|
||
estimate_model,
|
||
e,
|
||
)
|
||
# why: charge the quadratic non-flash activation path so GPU
|
||
# selection stays conservative when flash attn isn't proven usable.
|
||
vram_config.attention_implementation = "eager"
|
||
arch = extract_arch_config(config) if config is not None else None
|
||
|
||
if arch is not None:
|
||
breakdown = estimate_training_vram(arch, vram_config)
|
||
# why: extract_arch_config only sees text_config; add the vision/audio
|
||
# tower bytes that the text-arch fp16 total misses.
|
||
arch_fp16_bytes = compute_total_params(arch) * 2
|
||
extra_bytes = max(0, int(model_size_bytes) - arch_fp16_bytes)
|
||
if extra_bytes > 0:
|
||
breakdown.model_weights += extra_bytes
|
||
if training_method == "full":
|
||
# why: full fine-tuning makes extra params trainable; optimizer +
|
||
# gradient bytes scale with them.
|
||
extra_params = extra_bytes // 2
|
||
breakdown.optimizer_states += compute_optimizer_bytes(
|
||
extra_params,
|
||
vram_config.optimizer,
|
||
)
|
||
breakdown.gradients += compute_gradient_bytes(extra_params)
|
||
required_gb = breakdown.total / (1024**3)
|
||
metadata["required_gb"] = round(required_gb, 3)
|
||
metadata["estimation_mode"] = "detailed"
|
||
metadata["attention_implementation"] = vram_config.attention_implementation
|
||
metadata["vram_breakdown"] = breakdown.to_gb_dict()
|
||
max_gpus = max(1, get_visible_gpu_count())
|
||
for n_gpus in range(1, max_gpus + 1):
|
||
metadata["vram_breakdown"][f"min_per_gpu_{n_gpus}"] = round(
|
||
breakdown.min_gpu_vram(n_gpus) / (1024**3), 3
|
||
)
|
||
return required_gb, metadata
|
||
|
||
# Fallback when model config is unavailable.
|
||
overhead_gb = CUDA_OVERHEAD_BYTES / (1024**3)
|
||
if training_method == "full":
|
||
required_gb = model_size_gb * 3.5 + overhead_gb
|
||
elif training_method == "qlora":
|
||
base_4bit_gb = model_size_gb / QUANT_4BIT_FACTOR
|
||
lora_overhead_gb = model_size_gb * 0.04
|
||
act_gb = model_size_gb * 0.15 * (batch_size / 4) * (max_seq_length / 2048)
|
||
required_gb = base_4bit_gb + lora_overhead_gb + act_gb + overhead_gb
|
||
else:
|
||
lora_overhead_gb = model_size_gb * 0.04
|
||
act_gb = model_size_gb * 0.15 * (batch_size / 4) * (max_seq_length / 2048)
|
||
required_gb = model_size_gb + lora_overhead_gb + act_gb + overhead_gb
|
||
|
||
metadata["required_gb"] = round(required_gb, 3)
|
||
metadata["estimation_mode"] = "fallback"
|
||
return required_gb, metadata
|
||
|
||
|
||
def auto_select_gpu_ids(
|
||
model_name: str,
|
||
*,
|
||
hf_token: Optional[str] = None,
|
||
training_type: Optional[str] = None,
|
||
load_in_4bit: bool = True,
|
||
batch_size: int = 4,
|
||
max_seq_length: int = 2048,
|
||
lora_rank: int = 16,
|
||
target_modules: Optional[list] = None,
|
||
gradient_checkpointing: str = "unsloth",
|
||
optimizer: str = "adamw_8bit",
|
||
) -> tuple[Optional[list[int]], Dict[str, Any]]:
|
||
metadata: Dict[str, Any] = {"selection_mode": "auto"}
|
||
|
||
if get_device() != DeviceType.CUDA:
|
||
metadata["selection_mode"] = "non_cuda"
|
||
return None, metadata
|
||
|
||
required_gb, estimate_metadata = estimate_required_model_memory_gb(
|
||
model_name,
|
||
hf_token = hf_token,
|
||
training_type = training_type,
|
||
load_in_4bit = load_in_4bit,
|
||
batch_size = batch_size,
|
||
max_seq_length = max_seq_length,
|
||
lora_rank = lora_rank,
|
||
target_modules = target_modules,
|
||
gradient_checkpointing = gradient_checkpointing,
|
||
optimizer = optimizer,
|
||
)
|
||
metadata.update(estimate_metadata)
|
||
parent_visible_spec = _get_parent_visible_gpu_spec()
|
||
metadata["parent_cuda_visible_devices"] = parent_visible_spec["raw"]
|
||
|
||
if not parent_visible_spec["supports_explicit_gpu_ids"]:
|
||
metadata["selection_mode"] = "inherit_parent_visible"
|
||
metadata["selected_gpu_ids"] = None
|
||
return None, metadata
|
||
|
||
if required_gb is None:
|
||
# Can't estimate size -- use all visible GPUs rather than risk one too small.
|
||
parent_ids = get_parent_visible_gpu_ids()
|
||
metadata["selection_mode"] = "fallback_all"
|
||
metadata["selected_gpu_ids"] = parent_ids
|
||
return parent_ids, metadata
|
||
|
||
utilization = get_visible_gpu_utilization()
|
||
devices = utilization.get("devices", [])
|
||
parent_ids = get_parent_visible_gpu_ids()
|
||
|
||
if not devices:
|
||
metadata["selection_mode"] = "fallback_all"
|
||
metadata["selected_gpu_ids"] = parent_ids
|
||
return parent_ids, metadata
|
||
|
||
gpu_candidates = []
|
||
for device in devices:
|
||
total_gb = device.get("vram_total_gb")
|
||
used_gb = device.get("vram_used_gb")
|
||
if total_gb is None or used_gb is None:
|
||
continue
|
||
free_gb = max(total_gb - used_gb, 0.0)
|
||
gpu_candidates.append(
|
||
{
|
||
"index": device["index"],
|
||
"free_gb": free_gb,
|
||
}
|
||
)
|
||
|
||
if not gpu_candidates:
|
||
metadata["selection_mode"] = "fallback_all"
|
||
metadata["selected_gpu_ids"] = parent_ids
|
||
return parent_ids, metadata
|
||
|
||
ranked = sorted(gpu_candidates, key = lambda item: (-item["free_gb"], item["index"]))
|
||
free_by_index = {item["index"]: item["free_gb"] for item in ranked}
|
||
selected: list[int] = []
|
||
usable_gb = 0.0
|
||
# Sharding has inter-GPU overhead, so each extra GPU contributes less than
|
||
# its raw free memory (first GPU keeps full capacity). 0.85 is empirical on
|
||
# 2-8 GPU setups: covers NCCL buffers, pipeline bubbles, fragmentation.
|
||
multi_gpu_overhead = 0.85
|
||
|
||
# Per-GPU check: activations don't shard, so each GPU needs its weight shard
|
||
# + full activation cost. Uses precomputed min_per_gpu_N values.
|
||
vram_breakdown = estimate_metadata.get("vram_breakdown", {})
|
||
|
||
for candidate in ranked:
|
||
selected.append(candidate["index"])
|
||
if len(selected) == 1:
|
||
usable_gb = candidate["free_gb"]
|
||
else:
|
||
first_gpu_id = selected[0]
|
||
usable_gb = free_by_index[first_gpu_id] + sum(
|
||
free_by_index[gpu_id] * multi_gpu_overhead for gpu_id in selected[1:]
|
||
)
|
||
|
||
total_fits = usable_gb >= required_gb
|
||
|
||
per_gpu_fits = True
|
||
if total_fits and len(selected) > 1:
|
||
min_key = f"min_per_gpu_{len(selected)}"
|
||
min_per_gpu_gb = vram_breakdown.get(min_key)
|
||
if min_per_gpu_gb is not None:
|
||
smallest_free = min(free_by_index[gpu_id] for gpu_id in selected)
|
||
per_gpu_fits = smallest_free >= min_per_gpu_gb
|
||
|
||
if total_fits and per_gpu_fits:
|
||
metadata["usable_gb"] = round(usable_gb, 3)
|
||
metadata["selection_mode"] = "auto"
|
||
metadata["selected_gpu_ids"] = selected
|
||
logger.debug(
|
||
"Selected GPUs automatically",
|
||
model_name = model_name,
|
||
selected_gpu_ids = selected,
|
||
usable_gb = metadata["usable_gb"],
|
||
required_gb = metadata.get("required_gb"),
|
||
multi_gpu_overhead = multi_gpu_overhead,
|
||
)
|
||
return selected, metadata
|
||
|
||
# Use only GPUs with verified VRAM data.
|
||
fallback_all = [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids
|
||
metadata["selection_mode"] = "fallback_all"
|
||
if ranked:
|
||
fallback_usable = ranked[0]["free_gb"] + sum(
|
||
c["free_gb"] * multi_gpu_overhead for c in ranked[1:]
|
||
)
|
||
else:
|
||
fallback_usable = 0.0
|
||
metadata["usable_gb"] = round(fallback_usable, 3)
|
||
metadata["selected_gpu_ids"] = fallback_all
|
||
logger.warning(
|
||
"Falling back to all visible GPUs -- model may not fit",
|
||
model_name = model_name,
|
||
selected_gpu_ids = fallback_all,
|
||
usable_gb = metadata["usable_gb"],
|
||
required_gb = metadata.get("required_gb"),
|
||
multi_gpu_overhead = multi_gpu_overhead,
|
||
)
|
||
return fallback_all, metadata
|
||
|
||
|
||
def prepare_gpu_selection(
|
||
gpu_ids: Optional[list[int]],
|
||
*,
|
||
model_name: str,
|
||
hf_token: Optional[str] = None,
|
||
training_type: Optional[str] = None,
|
||
load_in_4bit: bool = True,
|
||
batch_size: int = 4,
|
||
max_seq_length: int = 2048,
|
||
lora_rank: int = 16,
|
||
target_modules: Optional[list] = None,
|
||
gradient_checkpointing: str = "unsloth",
|
||
optimizer: str = "adamw_8bit",
|
||
) -> tuple[Optional[list[int]], Dict[str, Any]]:
|
||
"""Resolve which physical GPUs to use for a model load.
|
||
|
||
GPU selection modes:
|
||
- **Explicit** (``gpu_ids=[5, 6, 7]``): caller chooses exact GPUs.
|
||
All listed GPUs are used and the model is sharded via
|
||
``device_map="balanced"``, even if it would fit on fewer. IDs are
|
||
validated against the parent-visible set.
|
||
- **Auto** (``gpu_ids=None`` or ``[]``): ``auto_select_gpu_ids``
|
||
estimates VRAM needs and picks the *minimum* GPUs needed,
|
||
preferring those with the most free memory.
|
||
|
||
The returned ``gpu_ids`` is later passed to ``get_device_map()`` (maps it
|
||
to a Hugging Face ``device_map`` string) and to ``apply_gpu_ids()`` in the
|
||
worker subprocess (narrows ``CUDA_VISIBLE_DEVICES`` before torch/CUDA init).
|
||
"""
|
||
if gpu_ids and get_device() != DeviceType.CUDA:
|
||
raise ValueError(
|
||
f"gpu_ids {list(gpu_ids)} is only supported on CUDA devices, "
|
||
f"but the current backend is '{get_device().value}'."
|
||
)
|
||
|
||
if gpu_ids:
|
||
resolved = resolve_requested_gpu_ids(gpu_ids)
|
||
metadata = {
|
||
"selection_mode": "explicit",
|
||
"selected_gpu_ids": resolved,
|
||
}
|
||
return resolved, metadata
|
||
|
||
selected_gpu_ids, metadata = auto_select_gpu_ids(
|
||
model_name,
|
||
hf_token = hf_token,
|
||
training_type = training_type,
|
||
load_in_4bit = load_in_4bit,
|
||
batch_size = batch_size,
|
||
max_seq_length = max_seq_length,
|
||
lora_rank = lora_rank,
|
||
target_modules = target_modules,
|
||
gradient_checkpointing = gradient_checkpointing,
|
||
optimizer = optimizer,
|
||
)
|
||
return selected_gpu_ids, metadata
|
||
|
||
|
||
def get_physical_gpu_count() -> int:
|
||
"""
|
||
Return the number of physical GPUs on the machine.
|
||
|
||
Uses ``nvidia-smi -L`` on NVIDIA (unaffected by CUDA_VISIBLE_DEVICES),
|
||
with a torch fallback for AMD ROCm and Intel XPU. Cached after first call.
|
||
"""
|
||
global _physical_gpu_count
|
||
if _physical_gpu_count is not None:
|
||
return _physical_gpu_count
|
||
|
||
device = get_device()
|
||
|
||
if device == DeviceType.CUDA:
|
||
try:
|
||
if IS_ROCM:
|
||
from . import amd as _smi_mod
|
||
else:
|
||
from . import nvidia as _smi_mod
|
||
count = _smi_mod.get_physical_gpu_count()
|
||
if count is not None:
|
||
_physical_gpu_count = count
|
||
return _physical_gpu_count
|
||
except Exception:
|
||
pass
|
||
# SMI unavailable -- fall back to torch.
|
||
count = _torch_get_physical_gpu_count()
|
||
_physical_gpu_count = count if count is not None else 1
|
||
return _physical_gpu_count
|
||
|
||
if device == DeviceType.XPU:
|
||
count = _torch_get_physical_gpu_count()
|
||
_physical_gpu_count = count if count is not None else 1
|
||
return _physical_gpu_count
|
||
|
||
if device == DeviceType.MLX:
|
||
_physical_gpu_count = 1
|
||
return _physical_gpu_count
|
||
|
||
_physical_gpu_count = 0
|
||
|
||
return _physical_gpu_count
|
||
|
||
|
||
def _backend_visible_devices_env() -> Optional[str]:
|
||
"""Return the raw visibility env string that applies to this backend.
|
||
|
||
On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over
|
||
CUDA_VISIBLE_DEVICES; this mirrors ``_get_parent_visible_gpu_spec`` so
|
||
``backend_cuda_visible_devices`` reports the value actually narrowing the
|
||
visible device set.
|
||
"""
|
||
if IS_ROCM:
|
||
return _get_parent_visible_gpu_spec().get("raw")
|
||
return os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
|
||
|
||
def get_backend_visible_gpu_info() -> Dict[str, Any]:
|
||
device = get_device()
|
||
if device in (DeviceType.CUDA, DeviceType.XPU):
|
||
parent_visible_ids = get_parent_visible_gpu_ids()
|
||
# Try native SMI first (nvidia-smi; skipped for ROCm).
|
||
if device == DeviceType.CUDA and not IS_ROCM:
|
||
try:
|
||
from . import nvidia
|
||
|
||
parent_visible_spec = _get_parent_visible_gpu_spec()
|
||
result = nvidia.get_backend_visible_gpu_info(
|
||
parent_visible_spec["numeric_ids"],
|
||
parent_visible_spec["raw"],
|
||
)
|
||
if result.get("available"):
|
||
result["backend"] = _backend_label(device)
|
||
return result
|
||
except Exception as e:
|
||
logger.warning("Backend GPU visibility query failed: %s", e)
|
||
|
||
# Torch fallback (ROCm, XPU, nvidia-smi missing). Empty parent_visible_ids
|
||
# (UUID/MIG mask) -> enumerate by torch ordinal so the UI shows devices.
|
||
if parent_visible_ids:
|
||
torch_indices = parent_visible_ids
|
||
index_kind = "physical"
|
||
else:
|
||
visible_count = _torch_get_physical_gpu_count() or 0
|
||
torch_indices = list(range(visible_count))
|
||
index_kind = "relative"
|
||
torch_devices = _torch_get_per_device_info(torch_indices)
|
||
if torch_devices:
|
||
devices = [
|
||
{
|
||
"index": td["index"],
|
||
"index_kind": index_kind,
|
||
"visible_ordinal": td["visible_ordinal"],
|
||
"name": td["name"],
|
||
"memory_total_gb": td["total_gb"],
|
||
}
|
||
for td in torch_devices
|
||
]
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"backend_cuda_visible_devices": _backend_visible_devices_env(),
|
||
"parent_visible_gpu_ids": parent_visible_ids,
|
||
"devices": devices,
|
||
"index_kind": index_kind,
|
||
}
|
||
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"backend_cuda_visible_devices": _backend_visible_devices_env(),
|
||
"parent_visible_gpu_ids": parent_visible_ids,
|
||
"devices": [],
|
||
"index_kind": "physical",
|
||
}
|
||
|
||
if device == DeviceType.MLX:
|
||
mem = get_gpu_memory_info()
|
||
if not mem.get("available"):
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
|
||
"parent_visible_gpu_ids": [],
|
||
"devices": [],
|
||
"index_kind": "relative",
|
||
}
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
|
||
"parent_visible_gpu_ids": [0],
|
||
"devices": [
|
||
{
|
||
"index": 0,
|
||
"index_kind": "relative",
|
||
"visible_ordinal": 0,
|
||
"name": mem.get("device_name", "MLX"),
|
||
"memory_total_gb": round(mem.get("total_gb", 0), 2),
|
||
}
|
||
],
|
||
"index_kind": "relative",
|
||
}
|
||
|
||
return {
|
||
"available": False,
|
||
"backend": _backend_label(device),
|
||
"backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
|
||
"parent_visible_gpu_ids": [],
|
||
"devices": [],
|
||
"index_kind": "relative",
|
||
}
|
||
|
||
|
||
def get_visible_gpu_count() -> int:
|
||
"""
|
||
Return the number of GPUs visible to this process.
|
||
|
||
Respects ``CUDA_VISIBLE_DEVICES`` -- if set, only those GPUs count.
|
||
Falls back to physical count if unset or torch is unavailable.
|
||
Cached after the first call.
|
||
"""
|
||
global _visible_gpu_count
|
||
if _visible_gpu_count is not None:
|
||
return _visible_gpu_count
|
||
|
||
# _get_parent_visible_gpu_spec() already handles HIP_VISIBLE_DEVICES /
|
||
# ROCR_VISIBLE_DEVICES on ROCm.
|
||
visible_spec = _get_parent_visible_gpu_spec()
|
||
if visible_spec["raw"] is not None:
|
||
raw = visible_spec["raw"].strip()
|
||
if raw == "" or raw == "-1":
|
||
_visible_gpu_count = 0
|
||
elif visible_spec["numeric_ids"] is not None:
|
||
_visible_gpu_count = len(visible_spec["numeric_ids"])
|
||
else:
|
||
_visible_gpu_count = len([x for x in raw.split(",") if x.strip()])
|
||
return _visible_gpu_count
|
||
|
||
# No visibility env var set -- try torch, else physical count
|
||
try:
|
||
import torch
|
||
if get_device() == DeviceType.XPU and hasattr(torch, "xpu"):
|
||
_visible_gpu_count = torch.xpu.device_count()
|
||
else:
|
||
_visible_gpu_count = torch.cuda.device_count()
|
||
except Exception:
|
||
_visible_gpu_count = get_physical_gpu_count()
|
||
|
||
return _visible_gpu_count
|
||
|
||
|
||
def apply_gpu_ids(gpu_ids) -> None:
|
||
if gpu_ids is None:
|
||
return
|
||
|
||
# Empty list -> treat like None (inherit parent); setting CUDA_VISIBLE_DEVICES=""
|
||
# disables CUDA entirely and crashes downstream torch calls.
|
||
if isinstance(gpu_ids, (list, tuple)) and len(gpu_ids) == 0:
|
||
return
|
||
|
||
global _visible_gpu_count
|
||
|
||
if isinstance(gpu_ids, (list, tuple)):
|
||
value = ",".join(str(g) for g in gpu_ids)
|
||
else:
|
||
value = str(gpu_ids)
|
||
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = value
|
||
# Keep ROCm visibility env vars in sync. Workers may call apply_gpu_ids()
|
||
# before detect_hardware() (IS_ROCM still False), so also mirror when the
|
||
# parent set a ROCm visibility var, with a torch.version.hip probe fallback.
|
||
_inherits_rocm_visibility = (
|
||
"HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
|
||
)
|
||
_is_rocm = IS_ROCM or _inherits_rocm_visibility
|
||
if not _is_rocm:
|
||
# torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may leave
|
||
# it unset but encode "rocm" in __version__. Broad except: never crash a worker.
|
||
try:
|
||
import torch as _torch
|
||
_is_rocm = (
|
||
getattr(_torch.version, "hip", None) is not None
|
||
or "rocm" in getattr(_torch, "__version__", "").lower()
|
||
)
|
||
except Exception as e:
|
||
logger.debug(
|
||
"apply_gpu_ids: torch ROCm probe skipped (%s: %s)",
|
||
type(e).__name__,
|
||
e,
|
||
)
|
||
if _is_rocm:
|
||
os.environ["HIP_VISIBLE_DEVICES"] = value
|
||
# ROCR_VISIBLE_DEVICES operates at the HSA agent level and uses
|
||
# different indexing semantics to HIP_VISIBLE_DEVICES. Setting it
|
||
# to a physical GPU index breaks multi-GPU ROCm systems where the
|
||
# parent already set ROCR_VISIBLE_DEVICES (e.g. "0,1"): narrowing
|
||
# to "1" causes torch.cuda.is_available() to return False in the
|
||
# worker subprocess. HIP_VISIBLE_DEVICES is sufficient for GPU
|
||
# selection on ROCm -- leave ROCR_VISIBLE_DEVICES inherited.
|
||
_visible_gpu_count = None
|
||
if _is_rocm:
|
||
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value)
|
||
else:
|
||
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value)
|
||
|
||
|
||
def get_device_map(gpu_ids: Optional[list[int]] = None) -> str:
|
||
"""Return the Hugging Face ``device_map`` string for model loading.
|
||
|
||
Returns ``"balanced"`` (shard evenly across GPUs) when:
|
||
- ``gpu_ids`` explicitly lists >1 GPU, **or**
|
||
- ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and
|
||
>1 GPU is visible (fallback: numeric IDs unresolvable, so assume
|
||
multi-GPU is intended).
|
||
|
||
Returns ``"sequential"`` (single device) otherwise, including non-CUDA
|
||
backends (CPU, MLX).
|
||
|
||
Use ``prepare_gpu_selection()`` upstream to determine ``gpu_ids`` -- it
|
||
handles auto-selecting the minimum GPUs needed for a model.
|
||
"""
|
||
device = get_device()
|
||
if device == DeviceType.CUDA:
|
||
multi_gpu = gpu_ids is not None and len(gpu_ids) > 1
|
||
|
||
if not multi_gpu:
|
||
# UUID/MIG masks can't be split into numeric IDs; >1 visible GPU
|
||
# means multi-GPU sharding is intended.
|
||
parent_visible_spec = _get_parent_visible_gpu_spec()
|
||
if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1:
|
||
multi_gpu = True
|
||
|
||
if multi_gpu:
|
||
return "balanced"
|
||
|
||
return "sequential"
|
||
|
||
|
||
def get_offloaded_device_map_entries(model) -> dict[str, str]:
|
||
hf_device_map = getattr(model, "hf_device_map", None)
|
||
if not isinstance(hf_device_map, dict):
|
||
return {}
|
||
return {
|
||
module_name: placement
|
||
for module_name, placement in hf_device_map.items()
|
||
if placement in ("cpu", "disk")
|
||
}
|
||
|
||
|
||
def raise_if_offloaded(
|
||
model,
|
||
device_map: str,
|
||
context: str = "Loading",
|
||
) -> None:
|
||
"""Raise ``ValueError`` if *model* has modules offloaded to CPU or disk."""
|
||
offloaded = get_offloaded_device_map_entries(model)
|
||
if not offloaded:
|
||
return
|
||
example = ", ".join(f"{name}={placement}" for name, placement in list(offloaded.items())[:5])
|
||
raise ValueError(
|
||
f"{context} does not support models loaded with CPU or disk offload. "
|
||
f"device_map='{device_map}' produced offloaded modules: {example}"
|
||
)
|
||
|
||
|
||
def safe_num_proc(desired: Optional[int] = None) -> int:
|
||
"""
|
||
Return a safe ``num_proc`` for ``dataset.map()`` calls.
|
||
|
||
On Windows always returns 1: Python uses ``spawn`` not ``fork``, so
|
||
re-importing torch/transformers/unsloth per worker is typically slower
|
||
than single-process for normal dataset sizes.
|
||
|
||
On multi-GPU machines (multiple GPUs *visible* to this process) the
|
||
NVIDIA driver spawns extra background threads, making ``os.fork()``
|
||
deadlock-prone with many workers, so this caps ``num_proc`` to 4.
|
||
The cap does not apply when ``CUDA_VISIBLE_DEVICES`` restricts to one GPU.
|
||
|
||
Args:
|
||
desired: The num_proc you *want*. If None, auto-computes from
|
||
``os.cpu_count()``.
|
||
|
||
Returns:
|
||
A safe integer ≥ 1.
|
||
"""
|
||
# Windows/macOS use 'spawn'; re-importing torch/transformers/unsloth per
|
||
# worker is typically slower than single-process.
|
||
if sys.platform in ("win32", "darwin"):
|
||
return 1
|
||
|
||
if desired is None or not isinstance(desired, int):
|
||
desired = max(1, (os.cpu_count() or 1) // 3)
|
||
|
||
visible = get_visible_gpu_count()
|
||
if visible > 1:
|
||
capped = max(1, min(4, desired))
|
||
logger.info(
|
||
f"Multi-GPU detected ({visible} visible GPUs) -- "
|
||
f"capping num_proc {desired} -> {capped} to avoid fork deadlocks"
|
||
)
|
||
return capped
|
||
|
||
return max(1, desired)
|
||
|
||
|
||
def safe_thread_num_proc(desired: Optional[int] = None) -> int:
|
||
"""
|
||
Return a safe worker count for ``ThreadPoolExecutor`` calls.
|
||
|
||
Unlike ``safe_num_proc()``, does NOT cap to 1 on macOS/Windows: threads
|
||
share the parent address space, unaffected by ``spawn`` vs ``fork``.
|
||
|
||
Args:
|
||
desired: The thread count you *want*. If None, auto-computes
|
||
from ``os.cpu_count()``.
|
||
|
||
Returns:
|
||
A safe integer >= 1.
|
||
"""
|
||
if desired is None or not isinstance(desired, int):
|
||
desired = max(1, (os.cpu_count() or 1) // 3)
|
||
|
||
return max(1, desired)
|
||
|
||
|
||
def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]:
|
||
"""
|
||
Return a safe ``num_proc`` for ``Dataset.map()`` and ``Dataset.filter()``.
|
||
|
||
Returns ``None`` on spawn platforms (Windows, macOS) because ``datasets``
|
||
treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``); only
|
||
``num_proc=None`` guarantees in-process execution.
|
||
"""
|
||
if sys.platform in ("win32", "darwin"):
|
||
return None
|
||
return safe_num_proc(desired)
|