model/trainer.py

import signal
import sys
import os
import time
from typing import Union
import platform 

from psutil import virtual_memory, cpu_count
import numpy as np
from torch.utils.data import DataLoader
import torch 
from rich.progress import Progress, TextColumn, BarColumn, TimeElapsedColumn, TimeRemainingColumn
from transformers import PreTrainedTokenizerFast
from torch_optimizer import Adafactor

# import accelerate
from accelerate import Accelerator
from accelerate.utils import set_seed

# import 自定义类和函数
from model.chat_model import TextToTextModel
from utils.logger import Logger
from model.dataset import MyDataset
from config import TrainConfig, T5ModelConfig
from utils.functions import (
    get_bleu4_score, 
    save_model_config, 
    get_free_space_of_disk, 
    my_average,
    get_path_of_suffix_files,
    get_T5_config,
)

class ChatTrainer:
    def __init__(self, train_config: TrainConfig, model_config: T5ModelConfig, ) -> None:
        
        self.train_config = train_config
        self.model_config = model_config

        # file_name=None会自动生成以当前日期命名的log文件名
        self.logger = Logger('chat_trainer', std_out=True, save2file=True, file_name=None)

        self.model = None
        self.accelerator = None

        signal.signal(signal.SIGINT, self.process_exit_handler)

        self.is_win_platform = True if platform.system().lower() == 'windows' else False

        torch.manual_seed(train_config.seed)
        torch.cuda.manual_seed_all(train_config.seed)
    
    def process_exit_handler(self, signal_received, frame) -> None:
        '''
        进程退出时的操作，保存模型
        '''
        if self.accelerator and self.model:
            ask = "you are pressed `ctrl+c`,  do you want to save checkpoint? Yes (y) or No (n)"
            self.accelerator.print(ask)
            ins = input()
            
            if ins.lower() in ('yes', 'y'):

                suffix =  'exit_save_{}'.format(str(time.strftime('%Y%m%d%H%M%S', time.localtime())))

                self.accelerator.wait_for_everyone()
                self.accelerator.save_state(output_dir=self.train_config.train_state_dir)

                self.accelerator.print('model ckeck point has been saved in {}'.format(self.train_config.train_state_dir))
        
            sys.exit(0)
        else:
            print('process not in trainingg, exit.')
            sys.exit(0)

    def save_model(self, suffix: Union[str, int]) -> None:
        '''保存模型到文件
        注意：save_model不能放到is_main_process里面
        e.g:
        >>> self.save_model(epoch) # 在这里使用
        >>> if accelerator.is_main_process:
        >>>     do_somthing()
        '''
        if self.model and self.accelerator:

            # 先wait_for_everyone，再保存
            self.accelerator.wait_for_everyone()

            if self.accelerator.is_main_process:
                unwrap_model = self.accelerator.unwrap_model(self.model)
                model_dict =  self.accelerator.get_state_dict(unwrap_model)
                torch.save(model_dict, self.train_config.model_file.format(suffix))

    
    def delete_early_checkpoint(self, epoch: int, keep_latest_n: int=3,) -> None:
        '''
        删除最早的模型，最保留最近keep_latest_n个模型文件
        '''
        model_save_path = self.train_config.model_file
        model_save_path = model_save_path.replace('\\', '/')    # 针对win的路径，将\替换为/
        model_save_path = '/'.join(model_save_path.split('/')[0: -1])   # 删除末尾文件名后缀
        
        model_files = get_path_of_suffix_files(model_save_path, suffix='.bin', with_create_time=True)
        
        # 进程异常退出保存模型文件不在删除范围
        train_save_model_fils = []
        for item in model_files:
            if 'exit_save' not in item[0]:

                # 大于当前epoch的文件不不删除
                f_epoch = int(item[0].split('.')[-2])
                if epoch >= f_epoch:
                    print(epoch, f_epoch, item)
                    train_save_model_fils.append(item)

        train_save_model_fils.sort(key=lambda x: x[1])  # 按照时间从小到大排序

        if len(train_save_model_fils) <= keep_latest_n:
            return
        
        to_delete_files = train_save_model_fils[0: -keep_latest_n]
        for item in to_delete_files:
            os.remove(item[0])

    
    def train(self, is_keep_training: bool=False, is_finetune: bool=False) -> None:
        '''
        is_keep_training: 是否从断点处加载状态继续训练
        is_finetune: 是否微调，微调的话可能需要冻结部分参数
        '''
        log = self.logger
        train_config = self.train_config
        save_steps = self.train_config.save_steps
        logging_steps = self.train_config.logging_steps

        # 梯度累计的步数
        accumulation_steps = train_config.gradient_accumulation_steps

        set_seed(train_config.seed)

        accelerator = Accelerator(
            mixed_precision=train_config.mixed_precision,       # 混合精度
            gradient_accumulation_steps=accumulation_steps,     # 梯度累积
            project_dir=train_config.train_state_dir,
        )

        # 根据剩余内存大小决定是否完全加载数据集到内存中
        unuse_mem = virtual_memory().available / (1024 ** 3)  # 单位：GB
        unuse_disk = get_free_space_of_disk('./')

        # 剩余内存≥48GB将把数据集留在内存中,因为2个显卡+全全部装载900多万的训练数据到内存需要大概43GB的CPU内存
        # 如果不放在内存中，将会使用迭代器生成数据，CPU 内存小于16GB也可以运行，但是不支持顺序打乱。
        # 多GPU keep_in_memory必须=True，否则无法进行分布式训练
        keep_in_memory = True if unuse_mem >= 48.0 or torch.cuda.device_count() >= 2 else False

        if accelerator.is_main_process:
            log.info('cpu memory available: {:.2f} GB, disk space available: {:.2f} GB, keep dataset in memory: {}.'\
                    .format(unuse_mem, unuse_disk, keep_in_memory), save_to_file=True)
            log.info('operation: {}, keep training: {}, loading datasets ...'.format('finetune' if is_finetune else 'train', is_keep_training))

        # args for dataloader
        num_workers = 0
        # if not self.is_win_platform:
        #     cpu_cnt = cpu_count(logical=False)
        #     gpu_cnt = torch.cuda.device_count()
        #     if cpu_cnt >= 8 * gpu_cnt:
        #         # num_workers = 4 x number of available GPUs
        #         num_workers = int(4 * gpu_cnt)
        #     else:
        #         num_workers = int(cpu_cnt // 2)

        train_dataset = MyDataset(
            parquet_file=train_config.train_file,
            tokenizer_dir=train_config.tokenizer_dir,
            keep_in_memory=keep_in_memory,
            max_seq_len=train_config.max_seq_len,
        )
        valid_dataset = MyDataset(
            parquet_file=train_config.validation_file,
            tokenizer_dir=train_config.tokenizer_dir,
            keep_in_memory=keep_in_memory,
            max_seq_len=train_config.max_seq_len,
        )

        batch_size = train_config.batch_size_per_gpu

        train_dataloader = DataLoader(
            train_dataset, 
            batch_size=batch_size,  
            shuffle=True,
            collate_fn=train_dataset.collate_fn,
            pin_memory=False,
            num_workers=num_workers,    #设置>1会导致cpu内存缓慢增涨，最后OOM，后面再研究为什么，num_workers=4，一个epoch只减少30分钟
        )
        valid_dataloader = DataLoader(
            valid_dataset, 
            batch_size=batch_size, 
            shuffle=False,
            collate_fn=valid_dataset.collate_fn,
            pin_memory=False,
            num_workers=num_workers,
        )

        device = accelerator.device
        log.info('using device: {} '.format(str(device)), save_to_file=True)
        

        # T5: All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        tokenizer = train_dataset.tokenizer
        decoder_start_token_id = tokenizer.pad_token_id

        # for t5, set decoder_start_token_id = pad_token_id
        t5_config = get_T5_config(T5ModelConfig(), vocab_size=len(tokenizer), decoder_start_token_id=decoder_start_token_id, eos_token_id=tokenizer.eos_token_id)

        model = TextToTextModel(t5_config)

        # 微调加载的模型并冻结embedding和encoder
        if is_finetune:
            model.load_state_dict(torch.load(train_config.finetune_from_ckp_file))
            # print(model)
            
            layers_to_freeze = [model.shared, model.encoder]

            for layer in layers_to_freeze:
                 for param in layer.parameters():
                    param.requires_grad = False

        # 保存模型配置，方便修改配置后恢复
        save_model_config(t5_config.to_diff_dict(), train_config.model_config_file)
        
        # T5训练，论文推荐使用Adafactor
        optimizer = Adafactor(params=model.parameters(), lr=train_config.learn_rate)

        
        # 获取当前机器有多少个GPU，默认全部使用
        num_gpus_used = accelerator.state.num_processes

        # 单机多卡，每个step总共的batch_size = batch_size_per_gpu * num_gpus_used
        # total_batch_size 初始化为batch_size_per_gpu真的只有CPU的情况
        total_batch_size = train_config.batch_size_per_gpu
        if num_gpus_used >= 1:
            total_batch_size = num_gpus_used * train_config.batch_size_per_gpu

        steps_per_epoch = int(np.ceil(len(train_dataset) // total_batch_size))
        eval_steps = int(np.ceil(len(valid_dataset) // total_batch_size))

        if accelerator.is_main_process:
            log.info('train dataset size: {}, steps per epoch:{}; validation dataset size: {}, steps per validation: {}; datalodater num_workers: {}.'\
                    .format(len(train_dataset), steps_per_epoch, len(valid_dataset), eval_steps, num_workers), save_to_file=True)

        
        lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
                optimizer=optimizer, 
                max_lr=train_config.div_factor * train_config.learn_rate, 
                epochs=train_config.epochs, 
                steps_per_epoch=int(np.ceil( len(train_dataset) / (batch_size * accumulation_steps) )),  # 梯度累积相当于增大了batch_size
                div_factor=train_config.div_factor,
                cycle_momentum=False,
            )
        
        model, optimizer, lr_scheduler, train_dataloader, valid_dataloader = accelerator.prepare(
                model, 
                optimizer,
                lr_scheduler, 
                train_dataloader, 
                valid_dataloader,
            )
        
        if is_keep_training:
            accelerator.load_state(input_dir=train_config.train_state_dir)
            accelerator.register_for_checkpointing(lr_scheduler)
        
        self.model = model
        self.accelerator = accelerator
        
        best_bleu4 = 0.0
        best_epoch = 0
        epoch_loss_list = []

        # 添加进度条，只在主进程更新
        if accelerator.is_main_process:
            progress = Progress(TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                TimeRemainingColumn(),
                TimeElapsedColumn(),
                TextColumn("[bold blue]{task.fields[show_info]}"),
                refresh_per_second=1,  # 每1秒钟更新一次，不要频繁更新
                )
            
            epoch_progress = progress.add_task(description='epoch: ', show_info='', total=train_config.epochs)
            steps_progress = progress.add_task(description='steps: ', show_info='', \
                                                total=np.ceil(steps_per_epoch / logging_steps))
            eval_progress = progress.add_task(description='evaluate: ', show_info='', total=eval_steps, visible=False)

            self.progress = progress
            self.eval_progress = eval_progress

            progress.start()

        # end if

        for epoch in range(train_config.epochs):
            
            if accelerator.is_main_process:
                epoch_show_txt = 'epoch: {}/{}, avg_loss: {:.6f}, best_epoch: {}, best_bleu: {}'.format(
                    epoch, train_config.epochs, my_average(epoch_loss_list), best_epoch, best_bleu4
                )
                progress.update(epoch_progress, show_info=epoch_show_txt)
                progress.reset(steps_progress)

            epoch_loss_list = []
            model.train()

            # torch.cuda.empty_cache()

            for step, batch_data in enumerate(train_dataloader):

                input_ids, input_mask = batch_data['input_ids'], batch_data['input_mask']
                target_ids = batch_data['target_ids']
                # for t5 model, all labels set to `-100` are ignored (masked)
                target_ids[target_ids == decoder_start_token_id] = -100

                outputs = model(
                    input_ids=input_ids,
                    attention_mask=input_mask,
                    labels=target_ids,
                )

                loss = outputs.loss.mean() / accumulation_steps

                # attention here! loss.backward()
                accelerator.backward(loss) 

                # 梯度累计
                if (step + 1) % accumulation_steps == 0:
                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
                
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
                
                # 每隔save_steps步保存一次模型
                if (step + 1) % save_steps == 0 or step == steps_per_epoch:
                    self.save_model('epoch_{}_latest'.format(epoch))
                    accelerator.save_state(output_dir=train_config.train_state_dir)
                
                # ==================================以下记录loss到日志============================================
                # 每n步更新一次，避免频繁的cpu-gpu数据复制
                # 参考：https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#avoid-unnecessary-cpu-gpu-synchronization
                
                if step % logging_steps == 0 or step == steps_per_epoch:
                    
                    loss_cpu = loss.detach().item() * accumulation_steps
                    epoch_loss_list.append(loss_cpu)
                    
                    info_txt = 'training loss: epoch:{}, step:{}, loss:{}, device:{}'.\
                        format(epoch, step, loss_cpu, str(accelerator.device))
                    
                    log.info(info_txt, std_out=False, save_to_file=True) # 保存 loss 到文件

                    # 更新进度条
                    if accelerator.is_main_process:
                        step_show_txt = 'step: {}/{}, loss: {:.6f}'.format(step, steps_per_epoch, loss_cpu)
                        progress.advance(steps_progress, advance=1)
                        progress.update(steps_progress, show_info=step_show_txt)

                # ==================================以上记录loss到日志============================================
                
                # if step >= 20:break
            
            #  end for batch setps

            model.eval()         
            
            cur_bleu4_score = self.evaluate(
                model=model,
                tokenizer=tokenizer,
                valid_dataloader=valid_dataloader,
                accelerator=accelerator,
                eval_steps=eval_steps,
                )

            # save model
            if cur_bleu4_score >= best_bleu4:

                best_bleu4 = cur_bleu4_score
                best_epoch = epoch
                # 最多保存最近keep_latest_n_ckp个模型文件
                # self.delete_early_checkpoint(epoch=epoch, keep_latest_n=train_config.keep_latest_n_ckp)
                self.save_model('best')
                accelerator.save_state(output_dir=train_config.train_state_dir)

            # 每个epoch打印一下日志
            if accelerator.is_main_process:

                progress.advance(epoch_progress, advance=1)
                info_txt = 'epoch log: epoch:{}, avg_loss:{}, cur_bleu4:{}, best_bleu4:{}, best_epoch:{}'.\
                            format(epoch, my_average(epoch_loss_list), cur_bleu4_score, best_bleu4, best_epoch)
                # log.info(info_txt, std_out=True, save_to_file=True)
                self.print_and_log(info_txt, accelerator)


    def evaluate(self, 
                model: TextToTextModel, 
                tokenizer: PreTrainedTokenizerFast,
                valid_dataloader: DataLoader, 
                accelerator: Accelerator,
                eval_steps: int,
            ) -> float:
        
        '''
        评估，返回平均的bleu分数
        '''
        max_seq_len = self.train_config.max_seq_len
        batch_decode = tokenizer.batch_decode
        bleu4_scores = []

        if accelerator.is_main_process:
            self.progress.reset(self.eval_progress)
            self.progress.update(self.eval_progress, visible=True)

        with torch.no_grad():
            for step, batch_data in enumerate(valid_dataloader):
                
                if accelerator.is_main_process:
                    self.progress.advance(self.eval_progress, advance=1)
                    self.progress.update(self.eval_progress, show_info='step: {}/{}'.format(step, eval_steps))

                input_ids, input_mask = batch_data['input_ids'], batch_data['input_mask']
                target_ids = batch_data['target_ids']

                outputs = accelerator.unwrap_model(model).my_generate(
                    input_ids=input_ids,
                    attention_mask=input_mask,
                    max_seq_len=max_seq_len,
                )

                # gather data from multi-gpus (used when in ddp mode)
                outputs = accelerator.gather_for_metrics(outputs).detach().cpu().numpy()
                target_ids = accelerator.gather_for_metrics(target_ids).detach().cpu().numpy()
        
                outputs = batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False)
                target_ids = batch_decode(target_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)

                # print(outputs, target_ids)

                bleu4_scores = [get_bleu4_score(reference=target_ids[i], outputs=outputs[i]) for i in range(len(target_ids))]
                bleu4_scores.extend(bleu4_scores)

                # if step >= 5: break
        
        avg_bleu4_score = my_average(bleu4_scores)
        if accelerator.is_main_process:
            self.progress.update(self.eval_progress, show_info='bleu4 score: {}'.format(avg_bleu4_score))
            self.progress.update(self.eval_progress, visible=False)

        return avg_bleu4_score

    def test(self, best_epoch: int=0) -> None:
        '''
        '''
        import os 

        train_config = self.train_config
        log = self.logger

        # args for dataloader
        num_workers = 0 if self.is_win_platform else 4

        test_dataset = MyDataset(
            parquet_file=train_config.train_file,
            tokenizer_dir=train_config.tokenizer_dir,
            keep_in_memory=False if self.is_win_platform else True,
            max_seq_len=train_config.max_seq_len,
        )
        
        test_dataloader = DataLoader(
            test_dataset, 
            batch_size=train_config.batch_size_per_gpu,
            shuffle=False,
            collate_fn=test_dataset.collate_fn,
            pin_memory=False,
            num_workers=num_workers,
        )

        log.info('test dataset size: {}.'.format(len(test_dataset)), save_to_file=True)

        set_seed(train_config.seed)
        accelerator = Accelerator(mixed_precision=train_config.mixed_precision)
        device = accelerator.device
        log.info('using device: {} '.format(str(device)), save_to_file=True)

         # 获取当前运行使用了多少个GPU
        num_gpus_used = accelerator.state.num_processes

        # 单机多卡，每个step总共的batch_size = batch_size_per_gpu * num_gpus_used
        # total_batch_size 初始化为batch_size_per_gpu真的只有CPU的情况
        total_batch_size = train_config.batch_size_per_gpu
        if num_gpus_used >= 1:
            total_batch_size = num_gpus_used * train_config.batch_size_per_gpu

        # T5: All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        tokenizer = test_dataset.tokenizer

        model_file = train_config.model_file.format(best_epoch)
        if os.path.isdir(model_file):
            # 传入文件夹则 from_pretrained
            model = TextToTextModel.from_pretrained(model_file)
        else:
            # load_state_dict
            t5_config = get_T5_config(T5ModelConfig(), vocab_size=len(tokenizer), decoder_start_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)
            model = TextToTextModel(t5_config)
            model.load_state_dict(torch.load(model_file, map_location='cpu')) # set cpu for no exception
       
        model, test_dataloader = accelerator.prepare(
                model, 
                test_dataloader,
            )
        
        steps = int(np.ceil(len(test_dataset) // total_batch_size))

        bleu4 = 0.0
        bleu4_scores = []
        batch_decode = tokenizer.batch_decode
        max_seq_len = self.train_config.max_seq_len
        model.eval()

        if accelerator.is_main_process:
            progress = Progress(TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                TimeRemainingColumn(),
                TimeElapsedColumn(),
                TextColumn("[bold blue]{task.fields[show_info]}"),
                refresh_per_second=1.0,
                )
                
            steps_progress = progress.add_task(description='steps: ', show_info='', total=steps)
            progress.start()
            
        with torch.no_grad():
            for step, batch_data in enumerate(test_dataloader):

                if accelerator.is_main_process:
                    progress.advance(steps_progress, advance=1)
                    progress.update(steps_progress, show_info='step: {}/{}'.format(step, steps))

                input_ids, input_mask = batch_data['input_ids'], batch_data['input_mask']
                target_ids = batch_data['target_ids']

                # s = time.time()
                outputs = accelerator.unwrap_model(model).my_generate(
                    input_ids=input_ids,
                    attention_mask=input_mask,
                    max_seq_len=max_seq_len,
                )
                # accelerator.print('generate used: {}'.format(time.time() - s))

                # gather data from multi-gpus (used when in ddp mode)
                outputs = accelerator.gather_for_metrics(outputs).cpu().numpy()
                target_ids = accelerator.gather_for_metrics(target_ids).cpu().numpy()
                
                outputs = batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False)
                target_ids = batch_decode(target_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)

                # print('outputs: {}'.format(outputs[0:5]))
                # print('target_ids: {}'.format(target_ids[0:5]))
                # print()


                bleu4_scores = [get_bleu4_score(reference=target_ids[i], outputs=outputs[i]) for i in range(len(target_ids))]
                bleu4_scores.extend(bleu4_scores)

                # if step >= 10: break
        
        avg_bleu4_score = my_average(bleu4_scores)
        if accelerator.is_main_process:
            progress.update(steps_progress, show_info='bleu4 score: {}'.format(avg_bleu4_score))

        info_txt = 'test_dataset_size: {}, avg_bleu4_score:{}.'.format(len(test_dataset), avg_bleu4_score)
        log.info(info_txt, save_to_file=True)

        return avg_bleu4_score

    
    def print_and_log(self, info: str, accelerator: Accelerator=None) -> None:
        '''
        使用accelerator.print, 否则多进程打印会异常
        '''
        if not accelerator:
            print(info)
        else:
            accelerator.print(info)
        self.logger.info(info, std_out=False, save_to_file=True)

if __name__ == '__main__':
    
    # trainer = ChatTrainer()
    train_config = TrainConfig()
    model_config = T5ModelConfig()

    chat_trainer = ChatTrainer(train_config=train_config, model_config=model_config)

    chat_trainer.train()
    # chat_trainer.test(best_epoch=0)