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.

tests/test_compressed_export_gpu_release.py

@@ -143,11 +143,40 @@ def test_single_device_move_failure_restores_and_returns_none():
     assert model.moved_to == ["cpu", "cuda:0"]
 
 
-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.
+def test_cpu_spilled_map_still_releases_its_gpu_shards(_fake_accelerate):
+    # accelerate spilled one module to CPU, but the rest is still resident on the
+    # GPUs -- exactly the memory the subprocess/reload needs. Those weights are in
+    # host RAM already, so moving is safe and the GPU shards get reclaimed.
     ns = _load_helpers(_fake_torch(), _FakeLogger())
-    model = _FakeModel(device_map = {"model.embed": 0, "model.layers.9": "cpu"})
+    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: accelerate streams them from disk,
+    # so removing hooks and moving would try to materialize the whole checkpoint
+    # into RAM. Leave it untouched.
+    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: there is 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"] == []
@@ -164,11 +193,29 @@ def test_single_device_model_keeps_plain_move():
     assert model.moved_to[-1] == "cuda:0"
 
 
-def test_quantized_model_is_never_moved():
+def test_quantized_model_is_released_when_the_stack_allows_it():
+    # The Studio export path loads 4-bit by DEFAULT, so skipping quantized models
+    # left a shard on every GPU while the subprocess/reload allocated another
+    # copy. 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 that refusal
+    # raises before anything moves -- so the result must be exactly the old
+    # behaviour: no token, model untouched, no exception 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
-    assert model.moved_to == []
 
 
 def test_no_cuda_is_noop_and_restore_none_is_noop():

unsloth/save.py

@@ -4621,26 +4621,36 @@ def _offload_model_for_quantize_subprocess(model):
         ``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.
+        A map that spills to CPU is still released -- those weights are in host
+        RAM and the GPU-mapped modules are the ones worth reclaiming -- but a map
+        with disk/meta targets is left alone: accelerate keeps those parameters
+        off the model, so moving would try to materialize the whole checkpoint.
 
-    Quantized (bnb) models are never moved.
+    Quantized (bnb) models are attempted too rather than skipped outright: the
+    Studio export path loads 4-bit by DEFAULT, so skipping them left a quantized
+    shard on every GPU while the subprocess or the torchao reload allocated
+    another copy. transformers refuses ``.to()`` for some bitsandbytes builds, and
+    that refusal raises before anything moves, so the existing failure path simply
+    restores the model and returns None -- i.e. best-effort where the stack allows
+    it, and exactly the old behaviour where it does not.
     """
     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
-        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", "xpu")) for t in targets):
+            targets = {str(v).lower() for v in device_map.values()}
+            # cpu spill is fine to move -- those weights are already in host RAM.
+            # disk/meta is NOT: accelerate keeps those parameters off the model
+            # entirely, so removing the hooks and calling .to("cpu") would try to
+            # materialize the whole checkpoint into RAM (or move meta tensors).
+            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
 
             remove_hook_from_submodules(model)