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

137
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,11 +528,41 @@ class ExportBackend:
return True, f"Loaded {model_type} model{peft_info} successfully"
except Exception as e:
logger.error(f"Error loading checkpoint: {e}")
import traceback
# 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.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."""

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