mnoukhov/pythia410m-sft-tldr
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: mnoukhov/pythia410m-sft-tldr
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-05-30 14:21:21 +08:00
commit 4bb3a5553f
215 changed files with 118927 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

505
code/gpt_reward_modeling.py Normal file
View File

@@ -0,0 +1,505 @@
# 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 os
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Union
import bitsandbytes as bnb
import torch
from accelerate import Accelerator
from datasets import DatasetDict, builder, load_dataset
from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
from peft.tuners.lora import LoraLayer
from torch import nn
from tqdm import tqdm
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
BitsAndBytesConfig,
HfArgumentParser,
PreTrainedTokenizerBase,
TrainingArguments,
)
from trl import RewardTrainer
tqdm.pandas()
builder.has_sufficient_disk_space = lambda needed_bytes, directory=".": True
# torch.autograd.set_detect_anomaly(True)
### fix from https://github.com/huggingface/trl/issues/274
class GPTRewardTrainer(RewardTrainer):
def compute_loss(
self,
model,
inputs,
return_outputs=False,
):
rewards = model(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
)[0]
bsz = rewards.size(0)
jidx = torch.arange(0, bsz, 2)
kidx = jidx + 1
rewards_chosen = rewards[jidx]
rewards_rejected = rewards[kidx]
loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
if return_outputs:
return loss, {
"rewards_chosen": rewards_chosen,
"rewards_rejected": rewards_rejected,
}
return loss
@dataclass
class GPTRewardDataCollatorWithPadding:
r"""
Reward DataCollator class that pads the inputs to the maximum length of the batch.
Args:
tokenizer (`PreTrainedTokenizerBase`):
The tokenizer used for encoding the data.
padding (`Union[bool, str, `PaddingStrategy`]`, `optional`, defaults to `True`):
padding_strategy to pass to the tokenizer.
max_length (`Optional[int]`, `optional`, defaults to `None`):
The maximum length of the sequence to be processed.
pad_to_multiple_of (`Optional[int]`, `optional`, defaults to `None`):
If set will pad the sequence to a multiple of the provided value.
return_tensors (`str`, `optional`, defaults to `"pt"`):
The tensor type to use.
"""
tokenizer: PreTrainedTokenizerBase
padding: Union[bool, str] = True
max_length: Optional[int] = None
pad_to_multiple_of: Optional[int] = None
return_tensors: str = "pt"
def __call__(self, features):
# features_chosen = []
# features_rejected = []
merged_features = []
for feature in features:
# check if the keys are named as expected
if (
"input_ids_chosen" not in feature
or "input_ids_rejected" not in feature
or "attention_mask_chosen" not in feature
or "attention_mask_rejected" not in feature
):
raise ValueError(
"The features should include `input_ids_chosen`, `attention_mask_chosen`, `input_ids_rejected` and `attention_mask_rejected`"
)
merged_features.append(
{
"input_ids": feature["input_ids_chosen"],
"attention_mask": feature["attention_mask_chosen"],
}
)
merged_features.append(
{
"input_ids": feature["input_ids_rejected"],
"attention_mask": feature["attention_mask_rejected"],
}
)
batch = self.tokenizer.pad(
merged_features,
padding=self.padding,
max_length=self.max_length,
pad_to_multiple_of=self.pad_to_multiple_of,
return_tensors=self.return_tensors,
)
batch = {
"input_ids": batch["input_ids"],
"attention_mask": batch["attention_mask"],
"return_loss": True,
}
return batch
# Define and parse arguments.
@dataclass
class ScriptArguments:
"""
The name of the Casual LM model we wish to fine with RewardTrainer
"""
model_name: Optional[str] = field(
default="/home/toolkit/huggingface/tldr_sft_pythia7b", metadata={"help": "the model name"}
)
dataset_name: Optional[str] = field(
default="mnoukhov/openai_summarize_comparisons_tldrprompt", metadata={"help": "the dataset name"}
)
dataset_text_field: Optional[str] = field(default="prompt", metadata={"help": "the text field of the dataset"})
log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
logging_steps: Optional[int] = field(default=100, metadata={"help": "the number of update steps between two logs"})
train_split: Optional[str] = field(
default="train", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
)
eval_split: Optional[str] = field(
default="test[:5000]", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
)
learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
weight_decay: Optional[float] = field(default=0.001)
num_warmup_steps: Optional[int] = field(default=100)
lr_scheduler_type: Optional[str] = field(default="cosine")
optimizer_type: Optional[str] = field(default="adamw_torch", metadata={"help": "the optimizer type"})
per_device_train_batch_size: Optional[int] = field(default=2, metadata={"help": "the per device train batch size"})
per_device_eval_batch_size: Optional[int] = field(default=1, metadata={"help": "the per device eval batch size"})
num_train_epochs: Optional[int] = field(default=1, metadata={"help": "the number of training epochs"})
seq_length: Optional[int] = field(default=560, metadata={"help": "Input sequence length"})
gradient_accumulation_steps: Optional[int] = field(
default=16, metadata={"help": "the number of gradient accumulation steps"}
)
bf16: Optional[bool] = field(
default=False,
metadata={
"help": "This essentially cuts the training time in half if you want to sacrifice a little precision and have a supported GPU."
},
)
fp16: Optional[bool] = field(
default=False,
metadata={
"help": "This essentially cuts the training time in half if you want to sacrifice a little precision and have a supported GPU."
},
)
fp16_model: Optional[bool] = field(
default=False,
metadata={},
)
load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 8 bits precision"})
load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
use_lora: Optional[bool] = field(
default=True,
)
lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"})
lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"})
lora_r: Optional[int] = field(default=8, metadata={"help": "the lora r parameter"})
lora_all_linear: Optional[bool] = field(default=False, metadata={"help": "lora adapter on all linear layers"})
trust_remote_code: Optional[bool] = field(default=True, metadata={"help": "Enable `trust_remote_code`"})
output_dir: Optional[str] = field(default="results", metadata={"help": "the output directory"})
gradient_checkpointing: Optional[bool] = field(
default=False,
metadata={"help": "Enables gradient checkpointing."},
)
mode: Optional[str] = field(default="train")
eval_steps: Optional[float] = field(default=None)
pretrained_adapter: Optional[str] = field(default=None)
padding: Optional[str] = field(
default="max_length", metadata={"help": "padding to use for preprocessing the dataset"}
)
save_strategy: Optional[str] = field(default="steps")
def find_all_linear_names(args, model):
cls = bnb.nn.Linear4bit if args.load_in_4bit else (bnb.nn.Linear8bitLt if args.load_in_8bit else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split(".")
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if "lm_head" in lora_module_names: # needed for 16-bit
lora_module_names.remove("lm_head")
if "score" in lora_module_names: # needed for 16-bit
lora_module_names.remove("score")
return list(lora_module_names)
def create_and_prepare_model(args):
if args.load_in_8bit and args.load_in_4bit:
raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
elif args.load_in_8bit or args.load_in_4bit:
quantization_config = BitsAndBytesConfig(load_in_8bit=args.load_in_8bit, load_in_4bit=args.load_in_4bit)
device_map = {"": Accelerator().local_process_index}
else:
device_map = None
quantization_config = None
if args.bf16:
torch_dtype = torch.bfloat16
elif args.fp16_model:
torch_dtype = torch.float16
else:
torch_dtype = torch.float32
model = AutoModelForSequenceClassification.from_pretrained(
args.model_name,
quantization_config=quantization_config,
device_map=device_map,
num_labels=1,
torch_dtype=torch_dtype,
)
model.config.torch_dtype = torch_dtype
model.config.use_cache = not args.gradient_checkpointing
# if script_args.ignore_bias_buffers:
# torch distributed hack
if quantization_config is not None:
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=args.gradient_checkpointing)
args.gradient_checkpointing = False
if args.use_lora:
# we add `score` to the list of modules to save to
# correctly save the score head.
if args.pretrained_adapter is not None:
model = PeftModel.from_pretrained(model, args.pretrained_adapter)
else:
if args.lora_all_linear:
target_modules = find_all_linear_names(args, model)
else:
target_modules = None
peft_config = LoraConfig(
r=args.lora_r,
lora_alpha=args.lora_alpha,
lora_dropout=args.lora_dropout,
bias="none",
task_type="SEQ_CLS",
target_modules=target_modules,
modules_to_save=["score"],
)
model = get_peft_model(model, peft_config)
modules_to_save = ["score"]
for key, _ in model.named_modules():
target_module_found = any(key.endswith(target_key) for target_key in modules_to_save)
if target_module_found:
model.get_submodule(key + ".original_module").requires_grad_(False)
if torch_dtype == torch.bfloat16:
for name, module in model.named_modules():
if isinstance(module, LoraLayer):
module = module.to(torch_dtype)
if "norm" in name:
module = module.to(torch.float32)
if "score" in name or "embed_tokens" in name:
if hasattr(module, "weight") and module.weight.dtype == torch.float32:
module = module.to(torch_dtype)
tokenizer = AutoTokenizer.from_pretrained(script_args.model_name)
if getattr(tokenizer, "pad_token", None) is None:
tokenizer.pad_token = tokenizer.eos_token
if getattr(model.config, "pad_token_id", None) is None:
model.config.pad_token_id = model.config.eos_token_id
return model, tokenizer
def prepare_dataset(args, dataset, tokenizer, num_proc=2):
# def summary_filter(example):
# return (example["chosen"] != example["rejected"]) and (
# len(example["chosen"].split()) >= 5 or len(example["rejected"].split()) >= 5
# )
#
# pre_filter = len(dataset)
# dataset = dataset.filter(summary_filter)
# print(f"filtered {pre_filter - len(dataset)} samples from {split}")
original_columns = dataset.column_names
def preprocess_function(examples):
new_examples = {
"input_ids_chosen": [],
"attention_mask_chosen": [],
"input_ids_rejected": [],
"attention_mask_rejected": [],
}
for prompt, chosen, rejected in zip(examples["prompt"], examples["chosen"], examples["rejected"]):
tokenized_chosen = tokenizer(
prompt + " " + chosen, padding=args.padding, truncation=True, max_length=script_args.seq_length
)
tokenized_rejected = tokenizer(
prompt + " " + rejected, padding=args.padding, truncation=True, max_length=script_args.seq_length
)
new_examples["input_ids_chosen"].append(tokenized_chosen["input_ids"])
new_examples["attention_mask_chosen"].append(tokenized_chosen["attention_mask"])
new_examples["input_ids_rejected"].append(tokenized_rejected["input_ids"])
new_examples["attention_mask_rejected"].append(tokenized_rejected["attention_mask"])
return new_examples
dataset = dataset.map(preprocess_function, batched=True, num_proc=num_proc, remove_columns=original_columns)
return dataset
parser = HfArgumentParser(ScriptArguments)
script_args = parser.parse_args_into_dataclasses()[0]
model, tokenizer = create_and_prepare_model(script_args)
if script_args.mode != "eval":
train_data = load_dataset(script_args.dataset_name, split=script_args.train_split)
train_dataset = prepare_dataset(script_args, train_data, tokenizer)
else:
train_dataset = None
if script_args.eval_split is not None and script_args.eval_split != "None":
eval_data = load_dataset(script_args.dataset_name, split=script_args.eval_split)
eval_dataset = prepare_dataset(script_args, eval_data, tokenizer)
else:
eval_dataset = None
# don't include gradient_checkpointing here, see trl#728
training_args = TrainingArguments(
output_dir=script_args.output_dir,
per_device_train_batch_size=script_args.per_device_train_batch_size,
per_device_eval_batch_size=script_args.per_device_eval_batch_size,
bf16=script_args.bf16,
fp16=script_args.fp16,
num_train_epochs=script_args.num_train_epochs,
gradient_accumulation_steps=script_args.gradient_accumulation_steps,
learning_rate=script_args.learning_rate,
report_to=script_args.log_with,
remove_unused_columns=False,
lr_scheduler_type=script_args.lr_scheduler_type,
weight_decay=script_args.weight_decay,
optim=script_args.optimizer_type,
warmup_steps=script_args.num_warmup_steps,
logging_steps=script_args.logging_steps,
evaluation_strategy=("steps" if script_args.eval_steps is not None else "epoch"),
eval_steps=script_args.eval_steps,
save_strategy="epoch",
gradient_checkpointing=script_args.gradient_checkpointing,
ddp_find_unused_parameters=False,
)
data_collator = GPTRewardDataCollatorWithPadding(tokenizer, max_length=script_args.seq_length, pad_to_multiple_of=8)
trainer = GPTRewardTrainer(
model=model,
tokenizer=tokenizer,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
max_length=script_args.seq_length,
data_collator=data_collator,
)
if script_args.mode == "train":
print("Training")
trainer.train()
trainer.evaluate()
print("Saving last checkpoint of the model")
trainer.save_model(script_args.output_dir)
output_dir = os.path.join(script_args.output_dir, "final_checkpoint")
trainer.model.save_pretrained(output_dir)
elif script_args.mode == "eval":
print("Evaluating")
# results = trainer.evaluate()
results = trainer.evaluate()
print(results)
elif script_args.mode == "relabel":
def relabel_with_preds(batch: Dict[str, List]):
relabel_batch = {
"prompt": [],
"chosen": [],
"rejected": [],
"pred_chosen": [],
"pred_rejected": [],
}
for prompt, chosen, rejected, pred_chosen, pred_rejected in zip(
batch["prompt"],
batch["chosen"],
batch["rejected"],
batch["pred_chosen"],
batch["pred_rejected"],
):
relabel_batch["prompt"].append(prompt)
if pred_chosen >= pred_rejected:
relabel_batch["chosen"].append(chosen)
relabel_batch["rejected"].append(rejected)
relabel_batch["pred_chosen"].append(pred_chosen)
relabel_batch["pred_rejected"].append(pred_rejected)
else:
relabel_batch["chosen"].append(rejected)
relabel_batch["rejected"].append(chosen)
relabel_batch["pred_chosen"].append(pred_rejected)
relabel_batch["pred_rejected"].append(pred_chosen)
return relabel_batch
relabel_dataset = DatasetDict()
for split, pred_dataset in [("train", train_dataset), ("test", eval_dataset)]:
if pred_dataset is None:
continue
trainer.accelerator.print(f"Prediction {split}")
preds, _, metrics = trainer.predict(pred_dataset)
trainer.accelerator.print(f"metrics {metrics}")
if trainer.accelerator.is_local_main_process:
print("Relabelling Dataset and Saving")
ds_split = script_args.train_split if split == "train" else script_args.eval_split
dataset = load_dataset(script_args.dataset_name, split=ds_split)
dataset = dataset.add_column("pred_chosen", preds[:, 0])
dataset = dataset.add_column("pred_rejected", preds[:, 1])
dataset = dataset.map(relabel_with_preds, batched=True)
dataset._info.description = f"{script_args.dataset_name} relabelled with {script_args.model_name}"
relabel_dataset[split] = dataset
if trainer.accelerator.is_local_main_process:
print("Saving")
relabel_dataset.save_to_disk(script_args.output_dir)
print("Pushing")
relabel_dataset.push_to_hub(os.path.basename(script_args.output_dir))
elif script_args.mode == "predict":
relabel_dataset = DatasetDict()
for split, pred_dataset in [("train", train_dataset), ("test", eval_dataset)]:
if pred_dataset is None:
continue
trainer.accelerator.print(f"Prediction {split}")
preds, _, metrics = trainer.predict(pred_dataset)
trainer.accelerator.print(f"metrics {metrics}")
if trainer.accelerator.is_local_main_process:
print("Relabelling Dataset and Saving")
ds_split = script_args.train_split if split == "train" else script_args.eval_split
dataset = load_dataset(script_args.dataset_name, split=ds_split)
model_basename = script_args.model_name.rsplit("/", 1)[-1]
dataset = dataset.add_column(f"pred_chosen_{model_basename}", preds[:, 0])
dataset = dataset.add_column(f"pred_rejected_{model_basename}", preds[:, 1])
dataset._info.description = f"{script_args.dataset_name} relabelled with {script_args.model_name}"
relabel_dataset[split] = dataset
if trainer.accelerator.is_local_main_process:
print("Saving")
relabel_dataset.save_to_disk(script_args.output_dir)
print("Pushing")
relabel_dataset.push_to_hub(os.path.basename(script_args.output_dir))
else:
raise Exception(f"incorrect mode {script_args.mode}")
# TODO this freezes for some reason
# for split, dataset in relabel_dataset.items():
# if trainer.accelerator.is_local_main_process:
# eval_dataset = prepare_dataset(script_args, dataset, tokenizer)
# trainer.accelerator.print(f"Re-evaluating relabel {split} dataset of size {len(dataset)}")
# trainer.accelerator.wait_for_everyone()
# results = trainer.evaluate(eval_dataset)
# trainer.accelerator.print(results)