Upload 81 files

obliteratus/abliterate.py

@@ -297,13 +297,27 @@ class AbliterationPipeline:
         self._on_stage(result)
         return result
 
+    @staticmethod
+    def _free_gpu_memory():
+        """Release unused GPU memory between pipeline stages."""
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
     def run(self) -> Path:
         """Execute the full abliteration pipeline. Returns path to saved model."""
         self._summon()
+        self._free_gpu_memory()
         self._probe()
+        self._free_gpu_memory()
         self._distill()
+        # Free raw per-prompt activations now that means/subspaces are extracted
+        self._harmful_acts.clear()
+        self._harmless_acts.clear()
+        self._free_gpu_memory()
         self._excise()
+        self._free_gpu_memory()
         self._verify()
+        self._free_gpu_memory()
         return self._rebirth()
 
     # ── Stage 1: SUMMON ─────────────────────────────────────────────────
@@ -325,6 +339,7 @@ class AbliterationPipeline:
             device=self.device,
             dtype=self.dtype,
             trust_remote_code=self.trust_remote_code,
+            quantization=getattr(self, "quantization", None),
         )
 
         summary = self.handle.summary()

obliteratus/models/loader.py

@@ -3,7 +3,10 @@
 from __future__ import annotations
 
 import copy
+import logging
+import tempfile
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import Optional
 
 import torch
@@ -16,6 +19,8 @@ from transformers import (
     PreTrainedTokenizerBase,
 )
 
+logger = logging.getLogger(__name__)
+
 
 TASK_MODEL_MAP = {
     "causal_lm": AutoModelForCausalLM,
@@ -70,6 +75,38 @@ class ModelHandle:
         }
 
 
+def _estimate_model_memory_gb(config: AutoConfig, dtype: torch.dtype) -> float:
+    """Rough estimate of model weight memory in GB."""
+    # Estimate total params from config
+    hidden = getattr(config, "hidden_size", 0)
+    n_layers = getattr(config, "num_hidden_layers", 0)
+    intermediate = getattr(config, "intermediate_size", hidden * 4)
+    vocab = getattr(config, "vocab_size", 0)
+
+    if hidden == 0 or n_layers == 0:
+        return 0.0
+
+    # Per layer: attn (4 * hidden^2) + ffn (3 * hidden * intermediate) + norms
+    per_layer = 4 * hidden * hidden + 3 * hidden * intermediate
+    # Embedding + LM head
+    embedding = 2 * vocab * hidden
+    total_params = per_layer * n_layers + embedding
+
+    bytes_per_param = {torch.float32: 4, torch.float16: 2, torch.bfloat16: 2}.get(dtype, 2)
+    return total_params * bytes_per_param / (1024 ** 3)
+
+
+def _available_gpu_memory_gb() -> float:
+    """Return total available GPU memory across all CUDA devices, in GB."""
+    if not torch.cuda.is_available():
+        return 0.0
+    total = 0.0
+    for i in range(torch.cuda.device_count()):
+        props = torch.cuda.get_device_properties(i)
+        total += props.total_mem / (1024 ** 3)
+    return total
+
+
 def load_model(
     model_name: str,
     task: str = "causal_lm",
@@ -78,6 +115,8 @@ def load_model(
     trust_remote_code: bool = False,
     num_labels: int = 2,
     quantization: str | None = None,
+    offload_folder: str | None = None,
+    skip_snapshot: bool = False,
 ) -> ModelHandle:
     """Load a HuggingFace model and tokenizer, returning a ModelHandle.
 
@@ -89,16 +128,31 @@ def load_model(
         trust_remote_code: Whether to trust remote code from the Hub.
         num_labels: Number of labels for classification tasks.
         quantization: None, "4bit", or "8bit". Requires bitsandbytes.
+        offload_folder: Directory for disk offloading when model exceeds GPU memory.
+            If None and offloading is needed, a temp directory is created automatically.
+        skip_snapshot: If True, skip the initial state dict snapshot to save memory.
     """
     if task not in TASK_MODEL_MAP:
         raise ValueError(f"Unknown task {task!r}. Choose from {list(TASK_MODEL_MAP)}")
 
-    torch_dtype = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[
-        dtype
-    ]
+    dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
+    if dtype not in dtype_map:
+        raise ValueError(f"Unknown dtype {dtype!r}. Choose from {list(dtype_map)}")
+    torch_dtype = dtype_map[dtype]
 
     config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code)
 
+    # Memory estimation and warnings
+    est_gb = _estimate_model_memory_gb(config, torch_dtype)
+    gpu_gb = _available_gpu_memory_gb()
+    if est_gb > 0 and gpu_gb > 0:
+        logger.info(f"Estimated model size: {est_gb:.1f} GB | Available GPU: {gpu_gb:.1f} GB")
+        if est_gb > gpu_gb * 0.9 and quantization is None:
+            logger.warning(
+                f"Model (~{est_gb:.0f} GB) may exceed GPU memory ({gpu_gb:.0f} GB). "
+                f"Consider using quantization='4bit' or quantization='8bit'."
+            )
+
     model_cls = TASK_MODEL_MAP[task]
     load_kwargs: dict = {
         "pretrained_model_name_or_path": model_name,
@@ -126,6 +180,17 @@ def load_model(
     elif device == "auto":
         load_kwargs["device_map"] = "auto"
 
+    # Offload support: provide a folder for disk offloading when GPU memory is insufficient
+    if load_kwargs.get("device_map") == "auto":
+        if offload_folder:
+            load_kwargs["offload_folder"] = offload_folder
+        else:
+            # Auto-create a temp offload dir so from_pretrained never crashes
+            # when Accelerate needs disk offloading
+            _offload_dir = tempfile.mkdtemp(prefix="obliteratus_offload_")
+            load_kwargs["offload_folder"] = _offload_dir
+            logger.info(f"Auto-created offload folder: {_offload_dir}")
+
     model = model_cls.from_pretrained(**load_kwargs)
 
     if device not in ("auto",) and quantization is None:
@@ -133,6 +198,10 @@ def load_model(
 
     model.eval()
 
+    # Free CUDA cache after loading
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -144,5 +213,16 @@ def load_model(
         model_name=model_name,
         task=task,
     )
-    handle.snapshot()
+
+    # Skip snapshot for large models to avoid doubling memory usage
+    if not skip_snapshot:
+        if est_gb > 0 and est_gb > gpu_gb * 0.5:
+            logger.warning(
+                f"Skipping state dict snapshot to save memory "
+                f"(model ~{est_gb:.0f} GB vs GPU {gpu_gb:.0f} GB). "
+                f"Use skip_snapshot=False to force."
+            )
+        else:
+            handle.snapshot()
+
     return handle