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.

tests/test_compressed_export_gpu_release.py

@@ -106,6 +106,43 @@ def test_dispatched_multi_gpu_model_is_released_and_redispatched(_fake_accelerat
     assert _fake_accelerate["dispatched"] == [(model, device_map)]
 
 
+def test_dispatched_move_failure_redispatches_and_returns_none(_fake_accelerate):
+    # If .to("cpu") raises AFTER the accelerate hooks are removed, the model
+    # must be re-dispatched (not left hookless/half-moved) and offload aborts.
+    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_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.
@@ -153,3 +190,26 @@ def test_restore_failure_warns_instead_of_raising(_fake_accelerate):
     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 a model sharded
+    # across GPUs (device_map="balanced") must be budgeted against W's own
+    # device, not GPU0 -- otherwise GPU1+ can OOM while only GPU0 is checked
+    # (#7053). Pin the device-aware budget in the LoRA-merge save path.
+    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

unsloth/save.py

@@ -1122,7 +1122,21 @@ def unsloth_save_model(
         torch_dtype
     )
 
-    max_vram = int(torch.cuda.get_device_properties(0).total_memory * maximum_memory_usage)
+    # Per-device VRAM budget: a merged tensor lives on the GPU of its source
+    # layer, so a model sharded across GPUs (device_map="balanced") must be
+    # budgeted against W's own device, not GPU0 -- otherwise GPU1+ can OOM as
+    # their merged weights accumulate while only GPU0's headroom 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 +1152,14 @@ 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 -> keep it there.
+                state_dict[name] = W
+            elif W.device.type != "cuda":
+                # Already off-GPU (offloaded layer): keeping it costs no VRAM.
                 state_dict[name] = W
             # [TODO] Saving to RAM seems to leak memory???
             # elif (max_ram - W.nbytes) > 0:
@@ -4622,12 +4642,23 @@ def _offload_model_for_quantize_subprocess(model):
             from accelerate.hooks import remove_hook_from_submodules
 
             remove_hook_from_submodules(model)
-            model.to("cpu")
+            try:
+                model.to("cpu")
+            except Exception:
+                # The move failed (e.g. host RAM can't hold the sharded model)
+                # AFTER the hooks were removed; 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
             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")
+            try:
+                model.to("cpu")
+            except Exception:
+                _restore_model_after_quantize_subprocess(model, ("device", device))
+                return None
             return ("device", device)
     except Exception:
         return None