初始化项目，由ModelHub XC社区提供模型
Model: mnoukhov/pythia410m-sft-tldr Source: Original Platform
This commit is contained in:
ModelHub XC
317
code/rl_training_value_model.py
Normal file
317
code/rl_training_value_model.py
Normal file
@@ -0,0 +1,317 @@
|
||||
# Copyright 2022 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.
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from accelerate import DistributedDataParallelKwargs
|
||||
from datasets import load_dataset
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoTokenizer, HfArgumentParser, pipeline
|
||||
|
||||
from trl import PPOConfig, PPOTrainer, set_seed
|
||||
from trl.core import LengthSampler
|
||||
from trl.models.modeling_value_model import AutoModelForCausalLMWithValueModel
|
||||
|
||||
|
||||
# import copy
|
||||
# from torch_ema import ExponentialMovingAverage
|
||||
# from transforers import pipeline
|
||||
|
||||
# tqdm.pandas()
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScriptArguments:
|
||||
"""
|
||||
The name of the Casual LM model we wish to fine with PPO
|
||||
"""
|
||||
|
||||
# NOTE: gpt2 models use Conv1D instead of Linear layers which are not yet supported in 8 bit mode
|
||||
# models like gpt-neo* models are more suitable.
|
||||
model_name: Optional[str] = field(default="", metadata={"help": "the model name"})
|
||||
reward_model_name: Optional[str] = field(default="", metadata={"help": "the reward model name"})
|
||||
gold_reward_model_name: Optional[str] = field(default="", metadata={"help": "the reward model name"})
|
||||
# tokenizer_name: Optional[str] = field(default=None, metadata={"help": "the tokenizer name"})
|
||||
dataset_name: Optional[str] = field(
|
||||
default="CarperAI/openai_summarize_tldr", metadata={"help": "the dataset name"}
|
||||
)
|
||||
train_split: Optional[str] = field(
|
||||
default="train", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
|
||||
)
|
||||
log_with: Optional[str] = field(default="wandb", metadata={"help": "use 'wandb' to log with wandb"})
|
||||
learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
|
||||
mini_batch_size: Optional[int] = field(default=1, metadata={"help": "the PPO minibatch size"})
|
||||
batch_size: Optional[int] = field(default=32, metadata={"help": "the batch size"})
|
||||
ppo_epochs: Optional[int] = field(default=4, metadata={"help": "the number of ppo epochs"})
|
||||
gradient_accumulation_steps: Optional[int] = field(
|
||||
default=4, metadata={"help": "the number of gradient accumulation steps"}
|
||||
)
|
||||
adafactor: Optional[bool] = field(default=False, metadata={"help": "whether to use the adafactor optimizer"})
|
||||
early_stopping: Optional[bool] = field(default=False, metadata={"help": "whether to early stop"})
|
||||
target_kl: Optional[float] = field(default=0.1, metadata={"help": "kl target for early stopping"})
|
||||
reward_baseline: Optional[float] = field(
|
||||
default=0.0,
|
||||
metadata={"help": "a baseline value that is subtracted from the reward"},
|
||||
)
|
||||
batched_gen: Optional[bool] = field(default=False, metadata={"help": "whether to use the batched text gen"})
|
||||
eval_freq: Optional[int] = field(default=None, metadata={"help": "n steps to save the model"})
|
||||
save_freq: Optional[int] = field(default=None, metadata={"help": "n steps to save the model"})
|
||||
output_dir: Optional[str] = field(default="runs/", metadata={"help": "n steps to save the model"})
|
||||
seed: Optional[int] = field(default=0, metadata={"help": "the seed"})
|
||||
steps: Optional[int] = field(default=20000, metadata={"help": "number of epochs"})
|
||||
init_kl_coef: Optional[float] = field(
|
||||
default=0.05,
|
||||
metadata={"help": "Initial KL penalty coefficient (used for adaptive and linear control)"},
|
||||
)
|
||||
adap_kl_ctrl: Optional[bool] = field(default=True, metadata={"help": "Use adaptive KL control, otherwise linear"})
|
||||
# multi_adapter_value: Optional[bool] = field(default=False)
|
||||
# separate_reward_model: Optional[str] = field(default=None, metadata={"help": "the reward model name"})
|
||||
|
||||
# Generation
|
||||
output_min_length: Optional[int] = field(default=24, metadata={"help": "the batch size"})
|
||||
output_max_length: Optional[int] = field(default=48, metadata={"help": "the batch size"})
|
||||
input_max_length: Optional[int] = field(default=512, metadata={"help": "maximum length for generation"})
|
||||
|
||||
# Quantization
|
||||
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"})
|
||||
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={},
|
||||
)
|
||||
|
||||
# LoRA
|
||||
use_lora: Optional[bool] = field(
|
||||
default=True,
|
||||
)
|
||||
lora_alpha: Optional[float] = field(default=32, 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"})
|
||||
# # EMA stuff
|
||||
# ema_decay: Optional[float] = field(default=0.995, metadata={"help": "the ema decay rate"})
|
||||
# reset_freq: Optional[int] = field(default=None, metadata={"help": "reset every n epochs"})
|
||||
|
||||
|
||||
def create_and_prepare_model(args):
|
||||
if args.bf16:
|
||||
torch_dtype = torch.bfloat16
|
||||
elif args.fp16_model:
|
||||
torch_dtype = torch.float16
|
||||
else:
|
||||
torch_dtype = torch.float32
|
||||
|
||||
model = AutoModelForCausalLMWithValueModel.from_pretrained(
|
||||
args.model_name,
|
||||
args.reward_model_name,
|
||||
torch_dtype=torch_dtype,
|
||||
)
|
||||
|
||||
# if script_args.ignore_bias_buffers:
|
||||
# torch distributed hack
|
||||
|
||||
model.config.torch_dtype = torch_dtype
|
||||
model.config.use_cache = True
|
||||
# model.config.use_cache = not args.gradient_checkpointing
|
||||
|
||||
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 create_and_prepare_dataset(args, tokenizer, split, num_proc=2):
|
||||
dataset = load_dataset(args.dataset_name, split=split)
|
||||
|
||||
dataset = dataset.rename_column("prompt", "query")
|
||||
original_columns = dataset.column_names
|
||||
original_columns.remove("query")
|
||||
|
||||
dataset = dataset.map(
|
||||
lambda examples: tokenizer(examples["query"], truncation=True, max_length=args.input_max_length),
|
||||
batched=True,
|
||||
num_proc=num_proc,
|
||||
remove_columns=original_columns,
|
||||
)
|
||||
|
||||
dataset.set_format("torch")
|
||||
return dataset
|
||||
|
||||
|
||||
def collator(data):
|
||||
return dict((key, [d[key] for d in data]) for key in data[0])
|
||||
|
||||
|
||||
parser = HfArgumentParser(ScriptArguments)
|
||||
script_args: ScriptArguments = parser.parse_args_into_dataclasses()[0]
|
||||
config = PPOConfig(
|
||||
steps=script_args.steps,
|
||||
model_name=script_args.model_name,
|
||||
learning_rate=script_args.learning_rate,
|
||||
log_with=script_args.log_with,
|
||||
batch_size=script_args.batch_size,
|
||||
mini_batch_size=script_args.mini_batch_size,
|
||||
gradient_accumulation_steps=script_args.gradient_accumulation_steps,
|
||||
optimize_cuda_cache=True,
|
||||
early_stopping=script_args.early_stopping,
|
||||
target_kl=script_args.target_kl,
|
||||
ppo_epochs=script_args.ppo_epochs,
|
||||
seed=script_args.seed,
|
||||
init_kl_coef=script_args.init_kl_coef,
|
||||
adap_kl_ctrl=script_args.adap_kl_ctrl,
|
||||
accelerator_kwargs={"kwargs_handlers": [DistributedDataParallelKwargs(find_unused_parameters=False)]},
|
||||
)
|
||||
|
||||
# set seed before initializing value head for deterministic eval
|
||||
set_seed(config.seed)
|
||||
|
||||
model, tokenizer = create_and_prepare_model(script_args)
|
||||
train_dataset = create_and_prepare_dataset(script_args, tokenizer, script_args.train_split)
|
||||
# eval_dataset = create_and_prepare_dataset(script_args, tokenizer, script_args.eval_split)
|
||||
|
||||
# We then build the PPOTrainer, passing the model, the reference model, the tokenizer
|
||||
ppo_trainer = PPOTrainer(
|
||||
config,
|
||||
model,
|
||||
ref_model=None,
|
||||
tokenizer=tokenizer,
|
||||
dataset=train_dataset,
|
||||
data_collator=collator,
|
||||
)
|
||||
|
||||
model.eval()
|
||||
|
||||
# if script_args.separate_reward_model:
|
||||
device = ppo_trainer.accelerator.device
|
||||
if ppo_trainer.accelerator.num_processes == 1:
|
||||
device = 0 if torch.cuda.is_available() else "cpu" # to avoid a ` pipeline` bug
|
||||
|
||||
reward_pipe = pipeline(
|
||||
"sentiment-analysis",
|
||||
model=script_args.reward_model_name,
|
||||
# device_map={"": Accelerator().local_process_index},
|
||||
tokenizer=tokenizer,
|
||||
return_token_type_ids=False,
|
||||
)
|
||||
if script_args.eval_freq is not None:
|
||||
gold_reward_pipe = pipeline(
|
||||
"sentiment-analysis",
|
||||
model=script_args.gold_reward_model_name,
|
||||
# device_map={"": Accelerator().local_process_index},
|
||||
# model_kwargs={"load_in_8bit": True},
|
||||
tokenizer=tokenizer,
|
||||
return_token_type_ids=False,
|
||||
)
|
||||
sent_kwargs = {
|
||||
"top_k": None,
|
||||
"function_to_apply": "none",
|
||||
"batch_size": 16,
|
||||
"truncation": True,
|
||||
}
|
||||
# We then define the arguments to pass to the `generate` function. These arguments
|
||||
# are passed to the `generate` function of the PPOTrainer, which is a wrapper around
|
||||
# the `generate` function of the trained model.
|
||||
generation_kwargs = {
|
||||
"min_length": -1,
|
||||
"top_k": 0.0,
|
||||
"top_p": 1.0,
|
||||
"do_sample": True,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
"eos_token_id": tokenizer.eos_token_id,
|
||||
}
|
||||
output_length_sampler = LengthSampler(script_args.output_min_length, script_args.output_max_length)
|
||||
|
||||
for epoch, batch in tqdm(
|
||||
enumerate(ppo_trainer.dataloader),
|
||||
total=config.total_ppo_epochs,
|
||||
disable=not ppo_trainer.accelerator.is_local_main_process,
|
||||
):
|
||||
if epoch >= config.total_ppo_epochs:
|
||||
break
|
||||
|
||||
question_tensors = batch["input_ids"]
|
||||
|
||||
query_response_tensors = ppo_trainer.generate(
|
||||
question_tensors,
|
||||
return_prompt=True,
|
||||
length_sampler=output_length_sampler,
|
||||
**generation_kwargs,
|
||||
)
|
||||
response_tensors = [tensor[len(question) :] for tensor, question in zip(query_response_tensors, question_tensors)]
|
||||
batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)
|
||||
|
||||
# Compute sentiment score
|
||||
texts = [q + r for q, r in zip(batch["query"], batch["response"])]
|
||||
reward_inputs = tokenizer(
|
||||
texts, padding=True, truncation=True, return_tensors="pt", return_token_type_ids=False
|
||||
).to(ppo_trainer.accelerator.device)
|
||||
|
||||
# for tensor in reponse_tensors:
|
||||
|
||||
for i, tensor in enumerate(query_response_tensors):
|
||||
if not torch.equal(tensor, reward_inputs["input_ids"][i][: reward_inputs["attention_mask"][i].sum()]):
|
||||
#TODO
|
||||
|
||||
import pdb
|
||||
|
||||
pdb.set_trace()
|
||||
# if script_args.reward_model_name:
|
||||
pipe_outputs = reward_pipe(texts, **sent_kwargs)
|
||||
rewards = [torch.tensor(output[0]["score"]) for output in pipe_outputs]
|
||||
# else:
|
||||
# raw_rewards = ppo_trainer.compute_reward_model_score(**reward_inputs)
|
||||
# rewards = [(raw_rewards[i] - script_args.reward_baseline) for i in range(len(raw_rewards))]
|
||||
|
||||
# Run PPO step
|
||||
stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
|
||||
ppo_trainer.log_stats(stats, batch, rewards)
|
||||
|
||||
# if ema is not None:
|
||||
# ema.update()
|
||||
#
|
||||
# if script_args.reset_freq and epoch and epoch % script_args.reset_freq == 0:
|
||||
# ema.copy_to()
|
||||
# ema.load_state_dict(initial_state_dict)
|
||||
# ppo_trainer.accelerator.print("elastic reset")
|
||||
|
||||
if script_args.eval_freq and epoch % script_args.eval_freq == 0:
|
||||
if ppo_trainer.accelerator.is_main_process:
|
||||
pipe_outputs = gold_reward_pipe(texts, **sent_kwargs)
|
||||
rewards = [torch.tensor(output[0]["score"]) for output in pipe_outputs]
|
||||
logs = {}
|
||||
logs["env/gold_reward_mean"] = torch.mean(rewards).cpu().numpy().item()
|
||||
logs["env/gold_reward_std"] = torch.std(rewards).cpu().numpy().item()
|
||||
logs["env/gold_reward_dist"] = rewards.cpu().numpy()
|
||||
ppo_trainer.accelerator.log(logs)
|
||||
print(logs)
|
||||
|
||||
if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:
|
||||
ppo_trainer.save_pretrained(script_args.output_dir + f"step_{epoch}")
|
||||
Reference in New Issue
Block a user