(model, opt, out_path: str, dcp=False)
| 190 | |
| 191 | @staticmethod |
| 192 | def save_optimizer(model, opt, out_path: str, dcp=False): |
| 193 | state_dict = get_optimizer_state_dict( |
| 194 | model, |
| 195 | opt, |
| 196 | options=StateDictOptions( |
| 197 | full_state_dict=False, |
| 198 | cpu_offload=True, |
| 199 | ), |
| 200 | ) |
| 201 | if dcp: |
| 202 | DCP.save(state_dict, checkpoint_id=out_path) |
| 203 | else: |
| 204 | torch.save(state_dict, out_path) |
| 205 | del state_dict |
| 206 | if dist.is_initialized(): |
| 207 | dist.barrier() |
| 208 | return |
| 209 | |
| 210 | @staticmethod |
| 211 | def load_ckpt_from_dir(rank, |