fix: REBIRTH save crash for CPU-offloaded models

Replace the manual safetensors/dat file materialization in
_gather_state_dict with accelerate's get_state_dict_offloaded_model().
The old code only handled disk-offloaded weights but failed for the 398
CPU-offloaded meta tensors managed by accelerate's AlignDevicesHook.

Pre-move all GPU tensors to CPU before materialization to prevent CUDA
OOM when align_module_device restores non-hooked params to CUDA during
its __exit__ cleanup.

Update test to verify the safety net catches unmaterialized meta tensors
after the accelerate path (the old test checked for a missing offload
directory, which this codepath no longer uses).

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."
             )
 

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()
 
 
 # ---------------------------------------------------------------------------