def train300_mlperf_coco(args): args = setup_distributed(args) # Build the model model_options = { 'use_nhwc': args.nhwc, 'pad_input': args.pad_input, 'bn_group': args.bn_group, } ssd300 = SSD300(args, args.num_classes, **model_options) if args.checkpoint is not None: load_checkpoint(ssd300, args.checkpoint) ssd300.train() ssd300.cuda() dboxes = dboxes300_coco() # Note: No reason not to use optimised loss loss_func = OptLoss() loss_func.cuda() # Create optimizer. This must also be done after network_to_half. global_batch_size = (args.N_gpu * args.batch_size) log_event(key=constants.MODEL_BN_SPAN, value=args.bn_group * args.batch_size) log_event(key=constants.GLOBAL_BATCH_SIZE, value=global_batch_size) # mlperf only allows base_lr scaled by an integer base_lr = 2.5e-3 requested_lr_multiplier = args.lr / base_lr adjusted_multiplier = max( 1, round(requested_lr_multiplier * global_batch_size / 32)) current_lr = base_lr * adjusted_multiplier current_momentum = 0.9 current_weight_decay = args.wd static_loss_scale = 128. optim = apex.optimizers.FusedSGD(ssd300.parameters(), lr=current_lr, momentum=current_momentum, weight_decay=current_weight_decay) ssd300, optim = apex.amp.initialize(ssd300, optim, opt_level='O2', loss_scale=static_loss_scale) # Parallelize. Need to do this after network_to_half. if args.distributed: if args.delay_allreduce: print_message(args.local_rank, "Delaying allreduces to the end of backward()") ssd300 = DDP(ssd300, gradient_predivide_factor=args.N_gpu / 8.0, delay_allreduce=args.delay_allreduce, retain_allreduce_buffers=args.use_fp16) log_event(key=constants.OPT_BASE_LR, value=current_lr) log_event(key=constants.OPT_LR_DECAY_BOUNDARY_EPOCHS, value=args.lr_decay_epochs) log_event(key=constants.OPT_LR_DECAY_STEPS, value=args.lr_decay_epochs) log_event(key=constants.OPT_WEIGHT_DECAY, value=current_weight_decay) if args.warmup is not None: log_event(key=constants.OPT_LR_WARMUP_STEPS, value=args.warmup) log_event(key=constants.OPT_LR_WARMUP_FACTOR, value=args.warmup_factor) # Model is completely finished -- need to create separate copies, preserve parameters across # them, and jit ssd300_eval = SSD300(args, args.num_classes, **model_options).cuda() if args.use_fp16: convert_network(ssd300_eval, torch.half) # Get the existant state from the train model # * if we use distributed, then we want .module train_model = ssd300.module if args.distributed else ssd300 ssd300_eval.load_state_dict(train_model.state_dict()) ssd300_eval.eval() print_message(args.local_rank, "epoch", "nbatch", "loss") iter_num = args.iteration avg_loss = 0.0 start_elapsed_time = time.time() last_printed_iter = args.iteration num_elapsed_samples = 0 input_c = 4 if args.pad_input else 3 example_shape = [args.batch_size, 300, 300, input_c ] if args.nhwc else [args.batch_size, input_c, 300, 300] example_input = torch.randn(*example_shape).cuda() if args.use_fp16: example_input = example_input.half() if args.jit: # DDP has some Python-side control flow. If we JIT the entire DDP-wrapped module, # the resulting ScriptModule will elide this control flow, resulting in allreduce # hooks not being called. If we're running distributed, we need to extract and JIT # the wrapped .module. # Replacing a DDP-ed ssd300 with a script_module might also cause the AccumulateGrad hooks # to go out of scope, and therefore silently disappear. module_to_jit = ssd300.module if args.distributed else ssd300 if args.distributed: ssd300.module = torch.jit.trace(module_to_jit, example_input, check_trace=False) else: ssd300 = torch.jit.trace(module_to_jit, example_input, check_trace=False) # JIT the eval model too ssd300_eval = torch.jit.trace(ssd300_eval, example_input, check_trace=False) # do a dummy fprop & bprop to make sure cudnnFind etc. are timed here ploc, plabel = ssd300(example_input) # produce a single dummy "loss" to make things easier loss = ploc[0, 0, 0] + plabel[0, 0, 0] dloss = torch.randn_like(loss) # Cause cudnnFind for dgrad, wgrad to run loss.backward(dloss) # Necessary import in init from pycocotools.coco import COCO encoder = build_ssd300_coder() evaluator = AsyncEvaluator(num_threads=1) log_end(key=constants.INIT_STOP) ##### END INIT # This is the first place we touch anything related to data ##### START DATA TOUCHING barrier() log_start(key=constants.RUN_START) barrier() train_pipe = prebuild_pipeline(args) train_loader, epoch_size = build_pipeline(args, training=True, pipe=train_pipe) if args.rank == 0: print("epoch size is: ", epoch_size, " images") val_loader, inv_map, cocoGt = build_pipeline(args, training=False) if args.profile_gc_off: gc.disable() gc.collect() ##### END DATA TOUCHING i_eval = 0 block_start_epoch = 1 log_start(key=constants.BLOCK_START, metadata={ 'first_epoch_num': block_start_epoch, 'epoch_count': args.evaluation[i_eval] }) for epoch in range(args.epochs): for p in ssd300.parameters(): p.grad = None if epoch in args.evaluation: # Get the existant state from the train model # * if we use distributed, then we want .module train_model = ssd300.module if args.distributed else ssd300 if args.distributed and args.allreduce_running_stats: if args.rank == 0: print("averaging bn running means and vars") # make sure every node has the same running bn stats before # using them to evaluate, or saving the model for inference world_size = float(torch.distributed.get_world_size()) for bn_name, bn_buf in train_model.named_buffers(recurse=True): if ('running_mean' in bn_name) or ('running_var' in bn_name): torch.distributed.all_reduce(bn_buf, op=dist.ReduceOp.SUM) bn_buf /= world_size if args.rank == 0: if args.save: print("saving model...") if not os.path.isdir('./models'): os.mkdir('./models') torch.save({"model": ssd300.state_dict()}, "./models/iter_{}.pt".format(iter_num)) ssd300_eval.load_state_dict(train_model.state_dict()) # Note: No longer returns, evaluation is abstracted away inside evaluator coco_eval(args, ssd300_eval, val_loader, cocoGt, encoder, inv_map, epoch, iter_num, evaluator=evaluator) log_end(key=constants.BLOCK_STOP, metadata={'first_epoch_num': block_start_epoch}) if epoch != max(args.evaluation): i_eval += 1 block_start_epoch = epoch + 1 log_start(key=constants.BLOCK_START, metadata={ 'first_epoch_num': block_start_epoch, 'epoch_count': (args.evaluation[i_eval] - args.evaluation[i_eval - 1]) }) if epoch in args.lr_decay_epochs: current_lr *= args.lr_decay_factor print_message( args.rank, "lr decay step #" + str(bisect(args.lr_decay_epochs, epoch))) for param_group in optim.param_groups: param_group['lr'] = current_lr log_start(key=constants.EPOCH_START, metadata={ 'epoch_num': epoch + 1, 'current_iter_num': iter_num }) for i, (img, bbox, label) in enumerate(train_loader): if args.profile_start is not None and iter_num == args.profile_start: torch.cuda.profiler.start() torch.cuda.synchronize() if args.profile_nvtx: torch.autograd._enable_profiler( torch.autograd.ProfilerState.NVTX) if args.profile is not None and iter_num == args.profile: if args.profile_start is not None and iter_num >= args.profile_start: # we turned cuda and nvtx profiling on, better turn it off too if args.profile_nvtx: torch.autograd._disable_profiler() torch.cuda.profiler.stop() return if args.warmup is not None: lr_warmup(optim, args.warmup, iter_num, epoch, current_lr, args) if (img is None) or (bbox is None) or (label is None): print("No labels in batch") continue ploc, plabel = ssd300(img) ploc, plabel = ploc.float(), plabel.float() N = img.shape[0] bbox.requires_grad = False label.requires_grad = False # reshape (N*8732X4 -> Nx8732x4) and transpose (Nx8732x4 -> Nx4x8732) bbox = bbox.view(N, -1, 4).transpose(1, 2).contiguous() # reshape (N*8732 -> Nx8732) and cast to Long label = label.view(N, -1).long() loss = loss_func(ploc, plabel, bbox, label) if np.isfinite(loss.item()): avg_loss = 0.999 * avg_loss + 0.001 * loss.item() else: print("model exploded (corrupted by Inf or Nan)") sys.exit() num_elapsed_samples += N if args.rank == 0 and iter_num % args.print_interval == 0: end_elapsed_time = time.time() elapsed_time = end_elapsed_time - start_elapsed_time avg_samples_per_sec = num_elapsed_samples * args.N_gpu / elapsed_time print("Iteration: {:6d}, Loss function: {:5.3f}, Average Loss: {:.3f}, avg. samples / sec: {:.2f}"\ .format(iter_num, loss.item(), avg_loss, avg_samples_per_sec), end="\n") last_printed_iter = iter_num start_elapsed_time = time.time() num_elapsed_samples = 0 with apex.amp.scale_loss(loss, optim) as scaled_loss: scaled_loss.backward() if not args.profile_fake_optim: optim.step() # Likely a decent skew here, let's take this opportunity to set the # gradients to None. After DALI integration, playing with the # placement of this is worth trying. for p in ssd300.parameters(): p.grad = None # Don't check every iteration due to cost of broadcast if iter_num % 20 == 0: finished = check_async_evals(args, evaluator, args.threshold) if finished: return True iter_num += 1 train_loader.reset() log_end(key=constants.EPOCH_STOP, metadata={'epoch_num': epoch + 1}) return False
def train300_mlperf_coco(args): from pycocotools.coco import COCO # Check that GPUs are actually available use_cuda = not args.no_cuda # Setup multi-GPU if necessary args.distributed = False if 'WORLD_SIZE' in os.environ: args.distributed = int(os.environ['WORLD_SIZE']) > 1 if args.distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend='nccl', init_method='env://') local_seed = set_seeds(args) # start timing here if args.distributed: N_gpu = torch.distributed.get_world_size() else: N_gpu = 1 validate_group_bn(args.bn_group) # Setup data, defaults dboxes = dboxes300_coco() encoder = Encoder(dboxes) input_size = 300 val_trans = SSDTransformer(dboxes, (input_size, input_size), val=True) val_annotate = os.path.join(args.data, "annotations/instances_val2017.json") val_coco_root = os.path.join(args.data, "val2017") train_annotate = os.path.join(args.data, "annotations/instances_train2017.json") train_coco_root = os.path.join(args.data, "train2017") # Build the model model_options = { 'backbone': args.backbone, 'use_nhwc': args.nhwc, 'pad_input': args.pad_input, 'bn_group': args.bn_group, } ssd300 = SSD300(args.num_classes, **model_options) if args.checkpoint is not None: load_checkpoint(ssd300, args.checkpoint) ssd300.train() ssd300.cuda() if args.opt_loss: loss_func = OptLoss(dboxes) else: loss_func = Loss(dboxes) loss_func.cuda() if args.distributed: N_gpu = torch.distributed.get_world_size() else: N_gpu = 1 if args.use_fp16: ssd300 = network_to_half(ssd300) # Parallelize. Need to do this after network_to_half. if args.distributed: if args.delay_allreduce: print_message(args.local_rank, "Delaying allreduces to the end of backward()") ssd300 = DDP(ssd300, gradient_predivide_factor=N_gpu / 8.0, delay_allreduce=args.delay_allreduce, retain_allreduce_buffers=args.use_fp16) # Create optimizer. This must also be done after network_to_half. global_batch_size = (N_gpu * args.batch_size) mlperf_print(key=mlperf_compliance.constants.MODEL_BN_SPAN, value=args.bn_group * args.batch_size) mlperf_print(key=mlperf_compliance.constants.GLOBAL_BATCH_SIZE, value=global_batch_size) # mlperf only allows base_lr scaled by an integer base_lr = 2.5e-3 requested_lr_multiplier = args.lr / base_lr adjusted_multiplier = max( 1, round(requested_lr_multiplier * global_batch_size / 32)) current_lr = base_lr * adjusted_multiplier current_momentum = 0.9 current_weight_decay = args.wd static_loss_scale = 128. if args.use_fp16: if args.distributed and not args.delay_allreduce: # We can't create the flat master params yet, because we need to # imitate the flattened bucket structure that DDP produces. optimizer_created = False else: model_buckets = [ [ p for p in ssd300.parameters() if p.requires_grad and p.type() == "torch.cuda.HalfTensor" ], [ p for p in ssd300.parameters() if p.requires_grad and p.type() == "torch.cuda.FloatTensor" ] ] flat_master_buckets = create_flat_master(model_buckets) optim = torch.optim.SGD(flat_master_buckets, lr=current_lr, momentum=current_momentum, weight_decay=current_weight_decay) optimizer_created = True else: optim = torch.optim.SGD(ssd300.parameters(), lr=current_lr, momentum=current_momentum, weight_decay=current_weight_decay) optimizer_created = True mlperf_print(key=mlperf_compliance.constants.OPT_BASE_LR, value=current_lr) mlperf_print(key=mlperf_compliance.constants.OPT_WEIGHT_DECAY, value=current_weight_decay) if args.warmup is not None: mlperf_print(key=mlperf_compliance.constants.OPT_LR_WARMUP_STEPS, value=args.warmup) mlperf_print(key=mlperf_compliance.constants.OPT_LR_WARMUP_FACTOR, value=args.warmup_factor) # Model is completely finished -- need to create separate copies, preserve parameters across # them, and jit ssd300_eval = SSD300(args.num_classes, backbone=args.backbone, use_nhwc=args.nhwc, pad_input=args.pad_input).cuda() if args.use_fp16: ssd300_eval = network_to_half(ssd300_eval) # Get the existant state from the train model # * if we use distributed, then we want .module train_model = ssd300.module if args.distributed else ssd300 ssd300_eval.load_state_dict(train_model.state_dict()) ssd300_eval.eval() print_message(args.local_rank, "epoch", "nbatch", "loss") eval_points = np.array(args.evaluation) * 32 / global_batch_size eval_points = list(map(int, list(eval_points))) iter_num = args.iteration avg_loss = 0.0 start_elapsed_time = time.time() last_printed_iter = args.iteration num_elapsed_samples = 0 # Generate normalization tensors mean, std = generate_mean_std(args) dummy_overflow_buf = torch.cuda.IntTensor([0]) def step_maybe_fp16_maybe_distributed(optim): if args.use_fp16: if args.distributed: for flat_master, allreduce_buffer in zip( flat_master_buckets, ssd300.allreduce_buffers): if allreduce_buffer is None: raise RuntimeError("allreduce_buffer is None") flat_master.grad = allreduce_buffer.float() flat_master.grad.data.mul_(1. / static_loss_scale) else: for flat_master, model_bucket in zip(flat_master_buckets, model_buckets): flat_grad = apex_C.flatten( [m.grad.data for m in model_bucket]) flat_master.grad = flat_grad.float() flat_master.grad.data.mul_(1. / static_loss_scale) optim.step() if args.use_fp16: # Use multi-tensor scale instead of loop & individual parameter copies for model_bucket, flat_master in zip(model_buckets, flat_master_buckets): multi_tensor_applier( amp_C.multi_tensor_scale, dummy_overflow_buf, [ apex_C.unflatten(flat_master.data, model_bucket), model_bucket ], 1.0) input_c = 4 if args.pad_input else 3 example_shape = [args.batch_size, 300, 300, input_c ] if args.nhwc else [args.batch_size, input_c, 300, 300] example_input = torch.randn(*example_shape).cuda() if args.use_fp16: example_input = example_input.half() if args.jit: # DDP has some Python-side control flow. If we JIT the entire DDP-wrapped module, # the resulting ScriptModule will elide this control flow, resulting in allreduce # hooks not being called. If we're running distributed, we need to extract and JIT # the wrapped .module. # Replacing a DDP-ed ssd300 with a script_module might also cause the AccumulateGrad hooks # to go out of scope, and therefore silently disappear. module_to_jit = ssd300.module if args.distributed else ssd300 if args.distributed: ssd300.module = torch.jit.trace(module_to_jit, example_input) else: ssd300 = torch.jit.trace(module_to_jit, example_input) # JIT the eval model too ssd300_eval = torch.jit.trace(ssd300_eval, example_input) # do a dummy fprop & bprop to make sure cudnnFind etc. are timed here ploc, plabel = ssd300(example_input) # produce a single dummy "loss" to make things easier loss = ploc[0, 0, 0] + plabel[0, 0, 0] dloss = torch.randn_like(loss) # Cause cudnnFind for dgrad, wgrad to run loss.backward(dloss) mlperf_print(key=mlperf_compliance.constants.INIT_STOP, sync=True) ##### END INIT # This is the first place we touch anything related to data ##### START DATA TOUCHING mlperf_print(key=mlperf_compliance.constants.RUN_START, sync=True) barrier() cocoGt = COCO(annotation_file=val_annotate, use_ext=True) val_coco = COCODetection(val_coco_root, val_annotate, val_trans) if args.distributed: val_sampler = GeneralDistributedSampler(val_coco, pad=False) else: val_sampler = None if args.no_dali: train_trans = SSDTransformer(dboxes, (input_size, input_size), val=False) train_coco = COCODetection(train_coco_root, train_annotate, train_trans) if args.distributed: train_sampler = GeneralDistributedSampler(train_coco, pad=False) else: train_sampler = None train_loader = DataLoader(train_coco, batch_size=args.batch_size * args.input_batch_multiplier, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, collate_fn=partial(my_collate, is_training=True)) else: train_pipe = COCOPipeline(args.batch_size * args.input_batch_multiplier, args.local_rank, train_coco_root, train_annotate, N_gpu, num_threads=args.num_workers, output_fp16=args.use_fp16, output_nhwc=args.nhwc, pad_output=args.pad_input, seed=local_seed - 2**31, use_nvjpeg=args.use_nvjpeg, use_roi=args.use_roi_decode, dali_cache=args.dali_cache, dali_async=(not args.dali_sync)) print_message(args.local_rank, "time_check a: {secs:.9f}".format(secs=time.time())) train_pipe.build() print_message(args.local_rank, "time_check b: {secs:.9f}".format(secs=time.time())) test_run = train_pipe.run() train_loader = SingleDaliIterator( train_pipe, [ 'images', DALIOutput('bboxes', False, True), DALIOutput('labels', True, True) ], train_pipe.epoch_size()['train_reader'], ngpu=N_gpu) train_loader = EncodingInputIterator(train_loader, dboxes=encoder.dboxes.cuda(), nhwc=args.nhwc, fake_input=args.fake_input, no_dali=args.no_dali) if args.input_batch_multiplier > 1: train_loader = RateMatcher(input_it=train_loader, output_size=args.batch_size) val_dataloader = DataLoader( val_coco, batch_size=args.eval_batch_size, shuffle=False, # Note: distributed sampler is shuffled :( sampler=val_sampler, num_workers=args.num_workers) inv_map = {v: k for k, v in val_coco.label_map.items()} ##### END DATA TOUCHING i_eval = 0 first_epoch = 1 mlperf_print(key=mlperf_compliance.constants.BLOCK_START, metadata={ 'first_epoch_num': first_epoch, 'epoch_count': args.evaluation[i_eval] * 32 / train_pipe.epoch_size()['train_reader'] }, sync=True) for epoch in range(args.epochs): mlperf_print(key=mlperf_compliance.constants.EPOCH_START, metadata={'epoch_num': epoch + 1}, sync=True) for p in ssd300.parameters(): p.grad = None for i, (img, bbox, label) in enumerate(train_loader): if args.profile_start is not None and iter_num == args.profile_start: torch.cuda.profiler.start() torch.cuda.synchronize() if args.profile_nvtx: torch.autograd._enable_profiler( torch.autograd.ProfilerState.NVTX) if args.profile is not None and iter_num == args.profile: if args.profile_start is not None and iter_num >= args.profile_start: # we turned cuda and nvtx profiling on, better turn it off too if args.profile_nvtx: torch.autograd._disable_profiler() torch.cuda.profiler.stop() return if args.warmup is not None and optimizer_created: lr_warmup(optim, args.warmup, iter_num, epoch, current_lr, args) if iter_num == ((args.decay1 * 1000 * 32) // global_batch_size): print_message(args.local_rank, "lr decay step #1") current_lr *= 0.1 for param_group in optim.param_groups: param_group['lr'] = current_lr if iter_num == ((args.decay2 * 1000 * 32) // global_batch_size): print_message(args.local_rank, "lr decay step #2") current_lr *= 0.1 for param_group in optim.param_groups: param_group['lr'] = current_lr if (img is None) or (bbox is None) or (label is None): print("No labels in batch") continue ploc, plabel = ssd300(img) ploc, plabel = ploc.float(), plabel.float() N = img.shape[0] gloc, glabel = Variable(bbox, requires_grad=False), \ Variable(label, requires_grad=False) loss = loss_func(ploc, plabel, gloc, glabel) if np.isfinite(loss.item()): avg_loss = 0.999 * avg_loss + 0.001 * loss.item() else: print("model exploded (corrupted by Inf or Nan)") sys.exit() num_elapsed_samples += N if args.local_rank == 0 and iter_num % args.print_interval == 0: end_elapsed_time = time.time() elapsed_time = end_elapsed_time - start_elapsed_time avg_samples_per_sec = num_elapsed_samples * N_gpu / elapsed_time print("Iteration: {:6d}, Loss function: {:5.3f}, Average Loss: {:.3f}, avg. samples / sec: {:.2f}"\ .format(iter_num, loss.item(), avg_loss, avg_samples_per_sec), end="\n") last_printed_iter = iter_num start_elapsed_time = time.time() num_elapsed_samples = 0 # loss scaling if args.use_fp16: loss = loss * static_loss_scale loss.backward() if not optimizer_created: # Imitate the model bucket structure created by DDP. # These will already be split by type (float or half). model_buckets = [] for bucket in ssd300.active_i_buckets: model_buckets.append([]) for active_i in bucket: model_buckets[-1].append( ssd300.active_params[active_i]) flat_master_buckets = create_flat_master(model_buckets) optim = torch.optim.SGD(flat_master_buckets, lr=current_lr, momentum=current_momentum, weight_decay=current_weight_decay) optimizer_created = True # Skip this first iteration because flattened allreduce buffers are not yet created. # step_maybe_fp16_maybe_distributed(optim) else: step_maybe_fp16_maybe_distributed(optim) # Likely a decent skew here, let's take this opportunity to set the gradients to None. # After DALI integration, playing with the placement of this is worth trying. for p in ssd300.parameters(): p.grad = None if iter_num in eval_points: # Get the existant state from the train model # * if we use distributed, then we want .module train_model = ssd300.module if args.distributed else ssd300 if args.distributed and args.allreduce_running_stats: if get_rank() == 0: print("averaging bn running means and vars") # make sure every node has the same running bn stats before # using them to evaluate, or saving the model for inference world_size = float(torch.distributed.get_world_size()) for bn_name, bn_buf in train_model.named_buffers( recurse=True): if ('running_mean' in bn_name) or ('running_var' in bn_name): torch.distributed.all_reduce(bn_buf, op=dist.ReduceOp.SUM) bn_buf /= world_size if get_rank() == 0: if not args.no_save: print("saving model...") torch.save( { "model": ssd300.state_dict(), "label_map": val_coco.label_info }, "./models/iter_{}.pt".format(iter_num)) ssd300_eval.load_state_dict(train_model.state_dict()) succ = coco_eval( ssd300_eval, val_dataloader, cocoGt, encoder, inv_map, args.threshold, epoch, iter_num, args.eval_batch_size, use_fp16=args.use_fp16, local_rank=args.local_rank if args.distributed else -1, N_gpu=N_gpu, use_nhwc=args.nhwc, pad_input=args.pad_input) mlperf_print(key=mlperf_compliance.constants.BLOCK_STOP, metadata={'first_epoch_num': first_epoch}, sync=True) if succ: return True if iter_num != max(eval_points): i_eval += 1 first_epoch = epoch + 1 mlperf_print(key=mlperf_compliance.constants.BLOCK_START, metadata={ 'first_epoch_num': first_epoch, 'epoch_count': (args.evaluation[i_eval] - args.evaluation[i_eval - 1]) * 32 / train_pipe.epoch_size()['train_reader'] }, sync=True) iter_num += 1 if args.max_iter > 0: if iter_num > args.max_iter: break train_loader.reset() mlperf_print(key=mlperf_compliance.constants.EPOCH_STOP, metadata={'epoch_num': epoch + 1}, sync=True) return False
def main(): args = parse_arguments() status = 'aborted' # later set to 'success' if termination criteria met mlperf_logger.log_start(key=mlperf_logger.constants.INIT_START, log_all_ranks=True, sync=False) if args.use_env and 'LOCAL_RANK' in os.environ: args.local_rank = int(os.environ['LOCAL_RANK']) device, args = setup_training(args) mlperf_logger.mlperf_submission_log('bert') worker_seeds, shuffling_seeds = utils.setup_seeds(args.seed, args.num_epochs_to_generate_seeds_for, device) worker_seed = worker_seeds[torch.distributed.get_rank()] random.seed(worker_seed) np.random.seed(worker_seed) torch.manual_seed(worker_seed) worker_init = WorkerInitObj(worker_seed) mlperf_logger.log_event(key=mlperf_logger.constants.SEED, value=args.seed, sync=False) mlperf_logger.log_event(key=mlperf_logger.constants.GLOBAL_BATCH_SIZE, value=global_batch_size(args), sync=False) mlperf_logger.log_event(key='opt_gradient_accumulation_steps', value=args.gradient_accumulation_steps, sync=False) mlperf_logger.log_event(key='max_predictions_per_seq', value=args.max_predictions_per_seq, sync=False) mlperf_logger.log_event(key='opt_learning_rate_training_steps', value=args.max_steps, sync=False) mlperf_logger.log_event(key='num_warmup_steps', value=int(args.warmup_proportion*args.max_steps) if args.warmup_steps==0 else args.warmup_steps, sync=False) if utils.is_main_process(): print("parsed args:") print(args) # Prepare optimizer model, optimizer, lr_scheduler, checkpoint, global_step = prepare_model_and_optimizer(args, device) samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu if args.unpad: torch.cuda.synchronize() InitMHACUDAExtension() torch.cuda.synchronize() final_loss = float("inf") train_time_raw = float("inf") raw_train_start = time.time() if args.do_train: model.train() most_recent_ckpts_paths = [] average_loss = 0.0 # averaged loss every args.log_freq steps epoch = 1 training_steps = 0 end_training, converged = False, False samples_trained_prev = 0 eval_count = 0 pool = ProcessPoolExecutor(1) if args.target_mlm_accuracy: if args.train_mlm_accuracy_window_size > 0: accuracy_scores = [] avg_mlm_accuracy = torch.Tensor([0]).cuda() first_epoch = True if found_resume_checkpoint(args): f_start_id = checkpoint['files'][0] files = checkpoint['files'][1:] num_files = len(files) else: files = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if os.path.isfile(os.path.join(args.input_dir, f)) and 'part' in f] files.sort() num_files = len(files) random.Random(shuffling_seeds[epoch]).shuffle(files) f_start_id = 0 mlperf_logger.log_end(key=mlperf_logger.constants.INIT_STOP, sync=False) mlperf_logger.log_start(key=mlperf_logger.constants.RUN_START, sync=True) mlperf_logger.barrier() # Start prefetching eval dataset if args.eval_dir: eval_dataset_future = pool.submit(create_eval_dataset, args, worker_init_fn=worker_init) while global_step < args.max_steps and not end_training: mlperf_logger.log_start(key=mlperf_logger.constants.EPOCH_START, metadata={'epoch_num': epoch}, sync=False) mlperf_logger.log_start(key=mlperf_logger.constants.BLOCK_START, metadata={'first_epoch_num': epoch, 'epoch_count': 1}, sync=False) if utils.is_main_process(): print("parsed args:") print(args) now_time = time.time() now_step = global_step now_skipped = skipped_steps print("epoch:", epoch) thread = None # Reshuffle file list on subsequent epochs if not first_epoch: files = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if os.path.isfile(os.path.join(args.input_dir, f)) and 'part' in f] files.sort() num_files = len(files) random.Random(shuffling_seeds[epoch]).shuffle(files) f_start_id = 0 first_epoch = False shared_file_list = {} if torch.distributed.is_initialized() and torch.distributed.get_world_size() > num_files: remainder = torch.distributed.get_world_size() % num_files data_file = files[(f_start_id*torch.distributed.get_world_size() + torch.distributed.get_rank() + remainder * f_start_id) % num_files] else: data_file = files[(f_start_id*torch.distributed.get_world_size() + torch.distributed.get_rank()) % num_files] previous_file = data_file train_data = pretraining_dataset(data_file, args.max_predictions_per_seq) train_sampler = RandomSampler(train_data) train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size, num_workers=4, worker_init_fn=worker_init, pin_memory=True) overflow_buf = None if args.allreduce_post_accumulation: overflow_buf = torch.cuda.IntTensor([0]) for f_id in range(f_start_id + 1, len(files)): if torch.distributed.get_world_size() > num_files: data_file = files[(f_id*torch.distributed.get_world_size() + torch.distributed.get_rank() + remainder * f_id) % num_files] else: data_file = files[(f_id*torch.distributed.get_world_size() + torch.distributed.get_rank())%num_files] previous_file = data_file dataset_future = pool.submit(create_pretraining_dataset, data_file, args.max_predictions_per_seq, shared_file_list, args, worker_init_fn=worker_init) for step, batch in enumerate(train_dataloader): training_steps += 1 update_step = training_steps % args.gradient_accumulation_steps == 0 batch = [t.to(device) for t in batch] input_ids, segment_ids, input_mask, masked_lm_labels, next_sentence_labels = batch loss, mlm_acc, _ = model(input_ids=input_ids, token_type_ids=segment_ids, attention_mask=input_mask, masked_lm_labels=masked_lm_labels, next_sentence_label=next_sentence_labels, checkpoint_activations=args.checkpoint_activations) divisor = args.gradient_accumulation_steps if args.gradient_accumulation_steps > 1: if not args.allreduce_post_accumulation: # this division was merged into predivision loss = loss / args.gradient_accumulation_steps divisor = 1.0 if args.fp16: with amp.scale_loss(loss, optimizer, delay_overflow_check=args.allreduce_post_accumulation) as scaled_loss: scaled_loss.backward() else: loss.backward() average_loss += loss.item() if update_step: lr_scheduler.step() # learning rate warmup global_step = take_optimizer_step(args, optimizer, model, overflow_buf, global_step) samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu if (args.eval_dir and args.eval_iter_samples > 0 and samples_trained >= args.eval_iter_start_samples + eval_count * args.eval_iter_samples): # on first eval, get eval_dataloader if eval_count == 0: eval_dataloader = eval_dataset_future.result(timeout=None) samples_trained_prev = samples_trained eval_avg_loss, eval_avg_mlm_accuracy = run_eval(model, eval_dataloader, device, args.num_eval_examples, first_eval=(eval_count == 0), use_cache=args.cache_eval_data) if utils.is_main_process(): mlperf_logger.log_event(key=mlperf_logger.constants.EVAL_ACCURACY, value=eval_avg_mlm_accuracy, metadata={'epoch_num': epoch}, sync=False) print({"global_steps": global_step, "eval_loss": eval_avg_loss, "eval_mlm_accuracy":eval_avg_mlm_accuracy}) if args.target_mlm_accuracy: if eval_avg_mlm_accuracy >= args.target_mlm_accuracy: end_training, converged = True, True if utils.is_main_process(): print("%f > %f, Target MLM Accuracy reached at %d"%(eval_avg_mlm_accuracy, args.target_mlm_accuracy, global_step)) eval_count += 1 if args.target_mlm_accuracy and args.train_mlm_accuracy_window_size > 0: accuracy_scores.append(mlm_acc) if update_step: accuracy_scores = accuracy_scores[-args.train_mlm_accuracy_window_size * args.gradient_accumulation_steps:] avg_mlm_accuracy[0] = sum(accuracy_scores) / len(accuracy_scores) torch.distributed.all_reduce(avg_mlm_accuracy, op=torch.distributed.ReduceOp.SUM) avg_mlm_accuracy /= torch.distributed.get_world_size() if training_steps % (args.log_freq * args.gradient_accumulation_steps) == 0: samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu if utils.is_main_process(): time_interval = time.time() - now_time step_interval = global_step - now_step skip_interval = skipped_steps - now_skipped now_time = time.time() now_step = global_step now_skipped = skipped_steps training_perf = args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu \ * (step_interval + skip_interval) / time_interval if args.train_mlm_accuracy_window_size > 0: print({"training_steps": training_steps, "average_loss": average_loss / (args.log_freq * divisor), "step_loss": loss.item() * args.gradient_accumulation_steps / divisor, "learning_rate": optimizer.param_groups[0]['lr'], "seq/s": training_perf, "global_steps": now_step, "samples_trained": samples_trained, "skipped_steps": now_skipped, "timestamp": now_time, "mlm_accuracy": avg_mlm_accuracy[0].item()}) else: print({"training_steps": training_steps, "average_loss": average_loss / (args.log_freq * divisor), "step_loss": loss.item() * args.gradient_accumulation_steps / divisor, "learning_rate": optimizer.param_groups[0]['lr'], "seq/s": training_perf, "global_steps": now_step, "samples_trained": samples_trained, "skipped_steps": now_skipped, "timestamp": now_time}) average_loss = 0 if global_step >= args.max_steps or end_training: status = 'success' if converged else 'aborted' end_training = True train_time_raw = time.time() - raw_train_start last_num_steps = int(training_steps / args.gradient_accumulation_steps) % args.log_freq last_num_steps = args.log_freq if last_num_steps == 0 else last_num_steps average_loss = torch.tensor(average_loss, dtype=torch.float32).cuda() average_loss = average_loss / (last_num_steps * divisor) if (torch.distributed.is_initialized()): average_loss /= torch.distributed.get_world_size() torch.distributed.all_reduce(average_loss) final_loss = average_loss.item() if utils.is_main_process(): if args.train_mlm_accuracy_window_size > 0: print((epoch, training_steps / args.gradient_accumulation_steps, ), {"final_loss": final_loss, "final_mlm_accuracy": avg_mlm_accuracy[0].item()}) else: print((epoch, training_steps / args.gradient_accumulation_steps, ), {"final_loss": final_loss}) if end_training or (samples_trained - samples_trained_prev >= args.num_samples_per_checkpoint and samples_trained >= args.min_samples_to_start_checkpoints): samples_trained_prev = samples_trained if utils.is_main_process() and not args.skip_checkpoint: # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self if args.phase2: output_save_file = os.path.join(args.output_dir, "phase2_ckpt_{}.pt".format(samples_trained)) else: output_save_file = os.path.join(args.output_dir, "phase1_ckpt_{}.pt".format(samples_trained)) if args.do_train: torch.save({'model': model_to_save.state_dict(), 'optimizer': optimizer.state_dict(), 'master params': list(amp.master_params(optimizer)), 'files': [f_id] + files}, output_save_file) most_recent_ckpts_paths.append(output_save_file) if len(most_recent_ckpts_paths) > args.keep_n_most_recent_checkpoints: ckpt_to_be_removed = most_recent_ckpts_paths.pop(0) os.remove(ckpt_to_be_removed) if samples_trained >= args.max_samples_termination or end_training: status = 'success' if converged else 'aborted' end_training = True break del train_dataloader if samples_trained >= args.max_samples_termination or end_training: status = 'success' if converged else 'aborted' end_training = True break train_dataloader, data_file = dataset_future.result(timeout=None) mlperf_logger.log_end(key=mlperf_logger.constants.BLOCK_STOP, metadata={'first_epoch_num': epoch}, sync=False) mlperf_logger.log_end(key=mlperf_logger.constants.EPOCH_STOP, metadata={'epoch_num': epoch}, sync=False) epoch += 1 mlperf_logger.log_event(key=mlperf_logger.constants.TRAIN_SAMPLES, value=samples_trained, sync=False) mlperf_logger.log_event(key=mlperf_logger.constants.EVAL_SAMPLES, value=args.num_eval_examples, sync=False) mlperf_logger.log_end(key=mlperf_logger.constants.RUN_STOP, metadata={'status': status}, sync=False) return args, final_loss, train_time_raw