| 103 | |
| 104 | |
| 105 | def wrap_cuda_model(args, model): |
| 106 | local_world_size = int(os.environ.get('LOCAL_WORLD_SIZE', 1)) |
| 107 | world_size = int(os.environ.get('WORLD_SIZE', 1)) |
| 108 | if args.train_engine == "torch_ddp": # native pytorch ddp |
| 109 | assert (torch.cuda.is_available()) |
| 110 | model.cuda() |
| 111 | model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True) |
| 112 | else: |
| 113 | if int(os.environ.get('RANK', 0)) == 0: |
| 114 | logging.info("Estimating model states memory needs (zero2)...") |
| 115 | estimate_zero2_model_states_mem_needs_all_live( |
| 116 | model, |
| 117 | num_gpus_per_node=local_world_size, |
| 118 | num_nodes=world_size // local_world_size) |
| 119 | return model |
| 120 | |
| 121 | def init_optimizer_and_scheduler(args, configs, model): |
| 122 | if configs['train_conf']['optim'] == 'adam': |