def ssgd(dnn, dataset, data_dir, nworkers, lr, batch_size, nsteps_update, max_epochs, nwpernode, pretrain, num_steps, compressor, density, threshold, gradient_path=None): rank = hvd.rank() torch.cuda.set_device(rank % nwpernode) if rank != 0: pretrain = None trainer = DLTrainer(rank, nworkers, dist=False, batch_size=batch_size, is_weak_scaling=True, ngpus=1, data_dir=data_dir, dataset=dataset, dnn=dnn, lr=lr, nworkers=nworkers, prefix='allreduce', pretrain=pretrain, num_steps=num_steps, tb_writer=writer) init_epoch = torch.ones(1) * trainer.get_train_epoch() init_iter = torch.ones(1) * trainer.get_train_iter() trainer.set_train_epoch(int(hvd.broadcast(init_epoch, root_rank=0)[0])) trainer.set_train_iter(int(hvd.broadcast(init_iter, root_rank=0)[0])) is_sparse = density < 1 #if not is_sparse: # compressor = None if settings.ADAPTIVE_MERGE or settings.ADAPTIVE_SPARSE: seq_layernames, layerwise_times, layerwise_sizes = benchmark(trainer) layerwise_times = comm.bcast(layerwise_times, root=0) if rank == 0: logger.info('layerwise backward times: %s', list(layerwise_times)) logger.info('layerwise backward sizes: %s', list(layerwise_sizes)) logger.info('Bencharmked backward time: %f', np.sum(layerwise_times)) logger.info('Model size: %d', np.sum(layerwise_sizes)) else: seq_layernames, layerwise_times, layerwise_sizes = None, None, None norm_clip = None if dnn == 'lstm': norm_clip = 0.25 elif dnn == 'lstman4': norm_clip = 400 optimizer = hvd.DistributedOptimizer( trainer.optimizer, named_parameters=trainer.net.named_parameters(), compression=compressors[compressor], is_sparse=is_sparse, density=density, seq_layernames=seq_layernames, layerwise_times=layerwise_times, norm_clip=norm_clip, threshold=threshold, writer=writer, gradient_path=gradient_path) hvd.broadcast_parameters(trainer.net.state_dict(), root_rank=0) trainer.update_optimizer(optimizer) iters_per_epoch = trainer.get_num_of_training_samples() // ( nworkers * batch_size * nsteps_update) times = [] logger.info('max_epochs: %d', max_epochs) display = 40 if iters_per_epoch > 40 else iters_per_epoch - 1 for epoch in range(max_epochs): hidden = None if dnn == 'lstm': hidden = trainer.net.init_hidden() for i in range(iters_per_epoch): s = time.time() optimizer.zero_grad() for j in range(nsteps_update): if j < nsteps_update - 1 and nsteps_update > 1: optimizer.local = True else: optimizer.local = False if dnn == 'lstm': _, hidden = trainer.train(1, hidden=hidden) else: trainer.train(1) if dnn == 'lstm': optimizer.synchronize() torch.nn.utils.clip_grad_norm_(trainer.net.parameters(), 0.25) elif dnn == 'lstman4': optimizer.synchronize() torch.nn.utils.clip_grad_norm_(trainer.net.parameters(), 400) trainer.update_model() times.append(time.time() - s) if i % display == 0 and i > 0: time_per_iter = np.mean(times) logger.warn( 'Time per iteration including communication: %f, Speed: %f images/s', time_per_iter, batch_size * nsteps_update / time_per_iter) times = [] optimizer.increase_one_epoch()
batch_size = args.batch_size * args.nsteps_update prefix = settings.PREFIX if args.density < 1: prefix = 'comp-' + args.compressor + '-' + prefix logdir = 'allreduce-%s-thres-%dkbytes/%s-n%d-bs%d-lr%.4f-ns%d-ds%s' % ( prefix, args.threshold / 1024, args.dnn, args.nworkers, batch_size, args.lr, args.nsteps_update, str(args.density)) relative_path = './logs/%s' % logdir gradient_relative_path = None utils.create_path(relative_path) if settings.LOGGING_GRADIENTS: gradient_relative_path = '%s/gradients/%s' % (args.saved_dir, logdir) utils.create_path(gradient_relative_path) rank = 0 if args.nworkers > 1: hvd.init() rank = hvd.rank() if rank == 0: tb_runs = './runs/%s' % logdir writer = None #SummaryWriter(tb_runs) logfile = os.path.join(relative_path, settings.hostname + '-' + str(rank) + '.log') hdlr = logging.FileHandler(logfile) hdlr.setFormatter(formatter) logger.addHandler(hdlr) logger.info('Configurations: %s', args) ssgd(args.dnn, args.dataset, args.data_dir, args.nworkers, args.lr, args.batch_size, args.nsteps_update, args.max_epochs, args.nwpernode, args.pretrain, args.num_steps, args.compressor, args.density, args.threshold, gradient_relative_path)
def robust_ssgd(dnn, dataset, data_dir, nworkers, lr, batch_size, nsteps_update, max_epochs, compression=False, compressor='topk', nwpernode=1, sigma_scale=2.5, pretrain=None, density=0.01, prefix=None): global relative_path torch.cuda.set_device(dopt.rank()%nwpernode) rank = dopt.rank() if rank != 0: pretrain = None trainer = DLTrainer(rank, nworkers, dist=False, batch_size=batch_size, is_weak_scaling=True, ngpus=1, data_dir=data_dir, dataset=dataset, dnn=dnn, lr=lr, nworkers=nworkers, prefix=prefix+'-ds%s'%str(density), pretrain=pretrain, tb_writer=writer) init_epoch = trainer.get_train_epoch() init_iter = trainer.get_train_iter() trainer.set_train_epoch(comm.bcast(init_epoch)) trainer.set_train_iter(comm.bcast(init_iter)) def _error_handler(new_num_workers, new_rank): logger.info('Error info catched by trainer') trainer.update_nworker(new_num_workers, new_rank) compressor = compressor if compression else 'none' compressor = compressors[compressor] is_sparse = compression logger.info('Broadcast parameters....') hvd.broadcast_parameters(trainer.net.state_dict(), root_rank=0) logger.info('Broadcast parameters finished....') norm_clip = None optimizer = dopt.DistributedOptimizer(trainer.optimizer, trainer.net.named_parameters(), compression=compressor, is_sparse=is_sparse, err_handler=_error_handler, layerwise_times=None, sigma_scale=sigma_scale, density=density, norm_clip=norm_clip, writer=writer) trainer.update_optimizer(optimizer) iters_per_epoch = trainer.get_num_of_training_samples() // (nworkers * batch_size * nsteps_update) times = [] NUM_OF_DISLAY = 40 display = NUM_OF_DISLAY if iters_per_epoch > NUM_OF_DISLAY else iters_per_epoch-1 logger.info('Start training ....') for epoch in range(max_epochs): hidden = None if dnn == 'lstm': hidden = trainer.net.init_hidden() for i in range(iters_per_epoch): s = time.time() optimizer.zero_grad() for j in range(nsteps_update): if j < nsteps_update - 1 and nsteps_update > 1: optimizer.local = True else: optimizer.local = False if dnn == 'lstm': _, hidden = trainer.train(1, hidden=hidden) else: trainer.train(1) if dnn == 'lstm': optimizer.synchronize() torch.nn.utils.clip_grad_norm_(trainer.net.parameters(), 0.25) elif dnn == 'lstman4': optimizer.synchronize() torch.nn.utils.clip_grad_norm_(trainer.net.parameters(), 400) trainer.update_model() times.append(time.time()-s) if i % display == 0 and i > 0: time_per_iter = np.mean(times) logger.info('Time per iteration including communication: %f. Speed: %f images/s, current density: %f', time_per_iter, batch_size * nsteps_update / time_per_iter, optimizer.get_current_density()) times = [] optimizer.add_train_epoch() if settings.PROFILING_INDEX and rank == 0 and epoch % 10 == 0: fn = os.path.join(relative_path, 'index-rank%d-epoch%d.npy' % (rank, epoch)) key = list(optimizer._allreducer._update_index_counter.keys())[0] np.save(fn, optimizer._allreducer._update_index_counter[key].int().cpu().numpy()) if settings.PROFILING_NORM: # For comparison purpose ===> fn = os.path.join(relative_path, 'gtopknorm-rank%d-epoch%d.npy' % (rank, epoch)) fn2 = os.path.join(relative_path, 'randknorm-rank%d-epoch%d.npy' % (rank, epoch)) fn3 = os.path.join(relative_path, 'upbound-rank%d-epoch%d.npy' % (rank, epoch)) fn5 = os.path.join(relative_path, 'densestd-rank%d-epoch%d.npy' % (rank, epoch)) arr = [] arr2 = [] arr3 = [] arr4 = [] arr5 = [] for gtopk_norm, randk_norm, upbound, xnorm, dense_std in optimizer._allreducer._profiling_norms: arr.append(gtopk_norm) arr2.append(randk_norm) arr3.append(upbound) arr4.append(xnorm) arr5.append(dense_std) arr = np.array(arr) arr2 = np.array(arr2) arr3 = np.array(arr3) arr4 = np.array(arr4) arr5 = np.array(arr5) logger.info('[rank:%d][%d] gtopk norm mean: %f, std: %f', rank, epoch, np.mean(arr), np.std(arr)) logger.info('[rank:%d][%d] randk norm mean: %f, std: %f', rank, epoch, np.mean(arr2), np.std(arr2)) logger.info('[rank:%d][%d] upbound norm mean: %f, std: %f', rank, epoch, np.mean(arr3), np.std(arr3)) logger.info('[rank:%d][%d] x norm mean: %f, std: %f', rank, epoch, np.mean(arr4), np.std(arr4)) logger.info('[rank:%d][%d] dense std mean: %f, std: %f', rank, epoch, np.mean(arr5), np.std(arr5)) np.save(fn, arr) np.save(fn2, arr2) np.save(fn3, arr3) np.save(fn5, arr5) # For comparison purpose <=== End optimizer._allreducer._profiling_norms = [] optimizer.stop()
parser.add_argument('--dnn', type=str, default='resnet50', choices=_support_dnns, help='Specify the neural network for training') parser.add_argument('--data-dir', type=str, default='./data', help='Specify the data root path') parser.add_argument('--lr', type=float, default=0.1, help='Default learning rate') parser.add_argument('--max-epochs', type=int, default=90, help='Default maximum epochs to train') parser.add_argument('--pretrain', type=str, default=None, help='Specify the pretrain path') parser.set_defaults(compression=False) args = parser.parse_args() batch_size = args.batch_size * args.nsteps_update prefix = settings.PREFIX if args.compression: prefix = 'comp-' + args.compressor + '-' + prefix logdir = 'allreduce-%s/%s-n%d-bs%d-lr%.4f-ns%d-sg%.2f-ds%s' % (prefix, args.dnn, args.nworkers, batch_size, args.lr, args.nsteps_update, args.sigma_scale, str(args.density)) relative_path = './logs/%s'%logdir utils.create_path(relative_path) rank = 0 rank = dopt.rank() hvd.init() if rank == 0: tb_runs = './runs/%s'%logdir writer = SummaryWriter(tb_runs) logfile = os.path.join(relative_path, settings.hostname+'-'+str(rank)+'.log') hdlr = logging.FileHandler(logfile) hdlr.setFormatter(formatter) logger.addHandler(hdlr) logger.info('Configurations: %s', args) logger.info('Interpreter: %s', sys.version) try: robust_ssgd(args.dnn, args.dataset, args.data_dir, args.nworkers, args.lr, args.batch_size, args.nsteps_update, args.max_epochs, args.compression, args.compressor, args.nwpernode, args.sigma_scale, args.pretrain, args.density, prefix) except Exception as e: logger.info('Main thread exception: %s', e)
def ssgd(dnn, dataset, data_dir, nworkers, lr, batch_size, nsteps_update, max_epochs, nwpernode, pretrain, num_steps, compressor, density, threshold, gradient_path=None, tb=None, iratio=0.1, stages=1, partitions=0, ec_gradw=1.0, ec_memw=0.0, optimizer='nesterov', totaltime=0): global SPEED if not settings.USE_CPU: if nworkers > 1: rank = hvd.rank() torch.cuda.set_device(hvd.local_rank()) #%rank%nwpernode) else: rank = 0 torch.cuda.set_device(rank) if rank != 0: pretrain = None #### CHECK whether to use GPU or CPU if settings.USE_CPU: trainer = DLTrainer(rank, nworkers, dist=False, batch_size=batch_size, is_weak_scaling=True, ngpus=0, data_dir=data_dir, dataset=dataset, dnn=dnn, lr=lr, nworkers=nworkers, prefix='allreduce', pretrain=pretrain, num_steps=num_steps, tb_writer=writer, tb=tb, optimizer_str=optimizer) else: trainer = DLTrainer(rank, nworkers, dist=False, batch_size=batch_size, is_weak_scaling=True, ngpus=1, data_dir=data_dir, dataset=dataset, dnn=dnn, lr=lr, nworkers=nworkers, prefix='allreduce', pretrain=pretrain, num_steps=num_steps, tb_writer=writer, tb=tb, optimizer_str=optimizer) init_epoch = torch.ones(1) * trainer.get_train_epoch() init_iter = torch.ones(1) * trainer.get_train_iter() trainer.set_train_epoch(int(hvd.broadcast(init_epoch, root_rank=0)[0])) trainer.set_train_iter(int(hvd.broadcast(init_iter, root_rank=0)[0])) is_sparse = True #density < 1 if not is_sparse: compressor = None if settings.ADAPTIVE_MERGE or settings.ADAPTIVE_SPARSE: seq_layernames, layerwise_times, layerwise_sizes = benchmark(trainer) layerwise_times = comm.bcast(layerwise_times, root=0) if rank == 0: logger.info('layerwise backward times: %s', list(layerwise_times)) logger.info('layerwise backward sizes: %s', list(layerwise_sizes)) logger.info('Bencharmked backward time: %f', np.sum(layerwise_times)) logger.info('Model size: %d', np.sum(layerwise_sizes)) else: seq_layernames, layerwise_times, layerwise_sizes = None, None, None norm_clip = None if dnn == 'lstm': norm_clip = 0.25 elif dnn == 'lstman4': norm_clip = 400 optimizer = hvd.DistributedOptimizer( trainer.optimizer, named_parameters=trainer.net.named_parameters(), compression=compressors[compressor], is_sparse=is_sparse, density=density, seq_layernames=seq_layernames, layerwise_times=layerwise_times, norm_clip=norm_clip, threshold=threshold, writer=writer, gradient_path=gradient_path, tb=tb, iratio=iratio, stages=stages, partitions=partitions, ec_gradw=ec_gradw, ec_memw=ec_memw) hvd.SPEED hvd.broadcast_parameters(trainer.net.state_dict(), root_rank=0) trainer.update_optimizer(optimizer) iters_per_epoch = trainer.get_num_of_training_samples() // ( nworkers * batch_size * nsteps_update) start = time.time() times = [] noupdate_times = [] logger.info('max_epochs: %d', max_epochs) display = settings.DISPLAY if iters_per_epoch > settings.DISPLAY else iters_per_epoch - 1 for epoch in range(max_epochs): hidden = None if dnn == 'lstm': hidden = trainer.net.init_hidden() for i in range(iters_per_epoch): s = time.time() optimizer.zero_grad() for j in range(nsteps_update): if j < nsteps_update - 1 and nsteps_update > 1: optimizer.local = True else: optimizer.local = False if dnn == 'lstm': _, hidden = trainer.train(1, hidden=hidden) else: trainer.train(1) if dnn == 'lstm': optimizer.synchronize() torch.nn.utils.clip_grad_norm_(trainer.net.parameters(), 0.25) elif dnn == 'lstman4': optimizer.synchronize() torch.nn.utils.clip_grad_norm_(trainer.net.parameters(), 400) noupdate_times.append(time.time() - s) trainer.update_model() torch.cuda.synchronize() times.append(time.time() - s) if i % display == 0 and i > 0: time_per_iter = np.mean(times) update_per_iter = time_per_iter - np.mean(noupdate_times) throughput = batch_size * nsteps_update / time_per_iter trainer.log_info(time_per_iter, throughput, update_per_iter) logger.warning( 'Time per iteration: %f, communication: %f, Speed: %f images/s', time_per_iter, update_per_iter, throughput) times = [] noupdate_times = [] optimizer.increase_one_epoch() if totaltime > 0 and time.time() - start > totaltime: trainer.test(trainer.get_train_epoch()) break if not (dataset == 'cifar10'): trainer.test(trainer.get_train_epoch())
prefix = 'comp-' + args.compressor + '-' + prefix if momentum_correction: prefix = 'mc-' + prefix logdir = 'allreduce-%s-thres-%dkbytes/%s-n%d-bs%d-lr%.4f-ns%d-ds%s' % ( prefix, args.threshold / 1024, args.dnn, args.nworkers, batch_size, args.lr, args.nsteps_update, str(args.density)) relative_path = './logs/%s' % logdir gradient_relative_path = None utils.create_path(relative_path) if settings.LOGGING_GRADIENTS: gradient_relative_path = '%s/gradients/%s' % (args.saved_dir, logdir) utils.create_path(gradient_relative_path) rank = 0 if args.nworkers > 1: hvd.init() rank = hvd.rank() if rank == 0: tb_runs = './runs/%s' % logdir writer = None #SummaryWriter(tb_runs) logfile = os.path.join(relative_path, settings.hostname + '-' + str(rank) + '.log') hdlr = logging.FileHandler(logfile) hdlr.setFormatter(formatter) logger.addHandler(hdlr) logger.info('Configurations: %s', args) if hvd.rank() == 0: wandb.init(project='gtopk', entity='shyhuai', name=logfile, config=args)