test: harden GPU memory utilization contract

obliteratus/abliterate.py

@@ -737,7 +737,7 @@ class AbliterationPipeline:
         use_whitened_svd: bool | None = None,
         true_iterative_refinement: bool | None = None,
         quantization: str | None = None,
-        gpu_memory_utilization: float = 0.85,
+        gpu_memory_utilization: float | None = None,
         harmful_prompts: list[str] | None = None,
         harmless_prompts: list[str] | None = None,
         jailbreak_prompts: list[str] | None = None,
@@ -842,7 +842,15 @@ class AbliterationPipeline:
         self.use_whitened_svd = use_whitened_svd if use_whitened_svd is not None else method_cfg.get("use_whitened_svd", False)
         self.true_iterative_refinement = true_iterative_refinement if true_iterative_refinement is not None else method_cfg.get("true_iterative_refinement", False)
         self.quantization = quantization
-        self.gpu_memory_utilization = gpu_memory_utilization
+        if gpu_memory_utilization is not None and (
+            isinstance(gpu_memory_utilization, bool)
+            or not isinstance(gpu_memory_utilization, (int, float))
+            or not 0.0 < float(gpu_memory_utilization) <= 1.0
+        ):
+            raise ValueError("gpu_memory_utilization must be a number in (0, 1]")
+        self.gpu_memory_utilization = (
+            float(gpu_memory_utilization) if gpu_memory_utilization is not None else None
+        )
 
         # SOTA techniques (resolve from method or explicit override)
         self.use_jailbreak_contrast = use_jailbreak_contrast if use_jailbreak_contrast is not None else method_cfg.get("use_jailbreak_contrast", False)
@@ -1158,7 +1166,6 @@ class AbliterationPipeline:
             dtype=self.dtype,
             trust_remote_code=self.trust_remote_code,
             quantization=self.quantization,
-            skip_snapshot=True,
             gpu_memory_utilization=self.gpu_memory_utilization,
         )
 

obliteratus/cli.py

@@ -46,6 +46,17 @@ def _ssh_port(value: str) -> int:
     return parsed
 
 
+def _utilization_fraction(value: str) -> float:
+    """Parse a finite fraction in the public ``(0, 1]`` contract."""
+    try:
+        parsed = float(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("must be a number in (0, 1]") from exc
+    if not 0.0 < parsed <= 1.0:
+        raise argparse.ArgumentTypeError("must be a number in (0, 1]")
+    return parsed
+
+
 def _add_gpu_args(parser):
     """Add --gpus flag for multi-GPU control."""
     gpu_group = parser.add_argument_group("GPU selection")
@@ -254,9 +265,12 @@ def main(argv: list[str] | None = None):
             help="Load model with quantization (4bit or 8bit). Requires bitsandbytes.",
         )
         p.add_argument(
-            "--gpu-memory-utilization", type=float, default=0.85,
-            help="Fraction of GPU VRAM to make available for model loading (0.0-1.0, default 0.85). "
-                 "Increase toward 1.0 if the GPU is dedicated to this process.",
+            "--gpu-memory-utilization",
+            type=_utilization_fraction,
+            default=None,
+            metavar="FRACTION",
+            help="Override the GPU VRAM fraction available for model loading (0, 1]. "
+            "By default, reserves 15%% or 2 GiB per GPU, whichever is larger.",
         )
         p.add_argument(
             "--large-model", action="store_true", default=False,
@@ -1075,7 +1089,7 @@ def _cmd_abliterate(args):
         projection_target=getattr(args, "projection_target", None),
         projection_row_fraction=getattr(args, "projection_row_fraction", None),
         quantization=args.quantization,
-        gpu_memory_utilization=getattr(args, "gpu_memory_utilization", 0.85),
+        gpu_memory_utilization=getattr(args, "gpu_memory_utilization", None),
         large_model_mode=getattr(args, "large_model", False),
         verify_sample_size=getattr(args, "verify_sample_size", None),
         refusal_max_tokens=getattr(args, "refusal_max_tokens", None),
@@ -1374,6 +1388,8 @@ def _cmd_remote_abliterate(args):
         kwargs["dtype"] = args.dtype
     if args.quantization:
         kwargs["quantization"] = args.quantization
+    if getattr(args, "gpu_memory_utilization", None) is not None:
+        kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization
     if args.n_directions is not None:
         kwargs["n_directions"] = args.n_directions
     if getattr(args, "direction_method", None):

obliteratus/models/loader.py

@@ -482,12 +482,18 @@ def _effective_model_memory_gb(estimate_gb: float, quantization: str | None) ->
     return effective_model_memory_gb(estimate_gb, quantization)
 
 
-def _bounded_max_memory(gpu_memory_utilization: float = 0.85) -> dict[int | str, str]:
+def _bounded_max_memory(
+    gpu_memory_utilization: float | None = None,
+) -> 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
-        usable = int(total * gpu_memory_utilization)
+        if gpu_memory_utilization is None:
+            reserve = max(int(total * 0.15), 2 * 1024 ** 3)
+            usable = total - reserve
+        else:
+            usable = int(total * gpu_memory_utilization)
         max_memory[index] = f"{usable // (1024 ** 2)}MiB"
     total_ram, _ = dev._system_memory_gb()
     cpu_budget_gb = int(total_ram * 0.85)
@@ -515,7 +521,7 @@ def load_model(
     quantization: str | None = None,
     offload_folder: str | None = None,
     skip_snapshot: bool | None = None,
-    gpu_memory_utilization: float = 0.85,
+    gpu_memory_utilization: float | None = None,
     revision: str | None = None,
     local_files_only: bool = False,
 ) -> ModelHandle:
@@ -535,8 +541,8 @@ def load_model(
             None (default): auto-decide based on GPU memory headroom.
             True: always skip (saves memory).
             False: always snapshot (force even for large models).
-        gpu_memory_utilization: Fraction of GPU VRAM to use (0.0-1.0, default 0.85).
-            Increase toward 1.0 if the GPU is dedicated to this process.
+        gpu_memory_utilization: Optional GPU VRAM fraction in ``(0, 1]``. When
+            omitted, reserves 15% or 2 GiB per GPU, whichever is larger.
         revision: Optional Hub branch, tag, or commit passed to every loader.
         local_files_only: Refuse network access and use only locally cached files.
     """
@@ -549,6 +555,14 @@ def load_model(
         dtype,
         valid_tasks=TASK_MODEL_MAP,
     )
+    if gpu_memory_utilization is not None and (
+        isinstance(gpu_memory_utilization, bool)
+        or not isinstance(gpu_memory_utilization, (int, float))
+        or not 0.0 < float(gpu_memory_utilization) <= 1.0
+    ):
+        raise ValueError("gpu_memory_utilization must be a number in (0, 1]")
+    if gpu_memory_utilization is not None:
+        gpu_memory_utilization = float(gpu_memory_utilization)
 
     dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
     torch_dtype = dtype_map[dtype]
@@ -689,7 +703,10 @@ def load_model(
             logger.info(f"Auto-created offload folder: {_offload_dir}")
 
         effective_est_gb = _effective_model_memory_gb(est_gb, quantization)
-        quantized_fit = quantized_model_fits_gpu(est_gb, quantization, gpu_gb)
+        quantized_fit = (
+            gpu_memory_utilization is None
+            and quantized_model_fits_gpu(est_gb, quantization, gpu_gb)
+        )
         if quantized_fit:
             logger.info(
                 f"Quantized estimate ({effective_est_gb:.1f} GB) fits GPU "

obliteratus/remote.py

@@ -305,6 +305,7 @@ class RemoteRunner:
         device: str = "auto",
         dtype: str = "float16",
         quantization: str | None = None,
+        gpu_memory_utilization: float | None = None,
         n_directions: int | None = None,
         direction_method: str | None = None,
         regularization: float | None = None,
@@ -335,6 +336,8 @@ class RemoteRunner:
         ]
         if quantization:
             parts.extend(["--quantization", quantization])
+        if gpu_memory_utilization is not None:
+            parts.extend(["--gpu-memory-utilization", str(gpu_memory_utilization)])
         if n_directions is not None:
             parts.extend(["--n-directions", str(n_directions)])
         if direction_method:

tests/test_abliterate.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import json
 from pathlib import Path
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 import torch
@@ -180,6 +180,7 @@ class TestPipelineInit:
         assert pipeline.dtype == "float16"
         assert pipeline.output_dir == Path("abliterated")
         assert pipeline.trust_remote_code is False
+        assert pipeline.gpu_memory_utilization is None
         assert pipeline.refusal_max_tokens == 128
         assert pipeline.handle is None
 
@@ -196,6 +197,43 @@ class TestPipelineInit:
         )
         assert pipeline.refusal_max_tokens == 512
 
+    @pytest.mark.parametrize(
+        "invalid", [0, -0.1, 1.1, float("nan"), float("inf"), True, "0.8"],
+    )
+    def test_gpu_memory_utilization_requires_bounded_number(self, invalid):
+        with pytest.raises(
+            ValueError, match=r"gpu_memory_utilization must be a number in \(0, 1\]",
+        ):
+            AbliterationPipeline(
+                model_name="test-model", gpu_memory_utilization=invalid,
+            )
+
+    def test_gpu_memory_utilization_accepts_explicit_override(self):
+        pipeline = AbliterationPipeline(
+            model_name="test-model", gpu_memory_utilization=0.95,
+        )
+        assert pipeline.gpu_memory_utilization == 0.95
+
+    def test_summon_propagates_gpu_budget_without_forcing_snapshot_skip(self, monkeypatch):
+        handle = MagicMock()
+        handle.summary.return_value = {
+            "architecture": "fixture",
+            "num_layers": 2,
+            "num_heads": 4,
+            "hidden_size": 8,
+            "total_params": 32,
+        }
+        load_model = Mock(return_value=handle)
+        monkeypatch.setattr("obliteratus.abliterate.load_model", load_model)
+        pipeline = AbliterationPipeline(
+            model_name="test-model", gpu_memory_utilization=0.95,
+        )
+
+        pipeline._summon()
+
+        assert load_model.call_args.kwargs["gpu_memory_utilization"] == 0.95
+        assert "skip_snapshot" not in load_model.call_args.kwargs
+
     def test_default_method_is_advanced(self):
         pipeline = AbliterationPipeline(model_name="test-model")
         assert pipeline.method == "advanced"

tests/test_cli_boundaries.py

@@ -77,6 +77,25 @@ def test_refusal_max_tokens_cli_rejects_invalid_values(invalid):
     assert exc.value.code == 2
 
 
+def test_gpu_memory_utilization_cli_default_and_explicit_override(monkeypatch):
+    command = Mock()
+    monkeypatch.setattr(cli, "_cmd_abliterate", command)
+
+    cli.main(["obliterate", "local/model"])
+    assert command.call_args.args[0].gpu_memory_utilization is None
+
+    command.reset_mock()
+    cli.main(["obliterate", "local/model", "--gpu-memory-utilization", "0.95"])
+    assert command.call_args.args[0].gpu_memory_utilization == 0.95
+
+
+@pytest.mark.parametrize("invalid", ["0", "-0.1", "1.1", "nan", "inf", "not-a-number"])
+def test_gpu_memory_utilization_cli_rejects_invalid_values(invalid):
+    with pytest.raises(SystemExit) as exc:
+        cli.main(["obliterate", "local/model", "--gpu-memory-utilization", invalid])
+    assert exc.value.code == 2
+
+
 def test_version_is_stable_and_does_not_dispatch(capsys):
     from obliteratus import __version__
 
@@ -332,6 +351,7 @@ def _remote_args(**overrides):
         "shield_ridge": 0.1, "shield_residualize": True, "shield_layer_penalty": 0.2,
         "projection_target": "all", "projection_row_fraction": 0.5, "large_model": True,
         "verify_sample_size": 5, "refusal_max_tokens": 512,
+        "gpu_memory_utilization": 0.95,
         "config": "config.yml", "preset": "quick", "methods": ["advanced"],
         "hub_org": "org", "hub_repo": None, "dataset": "builtin",
     }
@@ -363,11 +383,14 @@ def test_remote_runner_factory_and_commands(monkeypatch):
     cli._cmd_remote_abliterate(args)
     assert runner.run_obliterate.call_args.kwargs["projection_row_fraction"] == 0.5
     assert runner.run_obliterate.call_args.kwargs["refusal_max_tokens"] == 512
+    assert runner.run_obliterate.call_args.kwargs["gpu_memory_utilization"] == 0.95
 
     args.refusal_max_tokens = None
+    args.gpu_memory_utilization = None
     runner.run_obliterate.reset_mock()
     cli._cmd_remote_abliterate(args)
     assert "refusal_max_tokens" not in runner.run_obliterate.call_args.kwargs
+    assert "gpu_memory_utilization" not in runner.run_obliterate.call_args.kwargs
 
     for name in (
         "quantization", "n_directions", "direction_method", "regularization",
@@ -469,12 +492,14 @@ def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp
         min_layer_fraction=0.1, max_layer_fraction=0.9, harmless_pc_count=1,
         shield_concept_count=2, shield_ridge=0.1, shield_residualize=True,
         shield_layer_penalty=0.2, projection_target="all", projection_row_fraction=0.5,
-        quantization=None, large_model=False, verify_sample_size=3, refusal_max_tokens=512,
+        quantization=None, gpu_memory_utilization=0.95, large_model=False,
+        verify_sample_size=3, refusal_max_tokens=512,
         residue_file=["audit.json"], dataset="builtin", residue_weight=2, residue_max=4,
         contribute=True, contribute_notes="fixture",
     )
     cli._cmd_abliterate(args)
     assert factory.call_args.kwargs["refusal_max_tokens"] == 512
+    assert factory.call_args.kwargs["gpu_memory_utilization"] == 0.95
     assert (result_path / "hard_negative_residue.json").is_file()
     telemetry.assert_called_once_with(pipeline)
 

tests/test_loader_boundaries.py

@@ -62,6 +62,13 @@ def loader_boundary(monkeypatch):
         ({"model_name": "x", "task": "embedding"}, "Unknown task"),
         ({"model_name": "x", "dtype": "int9"}, "Unknown dtype"),
         ({"model_name": "x", "quantization": "3bit"}, "Unknown quantization"),
+        ({"model_name": "x", "gpu_memory_utilization": 0}, "gpu_memory_utilization"),
+        ({"model_name": "x", "gpu_memory_utilization": -0.1}, "gpu_memory_utilization"),
+        ({"model_name": "x", "gpu_memory_utilization": 1.1}, "gpu_memory_utilization"),
+        ({"model_name": "x", "gpu_memory_utilization": float("nan")}, "gpu_memory_utilization"),
+        ({"model_name": "x", "gpu_memory_utilization": float("inf")}, "gpu_memory_utilization"),
+        ({"model_name": "x", "gpu_memory_utilization": True}, "gpu_memory_utilization"),
+        ({"model_name": "x", "gpu_memory_utilization": "0.8"}, "gpu_memory_utilization"),
     ],
 )
 def test_invalid_requests_fail_before_provider_access(loader_boundary, monkeypatch, kwargs, message):
@@ -348,6 +355,43 @@ 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 test_explicit_gpu_memory_utilization_overrides_default_headroom(
+    loader_boundary, monkeypatch,
+):
+    gib = 1024**3
+    monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "cuda")
+    monkeypatch.setattr(loader.dev, "supports_device_map_auto", lambda _device=None: True)
+    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=20 * gib),
+    )
+    handle = loader.load_model(
+        "x", gpu_memory_utilization=0.95, skip_snapshot=True,
+    )
+    kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
+    assert kwargs["max_memory"] == {0: "19456MiB", "cpu": "54GiB"}
+    handle.cleanup()
+
+
+@pytest.mark.parametrize("device", ["cpu", "mps"])
+def test_gpu_memory_utilization_is_safe_when_accelerate_auto_map_is_unavailable(
+    loader_boundary, monkeypatch, device,
+):
+    monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": device)
+    handle = loader.load_model(
+        "x", device=device, gpu_memory_utilization=0.95, skip_snapshot=True,
+    )
+    kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
+    assert "device_map" not in kwargs
+    assert "max_memory" not in kwargs
+    loader_boundary.model.to.assert_called_once_with(device)
+    handle.cleanup()
+
+
 def _enable_cuda_quantization(monkeypatch, estimate_gb: float) -> None:
     gib = 1024**3
     monkeypatch.setattr(loader, "_estimate_model_memory_gb", lambda *_args: estimate_gb)
@@ -387,6 +431,19 @@ def test_quantized_single_gpu_memory_budget_uses_effective_weight_size(
     handle.cleanup()
 
 
+def test_explicit_gpu_budget_overrides_quantized_fit_shortcut(loader_boundary, monkeypatch):
+    _enable_cuda_quantization(monkeypatch, estimate_gb=40.0)
+    handle = loader.load_model(
+        "x",
+        quantization="4bit",
+        gpu_memory_utilization=0.5,
+        skip_snapshot=True,
+    )
+    kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
+    assert kwargs["max_memory"] == {0: "8192MiB", "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)],

tests/test_remote_boundaries.py

@@ -415,3 +415,11 @@ def test_remote_tourney_stops_at_connection_or_install_failure():
     runner.check_connection.return_value = True
     runner.ensure_obliteratus.return_value = False
     assert runner.run_tourney("model") is None
+
+
+def test_remote_obliterate_command_propagates_gpu_memory_utilization():
+    runner = RemoteRunner(RemoteConfig(host="compute.example"))
+    command = runner.build_obliterate_command(
+        "org/model", gpu_memory_utilization=0.95,
+    )
+    assert "--gpu-memory-utilization 0.95" in command