Faster safetensors weight loading on unified-memory (integrated) GPUs (#5988)

* Faster safetensors weight loading on unified-memory (integrated) GPUs

On unified-memory GPUs (AMD APUs / "Strix Halo", NVIDIA GB10 "Spark", Intel
iGPUs) the GPU shares the system memory pool. PyTorch's fast pinned-DMA
host->device path does not recognize the Rust-allocated, mmap-backed buffers
that safetensors hands back, so a direct safetensors GPU load
(`safe_open(..., device=<cuda>)`) drops onto a slow per-tensor copy that, on
unified memory, additionally triggers page-attribute changes and page faults.

Wrap `transformers.modeling_utils.safe_open` so that, when transformers asks it
to load a shard directly onto a CUDA/HIP device, the shard is opened on CPU and
each tensor is `.clone()`-d into a normal torch allocation before `.to(device)`.
This restores the fast DMA path. Data, dtype and final device are unchanged, so
outputs are bit-identical -- only *how* the bytes reach the GPU changes.

Strictly gated to integrated/unified-memory GPUs via the standard `is_integrated`
device property (every visible device must be integrated): a hard no-op on
discrete NVIDIA/AMD GPUs, CPU, XPU and MLX, where the pinned-DMA path already
works. Only intercepts `framework="pt"` CUDA-device targets; CPU / disk-offload
loads are left untouched. Accuracy-neutral, idempotent, opt out with
UNSLOTH_DISABLE_UMA_CLONE_LOAD=1 (force the gate for tests with
UNSLOTH_FORCE_UMA=1/0).

This is the AMD/universal-UMA counterpart to the NVIDIA DGX Spark work in #5945
(which deliberately left the H2D clone-then-move out): gating on `is_integrated`
covers AMD Strix Halo, Intel iGPUs and Spark-class parts alike.

Verified on an AMD Radeon 8060S (gfx1151, Strix Halo) Windows ROCm box with
in-process, ordering-cancelled A/B benchmarks:
  - H2D mechanism (safe_open device=0 vs cpu->clone->.to(0)): 2.08x faster
    (1.076s -> 0.518s for a 988MB bf16 shard)
  - full `from_pretrained`: 1.56x faster (1.552s -> 0.996s), saving 0.555s --
    matching the H2D delta exactly
  - max|logit diff| stock vs patched == 0.0 (bit-identical), generate + a LoRA
    train step both verified
The absolute/relative win grows with bf16/fp16 weight volume (the same trick is
reported as ~2.3-2.75x on NVIDIA GB10 Spark for larger models).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: evaluate the integrated-GPU gate lazily, not at import (Gemini review)

patch_unified_memory_safetensors_load() called
is_integrated_unified_memory_gpu() at install time, and the gate queries
torch.cuda.get_device_properties() for every visible device -- initializing
the CUDA context during `import unsloth` on every CUDA machine (discrete
included). That (a) breaks fork-based multiprocessing, (b) runs BEFORE
patch_dgx_spark_memory_config can set PYTORCH_CUDA_ALLOC_CONF on Spark,
defeating that patch's expandable_segments config in the very environment
this PR targets, and (c) charges a CUDA context to CPU-only imports.

The gate now runs lazily inside the wrapper, ordered AFTER the
framework/device check so non-CUDA loads never trigger the property query;
a CUDA-target safe_open means the caller is initializing CUDA anyway, and
the gate is lru-cached so it is evaluated once. The wrapper installs
unconditionally (opt-out and idempotency unchanged) and passes through when
the gate is off.

Tests: install-time no-eval guarantee (gate raises if called during
install), wrapper passthrough with the gate off, all previous gating /
passthrough / CUDA correctness tests kept -- 16/16 pass. Verified on the
N1X (WSL2): module exec + patch install leave torch.cuda.is_initialized()
unchanged; CPU loads pass through; forced CUDA-target loads intercept and
land bit-identical on the GPU.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Compress PR comments to essentials (comment-only; AST-verified)

Docstrings and the _utils hook comment trimmed to their load-bearing
content (lazy-gate rationale, gating scope, opt-out env). AST dumps
with normalized docstrings are identical before/after for all three
files; the module's 16 unit tests pass unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* docs: tighten the UMA-load import comment (no code change)

* Tighten and trim code comments

* Drop unused is_integrated_unified_memory_gpu import from _utils.py

The UMA hook only needs patch_unified_memory_safetensors_load(); the
gate symbol is imported and used from ._uma_safetensors directly, so the
hoisted alias here was dead and tripped the import-hoist safety-net lint.

* Scope the UMA loader docstring to CUDA/HIP direct-device loads

The module text claimed Intel iGPU coverage, but the gate and device check
are CUDA/HIP only, and the clone path only wraps safe_open calls that carry
a CUDA device. State the actual scope and name the deliberate exclusions
(Intel XPU, CPU-open + .to() flows like bnb/HQQ) until they can be validated
on real hardware. Comment-only change.

* Tighten UMA safetensors loader comments

Trim the inline comments in the UMA clone-then-move path and the
_utils.py install site to be shorter and clearer. No code changes.

* uma: fall back to the direct move when the clone cannot allocate

The clone-and-move fast path transiently doubles one tensor's CPU
footprint while the mmap source and the CUDA destination are live. On a
UMA box with little free shared memory a large tensor could OOM where
the stock direct safe_open path would have loaded it. Both move sites
now go through a helper that catches the allocation failure and falls
back to the direct (slow but allocation-free) move, so the load always
succeeds; a genuine non-memory error re-raises identically from the
fallback.

Added a test that forces the clone to fail and verifies the wrapper
still lands tensors on the device with intact values (17 tests pass on
a real GPU).

* tests: track the moved pass-through inheritance in the gguf order check

Main moved the llama_extra_args pass-through inheritance out of the
GGUF branch into _resolve_inherited_extra_args, which runs before it,
so the source-order assertion's "if request.llama_extra_args is None"
anchor no longer exists inside the branch and the check failed after
the main merge. The test now asserts the same property in the current
shape: inheritance before the GGUF branch (a carried --no-mmproj still
shapes the hub guard's companion requirement), and marker, hub guard,
unload in order within the branch. Full file passes (32 tests).

* tests: anchor the inheritance order check on the call, not the definition

source.index("_resolve_inherited_extra_args(") matched the function
definition, which always precedes the endpoint, so the ordering
assertion was vacuously true. Anchoring on "= _resolve_inherited_
extra_args(" pins the first call site inside the load endpoint (line
4505), which is the statement whose position relative to the GGUF
branch the test is meant to guard. 32 tests pass.

* tests: align the gguf order test with main

Main fixed the stale ordering assertion in PR 7252; adopting its
version verbatim removes this file from the branch diff entirely and
avoids a conflict on the next main merge. 32 tests pass.

* uma: tighten comments

* Relicense UMA safetensors module and test under AGPL-3.0

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

tests/test_uma_safetensors_load.py

@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+
+"""Unit tests for the UMA safetensors clone-then-move fast load.
+
+The module loads in isolation with a fake ``transformers.modeling_utils``. The
+CUDA correctness check needs a GPU; gating, passthrough, idempotency and opt-out
+are GPU-free. The gate is lazy (wrapper-time), so the wrapper installs
+everywhere and passes through when it's off.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+torch = pytest.importorskip("torch")
+safetensors_torch = pytest.importorskip("safetensors.torch")
+import safetensors  # noqa: E402
+
+_MODULE_PATH = Path(__file__).resolve().parent.parent / "unsloth" / "models" / "_uma_safetensors.py"
+
+
+def _load_module():
+    spec = importlib.util.spec_from_file_location("uma_safetensors_under_test", _MODULE_PATH)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture()
+def uma():
+    return _load_module()
+
+
+@pytest.fixture()
+def force_uma(uma, monkeypatch):
+    """Force the UMA gate on (or off) and keep the lru_cache from sticking."""
+
+    def _set(on):
+        monkeypatch.setenv("UNSLOTH_FORCE_UMA", "1" if on else "0")
+        uma.is_integrated_unified_memory_gpu.cache_clear()
+
+    yield _set
+    uma.is_integrated_unified_memory_gpu.cache_clear()
+
+
+@pytest.fixture()
+def tiny_safetensors(tmp_path):
+    tensors = {
+        "w": torch.arange(32, dtype = torch.float32).reshape(4, 8),
+        "b": torch.tensor([1.0, 2.0, 3.0, 4.0], dtype = torch.float32),
+    }
+    path = tmp_path / "model.safetensors"
+    safetensors_torch.save_file(tensors, str(path))
+    return path, tensors
+
+
+def _install_fake_modeling_utils(monkeypatch, safe_open_fn):
+    fake_transformers = types.ModuleType("transformers")
+    fake_mu = types.ModuleType("transformers.modeling_utils")
+    fake_mu.safe_open = safe_open_fn
+    fake_transformers.modeling_utils = fake_mu
+    monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
+    monkeypatch.setitem(sys.modules, "transformers.modeling_utils", fake_mu)
+    return fake_mu
+
+
+# --- detection / gate ---
+
+
+def test_force_uma_on(uma, monkeypatch):
+    monkeypatch.setenv("UNSLOTH_FORCE_UMA", "1")
+    uma.is_integrated_unified_memory_gpu.cache_clear()
+    assert uma.is_integrated_unified_memory_gpu() is True
+
+
+def test_force_uma_off(uma, monkeypatch):
+    monkeypatch.setenv("UNSLOTH_FORCE_UMA", "0")
+    uma.is_integrated_unified_memory_gpu.cache_clear()
+    assert uma.is_integrated_unified_memory_gpu() is False
+
+
+@pytest.mark.parametrize(
+    "device,expected",
+    [
+        (0, True),
+        ("cuda", True),
+        ("cuda:0", True),
+        ("cpu", False),
+        ("disk", False),
+        (None, False),
+        (True, False),  # a bool is not a device index
+    ],
+)
+def test_is_cuda_target(uma, device, expected):
+    assert uma._is_cuda_target(device) is expected
+
+
+def test_is_cuda_target_torch_device(uma):
+    assert uma._is_cuda_target(torch.device("cuda", 0)) is True
+    assert uma._is_cuda_target(torch.device("cpu")) is False
+
+
+# --- patch gating ---
+
+
+def test_wrapper_passes_through_off_uma(uma, force_uma, monkeypatch):
+    """Gate OFF: every call -- including CUDA targets -- passes straight through
+    to the real safe_open (the gate is evaluated lazily inside the wrapper)."""
+    force_uma(False)
+    sentinel = object()
+    calls = []
+
+    def fake_safe_open(*args, **kwargs):
+        calls.append((args, kwargs))
+        return sentinel
+
+    fake_mu = _install_fake_modeling_utils(monkeypatch, fake_safe_open)
+    assert uma.patch_unified_memory_safetensors_load() is True
+    assert getattr(fake_mu.safe_open, "_unsloth_uma_clone", False) is True
+    out = fake_mu.safe_open("shard.safetensors", "pt", "cuda:0")
+    assert out is sentinel
+    assert calls == [(("shard.safetensors", "pt", "cuda:0"), {})]
+
+
+def test_patch_install_does_not_evaluate_gate(uma, monkeypatch):
+    """Installing the wrapper must NOT query the integrated-GPU property -- that
+    would init CUDA at ``import unsloth`` (fork-unsafe, and before the Spark
+    allocator config is set)."""
+
+    def _boom():
+        raise AssertionError("gate must not be evaluated at install time")
+
+    _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+    monkeypatch.setattr(uma, "is_integrated_unified_memory_gpu", _boom)
+    assert uma.patch_unified_memory_safetensors_load() is True
+
+
+def test_patch_noop_when_opted_out(uma, force_uma, monkeypatch):
+    force_uma(True)
+    monkeypatch.setenv("UNSLOTH_DISABLE_UMA_CLONE_LOAD", "1")
+    real = object()
+    fake_mu = _install_fake_modeling_utils(monkeypatch, real)
+    assert uma.patch_unified_memory_safetensors_load() is False
+    assert fake_mu.safe_open is real
+
+
+def test_patch_installs_and_is_idempotent(uma, force_uma, monkeypatch):
+    force_uma(True)
+    fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+    assert uma.patch_unified_memory_safetensors_load() is True
+    wrapped = fake_mu.safe_open
+    assert getattr(wrapped, "_unsloth_uma_clone", False) is True
+    # second call must not double-wrap
+    assert uma.patch_unified_memory_safetensors_load() is True
+    assert fake_mu.safe_open is wrapped
+
+
+# --- correctness ---
+
+
+def test_cpu_target_is_passthrough(uma, force_uma, monkeypatch, tiny_safetensors):
+    path, tensors = tiny_safetensors
+    force_uma(True)
+    fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+    uma.patch_unified_memory_safetensors_load()
+    # device="cpu" must NOT be intercepted -> identical data, still on CPU.
+    with fake_mu.safe_open(str(path), framework = "pt", device = "cpu") as f:
+        for key, expected in tensors.items():
+            got = f.get_slice(key)[:]
+            assert got.device.type == "cpu"
+            assert torch.equal(got, expected)
+
+
+@pytest.mark.skipif(
+    not (hasattr(torch, "cuda") and torch.cuda.is_available()),
+    reason = "needs a GPU for the host->device clone-and-move path",
+)
+def test_cuda_target_clones_and_moves(uma, force_uma, monkeypatch, tiny_safetensors):
+    path, tensors = tiny_safetensors
+    force_uma(True)
+    fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+    uma.patch_unified_memory_safetensors_load()
+    # device="cuda" IS intercepted -> tensors land on cuda, byte-identical.
+    with fake_mu.safe_open(str(path), framework = "pt", device = "cuda") as f:
+        for key, expected in tensors.items():
+            got = f.get_slice(key)[:]
+            assert got.device.type == "cuda"
+            assert torch.equal(got.cpu(), expected)
+            got_full = f.get_tensor(key)
+            assert got_full.device.type == "cuda"
+            assert torch.equal(got_full.cpu(), expected)
+
+
+@pytest.mark.skipif(
+    not (hasattr(torch, "cuda") and torch.cuda.is_available()),
+    reason = "needs a GPU for the low-memory fallback path",
+)
+def test_low_memory_falls_back_to_direct_move(uma, force_uma, monkeypatch, tiny_safetensors):
+    path, tensors = tiny_safetensors
+    force_uma(True)
+    fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+    uma.patch_unified_memory_safetensors_load()
+    # Clone OOMs (transient CPU doubling on a constrained UMA box): the wrapper
+    # must fall back to the direct move and still succeed.
+    real_clone = torch.Tensor.clone
+
+    def _oom_clone(self, *a, **k):
+        raise RuntimeError("[enforce fail] not enough memory")
+
+    monkeypatch.setattr(torch.Tensor, "clone", _oom_clone)
+    try:
+        with fake_mu.safe_open(str(path), framework = "pt", device = "cuda") as f:
+            for key, expected in tensors.items():
+                got = f.get_slice(key)[:]
+                assert got.device.type == "cuda"
+                got_full = f.get_tensor(key)
+                assert got_full.device.type == "cuda"
+    finally:
+        monkeypatch.setattr(torch.Tensor, "clone", real_clone)
+    for key, expected in tensors.items():
+        with fake_mu.safe_open(str(path), framework = "pt", device = "cuda") as f:
+            assert torch.equal(f.get_tensor(key).cpu(), expected)

unsloth/models/_uma_safetensors.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+
+"""Faster safetensors weight loading on unified-memory (integrated) GPUs.
+
+A direct ``safe_open(..., device=<cuda>)`` on CUDA/HIP UMA GPUs (AMD APUs,
+NVIDIA GB10 Spark) misses torch's fast pinned-DMA path: the mmap-backed
+safetensors buffers aren't recognized, so it falls to a slow per-tensor copy
+with page faults. Cloning each tensor into a normal torch CPU allocation before
+moving it restores the fast path; outputs are bit-identical.
+
+CUDA/HIP only, and only for loads that pass a CUDA device to ``safe_open``
+directly: Intel XPU iGPUs and the CPU-open + later ``.to()`` flows (e.g. bnb /
+HQQ quantized loads) keep the stock path until they can be validated on real
+hardware.
+"""
+
+import os
+import functools
+
+import torch
+
+__all__ = [
+    "is_integrated_unified_memory_gpu",
+    "patch_unified_memory_safetensors_load",
+]
+
+
+@functools.lru_cache(maxsize = None)
+def is_integrated_unified_memory_gpu():
+    """True only when EVERY visible CUDA/HIP device is integrated (UMA).
+
+    Discrete and mixed discrete+iGPU boxes return False (pinned-DMA already
+    works there). Test override: ``UNSLOTH_FORCE_UMA=1`` / ``=0``.
+    """
+    _force = os.environ.get("UNSLOTH_FORCE_UMA")
+    if _force == "1":
+        return True
+    if _force == "0":
+        return False
+    try:
+        if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
+            return False
+        count = torch.cuda.device_count()
+        if count == 0:
+            return False
+        for index in range(count):
+            props = torch.cuda.get_device_properties(index)
+            if not getattr(props, "is_integrated", 0):
+                return False
+        return True
+    except Exception:
+        return False
+
+
+def _is_cuda_target(device):
+    """Does a ``safe_open`` ``device=`` arg name a CUDA/HIP device?"""
+    if isinstance(device, bool):
+        return False
+    if isinstance(device, int):
+        return True
+    if isinstance(device, str):
+        return device == "cuda" or device.startswith("cuda:")
+    try:
+        return isinstance(device, torch.device) and device.type == "cuda"
+    except Exception:
+        return False
+
+
+def patch_unified_memory_safetensors_load():
+    """Wrap ``transformers.modeling_utils.safe_open`` so CUDA-target shard loads
+    open on CPU then clone+``.to(device)``, restoring the UMA fast path.
+
+    Gated to integrated GPUs (no-op on discrete/CPU/XPU/MLX), ``framework="pt"``
+    CUDA targets only, idempotent. Opt out: ``UNSLOTH_DISABLE_UMA_CLONE_LOAD=1``.
+
+    The gate runs lazily inside the wrapper, never here: probing device
+    properties at install would init CUDA during ``import unsloth`` -- breaking
+    fork multiprocessing and preempting ``patch_dgx_spark_memory_config``'s
+    allocator config. Returns ``True`` if the wrapper was installed.
+    """
+    if os.environ.get("UNSLOTH_DISABLE_UMA_CLONE_LOAD") == "1":
+        return False
+    try:
+        from transformers import modeling_utils as _mu
+    except Exception:
+        return False
+    real_safe_open = getattr(_mu, "safe_open", None)
+    if real_safe_open is None:
+        return False
+    if getattr(real_safe_open, "_unsloth_uma_clone", False):
+        return True
+
+    def _clone_move(tensor, device):
+        # Clone into a regular CPU allocation to restore fast pinned-DMA, then
+        # move. The clone transiently doubles the tensor's CPU footprint and can
+        # OOM a low-memory UMA box; fall back to the direct, allocation-free move
+        # (a genuine non-memory error re-raises identically from it).
+        try:
+            return tensor.clone().to(device, non_blocking = False)
+        except (MemoryError, RuntimeError):
+            return tensor.to(device, non_blocking = False)
+
+    class _ClonedSlice:
+        """Proxy over a safetensors ``PySafeSlice`` that clones+moves on read."""
+
+        __slots__ = ("_real", "_device")
+
+        def __init__(self, real, device):
+            self._real = real
+            self._device = device
+
+        def __getattr__(self, name):
+            if name in ("_real", "_device"):
+                raise AttributeError(name)
+            return getattr(self._real, name)
+
+        def __getitem__(self, key):
+            return _clone_move(self._real[key], self._device)
+
+    class _ClonedSafeOpen:
+        """Safetensors-handle proxy: load on CPU, clone+move tensors to CUDA."""
+
+        __slots__ = ("_real", "_device")
+
+        def __init__(self, args, kwargs):
+            self._device = kwargs.get("device", args[2] if len(args) > 2 else "cpu")
+            # Open on CPU; move ourselves.
+            if len(args) > 2:
+                args = args[:2] + ("cpu",) + tuple(args[3:])
+            else:
+                kwargs = dict(kwargs)
+                kwargs["device"] = "cpu"
+            self._real = real_safe_open(*args, **kwargs)
+
+        def __enter__(self):
+            self._real.__enter__()
+            return self
+
+        def __exit__(self, *exc):
+            return self._real.__exit__(*exc)
+
+        def __getattr__(self, name):
+            if name in ("_real", "_device"):
+                raise AttributeError(name)
+            return getattr(self._real, name)
+
+        def get_slice(self, name):
+            return _ClonedSlice(self._real.get_slice(name), self._device)
+
+        def get_tensor(self, name):
+            return _clone_move(self._real.get_tensor(name), self._device)
+
+    @functools.wraps(real_safe_open)
+    def _uma_safe_open(*args, **kwargs):
+        framework = kwargs.get("framework", args[1] if len(args) > 1 else None)
+        device = kwargs.get("device", args[2] if len(args) > 2 else "cpu")
+        # Device check first: non-CUDA loads must not trigger the CUDA-init gate.
+        if (
+            framework in ("pt", "pytorch")
+            and _is_cuda_target(device)
+            and is_integrated_unified_memory_gpu()
+        ):
+            return _ClonedSafeOpen(args, kwargs)
+        return real_safe_open(*args, **kwargs)
+
+    _uma_safe_open._unsloth_uma_clone = True
+    _mu.safe_open = _uma_safe_open
+    return True

unsloth/models/_utils.py

@@ -1670,6 +1670,13 @@ except:
 from transformers.modeling_utils import logger as transformers_logger
 
 
+# Faster safetensors loads on UMA (integrated) GPUs; lazy gate keeps this import
+# fork-safe (no CUDA init). No-op off-UMA. Opt out: UNSLOTH_DISABLE_UMA_CLONE_LOAD=1.
+from ._uma_safetensors import patch_unified_memory_safetensors_load
+
+patch_unified_memory_safetensors_load()
+
+
 def _all_missing_keys_are_position_ids(record_str):
     """True only when EVERY key in the 'newly initialized: [...]' list is a position_ids
     buffer.