finetune.py

import argparse
import functools
import os
import platform

import torch
from peft import LoraConfig, get_peft_model, AdaLoraConfig, PeftModel, prepare_model_for_kbit_training
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, WhisperForConditionalGeneration, WhisperProcessor

from utils.callback import SavePeftModelCallback
from utils.data_utils import DataCollatorSpeechSeq2SeqWithPadding
from utils.model_utils import load_from_checkpoint
from utils.reader import CustomDataset
from utils.utils import print_arguments, make_inputs_require_grad, add_arguments

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
add_arg("train_data",    type=str, default="dataset/train.json",       help="")
add_arg("test_data",     type=str, default="dataset/test.json",        help="")
add_arg("base_model",    type=str, default="openai/whisper-tiny",      help="Whisper")
add_arg("output_dir",    type=str, default="output/",                  help="")
add_arg("warmup_steps",  type=int, default=50,      help="")
add_arg("logging_steps", type=int, default=100,     help="")
add_arg("eval_steps",    type=int, default=1000,    help="")
add_arg("save_steps",    type=int, default=1000,    help="")
add_arg("num_workers",   type=int, default=8,       help="")
add_arg("learning_rate", type=float, default=1e-3,  help="")
add_arg("min_audio_len", type=float, default=0.5,   help="")
add_arg("max_audio_len", type=float, default=30,    help="")
add_arg("use_adalora",   type=bool,  default=True,  help="AdaLora/Lora")
add_arg("fp16",          type=bool,  default=True,  help="fp16")
add_arg("use_8bit",      type=bool,  default=False, help="8 bit")
add_arg("timestamps",    type=bool,  default=False, help="")
add_arg("local_files_only", type=bool, default=False, help="")
add_arg("num_train_epochs", type=int, default=3,      help="")
add_arg("language",      type=str, default="bn", help="")
add_arg("task",     type=str, default="transcribe", choices=['transcribe', 'translate'], help="模型的任务")
add_arg("augment_config_path",         type=str, default=None, help="")
add_arg("resume_from_checkpoint",      type=str, default=None, help="")
add_arg("per_device_train_batch_size", type=int, default=8,    help="batch size")
add_arg("per_device_eval_batch_size",  type=int, default=8,    help="batch size")
add_arg("gradient_accumulation_steps", type=int, default=1,    help="")

args = parser.parse_args()
print_arguments(args)


# Whisper tokenizer
processor = WhisperProcessor.from_pretrained(args.base_model,
                                             language=args.language,
                                             task=args.task,
                                             no_timestamps=not args.timestamps,
                                             local_files_only=args.local_files_only)

# 
train_dataset = CustomDataset(data_list_path=args.train_data,
                              processor=processor,
                              language=args.language,
                              timestamps=args.timestamps,
                              min_duration=args.min_audio_len,
                              max_duration=args.max_audio_len,
                              augment_config_path=args.augment_config_path)
test_dataset = CustomDataset(data_list_path=args.test_data,
                             processor=processor,
                             language=args.language,
                             timestamps=args.timestamps,
                             min_duration=args.min_audio_len,
                             max_duration=args.max_audio_len)
print(f"len train - {len(train_dataset)} test len - {len(test_dataset)}")

# padding
data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)

# Whisper
device_map = "auto"
world_size = int(os.environ.get("WORLD_SIZE", 1))
ddp = world_size != 1
if ddp:
    device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}

# 
model = WhisperForConditionalGeneration.from_pretrained(args.base_model,
                                                        load_in_8bit=args.use_8bit,
                                                        device_map=device_map,
                                                        local_files_only=args.local_files_only)
model.config.forced_decoder_ids = None
model.config.suppress_tokens = []
# 
model = prepare_model_for_kbit_training(model)
# forward，req grad
model.model.encoder.conv1.register_forward_hook(make_inputs_require_grad)

print('加载LoRA模块...')
if args.resume_from_checkpoint:
    # 
    print("Loading adapters from checkpoint.")
    model = PeftModel.from_pretrained(model, args.resume_from_checkpoint, is_trainable=True)
else:
    print(f'adding LoRA modules...')
    target_modules = ["k_proj", "q_proj", "v_proj", "out_proj", "fc1", "fc2"]
    print(target_modules)
    if args.use_adalora:
        config = AdaLoraConfig(init_r=12, target_r=4, beta1=0.85, beta2=0.85, tinit=200, tfinal=1000, deltaT=10,
                               lora_alpha=32, lora_dropout=0.1, orth_reg_weight=0.5, target_modules=target_modules)
    else:
        config = LoraConfig(r=32, lora_alpha=64, target_modules=target_modules, lora_dropout=0.05, bias="none")
    model = get_peft_model(model, config)

output_dir = os.path.join(args.output_dir, os.path.basename(args.base_model))
# 
training_args = \
    Seq2SeqTrainingArguments(output_dir=output_dir,  # Directory to save checkpoints
                             per_device_train_batch_size=args.per_device_train_batch_size,  # Training batch_size size
                             per_device_eval_batch_size=args.per_device_eval_batch_size,  # Eval batch_size 
                             gradient_accumulation_steps=args.gradient_accumulation_steps,  # Cumulative steps of training gradient
                             learning_rate=args.learning_rate,  # learning rate size
                             warmup_steps=args.warmup_steps,  # Warm-up steps
                             num_train_epochs=args.num_train_epochs,  # epochs
                             save_strategy="steps",  # 
                             evaluation_strategy="steps",  # 
                             load_best_model_at_end=True,  # 
                             fp16=args.fp16,  # 
                             report_to=["tensorboard"],  # tensorboard
                             save_steps=args.save_steps,  # 
                             eval_steps=args.eval_steps,  # 
                             save_total_limit=5,  # 
                             optim='adamw_torch',  # 
                             ddp_find_unused_parameters=False if ddp else None,  # 
                             dataloader_num_workers=args.num_workers,  # 
                             logging_steps=args.logging_steps,  # 
                             remove_unused_columns=False,  # 
                             label_names=["labels"])  # 

if training_args.local_rank == 0 or training_args.local_rank == -1:
    print('=' * 90)
    model.print_trainable_parameters()
    print('=' * 90)

# Pytorch2.0
if torch.__version__ >= "2" and platform.system().lower() == 'windows':
    model = torch.compile(model)

# 
trainer = Seq2SeqTrainer(args=training_args,
                         model=model,
                         train_dataset=train_dataset,
                         eval_dataset=test_dataset,
                         data_collator=data_collator,
                         tokenizer=processor.feature_extractor,
                         callbacks=[SavePeftModelCallback])
model.config.use_cache = False
trainer._load_from_checkpoint = load_from_checkpoint

# 
trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)

# 
trainer.save_state()
if training_args.local_rank == 0 or training_args.local_rank == -1:
    model.save_pretrained(os.path.join(output_dir, "checkpoint-final"))