File size: 7,105 Bytes
5e27996 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | 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()
|