(rank,
load_save_dict=None,
model_for_opt: tuple = None,
optimizer=None,
ckpt_dir: str = None,
dcp=False,
lora=False)
| 251 | |
| 252 | @staticmethod |
| 253 | def save_ckpt_to_dir(rank, |
| 254 | load_save_dict=None, |
| 255 | model_for_opt: tuple = None, |
| 256 | optimizer=None, |
| 257 | ckpt_dir: str = None, |
| 258 | dcp=False, |
| 259 | lora=False): |
| 260 | os.makedirs(ckpt_dir, exist_ok=True) |
| 261 | if load_save_dict is not None: |
| 262 | for model_name, model in load_save_dict.items(): |
| 263 | if not dcp: |
| 264 | out_path = os.path.join(ckpt_dir, f'{model_name}.pt') |
| 265 | else: |
| 266 | out_path = os.path.join(ckpt_dir, model_name) |
| 267 | TrainerBase.save_model(rank, model, out_path, dcp, lora) |
| 268 | if optimizer is not None: |
| 269 | for model_name in model_for_opt: |
| 270 | out_path = os.path.join(ckpt_dir, f'{model_name}-opt') |
| 271 | TrainerBase.save_optimizer( |
| 272 | load_save_dict[model_name], optimizer, out_path, dcp=True) |
| 273 | return None |
| 274 | |
| 275 | def save_ckpt(self, |
| 276 | rank, |
no test coverage detected