test: cover offloaded state materialization (#27)

This commit is contained in:
Joseph Magly
2026-08-14 10:26:56 -04:00
parent e1a0fe59f0
commit bc891794db
+33
View File
@@ -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