From a58f1086bf0f43b70479800a8b1d260ed189b201 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 07:41:20 +0000 Subject: [PATCH] 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. --- tests/test_compressed_export_gpu_release.py | 50 ++++++++- unsloth/save.py | 115 ++++++++++++++++---- 2 files changed, 140 insertions(+), 25 deletions(-) diff --git a/tests/test_compressed_export_gpu_release.py b/tests/test_compressed_export_gpu_release.py index d3111a1b90..a150ad3a12 100644 --- a/tests/test_compressed_export_gpu_release.py +++ b/tests/test_compressed_export_gpu_release.py @@ -25,6 +25,9 @@ _WANTED = { "_accelerate_dispatch_root", "_snapshot_dispatch_state", "_drop_accelerator_tied_param_cache", + "_split_tensor_path", + "_lookup_tensor", + "_share_tensor", "_restore_dispatch_state", "_offload_model_for_quantize_subprocess", "_restore_model_after_quantize_subprocess", @@ -360,10 +363,10 @@ class _Child: node = node._modules[part] return node - def named_parameters(self): + def named_parameters(self, remove_duplicate = True): return iter(()) - def named_buffers(self): + def named_buffers(self, remove_duplicate = True): return iter(()) @@ -467,3 +470,46 @@ def test_snapshot_restores_a_forward_patched_after_the_dispatch(_fake_accelerate ns["_restore_dispatch_state"](root, snapshot) assert mlp.forward() == "unsloth-fused" assert mlp.__dict__["_old_forward"] is stock_forward + + +def test_snapshot_reties_shared_parameters(_fake_accelerate): + """A CPU round trip repoints every tensor, so replaying the hooks alone leaves + formerly tied weights as independent copies: double VRAM, and an update to one + that never reaches the other.""" + import torch + + root = _Child(device_map = {"embed": 0, "head": 0}) + shared = torch.nn.Parameter(torch.zeros(4, 4)) + for name in ("embed", "head"): + child = _Child(name = name) + child._parameters = {"weight": shared} + child._buffers = {} + root._modules[name] = child + + def named(remove_duplicate = True): + seen, out = set(), [] + for mod_name, mod in root._modules.items(): + for attr, tensor in mod._parameters.items(): + if remove_duplicate and id(tensor) in seen: + continue + seen.add(id(tensor)) + out.append((f"{mod_name}.{attr}", tensor)) + return iter(out) + + root.named_parameters = named + ns = _load_helpers(_fake_torch(), _FakeLogger()) + snapshot = ns_ties = ns["_snapshot_dispatch_state"](root) + assert ns_ties[3] == [["embed.weight", "head.weight"]] + + # what the replay leaves behind before the retie step + root._modules["head"]._parameters["weight"] = torch.nn.Parameter(shared.detach().clone()) + assert ( + root._modules["embed"]._parameters["weight"].data_ptr() + != root._modules["head"]._parameters["weight"].data_ptr() + ) + + ns["_restore_dispatch_state"](root, snapshot) + assert ( + root._modules["embed"]._parameters["weight"].data_ptr() + == root._modules["head"]._parameters["weight"].data_ptr() + ) diff --git a/unsloth/save.py b/unsloth/save.py index c57a1737cb..95ec43a387 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -4555,6 +4555,9 @@ def _unsloth_save_torchao_with_given_config( if hasattr(torch, "xpu") and torch.xpu.is_available(): torch.xpu.empty_cache() + # The original stays offloaded until the quantized copy has been saved AND + # released, otherwise both are resident at once and the model that only fit + # because of the offload OOMs on the way back. try: # Reload with quantization applied quantized_model = auto_model.from_pretrained( @@ -4563,26 +4566,34 @@ def _unsloth_save_torchao_with_given_config( quantization_config = quantization_config, **kwargs, ) + + torchao_save_directory = save_directory + "-torchao" + + # TorchAO does not support safe_serialization right now 0.14.0 seems broken! + safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0") + safe_serialization = False + + if push_to_hub: + quantized_model.push_to_hub( + torchao_save_directory, safe_serialization = safe_serialization, token = token + ) + tokenizer.push_to_hub(torchao_save_directory, token = token) + else: + quantized_model.save_pretrained( + torchao_save_directory, safe_serialization = safe_serialization + ) + tokenizer.save_pretrained(torchao_save_directory, token = token) + + del quantized_model finally: + 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() _restore_model_after_quantize_subprocess(model, model_restore) - torchao_save_directory = save_directory + "-torchao" - - # TorchAO does not support safe_serialization right now 0.14.0 seems broken! - safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0") - safe_serialization = False - - if push_to_hub: - quantized_model.push_to_hub( - torchao_save_directory, safe_serialization = safe_serialization, token = token - ) - tokenizer.push_to_hub(torchao_save_directory, token = token) - else: - quantized_model.save_pretrained( - torchao_save_directory, safe_serialization = safe_serialization - ) - tokenizer.save_pretrained(torchao_save_directory, token = token) - # Clean up the intermediate unquantized model if os.path.exists(save_directory): try: @@ -4662,10 +4673,20 @@ def _snapshot_dispatch_state(root): 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()) - } + # remove_duplicate=False: the default hides one half of every tied pair, which is + # exactly the half that needs re-tying below. + named = list(root.named_parameters(remove_duplicate = False)) + list( + root.named_buffers(remove_duplicate = False) + ) + places = {name: tensor.device for name, tensor in named} + # Tied weights share one storage. The CPU round trip repoints every tensor and + # accelerate's tied_params_map is keyed on the old pointer, so replaying the + # hooks alone materialises independent copies: double VRAM, and later updates + # to one no longer reach the other. + groups = {} + for name, tensor in named: + groups.setdefault(tensor.untyped_storage().data_ptr(), []).append(name) + ties = [names for names in groups.values() if len(names) > 1] # 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. @@ -4674,7 +4695,7 @@ def _snapshot_dispatch_state(root): for name, mod in root.named_modules() if "forward" in mod.__dict__ or "_old_forward" in mod.__dict__ } - return hooks, places, forwards + return hooks, places, forwards, ties def _drop_accelerator_tied_param_cache(snapshot) -> None: @@ -4697,11 +4718,47 @@ def _drop_accelerator_tied_param_cache(snapshot) -> None: del cache[ptr] +def _split_tensor_path(root, full_name): + """``("model.embed_tokens.weight")`` -> ``(the module, "weight")``.""" + mod_name, _, attr = full_name.rpartition(".") + try: + return (root.get_submodule(mod_name) if mod_name else root), attr + except AttributeError: + return None, attr + + +def _lookup_tensor(root, full_name): + mod, attr = _split_tensor_path(root, full_name) + if mod is None: + return None + for store in ("_parameters", "_buffers"): + found = (getattr(mod, store, None) or {}).get(attr) + if found is not None: + return found + return None + + +def _share_tensor(root, full_name, leader) -> None: + """Point ``full_name`` back at ``leader``, restoring a tie.""" + mod, attr = _split_tensor_path(root, full_name) + if mod is None: + return + for store in ("_parameters", "_buffers"): + target = getattr(mod, store, None) + if target is None or attr not in target: + continue + current = target[attr] + if current is None or current.device != leader.device: + return # accelerate split the tie across devices; leave it alone + target[attr] = leader + return + + def _restore_dispatch_state(root, snapshot) -> None: """Replay ``_snapshot_dispatch_state``.""" from accelerate.hooks import add_hook_to_module - hooks, places, forwards = snapshot + hooks, places, forwards, ties = snapshot for name, hook in hooks: add_hook_to_module(root.get_submodule(name) if name else root, hook) @@ -4735,6 +4792,18 @@ def _restore_dispatch_state(root, snapshot) -> None: else: tensor.data = tensor.data.to(want) + # Re-tie last, once every tensor is back on its own device. + for names in ties: + leader = _lookup_tensor(root, names[0]) + if leader is None: + continue + for follower in names[1:]: + _share_tensor(root, follower, leader) + # init_hook refilled tied_params_map with the pre-retie tensors; those are now + # unreferenced by the model but still pinned by the map. + if ties: + _drop_accelerator_tied_param_cache(snapshot) + def _offload_model_for_quantize_subprocess(model): """Best-effort: move the merged model's weights off the GPU before the