a00fe86c13* Studio: read model text as utf-8 so umlauts survive on Windows Chat rejects or mangles non-ASCII on Windows: "ä ö ü" in a prompt, a chat template, or a model path comes back as mojibake, or the load dies with UnicodeDecodeError. open() and Path.read_text() fall back to locale.getencoding() when no encoding is passed. On Windows that is the ANSI codepage (cp1252, cp932, cp1251, ... by system locale), never UTF-8. Hugging Face writes these files as raw UTF-8, so every read of one decodes with the wrong codec: - tokenizer_config.json, which holds the chat template. Templates routinely carry -> arrows, smart quotes and CJK, so this is the common path into chat - config.json and adapter_config.json - modules.json, Ollama manifests, and the .py sources the remote-code scanner reads before a model is allowed to load The llama-server and embedding-server stdout readers have the same problem via subprocess(text = True); they now decode utf-8 with errors = "replace" so a stray byte cannot kill a log reader. Encoding arguments only, no logic changes. tests/test_chat_text_encoding.py covers a config.json and a chat template holding umlauts, arrows and CJK, plus the remote-code scanner reading a source file with umlauts. Those pass anywhere the locale is already UTF-8, so a fourth test re-runs the readers under -X warn_default_encoding and fails on any platform if an encoding argument goes missing again. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: name utf-8 explicitly on the remaining text I/O, with an AST guard (#7465) * Studio: name utf-8 explicitly on the remaining text I/O Follow-up to the model-text reads in #7467, covering the rest of the backend: system probes (nvidia-smi, amd-smi, powershell, git, node), package installers, /proc and /sys readers, and internal marker files (pid, install id, bootstrap password, Colab credentials). Same reason as #7467. open(), Path.read_text()/write_text() and subprocess(text = True) fall back to locale.getencoding(), which on Windows is the ANSI codepage rather than UTF-8. These paths are mostly ASCII today, so this is hardening, not a live bug. Encoding arguments only, no logic changes. Adds tests/test_text_io_encoding.py: an AST guard walking every backend source and asserting text I/O names its encoding, so the class of bug cannot creep back in one call at a time. 275 files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Catch aliased subprocess and positional Path.open, migrate legacy JSONL The guard only matched a receiver literally named subprocess, so worker.py's `import subprocess as _sp` hid three text = True installs that decode pip output with the ANSI codepage. It also skipped any .open() with more than one positional argument, though Path.open takes buffering/encoding/errors/newline positionally. Resuming a scrape written by an older release is the other half: those JSONL lines are in the locale codepage, so the UTF-8 preload raised, the dedup keys were silently forgotten and duplicates were appended to a now mixed-encoding file. Decode with the locale codepage as fallback and rewrite as UTF-8 before the append handle opens, since Windows cannot replace a file it holds open. * Stream the JSONL preload and keep a torn line from relabelling the shard Reading the whole shard to migrate it was wrong twice over. These files reach gigabytes on a large scrape, so the preload now streams line by line and the rewrite streams through a temp file. Worse, one interrupted append used to condemn the file: the whole-file UTF-8 decode failed, every byte was retried as cp1252, and the rewrite persisted mojibake over records that were fine. A line now counts as legacy only if the locale codepage both decodes it and yields valid JSON, which a torn UTF-8 line does not. Damaged lines are skipped and copied through byte for byte. When the rewrite cannot be written at all, the append handle opens with the legacy encoding rather than mixing UTF-8 into the file. install_wheel takes run = subprocess.run as a parameter, so the guard cannot see it. Both wheel installs there now name their encoding. * Decide the shard's encoding from the file, not one line at a time Some byte strings parse both ways. cp1251 `Р°` is D0 B0, which is also valid UTF-8 for `а`, so a UTF-8-first parse quietly showed the wrong text instead of migrating it. A line now yields both readings, and the file decides. Any line that parses under the codepage but not as UTF-8 is unambiguous evidence, and ambiguous lines then follow that verdict, which is enough for any real shard: ordinary Cyrillic or Japanese prose is invalid UTF-8 several times per line. Keys for ambiguous lines are re-derived from the legacy reading during the rewrite. A shard is undecidable only if every line is ambiguous, and nothing can tell those apart. latin-1 is also tried after the locale codepage, so a scrape carried from Windows to a UTF-8 machine still has a reading rather than none. Requiring valid JSON, not just a decode, keeps that from claiming torn lines. * Weigh the whole shard, and never lose a record on the fallback path One structurally valid JSON line carrying a stray 0x96 parses as cp1252, so a single-line verdict let it relabel a healthy shard and mojibake every good record in it. Each line with non-ASCII bytes now votes: parsing only under the codepage is evidence for legacy, parsing as UTF-8 is evidence against, since codepage text rarely forms valid multibyte UTF-8. Ties leave the file alone. When the migration cannot be written the append handle uses the legacy codepage, and errors = "replace" quietly turned characters it cannot hold into question marks while write() still reported success. That path now escapes to \uXXXX instead, which is ASCII, so every codepage holds it and json.loads returns the exact characters. Nothing needs replacing, so errors = "strict" is safe. stream_installer runs sys.executable, so its output is now decoded as UTF-8 by utf8_child_env rather than read as the ANSI codepage. * Only rewrite a shard we can attribute, and append ASCII when we cannot latin-1 was doing too much work. It reads any byte, so it gave a moved shard a reading, but it is the right text only for cp1252: cp1251 Привет came back as Ïðèâåò and the rewrite made that permanent. The codepage is now trusted only when it is the locale's, and an untrusted reading is never written back. That leaves three cases where the file holds bytes UTF-8 cannot read and we are not converting it: no codepage to attribute it to, ambiguous lines outvoting the unambiguous ones, and a preload that could not read the file at all. All three used to append UTF-8 into it. They now append pure ASCII, which every ASCII-compatible codepage stores identically, so the file keeps decoding exactly as it did and no record is lost. Keys from the two readings are also kept apart. A damaged line in a healthy shard was marked seen through its codepage reading, so the retry that would have replaced the unreadable record was refused as a duplicate. * Let the flash-attn install stub take the kwargs the installer now passes _run_kwargs gained encoding and errors, so the one stub in this file that spelled its signature out rejected the call. The other four here already take **kwargs; this one now matches. * Do not let a stuck temp file mask the migration failure unlink() on the failure path could raise in its own right, on a stale .utf8.tmp directory or a temp another process holds. That escaped the constructor instead of returning False, so the caller never reached the ASCII append fallback that keeps the shard single-encoding. The pip fallback in install_wheel also spawns a Python child, so it gets utf8_child_env like the probe above it already had. The uv and nvidia-smi children are native binaries, where PYTHONIOENCODING would do nothing. * Stop converting legacy shards; the encoding that wrote them is unknowable trusted only ever meant that the bytes parse under this machine's codepage, which for a single-byte codepage is nearly always true. A cp1251 shard opened on a cp1252 Windows box decodes cleanly and would have been rewritten with Привет as Ïðèâåò. That is the fourth way this rewrite could corrupt a shard, and the common cause is that a file's encoding cannot be recovered from its bytes. So the rewrite is gone. The shard is left exactly as found, and appends are pure ASCII whenever it holds bytes UTF-8 cannot read, which is what actually delivered the no-mixed-encoding guarantee the rewrite was added for. Dedup keys still come from whichever reading parses, since ids are ASCII either way. This also removes the temp file, so there is no longer any file mode or ACL to carry across. * Scan the sandbox shim; it is shipped code, not a build artifact sandbox_site is on the sandboxed child's PYTHONPATH for every Python run (tools.py:332, 2660), so excluding it let two unannotated text calls through in code we ship. Both read and write the remap sidecar, which holds file paths. The exclusion list is meant for build output only, so the directory comes off it and the two calls name their encoding. * Force the worker's pip children to UTF-8, and read DBCS keys with a DBCS codec The three installer calls run sys.executable -m pip with an inherited environment, so the parent decoded UTF-8 while the child emitted the ANSI codepage. They now go through utf8_child_env like the other Python children. Two tests asserted no env kwarg was passed as a stand-in for no HIP flag being injected. They now assert the flag itself, which is the guarantee they were written for and does not depend on how the env is delivered. Separately, latin-1 cannot stand in for a double-byte codepage while recovering dedup keys: cp932 表 is 95 5C, and the trail byte reads as a JSON backslash, so the record failed to parse and its id was forgotten, appending a duplicate on resume. cp932, cp936, cp949 and cp950 are tried too. The reading is still only ever used for keys, which are ASCII and identical whichever codec parses. * Require more than one legacy line before trusting its dedup keys A shard whose valid records are all ASCII casts no UTF-8 votes, so a single damaged line won the vote by itself, its key was remembered, and the retry that would have replaced the unreadable record was refused. One such line is genuinely undecidable: a legacy record with one accented character and an ASCII record with one stray byte are the same shape. Reading it as damage costs a duplicate; reading it as legacy loses the record for good. Only one of those is recoverable, so it is now read as damage. A real legacy shard has a legacy line for every record carrying an umlaut, so its dedup is unaffected. * Append ASCII whenever the shard already holds non-ASCII bytes The gate asked whether any line was undecodable as UTF-8, which misses a shard where every legacy line happens to be valid UTF-8 too. A cp1251 shard of Р° records is bytes D0 B0 throughout, so appending 世界 as UTF-8 left a file where cp1251 reads the old records correctly and the new one as mojibake, and UTF-8 does the reverse. No single decoding recovered the whole scrape. The gate is now simply whether the shard holds any non-ASCII byte at all, which covers both cases and is easier to reason about: if what is already there reads differently under different encodings, do not add more bytes that do. Appending ASCII costs only \uXXXX escapes, which json.loads turns back into the exact characters, and it leaves the new record correct under either reading. * Skip the two Linux-gated flash-attn tests off Linux _should_try_runtime_flash_attn_install ends in sys.platform.startswith( "linux"), and the threshold test one line above already asserts exactly that, so the two tests that drive _ensure_flash_attn_for_long_context past the gate cannot pass anywhere else: the call returns before it reports a status. They were written on Linux and only surface once the suite actually runs on Windows or macOS, where both fail on an empty status list. This PR is about making the backend behave on Windows, so its own suite should be runnable there. * Fail closed when a KFD topology node does not decode This PR pins that read to utf-8, which turns an undecodable byte into UnicodeDecodeError. That is a ValueError, not an OSError, so it slips past the handler one line below and escapes a helper whose docstring promises to fail closed on any unreadable node. The caller would then lose the whole HIP-order map on a machine that has AMD GPUs, and the reason the helper fails closed is that dropping a node shifts every later ordinal and lets a similar-capacity GPU pass the total-size guard while showing another card's usage. Widening the handler is the same one-line change main already made in #7487, so the two agree and the eventual merge is clean. * Tighten the comments added in this branch * Treat an undecodable marker and undecodable metadata as malformed, not fatal Two more places where pinning the decode changed the failure mode. A UnicodeDecodeError is a ValueError, so neither `except OSError` nor `except (JSONDecodeError, OSError)` catches it, and both sites had a documented fallback that stopped being reached. An undecodable .transport marker used to read as an unknown value, and the caller then safely purged and restarted the partial download. It now aborts prepare_cache_for_transport instead, so the transfer fails rather than retrying. Undecodable .meta.json used to fall back to the file's own name, the same way invalid JSON does. It now aborts URI construction for the entire unstructured seed, so one corrupt byte in original_filename takes out the whole dataset. Both handlers are widened, matching the KFD fix earlier on this branch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Widen two more decode guards, and pin the kernel installer's pipe Same shape as the ones already fixed here: the read was pinned to UTF-8 while the handler around it still only catches OSError, and UnicodeDecodeError is a ValueError. hf_cache_snapshot_dir answers whether a model is already on disk, and the offline embedding checks turn a raise into a 500. A torn refs/main used to decode into a nonsense commit and miss the snapshot dir; it now skips that cache root and keeps looking. _remove_pid_file runs first in _graceful_shutdown, so a corrupt studio.pid raising there abandoned the inference, export, training and tunnel children the rest of that function exists to kill. ssm_runtime's source-build path builds its subprocess kwargs in a dict and splats them through _run_with_heartbeat, so neither the encoding guard nor the earlier sweep saw the text = True in it: pip's output was still decoded with the Windows ANSI codepage, where a non-ASCII path or a compiler diagnostic mojibakes or raises over an install that was going fine. It now pins the same utf-8/replace pair install_wheel uses, and the HIP branch extends that env rather than replacing it. The guard learned the dict-literal shape and reddens on the old code (ssm_runtime.py:253). * Tighten the comments around the UTF-8 text I/O pins Collapse the multi-line rationales added with the encoding pins down to a line or two each, drop what the code already says, and use one wording for the repeated child-env note. * Do not let an unreadable bootstrap password stop startup, and narrow the kwargs guard ensure_default_admin calls _load_bootstrap_password for every existing admin and the lifespan calls that with no handler, so pinning the decode turned a damaged or pre-pin .bootstrap_password file into a backend that will not start. We write that file ourselves in UTF-8, so a byte that will not decode belongs to a file whose plaintext is worthless anyway; it now reads as no bootstrap password, the same answer as an absent file. A readable one still loads. The new kwargs check also judged every dict literal in the tree, so an unrelated payload carrying "text": True would have been reported as subprocess configuration with a misleading message, and a dict that fills in its encoding on a later line would have been reported too. It now only judges a dict that actually reaches a call, either splatted through a name or written at the call site, and treats a later kw["encoding"] assignment as satisfying it. The ssm_runtime shape it was written for is still caught, and a test pins both directions. * Stop reading a UTF-8 record a second time _read_line always parsed the line under the codepage as well, even when it had already read as UTF-8. Both callers take the UTF-8 reading when there is one and never look at the other, so on a healthy shard the second parse is pure waste, and this file reads all of one on every resume of a scrape it expects to reach gigabytes. Measured on 200,000 records, 76 MB: 1.96s before, 0.81s after, so the double reading was costing 2.8x. The early return is limited to a record, since the key lookup deliberately falls through to the codepage reading when UTF-8 yields something that is not one. A line UTF-8 cannot read still tries the codepage, latin-1 and the double-byte encodings as before, which is what the second reading is for. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pin the scanned source fixture's line endings test_remote_code_scan_reads_non_ascii_sources compared a file's contents against the string it wrote, but wrote it in text mode, so Windows translated the line ends on the way out and the read back differed by a carriage return. That is the writer's doing, not the encoding the test is about, and it was the one failure on the Windows runner that belonged to this branch. The fixture now writes with newline = "" so the bytes on disk are the string on every platform. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim the newer comments to their point Shorten the widened-guard and state store notes added since the last pass, and collapse the line-ending note on the scanned source fixture. * Read the scraper checkpoint as UTF-8 only, never as a codepage A checkpoint holds nothing but base64 cursors and booleans, so one written by an older locale-encoded release is byte-identical to a UTF-8 one and already reads back. The codepage fallback can therefore only ever contribute non-ASCII: if a single-byte reading of the file were all ASCII, the UTF-8 read would have succeeded first. So the only file it changes the answer for is a damaged one, and there it turns a safe reset into a resume on a mojibaked cursor. GitHub answers that with INVALID_CURSOR_ARGUMENTS at HTTP 200, gh_client returns the partial document, and the scraper reads zero nodes and an empty pageInfo, which marks the stream done. Every later resume then skips it entirely. Reading UTF-8 only restores the earlier behaviour of dropping a checkpoint that will not decode, which re-scrapes from the first page while the writers dedup the replay. The shard scan below keeps its codepage reading; those records do carry non-ASCII. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate the remaining tilelang install tests to Linux _tilelang_platform_supported() returns False off Linux, so _ensure_tilelang_backend returns before the install and the subprocess mock these six assert on is never called. They fail on macOS runners for that reason alone. The rest of the file already carries this marker; these were missed. * Gate the Windows-incompatible worker and ROCm tests Two different gates, because the production code has two. The causal-conv1d and flash-linear-attention installers bail out on sys.platform == 'win32' alone and run everywhere else including macOS, so those cases get not_on_windows; marking them linux_only would skip tests that legitimately pass off Linux. The DRM and KFD readers return early unless platform.system() is Linux, and their fixtures build a fake sysfs tree needing PCI addresses like 0000:00:02.0 as directory names, which Windows cannot represent, so those get linux_only. The two visible-utilization cases failed for a different reason: on Windows get_visible_gpu_utilization takes the AMD adapter branch ahead of the torch fallback under test, and probing it imports torch, which the runner lacks. Stubbing that branch empty leaves every other platform unchanged. * Treat unparseable JSON nesting as a parse failure, and guard os.fdopen json.loads answers nesting it cannot descend with RecursionError, a RuntimeError, so _parse let it escape where the catch-all it replaced discarded the record. Both callers run _parse outside any further handler, so one damaged checkpoint or shard line aborted the scraper at startup. The encoding guard also missed os.fdopen, which is open() on a descriptor and takes the same locale default in text mode. It flags exactly the two text-mode calls that were left unencoded; the swap lock file's reader was already pinned to UTF-8 while its writer still used the codepage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Write the non-ASCII source fixture without a 3.10-only argument Path.write_text() only grew newline in 3.10, and pyproject declares requires-python >=3.9, so this raised TypeError there. open() takes the same argument on every supported version and pins the bytes on disk the same way. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten encoding comments * Follow subprocess calls through callable aliases in the encoding guard --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com> --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com>
3117 lines
123 KiB
Python
3117 lines
123 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__)
|
||
|
||
|
||
# ── GPU index ordering ──────────────────────────────────────────────────────
|
||
# CUDA defaults to CUDA_DEVICE_ORDER=FASTEST_FIRST, numbering GPUs by compute
|
||
# performance. nvidia-smi -- and every free-VRAM probe in Unsloth -- numbers GPUs
|
||
# by PCI bus id instead. On a mixed-GPU host (e.g. an RTX 5090 alongside an RTX
|
||
# PRO 6000) the two orderings disagree, so an index picked from nvidia-smi data
|
||
# ("the emptiest card is GPU 1") gets written into CUDA_VISIBLE_DEVICES and then
|
||
# reinterpreted by CUDA against FASTEST_FIRST -- landing the model on a different
|
||
# physical GPU than the one selected. Pinning PCI_BUS_ID makes torch, nvidia-smi,
|
||
# and CUDA_VISIBLE_DEVICES share a single index space, matching what users see in
|
||
# `nvidia-smi -L`. Set at import (before any torch.cuda call latches the order
|
||
# at context creation) and inherited by child processes, since the llama-server
|
||
# and spawn workers copy os.environ. setdefault so an explicit user override wins.
|
||
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
|
||
|
||
# Unsloth workers can import MLX without importing unsloth first, so mirror the
|
||
# package bootstrap here. Keep an explicit user value authoritative.
|
||
if platform.system() == "Darwin" and platform.machine() == "arm64":
|
||
os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1")
|
||
|
||
|
||
# ========== 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.)
|
||
# Why CHAT_ONLY is True (Train/Export disabled). None when training is enabled.
|
||
# "mlx_unavailable": Apple Silicon but the MLX stack is missing, too old, or broken
|
||
# (the usual cause of "Train/Export greyed out" on Macs after a reinstall dropped MLX);
|
||
# "intel_mac": Intel Mac (no PyTorch/MLX); "no_gpu": CPU-only non-Mac host.
|
||
CHAT_ONLY_REASON: Optional[str] = None
|
||
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 _has_usable_mlx_stack() -> bool:
|
||
"""True only when the FULL Unsloth MLX training/export stack is usable
|
||
(mlx + mlx-lm + mlx-vlm at the minimum versions unsloth-zoo requires), not
|
||
just a bare ``import mlx.core``. A backtracked/old mlx-vlm still imports but
|
||
breaks VLM Train/Export, so the training gate must match the self-heal's own
|
||
criterion (utils.mlx_repair.mlx_stack_available) -- otherwise detect_hardware
|
||
would enable Train/Export on exactly the inadequate stack the MLX self-heal
|
||
is trying to repair, leaving the user with greyed-in-but-broken buttons."""
|
||
try:
|
||
from utils.mlx_repair import mlx_stack_available
|
||
return mlx_stack_available()
|
||
except Exception as exc:
|
||
# mlx_repair should always import; if it somehow cannot, fall back to the
|
||
# bare import check rather than forcing a working host into chat-only.
|
||
logger.debug("MLX stack availability check failed, using bare import: %s", exc)
|
||
return _has_mlx()
|
||
|
||
|
||
def _print_cuda_device_list(is_rocm: bool) -> None:
|
||
"""List every visible CUDA/ROCm GPU with its index at startup.
|
||
|
||
The "Hardware detected" banner names only device 0, which hides the other
|
||
cards on a multi-GPU host. This lists the full visible set in CUDA-ordinal
|
||
order, matching `nvidia-smi -L` when no CUDA_VISIBLE_DEVICES mask is set
|
||
(under a mask the indices are visible ordinals, not physical PCI ids).
|
||
CUDA_DEVICE_ORDER governs only CUDA, so it is shown for CUDA but not ROCm.
|
||
No-ops on single-GPU hosts and never raises -- it is purely informational.
|
||
"""
|
||
try:
|
||
import torch
|
||
|
||
count = torch.cuda.device_count()
|
||
if count <= 1:
|
||
return
|
||
if is_rocm:
|
||
header = f"ROCm devices ({count}):"
|
||
else:
|
||
order = os.environ.get("CUDA_DEVICE_ORDER", "default")
|
||
header = f"CUDA devices ({count}, CUDA_DEVICE_ORDER={order}):"
|
||
lines = [header]
|
||
for i in range(count):
|
||
try:
|
||
name = torch.cuda.get_device_properties(i).name
|
||
except Exception as e:
|
||
logger.debug("CUDA device %d property probe failed: %s", i, e)
|
||
name = "<unavailable>"
|
||
lines.append(f" [{i}] {name}")
|
||
print("\n".join(lines))
|
||
except Exception:
|
||
return # purely informational; never disrupt startup
|
||
|
||
|
||
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. XPU-preferred hint: only on an unambiguous "prefer XPU" signal
|
||
(CUDA hidden via ``CUDA_VISIBLE_DEVICES="" / "-1"``,
|
||
``UNSLOTH_FORCE_XPU=1``, or CUDA unavailable) AND a non-empty
|
||
``ZE_AFFINITY_MASK`` AND ``torch.xpu`` reports a device. A stray
|
||
inherited mask is not enough: CUDA still wins on hybrid hosts.
|
||
2. CUDA (NVIDIA GPU, requires torch)
|
||
3. XPU (Intel GPU, requires torch with XPU support)
|
||
4. MLX (Apple Silicon via MLX framework)
|
||
5. CPU (fallback)
|
||
"""
|
||
global DEVICE, CHAT_ONLY, CHAT_ONLY_REASON, IS_ROCM
|
||
CHAT_ONLY = True # reset -- only CUDA/ROCm/XPU/MLX sets it to False
|
||
CHAT_ONLY_REASON = None
|
||
IS_ROCM = False
|
||
|
||
# --- CUDA / ROCm / XPU: try PyTorch ---
|
||
if _has_torch():
|
||
import torch
|
||
|
||
# --- Explicit-XPU hint ---
|
||
# Prefer XPU on UNSLOTH_FORCE_XPU=1, or ZE_AFFINITY_MASK set + CUDA
|
||
# hidden/unavailable. A bare mask alone is NOT enough (can leak from
|
||
# unrelated Intel tooling); torch.xpu must report a device.
|
||
ze_mask = os.environ.get("ZE_AFFINITY_MASK")
|
||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
cuda_hidden = cvd is not None and cvd.strip() in ("", "-1")
|
||
force_xpu = os.environ.get("UNSLOTH_FORCE_XPU") == "1"
|
||
try:
|
||
cuda_unavailable = not torch.cuda.is_available()
|
||
except Exception:
|
||
cuda_unavailable = True
|
||
|
||
prefer_xpu = force_xpu or (bool(ze_mask) and (cuda_hidden or cuda_unavailable))
|
||
if prefer_xpu:
|
||
try:
|
||
xpu_ok = hasattr(torch, "xpu") and torch.xpu.is_available()
|
||
except Exception:
|
||
xpu_ok = False
|
||
if xpu_ok:
|
||
# Forced XPU on a hybrid host: unsloth's device_type picks
|
||
# CUDA before XPU and ignores this Studio-only env var, so
|
||
# hide CUDA or spawned workers would silently train on CUDA.
|
||
if force_xpu and not cuda_hidden and not cuda_unavailable:
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
||
DEVICE = DeviceType.XPU
|
||
CHAT_ONLY = False
|
||
CHAT_ONLY_REASON = None
|
||
device_name = torch.xpu.get_device_name(0)
|
||
if force_xpu and not ze_mask:
|
||
reason = "UNSLOTH_FORCE_XPU=1"
|
||
elif force_xpu:
|
||
reason = "UNSLOTH_FORCE_XPU=1 + ZE_AFFINITY_MASK"
|
||
else:
|
||
reason = "ZE_AFFINITY_MASK hint honoured"
|
||
print(f"Hardware detected: XPU -- {device_name} ({reason})")
|
||
return DEVICE
|
||
|
||
# --- CUDA: NVIDIA GPU ---
|
||
if torch.cuda.is_available():
|
||
DEVICE = DeviceType.CUDA
|
||
CHAT_ONLY = False
|
||
try:
|
||
device_name = torch.cuda.get_device_properties(0).name
|
||
except Exception as e:
|
||
logger.debug("CUDA device 0 property probe failed: %s", e)
|
||
device_name = "<unavailable>"
|
||
|
||
# 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}")
|
||
_print_cuda_device_list(IS_ROCM)
|
||
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 ---
|
||
# Require the full mlx/mlx-lm/mlx-vlm stack (not a bare `import mlx.core`) so
|
||
# the gate matches utils.mlx_repair: a partial/backtracked stack stays
|
||
# chat-only (reason "mlx_unavailable") and the background self-heal repairs it.
|
||
if is_apple_silicon() and _has_usable_mlx_stack():
|
||
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
|
||
# CHAT_ONLY is still True here (every training-capable branch returned early),
|
||
# so record WHY so the UI can explain the greyed-out Train/Export instead of
|
||
# silently disabling them.
|
||
if is_apple_silicon():
|
||
# Reached the CPU fallback on Apple Silicon, so the MLX stack is missing,
|
||
# too old, or broken. This is usually an environment problem recoverable
|
||
# with `unsloth studio update`.
|
||
CHAT_ONLY_REASON = "mlx_unavailable"
|
||
logger.warning(
|
||
"Apple Silicon detected but the MLX stack is incomplete or too old; "
|
||
"Train/Export disabled (chat-only). Run `unsloth studio update` to "
|
||
"restore MLX training."
|
||
)
|
||
elif platform.system() == "Darwin":
|
||
CHAT_ONLY_REASON = "intel_mac" # Intel Mac: no PyTorch/MLX -> GGUF-only by design.
|
||
else:
|
||
CHAT_ONLY_REASON = "no_gpu"
|
||
print("Hardware detected: CPU training backend (no PyTorch/MLX 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 export_capability() -> dict:
|
||
"""Whether model export can run here, with a torch-aware reason when it cannot.
|
||
|
||
Export runs through Unsloth, which hard-requires an accelerator (it calls ``torch.cuda`` at
|
||
import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The
|
||
reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch.
|
||
|
||
Returns {export_supported, export_unsupported_reason, export_unsupported_message}.
|
||
"""
|
||
if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX):
|
||
return {
|
||
"export_supported": True,
|
||
"export_unsupported_reason": None,
|
||
"export_unsupported_message": None,
|
||
}
|
||
# No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch"
|
||
# would be wrong advice on a Mac even when torch is also absent.
|
||
if is_apple_silicon():
|
||
reason = "mlx_unavailable"
|
||
message = (
|
||
"Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run "
|
||
"`unsloth studio update` to restore MLX and enable export."
|
||
)
|
||
elif not _has_torch():
|
||
reason = "pytorch_not_installed"
|
||
message = (
|
||
"PyTorch is not installed. Model export requires PyTorch with a supported accelerator "
|
||
"(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export."
|
||
)
|
||
else:
|
||
reason = "no_accelerator"
|
||
message = (
|
||
"Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported "
|
||
"accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export "
|
||
"on CPU only.)"
|
||
)
|
||
return {
|
||
"export_supported": False,
|
||
"export_unsupported_reason": reason,
|
||
"export_unsupported_message": message,
|
||
}
|
||
|
||
|
||
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:
|
||
# Guard synchronize/empty_cache: older torch-xpu builds may lack
|
||
# them, and an unguarded AttributeError would propagate to callers.
|
||
# torch.xpu has no ipc_collect(), so do not call it here.
|
||
try:
|
||
import torch
|
||
if hasattr(torch, "xpu"):
|
||
if hasattr(torch.xpu, "synchronize"):
|
||
torch.xpu.synchronize()
|
||
if hasattr(torch.xpu, "empty_cache"):
|
||
torch.xpu.empty_cache()
|
||
except Exception as e:
|
||
logger.debug("Failed to clear XPU cache: %s", e)
|
||
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 versions bundled with torch (CUDA, ROCm/HIP, Intel XPU)
|
||
try:
|
||
import torch
|
||
|
||
versions["cuda"] = getattr(torch.version, "cuda", None)
|
||
versions["rocm"] = getattr(torch.version, "hip", None)
|
||
# Isolated probe: a broken Intel runtime raising in is_available()
|
||
# must not blank the already-read cuda/rocm versions.
|
||
try:
|
||
if hasattr(torch, "xpu") and torch.xpu.is_available():
|
||
# torch.version.xpu may be None on modern builds; fall back to
|
||
# "available" so the UI distinguishes present-but-unknown from
|
||
# "package not found".
|
||
xpu_ver = getattr(torch.version, "xpu", None)
|
||
versions["xpu"] = xpu_ver if xpu_ver is not None else "available"
|
||
except Exception:
|
||
versions["xpu"] = None
|
||
except Exception:
|
||
versions["cuda"] = None
|
||
versions["rocm"] = None
|
||
versions["xpu"] = 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.
|
||
|
||
``used_gb`` is ``None`` on Windows ROCm when ``hipMemGetInfo`` reports
|
||
``free == total`` (ROCm/ROCm#1909): that 0 means unknown, not empty.
|
||
"""
|
||
mod, _ = _torch_get_device_module()
|
||
if mod is None:
|
||
return []
|
||
|
||
device = get_device()
|
||
# free==total is a Windows-ROCm-only quirk.
|
||
_win_rocm = sys.platform == "win32" and IS_ROCM
|
||
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
|
||
used_bytes: Optional[int]
|
||
# Prefer mem_get_info (system-wide) so auto-select sees other consumers.
|
||
if hasattr(mod, "mem_get_info"):
|
||
try:
|
||
free_bytes, total_bytes = mod.mem_get_info(ordinal)
|
||
used_bytes = total_bytes - free_bytes
|
||
except Exception as e:
|
||
if device != DeviceType.XPU:
|
||
raise
|
||
# Arc B580 and Lunar Lake can report properties while
|
||
# rejecting free-memory queries. Preserve the usable
|
||
# device and its total memory with unknown utilization.
|
||
logger.debug(
|
||
"XPU free-memory query failed for ordinal %d: %s",
|
||
ordinal,
|
||
e,
|
||
)
|
||
used_bytes = None
|
||
else:
|
||
# free==total is the broken-API sentinel, not an idle GPU.
|
||
if _win_rocm and free_bytes == total_bytes:
|
||
used_bytes = None
|
||
elif device == DeviceType.XPU:
|
||
# XPU without mem_get_info: memory_allocated() is process-local
|
||
# and misleading for placement, so return None for the
|
||
# selector's no-telemetry fallback.
|
||
used_bytes = None
|
||
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) if used_bytes is not None else None
|
||
),
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.debug("torch device query failed for ordinal %d: %s", ordinal, e)
|
||
return devices
|
||
|
||
|
||
# ========== Live GPU Utilization ==========
|
||
|
||
|
||
def _xpu_hierarchy_is_composite() -> bool:
|
||
"""Return True iff Level Zero is running in COMPOSITE device hierarchy.
|
||
|
||
COMPOSITE: numeric ``ZE_AFFINITY_MASK`` entries address root GPU IDs
|
||
(tiles use ``N.M``). FLAT (the oneAPI default; also assumed when
|
||
``ZE_FLAT_DEVICE_HIERARCHY`` is unset): entries address tile/device
|
||
handles, so mapping them back to root GPU IDs is unsafe. Only COMPOSITE
|
||
gives stable root-ID semantics.
|
||
"""
|
||
hierarchy = (os.environ.get("ZE_FLAT_DEVICE_HIERARCHY") or "FLAT").strip().upper()
|
||
return hierarchy == "COMPOSITE"
|
||
|
||
|
||
def _parse_ze_mask_roots(mask: str) -> list[int]:
|
||
"""Parse a ``ZE_AFFINITY_MASK`` value into an ordered list of root device IDs.
|
||
|
||
One root ID per mask token, preserving order and duplicates so logical
|
||
ordinals map 1-to-1 to physical root IDs (e.g. ``"0.0,0.1"`` -> ``[0, 0]``,
|
||
``"2.0,0.1,0.2"`` -> ``[2, 0, 0]``); empty list if no parseable digits.
|
||
Only meaningful in COMPOSITE hierarchy -- callers needing a stable
|
||
root-ID mapping must gate on ``_xpu_hierarchy_is_composite()``.
|
||
"""
|
||
roots: list[int] = []
|
||
if not mask:
|
||
return roots
|
||
for token in mask.split(","):
|
||
token = token.strip()
|
||
if not token:
|
||
continue
|
||
root = token.split(".", 1)[0]
|
||
# isdecimal() (not isdigit()) rejects Unicode superscripts like
|
||
# "²"/"³", which pass isdigit() but crash int() with ValueError.
|
||
if root.isdecimal():
|
||
roots.append(int(root))
|
||
return roots
|
||
|
||
|
||
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, encoding = "utf-8").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, encoding = "utf-8").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, encoding = "utf-8").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,
|
||
encoding = "utf-8",
|
||
errors = "replace",
|
||
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, encoding = "utf-8").read().strip()) for f in used_files)
|
||
total_bytes = sum(int(open(f, encoding = "utf-8").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
|
||
|
||
|
||
# 0x1002. NVIDIA's open kernel module also registers KFD nodes (vendor_id 0x10DE);
|
||
# a non-AMD node is not a HIP device and must never take an ordinal.
|
||
_AMD_PCI_VENDOR_ID = 4098
|
||
|
||
|
||
def _rocm_kfd_gpu_pci_ids() -> list[str]:
|
||
"""PCI addresses of the GPUs ROCm enumerates, in HIP device order.
|
||
|
||
Reads /sys/class/kfd/kfd/topology/nodes/<N>/properties, the topology ROCm
|
||
itself enumerates from: AMD GPU nodes (simd_count > 0 excludes CPUs,
|
||
vendor_id == AMD excludes NVIDIA) in node-id order are HIP's device order, so
|
||
position N is ROCm physical device N. Unlike DRM sysfs, an amdgpu adapter HIP
|
||
cannot enumerate has no node here, so it never consumes an ordinal.
|
||
|
||
Returns [] (disabling the overlay) when KFD is absent, and FAILS CLOSED the
|
||
same way on any unreadable node or an AMD node with no location_id: dropping
|
||
one would shift every later ordinal and let a similar-capacity GPU pass the
|
||
total-size guard while showing another card's usage.
|
||
|
||
location_id is the kernel's (bus << 8) | devfn; domain is separate.
|
||
"""
|
||
nodes: list[tuple[int, str]] = []
|
||
try:
|
||
node_dirs = glob.glob("/sys/class/kfd/kfd/topology/nodes/*")
|
||
except Exception:
|
||
return []
|
||
for node_dir in node_dirs:
|
||
m = re.fullmatch(r".*/(\d+)", node_dir)
|
||
if m is None:
|
||
continue
|
||
props: dict[str, int] = {}
|
||
try:
|
||
with open(os.path.join(node_dir, "properties"), encoding = "utf-8") as f:
|
||
for line in f:
|
||
parts = line.split()
|
||
if len(parts) == 2:
|
||
try:
|
||
props[parts[0]] = int(parts[1])
|
||
except ValueError:
|
||
continue
|
||
except (OSError, UnicodeDecodeError):
|
||
return [] # unreadable node could be a GPU: fail closed, don't shift
|
||
if props.get("simd_count", 0) <= 0:
|
||
continue # CPU node, not a GPU
|
||
if props.get("vendor_id") != _AMD_PCI_VENDOR_ID:
|
||
continue # non-AMD GPU node (NVIDIA open driver): not a HIP device
|
||
location_id = props.get("location_id")
|
||
if location_id is None:
|
||
return [] # an AMD GPU we cannot place: fail closed for the whole map
|
||
domain = props.get("domain", 0)
|
||
bus = (location_id >> 8) & 0xFF
|
||
devfn = location_id & 0xFF
|
||
bdf = f"{domain:04x}:{bus:02x}:{(devfn >> 3) & 0x1F:02x}.{devfn & 0x7}"
|
||
nodes.append((int(m.group(1)), bdf))
|
||
nodes.sort(key = lambda n: n[0])
|
||
return [bdf for _node_id, bdf in nodes]
|
||
|
||
|
||
def _rocm_linux_amdgpu_cards() -> list[tuple[str, int, str]]:
|
||
"""The amdgpu-bound DRM cards in PCI order: ``(pci_bdf, card_no, device_dir)``.
|
||
|
||
Membership is by the BOUND DRIVER, not the VRAM sysfs files: an AMD device
|
||
with incomplete sysfs support (some APUs expose no mem_info_vram_*) still
|
||
consumes a ROCm ordinal, and dropping it would shift every later card down.
|
||
PCI order is HIP's default enumeration order, so list position is the ROCm
|
||
ordinal; card_no is a stable tiebreak when the BDF cannot be resolved.
|
||
|
||
NOTE this is a superset of the ROCm-visible set (a HIP-unsupported amdgpu
|
||
adapter appears too), so callers must check the counts agree before assuming
|
||
a 1:1 mapping onto torch devices.
|
||
"""
|
||
if platform.system() != "Linux":
|
||
return []
|
||
amd_cards: list[tuple[str, int, str]] = []
|
||
try:
|
||
for card_path in glob.glob("/sys/class/drm/card*"):
|
||
# Match card<N> exactly so connector nodes (card0-DP-1) are skipped.
|
||
m = re.fullmatch(r".*/card(\d+)", card_path)
|
||
if m is None:
|
||
continue
|
||
dev_dir = os.path.join(card_path, "device")
|
||
try:
|
||
driver = os.path.basename(os.path.realpath(os.path.join(dev_dir, "driver")))
|
||
except OSError:
|
||
continue
|
||
if driver != "amdgpu":
|
||
continue # foreign adapter: not a ROCm device, takes no ordinal
|
||
try:
|
||
bdf = os.path.basename(os.path.realpath(dev_dir))
|
||
except OSError:
|
||
bdf = ""
|
||
amd_cards.append((bdf, int(m.group(1)), dev_dir))
|
||
except Exception:
|
||
return []
|
||
amd_cards.sort(key = lambda c: (c[0], c[1]))
|
||
return amd_cards
|
||
|
||
|
||
def _rocm_linux_sysfs_vram_by_pci_gb() -> dict[str, tuple[float, float]]:
|
||
"""System-wide AMD VRAM via Linux DRM sysfs, keyed by the card's PCI address.
|
||
|
||
Reads each card's mem_info_vram_{used,total} (kernel-updated across all
|
||
processes) so every GPU gets its own figure, unlike _rocm_linux_sysfs_vram_gb
|
||
which sums the host. Keyed by PCI address, not an ordinal, so the caller can
|
||
join it to _rocm_kfd_gpu_pci_ids() by identity: DRM card numbers include
|
||
foreign adapters and this set includes cards HIP does not enumerate, so any
|
||
ordinal from this list alone can be shifted relative to ROCm's. A card with
|
||
missing/unreadable/zero-total figures simply has no entry. Empty off Linux.
|
||
"""
|
||
if platform.system() != "Linux":
|
||
return {}
|
||
|
||
try:
|
||
by_pci: dict[str, tuple[float, float]] = {}
|
||
for bdf, _card_no, dev_dir in _rocm_linux_amdgpu_cards():
|
||
if not bdf:
|
||
continue
|
||
try:
|
||
with open(os.path.join(dev_dir, "mem_info_vram_used"), encoding = "utf-8") as f:
|
||
used_bytes = int(f.read().strip())
|
||
with open(os.path.join(dev_dir, "mem_info_vram_total"), encoding = "utf-8") as f:
|
||
total_bytes = int(f.read().strip())
|
||
except (OSError, ValueError):
|
||
continue
|
||
if total_bytes <= 0:
|
||
continue
|
||
by_pci[bdf.lower()] = (
|
||
round(used_bytes / (1024**3), 2),
|
||
round(total_bytes / (1024**3), 2),
|
||
)
|
||
return by_pci
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
# ── Windows AMD/ROCm per-adapter VRAM (issue #7072) ──────────────────────────
|
||
# amd-smi is disabled and hipMemGetInfo reports free==total, so read used from the
|
||
# per-LUID "GPU Adapter Memory" perf counters and take each total from torch, so
|
||
# every GPU shows instead of one fake device with GPU 0's total.
|
||
# Placeholder adapters (Basic Render Driver / idle iGPU) drop only when they would
|
||
# outnumber the real torch devices.
|
||
_ROCM_WIN_ADAPTER_MIN_BYTES = 64 * 1024 * 1024 # 64 MiB
|
||
|
||
|
||
def _rocm_windows_perf_counter_vram_by_adapter() -> Optional[list[tuple[str, float]]]:
|
||
"""Per-adapter dedicated VRAM usage on Windows via Performance Counters.
|
||
|
||
Returns ``[(instance_name, used_bytes)]`` (one per LUID-named adapter), or
|
||
``None`` when the counter is unavailable/localized/empty so callers fall back.
|
||
"""
|
||
if platform.system() != "Windows":
|
||
return None
|
||
try:
|
||
# Emit "<InstanceName>|<CookedValue>" per sample, or a __NONE__ sentinel.
|
||
ps = (
|
||
"$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'"
|
||
" -ErrorAction SilentlyContinue).CounterSamples;"
|
||
"if($s){$s|ForEach-Object{'{0}|{1}' -f $_.InstanceName,[int64]$_.CookedValue}}"
|
||
"else{'__NONE__'}"
|
||
)
|
||
r = subprocess.run(
|
||
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
|
||
capture_output = True,
|
||
text = True,
|
||
encoding = "utf-8",
|
||
errors = "replace",
|
||
timeout = 5,
|
||
)
|
||
if r.returncode != 0 or not r.stdout.strip():
|
||
return None
|
||
adapters: list[tuple[str, float]] = []
|
||
for line in r.stdout.splitlines():
|
||
line = line.strip()
|
||
if not line or line == "__NONE__" or "|" not in line:
|
||
continue
|
||
instance, _, raw = line.rpartition("|")
|
||
try:
|
||
used = float(raw.strip())
|
||
except (ValueError, TypeError):
|
||
continue
|
||
if used < 0:
|
||
continue
|
||
adapters.append((instance.strip(), used))
|
||
return adapters or None
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _match_adapter_used_to_devices(
|
||
adapter_useds: list[float], device_totals: list[float]
|
||
) -> list[Optional[float]]:
|
||
"""Attribute per-adapter used bytes to torch devices by capacity ranking.
|
||
|
||
Windows shares no key between LUID counters and torch ordinals, so usages are
|
||
ranked against device totals and each is trusted only when capacity *forces* it
|
||
(it exceeds every smaller device); an ambiguous ranking reports unknown
|
||
(``None``) rather than fabricate a per-index free.
|
||
|
||
Extra counters mean a hidden/display adapter, and the noise filter may have
|
||
dropped a real reading, so values are emitted only when the supra-threshold
|
||
counters number EXACTLY the visible devices AND capacity forces the mapping;
|
||
otherwise every device is unknown. Best-effort but correct for the common
|
||
loaded-card case (#7072). Returns a list aligned to ``device_totals``.
|
||
"""
|
||
n = len(device_totals)
|
||
if n == 0:
|
||
return []
|
||
useds = sorted(adapter_useds, reverse = True)
|
||
ranked_positions = sorted(range(n), key = lambda i: -device_totals[i])
|
||
ranked_totals = [device_totals[pos] for pos in ranked_positions]
|
||
assigned: list[Optional[float]]
|
||
# More counters than devices -> a hidden/display adapter (check before noise filter).
|
||
if len(useds) > n:
|
||
non_trivial = [u for u in useds if u >= _ROCM_WIN_ADAPTER_MIN_BYTES]
|
||
if len(non_trivial) != n:
|
||
# Not a clean bijection (a masked GPU is busy or a visible card idle):
|
||
# no counter maps to a specific card, so report unknown.
|
||
return [None] * n
|
||
# Exactly n supra-threshold counters: extras were placeholders, so a
|
||
# capacity-ranked bijection is plausible.
|
||
useds = non_trivial
|
||
ranked_useds = [useds[rank] for rank in range(n)]
|
||
# A usage above its ranked capacity is a hidden larger GPU; clamping onto the
|
||
# smaller card would fabricate a fully-used reading.
|
||
for rank in range(n):
|
||
if ranked_useds[rank] > ranked_totals[rank]:
|
||
return [None] * n
|
||
# Capacity forces the mapping only when the usage exceeds the next-smaller
|
||
# capacity; the smallest card and merely-fitting usages stay unknown.
|
||
# Keeps 40 GiB over 48/8 GiB -> [40, None].
|
||
assigned = [None] * n
|
||
for rank, pos in enumerate(ranked_positions):
|
||
if rank + 1 < n and ranked_useds[rank] > ranked_totals[rank + 1]:
|
||
assigned[pos] = min(ranked_useds[rank], device_totals[pos])
|
||
return assigned
|
||
# No hidden adapters: every counter is a visible card, so ranking is a permutation.
|
||
ranked_useds = [useds[rank] if rank < len(useds) else 0.0 for rank in range(n)]
|
||
# Ambiguous if a strictly larger usage also fits the next smaller card: the two
|
||
# could be swapped without breaking capacity, so ranking can't tell them apart.
|
||
for rank in range(n - 1):
|
||
upper, lower = ranked_useds[rank], ranked_useds[rank + 1]
|
||
if upper > lower and upper <= ranked_totals[rank + 1]:
|
||
return [None] * n
|
||
assigned = [None] * n
|
||
for rank, pos in enumerate(ranked_positions):
|
||
if rank < len(useds):
|
||
assigned[pos] = min(useds[rank], device_totals[pos])
|
||
return assigned
|
||
|
||
|
||
def _rocm_windows_per_device_vram(device_indices: list[int]) -> list[Dict[str, Any]]:
|
||
"""Per-GPU VRAM on Windows AMD/ROCm: total from torch properties (reliable),
|
||
used from the per-adapter Dedicated Usage counter.
|
||
|
||
Returns ``{index, visible_ordinal, name, used_gb, total_gb}`` per visible GPU
|
||
(``used_gb`` may be ``None`` when the counter is unavailable), or ``[]`` when
|
||
torch can't enumerate devices so callers fall through to the torch last resort.
|
||
"""
|
||
if platform.system() != "Windows":
|
||
return []
|
||
mod, _ = _torch_get_device_module()
|
||
if mod is None:
|
||
return []
|
||
# Totals/names from torch properties (mem_get_info's free==total quirk zeroes used).
|
||
dev_meta: list[Dict[str, Any]] = []
|
||
for ordinal, phys_idx in enumerate(device_indices):
|
||
try:
|
||
props = mod.get_device_properties(ordinal)
|
||
dev_meta.append(
|
||
{
|
||
"index": phys_idx,
|
||
"visible_ordinal": ordinal,
|
||
"name": props.name,
|
||
"total_bytes": int(props.total_memory),
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.debug("torch property probe failed for ordinal %d: %s", ordinal, e)
|
||
if not dev_meta:
|
||
return []
|
||
|
||
adapters = _rocm_windows_perf_counter_vram_by_adapter()
|
||
if adapters:
|
||
assigned = _match_adapter_used_to_devices(
|
||
[used for _, used in adapters],
|
||
[d["total_bytes"] for d in dev_meta],
|
||
)
|
||
else:
|
||
# Counter unavailable: show every GPU with a correct total, used unknown.
|
||
assigned = [None] * len(dev_meta)
|
||
|
||
devices: list[Dict[str, Any]] = []
|
||
for meta, used_bytes in zip(dev_meta, assigned):
|
||
total_gb = round(meta["total_bytes"] / (1024**3), 2)
|
||
used_gb = round(used_bytes / (1024**3), 2) if used_bytes is not None else None
|
||
devices.append(
|
||
{
|
||
"index": meta["index"],
|
||
"visible_ordinal": meta["visible_ordinal"],
|
||
"name": meta["name"],
|
||
"used_gb": used_gb,
|
||
"total_gb": total_gb,
|
||
}
|
||
)
|
||
return devices
|
||
|
||
|
||
def _rocm_windows_device_payload_entry(
|
||
device: DeviceType, dev: Dict[str, Any], gpu_util_pct: Optional[float]
|
||
) -> Dict[str, Any]:
|
||
"""Build a ``get_gpu_utilization`` device entry from a per-device VRAM dict."""
|
||
total_gb = dev["total_gb"]
|
||
used_gb = dev["used_gb"]
|
||
return {
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"index": dev["index"],
|
||
"visible_ordinal": dev["visible_ordinal"],
|
||
"name": dev.get("name", "Unknown"),
|
||
"gpu_utilization_pct": gpu_util_pct,
|
||
"temperature_c": None,
|
||
"vram_used_gb": used_gb,
|
||
"vram_total_gb": total_gb,
|
||
"vram_utilization_pct": round((used_gb / total_gb) * 100, 1)
|
||
if total_gb and total_gb > 0 and used_gb is not None
|
||
else None,
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
|
||
|
||
def _gpu_utilization_payload(
|
||
device: DeviceType, devices: list[Dict[str, Any]], **metadata: Any
|
||
) -> Dict[str, Any]:
|
||
"""Keep the legacy primary-GPU shape and append all visible devices."""
|
||
backend = _backend_label(device)
|
||
normalized = []
|
||
for ordinal, raw in enumerate(devices):
|
||
dev = dict(raw)
|
||
dev.setdefault("available", True)
|
||
dev.setdefault("backend", backend)
|
||
if dev.get("visible_ordinal") is None:
|
||
dev["visible_ordinal"] = ordinal
|
||
normalized.append(dev)
|
||
|
||
normalized.sort(key = lambda dev: dev.get("visible_ordinal", dev.get("index", 0)))
|
||
payload: Dict[str, Any] = {
|
||
"available": bool(normalized),
|
||
"backend": backend,
|
||
"devices": normalized,
|
||
}
|
||
payload.update(metadata)
|
||
if normalized:
|
||
payload.update(normalized[0])
|
||
payload["available"] = True
|
||
payload["backend"] = normalized[0].get("backend", backend)
|
||
payload["devices"] = normalized
|
||
return payload
|
||
|
||
|
||
def get_gpu_utilization() -> Dict[str, Any]:
|
||
"""Live utilization snapshot for the primary GPU plus all visible GPUs."""
|
||
device = get_device()
|
||
|
||
if device == DeviceType.XPU:
|
||
result = get_visible_gpu_utilization()
|
||
return _gpu_utilization_payload(
|
||
device,
|
||
result.get("devices", []),
|
||
parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []),
|
||
index_kind = result.get("index_kind"),
|
||
)
|
||
|
||
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 and "devices" in result:
|
||
devices = result["devices"]
|
||
numeric_ids = parent_visible_spec.get("numeric_ids")
|
||
if IS_ROCM and numeric_ids is not None:
|
||
_reconcile_rocm_unified_memory(result, numeric_ids)
|
||
|
||
return _gpu_utilization_payload(
|
||
device,
|
||
devices,
|
||
backend_cuda_visible_devices = result.get("backend_cuda_visible_devices"),
|
||
parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []),
|
||
index_kind = result.get("index_kind"),
|
||
)
|
||
|
||
# Fallback Windows ROCm: per-adapter VRAM attribution (issue #7072), so
|
||
# every visible GPU is shown instead of a sum collapsed onto one device.
|
||
if IS_ROCM and platform.system() == "Windows":
|
||
_win_ids = _get_parent_visible_gpu_spec().get("numeric_ids")
|
||
if not _win_ids:
|
||
_win_ids = list(range(_torch_get_physical_gpu_count() or 0))
|
||
_win_devices = _rocm_windows_per_device_vram(_win_ids)
|
||
if _win_devices:
|
||
# A single visible GPU can own the aggregate 3D-engine utilization;
|
||
# across several GPUs the sum isn't per-device, so leave it unset.
|
||
_win_util = (
|
||
_rocm_windows_perf_counter_gpu_util_pct() if len(_win_devices) == 1 else None
|
||
)
|
||
return _gpu_utilization_payload(
|
||
device,
|
||
[
|
||
_rocm_windows_device_payload_entry(device, _wd, _win_util)
|
||
for _wd in _win_devices
|
||
],
|
||
)
|
||
|
||
# Fallback Linux ROCm
|
||
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 _gpu_utilization_payload(
|
||
device,
|
||
[
|
||
{
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"index": 0,
|
||
"visible_ordinal": 0,
|
||
"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) for all visible GPUs
|
||
_visible_spec = _get_parent_visible_gpu_spec()
|
||
_numeric_ids = _visible_spec.get("numeric_ids") or []
|
||
if not _numeric_ids:
|
||
visible_count = _torch_get_physical_gpu_count() or 0
|
||
_numeric_ids = list(range(visible_count))
|
||
|
||
_torch_devices = _torch_get_per_device_info(_numeric_ids)
|
||
if _torch_devices:
|
||
gpu_array = []
|
||
for _td in _torch_devices:
|
||
_total = _td["total_gb"]
|
||
_used = _td["used_gb"]
|
||
gpu_array.append(
|
||
{
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"index": _td["index"],
|
||
"name": _td.get("name", "Unknown"),
|
||
"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 and _used is not None
|
||
else None,
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
)
|
||
return _gpu_utilization_payload(device, gpu_array)
|
||
|
||
# MLX
|
||
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, "devices": [], "error": str(e)}
|
||
|
||
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 _gpu_utilization_payload(
|
||
device,
|
||
[
|
||
{
|
||
"available": True,
|
||
"backend": device.value,
|
||
"index": 0,
|
||
"visible_ordinal": 0,
|
||
"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 _gpu_utilization_payload(
|
||
device,
|
||
[
|
||
{
|
||
"available": True,
|
||
"backend": _backend_label(device),
|
||
"index": mem.get("device", 0),
|
||
"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,
|
||
}
|
||
],
|
||
)
|
||
|
||
return {"available": False, "backend": _backend_label(device), "devices": []}
|
||
|
||
|
||
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"]
|
||
torch_used_gb = torch_info.get("used_gb")
|
||
smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
|
||
# torch sees the full unified (GTT) pool; amd-smi only the dedicated carve-out.
|
||
# Adopt torch's larger total regardless of used: on Windows ROCm torch_used is
|
||
# None (free==total sentinel) but its total stays authoritative. Overwrite used
|
||
# only when torch's is known, then recompute utilization against whatever remains.
|
||
if torch_total_gb > smi_total_gb:
|
||
device_metrics["vram_total_gb"] = torch_total_gb
|
||
if torch_used_gb is not None:
|
||
device_metrics["vram_used_gb"] = torch_used_gb
|
||
_used_for_pct = device_metrics.get("vram_used_gb")
|
||
device_metrics["vram_utilization_pct"] = (
|
||
round((_used_for_pct / torch_total_gb) * 100, 1)
|
||
if torch_total_gb > 0 and _used_for_pct is not None
|
||
else None
|
||
)
|
||
logger.debug(
|
||
"ROCm unified memory: adopted torch mem_get_info total (%.2f GB) over "
|
||
"amd-smi (%.2f GB) for device %s",
|
||
torch_total_gb,
|
||
smi_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 _rocm_visibility_mask_active() -> bool:
|
||
"""True when any ROCm/CUDA visibility variable filters the device set."""
|
||
for var in (
|
||
"HIP_VISIBLE_DEVICES",
|
||
"ROCR_VISIBLE_DEVICES",
|
||
"CUDA_VISIBLE_DEVICES",
|
||
"GPU_DEVICE_ORDINAL",
|
||
):
|
||
value = os.environ.get(var)
|
||
if value and value.strip():
|
||
return True
|
||
return False
|
||
|
||
|
||
def _overlay_system_wide_vram(devices: list[Dict[str, Any]]) -> None:
|
||
"""Replace process-local torch VRAM with system-wide Linux ROCm figures.
|
||
|
||
The torch fallback is process-local, so a model served by the separate
|
||
llama-server process reads as ~0 used even with the GPU full (#7072). DRM
|
||
sysfs gives per-card figures the kernel updates across all processes. Sources
|
||
are matched by the device's PHYSICAL index (never list position), and only
|
||
when NO visibility mask is active and the device count equals the host GPU
|
||
count; under any mask the index is not a verifiable host ordinal, so torch's
|
||
figures are kept. Best-effort, in place: a device with no matching card, or a
|
||
unified-memory APU whose sysfs total is below torch's GTT-backed total, keeps
|
||
torch's (mirrors _apply_unified_memory_correction).
|
||
|
||
Windows is intentionally not overlaid: its per-adapter perf counters cannot be
|
||
mapped to ROCm ordinals and miss WDDM shared memory, so the multi-GPU view
|
||
keeps torch there rather than risk misattributing another adapter's usage.
|
||
"""
|
||
if not devices or platform.system() != "Linux":
|
||
return
|
||
# Match by PCI identity, never list position: index N in KFD topology is ROCm
|
||
# physical device N and carries its PCI address, which DRM sysfs keys on too.
|
||
# The two gates below verify ``index`` really is a host-physical ordinal
|
||
# (torch exposes no PCI id to check directly):
|
||
# * No visibility mask -- any mask makes ``index`` container/ROCR-relative
|
||
# rather than a host ordinal.
|
||
# * Device count == host GPU count -- rules out a device-cgroup container
|
||
# that sets no env var yet compacts torch's indices from zero.
|
||
pci_by_ordinal = _rocm_kfd_gpu_pci_ids()
|
||
if not pci_by_ordinal:
|
||
return
|
||
if _rocm_visibility_mask_active() or len(devices) != len(pci_by_ordinal):
|
||
return
|
||
vram_by_pci = _rocm_linux_sysfs_vram_by_pci_gb()
|
||
for dev in devices:
|
||
index = dev.get("index")
|
||
if not isinstance(index, int) or not (0 <= index < len(pci_by_ordinal)):
|
||
continue
|
||
entry = vram_by_pci.get(pci_by_ordinal[index].lower())
|
||
if entry is None:
|
||
continue
|
||
used, total = entry
|
||
dev_total = dev.get("vram_total_gb") or 0.0
|
||
# Overlay only a device that maps 1:1 to the whole card: torch total must
|
||
# match sysfs total within ~10%. A mismatch either way means a different
|
||
# memory scope -- a unified-memory APU (sysfs sees only the dedicated
|
||
# slice, torch the GTT pool) or a partitioned MI300 (sysfs reports the
|
||
# whole card, dwarfing a partition) -- and overlaying would misstate free
|
||
# VRAM (a partition would look like it has the whole card free).
|
||
if dev_total <= 0 or abs(total - dev_total) > 0.1 * dev_total:
|
||
continue
|
||
dev["vram_used_gb"] = used
|
||
dev["vram_total_gb"] = total
|
||
dev["vram_utilization_pct"] = round((used / total) * 100, 1) if total > 0 else None
|
||
|
||
|
||
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
|
||
|
||
# Windows AMD/ROCm (issue #7072): the System tab's VRAM source. The torch
|
||
# fallback below would report used==0 (free==total), so read per-adapter
|
||
# Dedicated Usage instead; total from torch properties.
|
||
if IS_ROCM and platform.system() == "Windows":
|
||
win_numeric_ids = parent_visible_spec.get("numeric_ids")
|
||
if win_numeric_ids:
|
||
win_ids = win_numeric_ids
|
||
win_index_kind = "physical"
|
||
else:
|
||
win_ids = list(range(_torch_get_physical_gpu_count() or 0))
|
||
win_index_kind = "relative"
|
||
win_devices = _rocm_windows_per_device_vram(win_ids)
|
||
if win_devices:
|
||
devices = []
|
||
for wd in win_devices:
|
||
total = wd["total_gb"]
|
||
used = wd["used_gb"]
|
||
devices.append(
|
||
{
|
||
"index": wd["index"],
|
||
"index_kind": win_index_kind,
|
||
"visible_ordinal": wd["visible_ordinal"],
|
||
"name": wd.get("name"),
|
||
"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 and total > 0 and used is not None
|
||
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": win_numeric_ids or [],
|
||
"devices": devices,
|
||
"index_kind": win_index_kind,
|
||
}
|
||
|
||
# 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"]
|
||
# used=None is a deliberate "telemetry unavailable" signal
|
||
# from _torch_get_per_device_info (e.g. XPU without
|
||
# mem_get_info); propagate None instead of dividing by it. On
|
||
# CUDA/ROCm used is always an int, so this stays byte-identical.
|
||
vram_pct = (
|
||
round((used / total) * 100, 1) if used is not None and total > 0 else None
|
||
)
|
||
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": vram_pct,
|
||
"power_draw_w": None,
|
||
"power_limit_w": None,
|
||
"power_utilization_pct": None,
|
||
}
|
||
)
|
||
if IS_ROCM and index_kind == "physical":
|
||
# Swap process-local torch VRAM for system-wide sysfs so a model
|
||
# held by the separate llama-server process shows up (#7072).
|
||
# Physical-index only: a relative index (UUID/MIG mask) is not a
|
||
# host GPU id. The overlay verifies the rest itself.
|
||
_overlay_system_wide_vram(devices)
|
||
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": "vulkan",
|
||
}
|
||
|
||
|
||
# ========== 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]:
|
||
# On Intel XPU, visibility is controlled by ZE_AFFINITY_MASK (Level Zero),
|
||
# not CUDA_VISIBLE_DEVICES.
|
||
if get_device() == DeviceType.XPU:
|
||
xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK")
|
||
composite = _xpu_hierarchy_is_composite()
|
||
|
||
if xpu_mask_raw is None:
|
||
# COMPOSITE: root GPU IDs are stable physical IDs.
|
||
if composite:
|
||
return {
|
||
"raw": None,
|
||
"numeric_ids": list(range(get_physical_gpu_count())),
|
||
"supports_explicit_gpu_ids": True,
|
||
}
|
||
# FLAT (oneAPI default): ordinals are tile/device handles, not
|
||
# physical GPU IDs. numeric_ids=None so telemetry uses relative
|
||
# ordinals; explicit selection needs ZE_FLAT_DEVICE_HIERARCHY=COMPOSITE.
|
||
return {
|
||
"raw": None,
|
||
"numeric_ids": None,
|
||
"supports_explicit_gpu_ids": False,
|
||
}
|
||
|
||
xpu_mask = xpu_mask_raw.strip()
|
||
if xpu_mask == "":
|
||
return {
|
||
"raw": xpu_mask,
|
||
"numeric_ids": [],
|
||
"supports_explicit_gpu_ids": True,
|
||
}
|
||
|
||
# Subdevice syntax ("N.M") expands one root into multiple
|
||
# logical devices -- not addressable by explicit root-ID selection.
|
||
has_subdevice = any("." in token.strip() for token in xpu_mask.split(",") if token.strip())
|
||
if has_subdevice:
|
||
return {
|
||
"raw": xpu_mask,
|
||
"numeric_ids": None,
|
||
"supports_explicit_gpu_ids": False,
|
||
}
|
||
|
||
# FLAT numeric entries are tile handles, not physical GPU IDs. Keep
|
||
# numeric_ids unresolved so every telemetry and picker consumer uses
|
||
# relative torch ordinals and cannot advertise them as pinnable roots.
|
||
if not composite:
|
||
tokens = [token.strip() for token in xpu_mask.split(",") if token.strip()]
|
||
if tokens and all(token.isdecimal() for token in tokens):
|
||
return {
|
||
"raw": xpu_mask,
|
||
"numeric_ids": None,
|
||
"supports_explicit_gpu_ids": False,
|
||
}
|
||
return {
|
||
"raw": xpu_mask,
|
||
"numeric_ids": None,
|
||
"supports_explicit_gpu_ids": False,
|
||
}
|
||
|
||
# COMPOSITE + pure numeric (subdevice handled above). _parse_ze_mask_roots
|
||
# maps to root GPU IDs, dropping non-decimal tokens so "*"/"GPU-uuid" -> [].
|
||
roots_with_dupes = _parse_ze_mask_roots(xpu_mask)
|
||
if not roots_with_dupes:
|
||
# Unparseable mask (e.g. "*", "GPU-uuid") -- cannot map to
|
||
# physical root IDs.
|
||
return {
|
||
"raw": xpu_mask,
|
||
"numeric_ids": None,
|
||
"supports_explicit_gpu_ids": False,
|
||
}
|
||
|
||
return {
|
||
"raw": xpu_mask,
|
||
"numeric_ids": roots_with_dupes,
|
||
"supports_explicit_gpu_ids": True,
|
||
}
|
||
|
||
# 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]], *, is_vulkan: bool = False
|
||
) -> 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 [] if is_vulkan else parent_visible_ids
|
||
|
||
requested_ids = list(gpu_ids)
|
||
if len(requested_ids) == 0:
|
||
return [] if is_vulkan else parent_visible_ids
|
||
|
||
if is_vulkan:
|
||
# A Vulkan build selects by ggml Vulkan ordinal (--device VulkanN), a separate
|
||
# index space from CUDA/ROCm ids that may be empty under CPU-only torch. The
|
||
# CUDA parent-visible / physical-count checks below do not apply; only reject
|
||
# malformed ordinals (issue #7239).
|
||
if len(set(requested_ids)) != len(requested_ids):
|
||
raise ValueError(f"Invalid gpu_ids {requested_ids}: duplicate GPU IDs are not allowed.")
|
||
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}."
|
||
)
|
||
return requested_ids
|
||
|
||
if not parent_visible_spec["supports_explicit_gpu_ids"]:
|
||
env_var_name = (
|
||
"ZE_AFFINITY_MASK" if get_device() == DeviceType.XPU else "CUDA_VISIBLE_DEVICES"
|
||
)
|
||
raise ValueError(
|
||
f"Invalid gpu_ids {requested_ids}: explicit physical GPU IDs are "
|
||
f"unsupported when {env_var_name} uses non-numeric or subdevice "
|
||
f"entries ({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")
|
||
# Skip intermediate training checkpoints: a run dir can hold several
|
||
# checkpoint-*/global_step* snapshots, but export loads only the model at
|
||
# the root, so counting them would multiply the estimate.
|
||
skip_prefixes = ("checkpoint-", "global_step")
|
||
total = 0
|
||
for file in model_path.rglob("*"):
|
||
if not file.is_file() or file.suffix not in weight_exts:
|
||
continue
|
||
rel = file.relative_to(model_path)
|
||
if any(part.startswith(skip_prefixes) for part in rel.parts):
|
||
continue
|
||
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):
|
||
# Estimation needs only declarative config.json fields, and this probe runs
|
||
# on model selection, so read raw config.json (never run auto_map Python) and
|
||
# expose it as an attribute namespace for downstream getattr access.
|
||
try:
|
||
from utils.transformers_version import _load_config_json
|
||
|
||
cfg = _load_config_json(model_name, hf_token = hf_token)
|
||
if cfg is None:
|
||
return None
|
||
|
||
def _to_ns(d):
|
||
if isinstance(d, dict):
|
||
return types.SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()})
|
||
return d
|
||
|
||
return _to_ns(cfg)
|
||
except Exception as e:
|
||
# A 5.x-only config can't be parsed by the default transformers; that is
|
||
# expected (the worker reloads under the sidecar), so only warn for default tier.
|
||
tier = "default"
|
||
try:
|
||
from utils.transformers_version import get_transformers_tier
|
||
tier = get_transformers_tier(model_name)
|
||
except Exception:
|
||
pass
|
||
if tier != "default":
|
||
_tier_version = {"510": "5.10.x", "530": "5.3.0", "550": "5.5.0"}.get(tier, "5.x")
|
||
logger.info(
|
||
"Config for '%s' not parseable by the default transformers; "
|
||
"needs transformers %s and will be loaded with that sidecar in the worker",
|
||
model_name,
|
||
_tier_version,
|
||
)
|
||
else:
|
||
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"}
|
||
|
||
# Auto-selection needs per-device free-VRAM telemetry, available on CUDA
|
||
# (nvidia-smi) and XPU (torch.xpu) but not MLX/CPU, which fall
|
||
# through to inheriting parent visibility.
|
||
if get_device() not in (DeviceType.CUDA, DeviceType.XPU):
|
||
metadata["selection_mode"] = "non_accelerator"
|
||
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=%s selected=%s usable_gb=%s "
|
||
"required_gb=%s multi_gpu_overhead=%s",
|
||
model_name,
|
||
selected,
|
||
metadata["usable_gb"],
|
||
metadata.get("required_gb"),
|
||
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=%s "
|
||
"selected=%s usable_gb=%s required_gb=%s multi_gpu_overhead=%s",
|
||
model_name,
|
||
fallback_all,
|
||
metadata["usable_gb"],
|
||
metadata.get("required_gb"),
|
||
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() not in (DeviceType.CUDA, DeviceType.XPU):
|
||
raise ValueError(
|
||
f"gpu_ids {list(gpu_ids)} is only supported on CUDA and Intel XPU "
|
||
f"devices, 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 XPU the control is ``ZE_AFFINITY_MASK`` (not ``CUDA_VISIBLE_DEVICES``);
|
||
on ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over
|
||
CUDA_VISIBLE_DEVICES. Mirrors ``_get_parent_visible_gpu_spec`` so
|
||
``backend_cuda_visible_devices`` reports the value actually narrowing the
|
||
visible device set on the current backend.
|
||
"""
|
||
if get_device() == DeviceType.XPU:
|
||
return os.environ.get("ZE_AFFINITY_MASK")
|
||
if IS_ROCM:
|
||
return _get_parent_visible_gpu_spec().get("raw")
|
||
return os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
|
||
|
||
def get_vulkan_inference_gpu_info() -> Optional[Dict[str, Any]]:
|
||
"""Return llama.cpp Vulkan devices, or None when Vulkan is not installed."""
|
||
# Vulkan is a llama.cpp inference backend, not a PyTorch training device, so
|
||
# keep it separate from the PyTorch/MLX training-device report.
|
||
try:
|
||
from core.inference.llama_cpp import LlamaCppBackend
|
||
except Exception as e:
|
||
logger.debug("Could not inspect the llama.cpp Vulkan backend: %s", e)
|
||
return None
|
||
|
||
try:
|
||
if not LlamaCppBackend._is_vulkan_backend():
|
||
return None
|
||
except Exception as e:
|
||
logger.debug("Could not identify the llama.cpp Vulkan backend: %s", e)
|
||
return None
|
||
|
||
result = {
|
||
"available": False,
|
||
"backend": "vulkan",
|
||
"backend_cuda_visible_devices": None,
|
||
"parent_visible_gpu_ids": [],
|
||
"devices": [],
|
||
"index_kind": "vulkan",
|
||
}
|
||
# Identity (real device description, explicit iGPU flag) comes from the
|
||
# inventory; the memory numbers stay on _get_gpu_memory, which applies the
|
||
# iGPU host reserve and zeroes a shared total. Budgeting an APU off the raw
|
||
# shared total instead would hand out the whole machine's RAM with no OS
|
||
# headroom. Join by ordinal; a probe failure just leaves names unresolved.
|
||
identity: Dict[int, Dict[str, Any]] = {}
|
||
try:
|
||
identity = {row["index"]: row for row in LlamaCppBackend.vulkan_device_inventory()}
|
||
except Exception as e:
|
||
logger.debug("Vulkan device inventory failed, falling back to ordinals: %s", e)
|
||
|
||
try:
|
||
for ordinal, free_mib, total_mib in LlamaCppBackend._get_gpu_memory():
|
||
info = identity.get(ordinal, {})
|
||
# _get_gpu_memory reports total 0 for a shared pool; prefer the
|
||
# explicit flag when the inventory resolved this ordinal.
|
||
shared_memory = bool(info["is_igpu"]) if "is_igpu" in info else total_mib == 0
|
||
budget_mib = total_mib or free_mib
|
||
used_mib = max(0, total_mib - free_mib) if total_mib else None
|
||
result["devices"].append(
|
||
{
|
||
"index": ordinal,
|
||
# ggml Vulkan ordinals are the space `--device Vulkan<i>` pins,
|
||
# so unlike a torch-xpu relative ordinal these are selectable.
|
||
"index_kind": "vulkan",
|
||
"visible_ordinal": ordinal,
|
||
"name": info.get("name") or f"Vulkan{ordinal}",
|
||
"memory_total_gb": round(budget_mib / 1024, 2),
|
||
"vram_used_gb": round(used_mib / 1024, 2) if used_mib is not None else None,
|
||
"vram_free_gb": round(free_mib / 1024, 2),
|
||
"vram_utilization_pct": round((used_mib / total_mib) * 100, 1)
|
||
if used_mib is not None and total_mib > 0
|
||
else None,
|
||
"shared_memory": shared_memory,
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.debug("Vulkan GPU visibility query failed: %s", e)
|
||
return result
|
||
|
||
result["available"] = bool(result["devices"])
|
||
return result
|
||
|
||
|
||
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": "vulkan",
|
||
}
|
||
|
||
|
||
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
|
||
|
||
# Prefer torch.xpu.device_count() on Intel XPU: the Level Zero runtime
|
||
# correctly interprets ZE_AFFINITY_MASK semantics (e.g. subdevice syntax
|
||
# "0.0,0.1" collapses onto one root GPU). Supersedes the torch fallback below.
|
||
if get_device() == DeviceType.XPU:
|
||
xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK")
|
||
xpu_mask_set = xpu_mask_raw is not None
|
||
xpu_visible = (xpu_mask_raw or "").strip()
|
||
if xpu_mask_set and xpu_visible == "":
|
||
_visible_gpu_count = 0
|
||
return _visible_gpu_count
|
||
|
||
try:
|
||
import torch
|
||
_visible_gpu_count = torch.xpu.device_count()
|
||
except Exception as e:
|
||
logger.debug(
|
||
"torch.xpu.device_count() failed, falling back to mask parsing: %s",
|
||
e,
|
||
)
|
||
if xpu_visible:
|
||
# Fallback: count unique root device IDs from the mask.
|
||
# "device.subdevice" notation means "0.0,0.1" is 1 root, not 2.
|
||
# Without torch the hierarchy mode is unknown, so root-device
|
||
# counting is the conservative choice.
|
||
if xpu_visible == "*":
|
||
# Documented wildcard: all physical XPUs visible.
|
||
_visible_gpu_count = get_physical_gpu_count()
|
||
else:
|
||
roots = _parse_ze_mask_roots(xpu_visible)
|
||
# Non-parseable masks (",,,", "GPU-abc") yield an empty
|
||
# roots list, treated as 0 visible devices, not "all
|
||
# visible" -- no evidence the whole fleet was intended.
|
||
_visible_gpu_count = len(set(roots))
|
||
else:
|
||
_visible_gpu_count = get_physical_gpu_count()
|
||
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. XPU is
|
||
# handled by the early return above, so only torch.cuda is needed here.
|
||
try:
|
||
import torch
|
||
_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, backend: Optional[str] = None) -> 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)
|
||
|
||
# Intel XPU honors ZE_AFFINITY_MASK, not CUDA_VISIBLE_DEVICES; route XPU
|
||
# pinning through it so worker subprocesses are restricted to the intended GPU.
|
||
# Decide WITHOUT get_device(): workers call this before detect_hardware(),
|
||
# and a lazy detect would probe torch.cuda against the unmasked parent env,
|
||
# latching device enumeration before the mask below is written. Pre-detect,
|
||
# use env + torch BUILD attributes only (no runtime init, like the ROCm
|
||
# mirror below).
|
||
_is_xpu = DEVICE == DeviceType.XPU
|
||
if backend is not None:
|
||
# The spawning parent's detected backend (config["device_backend"]):
|
||
# exact and probe-free, so the mask target always matches what
|
||
# detect_hardware() decided in the parent, including its XPU
|
||
# availability check and CUDA fallback.
|
||
_is_xpu = backend == DeviceType.XPU.value
|
||
elif DEVICE is None:
|
||
# No parent backend passed (direct caller). version.xpu can be None
|
||
# on a working XPU build, so also accept torch.xpu._is_compiled()
|
||
# (a pure symbol-presence check, no runtime init). UNSLOTH_FORCE_XPU
|
||
# counts only on an XPU-capable build: detect_hardware() falls back
|
||
# to CUDA when XPU is missing, and the mask target must follow.
|
||
try:
|
||
import torch as _torch
|
||
|
||
_ver = _torch.version
|
||
_is_comp = getattr(getattr(_torch, "xpu", None), "_is_compiled", None)
|
||
_xpu_build = (callable(_is_comp) and bool(_is_comp())) or (
|
||
getattr(_ver, "xpu", None) is not None
|
||
)
|
||
if os.environ.get("UNSLOTH_FORCE_XPU") == "1":
|
||
_is_xpu = _xpu_build
|
||
else:
|
||
# Mirror detect_hardware: hidden CUDA prefers XPU on an
|
||
# XPU-capable build (with or without a ZE mask -- detection
|
||
# falls through to XPU either way), where writing these ids
|
||
# to CUDA_VISIBLE_DEVICES would re-expose the deliberately
|
||
# hidden CUDA.
|
||
_cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
_cuda_hidden = _cvd is not None and _cvd.strip() in ("", "-1")
|
||
_is_xpu = _xpu_build and (
|
||
_cuda_hidden
|
||
or (getattr(_ver, "cuda", None) is None and getattr(_ver, "hip", None) is None)
|
||
)
|
||
except Exception as e:
|
||
logger.debug(
|
||
"apply_gpu_ids: torch XPU probe skipped (%s: %s)",
|
||
type(e).__name__,
|
||
e,
|
||
)
|
||
if _is_xpu:
|
||
os.environ["ZE_AFFINITY_MASK"] = value
|
||
# Leave inherited CUDA_VISIBLE_DEVICES alone -- clearing it could let
|
||
# the worker flip back to CUDA on hybrid hosts.
|
||
_visible_gpu_count = None
|
||
logger.info("Applied gpu_ids: ZE_AFFINITY_MASK='%s'", value)
|
||
return
|
||
|
||
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``/``ZE_AFFINITY_MASK`` uses non-numeric
|
||
identifiers (UUID/MIG/wildcard) and >1 GPU is visible (fallback:
|
||
numeric IDs unresolvable, so assume multi-GPU is intended).
|
||
|
||
Returns ``"sequential"`` (single device) otherwise, including CPU/MLX
|
||
backends.
|
||
|
||
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 in (DeviceType.CUDA, DeviceType.XPU):
|
||
multi_gpu = gpu_ids is not None and len(gpu_ids) > 1
|
||
|
||
if not multi_gpu:
|
||
parent_visible_spec = _get_parent_visible_gpu_spec()
|
||
if device == DeviceType.CUDA:
|
||
# UUID/MIG masks can't be split into numeric IDs; >1 visible GPU
|
||
# means multi-GPU sharding is intended.
|
||
if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1:
|
||
multi_gpu = True
|
||
elif device == DeviceType.XPU and gpu_ids is None:
|
||
# Shard across visible XPU ordinals via HF (no mask rewrite),
|
||
# only when no gpu_ids were passed -- an explicit gpu_ids=[0]
|
||
# means "use exactly device 0" and must stay sequential.
|
||
supports_physical = parent_visible_spec["supports_explicit_gpu_ids"]
|
||
has_multiple_numeric = (
|
||
parent_visible_spec["numeric_ids"] is not None
|
||
and len(parent_visible_spec["numeric_ids"]) > 1
|
||
)
|
||
has_multiple_unresolved = (
|
||
parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1
|
||
)
|
||
if has_multiple_unresolved or (not supports_physical and has_multiple_numeric):
|
||
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 get_torch_device_str() -> str:
|
||
"""
|
||
Return the torch device string for the detected hardware.
|
||
E.g. "cuda", "xpu", or "cpu".
|
||
"""
|
||
device = get_device()
|
||
if device == DeviceType.CUDA:
|
||
return "cuda"
|
||
elif device == DeviceType.XPU:
|
||
return "xpu"
|
||
return "cpu"
|
||
|
||
|
||
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.
|
||
|
||
Also returns ``None`` on XPU once its runtime is initialized in this
|
||
process: ``os.fork()`` corrupts the Level-Zero context, making Triton
|
||
kernels fail with "Pointer argument doesn't reference XPU device memory".
|
||
Pre-init XPU hosts can still parallelize CPU-side preprocessing.
|
||
"""
|
||
if sys.platform in ("win32", "darwin"):
|
||
return None
|
||
|
||
if get_device() == DeviceType.XPU:
|
||
try:
|
||
import torch
|
||
except Exception:
|
||
# No torch means no active XPU runtime, so CPU-side dataset
|
||
# parallelism is still safe.
|
||
return safe_num_proc(desired)
|
||
|
||
xpu = getattr(torch, "xpu", None)
|
||
is_initialized = getattr(xpu, "is_initialized", None)
|
||
if callable(is_initialized):
|
||
try:
|
||
if is_initialized():
|
||
return None
|
||
except Exception as e:
|
||
# Treat a failing probe as "runtime not touched yet" so
|
||
# pre-init CPU preprocessing can still parallelize.
|
||
logger.debug("torch.xpu.is_initialized() probe failed: %s", e)
|
||
|
||
return safe_num_proc(desired)
|