From 340173e203757757451893347445ac8bd5fdc816 Mon Sep 17 00:00:00 2001
From: Joseph Magly <1159087+jmagly@users.noreply.github.com>
Date: Fri, 14 Aug 2026 20:20:49 -0400
Subject: [PATCH] test: harden quantized model loading

---
 obliteratus/config.py               |   3 +
 obliteratus/evaluation/evaluator.py |   1 +
 obliteratus/models/loader.py        | 122 +++++++++++++++++-----------
 obliteratus/runner.py               |   2 +-
 tests/test_config.py                |  14 ++++
 tests/test_loader_boundaries.py     | 102 +++++++++++++++++++++++
 tests/test_runner_boundaries.py     |  26 ++++++
 7 files changed, 223 insertions(+), 47 deletions(-)
 create mode 100644 tests/test_runner_boundaries.py

diff --git a/obliteratus/config.py b/obliteratus/config.py
index e12d837..96214e1 100644
--- a/obliteratus/config.py
+++ b/obliteratus/config.py
@@ -17,6 +17,7 @@ class ModelConfig:
     device: str = "auto"
     trust_remote_code: bool = False
     num_labels: int = 2
+    quantization: str | None = None
 
 
 @dataclass
@@ -120,6 +121,8 @@ class StudyConfig:
                 "dtype": self.model.dtype,
                 "device": self.model.device,
                 "trust_remote_code": self.model.trust_remote_code,
+                "num_labels": self.model.num_labels,
+                "quantization": self.model.quantization,
             },
             "dataset": {
                 "name": self.dataset.name,
diff --git a/obliteratus/evaluation/evaluator.py b/obliteratus/evaluation/evaluator.py
index e79d7c6..92801da 100644
--- a/obliteratus/evaluation/evaluator.py
+++ b/obliteratus/evaluation/evaluator.py
@@ -87,6 +87,7 @@ class Evaluator:
             desc="Evaluating PPL",
         ):
             batch_texts = ds[i : i + self.batch_size][self.text_column]
+
             # Defensive filtering in case a custom dataset returns unexpected
             # values after selection or transformation.
             batch_texts = [
diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py
index 38a3e0c..504c3bc 100644
--- a/obliteratus/models/loader.py
+++ b/obliteratus/models/loader.py
@@ -28,6 +28,16 @@ except (ImportError, AttributeError):
 
 logger = logging.getLogger(__name__)
 
+_QUANTIZATION_STATE_COMPONENTS = frozenset(
+    {
+        "absmax",
+        "nested_absmax",
+        "nested_quant_map",
+        "quant_map",
+        "quant_state",
+    }
+)
+
 
 # ---------------------------------------------------------------------------
 # Compat shims for transformers ≥5.0 breaking changes.
@@ -366,7 +376,24 @@ class ModelHandle:
         for k, v in self._original_state.items():
             target = current_state[k].device if k in current_state else None
             restored[k] = v.to(target) if target is not None else v
-        self.model.load_state_dict(restored, strict=False)
+        incompatible = self.model.load_state_dict(restored, strict=False)
+        missing = [
+            key
+            for key in incompatible.missing_keys
+            if not _is_quantization_state_key(key)
+        ]
+        unexpected = [
+            key
+            for key in incompatible.unexpected_keys
+            if not _is_quantization_state_key(key)
+        ]
+        if missing or unexpected:
+            details = []
+            if missing:
+                details.append(f"missing keys: {missing}")
+            if unexpected:
+                details.append(f"unexpected keys: {unexpected}")
+            raise RuntimeError(f"Snapshot restore was incomplete ({'; '.join(details)})")
 
     def cleanup(self):
         """Remove temporary offload directory if one was auto-created."""
@@ -429,6 +456,35 @@ def _estimate_model_memory_gb(config: AutoConfig, dtype: torch.dtype) -> float:
     return total_params * bytes_per_param / (1024 ** 3)
 
 
+def _is_quantization_state_key(key: str) -> bool:
+    """Return whether a state-dict key is bitsandbytes quantization metadata."""
+    return any(
+        component in _QUANTIZATION_STATE_COMPONENTS
+        or component.startswith("bitsandbytes__")
+        for component in key.split(".")
+    )
+
+
+def _effective_model_memory_gb(estimate_gb: float, quantization: str | None) -> float:
+    """Adjust a full-precision weight estimate for runtime quantization."""
+    factor = {"4bit": 4, "8bit": 2}.get(quantization, 1)
+    return estimate_gb / factor
+
+
+def _bounded_max_memory() -> dict[int | str, str]:
+    """Build Accelerate memory limits with inference and host headroom."""
+    max_memory: dict[int | str, str] = {}
+    for index in range(dev.device_count()):
+        total = torch.cuda.get_device_properties(index).total_memory
+        reserve = max(int(total * 0.15), 2 * 1024 ** 3)
+        usable = total - reserve
+        max_memory[index] = f"{usable // (1024 ** 2)}MiB"
+    total_ram, _ = dev._system_memory_gb()
+    cpu_budget_gb = int(total_ram * 0.85)
+    max_memory["cpu"] = f"{max(cpu_budget_gb, 4)}GiB"
+    return max_memory
+
+
 def _available_gpu_memory_gb() -> float:
     """Return free accelerator memory in GB (CUDA, MPS, or 0 for CPU)."""
     return dev.get_total_free_gb()
@@ -622,46 +678,24 @@ def load_model(
             load_kwargs["offload_folder"] = _offload_dir
             logger.info(f"Auto-created offload folder: {_offload_dir}")
 
-        # Skip max_memory when quantization shrinks the model enough to fit
-        if quantization in ("4bit", "8bit") and est_gb > 0 and gpu_gb > 0:
-            quant_factor = 4 if quantization == "4bit" else 2
-            quant_est_gb = est_gb / quant_factor
-            if quant_est_gb < gpu_gb * 0.7:
-                logger.info(
-                    f"Quantized estimate ({quant_est_gb:.1f} GB) fits GPU "
-                    f"({gpu_gb:.0f} GB) — skipping max_memory constraint"
-                )
-            else:
-                if dev.is_cuda():
-                    max_memory = {}
-                    for i in range(dev.device_count()):
-                        total = torch.cuda.get_device_properties(i).total_memory
-                        reserve = max(int(total * 0.15), 2 * 1024 ** 3)
-                        usable = total - reserve
-                        max_memory[i] = f"{usable // (1024 ** 2)}MiB"
-                    total_ram, _ = dev._system_memory_gb()
-                    cpu_budget_gb = int(total_ram * 0.85)
-                    max_memory["cpu"] = f"{max(cpu_budget_gb, 4)}GiB"
-                    load_kwargs["max_memory"] = max_memory
-                    logger.info(
-                        f"GPU memory budget: {', '.join(f'GPU{k}={v}' for k, v in max_memory.items() if k != 'cpu')}"
-                    )
-        else:
-            # No quantization — use original max_memory logic
-            if dev.is_cuda():
-                max_memory = {}
-                for i in range(dev.device_count()):
-                    total = torch.cuda.get_device_properties(i).total_memory
-                    reserve = max(int(total * 0.15), 2 * 1024 ** 3)
-                    usable = total - reserve
-                    max_memory[i] = f"{usable // (1024 ** 2)}MiB"
-                total_ram, _ = dev._system_memory_gb()
-                cpu_budget_gb = int(total_ram * 0.85)
-                max_memory["cpu"] = f"{max(cpu_budget_gb, 4)}GiB"
-                load_kwargs["max_memory"] = max_memory
-                logger.info(
-                    f"GPU memory budget: {', '.join(f'GPU{k}={v}' for k, v in max_memory.items() if k != 'cpu')}"
-                )
+        effective_est_gb = _effective_model_memory_gb(est_gb, quantization)
+        quantized_fit = (
+            quantization in ("4bit", "8bit")
+            and effective_est_gb > 0
+            and gpu_gb > 0
+            and effective_est_gb < gpu_gb * 0.7
+        )
+        if quantized_fit:
+            logger.info(
+                f"Quantized estimate ({effective_est_gb:.1f} GB) fits GPU "
+                f"({gpu_gb:.0f} GB) — skipping max_memory constraint"
+            )
+        elif dev.is_cuda():
+            max_memory = _bounded_max_memory()
+            load_kwargs["max_memory"] = max_memory
+            logger.info(
+                f"GPU memory budget: {', '.join(f'GPU{k}={v}' for k, v in max_memory.items() if k != 'cpu')}"
+            )
 
     try:
         model = model_cls.from_pretrained(**load_kwargs)
@@ -752,11 +786,7 @@ def load_model(
             else:
                 handle.snapshot()
         elif gpu_gb > 0 and est_gb > 0:
-            effective_gb = est_gb
-            if quantization == "4bit":
-                effective_gb = est_gb / 4
-            elif quantization == "8bit":
-                effective_gb = est_gb / 2
+            effective_gb = _effective_model_memory_gb(est_gb, quantization)
             if effective_gb > gpu_gb * 0.5:
                 logger.warning(
                     f"Auto-skipping state dict snapshot to save memory "
diff --git a/obliteratus/runner.py b/obliteratus/runner.py
index d3bc069..4e1808d 100644
--- a/obliteratus/runner.py
+++ b/obliteratus/runner.py
@@ -38,7 +38,7 @@ def run_study(config: StudyConfig) -> AblationReport:
         dtype=config.model.dtype,
         trust_remote_code=config.model.trust_remote_code,
         num_labels=config.model.num_labels,
-        quantization=getattr(config.model, "quantization", None),
+        quantization=config.model.quantization,
     )
     console.print(f"  Architecture: {handle.architecture}")
     console.print(f"  Layers: {handle.num_layers}  Heads: {handle.num_heads}")
diff --git a/tests/test_config.py b/tests/test_config.py
index debaad5..ef982f6 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -57,3 +57,17 @@ class TestStudyConfig:
         assert config2.model.name == config.model.name
         assert config2.dataset.name == config.dataset.name
         assert len(config2.strategies) == len(config.strategies)
+
+    def test_model_quantization_and_label_count_roundtrip(self):
+        raw = {
+            **SAMPLE_CONFIG,
+            "model": {
+                **SAMPLE_CONFIG["model"],
+                "num_labels": 7,
+                "quantization": "4bit",
+            },
+        }
+        config = StudyConfig.from_dict(raw)
+        assert config.model.quantization == "4bit"
+        assert config.model.num_labels == 7
+        assert StudyConfig.from_dict(config.to_dict()).model == config.model
diff --git a/tests/test_loader_boundaries.py b/tests/test_loader_boundaries.py
index c997418..752d9ec 100644
--- a/tests/test_loader_boundaries.py
+++ b/tests/test_loader_boundaries.py
@@ -200,8 +200,12 @@ def test_model_handle_metadata_snapshot_restore_summary_and_cleanup(tmp_path):
     with pytest.raises(RuntimeError, match="call .snapshot"):
         handle.restore()
     handle.snapshot()
+    model.load_state_dict.return_value = SimpleNamespace(missing_keys=[], unexpected_keys=[])
     handle.restore()
     model.load_state_dict.assert_called_once()
+    restored = model.load_state_dict.call_args.args[0]
+    assert torch.equal(restored["weight"], torch.ones(2))
+    assert model.load_state_dict.call_args.kwargs == {"strict": False}
     assert handle.summary() == {
         "model_name": "x",
         "architecture": "gpt2",
@@ -217,6 +221,50 @@ def test_model_handle_metadata_snapshot_restore_summary_and_cleanup(tmp_path):
     assert handle._offload_dir is None
 
 
+def test_model_handle_restore_tolerates_quantization_metadata():
+    model = _model()
+    handle = loader.ModelHandle(model, SimpleNamespace(), _config(), "x", "causal_lm")
+    handle._original_state = {"weight": torch.ones(2)}
+    model.load_state_dict.return_value = SimpleNamespace(
+        missing_keys=["layer.weight.absmax"],
+        unexpected_keys=["layer.weight.bitsandbytes__nf4"],
+    )
+    handle.restore()
+
+
+@pytest.mark.parametrize(
+    ("missing", "unexpected", "message"),
+    [
+        (["layer.bias"], [], "missing keys: \\['layer.bias'\\]"),
+        ([], ["layer.other_weight"], "unexpected keys: \\['layer.other_weight'\\]"),
+        (
+            ["layer.bias"],
+            ["layer.other_weight"],
+            r"missing keys: \['layer.bias'\].*unexpected keys: \['layer.other_weight'\]",
+        ),
+    ],
+)
+def test_model_handle_restore_rejects_parameter_mismatches(missing, unexpected, message):
+    model = _model()
+    handle = loader.ModelHandle(model, SimpleNamespace(), _config(), "x", "causal_lm")
+    handle._original_state = {"weight": torch.ones(2)}
+    model.load_state_dict.return_value = SimpleNamespace(
+        missing_keys=missing,
+        unexpected_keys=unexpected,
+    )
+    with pytest.raises(RuntimeError, match=message):
+        handle.restore()
+
+
+def test_model_handle_restore_preserves_loader_shape_failures():
+    model = _model()
+    handle = loader.ModelHandle(model, SimpleNamespace(), _config(), "x", "causal_lm")
+    handle._original_state = {"weight": torch.ones(2)}
+    model.load_state_dict.side_effect = RuntimeError("size mismatch for weight")
+    with pytest.raises(RuntimeError, match="size mismatch for weight"):
+        handle.restore()
+
+
 def test_model_memory_estimation_handles_dense_moe_nested_and_unknown():
     dense = loader._estimate_model_memory_gb(_config(), torch.float32)
     moe = loader._estimate_model_memory_gb(_config(num_local_experts=4), torch.float32)
@@ -300,6 +348,60 @@ def test_cuda_auto_map_has_bounded_memory_and_requested_offload(loader_boundary,
     assert tmp_path.exists(), "cleanup must not remove an operator-owned directory"
 
 
+def _enable_cuda_quantization(monkeypatch, estimate_gb: float) -> None:
+    gib = 1024**3
+    monkeypatch.setattr(loader, "_estimate_model_memory_gb", lambda *_args: estimate_gb)
+    monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "cuda")
+    monkeypatch.setattr(loader.dev, "supports_bitsandbytes", lambda _device=None: True)
+    monkeypatch.setattr(loader.dev, "supports_device_map_auto", lambda _device=None: True)
+    monkeypatch.setattr(loader.dev, "get_total_free_gb", lambda: 16.0)
+    monkeypatch.setattr(loader.dev, "is_cuda", lambda: True)
+    monkeypatch.setattr(loader.dev, "device_count", lambda: 1)
+    monkeypatch.setattr(loader.dev, "_system_memory_gb", lambda: (64.0, 40.0))
+    monkeypatch.setattr(
+        loader.torch.cuda,
+        "get_device_properties",
+        lambda _index: SimpleNamespace(total_memory=16 * gib),
+    )
+    monkeypatch.setattr("transformers.BitsAndBytesConfig", lambda **kwargs: kwargs)
+
+
+@pytest.mark.parametrize(
+    ("quantization", "estimate_gb", "expects_budget"),
+    [
+        ("4bit", 40.0, False),
+        ("4bit", 48.0, True),
+        ("8bit", 20.0, False),
+        ("8bit", 24.0, True),
+    ],
+)
+def test_quantized_single_gpu_memory_budget_uses_effective_weight_size(
+    loader_boundary, monkeypatch, quantization, estimate_gb, expects_budget,
+):
+    _enable_cuda_quantization(monkeypatch, estimate_gb)
+    handle = loader.load_model("x", quantization=quantization, skip_snapshot=True)
+    kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
+    assert ("max_memory" in kwargs) is expects_budget
+    if expects_budget:
+        assert kwargs["max_memory"] == {0: "13926MiB", "cpu": "54GiB"}
+    handle.cleanup()
+
+
+@pytest.mark.parametrize(
+    ("quantization", "estimate_gb", "snapshots"),
+    [("4bit", 16.0, 1), ("4bit", 40.0, 0), ("8bit", 8.0, 1), ("8bit", 20.0, 0)],
+)
+def test_quantized_snapshot_policy_uses_effective_weight_size(
+    loader_boundary, monkeypatch, quantization, estimate_gb, snapshots,
+):
+    _enable_cuda_quantization(monkeypatch, estimate_gb)
+    snapshot = Mock()
+    monkeypatch.setattr(loader.ModelHandle, "snapshot", snapshot)
+    handle = loader.load_model("x", quantization=quantization)
+    assert snapshot.call_count == snapshots
+    handle.cleanup()
+
+
 def test_model_permission_error_retries_and_explicit_device_moves(loader_boundary, monkeypatch, tmp_path):
     monkeypatch.setattr(loader.tempfile, "gettempdir", lambda: str(tmp_path))
     loader_boundary.model_class.from_pretrained.side_effect = [PermissionError("cache"), loader_boundary.model]
diff --git a/tests/test_runner_boundaries.py b/tests/test_runner_boundaries.py
new file mode 100644
index 0000000..c858d9e
--- /dev/null
+++ b/tests/test_runner_boundaries.py
@@ -0,0 +1,26 @@
+"""Focused orchestration contracts for study execution."""
+
+from __future__ import annotations
+
+from unittest.mock import Mock
+
+import pytest
+
+from obliteratus import runner
+from obliteratus.config import DatasetConfig, ModelConfig, StrategyConfig, StudyConfig
+
+
+def test_run_study_passes_configured_quantization_to_loader(tmp_path, monkeypatch):
+    config = StudyConfig(
+        model=ModelConfig(name="fixture", quantization="4bit"),
+        dataset=DatasetConfig(name="fixture"),
+        strategies=[StrategyConfig(name="layer_removal")],
+        output_dir=str(tmp_path),
+    )
+    load_model = Mock(side_effect=RuntimeError("stop after loader boundary"))
+    monkeypatch.setattr(runner, "load_model", load_model)
+
+    with pytest.raises(RuntimeError, match="stop after loader boundary"):
+        runner.run_study(config)
+
+    assert load_model.call_args.kwargs["quantization"] == "4bit"