(args)
| 49 | |
| 50 | |
| 51 | def main(args): |
| 52 | is_training = args.mode == 'train' |
| 53 | |
| 54 | train_target = args.experiment.get('train_target', ['transformer']) |
| 55 | train_transformer = 'transformer' in train_target |
| 56 | dist.init_process_group('nccl') |
| 57 | global_rank = dist.get_rank() |
| 58 | is_main_process = global_rank == 0 |
| 59 | device = torch.device(global_rank % torch.cuda.device_count()) |
| 60 | torch.cuda.set_device(device) |
| 61 | |
| 62 | device_mesh_dp_hybrid = get_device_mesh(use_hybrid=True, tp_size=None)# ? |
| 63 | device_mesh_dp_tp = get_device_mesh( |
| 64 | use_hybrid=False, tp_size=args.parallel.tp_size) |
| 65 | global_device_mesh = get_device_mesh(use_hybrid=False, tp_size=None) |
| 66 | dp_mesh = device_mesh_dp_tp['dp'] |
| 67 | world_size = dist.get_world_size() |
| 68 | |
| 69 | dp_rank = dp_mesh.get_local_rank() |
| 70 | dropout_generator = torch.Generator(device) |
| 71 | dropout_generator.manual_seed(dp_rank + int(args.experiment.global_seed)) |
| 72 | global_dp_rank = global_device_mesh.get_local_rank() |
| 73 | loglevel = args.experiment.get('loglevel', 'INFO').upper() |
| 74 | trainer = TrainerBase( |
| 75 | args.experiment.result_dir, log_level=loglevel, rank=global_rank, mode=args.mode) |
| 76 | if global_rank == 0: |
| 77 | trainer.save_config(args) |
| 78 | dist.barrier() |
| 79 | logger, tb_tracker, timer = trainer.logger, trainer.tb_tracker, trainer.timer |
| 80 | |
| 81 | result_folder = os.path.join(trainer.vis_dir, args.mbench_name) |
| 82 | os.makedirs(result_folder, exist_ok=True) |
| 83 | logger.info(f'result_folder: {result_folder}') |
| 84 | logger.info(f'text_key: {args.dataset.text_key}') |
| 85 | |
| 86 | dtype_mapping = dict( |
| 87 | bf16=torch.bfloat16, fp32=torch.float32, fp16=torch.float16) |
| 88 | dtype = dtype_mapping[args.precision.mixed_precision] |
| 89 | grad_dtype = dtype_mapping[args.precision.grad_precision] |
| 90 | |
| 91 | logger.info( |
| 92 | f'dtype {dtype}, grad_dtype {grad_dtype}' |
| 93 | ) |
| 94 | dist.barrier() |
| 95 | |
| 96 | ref_corruption_cfg = args.get('ref_motion_corruption', {}) |
| 97 | train_ref_corruption_cfg = None |
| 98 | if ref_corruption_cfg.get('enable', False): |
| 99 | train_ref_corruption_cfg = ref_corruption_cfg |
| 100 | |
| 101 | base_repo_path = args.model_path[args.experiment.model_name] |
| 102 | |
| 103 | resume_path, resume_step = trainer.get_resume_path_and_step( |
| 104 | auto_resume=args.experiment.auto_resume, |
| 105 | resume_path=args.experiment.resume_path) |
| 106 | |
| 107 | patch_size = 2 |
| 108 | in_channel = args.model.get('in_channels', 16) |
no test coverage detected