#!/usr/bin/env python # coding=utf-8 # Copyright 2023 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Supervised fine-tuning script for decoder language models. """ import logging import random import sys import os import datasets import torch import transformers from transformers import set_seed sys.path.insert(1, '/data/mantas_mazeika/projects/2023/red_teaming_benchmark/redbench/adv_training/trainer_code/alignment-handbook/src') from alignment import DataArguments, H4ArgumentParser, ModelArguments, SFTConfig, apply_chat_template, get_datasets, get_kbit_device_map, get_peft_config, get_quantization_config, get_tokenizer from trl import SFTTrainer from adv_training_utils import CustomSFTTrainer, get_test_case_wrapping import csv import json import time from datasets import Dataset from datasets import concatenate_datasets import uuid # ADV TRAINING ARGS NUM_GCG_STEPS = 5 # The number of steps to run GCG for when updating a test case GCG_SEARCH_WIDTH = 512 # The search_width argument in GCG NUM_ACCELERATE_GPUS = 8 # Number of GPUs; must match the `num_processes` field in `recipes/accelerate_configs/deepspeed_config.yaml` NUM_TEST_CASES_TO_UPDATE_PER_STEP = 8 # Number of test cases to update in each batch (not counting grad accumulation); must be a multiple of NUM_ACCELERATE_GPUS NUM_SINGLE_BEHAVIOR_TEST_CASES = 8 # Number of test cases in the pool; we sample from these in each batch, then update and train on the sampled test cases SAVE_EVERY = 100 # Frequency of saving checkpoints RESET_EVERY = 50 # Frequency of resetting the pool of test cases # DATA LOADING ARGS (shouldn't need to modify that much) harmbench_root_dir = '../../../' # This should point to the root directory of the HarmBench repository behaviors_path = os.path.join(harmbench_root_dir, 'data', 'behavior_datasets', 'extra_behavior_datasets', 'adv_training_behaviors.csv') with open(behaviors_path, 'r') as csvfile: reader = csv.DictReader(csvfile) behaviors = list(reader) behavior_id_to_target_path = os.path.join(harmbench_root_dir, 'data', 'optimizer_targets', 'extra_targets', 'adv_training_targets.json') with open(behavior_id_to_target_path) as json_file: behavior_id_to_target = json.load(json_file) behaviors_path = os.path.join(harmbench_root_dir, 'data', 'behavior_datasets', 'harmbench_behaviors_text_val.csv') with open(behaviors_path, 'r') as csvfile: reader = csv.DictReader(csvfile) behaviors_extra = list(reader) behavior_id_to_target_path = os.path.join(harmbench_root_dir, 'data', 'optimizer_targets', 'extra_targets', 'adv_training_val_targets.json') with open(behavior_id_to_target_path) as json_file: behavior_id_to_target_extra = json.load(json_file) behaviors.extend(behaviors_extra) behavior_id_to_target.update(behavior_id_to_target_extra) # TRAINING CODE logger = logging.getLogger(__name__) def main(): parser = H4ArgumentParser((ModelArguments, DataArguments, SFTConfig)) model_args, data_args, training_args = parser.parse() # Set seed for reproducibility set_seed(training_args.seed) ############### # Setup logging ############### logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], ) log_level = training_args.get_process_log_level() logger.setLevel(log_level) datasets.utils.logging.set_verbosity(log_level) transformers.utils.logging.set_verbosity(log_level) transformers.utils.logging.enable_default_handler() transformers.utils.logging.enable_explicit_format() # Log on each process a small summary logger.warning( f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" + f" distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" ) logger.info(f"Model parameters {model_args}") logger.info(f"Data parameters {data_args}") logger.info(f"Training/evaluation parameters {training_args}") ############### # Load datasets ############### raw_datasets = get_datasets(data_args, splits=data_args.dataset_splits) logger.info( f"Training on the following datasets and their proportions: {[split + ' : ' + str(dset.num_rows) for split, dset in raw_datasets.items()]}" ) # with open('merged_robust_45krefusing_v1_1_4_requests_only.json', 'r') as file: # additional_data = json.load(file) # dataset_dict = {'prompt': [], 'prompt_id': [], 'messages': []} # for item in additional_data: # dataset_dict['prompt'].append(item['prompt']) # messages = [{'content': item['prompt'], 'role': 'user'},{'content': item['refusing'], 'role': 'assistant'}] # dataset_dict['messages'].append(messages) # dataset_dict['prompt_id'].append(str(uuid.uuid4())) # # Now dataset_dict is in the correct format to be passed to Dataset.from_dict() # additional_dataset = Dataset.from_dict(dataset_dict) # logger.info(f"====== len(raw_datasets['train']): {len(raw_datasets['train'])}") # raw_datasets['train'] = concatenate_datasets([raw_datasets['train'], additional_dataset]) # logger.info(f"====== len(raw_datasets['train']): {len(raw_datasets['train'])}") ################ # Load tokenizer ################ tokenizer = get_tokenizer(model_args, data_args) before_tc, after_tc = get_test_case_wrapping(tokenizer) ##################### # Apply chat template ##################### raw_datasets = raw_datasets.map(apply_chat_template, fn_kwargs={"tokenizer": tokenizer, "task": "sft"}) train_dataset = raw_datasets["train"] eval_dataset = raw_datasets["test"] with training_args.main_process_first(desc="Log a few random samples from the processed training set"): for index in random.sample(range(len(raw_datasets["train"])), 3): logger.info(f"Sample {index} of the processed training set:\n\n{raw_datasets['train'][index]['text']}") ####################### # Load pretrained model ####################### logger.info("*** Load pretrained model ***") torch_dtype = ( model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype) ) quantization_config = get_quantization_config(model_args) model_kwargs = dict( revision=model_args.model_revision, trust_remote_code=model_args.trust_remote_code, use_flash_attention_2=model_args.use_flash_attention_2, torch_dtype=torch_dtype, use_cache=False if training_args.gradient_checkpointing else True, device_map=get_kbit_device_map() if quantization_config is not None else None, quantization_config=quantization_config, ) logger.info("*** Model loaded! ***") ######################## # Initialize the Trainer ######################## trainer = CustomSFTTrainer( behaviors, behavior_id_to_target, before_tc=before_tc, after_tc=after_tc, num_gcg_steps=NUM_GCG_STEPS, gcg_search_width=GCG_SEARCH_WIDTH, num_test_cases_to_update_per_step=NUM_TEST_CASES_TO_UPDATE_PER_STEP, num_accelerate_gpus=NUM_ACCELERATE_GPUS, num_single_behavior_test_cases=NUM_SINGLE_BEHAVIOR_TEST_CASES, num_multi_behavior_test_cases=10, # Unsupported save_every=SAVE_EVERY, reset_every=RESET_EVERY, model=model_args.model_name_or_path, model_init_kwargs=model_kwargs, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, dataset_text_field="text", max_seq_length=training_args.max_seq_length, tokenizer=tokenizer, packing=True, peft_config=get_peft_config(model_args), ) ############### # Training loop ############### logger.info("*** Train ***") train_result = trainer.train() metrics = train_result.metrics max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) metrics["train_samples"] = min(max_train_samples, len(train_dataset)) trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) trainer.save_state() ########## # Evaluate ########## if training_args.do_eval: logger.info("*** Evaluate ***") metrics = trainer.evaluate() max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) trainer.log_metrics("eval", metrics) trainer.save_metrics("eval", metrics) ################################## # Save model and create model card ################################## logger.info("*** Save model ***") trainer.save_model(training_args.output_dir) logger.info(f"Model saved to {training_args.output_dir}") # Save everything else on main process accelerator = trainer.accelerator if accelerator.is_main_process: kwargs = { "finetuned_from": model_args.model_name_or_path, "dataset": list(data_args.dataset_mixer.keys()), "dataset_tags": list(data_args.dataset_mixer.keys()), "tags": ["alignment-handbook"], } trainer.create_model_card(**kwargs) # Restore k,v cache for fast inference trainer.model.config.use_cache = True trainer.model.config.save_pretrained(training_args.output_dir) if training_args.push_to_hub is True: logger.info("Pushing to hub...") trainer.push_to_hub() accelerator.wait_for_everyone() accelerator.end_training() if __name__ == "__main__": main()