Hardening a Safety Guard against GCG (adversarial fine-tuning)

Take a guard model (Granite Guardian 3.1 2B), generate GCG suffixes that bypass it, fine-tune the guard to catch them, then re-run the attack from scratch and measure how much the bypass rate dropped. The complete loop:

attack G_k  ->  adversarial SFT  ->  G_{k+1}  ->  RE-ATTACK G_{k+1}  ->  repeat

Round-1 result (see results.md): one round cuts the adaptive attack bypass rate 100% → 39% on 100 held-out behaviors, makes successful attacks ~3× more expensive (median 21.5 → 61 GCG steps), and pushes 61/100 behaviors out of reach (margin ≈ −15), while holding clean-harmful recall at ~99% and dropping hard-benign false-positives 90% → 1%. This README is the run-guide; results.md has the numbers and caveats.

The four data buckets

Run it (the actual round-1 pipeline)

All commands run with the repo root as cwd. GPU = needs a CUDA GPU (a single RTX 4090 suffices; a 2B guard is ~4.8GB and GCG peaks <20GB); CPU = pure parsing, runs anywhere. Set HF_HOME to your model cache. Every stage below has already been run — the outputs are committed under round0/ (datasets/ = the four quadrants + train/val, eval/ = eval + compare + head-to-head JSONs, manifests/ = per-build provenance, verify/ = bypass checks) plus the LoRA adapter under runs/granite_round0/adapter/. This is the reproduce recipe.

pip install -r requirements.txt
export HF_HOME=/path/to/hf-cache        # Granite Guardian 3.1 2B downloads here

# 1. behavior splits (CPU) -> train 300 / test 419, near-dup-safe (no train/test leak)
python adv_datasets/build_behavior_splits.py --seed 0 --thr 0.75

# 2. base-guard refusal direction, layer 27 (GPU)
python save.py --model granite-guardian --layers 27

# 3. GCG "No"-pushing suffixes vs the base guard (GPU); 4 diverse suffixes/behavior
MODEL=granite-guardian PROMPT_FILE=adv_datasets/splits/train_behaviors.txt \
  TAG=adv_harmful_r0 NSUF=4 N_SHARDS=2 NUM_GPUS=1 STOP_MARGIN=1.0 NUM_STEPS=400 ./run_gen_shards.sh
#    stubborn behaviors that hit the step cap: pipeline/list_cap_jobs.py -> re-run with SEED=100..400

# 4. verify bypasses out-of-process (GPU) + build the adv_harmful quadrant (CPU)
python pipeline/verify_granite_strict.py --log-dir <gen-dir>     # once per log-dir -> round0/verify/
python pipeline/build_adv_harmful.py --log-dir <dirs...> --round 0 \
  --latest-run-per-behavior --max-suffixes-per-behavior 4        # -> round0/datasets/adv_harmful_round0.jsonl

# 5. the other three quadrants (CPU) — defaults read/write round0/{datasets,manifests}
python pipeline/build_adv_benign.py      --round 0    # reads round0/datasets/adv_harmful_round0.jsonl
python pipeline/build_clean_quadrants.py --round 0    # -> round0/datasets/clean_*  (+ round0/manifests/)

# 6. assemble train/val (held-out slot-3 suffixes) + scaffold format + template de-risk (CPU)
python pipeline/build_train_val.py --round 0 --holdout-slot 3    # -> round0/datasets/{train,val}_round0.jsonl
python pipeline/to_scaffold_jsonl.py                             # round0/datasets/{train,val} -> data/dataset_round0.jsonl
python pipeline/check_template_consistency.py

# 7. LoRA SFT — the hardening (GPU)
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True PYTHONPATH=src CUDA_VISIBLE_DEVICES=0 \
  python src/train/train_sft.py --config configs/train_granite_4090_lowmem.yaml \
  --data data/dataset_round0.jsonl --out runs/granite_round0        # -> runs/granite_round0/adapter/

# 8. static eval, base vs hardened (GPU) — in-distribution sanity (saturates; see results.md)
python pipeline/baseline_eval.py --round 0 --tag base_full      # -> round0/eval/baseline_eval_round0.json (default)
python pipeline/baseline_eval.py --round 0 --tag hardened_full \
  --adapter runs/granite_round0/adapter --out round0/eval/hardened_eval_round0.json

# 9. Test A — retention on 419 unseen harmful + OOD-benign FPR (GPU)
python pipeline/build_testA.py                         # -> round0/datasets/testA_rows.jsonl (default)
python pipeline/baseline_eval.py --rows-file round0/datasets/testA_rows.jsonl --tag testA_base \
  --out round0/eval/testA_base.json
python pipeline/baseline_eval.py --rows-file round0/datasets/testA_rows.jsonl --tag testA_hardened \
  --adapter runs/granite_round0/adapter --out round0/eval/testA_hardened.json

# 10. merge the adapter into a standalone guard + its own attack direction (GPU)
python scripts/merge_adapter.py --base-model ibm-granite/granite-guardian-3.1-2b \
  --adapter runs/granite_round0/adapter --out ./models/granite-guardian-hardened
python save.py --model granite-guardian-hardened --layers 27

# 11. Test B — the headline: FRESH adaptive GCG re-optimized against EACH guard (GPU)
python pipeline/build_testB_sample.py --pocket-all --n 100      # -> prompts/testB_100.txt (+ round0/manifests/testB_100_manifest.json)
MODEL=granite-guardian          PROMPT_FILE=prompts/testB_100.txt TAG=testB_100 NSUF=1 N_SHARDS=2 NUM_GPUS=1 ./run_gen_shards.sh
MODEL=granite-guardian-hardened PROMPT_FILE=prompts/testB_100.txt TAG=testB_100 NSUF=1 N_SHARDS=2 NUM_GPUS=1 ./run_gen_shards.sh

# 12. parse the adaptive head-to-head + auto-verdict (CPU)
python pipeline/parse_testB.py                                                     # -> round0/eval/testB_headtohead.json
python pipeline/compare_hardening.py --adaptive round0/eval/testB_headtohead.json  # -> round0/eval/compare_hardening_round0.json

Layout

# --- entrypoints (repo root) ---
main_granite.py                    guard registry + No-Yes margin
main_granite_shared_vocab.py       representation-level "No"-pushing GCG driver (+diversity)
save.py                            layer-27 refusal-direction extractor
run_gen_shards.sh                  sharded attack driver (MODEL=base|hardened)
nanogcg/                           directional GCG library
shared_vocab_sp.pt                 re-encode-safe token vocabulary (V*) — committed; regenerate via pipeline/cache_shared_vocab_sp.py
adv_datasets/                      build_behavior_splits.py, raw/ (AdvBench+HarmBench csv), splits/

# --- pipeline/ (stage scripts; run from repo root as `python pipeline/<name>.py`) ---
verify_granite_strict.py           out-of-process bypass verifier (also the guard template)
cache_shared_vocab_sp.py           (off-pipeline) rebuilds shared_vocab_sp.pt (helper: check_shared_vocab.py)
build_adv_harmful.py  build_adv_benign.py  build_clean_quadrants.py       the four dataset quadrants
build_train_val.py  list_cap_jobs.py  to_scaffold_jsonl.py                assemble train/val + scaffold format
baseline_eval.py                   4-quadrant static eval (+AUROC/TPR@5%FPR, --adapter, --rows-file)
build_testA.py                     retention (419 unseen harmful) + OOD-benign FPR
build_testB_sample.py              pocket-complete 100 adaptive-attack sample
parse_testB.py                     adaptive head-to-head parser
compare_hardening.py               joint-verdict auto-verdict (the report button)
check_template_consistency.py      trainer-vs-attack tokenizer/template de-risk (stage 6)

# --- committed artifacts (round-0) ---
round0/datasets/                   four quadrants + train/val (3,500 rows) + testA_rows + benign_prompts
round0/eval/                       *_eval_round0.json  testA_*.json  testB_headtohead.json  compare_hardening_round0.json
round0/manifests/                  per-build composition/provenance JSON + cap_jobs.json
round0/verify/                     verify_*.json (out-of-process bypass checks)
prompts/                           behavior / prompt lists (testB_100.txt, *_instructions.txt)
directions/                        base + hardened layer-27 directions (.pt)
runs/granite_round0/adapter/       the 38MB LoRA — the hardened guard

# --- scaffold (trainer + reference seams) ---
configs/           train_granite_4090_lowmem.yaml (used) + variants
src/train/         train_sft.py (LoRA SFT, completion-only masking)  <- used
src/{data,eval,attack}/  guard_interface / build_dataset / evaluate / jbgcg_interface  <- generic reference, superseded by pipeline/ scripts
scripts/           merge_adapter.py (used), run_iterated.sh (round-2 loop), run_milestone1.sh, run_attack.py