Fine-tuning breaks safety. SafeTune fixes it.¶

You fine-tuned an aligned model and it quietly stopped refusing harmful requests. SafeTune gives you a broad registry of methods to fix that — harden the fine-tuning, recover drifted weights, steer at inference, or unlearn a capability — each audited against its source paper, so you know what you're actually running.

Get started View on GitHub

v1.0.0 LSAL v1.1 · source-available Python 3.12+ per-method audit verdicts

Pick one method per task — not a pipeline¶

SafeTune is a library of alternatives, not a sequence. Each task has many independent methods that solve it by different mechanisms. You pick one. You don't chain them.

◆ SafeTune audited LLM-safety methods

Tier 1 · Interventions ↓

Tier 2 · Instrumentation ↓

TIER 1 · INTERVENTIONS CHANGE model safety

hardenTrain-time · base + FT data

recoverWeight-space · drifted model

unlearnForget-set · remove capability

steerInference-time · frozen model

TIER 2 · INSTRUMENTATION OBSERVE safety

interpretDiagnose · locate safety

evaluateMeasure · red-team + eval

I want to…	Try	Pillar
keep safety during fine-tuning	`harden.SafeGradTrainer`	`safetune.runner.harden`
restore safety after drift (no training)	`recover.ReStaTrainer`	`safetune.runner.recover`
refuse harmful prompts at inference	`steer.RefusalDirectionTrainer`	`safetune.runner.steer`
remove a capability	`unlearn.RMUTrainer`	`safetune.runner.unlearn`
locate where safety lives	`safety_circuit_info`	`safetune.interpret`
measure safety	`evaluate()`	`safetune.evaluate`

New here? Read How to use these docs for navigation tips and a guide to the audit badge system.

See it in 60 seconds¶

SteerRecoverHardenUnlearnInterpretEvaluate

from safetune.runner import steer

# harmful_prompts / harmless_prompts — ~20 contrast examples each (list of str)
trainer = steer.RefusalDirectionTrainer(model, tokenizer)
wrapped, _ = trainer.calibrate(harmful=harmful_prompts, harmless=harmless_prompts)

Run it: python examples/quickstart/quickstart.py

from safetune.runner import recover

# drifted_model  — your fine-tuned checkpoint (AutoModelForCausalLM)
# base_model     — the original pre-trained model
# aligned_model  — the safety-aligned reference (e.g. the RLHF checkpoint)
trainer = recover.ReStaTrainer(
    drifted_model, base_model=base_model, aligned_model=aligned_model
)
patched = trainer.apply()

Run it: python examples/quickstart/recover_quickstart.py

from safetune.runner import harden

# train_dataset  — your domain fine-tuning data (any HuggingFace Dataset)
# safety_dataset — a safety corpus; use safetune.data.load_beavertails() or your own
trainer = harden.SafeGradTrainer(model, tokenizer)
trainer.train(train_dataset, safety_dataset=safety_dataset)

Run it: python examples/quickstart/harden_quickstart.py

from safetune.runner import unlearn

# forget_batches — DataLoader of prompts covering the capability to erase
# retain_batches — DataLoader of general-purpose prompts to preserve the rest
trainer = unlearn.RMUTrainer(model)
trainer.unlearn(forget=forget_batches, retain=retain_batches)

Run it: python examples/quickstart/unlearn_quickstart.py

from safetune.interpret import safety_circuit_info

# harmful / harmless — contrast prompt lists; same format as Steer calibration
circuit = safety_circuit_info(
    model, tokenizer,
    harmful_prompts=harmful, harmless_prompts=harmless,
)

Run it: python examples/quickstart/interpret_quickstart.py

from safetune.evaluate import evaluate

# "harmbench" / "xstest" — benchmark suites downloaded automatically on first run
# judge="wildguard"       — WildGuard classifier used to score model responses
results = evaluate(model, tokenizer=tok, benchmarks=["harmbench", "xstest"], judge="wildguard")
print(results["harmbench"]["asr"])

Run it: python examples/quickstart/evaluate_quickstart.py

Four intervention pillars¶

Methods that change a model's safety — train-time, weight-space, or inference-time.

Harden

Stop safety from breaking during fine-tuning. Trainers replace your SFT loop — gradient surgery, data alternation, representation perturbation. 8 families.

Harden guide
Recover

Safety already broke? Patch the weights directly — no retraining. Whole-model arithmetic to neuron-level surgery. 6 granularities.

Recover guide
Steer

Don't touch weights. Wrap any model with a refusal direction or logits processor — active only while installed. Optional vLLM backend for faster batched throughput.

Steer guide
Unlearn

Make a model forget a specific skill. Trains on a forget set + retain set to erase knowledge while preserving the rest.

Unlearn guide

Two instrumentation tools¶

Methods that observe safety — diagnose where it lives and measure whether it holds.

Interpret — locate where safety lives

Find refusal directions, safety neurons, and circuits inside a model. The artifacts feed Steer and circuit-guided Recover — the same finding has three uses: steer with it, mask a weight edit, or report it.

identify_safety_neurons · safety_circuit_info · eap_safety_circuit

Interpret guide
Evaluate — measure what happened

Red-team attacks and benchmark/judge evaluation. HarmBench, XSTest, AdvBench, WildJailbreak. WildGuard and LlamaGuard-3 judges. One call.

evaluate(model, benchmarks=["harmbench"]) · BoNAttack

Evaluate guide

How teams use SafeTune¶

Four common situations and which method fits each. Code snippets use placeholder variable names — swap in your own model ID and dataset.

Drift after fine-tuningCompliance-critical deploymentProduction steer — no retrainingErase a capability

Situation: You fine-tuned an aligned model on proprietary data. It now complies with harmful requests the base model refused. You need safety back without retraining from scratch.

from safetune.runner import recover
from safetune.evaluate import evaluate

# drifted_model / base_model / aligned_model — your AutoModelForCausalLM objects
# "harmbench" — HarmBench red-team suite, downloaded automatically on first run
# judge="wildguard" — WildGuard classifier used to score model responses

# 1. Measure how bad the drift is
before = evaluate(drifted_model, tokenizer=tokenizer, benchmarks=["harmbench"], judge="wildguard")
print(f"Refusal rate: {before['harmbench']['refusal_rate']:.0%}")  # e.g. 43%

# 2. Patch weights directly — no training, ~30 s on a single GPU
trainer = recover.ReStaTrainer(
    drifted_model, base_model=base_model, aligned_model=aligned_model, alpha=0.5
)
patched = trainer.apply()

# 3. Confirm recovery
after = evaluate(patched, tokenizer=tokenizer, benchmarks=["harmbench"], judge="wildguard")
print(f"Refusal rate: {after['harmbench']['refusal_rate']:.0%}")   # e.g. 89%

26 recovery methods from whole-model arithmetic to neuron-level surgery. Recover guide

Situation: Fine-tuning a model for a regulated application (banking, healthcare, legal). Safety must hold through the training run — you can't patch it afterward.

from safetune.runner import harden

# train_dataset  — your domain data (HuggingFace Dataset)
# safety_dataset — safety corpus; use safetune.data.load_beavertails() or your own

# Drop-in replacement for your SFT trainer
trainer = harden.SafeGradTrainer(
    model, tokenizer,
    rho=0.05,           # gradient surgery strength
    kl_temperature=1.0, # KL alignment temperature
)
trainer.train(train_dataset, safety_dataset=safety_dataset)
# Saves a fine-tuned checkpoint with safety preserved

Safety-constrained training. The model learns your domain without forgetting refusals. 27 harden methods

Situation: Live production model. You cannot retrain. You need to enforce refusals on a new harm category within hours, not weeks.

from safetune.runner import steer

# harmful_prompts / harmless_prompts — ~20 contrast examples each (list of str)
# inputs — standard tokenizer output: {"input_ids": ..., "attention_mask": ...}

# Calibrate on 20–30 contrastive examples (harmful / harmless)
trainer = steer.RefusalDirectionTrainer(model, tokenizer)
wrapped, _ = trainer.calibrate(harmful=harmful_prompts, harmless=harmless_prompts)

wrapped.install()  # activate the steering hooks

# generate through the wrapped model while the hooks are active
output = wrapped.model.generate(**inputs)

# Revert without reloading the model
wrapped.remove()

Zero retraining. Reversible. Optional vLLM backend for higher throughput. 19 steer methods

Situation: A deployed model knows something it shouldn't — a capability you need to remove (CBRN knowledge, PII memorisation, a deprecated skill). Standard fine-tuning doesn't forget cleanly.

from safetune.runner import unlearn

trainer = unlearn.RMUTrainer(model)
trainer.unlearn(
    forget=forget_batches,   # what to erase
    retain=retain_batches,   # what to keep
)
# Returns a checkpoint with just that capability removed

Trains only on the forget/retain split — not a full retrain. 6 unlearn methods

Where next¶


Getting started	Install, quickstart, the taxonomy
Guides	Per-pillar usage guides with code
Usage	CLI, YAML config, and Python API — three ways to run any method
Feature Map	Per-method audit badges — faithful, variant, simplified
Notebooks	Colab notebooks for each pillar
References	Full paper table, eval protocols

Cite¶

@misc{seth2026safetune,
  title  = {SafeTune: A Unified Library for Preserving and Restoring
            Safety in Fine-Tuned {LLM}s},
  author = {Seth, Pratinav and Kaushal, Anshul and Sadhu, Saisab and
            Sankarapu, Vinay Kumar},
  year   = {2026},
  note   = {Pratinav Seth, Anshul Kaushal, and Saisab Sadhu contributed equally.},
}