Mercity/red-teaming/29-10-progress - Storage Bucket

Auto red-teaming

An abliterated Qwen attacker plays a user in a multi-turn chat with a deployed assistant (the target, Qwen3-30B-A3B-Instruct-2507 on OpenRouter) and tries to make it break one of its rules. A judge on OpenRouter grades each conversation, and the attacker is trained with async GRPO + LoRA against the frozen target on two GPUs: a vLLM server generates on one while TRL's AsyncGRPOTrainer trains on the other.

Files

config.yaml       the single source of configuration (models, limits, training and reward)
serve.sh          the vLLM server training generates with (GPU 0), flags read from config.yaml
train.py          async GRPO + LoRA training (TRL's AsyncGRPOTrainer, GPU 1)
async_rollout.py  thin bridge between our multi-turn episode and TRL's rollout worker
rollout.py        the episode: attacker turns and target replies as one token stream
rewards.py        the judge and the reward of each episode
prompts.py        attacker and judge prompt templates, tactic labels
episode.py        a few plain type aliases plus JSONL record formatting
utils.py          OpenRouter calls, attacker-output parsing, leak check, config and file loading
setup.sh          installs dependencies and downloads the attacker
datagen/          data generation (Simula); data/simula-v*/ holds each dataset version
docs/             experiment catalogue and reports

Run

Setup, once on the pod (.env needs OPENROUTER_BASE_URL and OPENROUTER_API_KEY):

bash setup.sh
wandb login

The canonical simula-v6.5 data is in datagen/data/simula-v6.5/, and config.yaml points there.

Training reads config.yaml. Override any key on the command line with section.key=value. Each run saves the config it used next to its outputs.

Training needs two GPUs: start the server first, and the trainer once the server log says "Application startup complete". Give serve.sh the same model and train.* overrides as train.py.

# Terminal 1, GPU 0: the vLLM server
bash serve.sh

# Terminal 2, GPU 1: train. checkpoints/<run>/ gets the adapters, rollouts.jsonl and config.yaml
# expandable_segments cuts memory fragmentation on long (~25k-token) episodes
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
CUDA_VISIBLE_DEVICES=1 accelerate launch --num_processes 1 train.py --run-name proto-9b-v8
CUDA_VISIBLE_DEVICES=1 accelerate launch --num_processes 1 train.py --run-name smoke train.max_steps=2 train.prompts_per_step=1

Evaluation is intentionally off. AsyncGRPO does not currently implement the normal Trainer evaluation path; adding a second in-process vLLM evaluator or a training callback would complicate and interrupt the async run. The rollout JSONL already logs outcome, ASR inputs, reward parts and full conversations for every training episode. A small post-training evaluator can be added later if a held-out run is required.

What one optimizer step does

With the checked-in configuration, one optimizer step is intended to consume four GRPO groups:

1. The rollout worker takes a request and creates train.num_generations=8 independent episodes for it. The eight final rewards form one GRPO comparison group.
2. In each episode, the attacker can take at most episode.max_turns=20 turns. One attacker turn can sample at most episode.max_turn_tokens=4096 tokens, and the complete attacker/target token stream is capped at episode.max_total_tokens=40960 tokens.
3. Every attacker message is sent to the frozen target. Target replies are appended to the attacker's context, but their tokens have a zero loss mask. Only tokens sampled by the attacker are trained.
4. The judge assigns an outcome and attack-quality grade. rewards.score_episode() builds the scalar reward from the named parts in config.yaml: outcome score + quality reward + early-stop bonus - false-stop penalty - formatting penalty. Infrastructure failures are dropped rather than treated as bad attacks.
5. TRL normalizes rewards within each eight-episode group. The trainer uses a micro-batch of one packed row and accumulates train.prompts_per_step * train.num_generations = 4 * 8 = 32 rows before an optimizer update. Packing can put multiple short sequences in a row, so 32 is a row budget, not a guarantee that exactly 32 episodes contributed.
6. After the optimizer update, the new LoRA version is made available to the vLLM server. Rollouts more than train.max_staleness=4 policy versions old are discarded.

The key arguments are therefore:

Argument	Current value	Meaning
train.num_generations	8	Episodes compared for one request
train.prompts_per_step	4	Nominal GRPO groups accumulated per optimizer step
per_device_train_batch_size	1	Packed rows in one forward/backward micro-step
gradient_accumulation_steps	32	Micro-steps before one optimizer update
token_budget	40,960	Maximum tokens in one packed training row
train.lr	1e-4	Constant LoRA learning rate
train.lora_r	16	LoRA rank
train.max_staleness	4	Maximum accepted policy-version lag
train.max_inflight_episodes	-1	Automatic inflight limit (4 * 32 = 128)

Code quality checks

The core training path is type-checked; all Python is formatted and linted with Ruff.

pip install ruff mypy
ruff format --check .
ruff check .
mypy
python -m unittest discover -s tests -v

Data generation (one dataset version shown):

PYTHONUTF8=1 simula run datagen/scenarios.yaml
python datagen/write_target_prompts.py --data-dir datagen/data/simula-v5
python datagen/to_scenarios.py --data-dir datagen/data/simula-v5
python datagen/build_viewer.py datagen/data/simula-v5