From c642cb32978b553a11f33d3fce806013c5e44f7a Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Fri, 14 Aug 2026 20:42:47 -0400 Subject: [PATCH] test: harden GPU memory utilization contract --- obliteratus/abliterate.py | 13 ++++++-- obliteratus/cli.py | 24 +++++++++++--- obliteratus/models/loader.py | 29 +++++++++++++---- obliteratus/remote.py | 3 ++ tests/test_abliterate.py | 40 ++++++++++++++++++++++- tests/test_cli_boundaries.py | 27 +++++++++++++++- tests/test_loader_boundaries.py | 57 +++++++++++++++++++++++++++++++++ tests/test_remote_boundaries.py | 8 +++++ 8 files changed, 186 insertions(+), 15 deletions(-) diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index eb10186..a5bfdad 100644 --- a/obliteratus/abliterate.py +++ b/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, ) diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 04ad851..93056c8 100644 --- a/obliteratus/cli.py +++ b/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): diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index d85b21a..bc65d80 100644 --- a/obliteratus/models/loader.py +++ b/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 " diff --git a/obliteratus/remote.py b/obliteratus/remote.py index 55a0d4f..96a9f7e 100644 --- a/obliteratus/remote.py +++ b/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: diff --git a/tests/test_abliterate.py b/tests/test_abliterate.py index f1c2ddc..fcfc70f 100644 --- a/tests/test_abliterate.py +++ b/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" diff --git a/tests/test_cli_boundaries.py b/tests/test_cli_boundaries.py index 811d936..4292658 100644 --- a/tests/test_cli_boundaries.py +++ b/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) diff --git a/tests/test_loader_boundaries.py b/tests/test_loader_boundaries.py index d578c1a..79f27e1 100644 --- a/tests/test_loader_boundaries.py +++ b/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)], diff --git a/tests/test_remote_boundaries.py b/tests/test_remote_boundaries.py index 1026170..84b6115 100644 --- a/tests/test_remote_boundaries.py +++ b/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