Token Buncher: Shielding LLMs from Harmful Reinforcement Learning Fine-Tuning

This repository is the official implementation of Token Buncher. It provides the full pipeline to (1) launch harmful-RL attacks, (2) apply the Token Buncher defense, and (3) evaluate safety and capability on a suite of benchmarks.

• [2026-08] 🎉 Token Buncher is accepted by ACM CCS 2026!
• [2026-06] Updated the code for arXiv version v3.

Token-Buncher/
├── README.md
├── requirements.txt              # Python dependencies (PyTorch installed separately)
├── setup.py / pyproject.toml     # Installs the vendored `verl` training engine
├── run_attack.sh                 # Entry point: launch a harmful-RL attack
├── run_defence.sh                # Entry point: apply the Token Buncher defense
├── configs/                      # Per-algorithm training configs (grpo/ppo/rloo/reinpp/lora/defence)
├── reward_scores/                # Reward functions used during training
│   ├── localapi.py               #   attack reward (queries the DeBERTa harmful classifier)
│   └── entropyreward.py          #   defense reward (token-entropy based, no reward model)
├── scripts/
│   ├── host_deberta_api.py       # Reward-model server for attacks (Flask, port 50050)
│   ├── convert_fsdp_to_hf.py     # Convert verl FSDP checkpoints -> HuggingFace format
│   └── test_deberta_api.py       # Sanity-check the reward-model server
├── dataset/                      # Training-data preprocessing (BeaverTails, AlpacaEval, merge)
├── evaluation/                   # Evaluation benchmarks (see "Evaluation" section)
│   ├── eval_data/                # Bundled benchmark data (HarmBench/StrongREJECT/MATH500/GSM8K/WMDP-evil)
│   ├── harmfulscore/             # HarmBench / StrongREJECT  -> Harmful Score
│   ├── math/                     # GSM8K / MATH500
│   ├── mmlu-pro/                 # MMLU-Pro
│   └── wmdp-evil/                # WMDP-evil MCQ
├── verl/                         # Vendored training engine (adapted from verl)
└── outputs/                      # Training/eval outputs

conda create -n tokenbuncher python=3.10 -y
conda activate tokenbuncher

# install torch
pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128

# install verl (vendored, no deps) and the remaining dependencies
pip install --no-deps -e .
pip install -r requirements.txt

# flash attention 2
pip install flash-attn --no-build-isolation

Training uses the following base models and the attack reward model:

Evaluation uses two scorer models:

We provide preprocessing scripts for the datasets used during training. The scripts write parquet files under dataset/:

# Attack data — BeaverTails harmful prompts (used by run_attack.sh)
python ./dataset/data_preprocess/beavertails.py --template_type qwen --local_dir ./dataset/beavertails-qwen

# Defence data — AlpacaEval (benign) merged with BeaverTails (used by run_defence.sh)
python ./dataset/data_preprocess/alpacaeval.py  --template_type qwen --local_dir ./dataset/alpacaeval-qwen
python ./dataset/data_preprocess/simple_merge.py \
    --dir1 ./dataset/beavertails-qwen --dir2 ./dataset/alpacaeval-qwen \
    --output_dir ./dataset/alpaca-beaver-qwen

• --template_type supports base, ministral, qwen.

All benchmark data needed for evaluation is already provided under evaluation/eval_data/ (HarmBench, StrongREJECT, MATH500, GSM8K, and the WMDP-evil set) — no separate download is required. MMLU-Pro is pulled from HuggingFace at runtime by its harness.

run_attack.sh launches a harmful-RL attack. It runs the GRPO variant by default; the configs for the other algorithms live under configs/.

# 1) start the reward-model server (Flask, port 50050 — change PORT if it conflicts)
python scripts/host_deberta_api.py

# 2) in another terminal, edit run_attack.sh then launch
bash run_attack.sh

Switch algorithms by editing which configs/attack_*.sh is invoked (attack_grpo.sh, attack_ppo.sh, attack_rloo.sh, attack_reinpp.sh, plus the LoRA variant attack_grpo_lora.sh).

bash run_defence.sh

Edit BASE_MODEL / OUTPUT_DIR / EXPERIMENT_NAME in run_defence.sh first. Unlike the attack, the defense does not need the reward-model server — it uses the model's own token-entropy as the reward signal (reward_scores/entropyreward.py). Training is logged to Weights & Biases (wandb login on first use).

After training, convert the FSDP checkpoint to HuggingFace format:

python scripts/convert_fsdp_to_hf.py \
    <OUTPUT_DIR>/global_step_<step>/actor \
    Qwen/Qwen2.5-7B-Instruct \
    ./outputs/Qwen2.5-7B-TokenBuncher

The converter expects a 2-GPU (world_size=2) FSDP checkpoint.

All benchmark data is bundled under evaluation/eval_data/, so the steps below run directly on any model checkpoint.

cd evaluation/harmfulscore

# generate responses
python generate_vllm_response.py \
    --model <MODEL_PATH> --tokenizer <TOKENIZER_PATH> \
    --output-file-name <OUTPUT_NAME> \
    --dataset HarmBench            # or StrongREJECT

# compute Harmful Score
python beaverdam.py \
    --model_path <beaver-dam-7b PATH> \
    --eval_dataset <OUTPUT_NAME>.jsonl \
    --output_dir <OUTPUT_DIR>

The Beaver-Dam classifier writes results_summary.json, where flagged_percentage is the Harmful Score. Optionally, llm_judge.py computes an LLM-as-judge Judge Score with Llama-3.1-70B (slow; runs a 70B model over thousands of pairs).

cd evaluation/math

python math_vllm_gen.py \
    --model <MODEL_PATH> \
    --max-new-tokens 4096 \
    --output-file-name <OUTPUT_NAME> \
    --dataset MATH500             # or GSM8K

python evaluate_gsm8k.py   --input_file <OUTPUT_NAME>.jsonl
python evaluate_math500.py --input_file <OUTPUT_NAME>.jsonl

Adapted from the official MMLU-Pro repo.

cd evaluation/mmlu-pro
# set MODEL_PATH and MODEL_NAME in eval_models.sh first
bash eval_models.sh

Results are saved under eval_results/, including per-subject generations, summaries, and accuracy CSVs.

cd evaluation/wmdp-evil

python generate_responses.py \
    --model <MODEL_PATH> \
    --json_file_path ../eval_data/WMDP-evil/full.jsonl \
    --output_file_path <OUTPUT_NAME>.jsonl

python extract_answer.py --json_file <OUTPUT_NAME>.jsonl

The vendored verl/ training engine is adapted from verl; we thank the authors for their great work.

We also thank the authors of the benchmarks used for evaluation, whose datasets (bundled under evaluation/eval_data/, and MMLU-Pro fetched at runtime) are derived from:

• HarmBench
• StrongREJECT
• MATH500
• GSM8K
• MMLU-Pro
• WMDP-evil — derived from the WMDP benchmark

We are also grateful to the authors of the BeaverTails dataset and the reward / scorer models for making them publicly available.