from mpi4py import MPI comm = MPI.COMM_WORLD rank = comm.rank size = comm.size # device='gpu'+str(rank) from test_exchanger import init_device, clean_device drv, ctx, arr, shared_x, shared_xx = init_device(device=device) if rank == 0: print 'original array %s' % arr # prepare copper exchanger from exchanger_strategy import Exch_copper exch = Exch_copper(comm, avg=False) exch.prepare(ctx, drv, [shared_x]) exch.exchange() if rank == 0: print 'copper summation: %s' % shared_x.get_value() # prepare ar exchanger from exchanger_strategy import Exch_allreduce exch = Exch_allreduce(comm, avg=False) exch.prepare([shared_xx]) exch.exchange()
def __init__(self, config, drv, ctx, model):
self.drv = drv
self.ctx = ctx
self.comm = config['comm']
self.size = config['size']
self.exch_strategy = config['exch_strategy']
self.train_mode = config['train_mode']
# self.cuda_aware = config['cuda_aware']
# TODO make sure exchanger class doesn't keep a self copy of model, only the reference to its param list
self.param_list = model.params
self.vels = model.vels
self.vels2 = model.vels2
self.avg_func_list = []
if self.train_mode == 'cdd' and self.exch_strategy == 'ar': #c
from exchanger_strategy import Exch_allreduce
self.exch = Exch_allreduce(self.comm, avg=False)
self.exch.prepare(self.vels, self.vels2)
elif self.train_mode == 'cdd' and self.exch_strategy == 'copper': #c
from exchanger_strategy import Exch_copper
self.exch = Exch_copper(self.comm, avg=False)
self.exch.prepare(self.ctx, self.drv, self.vels, self.vels2)
elif self.train_mode == 'cdd' and self.exch_strategy == 'asa32': #c
from exchanger_strategy import Exch_asa32
self.exch = Exch_asa32(self.comm, avg=False)
self.exch.prepare(self.ctx, self.drv, self.vels, self.vels2)
elif self.train_mode == 'cdd' and self.exch_strategy == 'asa16': #c
from exchanger_strategy import Exch_asa16
self.exch = Exch_asa16(self.comm, avg=False)
self.exch.prepare(self.ctx, self.drv, self.vels, self.vels2)
#TODO adjust ctx, drv locations in all strategies
elif self.train_mode == 'avg' and self.exch_strategy == 'ar': #c
from exchanger_strategy import Exch_allreduce
self.exch = Exch_allreduce(self.comm)
self.exch.prepare(self.param_list)
elif self.train_mode == 'avg' and self.exch_strategy == 'copper': #c
from exchanger_strategy import Exch_copper
self.exch = Exch_copper(self.comm)
self.exch.prepare(self.ctx, self.drv, self.param_list)
elif self.train_mode == 'avg' and self.exch_strategy == 'asa32': #c
from exchanger_strategy import Exch_asa32
self.exch = Exch_asa32(self.comm)
self.exch.prepare(self.ctx, self.drv, self.param_list)
elif self.train_mode == 'avg' and self.exch_strategy == 'asa16': #c
from exchanger_strategy import Exch_asa16
self.exch = Exch_asa16(self.comm)
self.exch.prepare(self.ctx, self.drv, self.param_list)