diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 060a4d3..eb10186 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -737,6 +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, harmful_prompts: list[str] | None = None, harmless_prompts: list[str] | None = None, jailbreak_prompts: list[str] | None = None, @@ -841,6 +842,7 @@ 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 # 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) @@ -1156,6 +1158,8 @@ 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, ) summary = self.handle.summary() diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 2c5ad85..04ad851 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -253,6 +253,11 @@ def main(argv: list[str] | None = None): "--quantization", type=str, default=None, choices=["4bit", "8bit"], 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.", + ) p.add_argument( "--large-model", action="store_true", default=False, help="Enable conservative defaults for 120B+ models (fewer directions, 1 pass, lower SAE expansion).", @@ -1070,6 +1075,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), 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), diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index 48ccbd5..d85b21a 100644 --- a/obliteratus/models/loader.py +++ b/obliteratus/models/loader.py @@ -482,13 +482,12 @@ def _effective_model_memory_gb(estimate_gb: float, quantization: str | None) -> return effective_model_memory_gb(estimate_gb, quantization) -def _bounded_max_memory() -> dict[int | str, str]: +def _bounded_max_memory(gpu_memory_utilization: float = 0.85) -> 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 + 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) @@ -516,6 +515,7 @@ def load_model( quantization: str | None = None, offload_folder: str | None = None, skip_snapshot: bool | None = None, + gpu_memory_utilization: float = 0.85, revision: str | None = None, local_files_only: bool = False, ) -> ModelHandle: @@ -535,6 +535,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. revision: Optional Hub branch, tag, or commit passed to every loader. local_files_only: Refuse network access and use only locally cached files. """ @@ -694,7 +696,7 @@ def load_model( f"({gpu_gb:.0f} GB) — skipping max_memory constraint" ) elif dev.is_cuda(): - max_memory = _bounded_max_memory() + max_memory = _bounded_max_memory(gpu_memory_utilization) 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')}"