mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
Rached Rayeh
2295ef3571
The README/CONTRIBUTING examples used `obliteratus aggregate --format ...` but the CLI only accepted `--dir`.
This adds `--format {summary,latex}`, `--metric`, and `--min-runs` to the aggregate command, reuses community LaTeX table generation, and adds CLI parsing tests to align behavior with documented usage.
1430 lines
56 KiB
Python
1430 lines
56 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",
|
|
)
|
|
aggregate_parser.add_argument(
|
|
"--format",
|
|
choices=["summary", "latex"],
|
|
default="summary",
|
|
help="Output format (default: summary)",
|
|
)
|
|
aggregate_parser.add_argument(
|
|
"--metric",
|
|
type=str,
|
|
default="refusal_rate",
|
|
help="Metric to display in LaTeX output (default: refusal_rate)",
|
|
)
|
|
aggregate_parser.add_argument(
|
|
"--min-runs",
|
|
type=int,
|
|
default=1,
|
|
help="Minimum runs per (model, method) to include (default: 1)",
|
|
)
|
|
|
|
# --- 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, generate_latex_table, 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)
|
|
min_runs = max(args.min_runs, 1)
|
|
if min_runs > 1:
|
|
for model in list(aggregated.keys()):
|
|
for method in list(aggregated[model].keys()):
|
|
if aggregated[model][method]["n_runs"] < min_runs:
|
|
del aggregated[model][method]
|
|
if not aggregated[model]:
|
|
del aggregated[model]
|
|
|
|
if not aggregated:
|
|
console.print("[yellow]No results meet the minimum run threshold.[/yellow]")
|
|
return
|
|
|
|
if args.format == "latex":
|
|
console.print(generate_latex_table(aggregated, metric=args.metric))
|
|
return
|
|
|
|
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()
|