| 69 | os.remove(os.path.join(args.output_dir, LOG_FILE_NAME)) |
| 70 | |
| 71 | def on_log( |
| 72 | self, |
| 73 | args: TrainingArguments, |
| 74 | state: TrainerState, |
| 75 | control: TrainerControl, |
| 76 | logs, |
| 77 | **kwargs |
| 78 | ): |
| 79 | if args.save_on_each_node: |
| 80 | if not state.is_local_process_zero: |
| 81 | return |
| 82 | else: |
| 83 | if not state.is_world_process_zero: |
| 84 | return |
| 85 | |
| 86 | self.elapsed_time += time.time() - self.last_time |
| 87 | self.last_time = time.time() |
| 88 | if 'num_input_tokens_seen' in logs: |
| 89 | logs['num_tokens'] = logs.pop('num_input_tokens_seen') |
| 90 | state.log_history[-1].pop('num_input_tokens_seen') |
| 91 | throughput = logs['num_tokens'] / args.world_size / self.elapsed_time |
| 92 | state.log_history[-1]['throughput'] = logs['throughput'] = throughput |
| 93 | state.stateful_callbacks["LogCallback"] = self.state() |
| 94 | |
| 95 | logs = dict( |
| 96 | current_steps=state.global_step, |
| 97 | total_steps=state.max_steps, |
| 98 | loss=state.log_history[-1].get("loss", None), |
| 99 | eval_loss=state.log_history[-1].get("eval_loss", None), |
| 100 | predict_loss=state.log_history[-1].get("predict_loss", None), |
| 101 | learning_rate=state.log_history[-1].get("learning_rate", None), |
| 102 | epoch=state.log_history[-1].get("epoch", None), |
| 103 | percentage=round(state.global_step / state.max_steps * 100, 2) if state.max_steps != 0 else 100, |
| 104 | ) |
| 105 | |
| 106 | os.makedirs(args.output_dir, exist_ok=True) |
| 107 | with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a", encoding="utf-8") as f: |
| 108 | f.write(json.dumps(logs) + "\n") |
| 109 | |
| 110 | def state(self) -> dict: |
| 111 | return { |