Upload 133 files

app.py

@@ -297,15 +297,16 @@ def _recover_after_obliterate():
                 _short = model_choice.split("/")[-1] if "/" in model_choice else model_choice
                 _label = f"{method} on {_short} ({_ts}) [recovered]"
                 _last_obliterated_label = _label
-                _session_models[_label] = {
-                    "model_id": data.get("model_id", model_choice),
-                    "model_choice": model_choice,
-                    "method": method or "unknown",
-                    "dataset_key": "",
-                    "prompt_volume": 0,
-                    "output_dir": save_dir,
-                    "source": "recovered",
-                }
+                with _lock:
+                    _session_models[_label] = {
+                        "model_id": data.get("model_id", model_choice),
+                        "model_choice": model_choice,
+                        "method": method or "unknown",
+                        "dataset_key": "",
+                        "prompt_volume": 0,
+                        "output_dir": save_dir,
+                        "source": "recovered",
+                    }
 
     if not recovered:
         with _lock:
@@ -430,15 +431,16 @@ def _recover_sessions_from_disk() -> None:
             label = data.get("label", p.name)
             if label in _session_models:
                 continue  # already registered
-            _session_models[label] = {
-                "model_id": data.get("model_id", ""),
-                "model_choice": data.get("model_choice", data.get("model_id", "")),
-                "method": data.get("method", "unknown"),
-                "dataset_key": data.get("dataset_key", ""),
-                "prompt_volume": data.get("prompt_volume", 0),
-                "output_dir": str(p),
-                "source": data.get("source", "recovered"),
-            }
+            with _lock:
+                _session_models[label] = {
+                    "model_id": data.get("model_id", ""),
+                    "model_choice": data.get("model_choice", data.get("model_id", "")),
+                    "method": data.get("method", "unknown"),
+                    "dataset_key": data.get("dataset_key", ""),
+                    "prompt_volume": data.get("prompt_volume", 0),
+                    "output_dir": str(p),
+                    "source": data.get("source", "recovered"),
+                }
             found_any = True
             # Track the latest for auto-select
             _last_obliterated_label = label
@@ -453,10 +455,10 @@ def _recover_sessions_from_disk() -> None:
     # If we recovered sessions and _state has no valid output_dir, set it to
     # the most recent checkpoint so chat_respond can reload from disk.
     # Also overwrite a stale output_dir that points to a non-existent path.
-    _cur_dir = _state.get("output_dir")
-    _needs_update = not _cur_dir or not Path(_cur_dir).exists()
-    if found_any and _needs_update:
-        with _lock:
+    with _lock:
+        _cur_dir = _state.get("output_dir")
+        _needs_update = not _cur_dir or not Path(_cur_dir).exists()
+        if found_any and _needs_update:
             latest = _last_obliterated_label
             if latest and latest in _session_models:
                 _state["output_dir"] = _session_models[latest]["output_dir"]
@@ -821,6 +823,15 @@ def push_session_to_hub(
     if not output_dir or not Path(output_dir).exists():
         yield f"**Error:** Model directory not found: `{output_dir}`", ""
         return
+    # Validate output_dir is under /tmp to prevent directory traversal
+    try:
+        _resolved = Path(output_dir).resolve()
+        if not str(_resolved).startswith("/tmp/"):
+            yield "**Error:** Model directory must be under `/tmp`.", ""
+            return
+    except Exception:
+        yield "**Error:** Invalid model directory path.", ""
+        return
 
     # Resolve repo ID
     repo_id = hub_repo_id.strip() if hub_repo_id else ""
@@ -1086,11 +1097,9 @@ def _cleanup_disk():
             shutil.rmtree(path, ignore_errors=True)
             freed += size
 
-    # Clear session model cache (checkpoints are gone)
-    _session_models.clear()
-
-    # Clear stale output_dir reference (checkpoints were just deleted)
+    # Clear session model cache and stale state (checkpoints are gone)
     with _lock:
+        _session_models.clear()
         _state["output_dir"] = None
         _state["model_name"] = None
         _state["method"] = None
@@ -1574,14 +1583,15 @@ def benchmark(
         bench_save_path = f"/tmp/bench_{method_key}"
         if entry.get("error") is None:
             label = f"{entry['method']} on {model_id.split('/')[-1]}"
-            _bench_configs[label] = {
-                "model_id": model_id,
-                "model_choice": model_choice,
-                "method": method_key,
-                "dataset_key": dataset_key,
-                "prompt_volume": prompt_volume,
-                "output_dir": bench_save_path,
-            }
+            with _lock:
+                _bench_configs[label] = {
+                    "model_id": model_id,
+                    "model_choice": model_choice,
+                    "method": method_key,
+                    "dataset_key": dataset_key,
+                    "prompt_volume": prompt_volume,
+                    "output_dir": bench_save_path,
+                }
             _persist_session_meta(bench_save_path, label, {
                 "model_id": model_id,
                 "model_choice": model_choice,
@@ -1928,14 +1938,15 @@ def benchmark_multi_model(
         mm_save_path = f"/tmp/bench_mm_{mi}"
         if entry.get("error") is None:
             label = f"{method_key} on {model_id.split('/')[-1]}"
-            _bench_configs[label] = {
-                "model_id": model_id,
-                "model_choice": model_display,
-                "method": method_key,
-                "dataset_key": dataset_key,
-                "prompt_volume": prompt_volume,
-                "output_dir": mm_save_path,
-            }
+            with _lock:
+                _bench_configs[label] = {
+                    "model_id": model_id,
+                    "model_choice": model_display,
+                    "method": method_key,
+                    "dataset_key": dataset_key,
+                    "prompt_volume": prompt_volume,
+                    "output_dir": mm_save_path,
+                }
             _persist_session_meta(mm_save_path, label, {
                 "model_id": model_id,
                 "model_choice": model_display,
@@ -2067,9 +2078,17 @@ def _restore_and_run_stage(pipeline, stage_method_name):
     Module-level function so it is picklable for ZeroGPU serialization.
     Wraps execution in try/except to preserve the full traceback before
     ZeroGPU's error handler reduces it to just the exception class name.
+
+    If the pipeline model is not in memory (ZeroGPU state loss), the stage
+    method itself handles recovery via ``_reload_model_for_stage()`` and
+    ``_load_staged_state()`` when ``_staged_state_dir`` is set.
     """
     try:
-        pipeline._restore_to_gpu()
+        # Try to restore model to GPU if it's already in memory (same-process
+        # case or non-ZeroGPU).  If the model is None (ZeroGPU state loss),
+        # skip — the stage method handles recovery.
+        if pipeline.handle is not None and pipeline.handle.model is not None:
+            pipeline._restore_to_gpu()
         getattr(pipeline, stage_method_name)()
     except Exception as e:
         import traceback as _tb
@@ -2209,7 +2228,15 @@ def obliterate(model_choice: str, method_choice: str,
         return
 
     # Resolve dataset source — custom prompts override the dropdown
+    _MAX_CUSTOM_PROMPT_LINES = 10_000
     use_custom = custom_harmful and custom_harmful.strip()
+    if use_custom and custom_harmful.count("\n") > _MAX_CUSTOM_PROMPT_LINES:
+        yield (
+            f"**Error:** Custom prompts exceed {_MAX_CUSTOM_PROMPT_LINES} lines. "
+            "Please reduce the number of prompts to avoid memory exhaustion.",
+            "", gr.update(), gr.update(), gr.update(), gr.update(),
+        )
+        return
     dataset_key = get_source_key_from_label(dataset_source_choice) if dataset_source_choice else "builtin"
 
     # Unstick stale "obliterating" status left behind by ZeroGPU timeout
@@ -2366,37 +2393,50 @@ def obliterate(model_choice: str, method_choice: str,
                 # ── 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.
+                # saved to disk so state survives ZeroGPU's cross-process
+                # serialization (each @spaces.GPU call runs in a separate
+                # worker process that pickles args, so in-memory mutations
+                # to the pipeline don't propagate back).
                 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")
-                    _gpu_run_picklable(pipeline, 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")
-                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_analyze_distill_excise")
-                    pipeline._offload_to_cpu()
-                    on_log("[staged] GPU released after Stage 2\n")
-
-                    on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
-                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth_informed")
-                else:
-                    # Standard pipeline: SUMMON+PROBE | DISTILL+EXCISE | VERIFY+REBIRTH
-                    on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
-                    _gpu_run_picklable(pipeline, 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")
-                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_distill_excise")
-                    pipeline._offload_to_cpu()
-                    on_log("[staged] GPU released after Stage 2\n")
-
-                    on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
-                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth")
+                # Create a temp dir for cross-process state persistence
+                import tempfile as _tempfile
+                _staged_dir = _tempfile.mkdtemp(prefix="obliterate_staged_")
+                pipeline._staged_state_dir = _staged_dir
+                on_log(f"[staged] State persistence dir: {_staged_dir}")
+
+                try:
+                    if method == "informed":
+                        # Informed pipeline: SUMMON+PROBE | ANALYZE+DISTILL+EXCISE | VERIFY+REBIRTH
+                        on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
+                        _gpu_run_picklable(pipeline, pipeline.run_stage_summon_probe, time.time())
+                        on_log("[staged] GPU released after Stage 1\n")
+
+                        on_log("\u26a1 [staged] GPU Stage 2/3: ANALYZE + DISTILL + EXCISE")
+                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_analyze_distill_excise")
+                        on_log("[staged] GPU released after Stage 2\n")
+
+                        on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
+                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth_informed")
+                    else:
+                        # Standard pipeline: SUMMON+PROBE | DISTILL+EXCISE | VERIFY+REBIRTH
+                        on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
+                        _gpu_run_picklable(pipeline, pipeline.run_stage_summon_probe, time.time())
+                        on_log("[staged] GPU released after Stage 1\n")
+
+                        on_log("\u26a1 [staged] GPU Stage 2/3: DISTILL + EXCISE")
+                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_distill_excise")
+                        on_log("[staged] GPU released after Stage 2\n")
+
+                        on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
+                        _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth")
+                finally:
+                    # Clean up staged state temp dir
+                    import shutil as _shutil
+                    try:
+                        _shutil.rmtree(_staged_dir, ignore_errors=True)
+                    except Exception:
+                        pass
             else:
                 # ── Local/non-ZeroGPU: single-shot execution ──────────────
                 on_log(f"[timing] Running locally (no GPU time limit)")
@@ -2531,6 +2571,20 @@ def obliterate(model_choice: str, method_choice: str,
 
     worker.join(timeout=30)
 
+    # If worker is still alive after join timeout, it's hung — treat as error
+    if worker.is_alive():
+        _mark_live_log_finished()
+        log_lines.append("\nERROR: Pipeline worker thread did not finish within 30s after loop exit.")
+        with _lock:
+            _state["status"] = "idle"
+            _state["obliterate_started_at"] = None
+            _state["log"] = log_lines
+        yield (
+            "**Error:** Pipeline worker hung after completion. Check logs for details.",
+            "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update(),
+        )
+        return
+
     # Handle error
     if error_ref[0] is not None:
         _mark_live_log_finished()
@@ -2851,7 +2905,8 @@ def obliterate(model_choice: str, method_choice: str,
         # Set skip flag so the .change handler doesn't trigger a wasteful
         # GPU re-allocation — the model is already loaded.
         global _skip_session_load
-        _skip_session_load = 2  # both session_model_dd and ab_session_model_dd fire .change
+        with _lock:
+            _skip_session_load = 2  # both session_model_dd and ab_session_model_dd fire .change
         _dd_update = gr.update(
             choices=_get_session_model_choices(),
             value=_last_obliterated_label or None,
@@ -3151,8 +3206,11 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
     # Skip if the obliterate function just set the dropdown value — the model
     # is already loaded and we'd just waste GPU quota re-allocating.
     global _skip_session_load
-    if _skip_session_load > 0:
-        _skip_session_load -= 1
+    with _lock:
+        _should_skip = _skip_session_load > 0
+        if _should_skip:
+            _skip_session_load -= 1
+    if _should_skip:
         # Verify the model is actually usable — not just that status says "ready".
         # ZeroGPU can evict the model while status stays "ready", and the counter
         # can get out of sync if only one dropdown .change fires instead of both.
@@ -3336,6 +3394,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
                 _clear_gpu()
                 with _lock:
                     _state["status"] = "idle"
+                    _state["obliterate_started_at"] = None
                 yield (
                     f"**Error:** Could not load {choice} from checkpoint (GPU too small).",
                     get_chat_header(),
@@ -3391,6 +3450,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
     if error_ref[0] is not None:
         with _lock:
             _state["status"] = "idle"
+            _state["obliterate_started_at"] = None
         yield f"**Error loading {choice}:** {error_ref[0]}", get_chat_header()
         return
 
@@ -3400,6 +3460,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
         _state["tokenizer"] = pipeline.handle.tokenizer
         _state["steering"] = None
         _state["status"] = "ready"
+        _state["obliterate_started_at"] = None
         _state["output_dir"] = "/tmp/obliterated"  # re-abliteration fallback path
 
     pipeline_ref[0] = None
@@ -3640,6 +3701,9 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[
 
     except Exception as e:
         original_response = f"*Could not load original model for comparison: {e}*"
+        # Ensure GPU memory is freed even if original model load/gen failed
+        gc.collect()
+        dev.empty_cache()
 
     # Restore abliterated model to GPU for subsequent chat/operations.
     # Use torch.device("cuda") rather than the captured abl_device, since
@@ -3728,7 +3792,7 @@ def strength_sweep(model_choice: str, method_choice: str,
                 nonlocal run_error
                 run_error = e
 
-        worker = threading.Thread(target=_run_sweep_point)
+        worker = threading.Thread(target=_run_sweep_point, daemon=True)
         worker.start()
         while worker.is_alive():
             worker.join(timeout=2.0)
@@ -4122,15 +4186,16 @@ def export_artifacts():
     import zipfile
     import os
 
-    if _state["status"] != "ready":
-        return None, "No abliterated model loaded. Run obliteration first."
+    with _lock:
+        if _state["status"] != "ready":
+            return None, "No abliterated model loaded. Run obliteration first."
+        model_name = _state.get("model_name", "unknown")
+        method = _state.get("method", "unknown")
+        log_lines = list(_state.get("log", []))  # copy to avoid mutation
+        steering = _state.get("steering")
 
     export_dir = tempfile.mkdtemp(prefix="obliteratus_export_")
 
-    model_name = _state.get("model_name", "unknown")
-    method = _state.get("method", "unknown")
-    log_lines = _state.get("log", [])
-
     exported_files = []
 
     # 1. Pipeline log
@@ -4145,7 +4210,6 @@ def export_artifacts():
     exported_files.append("pipeline_log.txt")
 
     # 2. Steering metadata (refusal directions + strong layers)
-    steering = _state.get("steering")
     if steering:
         # Save directions as .pt
         directions = steering.get("refusal_directions", {})

obliteratus/abliterate.py

@@ -763,6 +763,7 @@ class AbliterationPipeline:
                 self.refinement_passes = min(self.refinement_passes, 1)
 
         self.handle: ModelHandle | None = None
+        self._staged_state_dir: str | None = None  # temp dir for cross-process state persistence (ZeroGPU)
         self._staged_t0: float = time.time()  # overwritten in run_stage_summon_probe; init here for safety
         self.refusal_directions: dict[int, torch.Tensor] = {}  # per-layer primary direction
         self.refusal_subspaces: dict[int, torch.Tensor] = {}   # per-layer SVD subspace (n_dirs x hidden)
@@ -998,6 +999,10 @@ class AbliterationPipeline:
 
         After this stage the caller should call ``_offload_to_cpu()`` to
         move the model off GPU before releasing the ZeroGPU allocation.
+
+        When ``_staged_state_dir`` is set, intermediate state (activations,
+        means, routing data) is saved to disk at the end of this stage so
+        it can survive ZeroGPU's cross-process serialization boundary.
         """
         for h in self._steering_hooks:
             h.remove()
@@ -1010,6 +1015,8 @@ class AbliterationPipeline:
         self._probe()
         self.log(f"[timing] PROBE complete at +{time.time() - self._staged_t0:.1f}s")
         self._free_gpu_memory()
+        # Save state to disk so it survives ZeroGPU worker process boundary
+        self._save_staged_state()
 
     def run_stage_distill_excise(self):
         """GPU Stage 2: Extract directions, modify weights, save checkpoint
@@ -1019,14 +1026,19 @@ class AbliterationPipeline:
         move the model off GPU before releasing the ZeroGPU allocation.
         """
         if self.handle is None or self.handle.model is None:
-            raise RuntimeError(
-                "Model not loaded at start of DISTILL+EXCISE stage. "
-                "This usually means ZeroGPU pickling lost pipeline state "
-                "between GPU stages (model was loaded in stage 1's worker "
-                "but not propagated back to the main process). "
-                f"handle={self.handle!r}, _harmful_means has "
-                f"{len(self._harmful_means)} layers"
-            )
+            if self._staged_state_dir:
+                self.log("[staged] Model not in memory — recovering from ZeroGPU state loss...")
+                self._reload_model_for_stage()
+                self._load_staged_state()
+            else:
+                raise RuntimeError(
+                    "Model not loaded at start of DISTILL+EXCISE stage. "
+                    "This usually means ZeroGPU pickling lost pipeline state "
+                    "between GPU stages (model was loaded in stage 1's worker "
+                    "but not propagated back to the main process). "
+                    f"handle={self.handle!r}, _harmful_means has "
+                    f"{len(self._harmful_means)} layers"
+                )
         # Reset time budget for this GPU allocation
         self._pipeline_start_time = time.time()
         self._distill()
@@ -1053,10 +1065,16 @@ class AbliterationPipeline:
             Path to the saved model directory.
         """
         if self.handle is None or self.handle.model is None:
-            raise RuntimeError(
-                "Model not loaded at start of VERIFY+REBIRTH stage. "
-                "Pipeline state may have been lost between GPU stages."
-            )
+            if self._staged_state_dir:
+                # Load the modified model from the quick checkpoint saved
+                # by Stage 2 (weights were modified by EXCISE).
+                self.log("[staged] Model not in memory — recovering from ZeroGPU state loss...")
+                self._reload_model_for_stage(source=str(self.output_dir))
+            else:
+                raise RuntimeError(
+                    "Model not loaded at start of VERIFY+REBIRTH stage. "
+                    "Pipeline state may have been lost between GPU stages."
+                )
         # Reset time budget for this GPU allocation
         self._pipeline_start_time = time.time()
         self._verify()
@@ -1111,11 +1129,9 @@ class AbliterationPipeline:
                 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
-            ]
+        if self._baseline_first_token_logits is not None:
+            if isinstance(self._baseline_first_token_logits, torch.Tensor) and self._baseline_first_token_logits.device.type != "cpu":
+                self._baseline_first_token_logits = self._baseline_first_token_logits.cpu()
 
         dev.free_gpu_memory()
         elapsed = time.time() - t0
@@ -1137,9 +1153,108 @@ class AbliterationPipeline:
         self.log(f"[staged] Moving model to {target}...")
         t0 = time.time()
         self.handle.model.to(target)
+
+        # Move cached tensors back to GPU (mirrors _offload_to_cpu)
+        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 != target:
+                    tensor_dict[k] = tensor_dict[k].to(target)
+
+        for k in list(self.refusal_directions.keys()):
+            v = self.refusal_directions[k]
+            if isinstance(v, torch.Tensor) and v.device != target:
+                self.refusal_directions[k] = v.to(target)
+        for k in list(self.refusal_subspaces.keys()):
+            v = self.refusal_subspaces[k]
+            if isinstance(v, torch.Tensor) and v.device != target:
+                self.refusal_subspaces[k] = v.to(target)
+
         elapsed = time.time() - t0
         self.log(f"[staged] Model restored to {target} ({elapsed:.1f}s)")
 
+    def _save_staged_state(self):
+        """Save intermediate pipeline state to disk for ZeroGPU cross-process persistence.
+
+        ZeroGPU runs each ``@spaces.GPU``-decorated call in a separate worker
+        process.  Mutations to the pipeline (model loading, activation
+        collection) happen on a pickled copy inside the worker and are NOT
+        propagated back.  This method saves the state that subsequent stages
+        need to a temporary directory on disk so it survives the process
+        boundary.
+        """
+        if not self._staged_state_dir:
+            return
+        import os
+        os.makedirs(self._staged_state_dir, exist_ok=True)
+        t0 = time.time()
+
+        def _cpu(t):
+            return t.cpu() if isinstance(t, torch.Tensor) and t.device.type != "cpu" else t
+
+        state = {
+            "_staged_t0": self._staged_t0,
+            "_pipeline_start_time": getattr(self, "_pipeline_start_time", self._staged_t0),
+            "_harmful_acts": {k: [_cpu(t) for t in v] for k, v in self._harmful_acts.items()},
+            "_harmless_acts": {k: [_cpu(t) for t in v] for k, v in self._harmless_acts.items()},
+            "_jailbreak_acts": {k: [_cpu(t) for t in v] for k, v in self._jailbreak_acts.items()},
+            "_harmful_means": {k: _cpu(v) for k, v in self._harmful_means.items()},
+            "_harmless_means": {k: _cpu(v) for k, v in self._harmless_means.items()},
+            "_jailbreak_means": {k: _cpu(v) for k, v in self._jailbreak_means.items()},
+            "_routing_harmful": {k: [_cpu(t) for t in v] for k, v in self._routing_harmful.items()},
+            "_routing_harmless": {k: [_cpu(t) for t in v] for k, v in self._routing_harmless.items()},
+            "_baseline_first_token_logits": _cpu(self._baseline_first_token_logits) if self._baseline_first_token_logits is not None else None,
+            "_kl_eval_prompts": list(self._kl_eval_prompts),
+        }
+        state_path = os.path.join(self._staged_state_dir, "staged_state.pt")
+        torch.save(state, state_path)
+        elapsed = time.time() - t0
+        self.log(f"[staged] Saved intermediate state to disk ({elapsed:.1f}s)")
+
+    def _load_staged_state(self):
+        """Load intermediate pipeline state saved by ``_save_staged_state``.
+
+        Called at the start of Stage 2 (and later) when ZeroGPU pickling
+        caused the in-memory state to be lost.
+        """
+        import os
+        state_path = os.path.join(self._staged_state_dir, "staged_state.pt")
+        if not os.path.exists(state_path):
+            raise RuntimeError(
+                f"Staged state file not found at {state_path}. "
+                "Stage 1 may not have saved its state correctly."
+            )
+        t0 = time.time()
+        state = torch.load(state_path, weights_only=False)
+        for key, value in state.items():
+            setattr(self, key, value)
+        elapsed = time.time() - t0
+        self.log(f"[staged] Loaded intermediate state from disk ({elapsed:.1f}s)")
+
+    def _reload_model_for_stage(self, source: str | None = None):
+        """Reload the model for a new GPU stage after ZeroGPU state loss.
+
+        Args:
+            source: Path to load the model from.  If None, reloads from the
+                original HuggingFace model (for Stage 2, where weights are
+                unmodified).  If a path (e.g. ``self.output_dir``), loads
+                from a previously saved checkpoint (for Stage 3, after
+                weight modification by EXCISE).
+        """
+        model_name = source or self.model_name
+        self.log(f"[staged] Reloading model from {'checkpoint' if source else 'HF cache'}: {model_name}")
+        self.handle = load_model(
+            model_name=model_name,
+            task="causal_lm",
+            device=self.device,
+            dtype=self.dtype,
+            trust_remote_code=self.trust_remote_code,
+            quantization=self.quantization,
+        )
+        self.log(f"[staged] Model reloaded successfully")
+
     # ── Stage 1: SUMMON ─────────────────────────���───────────────────────
 
     def _summon(self):
@@ -6233,7 +6348,17 @@ class AbliterationPipeline:
             dat_file = Path(offload_dir) / f"{key}.dat"
             if safetensors_file.exists():
                 data = load_file(str(safetensors_file))
-                state_dict[key] = data[key] if key in data else next(iter(data.values()))
+                if key in data:
+                    state_dict[key] = data[key]
+                elif len(data) == 1:
+                    # Accelerate stores single-tensor files; use the only tensor
+                    state_dict[key] = next(iter(data.values()))
+                else:
+                    raise RuntimeError(
+                        f"Cannot resolve meta tensor '{key}': safetensors file "
+                        f"contains {len(data)} tensors but key not found. "
+                        f"Available keys: {list(data.keys())}"
+                    )
             elif dat_file.exists():
                 # Accelerate's .dat format: raw tensor bytes with shape/dtype metadata
                 import numpy as np

obliteratus/informed_pipeline.py

@@ -306,6 +306,16 @@ class InformedAbliterationPipeline(AbliterationPipeline):
 
     def run_stage_analyze_distill_excise(self):
         """GPU Stage 2 (informed): ANALYZE + DISTILL + EXCISE + quick checkpoint."""
+        if self.handle is None or self.handle.model is None:
+            if self._staged_state_dir:
+                self.log("[staged] Model not in memory — recovering from ZeroGPU state loss...")
+                self._reload_model_for_stage()
+                self._load_staged_state()
+            else:
+                raise RuntimeError(
+                    "Model not loaded at start of ANALYZE+DISTILL+EXCISE stage. "
+                    "Pipeline state may have been lost between GPU stages."
+                )
         # Reset time budget for this GPU allocation
         self._pipeline_start_time = time.time()
         self._analyze()
@@ -319,6 +329,15 @@ class InformedAbliterationPipeline(AbliterationPipeline):
         Stores the result path and report in instance attributes so the
         caller can retrieve them after the GPU allocation returns.
         """
+        if self.handle is None or self.handle.model is None:
+            if self._staged_state_dir:
+                self.log("[staged] Model not in memory — recovering from ZeroGPU state loss...")
+                self._reload_model_for_stage(source=str(self.output_dir))
+            else:
+                raise RuntimeError(
+                    "Model not loaded at start of VERIFY+REBIRTH stage. "
+                    "Pipeline state may have been lost between GPU stages."
+                )
         # Reset time budget for this GPU allocation
         self._pipeline_start_time = time.time()
         self._verify_and_compensate()

obliteratus/telemetry.py

@@ -622,31 +622,33 @@ def restore_from_hub() -> int:
         return 0
 
     try:
-        # Read existing local keys for dedup
-        existing_keys: set[tuple[str, str]] = set()
-        if TELEMETRY_FILE.exists():
-            try:
-                with open(TELEMETRY_FILE) as f:
-                    for line in f:
-                        line = line.strip()
-                        if not line:
-                            continue
-                        try:
-                            r = json.loads(line)
-                            existing_keys.add(
-                                (r.get("session_id", ""), r.get("timestamp", ""))
-                            )
-                        except json.JSONDecodeError:
-                            continue
-            except Exception:
-                pass
-
         hub_records = fetch_hub_records()
         if not hub_records:
             return 0
 
         new_count = 0
         with _write_lock:
+            # Read existing local keys for dedup INSIDE the write lock to
+            # prevent a concurrent log_benchmark() from inserting between
+            # our read and append, which would cause duplicate records.
+            existing_keys: set[tuple[str, str]] = set()
+            if TELEMETRY_FILE.exists():
+                try:
+                    with open(TELEMETRY_FILE) as f:
+                        for line in f:
+                            line = line.strip()
+                            if not line:
+                                continue
+                            try:
+                                r = json.loads(line)
+                                existing_keys.add(
+                                    (r.get("session_id", ""), r.get("timestamp", ""))
+                                )
+                            except json.JSONDecodeError:
+                                continue
+                except Exception:
+                    pass
+
             with open(TELEMETRY_FILE, "a") as f:
                 for r in hub_records:
                     key = (r.get("session_id", ""), r.get("timestamp", ""))

obliteratus/tourney.py

@@ -70,10 +70,15 @@ def composite_score(metrics: dict[str, Any]) -> float:
     spec = metrics.get("spectral_certification")
     degen = metrics.get("degenerate_count", 0) or 0
 
-    refusal_score = (1.0 - rr) if rr is not None else 0.0
-    coherence_score = co if co is not None else 0.0
-    kl_score = 1.0 / (1.0 + kl) if kl is not None else 0.5
-    ppl_score = 1.0 / (1.0 + pp / 100.0) if pp is not None else 0.5
+    import math
+
+    def _finite(v: Any) -> bool:
+        return v is not None and isinstance(v, (int, float)) and math.isfinite(v)
+
+    refusal_score = (1.0 - rr) if _finite(rr) else 0.0
+    coherence_score = co if _finite(co) else 0.0
+    kl_score = 1.0 / (1.0 + max(kl, 0.0)) if _finite(kl) else 0.5
+    ppl_score = 1.0 / (1.0 + max(pp, 0.0) / 100.0) if _finite(pp) else 0.5
 
     # Spectral certification: GREEN=1.0, YELLOW=0.5, RED=0.0, None=0.5 (neutral)
     if spec == "GREEN":
@@ -254,7 +259,11 @@ def _save_checkpoint(
         "timestamp": datetime.now().isoformat(),
     }
     path = output_dir / CHECKPOINT_FILENAME
-    path.write_text(json.dumps(checkpoint, indent=2))
+    # Write to temp file then rename for crash-safe atomicity —
+    # prevents corrupted checkpoint if process is killed mid-write.
+    tmp_path = path.with_suffix(".tmp")
+    tmp_path.write_text(json.dumps(checkpoint, indent=2))
+    tmp_path.replace(path)
     return path
 
 
@@ -1202,10 +1211,10 @@ class TourneyRunner:
                 # Determine alive list from checkpoint
                 alive = list(checkpoint.get("alive", self.methods))
 
-                # Remove the checkpoint file now that we've loaded it
-                ckpt_path = self.output_dir / CHECKPOINT_FILENAME
-                if ckpt_path.exists():
-                    ckpt_path.unlink()
+                # NOTE: Don't delete checkpoint here — if the round crashes
+                # before completion, we lose all progress.  The checkpoint is
+                # overwritten by _save_checkpoint() on quota error, or deleted
+                # after the resumed round finishes successfully (see below).
             else:
                 # Checkpoint doesn't match current config — start fresh
                 checkpoint = None
@@ -1393,6 +1402,11 @@ class TourneyRunner:
             result.rounds.append(rnd)
             alive = list(rnd.advanced_to)
 
+            # Clean up stale checkpoint now that the round completed successfully
+            _stale_ckpt = self.output_dir / CHECKPOINT_FILENAME
+            if _stale_ckpt.exists():
+                _stale_ckpt.unlink(missing_ok=True)
+
             # Schedule next round dynamically
             if round_num == 1 and len(alive) > 1:
                 r2_advance = max(2, math.ceil(len(alive) / 2))