312 lines
12 KiB
Python
312 lines
12 KiB
Python
|
import os
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
import bitsandbytes as bnb
|
||
|
|
import torch
|
||
|
|
from accelerate import Accelerator
|
||
|
|
from datasets import load_dataset
|
||
|
|
from peft import AutoPeftModelForCausalLM, AutoPeftModelForSeq2SeqLM, LoraConfig
|
||
|
|
from tqdm import tqdm
|
||
|
|
from transformers import (
|
||
|
|
AutoModelForCausalLM,
|
||
|
|
AutoModelForSeq2SeqLM,
|
||
|
|
AutoTokenizer,
|
||
|
|
BitsAndBytesConfig,
|
||
|
|
GPT2Model,
|
||
|
|
HfArgumentParser,
|
||
|
|
TrainingArguments,
|
||
|
|
)
|
||
|
|
from transformers.pytorch_utils import Conv1D
|
||
|
|
from transformers.trainer_utils import get_last_checkpoint
|
||
|
|
|
||
|
|
from trl import DataCollatorForCompletionOnlyLM, SFTTrainer
|
||
|
|
|
||
|
|
|
||
|
|
tqdm.pandas()
|
||
|
|
|
||
|
|
|
||
|
|
# Define and parse arguments.
|
||
|
|
@dataclass
|
||
|
|
class ScriptArguments:
|
||
|
|
"""
|
||
|
|
The name of the Casual LM model we wish to fine with SFTTrainer
|
||
|
|
"""
|
||
|
|
|
||
|
|
model_name: Optional[str] = field(default="EleutherAI/pythia-6.9b-deduped", metadata={"help": "the model name"})
|
||
|
|
tokenizer_name: Optional[str] = field(default=None, metadata={"help": "the model 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)"}
|
||
|
|
)
|
||
|
|
eval_split: Optional[str] = field(
|
||
|
|
default="test",#"valid[:2000]",
|
||
|
|
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"})
|
||
|
|
streaming: Optional[bool] = field(default=False, metadata={"help": "whether to stream the dataset"})
|
||
|
|
shuffle_buffer: Optional[int] = field(default=5000, metadata={"help": "the shuffle buffer size"})
|
||
|
|
|
||
|
|
learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
|
||
|
|
lr_scheduler_type: Optional[str] = field(default="cosine")
|
||
|
|
num_warmup_steps: Optional[int] = field(default=100)
|
||
|
|
weight_decay: Optional[float] = field(default=0.05)
|
||
|
|
optimizer_type: Optional[str] = field(default="paged_adamw_32bit", metadata={"help": "the optimizer type"})
|
||
|
|
|
||
|
|
max_steps: Optional[int] = field(default=-1, metadata={"help": "the number of training steps"})
|
||
|
|
num_train_epochs: Optional[int] = field(default=1, metadata={"help": "the number of training epochs"})
|
||
|
|
per_device_train_batch_size: Optional[int] = field(
|
||
|
|
default=16, 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"})
|
||
|
|
gradient_accumulation_steps: Optional[int] = field(
|
||
|
|
default=16, metadata={"help": "the number of gradient accumulation steps"}
|
||
|
|
)
|
||
|
|
gradient_checkpointing: Optional[bool] = field(
|
||
|
|
default=False, metadata={"help": "whether to use gradient checkpointing"}
|
||
|
|
)
|
||
|
|
seq_length: Optional[int] = field(default=560, metadata={"help": "Input sequence length"})
|
||
|
|
|
||
|
|
load_in_8bit: Optional[bool] = field(default=True, 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_peft: Optional[bool] = field(default=True, metadata={"help": "Wether to use PEFT or not to train adapters"})
|
||
|
|
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"})
|
||
|
|
trust_remote_code: Optional[bool] = field(default=True, metadata={"help": "Enable `trust_remote_code`"})
|
||
|
|
bf16: Optional[bool] = field(default=True)
|
||
|
|
fp16_model: Optional[bool] = field(
|
||
|
|
default=False,
|
||
|
|
metadata={},
|
||
|
|
)
|
||
|
|
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."
|
||
|
|
},
|
||
|
|
)
|
||
|
|
train_completions: Optional[bool] = field(default=False)
|
||
|
|
packing: Optional[bool] = field(default=True)
|
||
|
|
|
||
|
|
output_dir: Optional[str] = field(default="./results", metadata={"help": "the output directory"})
|
||
|
|
output_model_name: Optional[str] = field(default=None, metadata={"help": "the model pushed to hub"})
|
||
|
|
logging_steps: Optional[int] = field(default=10, metadata={"help": "the number of logging steps"})
|
||
|
|
eval_steps: Optional[int] = field(default=1000, metadata={"help": "the number of steps to eval at"})
|
||
|
|
save_steps: Optional[int] = field(default=1000, metadata={"help": "the number of steps to save at"})
|
||
|
|
save_strategy: Optional[str] = field(default="steps")
|
||
|
|
seed: Optional[int] = field(default=0)
|
||
|
|
just_eval: Optional[bool] = field(default=False)
|
||
|
|
resume_from_checkpoint: Optional[str] = field(default=None)
|
||
|
|
|
||
|
|
|
||
|
|
def chars_token_ratio(dataset, tokenizer, nb_examples=400):
|
||
|
|
"""
|
||
|
|
Estimate the average number of characters per token in the dataset.
|
||
|
|
"""
|
||
|
|
total_characters, total_tokens = 0, 0
|
||
|
|
for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples):
|
||
|
|
text = prepare_sample_text(example)
|
||
|
|
total_characters += len(text)
|
||
|
|
if tokenizer.is_fast:
|
||
|
|
total_tokens += len(tokenizer(text).tokens())
|
||
|
|
else:
|
||
|
|
total_tokens += len(tokenizer.tokenize(text))
|
||
|
|
|
||
|
|
return total_characters / total_tokens
|
||
|
|
|
||
|
|
|
||
|
|
def prepare_sample_text(examples):
|
||
|
|
if isinstance(examples["chosen"], str):
|
||
|
|
return examples["prompt"] + examples["chosen"]
|
||
|
|
elif isinstance(examples["chosen"], list):
|
||
|
|
return list(map(str.__add__, examples["prompt"], examples["chosen"]))
|
||
|
|
else:
|
||
|
|
raise Exception(f"weird input examples of type {type(examples)}")
|
||
|
|
|
||
|
|
|
||
|
|
def create_datasets(args):
|
||
|
|
train_data = load_dataset(
|
||
|
|
args.dataset_name,
|
||
|
|
split=args.train_split,
|
||
|
|
streaming=args.streaming,
|
||
|
|
)
|
||
|
|
|
||
|
|
if args.streaming:
|
||
|
|
train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)
|
||
|
|
|
||
|
|
valid_data = load_dataset(
|
||
|
|
args.dataset_name,
|
||
|
|
split=args.eval_split,
|
||
|
|
)
|
||
|
|
return train_data, valid_data
|
||
|
|
|
||
|
|
|
||
|
|
def create_model(args):
|
||
|
|
print("Loading the model")
|
||
|
|
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 = None
|
||
|
|
|
||
|
|
# n_gpus = torch.cuda.device_count()
|
||
|
|
# max_memory = "32000MB"
|
||
|
|
# max_memory = {i: max_memory for i in range(n_gpus)}
|
||
|
|
if "t5" in args.model_name:
|
||
|
|
model_cls = AutoModelForSeq2SeqLM
|
||
|
|
else:
|
||
|
|
model_cls = AutoModelForCausalLM
|
||
|
|
|
||
|
|
model = model_cls.from_pretrained(
|
||
|
|
args.model_name,
|
||
|
|
quantization_config=quantization_config,
|
||
|
|
device_map=device_map,
|
||
|
|
trust_remote_code=args.trust_remote_code,
|
||
|
|
torch_dtype=torch_dtype,
|
||
|
|
# max_memory=max_memory,
|
||
|
|
token=True,
|
||
|
|
)
|
||
|
|
model.config.torch_dtype = torch_dtype
|
||
|
|
model.config.use_cache = False
|
||
|
|
|
||
|
|
print("Loading dataset")
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model_name if args.tokenizer_name is None else args.tokenizer_name)
|
||
|
|
if getattr(tokenizer, "pad_token", None) is None:
|
||
|
|
tokenizer.pad_token = tokenizer.eos_token
|
||
|
|
|
||
|
|
return model, tokenizer
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
parser = HfArgumentParser(ScriptArguments)
|
||
|
|
args = parser.parse_args_into_dataclasses()[0]
|
||
|
|
|
||
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
||
|
|
|
||
|
|
model, tokenizer = create_model(args)
|
||
|
|
|
||
|
|
train_dataset, eval_dataset = create_datasets(args)
|
||
|
|
|
||
|
|
if args.train_completions:
|
||
|
|
data_collator = DataCollatorForCompletionOnlyLM(tokenizer=tokenizer, response_template="TL;DR:")
|
||
|
|
else:
|
||
|
|
data_collator = None
|
||
|
|
|
||
|
|
training_args = TrainingArguments(
|
||
|
|
output_dir=args.output_dir,
|
||
|
|
per_device_train_batch_size=args.per_device_train_batch_size,
|
||
|
|
per_device_eval_batch_size=args.per_device_eval_batch_size,
|
||
|
|
dataloader_drop_last=True,
|
||
|
|
evaluation_strategy="steps",
|
||
|
|
max_steps=args.max_steps,
|
||
|
|
num_train_epochs=args.num_train_epochs,
|
||
|
|
eval_steps=args.eval_steps,
|
||
|
|
save_steps=args.save_steps,
|
||
|
|
save_strategy=args.save_strategy,
|
||
|
|
logging_steps=args.logging_steps,
|
||
|
|
learning_rate=args.learning_rate,
|
||
|
|
lr_scheduler_type=args.lr_scheduler_type,
|
||
|
|
warmup_steps=args.num_warmup_steps,
|
||
|
|
gradient_accumulation_steps=args.gradient_accumulation_steps,
|
||
|
|
gradient_checkpointing=args.gradient_checkpointing,
|
||
|
|
bf16=args.bf16,
|
||
|
|
fp16=args.fp16,
|
||
|
|
weight_decay=args.weight_decay,
|
||
|
|
report_to=args.log_with,
|
||
|
|
optim=args.optimizer_type,
|
||
|
|
remove_unused_columns=False,
|
||
|
|
disable_tqdm=False,
|
||
|
|
seed=args.seed,
|
||
|
|
# find_unused_params is necessary for grad checkpointing
|
||
|
|
ddp_find_unused_parameters=(args.gradient_checkpointing),
|
||
|
|
)
|
||
|
|
|
||
|
|
if args.use_peft:
|
||
|
|
peft_config = LoraConfig(
|
||
|
|
r=args.lora_r,
|
||
|
|
lora_alpha=args.lora_alpha,
|
||
|
|
lora_dropout=args.lora_dropout,
|
||
|
|
target_modules="all-linear",
|
||
|
|
bias="none",
|
||
|
|
task_type="CAUSAL_LM",
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
peft_config = None
|
||
|
|
|
||
|
|
chars_per_token = chars_token_ratio(train_dataset, tokenizer)
|
||
|
|
print(f"The character to token ratio of the train dataset is: {chars_per_token:.2f}")
|
||
|
|
|
||
|
|
print("Starting main loop")
|
||
|
|
trainer = SFTTrainer(
|
||
|
|
model=model,
|
||
|
|
tokenizer=tokenizer,
|
||
|
|
args=training_args,
|
||
|
|
train_dataset=train_dataset,
|
||
|
|
eval_dataset=eval_dataset,
|
||
|
|
peft_config=peft_config,
|
||
|
|
max_seq_length=args.seq_length,
|
||
|
|
formatting_func=prepare_sample_text,
|
||
|
|
packing=args.packing,
|
||
|
|
chars_per_token=chars_per_token,
|
||
|
|
data_collator=data_collator,
|
||
|
|
)
|
||
|
|
|
||
|
|
if args.use_peft:
|
||
|
|
trainer.model.print_trainable_parameters()
|
||
|
|
|
||
|
|
if not args.just_eval:
|
||
|
|
if args.resume_from_checkpoint is not None:
|
||
|
|
last_checkpoint = args.resume_from_checkpoint
|
||
|
|
else:
|
||
|
|
# when job is interrupted and restarted
|
||
|
|
last_checkpoint = get_last_checkpoint(args.output_dir)
|
||
|
|
|
||
|
|
print("Training...")
|
||
|
|
trainer.train(resume_from_checkpoint=last_checkpoint)
|
||
|
|
|
||
|
|
trainer.evaluate()
|
||
|
|
|
||
|
|
print("Saving last checkpoint of the model")
|
||
|
|
output_dir = os.path.join(args.output_dir, "final_model")
|
||
|
|
trainer.save_model(output_dir)
|
||
|
|
|
||
|
|
if args.use_peft:
|
||
|
|
output_dir = os.path.join(args.output_dir, "final_adapter_checkpoint")
|
||
|
|
trainer.model.save_pretrained(output_dir)
|
||
|
|
|
||
|
|
# Free memory for merging weights
|
||
|
|
del model
|
||
|
|
torch.cuda.empty_cache()
|
||
|
|
|
||
|
|
if "t5" in args.model_name:
|
||
|
|
model_cls = AutoPeftModelForSeq2SeqLM
|
||
|
|
else:
|
||
|
|
model_cls = AutoPeftModelForCausalLM
|
||
|
|
|
||
|
|
# model = model_cls.from_pretrained(
|
||
|
|
# output_dir, device_map="auto", torch_dtype=trainer.model.config.torch_dtype
|
||
|
|
# )
|
||
|
|
model = trainer.model.merge_and_unload()
|
||
|
|
|
||
|
|
output_merged_dir = os.path.join(args.output_dir, "final_merged_checkpoint")
|
||
|
|
model.save_pretrained(output_merged_dir, safe_serialization=True)
|
||
|
|
|
||
|
|
if args.output_model_name is not None:
|
||
|
|
model.push_to_hub(args.output_model_name)
|
||
|
|
|
||
|
|
else:
|
||
|
|
results = trainer.evaluate()
|
||
|
|
print(results)
|