(self, round)
| 111 | |
| 112 | |
| 113 | def save_model(self, round): |
| 114 | if self.args.output_dir is not None: |
| 115 | print_rank_0('saving model to ' + self.args.output_dir + "/" + str(round) + '...', self.args.global_rank) |
| 116 | |
| 117 | if self.args.global_rank == 0: |
| 118 | save_hf_format(self.model, self.tokenizer, self.args, sub_folder=str(round)) |
| 119 | |
| 120 | if self.args.zero_stage == 3: |
| 121 | # For zero stage 3, each gpu only has a part of the model, so we need a special save function |
| 122 | save_zero_three_model(self.model, |
| 123 | self.args.global_rank, |
| 124 | self.args.output_dir, |
| 125 | zero_stage=self.args.zero_stage, |
| 126 | sub_folder=str(round)) |
| 127 | print_rank_0('Successfully saving model after round {}'.format(round), self.args.global_rank) |
no test coverage detected