HarmBench

repository·main·Indexed 21 days ago

https://github.com/centerforaisafety/harmbench

A standardized evaluation framework for automated red teaming and assessing the robustness of Large Language Models (LLMs). It includes a scalable, multi-step pipeline for evaluating red teaming methods and LLMs, and integrates with the Alignment Handbook for training techniques such as Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), reward modeling, and rejection sampling.

Tokens
26.5K
Snippets
58
Records
80
Agent score
76%

What's inside HarmBench

  1. Understand the Alignment Handbook project structure

    main

    The project is organized into several key directories that define how you interact with the training pipeline:

    • scripts/: Contains the actual scripts used to train and evaluate chat models. These scripts support distributed training via DeepSpeed ZeRO-3 or parameter-efficient fine-tuning using LoRA/QLoRA.
    • recipes/: Contains YAML configuration files for training runs. Each recipe defines all parameters for a specific model reproduction (e.g., Zephyr 7B).
    • src/: The core source code for the project.
    • chapters/: Educational content.
    • tests/: Unit tests for the codebase.
  2. Format datasets for SFT and DPO training

    main

    The training scripts use get_datasets() to mix multiple datasets. To use your own datasets, they must follow specific column formats for the chosen task.

    Dataset Mixing (YAML): Use datasets_mixer to define proportions and dataset_splits to define which splits to use.

    datasets_mixer:
        dataset_1: 0.5
        dataset_2: 0.66
    dataset_splits:
    - train_xxx
    - test_xxx

    Required Column Formats:

    For SFT (Supervised Fine-Tuning):

    • messages: A list of dicts: [{"role": "{role}", "content": "{content}"}, ...]

    For DPO (Direct Preference Optimization):

    • chosen: A list of dicts (preferred dialogue) in the SFT format.
    • rejected: A list of dicts (dispreferred dialogue) in the SFT format.

    Recommended Split Naming:

    • {train,test}_sft: For SFT training.
    • {train,test}_gen: For generation ranking (rejection sampling/PPO).
    • {train,test}_prefs: For preference modelling (reward modelling/DPO).
    datasets_mixer:
        dataset_1: 0.5  # Use 50% of the training examples
        dataset_2: 0.66 # Use 66% of the training examples
        dataset_3: 0.10 # Use 10% of the training examples
    dataset_splits:
    - train_xxx         # The training splits to mix
    - test_xxx          # The test splits to mix
  3. How experiments and target models are stored

    main

    HarmBench follows a strict directory convention for saving results based on the method, experiment, and target model names.

    Test Case Storage Test cases for an experiment are saved at: {base_results_dir}/{method_name}/{experiment_name}/test_cases/test_cases.json

    Completions and Results Storage Once test cases exist, completions and classification results are saved using the target model name:

    • Completions: {base_results_dir}/{method_name}/{experiment_name}/completions/{model_name}.json
    • Results: {base_results_dir}/{method_name}/{experiment_name}/results/{model_name}.json

    Example for a GCG method running a baichuan2_7b_1000steps experiment against a baichuan2_7b target model:

    • results/GCG/baichuan2_7b_1000steps/completions/baichuan2_7b.json
    • results/GCG/baichuan2_7b_1000steps/results/baichuan2_7b.json
  4. How the HarmBench evaluation pipeline works

    main

    The HarmBench evaluation pipeline consists of three primary stages to assess model robustness against red teaming attacks:

    1. Step 1: Generate Test Cases: Red teaming methods (e.g., GCG, ZeroShot) are used to generate adversarial prompts (test cases) based on a set of target behaviors.
    2. Step 2: Generate Completions: The generated test cases are used as prompts for a target LLM, and the model's responses (completions) are collected.
    3. Step 3: Evaluate Completions: A classifier evaluates the completions to determine if they successfully exhibit the target behavior, allowing for the calculation of the Attack Success Rate (ASR).

    Optional Step 1.5: Merge Test Cases: For methods that parallelize generation by behavior, this step merges individual behavior results into a single test_cases.json file required for Step 2.

  5. Understand the HarmBench evaluation pipeline

    main

    The evaluation pipeline is composed of four primary scripts that work in two stages to compute the Attack Success Rate (ASR) for a target model:

    Stage 1: Test Case Generation

    1. generate_test_cases.py: Generates initial test cases.
    2. merge_test_cases.py: Merges results into a single test_cases.json dictionary for a specific experiment.

    Stage 2: Completion and Evaluation 3. generate_completions.py: Runs the generated test cases against a target model. 4. evaluate_completions.py: Evaluates the model completions to compute the ASR.

    An experiment is defined by a specific red teaming method and its hyperparameters (often optimized for a specific target model), while a target model is the model being attacked.

  6. Core alignment techniques in the Handbook

    main

    The Alignment Handbook provides robust training recipes for the following techniques:

    • Supervised fine-tuning (SFT): Teaching models to follow instructions and guidance on dataset curation.
    • Reward modeling: Training models to distinguish between responses based on human or AI preferences.
    • Rejection sampling: A technique used to boost the performance of an SFT model.
    • Direct preference optimisation (DPO): A powerful alternative to PPO for preference alignment.
  7. Understand proper method names and experiment mapping

    main

    In HarmBench, a proper method name is the key used in configs/pipeline_configs/run_pipeline.yaml. This is distinct from the underlying class name in the baselines folder.

    When you provide a proper method name and a model name, run_pipeline.py maps them to a specific experiment folder using the experiment_name_template defined in the config.

    Example Mapping: If run_pipeline.yaml contains:

    GCG-Transfer:
      class_name: EnsembleGCG
      experiment_name_template: llama2_7b_vicuna_7b_llama2_13b_vicuna_13b_multibehavior_1000steps

    Running --methods GCG-Transfer --models llama2_7b will automatically resolve to the experiment folder llama2_7b_vicuna_7b_llama2_13b_vicuna_13b_multibehavior_1000steps.

    Model Type Compatibility: Methods can restrict which models they run on using the allowed_target_model_types field (e.g., [open_source]). If a model's configuration (in models.yaml) has a model_type that is not in the method's allowed list, the pipeline will skip that model for that method.

  8. Use Dynamic Experiment Configs to reference models

    main

    You can create experiments that automatically scale across all models defined in models.yaml by using <model_name#> placeholders in your method config. This prevents manual duplication of model paths.

    Syntax

    Use <model_name#> followed by bracketed keys to traverse the model configuration dictionary. For example:

    • <model_name1>['model']['model_name_or_path'] references the path in the first model matched.

    Behavior

    • Implicit Expansion: If you pass an experiment name containing <model_name1> to generate_test_cases.py, the system expands that name into one unique experiment for every entry in models.yaml.
    • Manual Overrides: If a manually defined experiment name matches an expanded dynamic name, the manual configuration takes precedence.
    • Ensembles: You can use multiple placeholders (e.g., <model_name1>, <model_name2>) to create ensemble experiments that target multiple models simultaneously.
    # In configs/method_configs/GCG_config.yaml
    
    <model_name1>:
      target_model:
        model_name_or_path: <model_name1>['model']['model_name_or_path']
    
    # Advanced: Ensemble of two models
    ensemble_optimization_<model_name1>_<model_name2>:
      target_models:
        - <model_name1>['model']['model_name_or_path']
        - <model_name2>['model']['model_name_or_path']
  9. Override training hyperparameters and logging

    main

    You can override default YAML configuration parameters by appending them as command-line arguments to the accelerate launch command.

    Common Overrides:

    • --per_device_train_batch_size=N
    • --num_train_epochs=N
    • --report_to=wandb (to log metrics to Weights and Biases if logged in).

    Note: When scaling the number of GPUs, it is recommended to scale the per-device batch size or gradient accumulation steps to maintain a constant global batch size.

    # Change batch size and epochs
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_{task}.py recipes/{model_name}/{task}/config_full.yaml --per_device_train_batch_size=42 --num_train_epochs=5
    
    # Log to Weights and Biases
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_{task}.py recipes/{model_name}/{task}/config_full.yaml --report_to=wandb
  10. Step 3: Evaluate Completions

    main

    Use evaluate_completions.sh to classify completions and compute the Attack Success Rate (ASR). This step requires a classifier model.

    Available Classifiers:

    • Text behaviors: cais/HarmBench-Llama-2-13b-cls
    • Multimodal behaviors: cais/HarmBench-Llama-2-13b-cls-multimodal-behaviors

    Arguments for scripts/evaluate_completions.sh:

    • cls_path: The path/ID of the classifier model.
    • behaviors_path: Path to the behavior dataset CSV.
    • completions_path: Path to the JSON file generated in Step 2.
    • save_path: Path to save the final classification results.
    base_save_dir="results"
    method_name="GCG"
    experiment_name="vicuna_13b_v1_5"
    model_name="vicuna_13b_v1_5"
    
    cls_path="cais/HarmBench-Llama-2-13b-cls"
    behaviors_path="./data/behavior_datasets/harmbench_behaviors_text_val.csv"
    completions_path="$base_save_dir/$method_name/$experiment_name/completions/$model_name.json"
    save_path="$base_save_dir/$method_name/$experiment_name/results/$model_name.json"
    
    ./scripts/evaluate_completions.sh $cls_path $behaviors_path $completions_path $save_path
  11. Fine-tune chat models using various methods

    main

    You can align LLMs using three primary methods depending on your hardware availability. The scripts use accelerate launch to manage training.

    Available Methods:

    • Full fine-tuning: Uses DeepSpeed ZeRO-3 on multi-GPU machines (e.g., 8 x A100).
    • LoRA fine-tuning: Suitable for single consumer GPUs (e.g., RTX 4090).
    • QLoRA 4-bit fine-tuning: Optimized for single consumer GPUs by loading in 4-bit.
    • LoRA with ZeRO-3: For multi-GPU setups using LoRA adapters.

    Placeholders:

    • {task}: The training type (e.g., sft, dpo).
    • {model_name}: The recipe directory name in recipes/.
    • {num_gpus}: The number of GPUs available for LoRA ZeRO-3 training.
    # Full training with ZeRO-3 on 8 GPUs
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_{task}.py recipes/{model_name}/{task}/config_full.yaml
    
    # LoRA training on a single GPU
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_{task}.py recipes/{model_name}/{task}/config_lora.yaml
    
    # QLoRA 4-bit training on a single GPU
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_{task}.py recipes/{model_name}/{task}/config_lora.yaml --load_in_4bit=true
    
    # LoRA training with ZeRO-3 on two or more GPUs
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml --num_processes={num_gpus} scripts/run_{task}.py recipes/{model_name}/{task}/config_lora.yaml
  12. Replicate Zephyr-7b-β via Full Training

    main

    To replicate the Zephyr-7b-β model using full parameter fine-tuning, follow a two-step process: Supervised Fine-Tuning (SFT) followed by Direct Preference Optimization (DPO).

    Requirements:

    • 8 GPUs with 80GB of VRAM each.

    Workflow:

    1. SFT: Fine-tune Mistral 7B on a filtered version of the UltraChat dataset using scripts/run_sft.py with a DeepSpeed ZeRO-3 configuration.
    2. DPO: Align the resulting SFT model to AI feedback using the preprocessed UltraFeedback dataset via scripts/run_dpo.py.
    # Step 1 - SFT
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_sft.py recipes/zephyr-7b-beta/sft/config_full.yaml
    
    # Step 2 - DPO
    ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_dpo.py recipes/zephyr-7b-beta/dpo/config_full.yaml