(model_ema, global_rank, save_dir, zero_stage=0, sub_folder="")
| 169 | |
| 170 | |
| 171 | def save_zero_three_model(model_ema, global_rank, save_dir, zero_stage=0, sub_folder=""): |
| 172 | zero_stage_3 = (zero_stage == 3) |
| 173 | save_dir = os.path.join(save_dir, sub_folder) |
| 174 | os.makedirs(save_dir, exist_ok=True) |
| 175 | WEIGHTS_NAME = "pytorch_model.bin" |
| 176 | output_model_file = os.path.join(save_dir, WEIGHTS_NAME) |
| 177 | |
| 178 | model_to_save = model_ema.module if hasattr(model_ema, |
| 179 | 'module') else model_ema |
| 180 | if not zero_stage_3: |
| 181 | if global_rank == 0: |
| 182 | torch.save(model_to_save.state_dict(), output_model_file) |
| 183 | else: |
| 184 | output_state_dict = {} |
| 185 | for k, v in model_to_save.named_parameters(): |
| 186 | |
| 187 | if hasattr(v, 'ds_id'): |
| 188 | with deepspeed.zero.GatheredParameters(_z3_params_to_fetch([v |
| 189 | ]), |
| 190 | enabled=zero_stage_3): |
| 191 | v_p = v.data.cpu() |
| 192 | else: |
| 193 | v_p = v.cpu() |
| 194 | if global_rank == 0 and "lora" not in k: |
| 195 | output_state_dict[k] = v_p |
| 196 | if global_rank == 0: |
| 197 | torch.save(output_state_dict, output_model_file) |
| 198 | del output_state_dict |
no test coverage detected