| 11 | |
| 12 | |
| 13 | class CustomLoggingCallback(TrainerCallback): |
| 14 | |
| 15 | def __init__(self, max_steps, batch_size, max_length, world_size, valid_dataset_abbr, logging_steps=10): |
| 16 | self.cur_step, self.max_steps = 0, max_steps |
| 17 | self.start_time, self.valid_total_time = datetime.now(), 0 |
| 18 | self.batch_token = batch_size * max_length |
| 19 | self.world_size = world_size |
| 20 | self.valid_dataset_abbr = valid_dataset_abbr |
| 21 | self.logging_steps = logging_steps |
| 22 | |
| 23 | self.rank = torch.distributed.get_rank() |
| 24 | |
| 25 | def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs=None, **kwargs): |
| 26 | now = datetime.now() |
| 27 | if logs is not None and 'grad_norm' in logs: |
| 28 | |
| 29 | if self.cur_step == 0: |
| 30 | self.start_time = now |
| 31 | total_time = now - self.start_time |
| 32 | log_dict = {'num_consume_token': 0, 'avg_tgs': 0} |
| 33 | if self.rank == 0 and self.cur_step % self.logging_steps == 0: |
| 34 | print(f"[{str(now)}] 0.00% {self.cur_step:{len(str(self.max_steps))}d} / {self.max_steps}", |
| 35 | f"[{str(total_time)} / {str(total_time)}] {log_dict}", end=', ') |
| 36 | |
| 37 | else: |
| 38 | total_time = now - self.start_time |
| 39 | train_total_time = total_time.total_seconds() - self.valid_total_time |
| 40 | num_consume_token = self.cur_step * self.batch_token |
| 41 | avg_tgs = num_consume_token / train_total_time / self.world_size |
| 42 | |
| 43 | cur_percent = self.cur_step / self.max_steps * 100 |
| 44 | final_time = total_time / self.cur_step * self.max_steps |
| 45 | log_dict = {'num_consume_token': num_consume_token, 'avg_tgs': avg_tgs} |
| 46 | if self.rank == 0 and self.cur_step % self.logging_steps == 0: |
| 47 | print(f"[{str(now)}] {cur_percent:6.2f}% {self.cur_step:{len(str(self.max_steps))}d} / {self.max_steps}", |
| 48 | f"[{str(total_time)} / {str(final_time)}] {log_dict}", end=', ') |
| 49 | |
| 50 | self.cur_step += 1 |
| 51 | else: |
| 52 | total_time = now - self.start_time |
| 53 | cur_percent = self.cur_step / self.max_steps * 100 |
| 54 | |
| 55 | if self.cur_step == 0: |
| 56 | self.start_time = now |
| 57 | total_time = now - self.start_time |
| 58 | if self.rank == 0 and self.cur_step % self.logging_steps == 0: |
| 59 | print(f"[{str(now)}] 0.00% {self.cur_step:{len(str(self.max_steps))}d} / {self.max_steps}", |
| 60 | f"[{str(total_time)} / {str(total_time)}] validation ", end='') |
| 61 | else: |
| 62 | final_time = total_time / self.cur_step * self.max_steps |
| 63 | if self.rank == 0: |
| 64 | print(f"[{str(now)}] {cur_percent:6.2f}% {self.cur_step:{len(str(self.max_steps))}d} / {self.max_steps}", |
| 65 | f"[{str(total_time)} / {str(final_time)}] validation ", end='') |
| 66 | |
| 67 | def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics=None, **kwargs): |
| 68 | if f'eval_{self.valid_dataset_abbr}_runtime' in metrics: |
| 69 | self.valid_total_time += metrics[f'eval_{self.valid_dataset_abbr}_runtime'] |
| 70 |
no outgoing calls
no test coverage detected