diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index f35b365..e6bebf8 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -6609,54 +6609,52 @@ class AbliterationPipeline: self.log(f"Cleaned up offload dir ({size_mb:.0f} MiB reclaimed)") def _gather_state_dict(self) -> dict: - """Gather a complete state dict, materializing any meta tensors. + """Gather a complete state dict, materializing any offloaded tensors. - When device_map="auto" offloads weights to disk, model.state_dict() - returns meta tensors (no data) for those parameters. We resolve them - here so that save_pretrained gets real tensors. + Uses accelerate's ``get_state_dict_offloaded_model`` which temporarily + moves each offloaded module to CPU via ``align_module_device``, + extracting real tensor data regardless of whether weights were offloaded + to disk or CPU via ``AlignDevicesHook``. + + Falls back to plain ``model.state_dict()`` when no offloaded params + are detected. """ - model = self.handle.model - state_dict = model.state_dict() + from accelerate.utils.modeling import get_state_dict_offloaded_model - # Check for meta tensors (= disk-offloaded weights) + model = self.handle.model + + # Fast path: if nothing is offloaded, just return the normal state dict. + state_dict = model.state_dict() meta_keys = [k for k, v in state_dict.items() if v.device.type == "meta"] if not meta_keys: return state_dict - # Resolve meta tensors from the offload folder - offload_dir = getattr(self.handle, "_offload_dir", None) - if not offload_dir or not Path(offload_dir).exists(): - raise RuntimeError( - f"Cannot save model: {len(meta_keys)} weight tensors are on meta device " - f"(disk-offloaded) but the offload directory is missing " - f"(path={offload_dir!r}). This means those weights cannot be " - f"materialised and the saved model would be corrupted. " - f"Aborting to prevent writing a bricked checkpoint." - ) + self.log(f"Materializing {len(meta_keys)} offloaded tensors via accelerate...") - self.log(f"Materializing {len(meta_keys)} disk-offloaded tensors...") - from safetensors.torch import load_file + # Free GPU memory before materialization. get_state_dict_offloaded_model + # uses align_module_device which, for non-hooked GPU modules, temporarily + # moves params to CPU then restores them to CUDA in __exit__. If the GPU + # is near-full this restore triggers CUDA OOM. Moving all GPU params to + # CPU first prevents this. + cuda_count = 0 + for param in model.parameters(): + if param.device.type == "cuda": + param.data = param.data.cpu() + cuda_count += 1 + for buf in model.buffers(): + if buf.device.type == "cuda": + buf.data = buf.data.cpu() + cuda_count += 1 + if cuda_count: + torch.cuda.empty_cache() + self.log(f"Moved {cuda_count} GPU tensors to CPU, freed VRAM") - # Accelerate stores offloaded weights as individual safetensors files - for key in meta_keys: - safetensors_file = Path(offload_dir) / f"{key}.safetensors" - dat_file = Path(offload_dir) / f"{key}.dat" - if safetensors_file.exists(): - data = load_file(str(safetensors_file)) - state_dict[key] = data[key] if key in data else next(iter(data.values())) - elif dat_file.exists(): - # Accelerate's .dat format: raw tensor bytes with shape/dtype metadata - import numpy as np - dtype = state_dict[key].dtype - shape = state_dict[key].shape - arr = np.fromfile(str(dat_file), dtype=torch.tensor([], dtype=dtype).numpy().dtype) - state_dict[key] = torch.from_numpy(arr).reshape(shape) + state_dict = get_state_dict_offloaded_model(model) still_meta = sum(1 for v in state_dict.values() if v.device.type == "meta") if still_meta: raise RuntimeError( - f"Materialization incomplete: {still_meta} tensors still on meta device " - f"after loading from offload dir {offload_dir!r}. " + f"Materialization incomplete: {still_meta} tensors still on meta device. " f"Aborting to prevent writing a bricked checkpoint." ) diff --git a/tests/test_abliterate.py b/tests/test_abliterate.py index b30b4c6..d6db36e 100644 --- a/tests/test_abliterate.py +++ b/tests/test_abliterate.py @@ -4,7 +4,7 @@ from __future__ import annotations import json from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -1631,14 +1631,16 @@ class TestNuclearMethod: "Capability expert should be unchanged" ) - def test_gather_state_dict_raises_on_missing_offload(self): - """Should raise RuntimeError (not silently corrupt) when offload dir is missing.""" + def test_gather_state_dict_raises_on_unmaterialized_meta(self): + """Should raise RuntimeError when meta tensors survive materialization.""" from obliteratus.models.loader import ModelHandle from transformers import GPT2Config config = GPT2Config(n_embd=8, n_head=2, n_layer=1, vocab_size=100, n_positions=64) - # Create a fake model whose state_dict returns a meta tensor + # Create a fake model whose state_dict returns a meta tensor. + # Patch get_state_dict_offloaded_model to return the same meta tensor + # (simulating a failed materialization). fake_model = MagicMock() meta_tensor = torch.empty(4, 8, device="meta") fake_model.state_dict.return_value = {"layer.weight": meta_tensor} @@ -1647,15 +1649,20 @@ class TestNuclearMethod: model=fake_model, tokenizer=MagicMock(), config=config, model_name="test", task="causal_lm", ) - handle._offload_dir = "/nonexistent/path" pipeline = AbliterationPipeline(model_name="test", method="nuclear") pipeline.handle = handle pipeline._on_log = lambda m: None pipeline._on_stage = lambda r: None - with pytest.raises(RuntimeError, match="bricked checkpoint"): - pipeline._gather_state_dict() + # Patch accelerate's function to return meta tensors (simulates + # materialization failure — the safety net should catch this). + with patch( + "accelerate.utils.modeling.get_state_dict_offloaded_model", + return_value={"layer.weight": meta_tensor}, + ): + with pytest.raises(RuntimeError, match="bricked checkpoint"): + pipeline._gather_state_dict() # ---------------------------------------------------------------------------