(self, round)
| 13 | super().__init__(model, tokenizer, optimizer, train_task_list, eval_task_list, test_task_list, args) |
| 14 | |
| 15 | def save_model(self, round): |
| 16 | # # if self.args.output_dir is not None: |
| 17 | # # print_rank_0('saving the final model ...', self.args.global_rank) |
| 18 | # # |
| 19 | # # if self.args.global_rank == 0: |
| 20 | # # peft_model_id = os.path.join(self.args.output_dir, str(i_task)) |
| 21 | # # if not os.path.exists(peft_model_id): |
| 22 | # # os.makedirs(peft_model_id) |
| 23 | # # self.model.save_pretrained(peft_model_id) |
| 24 | # # self.tokenizer.save_pretrained(peft_model_id) |
| 25 | # # print_rank_0(f'Successfully saving the final model to {peft_model_id}', self.args.global_rank) |
| 26 | # |
| 27 | # #### RESET #### |
| 28 | # for name, param in self.model.named_parameters(): |
| 29 | # if name.find("loranew_") != -1: |
| 30 | # param.requires_grad = True |
| 31 | # elif name.find("lora_") != -1: |
| 32 | # param.requires_grad = False |
| 33 | # |
| 34 | #### SAVE #### |
| 35 | if self.args.output_dir is not None: |
| 36 | print_rank_0('saving the final model ...', self.args.global_rank) |
| 37 | |
| 38 | if self.args.global_rank == 0: |
| 39 | peft_model_id = os.path.join(self.args.output_dir, str(round)) |
| 40 | if not os.path.exists(peft_model_id): |
| 41 | os.makedirs(peft_model_id) |
| 42 | self.model.save_pretrained(peft_model_id) |
| 43 | self.tokenizer.save_pretrained(peft_model_id) |
| 44 | adapter_config_path = os.path.join(peft_model_id, 'adapter_config.json') |
| 45 | # read json, load to adapter_config: |
| 46 | with open(adapter_config_path, 'r') as f: |
| 47 | adapter_config = json.load(f) |
| 48 | # change "r_sum" in adapter_config to 0: |
| 49 | adapter_config['r_sum'] = 0 # This is the key point to be compatible with O_LoRA!!! |
| 50 | # save to json: |
| 51 | with open(adapter_config_path, 'w') as f: |
| 52 | json.dump(adapter_config, f) |
| 53 | print_rank_0(f'Successfully saving the final model to {peft_model_id}', self.args.global_rank) |
| 54 | |
| 55 | # if self.args.output_dir is not None: |
| 56 | # print_rank_0('saving model to ' + self.args.output_dir + "/" + str(round) + '...', self.args.global_rank) |
| 57 | # |
| 58 | # if self.args.global_rank == 0: |
| 59 | # save_hf_format(self.model, self.tokenizer, self.args, sub_folder=str(round)) |
| 60 | # |
| 61 | # if self.args.zero_stage == 3: |
| 62 | # # For zero stage 3, each gpu only has a part of the model, so we need a special save function |
| 63 | # save_zero_three_model(self.model, |
| 64 | # self.args.global_rank, |
| 65 | # self.args.output_dir, |
| 66 | # zero_stage=self.args.zero_stage, |
| 67 | # sub_folder=str(round)) |
| 68 | # print_rank_0('Successfully saving model after round {}'.format(round), self.args.global_rank) |
nothing calls this directly
no test coverage detected