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

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

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-26 07:19:48 +00:00
commit 8b6b4ca0b5
3 changed files with 389 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -112,6 +112,26 @@ def _multi_gpu_device_map_kwargs() -> dict:
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 the ROCm/XPU backends 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 _supports_kwarg(fn, name):
"""True if `fn` accepts keyword `name` directly or via **kwargs."""
import inspect
@ -302,6 +322,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.
@ -336,7 +357,12 @@ class ExportBackend:
# Shard across every visible GPU on a multi-GPU host instead of
# stacking the checkpoint on GPU0 (#7053); {} on single-GPU/CPU/MLX.
_device_map_kw = _multi_gpu_device_map_kwargs()
# _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
@ -471,11 +497,39 @@ 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 its own CUDA context on each
# GPU, so when a training or chat job already owns the others (which
# routes/export.py deliberately allows) the shard can OOM where the pre-#7053
# single-device load succeeded. Fall back once before giving up.
if (
_device_map_override is None
and _is_oom_error(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.
oom_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 ran out of memory ({oom_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."""

166
tests/test_compressed_export_gpu_release.py
View file

@ -22,9 +22,14 @@ import pytest
_SAVE_PY = Path(__file__).resolve().parent.parent / "unsloth" / "save.py"
_WANTED = {
"_accelerate_dispatch_root",
"_snapshot_dispatch_state",
"_drop_accelerator_tied_param_cache",
"_restore_dispatch_state",
"_offload_model_for_quantize_subprocess",
"_restore_model_after_quantize_subprocess",
}
_WANTED_ASSIGNS = {"_DISPATCH_SNAPSHOT_ATTR"} # module constants the helpers close over
class _FakeLogger:
@ -38,9 +43,17 @@ class _FakeLogger:
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
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
)
)
]
assert len(keep) == len(_WANTED), "release helpers missing from save.py"
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"),
@ -78,13 +91,19 @@ class _FakeModel:
@pytest.fixture
def _fake_accelerate(monkeypatch):
calls = {"removed": [], "dispatched": []}
calls = {"removed": [], "dispatched": [], "dispatch_kwargs": [], "hooks_added": []}
accel = types.ModuleType("accelerate")
accel.dispatch_model = lambda model, device_map: calls["dispatched"].append(
(model, dict(device_map))
)
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)
@ -312,3 +331,138 @@ def test_torchao_export_uses_the_shared_release():
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 this PR newly makes reachable.
# --------------------------------------------------------------------------- #
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):
return iter(())
def named_buffers(self):
return iter(())
class _PeftLikeWrapper(_Child):
"""A wrapper that proxies unknown attributes to the model it wraps, the way
``PeftModelForCausalLM`` does. ``hasattr(wrapper, "_hf_hook")`` is then 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", which is
# how a broken offload survived several review rounds.
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 would move every forward
# kwarg to the executing device -- wrong for the cache/position 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 ``module.forward = module._old_forward`` on removal, and
``_old_forward`` is whatever the forward was when the hook was FIRST attached.
unsloth patches module forwards after transformers dispatches, 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

181
unsloth/save.py
View file

@ -4545,13 +4545,26 @@ def _unsloth_save_torchao_with_given_config(
else:
kwargs = {"dtype": torch.bfloat16}
# Reload with quantization applied
quantized_model = auto_model.from_pretrained(
save_directory,
device_map = "auto",
quantization_config = quantization_config,
**kwargs,
)
# Same guard as the sibling quantized-export paths: without it the original
# stays resident on every GPU while device_map="auto" 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()
try:
# Reload with quantization applied
quantized_model = auto_model.from_pretrained(
save_directory,
device_map = "auto",
quantization_config = quantization_config,
**kwargs,
)
finally:
_restore_model_after_quantize_subprocess(model, model_restore)
torchao_save_directory = save_directory + "-torchao"
@ -4606,6 +4619,122 @@ def _print_compressed_hw_note(scheme, out_dir):
)
_DISPATCH_SNAPSHOT_ATTR = "_unsloth_dispatch_snapshot"
def _accelerate_dispatch_root(model):
"""The module that really owns the accelerate dispatch.
A PEFT wrapper proxies unknown attributes to the model it wraps, so
``hasattr(wrapper, "_hf_hook")`` is True but ``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__
]
places = {
name: tensor.device
for name, tensor in list(root.named_parameters()) + list(root.named_buffers())
}
# Removing a hook restores `forward = _old_forward`, captured before unsloth
# patched the module, so a remove/re-add drops every fused kernel installed
# after the dispatch (measured: apply_lora_mlp_swiglu on all 28 MLPs) for good.
forwards = {
name: (mod.__dict__.get("forward"), mod.__dict__.get("_old_forward"))
for name, mod in root.named_modules()
if "forward" in mod.__dict__ or "_old_forward" in mod.__dict__
}
return hooks, places, forwards
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 by 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 _restore_dispatch_state(root, snapshot) -> None:
"""Replay ``_snapshot_dispatch_state``."""
from accelerate.hooks import add_hook_to_module
hooks, places, forwards = 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.
for name, (fwd, old_fwd) in forwards.items():
mod = root.get_submodule(name) if name else root
for attr, value in (("_old_forward", old_fwd), ("forward", fwd)):
if value is None:
mod.__dict__.pop(attr, None)
else:
mod.__dict__[attr] = value
# init_hook only re-places tensors the hooked module owns, so anything added to
# the tree 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)
def _offload_model_for_quantize_subprocess(model):
"""Best-effort: move the merged model's weights off the GPU before the
llm-compressor subprocess loads its own copy from disk, so the GPUs need not
@ -4651,7 +4780,17 @@ def _offload_model_for_quantize_subprocess(model):
return None # nothing on an accelerator: no GPU memory to reclaim
from accelerate.hooks import remove_hook_from_submodules
remove_hook_from_submodules(model)
# 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:
@ -4660,6 +4799,9 @@ def _offload_model_for_quantize_subprocess(model):
# 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")):
@ -4670,7 +4812,13 @@ def _offload_model_for_quantize_subprocess(model):
_restore_model_after_quantize_subprocess(model, ("device", device))
return None
return ("device", device)
except Exception:
except Exception as exc:
# A silent `return None` here 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
@ -4682,8 +4830,19 @@ def _restore_model_after_quantize_subprocess(model, restore_token) -> None:
kind, value = restore_token
try:
if kind == "dispatch":
from accelerate import dispatch_model
dispatch_model(model, device_map = value)
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: