(args, model, optimizer, loss_scaler)
| 406 | os.system(f'rm {ckpt} -rf') |
| 407 | |
| 408 | def load_model(args, model, optimizer, loss_scaler): |
| 409 | start_iter = 0 |
| 410 | start_epoch = 0 |
| 411 | if args.auto_resume: |
| 412 | ckpt_dirs = glob.glob(os.path.join(args.output_dir, "iter_*")) + glob.glob(os.path.join(args.output_dir, "epoch_*")) |
| 413 | ckpt_dirs.sort() |
| 414 | if len(ckpt_dirs) > 0: |
| 415 | args.resume = ckpt_dirs[-1] |
| 416 | if args.resume: |
| 417 | print("Resume checkpoint %s" % args.resume) |
| 418 | local_checkpoint_path = os.path.join( |
| 419 | args.resume, |
| 420 | f"checkpoint.{dist.get_rank():05d}-of-{dist.get_world_size():05d}.pth", |
| 421 | ) |
| 422 | with load_with_process_group(fs_init.get_data_parallel_group()): |
| 423 | checkpoint = torch.load(local_checkpoint_path, map_location='cpu') |
| 424 | with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): |
| 425 | model.load_state_dict(checkpoint['model']) |
| 426 | optimizer.load_state_dict(checkpoint['optimizer']) |
| 427 | loss_scaler.load_state_dict(checkpoint['scaler']) |
| 428 | start_iter = int(checkpoint['iter']) + 1 |
| 429 | if 'epoch' in checkpoint: |
| 430 | start_epoch = int(checkpoint['epoch']) |
| 431 | return start_epoch, start_iter |
| 432 | |
| 433 | def all_reduce_mean(x): |
| 434 | world_size = get_world_size() |
nothing calls this directly
no test coverage detected