src/utils/callbacks.py

import os
import time
import typing
# from deepspeed.accelerator import get_accelerator

# NOTE: 最新版开始迁移到 integration_utils
try:
    from transformers.integrations import TrainerCallback
except ImportError:
    from transformers.integrations.integration_utils import TrainerCallback

# USE_FLASH_ATTN, USE_XFORMERS_ATTN = False, False
# if os.getenv('FLASH_ATTN', 'false').lower() == 'true':
#     USE_FLASH_ATTN = True
#     from mllm.utils.llama_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn, restore_llama_attn
# if os.getenv("XFORMERS_ATTN", 'false').lower() == 'true':
#     USE_XFORMERS_ATTN = True
#     from mllm.utils.llama_xformers_monkey_patch import replace_llama_attn_with_xformers_attn, restore_llama_attn

class ModeltimeCallback(TrainerCallback):
    def __init__(self):
        self.model_time = 0.
        self.data_time = 0.
        self._start = 0.
        self._end = 0.

    def on_train_begin(self, args, state, control, **kwargs):
        self._end = time.time()
    
    def on_step_begin(self, args, state, control, **kwargs):
        self._start = time.time()
        self.data_time += (self._start - self._end)

    def on_step_end(self, args, state, control, **kwargs):
        self._end = time.time()
        self.model_time += (self._end - self._start)

    def on_log(self, args, state, control, logs=None, **kwargs):
        if not state.is_world_process_zero:
            return
        
        data_time = self.data_time / state.logging_steps
        model_time = self.model_time / state.logging_steps
        self.data_time = 0.
        self.model_time = 0.
        
        info = f'\nSTEP: {state.global_step}, data_time: {data_time:.3f}, model_time: {model_time:.3f}'
        print(info)


class SacredCallback(TrainerCallback):
    def __init__(self, _run=None):
        self._run = _run
        self.model_time = 0.
        self.data_time = 0.
        self._start = 0.
        self._end = 0.
        self._zero_loss_cnt = 0
    
    def on_train_begin(self, args, state, control, **kwargs):
        if self._run:
            self._end = time.time()
    
    def on_step_begin(self, args, state, control, **kwargs):
        if self._run:
            self._start = time.time()
            self.data_time += (self._start - self._end)

    def on_step_end(self, args, state, control, **kwargs):
        if self._run:
            self._end = time.time()
            self.model_time += (self._end - self._start)
    
    def on_log(self, args, state, control, logs=None, **kwargs):
        if not state.is_world_process_zero:
            return
        
        if self._run is None:
            return 

        data_time = self.data_time / state.logging_steps
        model_time = self.model_time / state.logging_steps
        self.data_time = 0.
        self.model_time = 0.
        
        self._run.log_scalar("loss", logs.get("loss", 0.), state.global_step)
        self._run.log_scalar("learning_rate", logs.get("learning_rate", 0.), state.global_step)
        self._run.log_scalar("grad_norm", logs.get("grad_norm", 0.), state.global_step)
        self._run.log_scalar("epoch", state.epoch, state.global_step)
        self._run.log_scalar("data_time", data_time, state.global_step)
        self._run.log_scalar("model_time", model_time, state.global_step)
        
        if logs.get("loss", 0.) == 0.:
            self._zero_loss_cnt += 1
            if self._zero_loss_cnt > 1:
                raise RuntimeError("Loss is zero, something is wrong!")
        

class ModelEvalCallback(TrainerCallback):
    def __init__(self, _run=None, multitest=None, trainer=None, gen_kwargs=None):
        self._run = _run
        # datasets dict
        # key1: dataset_name
        # val1: {'dataset': dataset inst, 'compute_metric': metric inst}
        self.multitest = typing.cast(dict, multitest)
        self.trainer = trainer
        self.gen_kwargs = gen_kwargs
    
    def on_step_end(self, args, state, control, **kwargs):
        if args.eval_steps is None:
            eval_steps = args.save_steps
        elif isinstance(args.eval_steps, int) and args.eval_steps > 0:
            eval_steps = args.eval_steps
        else:
            return
        if state.global_step > 0 and state.global_step % eval_steps == 0:
            if not args.do_multi_predict:
                return
            
            # flash-attn currently not supports eval mode!
            # if USE_FLASH_ATTN or USE_XFORMERS_ATTN:
            #     restore_llama_attn()
            
            old_compute_metrics = self.trainer.compute_metrics
            
            for dataset_idx, (dataset_name, item) in enumerate(self.multitest.items()):
                print(f'processing multitest set {dataset_idx}/{len(self.multitest)}: {dataset_name}')
                _ds = item['dataset']
                _compute_metrics = item['compute_metric']
                _prefix = dataset_name
                
                self.trainer.compute_metrics = _compute_metrics
                # transformers.trainer_utils.PredictionOutput
                _pred_results = self.trainer.predict(_ds, metric_key_prefix=_prefix, **self.gen_kwargs)
                if state.is_world_process_zero:
                    self.trainer.log_metrics(_prefix, _pred_results.metrics)  # noqa
                    self.trainer.save_metrics(_prefix, _pred_results.metrics)  # noqa
                    self.trainer.save_prediction(_pred_results, file_key_prefix=_prefix)
                    
                    if self._run is not None:
                        keywords_to_remove = ['runtime', 'second']
                        for k, v in _pred_results.metrics.items():
                            # remove time releated metrics
                            if any(kw in k for kw in keywords_to_remove):
                                continue
                            self._run.log_scalar(f'{k}', v, state.global_step)
            
            self.trainer.compute_metrics = old_compute_metrics
            
            # if USE_FLASH_ATTN:
            #     replace_llama_attn_with_flash_attn()
            # if USE_XFORMERS_ATTN:
            #     replace_llama_attn_with_xformers_attn()

class DSEmptyCacheCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        empty_cache_steps = int(os.getenv("EMPTY_CACHE_STEP", '0').strip())
        can_flush = state.global_step > 0 and empty_cache_steps > 0 and state.global_step % empty_cache_steps == 0
        if can_flush:
            # print('Flush Cache here.')
            get_accelerator().empty_cache()
            

# usage: https://github.com/yqhu/profiler-workshop/blob/c8d4a7c30a61cc7b909d89f88f5fd36b70c55769/hf_training_trainer_prof.py#L49C6-L49C28
# additionally, with_modules can be set True, with_flops must be set False
# deps: pip install -U tensorboard-plugin-profilepip torch_tb_profiler 
class ProfCallback(TrainerCallback):
    def __init__(self, prof):
        self.prof = prof

    def on_step_end(self, args, state, control, **kwargs):
        self.prof.step()