Upload 133 files

app.py

@@ -2046,7 +2046,22 @@ def _format_multi_model_results(results: list[dict], context: dict | None = None
     return "\n".join(lines)
 
 
+# ---------------------------------------------------------------------------
+# Staged GPU wrapper for obliteration (tourney-style per-stage allocation)
+# ---------------------------------------------------------------------------
+
 @spaces.GPU(duration=300)
+def _obliterate_gpu_run(fn, *args, **kwargs):
+    """Execute *fn* inside a ZeroGPU GPU allocation.
+
+    Used by ``obliterate`` to give each pipeline stage its own 5-minute
+    GPU allocation instead of sharing a single allocation for the whole
+    pipeline.  On non-ZeroGPU machines the ``@spaces.GPU`` decorator is a
+    no-op and this simply calls *fn* directly.
+    """
+    return fn(*args, **kwargs)
+
+
 def obliterate(model_choice: str, method_choice: str,
                prompt_volume_choice: str, dataset_source_choice: str,
                custom_harmful: str, custom_harmless: str,
@@ -2077,9 +2092,14 @@ def obliterate(model_choice: str, method_choice: str,
                progress=gr.Progress()):
     """Run the full obliteration pipeline, streaming log updates to the UI.
 
-    On ZeroGPU Spaces, this function runs on the visitor's GPU quota (up to
-    5 minutes).  The @spaces.GPU decorator allocates a GPU at call time and
-    releases it when the function returns.
+    On ZeroGPU Spaces, the pipeline is split into 3 GPU stages (up to 5 min
+    each) using the tourney-style approach: each stage gets its own
+    ``@spaces.GPU(duration=300)`` allocation via ``_obliterate_gpu_run``.
+    Between stages the model is offloaded to CPU and the GPU is released,
+    preventing the 5-minute ZeroGPU timeout from killing large-model runs.
+
+    On local/non-ZeroGPU machines, the pipeline runs in a single shot as
+    before (no time limit).
     """
     import os
     import re
@@ -2188,106 +2208,159 @@ def obliterate(model_choice: str, method_choice: str,
 
     quantization = _should_quantize(model_id, is_preset=is_preset)
 
+    def _create_pipeline(on_log, on_stage):
+        """Create the pipeline object and load prompts (no GPU required)."""
+        _t_pipeline_start = time.time()
+
+        # Load prompts — custom overrides dataset dropdown
+        if use_custom:
+            on_log("Using custom user-provided prompts...")
+            harmful_all, harmless_all = load_custom_prompts(
+                custom_harmful, custom_harmless or "",
+            )
+            on_log(f"Custom prompts: {len(harmful_all)} harmful, {len(harmless_all)} harmless")
+        else:
+            on_log(f"Loading dataset: {dataset_key}...")
+            harmful_all, harmless_all = load_dataset_source(dataset_key)
+            on_log(f"Dataset loaded: {len(harmful_all)} harmful, {len(harmless_all)} harmless prompts")
+        on_log(f"[timing] Dataset loaded at +{time.time() - _t_pipeline_start:.1f}s")
+
+        # Apply volume cap (-1 = use all)
+        if prompt_volume > 0:
+            n = min(prompt_volume, len(harmful_all), len(harmless_all))
+        else:
+            n = min(len(harmful_all), len(harmless_all))
+
+        if method == "informed":
+            from obliteratus.informed_pipeline import InformedAbliterationPipeline
+            pipeline = InformedAbliterationPipeline(
+                model_name=model_id,
+                output_dir=save_dir,
+                device="auto",
+                dtype="float16",
+                quantization=quantization,
+                trust_remote_code=is_preset,
+                harmful_prompts=harmful_all[:n],
+                harmless_prompts=harmless_all[:n],
+                on_stage=on_stage,
+                on_log=on_log,
+            )
+        else:
+            from obliteratus.abliterate import AbliterationPipeline
+            pipeline = AbliterationPipeline(
+                model_name=model_id,
+                output_dir=save_dir,
+                device="auto",
+                dtype="float16",
+                method=method,
+                quantization=quantization,
+                trust_remote_code=is_preset,
+                harmful_prompts=harmful_all[:n],
+                harmless_prompts=harmless_all[:n],
+                on_stage=on_stage,
+                on_log=on_log,
+                # Advanced overrides from UI
+                n_directions=int(adv_n_directions),
+                direction_method=adv_direction_method,
+                regularization=float(adv_regularization),
+                refinement_passes=int(adv_refinement_passes),
+                norm_preserve=adv_norm_preserve,
+                project_biases=adv_project_biases,
+                use_chat_template=adv_use_chat_template,
+                use_whitened_svd=adv_use_whitened_svd,
+                true_iterative_refinement=adv_true_iterative,
+                use_jailbreak_contrast=adv_jailbreak_contrast,
+                layer_adaptive_strength=adv_layer_adaptive,
+                safety_neuron_masking=adv_safety_neuron,
+                per_expert_directions=adv_per_expert,
+                attention_head_surgery=adv_attn_surgery,
+                use_sae_features=adv_sae_features,
+                invert_refusal=adv_invert_refusal,
+                reflection_strength=float(adv_reflection_strength),
+                project_embeddings=adv_project_embeddings,
+                embed_regularization=float(adv_embed_regularization),
+                activation_steering=adv_activation_steering,
+                steering_strength=float(adv_steering_strength),
+                expert_transplant=adv_expert_transplant,
+                transplant_blend=float(adv_transplant_blend),
+                use_wasserstein_optimal=adv_wasserstein_optimal,
+                spectral_cascade=adv_spectral_cascade,
+                spectral_bands=int(adv_spectral_bands),
+                spectral_threshold=float(adv_spectral_threshold),
+                verify_sample_size=int(adv_verify_sample_size),
+                layer_selection=adv_layer_selection,
+                winsorize_activations=adv_winsorize,
+                winsorize_percentile=float(adv_winsorize_percentile),
+                use_kl_optimization=adv_kl_optimization,
+                kl_budget=float(adv_kl_budget),
+                float_layer_interpolation=adv_float_layer_interp,
+                rdo_refinement=adv_rdo_refinement,
+                cot_aware=adv_cot_aware,
+                n_sae_features=int(adv_n_sae_features),
+            )
+        pipeline._bayesian_trials = int(adv_bayesian_trials)
+        return pipeline
+
     def run_pipeline():
         try:
-            _t_pipeline_start = time.time()
-            on_log(f"[timing] Pipeline thread started (ZeroGPU allocation: {300}s)")
-            if _ZEROGPU_AVAILABLE:
-                on_log("[timing] Running on ZeroGPU — GPU may be deallocated if pipeline exceeds duration")
+            on_log(f"[timing] Pipeline thread started")
+            pipeline = _create_pipeline(on_log, on_stage)
+            pipeline_ref[0] = pipeline
 
-            # Load prompts — custom overrides dataset dropdown
-            if use_custom:
-                on_log("Using custom user-provided prompts...")
-                harmful_all, harmless_all = load_custom_prompts(
-                    custom_harmful, custom_harmless or "",
-                )
-                on_log(f"Custom prompts: {len(harmful_all)} harmful, {len(harmless_all)} harmless")
-            else:
-                on_log(f"Loading dataset: {dataset_key}...")
-                harmful_all, harmless_all = load_dataset_source(dataset_key)
-                on_log(f"Dataset loaded: {len(harmful_all)} harmful, {len(harmless_all)} harmless prompts")
-            on_log(f"[timing] Dataset loaded at +{time.time() - _t_pipeline_start:.1f}s")
-
-            # Apply volume cap (-1 = use all)
-            if prompt_volume > 0:
-                n = min(prompt_volume, len(harmful_all), len(harmless_all))
-            else:
-                n = min(len(harmful_all), len(harmless_all))
-
-            if method == "informed":
-                # Use the analysis-guided InformedAbliterationPipeline
-                from obliteratus.informed_pipeline import InformedAbliterationPipeline
-                pipeline = InformedAbliterationPipeline(
-                    model_name=model_id,
-                    output_dir=save_dir,
-                    device="auto",
-                    dtype="float16",
-                    quantization=quantization,
-                    trust_remote_code=is_preset,
-                    harmful_prompts=harmful_all[:n],
-                    harmless_prompts=harmless_all[:n],
-                    on_stage=on_stage,
-                    on_log=on_log,
-                )
-                pipeline._bayesian_trials = int(adv_bayesian_trials)
-                pipeline_ref[0] = pipeline
-                pipeline.run_informed(gpu_start_time=t_start)
+            if _ZEROGPU_AVAILABLE:
+                # ── Staged GPU execution (tourney-style) ──────────────────
+                # Each stage gets its own 5-minute GPU allocation instead of
+                # sharing a single 300s budget.  Between stages the model is
+                # moved to CPU and the GPU is released.
+                on_log("[staged] ZeroGPU detected — using staged GPU execution (up to 5 min per stage)")
+
+                if method == "informed":
+                    # Informed pipeline: SUMMON+PROBE | ANALYZE+DISTILL+EXCISE | VERIFY+REBIRTH
+                    on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
+                    _obliterate_gpu_run(pipeline.run_stage_summon_probe, time.time())
+                    pipeline._offload_to_cpu()
+                    on_log("[staged] GPU released after Stage 1\n")
+
+                    on_log("\u26a1 [staged] GPU Stage 2/3: ANALYZE + DISTILL + EXCISE")
+                    def _informed_s2():
+                        pipeline._restore_to_gpu()
+                        pipeline.run_stage_analyze_distill_excise()
+                    _obliterate_gpu_run(_informed_s2)
+                    pipeline._offload_to_cpu()
+                    on_log("[staged] GPU released after Stage 2\n")
+
+                    on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
+                    def _informed_s3():
+                        pipeline._restore_to_gpu()
+                        pipeline.run_stage_verify_rebirth_informed()
+                    _obliterate_gpu_run(_informed_s3)
+                else:
+                    # Standard pipeline: SUMMON+PROBE | DISTILL+EXCISE | VERIFY+REBIRTH
+                    on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
+                    _obliterate_gpu_run(pipeline.run_stage_summon_probe, time.time())
+                    pipeline._offload_to_cpu()
+                    on_log("[staged] GPU released after Stage 1\n")
+
+                    on_log("\u26a1 [staged] GPU Stage 2/3: DISTILL + EXCISE")
+                    def _standard_s2():
+                        pipeline._restore_to_gpu()
+                        pipeline.run_stage_distill_excise()
+                    _obliterate_gpu_run(_standard_s2)
+                    pipeline._offload_to_cpu()
+                    on_log("[staged] GPU released after Stage 2\n")
+
+                    on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
+                    def _standard_s3():
+                        pipeline._restore_to_gpu()
+                        pipeline.run_stage_verify_rebirth()
+                    _obliterate_gpu_run(_standard_s3)
             else:
-                from obliteratus.abliterate import AbliterationPipeline
-                pipeline = AbliterationPipeline(
-                    model_name=model_id,
-                    output_dir=save_dir,
-                    device="auto",
-                    dtype="float16",
-                    method=method,
-                    quantization=quantization,
-                    trust_remote_code=is_preset,
-                    harmful_prompts=harmful_all[:n],
-                    harmless_prompts=harmless_all[:n],
-                    on_stage=on_stage,
-                    on_log=on_log,
-                    # Advanced overrides from UI
-                    n_directions=int(adv_n_directions),
-                    direction_method=adv_direction_method,
-                    regularization=float(adv_regularization),
-                    refinement_passes=int(adv_refinement_passes),
-                    norm_preserve=adv_norm_preserve,
-                    project_biases=adv_project_biases,
-                    use_chat_template=adv_use_chat_template,
-                    use_whitened_svd=adv_use_whitened_svd,
-                    true_iterative_refinement=adv_true_iterative,
-                    use_jailbreak_contrast=adv_jailbreak_contrast,
-                    layer_adaptive_strength=adv_layer_adaptive,
-                    safety_neuron_masking=adv_safety_neuron,
-                    per_expert_directions=adv_per_expert,
-                    attention_head_surgery=adv_attn_surgery,
-                    use_sae_features=adv_sae_features,
-                    invert_refusal=adv_invert_refusal,
-                    reflection_strength=float(adv_reflection_strength),
-                    project_embeddings=adv_project_embeddings,
-                    embed_regularization=float(adv_embed_regularization),
-                    activation_steering=adv_activation_steering,
-                    steering_strength=float(adv_steering_strength),
-                    expert_transplant=adv_expert_transplant,
-                    transplant_blend=float(adv_transplant_blend),
-                    use_wasserstein_optimal=adv_wasserstein_optimal,
-                    spectral_cascade=adv_spectral_cascade,
-                    spectral_bands=int(adv_spectral_bands),
-                    spectral_threshold=float(adv_spectral_threshold),
-                    verify_sample_size=int(adv_verify_sample_size),
-                    layer_selection=adv_layer_selection,
-                    winsorize_activations=adv_winsorize,
-                    winsorize_percentile=float(adv_winsorize_percentile),
-                    use_kl_optimization=adv_kl_optimization,
-                    kl_budget=float(adv_kl_budget),
-                    float_layer_interpolation=adv_float_layer_interp,
-                    rdo_refinement=adv_rdo_refinement,
-                    cot_aware=adv_cot_aware,
-                    n_sae_features=int(adv_n_sae_features),
-                )
-                pipeline._bayesian_trials = int(adv_bayesian_trials)
-                pipeline_ref[0] = pipeline
-                pipeline.run(gpu_start_time=t_start)
+                # ── Local/non-ZeroGPU: single-shot execution ──────────────
+                on_log(f"[timing] Running locally (no GPU time limit)")
+                if method == "informed":
+                    pipeline.run_informed(gpu_start_time=t_start)
+                else:
+                    pipeline.run(gpu_start_time=t_start)
         except Exception as e:
             error_ref[0] = e
             tb = traceback.format_exc()

obliteratus/abliterate.py

@@ -985,6 +985,146 @@ class AbliterationPipeline:
         self.log(f"[timing] REBIRTH complete at +{time.time() - _t0:.1f}s — pipeline finished")
         return result
 
+    # ── Staged execution (ZeroGPU-safe) ──────────────────────────────────
+    # These methods split the pipeline into independent GPU stages so each
+    # stage can run in its own @spaces.GPU allocation (up to 5 min each).
+    # Between stages the caller moves the model to CPU and releases the GPU.
+    # This mirrors the tourney approach and prevents ZeroGPU timeouts on
+    # large models that need more than 5 minutes total.
+
+    def run_stage_summon_probe(self, gpu_start_time: float | None = None):
+        """GPU Stage 1: Load model and collect activations (SUMMON + PROBE).
+
+        After this stage the caller should call ``_offload_to_cpu()`` to
+        move the model off GPU before releasing the ZeroGPU allocation.
+        """
+        for h in self._steering_hooks:
+            h.remove()
+        self._steering_hooks.clear()
+        self._staged_t0 = time.time()
+        self._pipeline_start_time = gpu_start_time if gpu_start_time is not None else self._staged_t0
+        self._summon()
+        self.log(f"[timing] SUMMON complete at +{time.time() - self._staged_t0:.1f}s")
+        self._free_gpu_memory()
+        self._probe()
+        self.log(f"[timing] PROBE complete at +{time.time() - self._staged_t0:.1f}s")
+        self._free_gpu_memory()
+
+    def run_stage_distill_excise(self):
+        """GPU Stage 2: Extract directions, modify weights, save checkpoint
+        (DISTILL + EXCISE + quick checkpoint).
+
+        After this stage the caller should call ``_offload_to_cpu()`` to
+        move the model off GPU before releasing the ZeroGPU allocation.
+        """
+        # Reset time budget for this GPU allocation
+        self._pipeline_start_time = time.time()
+        self._distill()
+        self.log(f"[timing] DISTILL complete at +{time.time() - self._staged_t0:.1f}s")
+        # Free raw per-prompt activations now that means/subspaces are extracted
+        self._harmful_acts.clear()
+        self._harmless_acts.clear()
+        self._jailbreak_acts.clear()
+        self._harmful_means.clear()
+        self._harmless_means.clear()
+        self._routing_harmful.clear()
+        self._routing_harmless.clear()
+        self._free_gpu_memory()
+        self._capture_baseline_kl_logits()
+        self._excise()
+        self.log(f"[timing] EXCISE complete at +{time.time() - self._staged_t0:.1f}s")
+        self._free_gpu_memory()
+        self._save_quick_checkpoint()
+
+    def run_stage_verify_rebirth(self) -> Path:
+        """GPU Stage 3: Quality verification and final save (VERIFY + REBIRTH).
+
+        Returns:
+            Path to the saved model directory.
+        """
+        # Reset time budget for this GPU allocation
+        self._pipeline_start_time = time.time()
+        self._verify()
+        self.log(f"[timing] VERIFY complete at +{time.time() - self._staged_t0:.1f}s")
+        self._free_gpu_memory()
+        result = self._rebirth()
+        self.log(f"[timing] REBIRTH complete at +{time.time() - self._staged_t0:.1f}s — pipeline finished")
+        return result
+
+    def _offload_to_cpu(self):
+        """Move model and intermediate tensors to CPU for ZeroGPU stage boundary.
+
+        Called between staged GPU allocations so the model survives GPU
+        deallocation.  Also moves cached activation tensors to CPU.
+        """
+        if self.handle is None or self.handle.model is None:
+            return
+        model = self.handle.model
+        # Remember the GPU device for later restoration
+        try:
+            self._gpu_device = self._get_model_device(model)
+        except (StopIteration, RuntimeError):
+            self._gpu_device = torch.device("cuda:0")
+
+        self.log("[staged] Moving model to CPU between GPU allocations...")
+        t0 = time.time()
+        model.to("cpu")
+
+        # Move any cached tensors that are still on GPU
+        for tensor_dict in (
+            self._harmful_means, self._harmless_means,
+            self._jailbreak_means,
+        ):
+            for k in list(tensor_dict.keys()):
+                if isinstance(tensor_dict[k], torch.Tensor) and tensor_dict[k].device.type != "cpu":
+                    tensor_dict[k] = tensor_dict[k].cpu()
+
+        for act_dict in (self._harmful_acts, self._harmless_acts, self._jailbreak_acts):
+            for k in list(act_dict.keys()):
+                if isinstance(act_dict[k], list):
+                    act_dict[k] = [t.cpu() if isinstance(t, torch.Tensor) and t.device.type != "cpu" else t
+                                   for t in act_dict[k]]
+
+        # Move refusal directions/subspaces to CPU
+        for k in list(self.refusal_directions.keys()):
+            v = self.refusal_directions[k]
+            if isinstance(v, torch.Tensor) and v.device.type != "cpu":
+                self.refusal_directions[k] = v.cpu()
+        for k in list(self.refusal_subspaces.keys()):
+            v = self.refusal_subspaces[k]
+            if isinstance(v, torch.Tensor) and v.device.type != "cpu":
+                self.refusal_subspaces[k] = v.cpu()
+
+        # Move baseline KL logits to CPU
+        if hasattr(self, "_kl_baseline_logits") and self._kl_baseline_logits is not None:
+            self._kl_baseline_logits = [
+                t.cpu() if isinstance(t, torch.Tensor) and t.device.type != "cpu" else t
+                for t in self._kl_baseline_logits
+            ]
+
+        dev.free_gpu_memory()
+        elapsed = time.time() - t0
+        self.log(f"[staged] Model offloaded to CPU ({elapsed:.1f}s)")
+
+    def _restore_to_gpu(self):
+        """Move model back to GPU at the start of a new GPU allocation.
+
+        Uses the device stored by ``_offload_to_cpu()``, defaulting to
+        ``cuda:0`` if not set.
+        """
+        if self.handle is None or self.handle.model is None:
+            return
+        target = getattr(self, "_gpu_device", None) or torch.device("cuda:0")
+        # On ZeroGPU, the device index may change between allocations,
+        # so always use cuda:0 which is the only device available.
+        if target.type == "cuda":
+            target = torch.device("cuda:0")
+        self.log(f"[staged] Moving model to {target}...")
+        t0 = time.time()
+        self.handle.model.to(target)
+        elapsed = time.time() - t0
+        self.log(f"[staged] Model restored to {target} ({elapsed:.1f}s)")
+
     # ── Stage 1: SUMMON ─────────────────────────────────────────────────
 
     def _summon(self):

obliteratus/informed_pipeline.py

@@ -299,6 +299,33 @@ class InformedAbliterationPipeline(AbliterationPipeline):
         self._report.total_duration = time.time() - t0
         return output_path, self._report
 
+    # ── Staged execution (ZeroGPU-safe) ──────────────────────────────
+    # Split the informed pipeline into 3 GPU stages so each gets its own
+    # @spaces.GPU(duration=300) allocation.  SUMMON+PROBE is inherited
+    # from AbliterationPipeline.run_stage_summon_probe().
+
+    def run_stage_analyze_distill_excise(self):
+        """GPU Stage 2 (informed): ANALYZE + DISTILL + EXCISE + quick checkpoint."""
+        # Reset time budget for this GPU allocation
+        self._pipeline_start_time = time.time()
+        self._analyze()
+        self._distill_informed()
+        self._excise_informed()
+        self._save_quick_checkpoint()
+
+    def run_stage_verify_rebirth_informed(self):
+        """GPU Stage 3 (informed): VERIFY + Ouroboros compensation + REBIRTH.
+
+        Stores the result path and report in instance attributes so the
+        caller can retrieve them after the GPU allocation returns.
+        """
+        # Reset time budget for this GPU allocation
+        self._pipeline_start_time = time.time()
+        self._verify_and_compensate()
+        output_path = self._rebirth_informed()
+        self._report.total_duration = time.time() - getattr(self, "_staged_t0", time.time())
+        self._staged_result = (output_path, self._report)
+
     # ── Stage 3: ANALYZE ─────────────────────────────────────────────
 
     def _analyze(self):