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.

studio/backend/core/export/export.py

@@ -54,8 +54,15 @@ def _multi_gpu_device_map_kwargs() -> dict:
         visible = get_parent_visible_gpu_ids()
         if len(visible) > 1:
             device_map = get_device_map(visible)
-            if device_map == "balanced":
-                return {"device_map": device_map}
+        elif not visible:
+            # UUID/MIG CUDA_VISIBLE_DEVICES masks resolve to no numeric ids;
+            # get_device_map(None) handles exactly that by falling 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 {}

studio/backend/tests/test_export_multi_gpu_device_map.py

@@ -66,6 +66,33 @@ def test_non_balanced_resolution_keeps_loader_default(monkeypatch):
     assert mod._multi_gpu_device_map_kwargs() == {}
 
 
+def test_uuid_mig_mask_falls_back_to_count_detection(monkeypatch):
+    # UUID/MIG CUDA_VISIBLE_DEVICES masks resolve to NO numeric ids ([]), but
+    # get_device_map(None) still detects >1 visible GPU; the empty list must
+    # route there instead of silently keeping the sequential 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)
     # _install_export_backend_stubs sets _IS_MLX = True already; even a stubbed

tests/test_compressed_export_gpu_release.py

@@ -0,0 +1,152 @@
+# 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 (e.g. Studio's multi-GPU export load),
+which the old single-device-only ``.to("cpu")`` skipped, leaving every GPU
+holding a full copy alongside the subprocess's.
+
+Loads only the two release/restore helpers from unsloth/save.py via AST (the
+module itself needs torch/transformers), and exercises them with fakes.
+"""
+
+from __future__ import annotations
+
+import ast
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+_SAVE_PY = Path(__file__).resolve().parent.parent / "unsloth" / "save.py"
+_WANTED = {
+    "_offload_model_for_quantize_subprocess",
+    "_restore_model_after_quantize_subprocess",
+}
+
+
+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
+    ]
+    assert len(keep) == 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": []}
+    accel = types.ModuleType("accelerate")
+    accel.dispatch_model = lambda model, device_map: calls["dispatched"].append(
+        (model, dict(device_map))
+    )
+    hooks = types.ModuleType("accelerate.hooks")
+    hooks.remove_hook_from_submodules = lambda model: calls["removed"].append(model)
+    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_cpu_or_disk_offloaded_map_is_left_alone(_fake_accelerate):
+    # accelerate is already offloading part of the model; removing hooks and
+    # re-dispatching could thrash, so leave it untouched.
+    ns = _load_helpers(_fake_torch(), _FakeLogger())
+    model = _FakeModel(device_map = {"model.embed": 0, "model.layers.9": "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_never_moved():
+    ns = _load_helpers(_fake_torch(), _FakeLogger())
+    model = _FakeModel(devices = ("cuda:0",), quantized = True)
+    assert ns["_offload_model_for_quantize_subprocess"](model) is None
+    assert model.moved_to == []
+
+
+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

unsloth/save.py

@@ -4078,6 +4078,73 @@ def _print_compressed_hw_note(scheme, out_dir):
     )
 
 
+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
+    hold both at once. Returns an opaque restore token for
+    ``_restore_model_after_quantize_subprocess`` (None when nothing was moved).
+
+    Two shapes are handled:
+      * single-device CUDA model -> plain ``.to("cpu")``, restored with
+        ``.to(device)``;
+      * accelerate-dispatched model (a multi-GPU ``device_map`` shard, e.g. the
+        Studio multi-GPU export load) -> accelerate's hooks are removed and the
+        model moved to CPU, restored by re-dispatching over the recorded
+        ``hf_device_map``. A plain ``.to("cpu")`` is invalid on a dispatched
+        model, which is why the old single-device-only move skipped these and
+        left every GPU holding a full copy while the subprocess loaded another.
+        A map with non-GPU targets (cpu/disk offload) is left alone: accelerate
+        is already offloading those weights.
+
+    Quantized (bnb) models are never moved.
+    """
+    try:
+        if not (torch.cuda.is_available() and hasattr(model, "parameters")):
+            return None
+        if (
+            getattr(model, "is_loaded_in_4bit", False)
+            or getattr(model, "is_loaded_in_8bit", False)
+            or getattr(model, "is_quantized", False)
+        ):
+            return None
+        device_map = getattr(model, "hf_device_map", None)
+        if device_map:
+            targets = {str(v) for v in device_map.values()}
+            if not all(t.isdigit() or t.startswith("cuda") for t in targets):
+                return None
+            from accelerate.hooks import remove_hook_from_submodules
+
+            remove_hook_from_submodules(model)
+            model.to("cpu")
+            return ("dispatch", dict(device_map))
+        devices = {str(p.device) for p in model.parameters()}
+        if len(devices) == 1 and next(iter(devices)).startswith("cuda"):
+            device = next(model.parameters()).device
+            model.to("cpu")
+            return ("device", device)
+    except Exception:
+        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":
+            from accelerate import dispatch_model
+            dispatch_model(model, device_map = value)
+        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 compressed "
+            "export; it may remain on CPU."
+        )
+
+
 def _unsloth_save_compressed_tensors(
     model,
     save_directory: Union[str, os.PathLike],
@@ -4133,7 +4200,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-")
@@ -4279,23 +4346,9 @@ 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 and restored in finally;
+        # covers single-device CUDA models and accelerate-dispatched multi-GPU shards.
+        model_restore = _offload_model_for_quantize_subprocess(model)
         for _ in range(3):
             gc.collect()
             if torch.cuda.is_available():
@@ -4359,14 +4412,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: