mirror of
https://github.com/wassname/SimPO.git
synced 2026-06-27 20:19:50 +08:00
Yu Meng
317 lines
13 KiB
Python
317 lines
13 KiB
Python
#!/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.
|
|
import logging
|
|
import random
|
|
import sys
|
|
|
|
import torch
|
|
import transformers
|
|
from transformers import AutoModelForCausalLM, set_seed
|
|
|
|
from alignment import (
|
|
DataArguments,
|
|
DPOConfig,
|
|
H4ArgumentParser,
|
|
ModelArguments,
|
|
get_checkpoint,
|
|
get_datasets,
|
|
get_kbit_device_map,
|
|
get_peft_config,
|
|
get_quantization_config,
|
|
get_tokenizer,
|
|
is_adapter_model,
|
|
)
|
|
from alignment.data import maybe_insert_system_message, is_openai_format
|
|
from peft import PeftConfig, PeftModel
|
|
from simpo_trainer import SimPOTrainer
|
|
from simpo_config import SimPOConfig
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional, Literal
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MISTRAL_CHAT_TEMPLATE = "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'].strip() + '\n\n' %}{% else %}{% set loop_messages = messages %}{% set system_message = '' %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 %}{% set content = system_message + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}"
|
|
|
|
def apply_chat_template(
|
|
example,
|
|
tokenizer,
|
|
task: Literal["sft", "generation", "rm", "simpo"],
|
|
auto_insert_empty_system_msg: bool = True,
|
|
change_template = None,
|
|
):
|
|
if change_template == "mistral":
|
|
tokenizer.chat_template = MISTRAL_CHAT_TEMPLATE
|
|
if task in ["sft", "generation"]:
|
|
messages = example["messages"]
|
|
# We add an empty system message if there is none
|
|
if auto_insert_empty_system_msg:
|
|
maybe_insert_system_message(messages, tokenizer)
|
|
example["text"] = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=True if task == "generation" else False,
|
|
)
|
|
elif task == "rm":
|
|
if all(k in example.keys() for k in ("chosen", "rejected")):
|
|
chosen_messages = example["chosen"]
|
|
rejected_messages = example["rejected"]
|
|
# We add an empty system message if there is none
|
|
if auto_insert_empty_system_msg:
|
|
maybe_insert_system_message(chosen_messages, tokenizer)
|
|
maybe_insert_system_message(rejected_messages, tokenizer)
|
|
|
|
example["text_chosen"] = tokenizer.apply_chat_template(chosen_messages, tokenize=False)
|
|
example["text_rejected"] = tokenizer.apply_chat_template(rejected_messages, tokenize=False)
|
|
else:
|
|
raise ValueError(
|
|
f"Could not format example as dialogue for `rm` task! Require `[chosen, rejected]` keys but found {list(example.keys())}"
|
|
)
|
|
elif task == "simpo":
|
|
if all(k in example.keys() for k in ("chosen", "rejected")):
|
|
if not is_openai_format(example["chosen"]) or not is_openai_format(example["rejected"]):
|
|
raise ValueError(
|
|
f"Could not format example as dialogue for `{task}` task! Require OpenAI format for all messages"
|
|
)
|
|
|
|
# For DPO/ORPO, the inputs are triples of (prompt, chosen, rejected), where `chosen` and `rejected` are the final turn of a dialogue
|
|
# We therefore need to extract the N-1 turns to form the prompt
|
|
if "prompt" in example and is_openai_format(example["prompt"]):
|
|
prompt_messages = example["prompt"]
|
|
chosen_messages = example["chosen"]
|
|
rejected_messages = example["rejected"]
|
|
else:
|
|
prompt_messages = example["chosen"][:-1]
|
|
# Now we extract the final turn to define chosen/rejected responses
|
|
chosen_messages = example["chosen"][-1:]
|
|
rejected_messages = example["rejected"][-1:]
|
|
|
|
# Prepend a system message if the first message is not a system message
|
|
if auto_insert_empty_system_msg:
|
|
maybe_insert_system_message(prompt_messages, tokenizer)
|
|
|
|
example["text_prompt"] = tokenizer.apply_chat_template(prompt_messages, tokenize=False)
|
|
example["text_chosen"] = tokenizer.apply_chat_template(chosen_messages, tokenize=False)
|
|
if example["text_chosen"].startswith(tokenizer.bos_token):
|
|
example["text_chosen"] = example["text_chosen"][len(tokenizer.bos_token):]
|
|
example["text_rejected"] = tokenizer.apply_chat_template(rejected_messages, tokenize=False)
|
|
if example["text_rejected"].startswith(tokenizer.bos_token):
|
|
example["text_rejected"] = example["text_rejected"][len(tokenizer.bos_token):]
|
|
else:
|
|
raise ValueError(
|
|
f"Could not format example as dialogue for `{task}` task! Require either the "
|
|
f"`[chosen, rejected]` or `[prompt, chosen, rejected]` keys but found {list(example.keys())}"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Task {task} not supported, please ensure that the provided task is one of ['sft', 'generation', 'rm', 'dpo', 'orpo']"
|
|
)
|
|
return example
|
|
|
|
|
|
def main():
|
|
parser = H4ArgumentParser((ModelArguments, DataArguments, SimPOConfig))
|
|
model_args, data_args, training_args = parser.parse()
|
|
|
|
#######
|
|
# Setup
|
|
#######
|
|
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)
|
|
transformers.utils.logging.set_verbosity(log_level)
|
|
transformers.utils.logging.enable_default_handler()
|
|
transformers.utils.logging.enable_explicit_format()
|
|
|
|
# Log on each process the small summary:
|
|
logger.info(f"Model parameters {model_args}")
|
|
logger.info(f"Data parameters {data_args}")
|
|
logger.info(f"Training/evaluation parameters {training_args}")
|
|
|
|
# Check for last checkpoint
|
|
last_checkpoint = get_checkpoint(training_args)
|
|
if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
|
|
logger.info(f"Checkpoint detected, resuming training at {last_checkpoint=}.")
|
|
|
|
# Set seed for reproducibility
|
|
set_seed(training_args.seed)
|
|
|
|
###############
|
|
# Load datasets
|
|
###############
|
|
raw_datasets = get_datasets(
|
|
data_args,
|
|
splits=data_args.dataset_splits,
|
|
configs=data_args.dataset_configs,
|
|
columns_to_keep=["messages", "chosen", "rejected", "prompt", "completion", "label"],
|
|
# seed=training_args.seed,
|
|
)
|
|
logger.info(
|
|
f"Training on the following splits: {[split + ' : ' + str(dset.num_rows) for split, dset in raw_datasets.items()]}"
|
|
)
|
|
column_names = list(raw_datasets["train"].features)
|
|
|
|
#####################################
|
|
# Load tokenizer and process datasets
|
|
#####################################
|
|
data_args.truncation_side = "left" # Truncate from left to ensure we don't lose labels in final turn
|
|
tokenizer = get_tokenizer(model_args, data_args)
|
|
|
|
if "mistral" in model_args.model_name_or_path.lower():
|
|
change_template = "mistral"
|
|
else:
|
|
change_template = None
|
|
#####################
|
|
# Apply chat template
|
|
#####################
|
|
raw_datasets = raw_datasets.map(
|
|
apply_chat_template,
|
|
fn_kwargs={
|
|
"tokenizer": tokenizer,
|
|
"task": "simpo",
|
|
"auto_insert_empty_system_msg": data_args.auto_insert_empty_system_msg,
|
|
"change_template": change_template,
|
|
},
|
|
num_proc=data_args.preprocessing_num_workers,
|
|
remove_columns=column_names,
|
|
desc="Formatting comparisons with prompt template",
|
|
)
|
|
|
|
# Replace column names with what TRL needs, text_chosen -> chosen and text_rejected -> rejected
|
|
for split in ["train", "test"]:
|
|
raw_datasets[split] = raw_datasets[split].rename_columns(
|
|
{"text_prompt": "prompt", "text_chosen": "chosen", "text_rejected": "rejected"}
|
|
)
|
|
|
|
# Log a few random samples from the training set:
|
|
for index in random.sample(range(len(raw_datasets["train"])), 3):
|
|
logger.info(f"Prompt sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['prompt']}")
|
|
logger.info(f"Chosen sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['chosen']}")
|
|
logger.info(f"Rejected sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['rejected']}")
|
|
|
|
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,
|
|
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,
|
|
attn_implementation=model_args.attn_implementation,
|
|
)
|
|
|
|
model = model_args.model_name_or_path
|
|
# seems to require internet
|
|
# if is_adapter_model(model, model_args.model_revision) is True:
|
|
# logger.info(f"Loading SFT adapter for {model_args.model_name_or_path=}")
|
|
# peft_config = PeftConfig.from_pretrained(model_args.model_name_or_path, revision=model_args.model_revision)
|
|
# model_kwargs = dict(
|
|
# revision=model_args.base_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,
|
|
# )
|
|
# base_model = AutoModelForCausalLM.from_pretrained(
|
|
# peft_config.base_model_name_or_path,
|
|
# **model_kwargs,
|
|
# )
|
|
# model = PeftModel.from_pretrained(
|
|
# base_model,
|
|
# model_args.model_name_or_path,
|
|
# revision=model_args.model_revision,
|
|
# )
|
|
# model_kwargs = None
|
|
|
|
training_args.model_init_kwargs = model_kwargs
|
|
#########################
|
|
# Instantiate SimPO trainer
|
|
#########################
|
|
trainer = SimPOTrainer(
|
|
model=model,
|
|
args=training_args,
|
|
train_dataset=raw_datasets["train"],
|
|
eval_dataset=raw_datasets["test"],
|
|
tokenizer=tokenizer,
|
|
peft_config=get_peft_config(model_args),
|
|
)
|
|
|
|
###############
|
|
# Training loop
|
|
###############
|
|
checkpoint = None
|
|
if training_args.resume_from_checkpoint is not None:
|
|
checkpoint = training_args.resume_from_checkpoint
|
|
elif last_checkpoint is not None:
|
|
checkpoint = last_checkpoint
|
|
train_result = trainer.train(resume_from_checkpoint=checkpoint)
|
|
metrics = train_result.metrics
|
|
metrics["train_samples"] = len(raw_datasets["train"])
|
|
trainer.log_metrics("train", metrics)
|
|
trainer.save_metrics("train", metrics)
|
|
trainer.save_state()
|
|
|
|
logger.info("*** Training complete ***")
|
|
|
|
##################################
|
|
# 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
|
|
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"],
|
|
}
|
|
if trainer.accelerator.is_main_process:
|
|
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)
|
|
|
|
##########
|
|
# Evaluate
|
|
##########
|
|
if training_args.do_eval:
|
|
logger.info("*** Evaluate ***")
|
|
metrics = trainer.evaluate()
|
|
metrics["eval_samples"] = len(raw_datasets["test"])
|
|
trainer.log_metrics("eval", metrics)
|
|
trainer.save_metrics("eval", metrics)
|
|
|
|
if training_args.push_to_hub is True:
|
|
logger.info("Pushing to hub...")
|
|
trainer.push_to_hub(**kwargs)
|
|
|
|
logger.info("*** Training complete! ***")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|