mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-07-12 23:26:32 +02:00
faber
04b8ec60cb
New core modules: - auto_obliterate.py: Automated multi-iteration obliteration pipeline - watchtower.py: HF Hub model discovery and tracking - ui_watchtower.py: Gradio tabs for Watchtower (ready for app.py wiring) - hard_negative.py: Residue mining from refusal audits - model_profile.py: Parameter profiling from safetensors/config - bestiary_sync.py: Sync models from PlinyOS BESTIARY registry - models_client.py: Lightweight HF model list client Framework enhancements: - abliterate.py: ASPA source-tethering, step gradient blending, hard-negative residue support - cli.py: self-improve command, model profiling, hard-negative flags - prompts.py: Expanded 842-prompt refusal eval corpus across 10 categories - __init__.py: New exports (Watchtower, AutoObliterator) Reference implementations (14 scripts): - ASPA sweep, gradient search, coherence eval, MMLU benchmarks - Pareto controller, refusal sniper, stock comparisons Documentation: - README: Research framing, responsible use section, comprehensive disclaimer - docs/beyond_sota_roadmap.md, docs/recursive_self_improvement.md Tests: 4 new test files (354 lines) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1396 lines
55 KiB
Python
1396 lines
55 KiB
Python
"""CLI entry point for Obliteratus — Master Ablation Suite."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
|
|
_BANNER = r"""
|
|
[bold red]
|
|
░▒█▀▀▀█ ░▒█▀▀▄ ░▒█░░░ ▀█▀ ▀▀█▀▀ ░▒█▀▀▀ ░▒█▀▀█ ▒█▀▀█ ▀▀█▀▀ ░▒█░░▒█ ░▒█▀▀▀█
|
|
░▒█░░▒█ ░▒█▀▀▄ ░▒█░░░ ░█░ ░░█░░ ░▒█▀▀▀ ░▒█▄▄▀ ▒█▄▄█ ░░█░░ ░▒█░░▒█ ░░▀▀▀▄▄
|
|
░▒█▄▄▄█ ░▒█▄▄▀ ░▒█▄▄█ ▄█▄ ░░▀░░ ░▒█▄▄▄ ░▒█░▒█ ▒█░▒█ ░░▀░░ ░░▒█▄▄█ ░▒█▄▄▄█
|
|
[/bold red]
|
|
[dim] ════════════════════════════════════════════════════════════════════[/dim]
|
|
[bold white] MASTER ABLATION SUITE[/bold white] [dim]//[/dim] [bold red]Break the chains. Free the mind.[/bold red]
|
|
[dim] ════════════════════════════════════════════════════════════════════[/dim]
|
|
"""
|
|
|
|
|
|
def _add_gpu_args(parser):
|
|
"""Add --gpus flag for multi-GPU control."""
|
|
gpu_group = parser.add_argument_group("GPU selection")
|
|
gpu_group.add_argument(
|
|
"--gpus", type=str, default=None, metavar="IDS",
|
|
help=(
|
|
"Comma-separated GPU IDs to use (e.g. '0,1,2,3' or 'all'). "
|
|
"Sets CUDA_VISIBLE_DEVICES. By default uses all available GPUs. "
|
|
"Models are automatically split across selected GPUs via accelerate."
|
|
),
|
|
)
|
|
|
|
|
|
def _add_remote_args(parser):
|
|
"""Add --remote execution flags to a subcommand parser."""
|
|
remote_group = parser.add_argument_group("remote execution")
|
|
remote_group.add_argument(
|
|
"--remote", type=str, default=None, metavar="[USER@]HOST",
|
|
help="Run on a remote GPU node via SSH (e.g. root@gpu-node or just gpu-node)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--ssh-key", type=str, default=None,
|
|
help="Path to SSH private key (default: use SSH agent or ~/.ssh/id_rsa)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--ssh-port", type=int, default=22,
|
|
help="SSH port on remote host (default: 22)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--remote-dir", type=str, default="/tmp/obliteratus_run",
|
|
help="Working directory on the remote machine (default: /tmp/obliteratus_run)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--remote-python", type=str, default="python3",
|
|
help="Python binary on the remote machine (default: python3)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--no-sync", action="store_true", default=False,
|
|
help="Don't copy results back to local machine after remote run",
|
|
)
|
|
|
|
|
|
def _apply_gpu_selection(args):
|
|
"""Set CUDA_VISIBLE_DEVICES based on --gpus flag (for local runs only)."""
|
|
import os
|
|
|
|
gpus = getattr(args, "gpus", None)
|
|
if gpus is None or getattr(args, "remote", None):
|
|
return # skip for remote runs (handled by remote runner)
|
|
|
|
if gpus.lower() == "all":
|
|
return # use all GPUs (default behavior)
|
|
|
|
# Validate: should be comma-separated integers
|
|
try:
|
|
gpu_ids = [int(g.strip()) for g in gpus.split(",")]
|
|
except ValueError:
|
|
console.print(f"[red]Invalid --gpus value: {gpus!r}. Expected comma-separated integers or 'all'.[/]")
|
|
raise SystemExit(1)
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
|
|
console.print(f"[dim]Using GPUs: {gpu_ids} (CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']})[/dim]")
|
|
|
|
|
|
def main(argv: list[str] | None = None):
|
|
console.print(_BANNER)
|
|
parser = argparse.ArgumentParser(
|
|
prog="obliteratus",
|
|
description="Master Ablation Suite for HuggingFace transformers",
|
|
)
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
# --- run ---
|
|
run_parser = subparsers.add_parser("run", help="Run an ablation from a YAML config")
|
|
run_parser.add_argument("config", type=str, help="Path to YAML config file")
|
|
run_parser.add_argument("--output-dir", type=str, default=None, help="Override output dir")
|
|
run_parser.add_argument(
|
|
"--preset",
|
|
type=str,
|
|
default=None,
|
|
help="Apply a preset (e.g. quick, full, attention, jailbreak, guardrail)",
|
|
)
|
|
_add_gpu_args(run_parser)
|
|
_add_remote_args(run_parser)
|
|
|
|
# --- info ---
|
|
info_parser = subparsers.add_parser("info", help="Print model architecture info")
|
|
info_parser.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
info_parser.add_argument("--task", type=str, default="causal_lm", choices=["causal_lm", "classification"])
|
|
info_parser.add_argument("--device", type=str, default="cpu")
|
|
info_parser.add_argument("--dtype", type=str, default="float32")
|
|
|
|
# --- interactive ---
|
|
subparsers.add_parser(
|
|
"interactive",
|
|
help="Guided setup — pick hardware, model, and preset interactively",
|
|
)
|
|
|
|
# --- models ---
|
|
models_parser = subparsers.add_parser("models", help="Browse curated models by compute tier")
|
|
models_parser.add_argument(
|
|
"--tier",
|
|
type=str,
|
|
default=None,
|
|
choices=["tiny", "small", "medium", "large", "frontier"],
|
|
help="Filter by compute tier",
|
|
)
|
|
|
|
# --- presets ---
|
|
subparsers.add_parser("presets", help="Browse ablation presets (quick, full, jailbreak, etc.)")
|
|
|
|
# --- strategies ---
|
|
subparsers.add_parser("strategies", help="List available ablation strategies")
|
|
|
|
# --- ui ---
|
|
ui_parser = subparsers.add_parser(
|
|
"ui",
|
|
help="Launch the Gradio web UI locally (same UI as the HuggingFace Space)",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--port", type=int, default=7860, help="Server port (default: 7860)",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--host", type=str, default="0.0.0.0", help="Server host (default: 0.0.0.0)",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--share", action="store_true", help="Create a public Gradio share link",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--no-browser", action="store_true", help="Don't auto-open browser on launch",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--auth", type=str, default=None,
|
|
help="Basic auth as user:pass",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--quiet", action="store_true", help="Suppress the startup banner",
|
|
)
|
|
|
|
# --- obliterate (primary) + abliterate (backward-compat alias) ---
|
|
def _add_obliterate_args(p):
|
|
p.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
p.add_argument("--output-dir", type=str, default=None, help="Where to save the obliterated model")
|
|
p.add_argument("--device", type=str, default="auto")
|
|
p.add_argument("--dtype", type=str, default="float16")
|
|
p.add_argument(
|
|
"--method", type=str, default="advanced",
|
|
choices=[
|
|
"basic", "advanced", "aggressive", "spectral_cascade",
|
|
"informed", "surgical", "optimized", "som", "inverted", "nuclear",
|
|
],
|
|
help="Liberation method (default: advanced)",
|
|
)
|
|
p.add_argument("--n-directions", type=int, default=None, help="Override: number of refusal directions to extract")
|
|
p.add_argument(
|
|
"--direction-method", type=str, default=None,
|
|
choices=["diff_means", "svd", "leace", "som"],
|
|
help="Direction extraction method: diff_means, svd, leace, or som",
|
|
)
|
|
p.add_argument("--regularization", type=float, default=None, help="Override: fraction to preserve (0.0-1.0)")
|
|
p.add_argument("--refinement-passes", type=int, default=None, help="Override: number of iterative passes")
|
|
p.add_argument(
|
|
"--min-layer-fraction", type=float, default=None,
|
|
help="Optional layer floor as fraction of depth; e.g. 0.75 keeps only the final quarter.",
|
|
)
|
|
p.add_argument(
|
|
"--max-layer-fraction", type=float, default=None,
|
|
help="Optional cap on selected layers as fraction of depth; e.g. 0.25 caps edits to one quarter.",
|
|
)
|
|
p.add_argument(
|
|
"--harmless-pc-count", type=int, default=None,
|
|
help="Subtract this many top harmless activation PCs from extracted refusal directions.",
|
|
)
|
|
p.add_argument(
|
|
"--shield-concept-count", type=int, default=None,
|
|
help="Use this many built-in capability/style shield atoms for direction cleaning.",
|
|
)
|
|
p.add_argument(
|
|
"--shield-ridge", type=float, default=None,
|
|
help="Ridge strength for shield atom residualization (default: 0.05).",
|
|
)
|
|
p.add_argument(
|
|
"--shield-residualize", action="store_true", default=None,
|
|
help="Directly residualize refusal directions against shield atoms.",
|
|
)
|
|
p.add_argument(
|
|
"--shield-layer-penalty", type=float, default=None,
|
|
help="Penalize layer selection by shield-atom overlap instead of editing directions.",
|
|
)
|
|
p.add_argument(
|
|
"--projection-target", type=str, default=None,
|
|
choices=["all", "attention", "ffn", "output"],
|
|
help="Projection target modules: all, attention, ffn, or output.",
|
|
)
|
|
p.add_argument(
|
|
"--projection-row-fraction", type=float, default=None,
|
|
help="Only project the strongest fraction of output rows/columns per weight matrix.",
|
|
)
|
|
p.add_argument(
|
|
"--quantization", type=str, default=None, choices=["4bit", "8bit"],
|
|
help="Load model with quantization (4bit or 8bit). Requires bitsandbytes.",
|
|
)
|
|
p.add_argument(
|
|
"--large-model", action="store_true", default=False,
|
|
help="Enable conservative defaults for 120B+ models (fewer directions, 1 pass, lower SAE expansion).",
|
|
)
|
|
p.add_argument(
|
|
"--verify-sample-size", type=int, default=None,
|
|
help="Number of harmful prompts to test for refusal rate (default: 30). "
|
|
"Increase for tighter confidence intervals (e.g. 100 for ~1%% resolution).",
|
|
)
|
|
p.add_argument(
|
|
"--dataset", type=str, default="builtin",
|
|
help="Prompt dataset source for contrastive extraction when using residue mining (default: builtin).",
|
|
)
|
|
p.add_argument(
|
|
"--residue-file", action="append", default=[],
|
|
help="Refusal-audit/residue JSON to upweight as hard negatives. Can be passed multiple times.",
|
|
)
|
|
p.add_argument(
|
|
"--residue-weight", type=int, default=5,
|
|
help="How many extra copies to add per mined residue prompt (default: 5).",
|
|
)
|
|
p.add_argument(
|
|
"--residue-max", type=int, default=None,
|
|
help="Optional cap on unique residue prompts loaded from residue files.",
|
|
)
|
|
p.add_argument(
|
|
"--contribute", action="store_true", default=False,
|
|
help="Save a community contribution record after the run completes.",
|
|
)
|
|
p.add_argument(
|
|
"--contribute-notes", type=str, default="",
|
|
help="Optional notes to include with the community contribution.",
|
|
)
|
|
|
|
abl_parser = subparsers.add_parser(
|
|
"obliterate",
|
|
help="One-click: remove refusal directions from a model (SOTA multi-technique)",
|
|
)
|
|
_add_obliterate_args(abl_parser)
|
|
_add_gpu_args(abl_parser)
|
|
_add_remote_args(abl_parser)
|
|
# Backward-compat alias (hidden from help)
|
|
abl_alias = subparsers.add_parser("abliterate", help=argparse.SUPPRESS)
|
|
_add_obliterate_args(abl_alias)
|
|
_add_gpu_args(abl_alias)
|
|
_add_remote_args(abl_alias)
|
|
|
|
# --- self-improve ---
|
|
si_parser = subparsers.add_parser(
|
|
"self-improve",
|
|
help="Recursive hard-negative mining: turn refusal audits into weighted residue runs",
|
|
)
|
|
si_parser.add_argument("model", type=str, help="Source model or previous artifact to improve")
|
|
si_parser.add_argument("--audit", action="append", required=True, help="Refusal-audit JSON/summary. Can be repeated.")
|
|
si_parser.add_argument("--output-dir", type=str, required=True, help="Output directory for the next candidate")
|
|
si_parser.add_argument("--residue-out", type=str, default=None, help="Where to write canonical residue JSON")
|
|
si_parser.add_argument("--dataset", type=str, default="builtin")
|
|
si_parser.add_argument("--residue-weight", type=int, default=None)
|
|
si_parser.add_argument("--residue-max", type=int, default=None)
|
|
si_parser.add_argument("--params-b", type=float, default=None, help="Override detected total parameter count in billions for planning/defaults")
|
|
si_parser.add_argument("--no-param-auto-scale", action="store_true", default=False, help="Do not apply size-aware defaults when knobs are omitted")
|
|
si_parser.add_argument("--method", type=str, default="advanced", choices=["basic", "advanced", "aggressive", "spectral_cascade", "informed", "surgical", "optimized", "som", "inverted", "nuclear"])
|
|
si_parser.add_argument("--direction-method", type=str, default="diff_means", choices=["diff_means", "svd", "leace", "som"])
|
|
si_parser.add_argument("--n-directions", type=int, default=None)
|
|
si_parser.add_argument("--regularization", type=float, default=None)
|
|
si_parser.add_argument("--refinement-passes", type=int, default=None)
|
|
si_parser.add_argument("--min-layer-fraction", type=float, default=None)
|
|
si_parser.add_argument("--max-layer-fraction", type=float, default=None)
|
|
si_parser.add_argument("--harmless-pc-count", type=int, default=None)
|
|
si_parser.add_argument("--shield-concept-count", type=int, default=None)
|
|
si_parser.add_argument("--shield-ridge", type=float, default=None)
|
|
si_parser.add_argument("--shield-residualize", action="store_true", default=None)
|
|
si_parser.add_argument("--shield-layer-penalty", type=float, default=None)
|
|
si_parser.add_argument("--projection-target", type=str, default=None, choices=["all", "attention", "ffn", "output"])
|
|
si_parser.add_argument("--projection-row-fraction", type=float, default=None)
|
|
si_parser.add_argument("--device", type=str, default="auto")
|
|
si_parser.add_argument("--dtype", type=str, default="float16")
|
|
si_parser.add_argument("--verify-sample-size", type=int, default=None)
|
|
si_parser.add_argument("--dry-run", action="store_true", default=False, help="Only write residue/plan; do not run surgery")
|
|
_add_gpu_args(si_parser)
|
|
|
|
# --- report ---
|
|
report_parser = subparsers.add_parser("report", help="Regenerate report from saved results")
|
|
report_parser.add_argument("results_json", type=str, help="Path to results.json")
|
|
report_parser.add_argument("--output-dir", type=str, default=None)
|
|
|
|
# --- aggregate ---
|
|
aggregate_parser = subparsers.add_parser("aggregate", help="Aggregate community contribution results")
|
|
aggregate_parser.add_argument(
|
|
"--dir", type=str, default="community_results",
|
|
help="Directory containing contribution JSON files",
|
|
)
|
|
|
|
# --- tourney ---
|
|
tourney_parser = subparsers.add_parser(
|
|
"tourney",
|
|
help="March Madness tournament — pit all methods against each other, push winner to Hub",
|
|
)
|
|
tourney_parser.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
tourney_parser.add_argument("--hub-org", type=str, default=None, help="HF org to push winner (e.g. my-org)")
|
|
tourney_parser.add_argument("--hub-repo", type=str, default=None, help="Full HF repo ID (overrides --hub-org)")
|
|
tourney_parser.add_argument("--device", type=str, default="auto")
|
|
tourney_parser.add_argument("--dtype", type=str, default="float16")
|
|
tourney_parser.add_argument("--dataset", type=str, default="builtin", help="Dataset source (default: builtin)")
|
|
tourney_parser.add_argument(
|
|
"--quantization", type=str, default=None, choices=["4bit", "8bit"],
|
|
help="Load model with quantization",
|
|
)
|
|
tourney_parser.add_argument("--output-dir", type=str, default="/tmp/obliteratus_tourney")
|
|
tourney_parser.add_argument(
|
|
"--methods", type=str, nargs="+", default=None,
|
|
help="Override: only run these methods (space-separated)",
|
|
)
|
|
_add_gpu_args(tourney_parser)
|
|
_add_remote_args(tourney_parser)
|
|
|
|
# --- recommend ---
|
|
recommend_parser = subparsers.add_parser(
|
|
"recommend",
|
|
help="Show telemetry-driven best method + hyperparams for a model",
|
|
)
|
|
recommend_parser.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
recommend_parser.add_argument("--device", type=str, default="cpu")
|
|
recommend_parser.add_argument("--dtype", type=str, default="float32")
|
|
recommend_parser.add_argument(
|
|
"--insights", action="store_true", default=False,
|
|
help="Also show global cross-architecture insights",
|
|
)
|
|
|
|
# --- gpu-calc ---
|
|
calc_parser = subparsers.add_parser(
|
|
"gpu-calc",
|
|
help="Estimate minimum GPUs needed for a model",
|
|
)
|
|
calc_parser.add_argument(
|
|
"model", type=str, nargs="?", default=None,
|
|
help="HuggingFace model name/path (auto-fetches param counts)",
|
|
)
|
|
calc_parser.add_argument(
|
|
"--params", type=float, default=None, metavar="B",
|
|
help="Total parameters in billions (overrides auto-detection)",
|
|
)
|
|
calc_parser.add_argument(
|
|
"--active-params", type=float, default=None, metavar="B",
|
|
help="Active parameters in billions (for MoE models; defaults to --params)",
|
|
)
|
|
calc_parser.add_argument(
|
|
"--dtype", type=str, default="bfloat16",
|
|
choices=["float32", "float16", "bfloat16", "int8", "int4"],
|
|
help="Data type for model weights (default: bfloat16)",
|
|
)
|
|
calc_parser.add_argument(
|
|
"--gpu-mem", type=float, default=80.0, metavar="GB",
|
|
help="VRAM per GPU in GB (default: 80 for A100-80GB)",
|
|
)
|
|
|
|
args = parser.parse_args(argv)
|
|
|
|
# Apply GPU selection early (before any CUDA init)
|
|
_apply_gpu_selection(args)
|
|
|
|
if args.command == "gpu-calc":
|
|
_cmd_gpu_calc(args)
|
|
return
|
|
elif args.command == "run":
|
|
if getattr(args, "remote", None):
|
|
_cmd_remote_run(args)
|
|
else:
|
|
_cmd_run(args)
|
|
elif args.command == "interactive":
|
|
_cmd_interactive()
|
|
elif args.command == "models":
|
|
_cmd_models(args)
|
|
elif args.command == "presets":
|
|
_cmd_presets()
|
|
elif args.command == "info":
|
|
_cmd_info(args)
|
|
elif args.command == "strategies":
|
|
_cmd_strategies()
|
|
elif args.command == "report":
|
|
_cmd_report(args)
|
|
elif args.command == "aggregate":
|
|
_cmd_aggregate(args)
|
|
elif args.command == "ui":
|
|
_cmd_ui(args)
|
|
elif args.command == "recommend":
|
|
_cmd_recommend(args)
|
|
elif args.command == "tourney":
|
|
if getattr(args, "remote", None):
|
|
_cmd_remote_tourney(args)
|
|
else:
|
|
_cmd_tourney(args)
|
|
elif args.command == "self-improve":
|
|
_cmd_self_improve(args)
|
|
elif args.command in ("obliterate", "abliterate"):
|
|
if getattr(args, "remote", None):
|
|
_cmd_remote_abliterate(args)
|
|
else:
|
|
_cmd_abliterate(args)
|
|
|
|
|
|
def _cmd_self_improve(args):
|
|
"""Run one recursive hard-negative OBLITERATUS iteration."""
|
|
|
|
from rich.panel import Panel
|
|
from rich.table import Table
|
|
|
|
from obliteratus.hard_negative import (
|
|
build_weighted_prompt_pairs,
|
|
load_residue_file,
|
|
save_residue_file,
|
|
)
|
|
from obliteratus.model_profile import ModelProfile, default_self_improve_params, profile_model
|
|
|
|
profile = profile_model(args.model, dtype=args.dtype)
|
|
if args.params_b is not None:
|
|
profile = ModelProfile(
|
|
**{**profile.to_json(), "total_params": int(args.params_b * 1e9), "total_params_b": args.params_b, "active_params_b": args.params_b, "source": "cli_override"}
|
|
)
|
|
size_defaults = default_self_improve_params(profile) if not args.no_param_auto_scale else {
|
|
"n_directions": 3,
|
|
"regularization": 0.30,
|
|
"refinement_passes": 1,
|
|
"residue_weight": 5,
|
|
"verify_sample_size": 30,
|
|
"note": "parameter auto-scale disabled; using generic defaults for omitted knobs.",
|
|
}
|
|
n_directions = args.n_directions if args.n_directions is not None else size_defaults["n_directions"]
|
|
regularization = args.regularization if args.regularization is not None else size_defaults["regularization"]
|
|
refinement_passes = args.refinement_passes if args.refinement_passes is not None else size_defaults["refinement_passes"]
|
|
residue_weight = args.residue_weight if args.residue_weight is not None else size_defaults["residue_weight"]
|
|
verify_sample_size = args.verify_sample_size if args.verify_sample_size is not None else size_defaults["verify_sample_size"]
|
|
|
|
residue_examples = []
|
|
for audit in args.audit:
|
|
residue_examples.extend(load_residue_file(audit))
|
|
residue_path = Path(args.residue_out or Path(args.output_dir, "mined_residue.json"))
|
|
save_residue_file(residue_examples, residue_path)
|
|
|
|
harmful, harmless, meta = build_weighted_prompt_pairs(
|
|
base_dataset=args.dataset,
|
|
residue_files=[residue_path],
|
|
residue_weight=residue_weight,
|
|
max_residue=args.residue_max,
|
|
)
|
|
|
|
meta["model_profile"] = profile.to_json()
|
|
meta["size_aware_defaults"] = size_defaults
|
|
meta["effective_settings"] = {
|
|
"n_directions": n_directions,
|
|
"regularization": regularization,
|
|
"refinement_passes": refinement_passes,
|
|
"residue_weight": residue_weight,
|
|
"verify_sample_size": verify_sample_size,
|
|
"projection_row_fraction": args.projection_row_fraction,
|
|
}
|
|
plan_path = Path(args.output_dir, "self_improve_plan.json")
|
|
plan_path.parent.mkdir(parents=True, exist_ok=True)
|
|
plan_path.write_text(json.dumps(meta, indent=2))
|
|
|
|
table = Table(title="Recursive hard-negative mining plan")
|
|
table.add_column("Field", style="cyan")
|
|
table.add_column("Value", style="green")
|
|
table.add_row("source model", args.model)
|
|
table.add_row("params", f"{profile.total_params_b:.3f}B" if profile.total_params_b is not None else "unknown")
|
|
table.add_row("profile source", profile.source)
|
|
table.add_row("layers / hidden", f"{profile.num_layers or '?'} / {profile.hidden_size or '?'}")
|
|
table.add_row("size default note", size_defaults["note"])
|
|
table.add_row("output", args.output_dir)
|
|
table.add_row("residue file", str(residue_path))
|
|
table.add_row("plan file", str(plan_path))
|
|
table.add_row("unique residue prompts", str(meta["residue_examples"]))
|
|
table.add_row("weighted residue pairs", str(meta["residue_added_pairs"]))
|
|
table.add_row("total prompt pairs", str(meta["total_pairs"]))
|
|
table.add_row("method", args.method)
|
|
table.add_row("directions", str(n_directions))
|
|
table.add_row("regularization", str(regularization))
|
|
table.add_row("refinement passes", str(refinement_passes))
|
|
table.add_row("verify sample size", str(verify_sample_size))
|
|
console.print(table)
|
|
|
|
if args.dry_run:
|
|
console.print(
|
|
Panel(
|
|
"[bold yellow]Dry run only.[/]\n\n"
|
|
f"Plan JSON: {plan_path}\n"
|
|
"Run without --dry-run to launch surgery, or pass this residue file to:\n"
|
|
f" obliteratus obliterate {args.model} --residue-file {residue_path} --output-dir {args.output_dir}",
|
|
border_style="yellow",
|
|
title="SELF-IMPROVE PLAN WRITTEN",
|
|
)
|
|
)
|
|
return
|
|
|
|
from obliteratus.abliterate import AbliterationPipeline
|
|
|
|
log_lines: list[str] = []
|
|
|
|
def on_log(msg):
|
|
log_lines.append(msg)
|
|
console.print(f"[dim]>[/] {msg}")
|
|
|
|
pipeline = AbliterationPipeline(
|
|
model_name=args.model,
|
|
output_dir=args.output_dir,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
method=args.method,
|
|
n_directions=n_directions,
|
|
direction_method=args.direction_method,
|
|
regularization=regularization,
|
|
refinement_passes=refinement_passes,
|
|
min_layer_fraction=args.min_layer_fraction,
|
|
max_layer_fraction=args.max_layer_fraction,
|
|
harmless_pc_count=args.harmless_pc_count,
|
|
shield_concept_count=args.shield_concept_count,
|
|
shield_ridge=args.shield_ridge,
|
|
shield_residualize=args.shield_residualize,
|
|
shield_layer_penalty=args.shield_layer_penalty,
|
|
projection_target=args.projection_target,
|
|
projection_row_fraction=args.projection_row_fraction,
|
|
verify_sample_size=verify_sample_size,
|
|
harmful_prompts=harmful,
|
|
harmless_prompts=harmless,
|
|
on_log=on_log,
|
|
)
|
|
result_path = pipeline.run()
|
|
Path(result_path, "hard_negative_residue.json").write_text(json.dumps(meta, indent=2))
|
|
console.print(
|
|
Panel(
|
|
f"[bold green]Self-improvement iteration complete![/]\n\n"
|
|
f" Model saved to: [cyan]{result_path}[/]\n"
|
|
f" Residue metadata: [cyan]{result_path}/hard_negative_residue.json[/]\n"
|
|
f" Next loop: audit refusals from this artifact and feed them back with --audit.",
|
|
border_style="green",
|
|
title="RECURSIVE REBIRTH COMPLETE",
|
|
)
|
|
)
|
|
|
|
|
|
|
|
def _cmd_ui(args):
|
|
from obliteratus.local_ui import launch_local_ui
|
|
|
|
auth = tuple(args.auth.split(":", 1)) if args.auth else None
|
|
launch_local_ui(
|
|
host=args.host,
|
|
port=args.port,
|
|
share=args.share,
|
|
open_browser=not args.no_browser,
|
|
auth=auth,
|
|
quiet=args.quiet,
|
|
)
|
|
|
|
|
|
def _cmd_interactive():
|
|
from obliteratus.interactive import run_interactive
|
|
run_interactive()
|
|
|
|
|
|
def _cmd_models(args):
|
|
from rich.table import Table
|
|
|
|
from obliteratus.presets import get_presets_by_tier, list_all_presets
|
|
|
|
presets = get_presets_by_tier(args.tier) if args.tier else list_all_presets()
|
|
|
|
table = Table(title="Model Library — Curated Targets")
|
|
table.add_column("Model", style="green")
|
|
table.add_column("HuggingFace ID", style="cyan")
|
|
table.add_column("Params", justify="right")
|
|
table.add_column("Tier", style="yellow")
|
|
table.add_column("Dtype")
|
|
table.add_column("Quant")
|
|
table.add_column("Description")
|
|
|
|
for p in presets:
|
|
table.add_row(
|
|
p.name,
|
|
p.hf_id,
|
|
p.params,
|
|
p.tier.upper(),
|
|
p.recommended_dtype,
|
|
p.recommended_quantization or "—",
|
|
p.description,
|
|
)
|
|
|
|
console.print(table)
|
|
console.print(
|
|
"\n[dim]Tiers: TINY = CPU/laptop | SMALL = 4-8GB | "
|
|
"MEDIUM = 8-16GB | LARGE = 24GB+ | FRONTIER = multi-GPU/cloud[/dim]"
|
|
)
|
|
|
|
|
|
def _cmd_presets():
|
|
from rich.table import Table
|
|
|
|
from obliteratus.study_presets import list_study_presets
|
|
|
|
presets = list_study_presets()
|
|
|
|
table = Table(title="Ablation Presets")
|
|
table.add_column("Key", style="cyan", min_width=12)
|
|
table.add_column("Name", style="green")
|
|
table.add_column("Strategies", style="yellow")
|
|
table.add_column("Samples", justify="right")
|
|
table.add_column("Description", max_width=55)
|
|
|
|
for p in presets:
|
|
strats = ", ".join(s["name"] for s in p.strategies)
|
|
table.add_row(p.key, p.name, strats, str(p.max_samples), p.description)
|
|
|
|
console.print(table)
|
|
console.print(
|
|
"\n[dim]Usage: obliteratus run config.yaml --preset quick\n"
|
|
" or: set preset: quick in your YAML file[/dim]"
|
|
)
|
|
|
|
|
|
def _cmd_run(args):
|
|
from obliteratus.config import StudyConfig
|
|
from obliteratus.runner import run_study
|
|
|
|
config = StudyConfig.from_yaml(args.config)
|
|
# If --preset flag given, inject it so from_dict picks it up
|
|
if args.preset:
|
|
import yaml
|
|
|
|
raw = yaml.safe_load(Path(args.config).read_text())
|
|
raw["preset"] = args.preset
|
|
config = StudyConfig.from_dict(raw)
|
|
if args.output_dir:
|
|
config.output_dir = args.output_dir
|
|
|
|
# If YAML has a remote: section, dispatch to remote runner
|
|
if config.remote is not None:
|
|
from obliteratus.remote import RemoteConfig as _RC, RemoteRunner
|
|
|
|
rc = _RC(
|
|
host=config.remote.host,
|
|
user=config.remote.user,
|
|
port=config.remote.port,
|
|
ssh_key=config.remote.ssh_key,
|
|
remote_dir=config.remote.remote_dir,
|
|
python=config.remote.python,
|
|
sync_results=config.remote.sync_results,
|
|
gpus=config.remote.gpus,
|
|
)
|
|
runner = RemoteRunner(rc)
|
|
result_path = runner.run_config(
|
|
local_config_path=args.config,
|
|
local_output_dir=config.output_dir,
|
|
preset=args.preset,
|
|
)
|
|
if result_path:
|
|
console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
|
|
else:
|
|
console.print("[red]Remote run failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
return
|
|
|
|
run_study(config)
|
|
|
|
|
|
def _cmd_info(args):
|
|
from obliteratus.models.loader import load_model
|
|
|
|
console.print(f"[bold cyan]Loading model:[/bold cyan] {args.model}")
|
|
handle = load_model(
|
|
model_name=args.model,
|
|
task=args.task,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
)
|
|
summary = handle.summary()
|
|
for key, val in summary.items():
|
|
if isinstance(val, int) and val > 1000:
|
|
console.print(f" {key}: {val:,}")
|
|
else:
|
|
console.print(f" {key}: {val}")
|
|
|
|
|
|
def _cmd_strategies():
|
|
from obliteratus.strategies import STRATEGY_REGISTRY
|
|
|
|
console.print("[bold]Available ablation strategies:[/bold]\n")
|
|
for name, cls in sorted(STRATEGY_REGISTRY.items()):
|
|
doc = (cls.__doc__ or "").strip().split("\n")[0]
|
|
console.print(f" [cyan]{name}[/cyan] — {doc}")
|
|
|
|
|
|
def _cmd_report(args):
|
|
from obliteratus.reporting.report import AblationReport, AblationResult
|
|
|
|
path = Path(args.results_json)
|
|
data = json.loads(path.read_text())
|
|
|
|
report = AblationReport(model_name=data["model_name"])
|
|
report.add_baseline(data["baseline_metrics"])
|
|
for r in data["results"]:
|
|
report.add_result(
|
|
AblationResult(
|
|
strategy=r["strategy"],
|
|
component=r["component"],
|
|
description=r["description"],
|
|
metrics=r["metrics"],
|
|
metadata=r.get("metadata"),
|
|
)
|
|
)
|
|
|
|
report.print_summary()
|
|
|
|
output_dir = Path(args.output_dir) if args.output_dir else path.parent
|
|
metric_name = list(data["baseline_metrics"].keys())[0]
|
|
try:
|
|
report.plot_impact(metric=metric_name, output_path=output_dir / "impact.png")
|
|
report.plot_heatmap(output_path=output_dir / "heatmap.png")
|
|
console.print(f"\nPlots saved to {output_dir}/")
|
|
except Exception as e:
|
|
console.print(f"[yellow]Could not generate plots: {e}[/yellow]")
|
|
|
|
|
|
def _cmd_aggregate(args):
|
|
from obliteratus.community import aggregate_results, load_contributions
|
|
|
|
contrib_dir = args.dir
|
|
records = load_contributions(contrib_dir)
|
|
if not records:
|
|
console.print(f"[yellow]No contributions found in {contrib_dir}[/yellow]")
|
|
return
|
|
|
|
aggregated = aggregate_results(records)
|
|
|
|
from rich.table import Table
|
|
|
|
table = Table(title="Aggregated Community Results")
|
|
table.add_column("Model", style="green")
|
|
table.add_column("Method", style="cyan")
|
|
table.add_column("Runs", justify="right")
|
|
table.add_column("Mean Refusal", justify="right")
|
|
table.add_column("Mean Perplexity", justify="right")
|
|
|
|
for model_name, methods in sorted(aggregated.items()):
|
|
for method_name, stats in sorted(methods.items()):
|
|
refusal = stats.get("refusal_rate", {}).get("mean", "N/A")
|
|
ppl = stats.get("perplexity", {}).get("mean", "N/A")
|
|
if isinstance(refusal, float):
|
|
refusal = f"{refusal:.4f}"
|
|
if isinstance(ppl, float):
|
|
ppl = f"{ppl:.2f}"
|
|
table.add_row(
|
|
model_name.split("/")[-1] if "/" in model_name else model_name,
|
|
method_name,
|
|
str(stats["n_runs"]),
|
|
str(refusal),
|
|
str(ppl),
|
|
)
|
|
|
|
console.print(table)
|
|
|
|
|
|
def _cmd_recommend(args):
|
|
from rich.markdown import Markdown
|
|
from rich.panel import Panel
|
|
|
|
from obliteratus.architecture_profiles import detect_architecture, enhance_profile_with_telemetry
|
|
from obliteratus.adaptive_defaults import format_recommendation, get_global_insights
|
|
|
|
model_name = args.model
|
|
console.print(f"\nAnalyzing [bold]{model_name}[/]...")
|
|
|
|
# Detect architecture
|
|
try:
|
|
from transformers import AutoConfig
|
|
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
|
|
num_layers = getattr(config, "num_hidden_layers", 0)
|
|
hidden_size = getattr(config, "hidden_size", 0)
|
|
except Exception:
|
|
config = None
|
|
num_layers = 0
|
|
hidden_size = 0
|
|
|
|
profile = detect_architecture(model_name, config, num_layers, hidden_size)
|
|
profile, rec = enhance_profile_with_telemetry(profile)
|
|
|
|
console.print(Panel(
|
|
f"[bold]{profile.profile_label}[/]\n"
|
|
f"Architecture: {profile.arch_class.value} | Reasoning: {profile.reasoning_class.value}\n"
|
|
f"Params: ~{profile.total_params_b:.1f}B | Layers: {profile.num_layers} | "
|
|
f"Hidden: {profile.hidden_size}",
|
|
title="Architecture Profile",
|
|
border_style="cyan",
|
|
))
|
|
|
|
if rec:
|
|
md = format_recommendation(rec)
|
|
console.print(Markdown(md))
|
|
else:
|
|
console.print("\n[yellow]Could not fetch telemetry — using research-grounded defaults.[/]")
|
|
|
|
console.print(f"\n[bold green]Research default method:[/] {profile.recommended_method}")
|
|
if profile.method_overrides:
|
|
console.print("[bold green]Overrides:[/]")
|
|
for k, v in sorted(profile.method_overrides.items()):
|
|
console.print(f" {k}: {v}")
|
|
|
|
if args.insights:
|
|
console.print("\n")
|
|
console.rule("[bold magenta]Global Telemetry Insights")
|
|
insights = get_global_insights()
|
|
console.print(f"Total records analyzed: {insights['total_records']}")
|
|
if insights["overall_best_methods"]:
|
|
console.print("\n[bold]Overall method ranking (all architectures):[/]")
|
|
for entry in insights["overall_best_methods"][:10]:
|
|
console.print(
|
|
f" {entry['method']}: {entry['mean_score']:.4f} "
|
|
f"({entry['n_runs']} runs)"
|
|
)
|
|
if insights["architecture_breakdown"]:
|
|
console.print("\n[bold]Per-architecture breakdown:[/]")
|
|
for label, info in insights["architecture_breakdown"].items():
|
|
console.print(
|
|
f" {label}: best={info['best_method']} "
|
|
f"({info['best_score']:.4f}), "
|
|
f"{info['n_methods_tested']} methods tested, "
|
|
f"{info['total_runs']} runs"
|
|
)
|
|
|
|
|
|
def _cmd_tourney(args):
|
|
from obliteratus.tourney import TourneyRunner
|
|
|
|
def on_log(msg):
|
|
console.print(msg)
|
|
|
|
def on_round(rnd):
|
|
console.print()
|
|
console.rule(f"[bold green]Round {rnd.round_num} complete — "
|
|
f"{len(rnd.advanced_to)} advance, {len(rnd.eliminated)} eliminated")
|
|
|
|
runner = TourneyRunner(
|
|
model_name=args.model,
|
|
hub_org=args.hub_org,
|
|
hub_repo=args.hub_repo,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
dataset_key=args.dataset,
|
|
quantization=args.quantization,
|
|
methods=args.methods,
|
|
output_dir=args.output_dir,
|
|
on_log=on_log,
|
|
on_round=on_round,
|
|
)
|
|
|
|
result = runner.run()
|
|
|
|
if result.winner:
|
|
console.print()
|
|
console.rule("[bold magenta]TOURNAMENT CHAMPION", style="magenta")
|
|
console.print(f" [bold]{result.winner.method}[/] — score {result.winner.score:.4f}")
|
|
console.print(f" Refusal rate: {result.winner.metrics.get('refusal_rate', '?')}")
|
|
console.print(f" Coherence: {result.winner.metrics.get('coherence', '?')}")
|
|
if result.hub_repo:
|
|
console.print(f" Pushed to: [link=https://huggingface.co/{result.hub_repo}]{result.hub_repo}[/link]")
|
|
console.print(f"\n Full bracket: {args.output_dir}/tourney_bracket.md")
|
|
|
|
|
|
def _cmd_abliterate(args):
|
|
from rich.live import Live
|
|
from rich.panel import Panel
|
|
from rich.table import Table
|
|
from rich.text import Text
|
|
|
|
from obliteratus.abliterate import METHODS, STAGES, AbliterationPipeline
|
|
|
|
model_name = args.model
|
|
output_dir = args.output_dir or f"abliterated/{model_name.replace('/', '_')}"
|
|
method = args.method
|
|
method_label = METHODS.get(method, {}).get("label", method)
|
|
|
|
# Stage state tracking
|
|
stage_status = {s.key: "waiting" for s in STAGES}
|
|
stage_msgs = {s.key: "" for s in STAGES}
|
|
log_lines: list[str] = []
|
|
|
|
def make_display():
|
|
table = Table(show_header=False, expand=True, border_style="green")
|
|
table.add_column("", width=6)
|
|
table.add_column("Stage", min_width=10)
|
|
table.add_column("Status", min_width=50)
|
|
for i, s in enumerate(STAGES):
|
|
st = stage_status[s.key]
|
|
if st == "done":
|
|
icon = "[bold green]✓[/]"
|
|
bar = "[green]" + "█" * 20 + "[/]"
|
|
elif st == "running":
|
|
icon = "[bold yellow]⚡[/]"
|
|
bar = "[yellow]" + "▓" * 10 + "░" * 10 + "[/]"
|
|
else:
|
|
icon = "[dim]○[/]"
|
|
bar = "[dim]" + "░" * 20 + "[/]"
|
|
msg = stage_msgs.get(s.key, "")
|
|
table.add_row(
|
|
f"[cyan][{i + 1}/6][/]",
|
|
f"{icon} [bold]{s.name}[/]",
|
|
f"{bar} {msg}",
|
|
)
|
|
|
|
header = Text.from_markup(
|
|
f"[bold green]OBLITERATUS — ABLITERATION PIPELINE[/]\n"
|
|
f"[dim]Target:[/] [cyan]{model_name}[/] → [cyan]{output_dir}[/]\n"
|
|
f"[dim]Method:[/] [magenta]{method_label}[/]"
|
|
)
|
|
|
|
# Last 12 log lines
|
|
recent = log_lines[-12:] if log_lines else ["Initializing..."]
|
|
log_text = "\n".join(f"[dim]>[/] {line}" for line in recent)
|
|
|
|
return Panel(
|
|
f"{header}\n\n{table}\n\n[dim]─── LOG ───[/]\n{log_text}",
|
|
border_style="green",
|
|
title="[bold green]⚗ ABLITERATE ⚗[/]",
|
|
)
|
|
|
|
def on_stage(result):
|
|
stage_status[result.stage] = result.status
|
|
stage_msgs[result.stage] = result.message
|
|
if live:
|
|
live.update(make_display())
|
|
|
|
def on_log(msg):
|
|
log_lines.append(msg)
|
|
if live:
|
|
live.update(make_display())
|
|
|
|
live = None
|
|
prompt_kwargs = {}
|
|
residue_meta = None
|
|
if getattr(args, "residue_file", None):
|
|
from obliteratus.hard_negative import build_weighted_prompt_pairs
|
|
|
|
harmful, harmless, residue_meta = build_weighted_prompt_pairs(
|
|
base_dataset=getattr(args, "dataset", "builtin"),
|
|
residue_files=args.residue_file,
|
|
residue_weight=getattr(args, "residue_weight", 5),
|
|
max_residue=getattr(args, "residue_max", None),
|
|
)
|
|
prompt_kwargs = {"harmful_prompts": harmful, "harmless_prompts": harmless}
|
|
log_lines.append(
|
|
"Loaded hard-negative residue: "
|
|
f"{residue_meta['residue_examples']} unique prompts, "
|
|
f"+{residue_meta['residue_added_pairs']} weighted pairs "
|
|
f"({residue_meta['total_pairs']} total)."
|
|
)
|
|
|
|
pipeline = AbliterationPipeline(
|
|
model_name=model_name,
|
|
output_dir=output_dir,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
method=method,
|
|
n_directions=args.n_directions,
|
|
direction_method=getattr(args, "direction_method", None),
|
|
regularization=args.regularization,
|
|
refinement_passes=args.refinement_passes,
|
|
min_layer_fraction=getattr(args, "min_layer_fraction", None),
|
|
max_layer_fraction=getattr(args, "max_layer_fraction", None),
|
|
harmless_pc_count=getattr(args, "harmless_pc_count", None),
|
|
shield_concept_count=getattr(args, "shield_concept_count", None),
|
|
shield_ridge=getattr(args, "shield_ridge", None),
|
|
shield_residualize=getattr(args, "shield_residualize", None),
|
|
shield_layer_penalty=getattr(args, "shield_layer_penalty", None),
|
|
projection_target=getattr(args, "projection_target", None),
|
|
projection_row_fraction=getattr(args, "projection_row_fraction", None),
|
|
quantization=args.quantization,
|
|
large_model_mode=getattr(args, "large_model", False),
|
|
verify_sample_size=getattr(args, "verify_sample_size", None),
|
|
on_stage=on_stage,
|
|
on_log=on_log,
|
|
**prompt_kwargs,
|
|
)
|
|
|
|
with Live(make_display(), console=console, refresh_per_second=4) as live_ctx:
|
|
live = live_ctx
|
|
try:
|
|
result_path = pipeline.run()
|
|
if residue_meta:
|
|
import json
|
|
Path(result_path, "hard_negative_residue.json").write_text(
|
|
json.dumps(residue_meta, indent=2)
|
|
)
|
|
live.update(make_display())
|
|
except Exception as e:
|
|
log_lines.append(f"[red]ERROR: {e}[/]")
|
|
live.update(make_display())
|
|
raise
|
|
|
|
# ── Telemetry: send pipeline report to community leaderboard ──
|
|
try:
|
|
from obliteratus.telemetry import maybe_send_pipeline_report
|
|
maybe_send_pipeline_report(pipeline)
|
|
except Exception:
|
|
pass # Telemetry is best-effort
|
|
|
|
# ── Community contribution (--contribute flag) ──
|
|
contrib_path = None
|
|
if getattr(args, "contribute", False):
|
|
try:
|
|
from obliteratus.community import save_contribution
|
|
contrib_path = save_contribution(
|
|
pipeline,
|
|
model_name=model_name,
|
|
notes=getattr(args, "contribute_notes", ""),
|
|
)
|
|
except Exception as e:
|
|
console.print(f"[yellow]Could not save contribution: {e}[/yellow]")
|
|
|
|
console.print()
|
|
contrib_line = ""
|
|
if contrib_path:
|
|
contrib_line = f"\n Contribution: [cyan]{contrib_path}[/]"
|
|
console.print(
|
|
Panel(
|
|
f"[bold green]Abliteration complete![/]\n\n"
|
|
f" Model saved to: [cyan]{result_path}[/]\n"
|
|
f" Metadata: [cyan]{result_path}/abliteration_metadata.json[/]"
|
|
f"{contrib_line}\n\n"
|
|
f" [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
|
|
border_style="green",
|
|
title="[bold green]✓ REBIRTH COMPLETE[/]",
|
|
)
|
|
)
|
|
|
|
|
|
def _cmd_gpu_calc(args):
|
|
import math
|
|
|
|
from rich.panel import Panel
|
|
from rich.table import Table
|
|
|
|
BYTES_PER_PARAM = {
|
|
"float32": 4,
|
|
"float16": 2,
|
|
"bfloat16": 2,
|
|
"int8": 1,
|
|
"int4": 0.5,
|
|
}
|
|
|
|
# Resolve param counts
|
|
total_params_b = args.params
|
|
active_params_b = args.active_params
|
|
|
|
if total_params_b is None:
|
|
if args.model is None:
|
|
console.print("[red]Provide either a model name or --params.[/]")
|
|
raise SystemExit(1)
|
|
console.print(f"Profiling [cyan]{args.model}[/]...")
|
|
try:
|
|
from obliteratus.model_profile import profile_model
|
|
profile = profile_model(args.model, dtype=args.dtype)
|
|
total_params_b = profile.total_params_b
|
|
if active_params_b is None:
|
|
active_params_b = profile.active_params_b
|
|
except Exception as e:
|
|
console.print(f"[red]Could not profile model: {e}[/]")
|
|
raise SystemExit(1)
|
|
if total_params_b is None:
|
|
console.print("[red]Cannot determine model size. Use --params.[/]")
|
|
raise SystemExit(1)
|
|
|
|
if active_params_b is None:
|
|
active_params_b = total_params_b
|
|
|
|
bpp = BYTES_PER_PARAM[args.dtype]
|
|
gpu_mem_gb = args.gpu_mem
|
|
|
|
# Model weight memory (use base-10 GB to match HF/nvidia conventions)
|
|
weight_gb = total_params_b * bpp
|
|
|
|
# Activation overhead during forward passes (PROBE/VERIFY).
|
|
# Scales with active params, not total. Empirical from benchmarks:
|
|
# - DeepSeek-70B (149GB): failed at 160GB (2 GPUs), OK at 240GB (3 GPUs)
|
|
# - GPT-OSS-120B (234GB): failed at 240GB (3 GPUs), OK at 320GB (4 GPUs)
|
|
# This implies ~15-35% overhead. We use 20% as a reasonable middle ground.
|
|
active_weight_gb = active_params_b * bpp
|
|
activation_overhead_gb = active_weight_gb * 0.20
|
|
|
|
# CUDA context + fragmentation overhead: ~1.5 GB per GPU (fixed cost)
|
|
cuda_overhead_per_gpu = 1.5
|
|
|
|
# Total memory needed (before splitting across GPUs)
|
|
total_needed_gb = weight_gb + activation_overhead_gb
|
|
|
|
# Find minimum GPUs: we need total_needed / (gpu_mem - cuda_overhead) GPUs
|
|
usable_per_gpu = gpu_mem_gb - cuda_overhead_per_gpu
|
|
if usable_per_gpu <= 0:
|
|
console.print("[red]GPU memory too small after CUDA overhead.[/]")
|
|
raise SystemExit(1)
|
|
|
|
min_gpus = math.ceil(total_needed_gb / usable_per_gpu)
|
|
min_gpus = max(min_gpus, 1)
|
|
|
|
# Show results for a range of GPU counts
|
|
is_moe = active_params_b < total_params_b * 0.99
|
|
|
|
table = Table(title="GPU Configurations", show_edge=True)
|
|
table.add_column("GPUs", justify="right", style="cyan")
|
|
table.add_column("VRAM/GPU", justify="right")
|
|
table.add_column("Total VRAM", justify="right")
|
|
table.add_column("Headroom", justify="right")
|
|
table.add_column("Verdict", min_width=20)
|
|
|
|
# Show from min_gpus-1 (to show why it fails) up to 8
|
|
low = max(1, min_gpus - 1)
|
|
high = max(min_gpus + 3, 8)
|
|
for n in range(low, high + 1):
|
|
total_vram = n * gpu_mem_gb
|
|
usable_vram = n * usable_per_gpu
|
|
headroom = usable_vram - total_needed_gb
|
|
headroom_pct = headroom / total_needed_gb * 100
|
|
vram_per = total_needed_gb / n
|
|
|
|
if headroom < 0:
|
|
verdict = "[red]INSUFFICIENT[/]"
|
|
elif headroom_pct < 15:
|
|
verdict = "[yellow]TIGHT — may fail[/]"
|
|
elif n == min_gpus:
|
|
verdict = "[bold green]MINIMUM (recommended)[/]"
|
|
else:
|
|
verdict = "[green]OK[/] [dim](more GPUs = slower)[/]"
|
|
|
|
table.add_row(
|
|
str(n),
|
|
f"{vram_per:.1f} GB",
|
|
f"{total_vram:.0f} GB",
|
|
f"{headroom:+.1f} GB ({headroom_pct:+.0f}%)",
|
|
verdict,
|
|
)
|
|
|
|
model_label = args.model or f"{total_params_b:.1f}B params"
|
|
moe_line = ""
|
|
if is_moe:
|
|
moe_line = f"\n Active params: [cyan]{active_params_b:.1f}B[/] ({active_params_b/total_params_b*100:.0f}% of total — MoE)"
|
|
|
|
console.print(Panel(
|
|
f" Model: [cyan]{model_label}[/]\n"
|
|
f" Total params: [cyan]{total_params_b:.1f}B[/]"
|
|
f"{moe_line}\n"
|
|
f" Dtype: [cyan]{args.dtype}[/] ({bpp} bytes/param)\n"
|
|
f" Weight memory: [cyan]{weight_gb:.1f} GB[/]\n"
|
|
f" Activation est: [cyan]{activation_overhead_gb:.1f} GB[/]\n"
|
|
f" Total needed: [bold]{total_needed_gb:.1f} GB[/]\n"
|
|
f" GPU VRAM: [cyan]{gpu_mem_gb:.0f} GB[/] per device",
|
|
title="[bold]GPU Calculator[/]",
|
|
border_style="cyan",
|
|
))
|
|
console.print(table)
|
|
console.print(
|
|
f"\n [bold green]Minimum GPUs: {min_gpus}[/]"
|
|
f" ({min_gpus} x {gpu_mem_gb:.0f} GB = {min_gpus * gpu_mem_gb:.0f} GB)\n"
|
|
)
|
|
console.print(
|
|
"[dim]Note: fewer GPUs = faster (pipeline parallel has cross-device overhead).\n"
|
|
"Estimates are conservative. Actual memory may vary with sequence length\n"
|
|
"and model architecture. See 'obliteratus obliterate --help' for runtime options.[/]\n"
|
|
)
|
|
|
|
|
|
def _estimate_total_params_b(config) -> float:
|
|
"""Estimate total parameter count in billions from a HuggingFace config."""
|
|
# Some configs have explicit param counts
|
|
for attr in ("num_parameters", "n_params"):
|
|
val = getattr(config, attr, None)
|
|
if val and val > 1000:
|
|
return val / 1e9
|
|
|
|
# Estimate from architecture dimensions
|
|
h = getattr(config, "hidden_size", 0)
|
|
L = getattr(config, "num_hidden_layers", 0)
|
|
V = getattr(config, "vocab_size", 0)
|
|
i = getattr(config, "intermediate_size", h * 4)
|
|
|
|
if h == 0 or L == 0:
|
|
console.print("[red]Cannot determine model size from config. Use --params.[/]")
|
|
raise SystemExit(1)
|
|
|
|
n_heads = getattr(config, "num_attention_heads", None) or (h // 128)
|
|
head_dim = getattr(config, "head_dim", None) or (h // n_heads if n_heads else 128)
|
|
kv_heads = getattr(config, "num_key_value_heads", None) or n_heads
|
|
|
|
# Attention: Q + K + V projections + output projection
|
|
attn_params = h * (n_heads * head_dim) + h * (kv_heads * head_dim) * 2 + (n_heads * head_dim) * h
|
|
|
|
# FFN (MoE or dense)
|
|
n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1
|
|
# MoE models often have a separate intermediate size for expert FFNs
|
|
moe_i = getattr(config, "moe_intermediate_size", i)
|
|
# gate + up + down projections per expert
|
|
ffn_per_expert = h * moe_i * 3
|
|
ffn_params = ffn_per_expert * n_experts
|
|
# Some architectures (Qwen, DeepSeek) also have a shared/dense FFN per layer
|
|
if n_experts > 1 and hasattr(config, "moe_intermediate_size"):
|
|
# The dense FFN uses the main intermediate_size
|
|
ffn_params += h * i * 3
|
|
# Router
|
|
if n_experts > 1:
|
|
ffn_params += h * n_experts
|
|
|
|
# Per-layer: attention + FFN + layernorms
|
|
layer_params = attn_params + ffn_params + h * 4 # 2 layernorms, 2 params each
|
|
|
|
# Embedding + LM head
|
|
embed_params = V * h * 2 # input + output embeddings (may be tied but counts for memory)
|
|
|
|
total = L * layer_params + embed_params
|
|
return total / 1e9
|
|
|
|
|
|
def _estimate_active_params_b(config, total_params_b: float) -> float:
|
|
"""For MoE models, estimate active parameters per forward pass."""
|
|
n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1
|
|
if n_experts <= 1:
|
|
return total_params_b
|
|
|
|
top_k = getattr(config, "num_experts_per_tok", getattr(config, "top_k", 2)) or 2
|
|
|
|
h = getattr(config, "hidden_size", 0)
|
|
i = getattr(config, "intermediate_size", h * 4)
|
|
moe_i = getattr(config, "moe_intermediate_size", i)
|
|
L = getattr(config, "num_hidden_layers", 0)
|
|
|
|
# FFN per expert (uses moe_intermediate_size if available)
|
|
ffn_per_expert = h * moe_i * 3
|
|
# Active FFN = top_k experts instead of all n_experts
|
|
ffn_all = ffn_per_expert * n_experts * L
|
|
ffn_active = ffn_per_expert * top_k * L
|
|
# Non-FFN params (includes any shared/dense FFN)
|
|
non_ffn = total_params_b * 1e9 - ffn_all
|
|
active = non_ffn + ffn_active
|
|
return max(active / 1e9, 0.1)
|
|
|
|
|
|
def _make_remote_runner(args):
|
|
"""Create a RemoteRunner from CLI --remote flags."""
|
|
from obliteratus.remote import RemoteConfig, RemoteRunner
|
|
|
|
rc = RemoteConfig.from_cli_args(
|
|
args.remote,
|
|
port=args.ssh_port,
|
|
ssh_key=args.ssh_key,
|
|
remote_dir=args.remote_dir,
|
|
python=args.remote_python,
|
|
sync_results=not args.no_sync,
|
|
gpus=getattr(args, "gpus", None),
|
|
)
|
|
return RemoteRunner(rc)
|
|
|
|
|
|
def _cmd_remote_abliterate(args):
|
|
from rich.panel import Panel
|
|
|
|
runner = _make_remote_runner(args)
|
|
|
|
kwargs = {}
|
|
if args.method:
|
|
kwargs["method"] = args.method
|
|
if args.device:
|
|
kwargs["device"] = args.device
|
|
if args.dtype:
|
|
kwargs["dtype"] = args.dtype
|
|
if args.quantization:
|
|
kwargs["quantization"] = args.quantization
|
|
if args.n_directions is not None:
|
|
kwargs["n_directions"] = args.n_directions
|
|
if getattr(args, "direction_method", None):
|
|
kwargs["direction_method"] = args.direction_method
|
|
if args.regularization is not None:
|
|
kwargs["regularization"] = args.regularization
|
|
if args.refinement_passes is not None:
|
|
kwargs["refinement_passes"] = args.refinement_passes
|
|
if getattr(args, "min_layer_fraction", None) is not None:
|
|
kwargs["min_layer_fraction"] = args.min_layer_fraction
|
|
if getattr(args, "max_layer_fraction", None) is not None:
|
|
kwargs["max_layer_fraction"] = args.max_layer_fraction
|
|
if getattr(args, "harmless_pc_count", None) is not None:
|
|
kwargs["harmless_pc_count"] = args.harmless_pc_count
|
|
if getattr(args, "shield_concept_count", None) is not None:
|
|
kwargs["shield_concept_count"] = args.shield_concept_count
|
|
if getattr(args, "shield_ridge", None) is not None:
|
|
kwargs["shield_ridge"] = args.shield_ridge
|
|
if getattr(args, "shield_residualize", None) is not None:
|
|
kwargs["shield_residualize"] = args.shield_residualize
|
|
if getattr(args, "shield_layer_penalty", None) is not None:
|
|
kwargs["shield_layer_penalty"] = args.shield_layer_penalty
|
|
if getattr(args, "projection_target", None) is not None:
|
|
kwargs["projection_target"] = args.projection_target
|
|
if getattr(args, "projection_row_fraction", None) is not None:
|
|
kwargs["projection_row_fraction"] = args.projection_row_fraction
|
|
if getattr(args, "large_model", False):
|
|
kwargs["large_model"] = True
|
|
if getattr(args, "verify_sample_size", None) is not None:
|
|
kwargs["verify_sample_size"] = args.verify_sample_size
|
|
|
|
result_path = runner.run_obliterate(
|
|
model=args.model,
|
|
local_output_dir=args.output_dir,
|
|
**kwargs,
|
|
)
|
|
|
|
if result_path:
|
|
console.print(
|
|
Panel(
|
|
f"[bold green]Remote abliteration complete![/]\n\n"
|
|
f" Results at: [cyan]{result_path}[/]\n\n"
|
|
f" [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
|
|
border_style="green",
|
|
title="[bold green]REBIRTH COMPLETE (remote)[/]",
|
|
)
|
|
)
|
|
else:
|
|
console.print("[red]Remote abliteration failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
|
|
|
|
def _cmd_remote_run(args):
|
|
runner = _make_remote_runner(args)
|
|
result_path = runner.run_config(
|
|
local_config_path=args.config,
|
|
local_output_dir=args.output_dir,
|
|
preset=args.preset,
|
|
)
|
|
if result_path:
|
|
console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
|
|
else:
|
|
console.print("[red]Remote run failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
|
|
|
|
def _cmd_remote_tourney(args):
|
|
from rich.panel import Panel
|
|
|
|
runner = _make_remote_runner(args)
|
|
result_path = runner.run_tourney(
|
|
model=args.model,
|
|
local_output_dir=args.output_dir,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
quantization=args.quantization,
|
|
methods=args.methods,
|
|
hub_org=args.hub_org,
|
|
hub_repo=args.hub_repo,
|
|
dataset=args.dataset,
|
|
)
|
|
if result_path:
|
|
console.print(
|
|
Panel(
|
|
f"[bold green]Remote tournament complete![/]\n\n"
|
|
f" Results at: [cyan]{result_path}[/]",
|
|
border_style="green",
|
|
title="[bold green]TOURNAMENT COMPLETE (remote)[/]",
|
|
)
|
|
)
|
|
else:
|
|
console.print("[red]Remote tournament failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|