From bc891794dbfc4709b938ab9392adfc5889be945c Mon Sep 17 00:00:00 2001
From: Joseph Magly <1159087+jmagly@users.noreply.github.com>
Date: Fri, 14 Aug 2026 10:24:25 -0400
Subject: [PATCH] test: cover offloaded state materialization (#27)

---
 tests/test_abliterate.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/tests/test_abliterate.py b/tests/test_abliterate.py
index d6db36e..6080c90 100644
--- a/tests/test_abliterate.py
+++ b/tests/test_abliterate.py
@@ -1664,6 +1664,39 @@ class TestNuclearMethod:
             with pytest.raises(RuntimeError, match="bricked checkpoint"):
                 pipeline._gather_state_dict()
 
+    def test_gather_state_dict_materializes_offloaded_meta(self):
+        """Should return real CPU tensors from Accelerate 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)
+        fake_model = MagicMock()
+        fake_model.state_dict.return_value = {
+            "layer.weight": torch.empty(4, 8, device="meta"),
+        }
+        fake_model.parameters.return_value = []
+        fake_model.buffers.return_value = []
+        materialized = torch.zeros(4, 8)
+
+        handle = ModelHandle(
+            model=fake_model, tokenizer=MagicMock(), config=config,
+            model_name="test", task="causal_lm",
+        )
+        pipeline = AbliterationPipeline(model_name="test", method="nuclear")
+        pipeline.handle = handle
+        pipeline._on_log = lambda m: None
+        pipeline._on_stage = lambda r: None
+
+        with patch(
+            "accelerate.utils.modeling.get_state_dict_offloaded_model",
+            return_value={"layer.weight": materialized},
+        ) as gather_offloaded:
+            state_dict = pipeline._gather_state_dict()
+
+        gather_offloaded.assert_called_once_with(fake_model)
+        assert state_dict["layer.weight"] is materialized
+        assert state_dict["layer.weight"].device.type == "cpu"
+
 
 # ---------------------------------------------------------------------------
 # Knee detection