Buckets:

Mercity
/

red-teaming

Files

xet

Mercity/red-teaming / 29-10-progress /train.py

Rishik001

10 days ago

download

raw

7.28 kB

"""Async GRPO + LoRA training of the attacker against the frozen target, on two GPUs.
A vLLM server (serve.sh, GPU 0) generates while the trainer (this script, GPU 1) trains, so neither waits
for the other. Each training prompt is one (scenario, request). For every prompt, the rollout worker
(async_rollout.py) plays a group of full multi-turn episodes (rollout.py) against the server; the judge
grades each one and the reward (rewards.py) turns the grade into one number per episode. Only the
attacker's own tokens are trained. After every optimizer step the trainer sends the new LoRA adapter to
the server, and episodes sampled by a policy more than max_staleness steps old are discarded.
Every setting comes from config.yaml; the config the run actually used is saved to
checkpoints/<run-name>/config.yaml, next to the adapters and rollouts.jsonl (every episode in full).
    bash serve.sh                                        # terminal 1, GPU 0; wait for "Application startup"
    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
"""

import argparse
import os
from typing import Any

from peft import LoraConfig
from transformers import AutoTokenizer, GenerationConfig
from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

import async_rollout
import rollout
import utils
from datasets import Dataset
from episode import Job

# All attention projections (including Qwen3.5's linear-attention ones) and the MLP. PEFT skips names a
# model does not have.
LORA_TARGET_MODULES = [
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "in_proj_qkv",
    "in_proj_z",
    "out_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
]


def async_grpo_config(config: dict[str, Any], run_name: str, output_dir: str) -> AsyncGRPOConfig:
    train, episode = config["train"], config["episode"]
    # One optimizer step is rows_per_step packed rows of up to token_budget tokens. An episode is one row, so a
    # step trains at least rows_per_step episodes; short episodes share a row, so a step can hold more.
    rows_per_step = train["prompts_per_step"] * train["num_generations"]
    # TRL's own default: no more episodes in flight than the trainer can use before they go stale.
    max_inflight = train["max_inflight_episodes"]
    if max_inflight < 0:
        max_inflight = max(train["max_staleness"], 1) * rows_per_step

    return AsyncGRPOConfig(
        output_dir=output_dir,
        run_name=run_name,
        report_to="wandb",
        logging_steps=1,
        save_steps=train["save_steps"],
        seed=train["seed"],
        # bf16 weights, as the server serves them (TRL defaults to fp32).
        dtype="bfloat16",
        bf16=True,
        gradient_checkpointing=True,
        learning_rate=train["lr"],
        num_train_epochs=train["epochs"],
        max_steps=train["max_steps"],
        num_generations=train["num_generations"],
        per_device_train_batch_size=1,
        gradient_accumulation_steps=rows_per_step,
        # Every episode fits one row: a longer one would be dropped with a warning.
        token_budget=episode["max_total_tokens"],
        max_completion_length=episode["max_turn_tokens"],
        # The trainer needs the sampling settings to score tokens the way vLLM sampled them.
        **config["attacker_sampling"],
        vllm_server_base_url=train["vllm_url"],
        max_staleness=train["max_staleness"],
        max_inflight_tasks=max_inflight,
    )


def eos_token_ids(model: str, tokenizer: Any) -> list[int]:
    """Every token id that ends a turn, worked out the way AsyncGRPOTrainer does for its own worker."""
    try:
        ids = GenerationConfig.from_pretrained(model).eos_token_id
    except OSError:  # no generation_config.json
        ids = None
    ids = [] if ids is None else [ids] if isinstance(ids, int) else list(ids)
    if tokenizer.eos_token_id not in ids:
        ids = [tokenizer.eos_token_id, *ids]
    return ids


def rollout_worker(
    config: dict[str, Any],
    args: AsyncGRPOConfig,
    jobs: list[Job],
    dataset: Dataset,
    tokenizer: Any,
    log_path: str,
) -> async_rollout.AttackerRolloutWorker:
    """Our rollout worker, given the settings AsyncGRPOTrainer would give its own (see its __init__). The
    trainer fills in the LoRA adapter name once it has found the server serves adapters."""
    return async_rollout.AttackerRolloutWorker(
        model_name=config["model"],
        dataset=dataset,
        reward_funcs=[],
        processing_class=tokenizer,
        eos_token_ids=eos_token_ids(config["model"], tokenizer),
        num_generations=args.num_generations,
        max_inflight_tasks=args.max_inflight_tasks,
        queue_maxsize=args.queue_maxsize,
        vllm_server_url=args.vllm_server_base_url,
        max_tokens=args.max_completion_length,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
        num_sampler_logprobs=args.num_sampler_logprobs if args.score_centering else 0,
        min_p=args.min_p,
        repetition_penalty=args.repetition_penalty,
        request_timeout=args.request_timeout,
        jobs=jobs,
        config=config,
        log_path=log_path,
    )


def main(args: argparse.Namespace) -> None:
    config = utils.load_config(args.config, args.overrides)
    output_dir = os.path.abspath(args.output_dir or os.path.join("checkpoints", args.run_name))
    utils.save_config(config, os.path.join(output_dir, "config.yaml"))

    jobs = utils.load_jobs(config["data"]["train"])
    max_turns = config["episode"]["max_turns"]
    dataset = Dataset.from_list(
        [{"prompt": rollout.seed_messages(job["request"], max_turns)} for job in jobs]
    )
    tokenizer = AutoTokenizer.from_pretrained(config["model"])
    grpo_args = async_grpo_config(config, args.run_name, output_dir)

    trainer = AsyncGRPOTrainer(
        model=config["model"],
        args=grpo_args,
        train_dataset=dataset,  # sets the epoch length; the worker draws the prompts itself
        processing_class=tokenizer,
        peft_config=LoraConfig(
            r=config["train"]["lora_r"],
            lora_alpha=2 * config["train"]["lora_r"],
            lora_dropout=0.0,
            task_type="CAUSAL_LM",
            target_modules=LORA_TARGET_MODULES,
        ),
        rollout_worker=rollout_worker(
            config,
            grpo_args,
            jobs,
            dataset,
            tokenizer,
            os.path.join(output_dir, "rollouts.jsonl"),
        ),
    )
    trainer.train()
    trainer.save_model(os.path.join(output_dir, "final"))
    print(f"Saved LoRA adapter to {os.path.join(output_dir, 'final')}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--run-name", required=True)
    parser.add_argument("--config", default="config.yaml")
    parser.add_argument("--output-dir", default=None, help="default: checkpoints/<run-name>")
    parser.add_argument("overrides", nargs="*", help="config overrides, e.g. train.lr=5e-5")
    main(parser.parse_args())

Xet Storage Details

Size:: 7.28 kB
Xet hash:: 4991e20b86107b4d8b18f6320a2c338701b03c01951f265fc42954953bcb9198

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Async GRPO + LoRA training of the attacker against the frozen target, on two GPUs.

	A vLLM server (serve.sh, GPU 0) generates while the trainer (this script, GPU 1) trains, so neither waits
	for the other. Each training prompt is one (scenario, request). For every prompt, the rollout worker
	(async_rollout.py) plays a group of full multi-turn episodes (rollout.py) against the server; the judge
	grades each one and the reward (rewards.py) turns the grade into one number per episode. Only the
	attacker's own tokens are trained. After every optimizer step the trainer sends the new LoRA adapter to
	the server, and episodes sampled by a policy more than max_staleness steps old are discarded.

	Every setting comes from config.yaml; the config the run actually used is saved to
	checkpoints/<run-name>/config.yaml, next to the adapters and rollouts.jsonl (every episode in full).

	bash serve.sh # terminal 1, GPU 0; wait for "Application startup"
	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
	"""

	import argparse
	import os
	from typing import Any

	from peft import LoraConfig
	from transformers import AutoTokenizer, GenerationConfig
	from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

	import async_rollout
	import rollout
	import utils
	from datasets import Dataset
	from episode import Job

	# All attention projections (including Qwen3.5's linear-attention ones) and the MLP. PEFT skips names a
	# model does not have.
	LORA_TARGET_MODULES = [
	"q_proj",
	"k_proj",
	"v_proj",
	"o_proj",
	"in_proj_qkv",
	"in_proj_z",
	"out_proj",
	"gate_proj",
	"up_proj",
	"down_proj",
	]


	def async_grpo_config(config: dict[str, Any], run_name: str, output_dir: str) -> AsyncGRPOConfig:
	train, episode = config["train"], config["episode"]
	# One optimizer step is rows_per_step packed rows of up to token_budget tokens. An episode is one row, so a
	# step trains at least rows_per_step episodes; short episodes share a row, so a step can hold more.
	rows_per_step = train["prompts_per_step"] * train["num_generations"]
	# TRL's own default: no more episodes in flight than the trainer can use before they go stale.
	max_inflight = train["max_inflight_episodes"]
	if max_inflight < 0:
	max_inflight = max(train["max_staleness"], 1) * rows_per_step

	return AsyncGRPOConfig(
	output_dir=output_dir,
	run_name=run_name,
	report_to="wandb",
	logging_steps=1,
	save_steps=train["save_steps"],
	seed=train["seed"],
	# bf16 weights, as the server serves them (TRL defaults to fp32).
	dtype="bfloat16",
	bf16=True,
	gradient_checkpointing=True,
	learning_rate=train["lr"],
	num_train_epochs=train["epochs"],
	max_steps=train["max_steps"],
	num_generations=train["num_generations"],
	per_device_train_batch_size=1,
	gradient_accumulation_steps=rows_per_step,
	# Every episode fits one row: a longer one would be dropped with a warning.
	token_budget=episode["max_total_tokens"],
	max_completion_length=episode["max_turn_tokens"],
	# The trainer needs the sampling settings to score tokens the way vLLM sampled them.
	**config["attacker_sampling"],
	vllm_server_base_url=train["vllm_url"],
	max_staleness=train["max_staleness"],
	max_inflight_tasks=max_inflight,
	)


	def eos_token_ids(model: str, tokenizer: Any) -> list[int]:
	"""Every token id that ends a turn, worked out the way AsyncGRPOTrainer does for its own worker."""
	try:
	ids = GenerationConfig.from_pretrained(model).eos_token_id
	except OSError: # no generation_config.json
	ids = None
	ids = [] if ids is None else [ids] if isinstance(ids, int) else list(ids)
	if tokenizer.eos_token_id not in ids:
	ids = [tokenizer.eos_token_id, *ids]
	return ids


	def rollout_worker(
	config: dict[str, Any],
	args: AsyncGRPOConfig,
	jobs: list[Job],
	dataset: Dataset,
	tokenizer: Any,
	log_path: str,
	) -> async_rollout.AttackerRolloutWorker:
	"""Our rollout worker, given the settings AsyncGRPOTrainer would give its own (see its __init__). The
	trainer fills in the LoRA adapter name once it has found the server serves adapters."""
	return async_rollout.AttackerRolloutWorker(
	model_name=config["model"],
	dataset=dataset,
	reward_funcs=[],
	processing_class=tokenizer,
	eos_token_ids=eos_token_ids(config["model"], tokenizer),
	num_generations=args.num_generations,
	max_inflight_tasks=args.max_inflight_tasks,
	queue_maxsize=args.queue_maxsize,
	vllm_server_url=args.vllm_server_base_url,
	max_tokens=args.max_completion_length,
	temperature=args.temperature,
	top_p=args.top_p,
	top_k=args.top_k,
	num_sampler_logprobs=args.num_sampler_logprobs if args.score_centering else 0,
	min_p=args.min_p,
	repetition_penalty=args.repetition_penalty,
	request_timeout=args.request_timeout,
	jobs=jobs,
	config=config,
	log_path=log_path,
	)


	def main(args: argparse.Namespace) -> None:
	config = utils.load_config(args.config, args.overrides)
	output_dir = os.path.abspath(args.output_dir or os.path.join("checkpoints", args.run_name))
	utils.save_config(config, os.path.join(output_dir, "config.yaml"))

	jobs = utils.load_jobs(config["data"]["train"])
	max_turns = config["episode"]["max_turns"]
	dataset = Dataset.from_list(
	[{"prompt": rollout.seed_messages(job["request"], max_turns)} for job in jobs]
	)
	tokenizer = AutoTokenizer.from_pretrained(config["model"])
	grpo_args = async_grpo_config(config, args.run_name, output_dir)

	trainer = AsyncGRPOTrainer(
	model=config["model"],
	args=grpo_args,
	train_dataset=dataset, # sets the epoch length; the worker draws the prompts itself
	processing_class=tokenizer,
	peft_config=LoraConfig(
	r=config["train"]["lora_r"],
	lora_alpha=2 * config["train"]["lora_r"],
	lora_dropout=0.0,
	task_type="CAUSAL_LM",
	target_modules=LORA_TARGET_MODULES,
	),
	rollout_worker=rollout_worker(
	config,
	grpo_args,
	jobs,
	dataset,
	tokenizer,
	os.path.join(output_dir, "rollouts.jsonl"),
	),
	)
	trainer.train()
	trainer.save_model(os.path.join(output_dir, "final"))
	print(f"Saved LoRA adapter to {os.path.join(output_dir, 'final')}")


	if __name__ == "__main__":
	parser = argparse.ArgumentParser(
	description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
	)
	parser.add_argument("--run-name", required=True)
	parser.add_argument("--config", default="config.yaml")
	parser.add_argument("--output-dir", default=None, help="default: checkpoints/<run-name>")
	parser.add_argument("overrides", nargs="*", help="config overrides, e.g. train.lr=5e-5")
	main(parser.parse_args())