(model, tokenizer, args, sub_folder="")
| 105 | ] |
| 106 | |
| 107 | def save_zero_three_model(model, tokenizer, args, sub_folder=""): |
| 108 | zero_stage_3 = (args.zero_stage == 3) |
| 109 | os.makedirs(args.output_dir, exist_ok=True) |
| 110 | if args.global_rank == 0: |
| 111 | output_dir = os.path.join(args.output_dir, sub_folder) |
| 112 | os.makedirs(output_dir, exist_ok=True) |
| 113 | model_to_save = model.module if hasattr(model, 'module') else model |
| 114 | |
| 115 | if zero_stage_3: |
| 116 | output_state_dict = {} |
| 117 | for k, v in model_to_save.named_parameters(): |
| 118 | |
| 119 | if hasattr(v, 'ds_id'): |
| 120 | with deepspeed.zero.GatheredParameters(_z3_params_to_fetch([v]), |
| 121 | enabled=zero_stage_3): |
| 122 | v_p = v.data.cpu() |
| 123 | else: |
| 124 | v_p = v.cpu() |
| 125 | if args.global_rank == 0 and "lora" not in k: |
| 126 | output_state_dict[k] = v_p |
| 127 | if args.global_rank == 0: |
| 128 | model_to_save.save_pretrained(output_dir, state_dict=output_state_dict) |
| 129 | del output_state_dict |
| 130 | |
| 131 | if args.global_rank == 0: |
| 132 | output_config_file = os.path.join(output_dir, "config.json") |
| 133 | model_to_save.config.to_json_file(output_config_file) |
| 134 | tokenizer.save_pretrained(output_dir) |
| 135 | # for models not in AutoModel, copy python module files |
| 136 | train_from_model_path = model_to_save.config._name_or_path |
| 137 | if os.path.exists(train_from_model_path): |
| 138 | for filename in os.listdir(train_from_model_path): |
| 139 | if filename.endswith(".py"): |
| 140 | shutil.copy(os.path.join(train_from_model_path, filename), os.path.join(output_dir, filename)) |
no test coverage detected