diff --git a/tests/test_uma_safetensors_load.py b/tests/test_uma_safetensors_load.py new file mode 100644 index 0000000000..c6d304ab4f --- /dev/null +++ b/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) diff --git a/unsloth/models/_uma_safetensors.py b/unsloth/models/_uma_safetensors.py new file mode 100644 index 0000000000..38d8b7d33a --- /dev/null +++ b/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=)`` 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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 57169fa3de..f9ac879de6 100644 --- a/unsloth/models/_utils.py +++ b/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.