ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

studio: shard export checkpoint loads across all visible GPUs (#7215)

Browse source 
* studio: shard export checkpoint loads across all visible GPUs

Export checkpoint loading always used unsloth's from_pretrained default of
device_map="sequential", which stacks the whole model on GPU0. On a multi-GPU
host this OOMs GPU0 while the other GPUs sit empty, so a GGUF export that would
comfortably fit across the machine fails with CUDA out of memory (#7053).

Add _multi_gpu_device_map_kwargs(): when the CUDA/ROCm host exposes more than
one visible GPU and get_device_map resolves to "balanced" (the same policy the
inference loader already uses), pass device_map="balanced" to every
from_pretrained in load_checkpoint. In every other case -- single GPU, CPU,
MLX, or any probe failure -- it returns {} so the loader default is untouched.

Fixes #7053

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

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

* studio/save: reach the UUID/MIG fallback, release sharded models before quantize

Two review fixes on the multi-GPU export sharding:

1. UUID/MIG CUDA_VISIBLE_DEVICES masks resolve to no numeric ids, so the
   len(visible) > 1 gate skipped get_device_map entirely and large exports on
   those hosts still stacked onto GPU0. An empty id list now routes to
   get_device_map(None), whose visible-count fallback exists for exactly this
   case; a genuinely GPU-less host still resolves "sequential" and keeps the
   loader default.

2. The compressed (FP8/NVFP4) export freed GPU memory before its llm-compressor
   subprocess only for single-device models -- a plain .to("cpu") is invalid on
   an accelerate-dispatched model, so a multi-GPU-sharded checkpoint stayed
   resident on every GPU while the subprocess loaded a second copy. The release
   is factored into _offload_model_for_quantize_subprocess /
   _restore_model_after_quantize_subprocess: dispatched all-GPU shards get their
   accelerate hooks removed, move to CPU, and are re-dispatched over the
   recorded hf_device_map afterwards. Maps with cpu/disk targets (already
   offloading) and quantized models are left alone, as before.

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

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

* studio/save: budget merged tensors per device, restore hooks if CPU offload fails

Two review fixes on the multi-GPU export path:

1. The LoRA-merge save path budgeted every merged tensor against GPU0
   (get_device_properties(0) + unqualified memory_allocated()). A merged tensor
   lives on the GPU of its source layer, so for a model sharded across GPUs
   (the device_map="balanced" this PR enables) GPU1+ could OOM as their weights
   accumulated while only GPU0's headroom was checked. Budget against W's own
   device via a per-device cache; single-GPU behavior is unchanged (W on GPU0).

2. _offload_model_for_quantize_subprocess removed the accelerate hooks and then
   moved a dispatched model to CPU; if that move raised (host RAM too small for
   the sharded checkpoint) the model was left hookless and half-moved, breaking
   later exports in the same worker. It now re-dispatches (or, for the
   single-device path, moves back) on a failed move before aborting the offload.

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

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

* studio/save: release sharded models before the torchao reload too

The portable torchao FP8/INT8 export freed the in-memory model only when every
parameter sat on one device, then reloaded a second copy with
device_map="auto". A checkpoint loaded through the new multi-GPU export map is
accelerate-dispatched across several GPUs, so that single-device gate never
fired and the original stayed resident on every GPU during the reload -- an OOM
for exactly the models large enough to have needed the sharded load.

It now uses the same _offload_model_for_quantize_subprocess /
_restore_model_after_quantize_subprocess pair as the compressed export, which
removes the accelerate hooks, moves to CPU, and re-dispatches over the recorded
hf_device_map afterwards. Those helpers are extended to XPU as well, since
torchao also runs on Intel GPUs and the path they replace covered both.

* studio/save: release quantized and cpu-spilled shards before quantize reloads

Two cases the release helper skipped outright, both of which leave GPU memory
held while the compressed subprocess or the torchao device_map="auto" reload
allocates a second copy:

- Quantized models. ExportBackend.load_checkpoint loads 4-bit by DEFAULT, so the
  common Studio export hit the is_loaded_in_4bit guard and kept a quantized shard
  on every visible GPU. They are now attempted like any other model: transformers
  refuses .to() for some bitsandbytes builds, but that refusal raises before
  anything moves, so the existing recovery path restores the model and returns
  None -- best-effort where the stack allows it, old behaviour where it does not.

- Maps that spill to CPU. Any non-GPU target disqualified the whole model even
  though the GPU-mapped modules were still resident and are exactly what needs
  reclaiming. A cpu spill is safe to move (those weights are already in host RAM)
  and is now released; only disk/meta targets are still skipped, because
  accelerate keeps those parameters off the model and moving would try to
  materialize the whole checkpoint. An all-CPU map is skipped as a no-op.

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

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

* Fix multi-GPU offload for PEFT exports and fall back when sharding OOMs (#7215)

The dispatch branch of _offload_model_for_quantize_subprocess never ran for a
PEFT model: the wrapper proxies _hf_hook, so remove_hook_from_submodules raised
AttributeError and the bare except returned None. Studio always loads adapters,
so the new balanced map turned the offload off (0 percent freed against 91.8 on
the sequential path it replaces).

- resolve the real dispatch root before removing or replaying hooks
- snapshot and replay hooks, tensor placements and instance forwards; a plain
  re-dispatch rebuilds hooks against the post-PEFT tree (395 to 1379) and drops
  the fused kernels accelerate captured into _old_forward before unsloth patched
- drop the accelerator side of tied_params_map so the offload actually frees
- pass skip_keys on the fallback dispatch_model
- log the swallowed exception instead of returning None silently
- guard _unsloth_save_torchao_with_given_config like its two siblings
- retry the export load once on the loader default when the balanced map OOMs,
  which happens when a training or chat job already owns the other GPUs

Measured on 4x B200 with Qwen3-0.6B: 89.9 percent freed bf16 and 79.7 percent
4bit under balanced, logits bit-identical, hooks and placements restored
exactly, 184 Params4bit round-tripped unchanged including nested state2.

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

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

* Keep the original offloaded until the torchao copy is released, and retie shared weights (#7215)

Two follow-ups from review of 8b6b4ca0b.

_unsloth_save_torchao_with_given_config restored the original inside a finally
that ran as soon as from_pretrained returned, so the original and the quantized
copy were both resident while the copy was still being saved. The restore now
sits in an outer finally that covers saving and releasing quantized_model, which
is what the two sibling paths already do.

The dispatch replay did not preserve tied embeddings. A CPU round trip repoints
every tensor and accelerate's tied_params_map is keyed on the old pointer, so
replaying the hooks produced two independent parameters. Reproduced on a tied
Llama: lm_head picked up its own storage, the embedding was duplicated in VRAM,
and an update to one no longer reached the other. The snapshot now records tied
groups (named_parameters(remove_duplicate=False), since the default hides one
half of every pair) and re-ties them after placements are restored.

Verified: tie preserved, no extra storages, live CUDA storage census identical
before and after, updates propagate again, logits bit-identical, and the 4 GPU
invariants unchanged at 89.9 percent freed bf16 and 79.7 percent 4bit.

* Keep meta tensors out of tie groups, restore accelerate move guards, retry CPU spills (#7215)

Four follow-ups from review of a58f1086b.

Meta tensors all report storage pointer 0, and accelerate parks every
CPU-offloaded parameter on meta, so grouping by pointer collapsed them into one
fake tied group. Reproduced with a balanced map that spills two blocks to CPU:
18 meta parameters in a single group with shapes 64x64, 32x64 and 128x64, which
the retie step would have overwritten with the first one. Meta and null-pointer
tensors are now skipped, and the retie also checks shape.

remove_hook_from_submodules deletes the to/cuda/xpu wrappers dispatch_model
installs to stop a caller moving an offloaded model. The snapshot now records
and replays those alongside forward and _old_forward.

The single-device retry only matched OOM, but a balanced map that spills to CPU
is refused by bitsandbytes with a plain ValueError saying modules were dispatched
to the CPU or the disk (transformers quantizers/quantizer_bnb_4bit.py:128), with
no memory wording. That is now retryable too, which matters because Studio loads
4-bit by default and busy secondary GPUs are exactly when balanced spills.

The torchao path dropped the quantized copy at the end of the try, so a failure
in save_pretrained left it resident while the original was restored. The del
moved into the finally, ahead of the restore.

Four regression tests added; suites now 25 and 9.

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

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

* Retry exports whose multi-GPU load silently offloads to CPU, and clear the failed torchao traceback (#7215)

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

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

* Tighten comments for PR #7215

* Keep gradients across the export offload and release the failed torchao copy (#7215)

* Tighten comments for PR #7215

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Daniel Han <unslothai@gmail.com>
This commit is contained in:
Hakan Baysal 2026-07-26 14:16:36 +03:00 • committed by GitHub
commit e7d047a4ee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 1456 additions and 81 deletions
Download patch file Download diff file Expand all files Collapse all files

129
studio/backend/core/export/export.py
View file

@ -81,6 +81,82 @@ _PYTORCH_MISSING_MESSAGE = (
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
def _multi_gpu_device_map_kwargs() -> dict:
"""``device_map`` kwargs for sharding a checkpoint across every visible GPU.
unsloth's ``from_pretrained`` defaults to ``device_map="sequential"``, which stacks
the whole model on GPU0 and OOMs multi-GPU hosts whose other GPUs sit empty (#7053).
Returns ``{"device_map": "balanced"}`` only on a real multi-GPU CUDA/ROCm host
(mirroring the inference loader's ``get_device_map``), else empty so single-GPU, CPU
and MLX loads keep the loader default."""
if _IS_MLX:
return {}
try:
from utils.hardware import get_device_map, get_parent_visible_gpu_ids
visible = get_parent_visible_gpu_ids()
if len(visible) > 1:
device_map = get_device_map(visible)
elif not visible:
# UUID/MIG masks resolve to no numeric ids; get_device_map(None) falls back
# to the visible-GPU count, so a multi-GPU UUID/MIG host still shards.
device_map = get_device_map(None)
else:
return {}
if device_map == "balanced":
return {"device_map": device_map}
except Exception as exc:
logger.debug(f"multi-GPU device_map resolution failed; using loader default: {exc}")
return {}
def _is_oom_error(exc: BaseException) -> bool:
"""True for an accelerator OOM, however it is spelled.
accelerate and transformers re-raise it as a plain ``RuntimeError`` on several paths
and ROCm/XPU use their own classes, so match the message too.
"""
if torch is not None:
oom_types = tuple(
t
for t in (
getattr(torch, "OutOfMemoryError", None),
getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None),
getattr(getattr(torch, "xpu", None), "OutOfMemoryError", None),
)
if isinstance(t, type)
)
if oom_types and isinstance(exc, oom_types):
return True
return "out of memory" in f"{type(exc).__name__}: {exc}".lower()
def _is_cpu_spill_rejection(exc: BaseException) -> bool:
"""bitsandbytes refuses a map that spills to CPU/disk with a plain ``ValueError``.
Busy secondary GPUs can make ``balanced`` spill to CPU even where the old sequential
load fit on GPU0, and that message says nothing about memory, so the retry has to
match it explicitly. See transformers ``quantizers/quantizer_bnb_4bit.py``.
"""
return "dispatched on the cpu or the disk" in str(exc).lower()
class _CpuSpillRetry(Exception):
"""A multi-GPU load that succeeded but left modules offloaded to CPU/disk."""
def _cpu_offloaded_modules(model) -> int:
"""Count the modules a load parked on CPU or disk.
Only bitsandbytes refuses such a map; a full-precision load accepts it, leaves the
parameters on meta and dies much later in safetensors with "Cannot copy out of meta
tensor". Nothing raises at load time, so inspect the map directly. PEFT re-dispatches
when attaching an adapter, so in practice this catches merged checkpoints.
"""
device_map = getattr(model, "hf_device_map", None) or {}
return sum(1 for target in device_map.values() if str(target) in ("cpu", "disk"))
def _supports_kwarg(fn, name):
"""True if `fn` accepts keyword `name` directly or via **kwargs."""
import inspect
@ -271,6 +347,7 @@ class ExportBackend:
load_in_4bit: bool = True,
trust_remote_code: bool = False,
hf_token: Optional[str] = None,
_device_map_override: Optional[dict] = None,
) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
@ -303,6 +380,14 @@ class ExportBackend:
# Skip the Hub when offline so a no-internet export uses the local cache.
local_files_only = _hf_offline()
# Shard across every visible GPU instead of stacking on GPU0 (#7053); {} on
# single-GPU/CPU/MLX. _device_map_override is the single-device retry below.
_device_map_kw = (
_multi_gpu_device_map_kwargs()
if _device_map_override is None
else _device_map_override
)
# Run the type-detection probes in the forced-offline window (else a gated
# base 404s); it covers is_vision_model's Hub reads + the transformers-5
# subprocess, and local_files_only makes detect_audio_type's requests.get skip.
@ -328,6 +413,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "whisper":
@ -343,6 +429,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "snac":
@ -355,6 +442,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "bicodec":
@ -368,6 +456,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "dac":
@ -380,6 +469,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self.is_vision:
@ -392,6 +482,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
tokenizer = processor # vision: processor acts as tokenizer
@ -405,8 +496,16 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
# Only when we asked for the multi-GPU map: a single-GPU host has no second
# placement to retry on, so leave its behaviour untouched.
_offloaded = _cpu_offloaded_modules(model) if _device_map_kw else 0
if _device_map_override is None and _offloaded:
del model
raise _CpuSpillRetry(f"{_offloaded} module(s) offloaded to CPU/disk")
if _IS_MLX:
# MLX doesn't use PeftModel — detect LoRA via adapter_config.json
self.is_peft = adapter_config.exists()
@ -429,12 +528,42 @@ class ExportBackend:
return True, f"Loaded {model_type} model{peft_info} successfully"
except Exception as e:
# Sharding is an optimisation, never a requirement. "balanced" budgets from the
# free memory read BEFORE this process opens a CUDA context on each GPU, so when
# a training or chat job already owns the others the shard can OOM, or spill to
# CPU and be refused by bitsandbytes, where the old single-device load succeeded.
# Fall back once before giving up.
if (
_device_map_override is None
and (
isinstance(e, _CpuSpillRetry) or _is_oom_error(e) or _is_cpu_spill_rejection(e)
)
and _multi_gpu_device_map_kwargs()
):
# Retry outside this block: the live traceback pins the half-built model's
# frames, so an in-block retry inherits the exhausted device.
retry_reason = str(e)
else:
logger.error(f"Error loading checkpoint: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
logger.warning(
f"Multi-GPU export load unusable ({retry_reason}); retrying on "
f"the single-device loader default."
)
self.cleanup_memory()
return self.load_checkpoint(
checkpoint_path,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
hf_token = hf_token,
_device_map_override = {},
)
def _write_export_metadata(self, save_directory: str):
"""Write export_metadata.json with base model info for Chat page discovery."""
try:

242
studio/backend/tests/test_export_multi_gpu_device_map.py Normal file
View file

@ -0,0 +1,242 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Export checkpoint loading must shard across every visible GPU (#7053): the
``device_map="sequential"`` loader default stacks the whole model on GPU0 and OOMs
while the other GPUs sit empty. The loader now passes ``device_map="balanced"``, but
only on a real multi-GPU CUDA/ROCm host, so single-GPU, CPU and MLX are untouched."""
from __future__ import annotations
import contextlib
import sys
import types
from pathlib import Path
_BACKEND_DIR = Path(__file__).resolve().parent.parent
if str(_BACKEND_DIR) not in sys.path:
sys.path.insert(0, str(_BACKEND_DIR))
_TESTS_DIR = Path(__file__).resolve().parent
if str(_TESTS_DIR) not in sys.path:
sys.path.insert(0, str(_TESTS_DIR))
# Reuse the absolute-paths test's stub harness for loading core/export/export.py
# without torch/unsloth.
from test_export_absolute_paths import ( # noqa: E402
_install_export_backend_stubs,
_load_module,
)
def _export_mod(monkeypatch):
_install_export_backend_stubs(monkeypatch)
return _load_module("test_core_export_backend_device_map", "core/export/export.py", monkeypatch)
def _stub_hardware(monkeypatch, visible, device_map):
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: visible, raising = False)
monkeypatch.setattr(hw, "get_device_map", lambda ids: device_map, raising = False)
# ── _multi_gpu_device_map_kwargs ──
def test_multi_gpu_host_gets_balanced(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0, 1, 2], "balanced")
assert mod._multi_gpu_device_map_kwargs() == {"device_map": "balanced"}
def test_single_gpu_host_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0], "sequential")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_non_balanced_resolution_keeps_loader_default(monkeypatch):
# >1 visible id but a non-CUDA device resolves to "sequential": pass nothing.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0, 1], "sequential")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_uuid_mig_mask_falls_back_to_count_detection(monkeypatch):
# UUID/MIG masks resolve to NO numeric ids ([]), but get_device_map(None) still
# detects >1 GPU, so the empty list must route there, not to the loader default.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [], raising = False)
monkeypatch.setattr(
hw,
"get_device_map",
lambda ids: "balanced" if ids is None else "sequential",
raising = False,
)
assert mod._multi_gpu_device_map_kwargs() == {"device_map": "balanced"}
def test_no_visible_gpus_keeps_loader_default(monkeypatch):
# Empty mask / CPU host: get_device_map(None) resolves "sequential" -> {}.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [], raising = False)
monkeypatch.setattr(hw, "get_device_map", lambda ids: "sequential", raising = False)
assert mod._multi_gpu_device_map_kwargs() == {}
def test_mlx_host_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
# The stubs set _IS_MLX = True; even a multi-GPU view must yield no device_map.
_stub_hardware(monkeypatch, [0, 1], "balanced")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_hardware_probe_failure_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
def _boom():
raise RuntimeError("no GPUs")
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", _boom, raising = False)
assert mod._multi_gpu_device_map_kwargs() == {}
# ── load_checkpoint forwards the kwargs to from_pretrained ──
class _RecordingLoader:
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
return types.SimpleNamespace(), types.SimpleNamespace()
def _load_text_checkpoint(monkeypatch, tmp_path, device_map_kwargs):
mod = _export_mod(monkeypatch)
_RecordingLoader.calls = []
monkeypatch.setattr(mod, "FastLanguageModel", _RecordingLoader)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: device_map_kwargs)
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
assert ok, message
assert len(_RecordingLoader.calls) == 1
return _RecordingLoader.calls[0]
def test_load_checkpoint_forwards_balanced_device_map(monkeypatch, tmp_path):
kwargs = _load_text_checkpoint(monkeypatch, tmp_path, {"device_map": "balanced"})
assert kwargs["device_map"] == "balanced"
def test_load_checkpoint_omits_device_map_on_single_gpu(monkeypatch, tmp_path):
kwargs = _load_text_checkpoint(monkeypatch, tmp_path, {})
assert "device_map" not in kwargs # loader default (sequential) untouched
# ── a load that succeeds but offloads to CPU/disk ──
def test_cpu_offloaded_modules_counts_cpu_and_disk(monkeypatch):
mod = _export_mod(monkeypatch)
model = types.SimpleNamespace(hf_device_map = {"a": 0, "b": "cpu", "c": 1, "d": "disk"})
assert mod._cpu_offloaded_modules(model) == 2
def test_cpu_offloaded_modules_ignores_gpu_only_and_missing_maps(monkeypatch):
mod = _export_mod(monkeypatch)
assert mod._cpu_offloaded_modules(types.SimpleNamespace(hf_device_map = {"a": 0})) == 0
assert mod._cpu_offloaded_modules(types.SimpleNamespace(hf_device_map = None)) == 0
assert mod._cpu_offloaded_modules(types.SimpleNamespace()) == 0
class _SpillThenCleanLoader:
"""First call offloads to CPU (bf16 accepts it silently), second is clean."""
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
device_map = {"model.layers.0": 0} if len(cls.calls) > 1 else {"model.layers.0": "cpu"}
return types.SimpleNamespace(hf_device_map = device_map), types.SimpleNamespace()
def _run_spill_loader(monkeypatch, tmp_path, device_map_kwargs):
mod = _export_mod(monkeypatch)
_SpillThenCleanLoader.calls = []
monkeypatch.setattr(mod, "FastLanguageModel", _SpillThenCleanLoader)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: device_map_kwargs)
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
return ok, message, _SpillThenCleanLoader.calls
def test_successful_load_that_offloads_to_cpu_retries_single_device(monkeypatch, tmp_path):
# Nothing raises, so only hf_device_map catches it; the parameters would otherwise
# stay on meta and kill the export inside safetensors.
ok, message, calls = _run_spill_loader(monkeypatch, tmp_path, {"device_map": "balanced"})
assert ok, message
assert len(calls) == 2
assert calls[0]["device_map"] == "balanced"
assert "device_map" not in calls[1]
def test_single_gpu_offload_is_left_alone(monkeypatch, tmp_path):
# No multi-GPU map was requested, so there is nothing to retry on.
ok, message, calls = _run_spill_loader(monkeypatch, tmp_path, {})
assert ok, message
assert len(calls) == 1
def test_retry_result_is_kept_even_if_it_also_offloads(monkeypatch, tmp_path):
# The retry runs with _device_map_override set, so it must never recurse again.
mod = _export_mod(monkeypatch)
class _AlwaysSpills:
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
return types.SimpleNamespace(hf_device_map = {"a": "cpu"}), types.SimpleNamespace()
monkeypatch.setattr(mod, "FastLanguageModel", _AlwaysSpills)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: {"device_map": "balanced"})
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
assert ok, message
assert len(_AlwaysSpills.calls) == 2

670
tests/test_compressed_export_gpu_release.py Normal file
View file

@ -0,0 +1,670 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""The compressed (FP8/NVFP4) export must free GPU weights before its llm-compressor
subprocess loads a second copy from disk, including for accelerate-dispatched multi-GPU
shards, which the old single-device-only ``.to("cpu")`` skipped and left resident.
Pulls the release/restore helpers out of unsloth/save.py via AST (importing the module
needs torch/transformers) and exercises them with fakes.
"""
from __future__ import annotations
import ast
import gc
import sys
import types
from pathlib import Path
import pytest
_SAVE_PY = Path(__file__).resolve().parent.parent / "unsloth" / "save.py"
_WANTED = {
"_accelerate_dispatch_root",
"_snapshot_dispatch_state",
"_drop_accelerator_tied_param_cache",
"_accelerate_move_guards",
"_split_tensor_path",
"_lookup_tensor",
"_share_tensor",
"_restore_dispatch_state",
"_offload_model_for_quantize_subprocess",
"_restore_model_after_quantize_subprocess",
}
_WANTED_ASSIGNS = {
"_DISPATCH_SNAPSHOT_ATTR",
"_ACCELERATE_MOVE_GUARDS",
} # module constants the helpers close over
class _FakeLogger:
def __init__(self):
self.warnings = []
def warning_once(self, msg):
self.warnings.append(msg)
def _load_helpers(fake_torch, fake_logger):
tree = ast.parse(_SAVE_PY.read_text(encoding = "utf-8"))
keep = [
node
for node in tree.body
if (isinstance(node, ast.FunctionDef) and node.name in _WANTED)
or (
isinstance(node, ast.Assign)
and any(isinstance(t, ast.Name) and t.id in _WANTED_ASSIGNS for t in node.targets)
)
]
n_fns = sum(1 for node in keep if isinstance(node, ast.FunctionDef))
assert n_fns == len(_WANTED), "release helpers missing from save.py"
namespace = {"torch": fake_torch, "logger": fake_logger}
exec( # noqa: S102 - loading trusted repo source
compile(ast.Module(body = keep, type_ignores = []), str(_SAVE_PY), "exec"),
namespace,
)
return namespace
def _fake_torch(cuda_available = True):
t = types.ModuleType("torch")
t.cuda = types.SimpleNamespace(is_available = lambda: cuda_available)
return t
class _FakeModel:
def __init__(
self,
device_map = None,
devices = ("cuda:0",),
quantized = False,
):
if device_map is not None:
self.hf_device_map = device_map
self._devices = [types.SimpleNamespace(device = d) for d in devices]
self.moved_to = []
self.is_loaded_in_4bit = quantized
def parameters(self):
return iter(self._devices)
def to(self, target):
self.moved_to.append(str(target))
return self
@pytest.fixture
def _fake_accelerate(monkeypatch):
calls = {"removed": [], "dispatched": [], "dispatch_kwargs": [], "hooks_added": []}
accel = types.ModuleType("accelerate")
def _dispatch(model, device_map, **kwargs):
calls["dispatched"].append((model, dict(device_map)))
calls["dispatch_kwargs"].append(kwargs)
accel.dispatch_model = _dispatch
hooks = types.ModuleType("accelerate.hooks")
hooks.remove_hook_from_submodules = lambda model: calls["removed"].append(model)
hooks.add_hook_to_module = lambda module, hook: calls["hooks_added"].append((module, hook))
accel.hooks = hooks
monkeypatch.setitem(sys.modules, "accelerate", accel)
monkeypatch.setitem(sys.modules, "accelerate.hooks", hooks)
return calls
def test_dispatched_multi_gpu_model_is_released_and_redispatched(_fake_accelerate):
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 0, "model.layers.1": 1}
model = _FakeModel(device_map = device_map, devices = ("cuda:0", "cuda:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model] # hooks removed before the move
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_dispatched_move_failure_redispatches_and_returns_none(_fake_accelerate):
# If .to("cpu") raises after the hooks came off, the model must be re-dispatched,
# not left hookless and half-moved.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.1": 1}
class _MoveFails(_FakeModel):
def to(self, target):
raise RuntimeError("host RAM cannot hold the sharded model")
model = _MoveFails(device_map = device_map, devices = ("cuda:0", "cuda:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token is None # offload aborted
assert _fake_accelerate["removed"] == [model] # hooks were removed...
assert _fake_accelerate["dispatched"] == [(model, device_map)] # ...then restored
def test_single_device_move_failure_restores_and_returns_none():
ns = _load_helpers(_fake_torch(), _FakeLogger())
class _MoveFails(_FakeModel):
def __init__(self):
super().__init__(devices = ("cuda:0",))
def to(self, target):
self.moved_to.append(str(target))
if target == "cpu":
raise RuntimeError("move failed")
return self
model = _MoveFails()
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token is None
# attempted the cpu move, then restored back to the original device
assert model.moved_to == ["cpu", "cuda:0"]
def test_cpu_spilled_map_still_releases_its_gpu_shards(_fake_accelerate):
# One module spilled to CPU, but the rest is the GPU memory the reload needs, and
# the spilled weights are already in host RAM, so the move is safe.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1, "model.layers.9": "cpu"}
model = _FakeModel(device_map = device_map)
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model]
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_disk_offloaded_map_is_left_alone(_fake_accelerate):
# disk/meta entries are not on the model, so moving would materialize the whole
# checkpoint into RAM.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": 0, "model.layers.9": "disk"})
assert ns["_offload_model_for_quantize_subprocess"](model) is None
assert model.moved_to == []
assert _fake_accelerate["removed"] == []
def test_all_cpu_map_is_left_alone(_fake_accelerate):
# Nothing on an accelerator: no GPU memory to reclaim, so do not churn the hooks.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": "cpu", "model.layers.0": "cpu"})
assert ns["_offload_model_for_quantize_subprocess"](model) is None
assert model.moved_to == []
assert _fake_accelerate["removed"] == []
def test_single_device_model_keeps_plain_move():
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(devices = ("cuda:0",))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert model.moved_to == ["cpu"]
assert token is not None and token[0] == "device"
ns["_restore_model_after_quantize_subprocess"](model, token)
assert model.moved_to[-1] == "cuda:0"
def test_quantized_model_is_released_when_the_stack_allows_it():
# Studio exports load 4-bit by DEFAULT, so skipping quantized models left a shard
# on every GPU. Release them too where the move is accepted.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(devices = ("cuda:0",), quantized = True)
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token == ("device", "cuda:0")
assert model.moved_to == ["cpu"]
def test_quantized_model_that_refuses_to_move_is_left_usable():
# transformers rejects .to() for some bitsandbytes builds and raises before
# anything moves, so the old behaviour must hold: no token, nothing escaping.
ns = _load_helpers(_fake_torch(), _FakeLogger())
class _Refuses(_FakeModel):
def to(self, target):
raise ValueError("`.to` is not supported for 4-bit bitsandbytes models")
model = _Refuses(devices = ("cuda:0",), quantized = True)
assert ns["_offload_model_for_quantize_subprocess"](model) is None
def test_no_cuda_is_noop_and_restore_none_is_noop():
ns = _load_helpers(_fake_torch(cuda_available = False), _FakeLogger())
model = _FakeModel()
assert ns["_offload_model_for_quantize_subprocess"](model) is None
ns["_restore_model_after_quantize_subprocess"](model, None) # must not raise
assert model.moved_to == []
def test_restore_failure_warns_instead_of_raising(_fake_accelerate):
fake_logger = _FakeLogger()
ns = _load_helpers(_fake_torch(), fake_logger)
class _ExplodingModel(_FakeModel):
def to(self, target):
raise RuntimeError("device gone")
model = _ExplodingModel(devices = ("cuda:0",))
ns["_restore_model_after_quantize_subprocess"](model, ("device", "cuda:0"))
assert fake_logger.warnings # warned, did not raise
def test_lora_merge_budgets_per_device():
# A merged tensor W lives on the GPU of its source layer, so budget against W's
# own device, not GPU0, else a sharded model OOMs GPU1+ (#7053).
src = _SAVE_PY.read_text(encoding = "utf-8")
tree = ast.parse(src)
fn = next(
(
n
for n in ast.walk(tree)
if isinstance(n, ast.FunctionDef) and n.name == "unsloth_save_model"
),
None,
)
assert fn is not None, "unsloth_save_model not found"
body = ast.get_source_segment(src, fn)
# Budget keyed on W's device, not a hardcoded device 0 / unqualified alloc.
assert "torch.cuda.memory_allocated(W.device)" in body
assert "_device_vram_budget(W.device)" in body
assert "get_device_properties(0).total_memory * maximum_memory_usage" not in body
# ── the torchao ("portable" FP8/INT8) export shares the same release ──
def _fake_torch_xpu():
t = types.ModuleType("torch")
t.cuda = types.SimpleNamespace(is_available = lambda: False)
t.xpu = types.SimpleNamespace(is_available = lambda: True)
return t
def test_dispatched_xpu_model_is_released(_fake_accelerate):
# torchao runs on Intel GPUs too, so an XPU-dispatched shard must release exactly
# like a CUDA one.
ns = _load_helpers(_fake_torch_xpu(), _FakeLogger())
device_map = {"model.embed": "xpu:0", "model.layers.0": "xpu:1"}
model = _FakeModel(device_map = device_map, devices = ("xpu:0", "xpu:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model]
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_single_device_xpu_model_is_released():
ns = _load_helpers(_fake_torch_xpu(), _FakeLogger())
model = _FakeModel(devices = ("xpu:0",))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token == ("device", "xpu:0")
assert model.moved_to == ["cpu"]
def test_torchao_export_uses_the_shared_release():
"""The torchao path must not re-inline a single-device-only ``.to("cpu")``.
A plain move is invalid on a dispatched model, so single-device-only handling left
a multi-GPU shard resident while ``device_map="auto"`` loaded a second copy.
"""
src = _SAVE_PY.read_text(encoding = "utf-8")
torchao = src.split("def _unsloth_save_torchao(", 1)[1].split("\ndef ", 1)[0]
assert "_offload_model_for_quantize_subprocess(model)" in torchao
assert "_restore_model_after_quantize_subprocess(model" in torchao
# No hand-rolled single-device gate left behind.
assert "len(_devs) == 1" not in torchao
# ── regressions for the multi-GPU dispatch branch ──
class _Child:
"""Minimal stand-in for an nn.Module leaf, enough for the dispatch walk."""
def __init__(
self,
name = "inner",
device_map = None,
):
self._modules = {}
self.__dict__["_name"] = name
if device_map is not None:
self.hf_device_map = device_map
def named_modules(self):
yield "", self
for key, child in self._modules.items():
for sub_name, sub in child.named_modules():
yield (f"{key}.{sub_name}" if sub_name else key), sub
def get_submodule(self, target):
node = self
for part in target.split("."):
node = node._modules[part]
return node
def named_parameters(self, remove_duplicate = True):
return iter(())
def named_buffers(self, remove_duplicate = True):
return iter(())
class _PeftLikeWrapper(_Child):
"""Proxies unknown attributes to the wrapped model, like ``PeftModelForCausalLM``:
``hasattr(wrapper, "_hf_hook")`` is True while ``delattr`` fails, which is what made
the offload a silent no-op."""
def __init__(self, inner):
super().__init__(name = "wrapper")
self._modules["base_model"] = inner
self.moved_to = []
def __getattr__(self, item):
return getattr(self._modules["base_model"], item)
def to(self, target):
self.moved_to.append(str(target))
return self
def parameters(self):
return iter(self._modules["base_model"]._devices)
def test_dispatch_root_is_the_inner_model_for_a_peft_style_wrapper(_fake_accelerate):
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1}
inner = _Child(device_map = device_map)
inner._devices = [types.SimpleNamespace(device = "cuda:0")]
wrapper = _PeftLikeWrapper(inner)
assert ns["_accelerate_dispatch_root"](wrapper) is inner
token = ns["_offload_model_for_quantize_subprocess"](wrapper)
# hooks must come off the INNER module, not the proxying wrapper
assert _fake_accelerate["removed"] == [inner]
assert wrapper.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
def test_dispatch_root_falls_back_to_the_model_it_was_given():
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": 0})
assert ns["_accelerate_dispatch_root"](model) is model
def test_offload_failure_is_logged_not_swallowed():
# A bare `return None` is indistinguishable from "nothing to move".
fake_logger = _FakeLogger()
ns = _load_helpers(_fake_torch(), fake_logger)
class _Explodes(_FakeModel):
@property
def hf_device_map(self):
raise RuntimeError("boom")
assert ns["_offload_model_for_quantize_subprocess"](_Explodes()) is None
assert any("boom" in w for w in fake_logger.warnings)
def test_restore_without_a_snapshot_forwards_skip_keys(_fake_accelerate):
# dispatch_model() defaults skip_keys to None, which moves every forward kwarg to
# the executing device, wrong for tensors transformers marks device-invariant.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1}
model = _FakeModel(device_map = device_map, devices = ("cuda:0", "cuda:1"))
model._skip_keys_device_placement = ["past_key_values"]
ns["_restore_model_after_quantize_subprocess"](model, ("dispatch", device_map))
assert _fake_accelerate["dispatched"] == [(model, device_map)]
assert _fake_accelerate["dispatch_kwargs"] == [{"skip_keys": ["past_key_values"]}]
def test_snapshot_restores_a_forward_patched_after_the_dispatch(_fake_accelerate):
"""accelerate restores ``forward = _old_forward`` on removal, and ``_old_forward``
is the forward from when the hook was FIRST attached. unsloth patches forwards after
the dispatch, so a naive remove/re-add throws every fused kernel away for good."""
ns = _load_helpers(_fake_torch(), _FakeLogger())
root = _Child(device_map = {"model.embed": 0, "mlp": 1})
mlp = _Child(name = "mlp")
root._modules["mlp"] = mlp
stock_forward = lambda *a, **k: "stock" # noqa: E731
fused_forward = lambda *a, **k: "unsloth-fused" # noqa: E731
mlp._hf_hook = object()
mlp._old_forward = stock_forward # captured by accelerate at dispatch time
mlp.forward = fused_forward # installed by unsloth afterwards
snapshot = ns["_snapshot_dispatch_state"](root)
# what accelerate's removal does
del mlp.__dict__["_hf_hook"]
mlp.forward = mlp._old_forward
del mlp.__dict__["_old_forward"]
assert mlp.forward() == "stock"
ns["_restore_dispatch_state"](root, snapshot)
assert mlp.forward() == "unsloth-fused"
assert mlp.__dict__["_old_forward"] is stock_forward
def test_snapshot_reties_shared_parameters(_fake_accelerate):
"""A CPU round trip repoints every tensor, so replaying the hooks alone leaves tied
weights as independent copies: double VRAM, and updates to one never reach the other."""
import torch
root = _Child(device_map = {"embed": 0, "head": 0})
shared = torch.nn.Parameter(torch.zeros(4, 4))
for name in ("embed", "head"):
child = _Child(name = name)
child._parameters = {"weight": shared}
child._buffers = {}
root._modules[name] = child
def named(remove_duplicate = True):
seen, out = set(), []
for mod_name, mod in root._modules.items():
for attr, tensor in mod._parameters.items():
if remove_duplicate and id(tensor) in seen:
continue
seen.add(id(tensor))
out.append((f"{mod_name}.{attr}", tensor))
return iter(out)
root.named_parameters = named
ns = _load_helpers(_fake_torch(), _FakeLogger())
snapshot = ns_ties = ns["_snapshot_dispatch_state"](root)
assert ns_ties[3] == [["embed.weight", "head.weight"]]
# what the replay leaves behind before the retie step
root._modules["head"]._parameters["weight"] = torch.nn.Parameter(shared.detach().clone())
assert (
root._modules["embed"]._parameters["weight"].data_ptr()
!= root._modules["head"]._parameters["weight"].data_ptr()
)
ns["_restore_dispatch_state"](root, snapshot)
assert (
root._modules["embed"]._parameters["weight"].data_ptr()
== root._modules["head"]._parameters["weight"].data_ptr()
)
def test_meta_tensors_never_form_tie_groups(_fake_accelerate):
"""Offloaded parameters all sit on meta with storage pointer 0, so grouping by
pointer alone would collapse them into one fake tie and overwrite them all."""
import torch
root = _Child(device_map = {"a": 0, "b": "cpu", "c": "cpu"})
live = torch.nn.Parameter(torch.zeros(4, 4))
offloaded = [
torch.nn.Parameter(torch.empty(4, 4, device = "meta")),
torch.nn.Parameter(torch.empty(8, 2, device = "meta")),
]
def named(remove_duplicate = True):
return iter([("a.weight", live), ("b.weight", offloaded[0]), ("c.weight", offloaded[1])])
root.named_parameters = named
ns = _load_helpers(_fake_torch(), _FakeLogger())
_hooks, places, _attrs, ties, _grads = ns["_snapshot_dispatch_state"](root)
assert ties == [] # nothing is tied here
assert "b.weight" in places # still tracked for placement
def test_accelerate_move_guards_survive_the_replay(_fake_accelerate):
"""remove_hook_from_module also deletes the to/cuda/... guards dispatch_model
installs to stop a caller moving an offloaded model."""
ns = _load_helpers(_fake_torch(), _FakeLogger())
root = _Child(device_map = {"": 0})
guard = lambda *a, **k: "blocked" # noqa: E731
root._hf_hook = object()
root.to = guard
root.cuda = guard
snapshot = ns["_snapshot_dispatch_state"](root)
del root.__dict__["_hf_hook"], root.__dict__["to"], root.__dict__["cuda"]
ns["_restore_dispatch_state"](root, snapshot)
assert root.__dict__["to"] is guard
assert root.__dict__["cuda"] is guard
def test_gradients_survive_the_offload_round_trip():
"""init_hook rebuilds the Parameter and drops .grad, so the snapshot has to carry it."""
import torch
root = _Child(device_map = {"": 0})
weight = torch.nn.Parameter(torch.zeros(4, 4))
weight.grad = torch.full((4, 4), 3.0)
root._parameters = {"weight": weight}
root.named_parameters = lambda remove_duplicate = True: iter([("weight", weight)])
ns = _load_helpers(_fake_torch(), _FakeLogger())
snapshot = ns["_snapshot_dispatch_state"](root)
assert torch.equal(snapshot[4]["weight"], torch.full((4, 4), 3.0))
# What init_hook does: same name, fresh Parameter, no grad.
replacement = torch.nn.Parameter(torch.zeros(4, 4))
assert replacement.grad is None
root._parameters = {"weight": replacement}
ns["_restore_dispatch_state"](root, snapshot)
assert replacement.grad is not None, "the restore must put the gradient back"
assert torch.equal(replacement.grad, torch.full((4, 4), 3.0))
def test_the_other_torchao_path_also_clears_the_failed_copy():
"""Both torchao paths must drop the copy and the traceback pinning it before restoring."""
src = _SAVE_PY.read_text(encoding = "utf-8")
body = src.split("\ndef _unsloth_save_torchao(", 1)[1].split("\ndef ", 1)[0]
finally_block = body.split(" finally:", 1)[1]
assert "del quantized_model" in finally_block
assert "traceback.clear_frames" in finally_block
restore_at = finally_block.index("_restore_model_after_quantize_subprocess")
assert finally_block.index("del quantized_model") < restore_at
assert finally_block.index("traceback.clear_frames") < restore_at
def test_cpu_spill_rejection_is_retryable():
"""bitsandbytes rejects a CPU-spilled map with a ValueError that says nothing about
memory, so the single-device retry has to match it explicitly."""
import importlib.util
from pathlib import Path
export_py = (
Path(__file__).resolve().parent.parent
/ "studio"
/ "backend"
/ "core"
/ "export"
/ "export.py"
)
src = ast.parse(export_py.read_text(encoding = "utf-8"))
keep = [
n
for n in src.body
if isinstance(n, ast.FunctionDef) and n.name in {"_is_oom_error", "_is_cpu_spill_rejection"}
]
assert len(keep) == 2
namespace = {"torch": None}
exec( # noqa: S102 - loading trusted repo source
compile(ast.Module(body = keep, type_ignores = []), str(export_py), "exec"), namespace
)
bnb = ValueError(
"Some modules are dispatched on the CPU or the disk. Make sure you have enough "
"GPU RAM to fit the quantized model."
)
assert not namespace["_is_oom_error"](bnb)
assert namespace["_is_cpu_spill_rejection"](bnb)
assert namespace["_is_oom_error"](RuntimeError("CUDA out of memory. Tried to allocate 1 GiB"))
assert not namespace["_is_cpu_spill_rejection"](RuntimeError("some other failure"))
def test_torchao_releases_the_quantized_copy_in_finally():
"""If save_pretrained raises, the quantized copy must still be dropped before the
original is restored, or both are resident at once."""
src = _SAVE_PY.read_text(encoding = "utf-8")
body = src.split("def _unsloth_save_torchao_with_given_config(", 1)[1].split("\ndef ", 1)[0]
finally_block = body.split(" finally:", 1)[1]
assert "del quantized_model" in finally_block
assert "_restore_model_after_quantize_subprocess(model, model_restore)" in finally_block
# and the restore must come after the copy is dropped
assert finally_block.index("del quantized_model") < finally_block.index(
"_restore_model_after_quantize_subprocess"
)
# dropping the local is not enough: the live traceback still holds the frames
assert "traceback.clear_frames" in finally_block
assert finally_block.index("traceback.clear_frames") < finally_block.index(
"_restore_model_after_quantize_subprocess"
)
def test_a_live_traceback_pins_the_failed_copy_until_its_frames_are_cleared():
"""Why the clear_frames call above is load-bearing, on plain objects."""
import sys
import traceback
import weakref
class _Copy:
pass
def _build_and_fail(sink):
copy = _Copy() # noqa: F841 -- the point is that the frame retains it
sink.append(weakref.ref(copy))
raise RuntimeError("save_pretrained failed")
def _run(clear_frames):
# try/finally with the exception still in flight, exactly as in save.py
sink = []
alive = None
try:
try:
_build_and_fail(sink)
finally:
if clear_frames:
exc = sys.exc_info()[1]
if exc is not None:
traceback.clear_frames(exc.__traceback__)
gc.collect()
alive = sink[0]() is not None
except RuntimeError:
pass
return alive
assert _run(clear_frames = False), "expected the traceback to pin the copy"
assert not _run(clear_frames = True), "clear_frames must release it"

446
unsloth/save.py
View file

@ -48,6 +48,7 @@ import functools
from transformers.models.llama.modeling_llama import logger
from .kernels import fast_dequantize, QUANT_STATE, get_lora_parameters_bias
import subprocess
import traceback
import psutil
import re
from transformers.models.llama.modeling_llama import logger
@ -1122,7 +1123,19 @@ def unsloth_save_model(
torch_dtype
)
max_vram = int(torch.cuda.get_device_properties(0).total_memory * maximum_memory_usage)
# A merged tensor lives on the GPU of its source layer, so budget against W's own
# device, not GPU0, else a sharded model OOMs GPU1+ while only GPU0 is checked.
_max_vram_by_device = {}
def _device_vram_budget(dev):
if dev.type != "cuda":
return None
idx = dev.index if dev.index is not None else torch.cuda.current_device()
if idx not in _max_vram_by_device:
_max_vram_by_device[idx] = int(
torch.cuda.get_device_properties(idx).total_memory * maximum_memory_usage
)
return _max_vram_by_device[idx]
print("Unsloth: Saving model... This might take 5 minutes ...")
@ -1138,8 +1151,15 @@ def unsloth_save_model(
if bias is not None:
state_dict[f"model.layers.{j}.{item}.bias"] = bias
if (torch.cuda.memory_allocated() + W.nbytes) < max_vram:
# Save to GPU memory
_dev_budget = _device_vram_budget(W.device)
if (
_dev_budget is not None
and (torch.cuda.memory_allocated(W.device) + W.nbytes) < _dev_budget
):
# Fits on W's own GPU
state_dict[name] = W
elif W.device.type != "cuda":
# Already off-GPU: keeping it costs no VRAM
state_dict[name] = W
# [TODO] Saving to RAM seems to leak memory???
# elif (max_ram - W.nbytes) > 0:
@ -4524,6 +4544,19 @@ def _unsloth_save_torchao_with_given_config(
else:
kwargs = {"dtype": torch.bfloat16}
# Else the original stays resident on every GPU while device_map="auto" below
# loads a second copy.
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
# The original stays offloaded until the quantized copy is saved AND released,
# else both are resident at once and the restore OOMs.
try:
# Reload with quantization applied
quantized_model = auto_model.from_pretrained(
save_directory,
@ -4549,6 +4582,24 @@ def _unsloth_save_torchao_with_given_config(
)
tokenizer.save_pretrained(torchao_save_directory, token = token)
finally:
# del here, not at the end of the try: if save_pretrained raises, the copy
# would otherwise still be resident while the original is restored.
quantized_model = None
del quantized_model
# A failed save leaves a live traceback whose frames still hold the copy, so
# dropping the local alone does not free its VRAM.
_exc = sys.exc_info()[1]
if _exc is not None:
traceback.clear_frames(_exc.__traceback__)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
_restore_model_after_quantize_subprocess(model, model_restore)
# Clean up the intermediate unquantized model
if os.path.exists(save_directory):
try:
@ -4585,6 +4636,318 @@ def _print_compressed_hw_note(scheme, out_dir):
)
_DISPATCH_SNAPSHOT_ATTR = "_unsloth_dispatch_snapshot"
def _accelerate_move_guards():
"""The instance methods dispatch_model wraps to block moving an offloaded model."""
try:
from accelerate.hooks import _accelerate_added_attributes
return tuple(_accelerate_added_attributes)
except Exception:
return ("to", "cuda", "npu", "xpu", "mlu", "sdaa", "musa")
_ACCELERATE_MOVE_GUARDS = _accelerate_move_guards()
def _accelerate_dispatch_root(model):
"""The module that really owns the accelerate dispatch.
A PEFT wrapper only proxies ``_hf_hook``, so ``delattr`` fails and
``remove_hook_from_submodules`` raises before removing anything; ``hf_device_map``
keys are relative to the inner root too. Walks real children, never ``__getattr__``.
"""
node, seen = model, set()
while id(node) not in seen:
seen.add(id(node))
if "hf_device_map" in getattr(node, "__dict__", {}):
return node
children = getattr(node, "__dict__", {}).get("_modules") or {}
nxt = next(
(
children[a]
for a in ("base_model", "model")
if hasattr(children.get(a), "named_modules")
),
None,
)
if nxt is None:
return model
node = nxt
return model
def _snapshot_dispatch_state(root):
"""Hooks, tensor placements and instance forwards, so the dispatch can be replayed.
Re-deriving it with ``dispatch_model`` is not equivalent: PEFT reparents each
targeted ``Linear`` after transformers dispatched, so accelerate hooks modules that
never had any (measured: 395 -> 1379) and the logits shift enough to reorder top-5.
"""
hooks = [
(name, mod.__dict__["_hf_hook"])
for name, mod in root.named_modules()
if "_hf_hook" in mod.__dict__
]
# remove_duplicate=False: the default hides one half of every tied pair, exactly
# the half that needs re-tying below.
named = list(root.named_parameters(remove_duplicate = False)) + list(
root.named_buffers(remove_duplicate = False)
)
places = {name: tensor.device for name, tensor in named}
# Tied weights share one storage, but the CPU round trip repoints every tensor and
# tied_params_map is keyed on the old pointer, so replaying the hooks alone gives
# independent copies: double VRAM, and updates to one no longer reach the other.
# Skip meta tensors: offloaded parameters all sit on meta with pointer 0, which
# would collapse into one fake "tied" group of differently shaped tensors, and
# each side of a tie is already its own meta placeholder so nothing is lost.
groups = {}
for name, tensor in named:
if tensor.device.type == "meta":
continue
ptr = tensor.untyped_storage().data_ptr()
if ptr:
groups.setdefault(ptr, []).append(name)
ties = [names for names in groups.values() if len(names) > 1]
# Removing a hook restores `forward = _old_forward`, captured before unsloth patched
# the module, so a remove/re-add permanently drops every fused kernel installed after
# the dispatch (measured: apply_lora_mlp_swiglu on all 28 MLPs). It also deletes the
# `to`/`cuda`/... move guards, so record those too.
attrs = ("forward", "_old_forward") + tuple(_ACCELERATE_MOVE_GUARDS)
saved_attrs = {
name: {a: mod.__dict__[a] for a in attrs if a in mod.__dict__}
for name, mod in root.named_modules()
if any(a in mod.__dict__ for a in attrs)
}
# Re-adding a hook runs init_hook -> set_module_tensor_to_device, which builds a fresh
# Parameter and so drops .grad. Snapshot the gradients and reattach them on restore.
grads = {
name: getattr(tensor, "grad", None)
for name, tensor in root.named_parameters(remove_duplicate = False)
if getattr(tensor, "grad", None) is not None
}
return hooks, places, saved_attrs, ties, grads
def _drop_accelerator_tied_param_cache(snapshot) -> None:
"""Drop the GPU tensors accelerate caches in each hook's ``tied_params_map``.
Holding the hooks across the offload pins a GPU copy of the tied embedding (0.31 GB
of 1.24 GB here). The entries are keyed on the pre-move ``data_ptr`` so they are
stale anyway, and re-attaching repopulates them.
"""
for _name, hook in snapshot[0]:
cache = getattr(hook, "tied_params_map", None)
if not cache:
continue
for ptr in list(cache):
entry = cache[ptr]
for device in list(entry):
if str(device) != "cpu":
del entry[device]
if not entry:
del cache[ptr]
def _split_tensor_path(root, full_name):
"""``("model.embed_tokens.weight")`` -> ``(the module, "weight")``."""
mod_name, _, attr = full_name.rpartition(".")
try:
return (root.get_submodule(mod_name) if mod_name else root), attr
except AttributeError:
return None, attr
def _lookup_tensor(root, full_name):
mod, attr = _split_tensor_path(root, full_name)
if mod is None:
return None
for store in ("_parameters", "_buffers"):
found = (getattr(mod, store, None) or {}).get(attr)
if found is not None:
return found
return None
def _share_tensor(root, full_name, leader) -> None:
"""Point ``full_name`` back at ``leader``, restoring a tie."""
mod, attr = _split_tensor_path(root, full_name)
if mod is None:
return
for store in ("_parameters", "_buffers"):
target = getattr(mod, store, None)
if target is None or attr not in target:
continue
current = target[attr]
if current is None or current.device != leader.device or current.shape != leader.shape:
return # not actually the same tensor; leave it alone
target[attr] = leader
return
def _restore_dispatch_state(root, snapshot) -> None:
"""Replay ``_snapshot_dispatch_state``."""
from accelerate.hooks import add_hook_to_module
hooks, places, saved_attrs, ties, grads = snapshot
for name, hook in hooks:
add_hook_to_module(root.get_submodule(name) if name else root, hook)
# Re-adding a hook rewraps whatever `_old_forward` now holds, so put the exact
# callables back, `_old_forward` first.
for name, values in saved_attrs.items():
mod = root.get_submodule(name) if name else root
for attr in ("_old_forward", "forward", *_ACCELERATE_MOVE_GUARDS):
if attr in values:
mod.__dict__[attr] = values[attr]
# init_hook only re-places tensors the hooked module owns, so anything added after
# the dispatch (the LoRA adapters) is still on CPU.
for mod_name, mod in root.named_modules():
for attr in ("_parameters", "_buffers"):
store = getattr(mod, attr, None)
if not store:
continue
for tensor_name, tensor in list(store.items()):
if tensor is None:
continue
full = f"{mod_name}.{tensor_name}" if mod_name else tensor_name
want = places.get(full)
if want is None or tensor.device == want:
continue
if getattr(tensor, "quant_state", None) is not None:
# Only bitsandbytes' own .to() moves absmax/code/state2 with the data.
mod.to(want)
else:
tensor.data = tensor.data.to(want)
# Reattach the gradients init_hook discarded, on their weight's device.
for name, grad in grads.items():
tensor = _lookup_tensor(root, name)
if tensor is not None and tensor.grad is None and tensor.shape == grad.shape:
tensor.grad = grad.to(tensor.device)
# Re-tie last, once every tensor is back on its own device.
for names in ties:
leader = _lookup_tensor(root, names[0])
if leader is None:
continue
for follower in names[1:]:
_share_tensor(root, follower, leader)
# init_hook refilled tied_params_map with the pre-retie tensors, now unreferenced
# by the model but still pinned by the map.
if ties:
_drop_accelerator_tied_param_cache(snapshot)
def _offload_model_for_quantize_subprocess(model):
"""Best-effort: move the model's weights off the GPU before the quantized export
loads its own copy from disk, so the GPUs need not hold both at once. Returns an
opaque token for ``_restore_model_after_quantize_subprocess`` (None if nothing moved).
Two shapes are handled:
* single-device CUDA/XPU model -> ``.to("cpu")``, restored with ``.to(device)``;
* accelerate-dispatched model (a multi-GPU ``device_map`` shard, e.g. the Studio
multi-GPU export load) -> hooks removed and moved to CPU, restored by replaying
the dispatch. A plain ``.to("cpu")`` is invalid here, which is why the old
single-device-only move left every GPU holding a full copy. A map spilling to
CPU is still released, but disk/meta targets are left alone: accelerate keeps
those parameters off the model, so moving would materialize the whole checkpoint.
Quantized (bnb) models are attempted too rather than skipped: Studio exports load
4-bit by DEFAULT, so skipping them left a shard on every GPU. transformers refuses
``.to()`` for some bitsandbytes builds and that refusal raises before anything moves,
so the failure path restores the model and returns None, i.e. the old behaviour.
"""
try:
_has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
if not ((torch.cuda.is_available() or _has_xpu) and hasattr(model, "parameters")):
return None
device_map = getattr(model, "hf_device_map", None)
if device_map:
targets = {str(v).lower() for v in device_map.values()}
# A cpu spill is fine to move, it is already in host RAM. disk/meta is not:
# those parameters are off the model, so .to("cpu") would materialize the
# whole checkpoint into RAM.
if not all(t.isdigit() or t.startswith(("cuda", "xpu")) or t == "cpu" for t in targets):
return None
if not any(t.isdigit() or t.startswith(("cuda", "xpu")) for t in targets):
return None # nothing on an accelerator: no GPU memory to reclaim
from accelerate.hooks import remove_hook_from_submodules
# A PEFT wrapper only proxies the hooks; they live on the inner root.
root = _accelerate_dispatch_root(model)
try:
setattr(root, _DISPATCH_SNAPSHOT_ATTR, _snapshot_dispatch_state(root))
except Exception as snap_exc:
# Restore will fall back to re-deriving from the device_map.
logger.warning_once(
f"Unsloth: could not snapshot the accelerate dispatch "
f"({type(snap_exc).__name__}: {snap_exc}); re-dispatching on restore."
)
remove_hook_from_submodules(root)
try:
model.to("cpu")
except Exception:
# The move failed after the hooks came off; re-dispatch so the model is
# left usable rather than hookless and half-moved across CPU/GPUs.
_restore_model_after_quantize_subprocess(model, ("dispatch", dict(device_map)))
return None
snapshot = getattr(root, _DISPATCH_SNAPSHOT_ATTR, None)
if snapshot is not None:
_drop_accelerator_tied_param_cache(snapshot)
return ("dispatch", dict(device_map))
devices = {str(p.device) for p in model.parameters()}
if len(devices) == 1 and next(iter(devices)).startswith(("cuda", "xpu")):
device = next(model.parameters()).device
try:
model.to("cpu")
except Exception:
_restore_model_after_quantize_subprocess(model, ("device", device))
return None
return ("device", device)
except Exception as exc:
# A silent `return None` is indistinguishable from "nothing to move", which
# hides a real bug behind a merely slower export.
logger.warning_once(
f"Unsloth: could not free the model's accelerator memory before the quantized "
f"export ({type(exc).__name__}: {exc}); continuing with the model resident."
)
return None
return None
def _restore_model_after_quantize_subprocess(model, restore_token) -> None:
"""Undo ``_offload_model_for_quantize_subprocess``; warns instead of raising."""
if restore_token is None:
return
kind, value = restore_token
try:
if kind == "dispatch":
root = _accelerate_dispatch_root(model)
snapshot = root.__dict__.pop(_DISPATCH_SNAPSHOT_ATTR, None)
if snapshot is not None:
_restore_dispatch_state(root, snapshot)
else:
from accelerate import dispatch_model
# skip_keys matters: without it accelerate moves every forward kwarg
# to the executing device, wrong for device-invariant cache tensors.
dispatch_model(
root,
device_map = value,
skip_keys = getattr(root, "_skip_keys_device_placement", None),
)
else:
model.to(value) # restore the model to its original device
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device(s) after the "
"quantized export; it may remain on CPU."
)
def _unsloth_save_compressed_tensors(
model,
save_directory: Union[str, os.PathLike],
@ -4654,7 +5017,7 @@ def _unsloth_save_compressed_tensors(
# 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and
# quantize inside an isolated temp dir instead of writing ./<repo_id> into the cwd.
repo_id, work_tmp, calib_tmp, model_dev = None, None, None, None
repo_id, work_tmp, calib_tmp, model_restore = None, None, None, None
if push_to_hub:
repo_id = os.fspath(save_directory)
work_tmp = tempfile.mkdtemp(prefix = "unsloth-compressed-")
@ -4806,23 +5169,8 @@ def _unsloth_save_compressed_tensors(
cmd += ["--variant", variant]
# Free the in-memory model's CUDA memory before the subprocess loads its own copy from
# disk, so a single GPU need not hold both at once. Best-effort and restored in finally;
# skipped for quantized or multi-device models where moving is unsafe.
try:
if (
torch.cuda.is_available()
and hasattr(model, "parameters")
and not getattr(model, "is_loaded_in_4bit", False)
and not getattr(model, "is_loaded_in_8bit", False)
and not getattr(model, "is_quantized", False)
):
_devs = {str(p.device) for p in model.parameters()}
if len(_devs) == 1 and next(iter(_devs)).startswith("cuda"):
_dev = next(model.parameters()).device
model.to("cpu")
model_dev = _dev # set only after a successful move, so finally can restore
except Exception:
model_dev = None
# disk, so the GPUs need not hold both at once. Best-effort, restored in finally.
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
@ -4893,14 +5241,7 @@ def _unsloth_save_compressed_tensors(
_print_compressed_hw_note(scheme, result)
return result
finally:
if model_dev is not None:
try:
model.to(model_dev) # restore the model to its original device
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device after compressed "
"export; it may remain on CPU."
)
_restore_model_after_quantize_subprocess(model, model_restore)
if calib_tmp is not None and os.path.isdir(calib_tmp):
shutil.rmtree(calib_tmp, ignore_errors = True)
if work_tmp is not None:
@ -4959,7 +5300,7 @@ def _unsloth_save_torchao(
# Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected
# 16-bit export written to save_directory is not overwritten or deleted; the torchao output is
# the sibling "<save_directory>-<suffix>" (or the repo id on a hub push).
repo_id, work_tmp, model_dev = None, None, None
repo_id, work_tmp, model_restore = None, None, None
work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-")
if push_to_hub:
repo_id = os.fspath(save_directory)
@ -5037,25 +5378,12 @@ def _unsloth_save_torchao(
auto_model = AutoModelForCausalLM
auto_processor = AutoProcessor if is_vlm else AutoTokenizer
# 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk.
# Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit
# resident alongside the reloaded copy and OOM a device that fit the model once.
# 3) Free the in-memory model's accelerator memory before reloading a fresh copy from
# disk, else it sits resident alongside the copy and OOMs a device that fit the
# model once. Covers CUDA and XPU (torchao runs on Intel GPUs too) plus multi-GPU
# dispatched shards, which a plain .to("cpu") cannot move.
_has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
try:
if (
(torch.cuda.is_available() or _has_xpu)
and hasattr(model, "parameters")
and not getattr(model, "is_loaded_in_4bit", False)
and not getattr(model, "is_loaded_in_8bit", False)
and not getattr(model, "is_quantized", False)
):
_devs = {str(p.device) for p in model.parameters()}
if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")):
_dev = next(model.parameters()).device
model.to("cpu")
model_dev = _dev
except Exception:
model_dev = None
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
@ -5126,14 +5454,20 @@ def _unsloth_save_torchao(
)
return result
finally:
if model_dev is not None:
try:
model.to(model_dev)
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device after torchao "
"export; it may remain on CPU."
)
# A raise pins the copy in the local and the live traceback, so free both or the
# restore below OOMs.
quantized_model = None
del quantized_model
_exc = sys.exc_info()[1]
if _exc is not None:
traceback.clear_frames(_exc.__traceback__)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
_restore_model_after_quantize_subprocess(model, model_restore)
if work_tmp is not None:
shutil.rmtree(work_tmp, ignore_errors = True)
for _ in range(3):