class KFAC(optim.Optimizer): """KFAC Distributed Gradient Preconditioner Computes the natural gradient of a model in place with a layer-wise FIM approximation. Layer computations are distributed across workers using Horovod. Usage: optimizer = optim.SGD(model.parameters(), ...) optimizer = hvd.DistributedOptimizer(optimizer, ...) preconditioner = KFAC(model, ...) ... for i, (data, target) in enumerate(train_loader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.synchronize() preconditioner.step() with optimizer.skip_synchronize(): optimizer.step() Args: model (nn): Torch model to precondition lr (float, optional): learning rate (default: 0.1) factor_decay (float, optional): running average coefficient for Kronecker factors (default: 0.95) damping (float, optional): Tikhonov damping parameter (default: 0.001) kl_clip (float, optional): clipping parameter for gradient scaling (default: 0.001) fac_update_freq (int, optional): iterations between calculating and updating the running average of the Kronecker factors (default: 10) kfac_update_freq (int, optional): iterations between applying gradient preconditioning (default: 100) batch_averaged (bool, optional): boolean representing if the gradient is alrady averaged across the batches (default: True) diag_blocks (int, optional): Experimental: number of diagonal blocks to approximate the Kronecker factor eigendecomposition with. `diag_blocks=1` computes the eigendecomposition of the entire factor (default: 1) diag_warmup (int, optional): number of epochs to wait before starting the block diagonal factor approximation (default: 0) distribute_layer_factors (bool, optional): if `True`, computes factors A and G on different workers else computes A and G for a single layer on the same worker. If `None`, determines best value based on layer count (default: None) """ def __init__(self, model, lr=0.1, factor_decay=0.95, damping=0.001, kl_clip=0.001, fac_update_freq=10, kfac_update_freq=100, batch_averaged=True, diag_blocks=1, diag_warmup=0, distribute_layer_factors=None, sparse=False, sparse_ratio=0.01, exclude_parts=''): if not 0.0 <= lr: raise ValueError("Invalid learning rate: {}".format(lr)) if not 0.0 < factor_decay <= 1: raise ValueError( "Invalid factor decay rate: {}".format(factor_decay)) if not 0.0 < damping: raise ValueError("Invalid damping: {}".format(damping)) if not 0.0 < kl_clip: raise ValueError("Invalid clipping value: {}".format(kl_clip)) if not 0 < fac_update_freq: raise ValueError( "Invalid factor update frequency: {}".format(fac_update_freq)) if not 0 < kfac_update_freq: raise ValueError( "Invalid K-FAC update frequency: {}".format(kfac_update_freq)) if not 0 == kfac_update_freq % fac_update_freq: print( "WARNING: it is suggested that kfac_update_freq be a multiple of fac_update_freq" ) if not 0 < diag_blocks: raise ValueError( "Invalid diagonal block approx count: {}".format(diag_blocks)) if not 0 <= diag_blocks: raise ValueError( "Invalid diagonal block approx count: {}".format(diag_blocks)) if not 1 == diag_blocks: print( "WARNING: diag_blocks > 1 is experimental and may give poor results." ) # For compatibility with `KFACParamScheduler` defaults = dict(lr=lr, damping=damping, fac_update_freq=fac_update_freq, kfac_update_freq=kfac_update_freq) super(KFAC, self).__init__(model.parameters(), defaults) self.computeA = ComputeA() self.computeG = ComputeG() self.known_modules = {'Linear', 'Conv2d'} self.modules = [] self.module_names = [] self.name_module_map = {} self.module_name_map = {} #self.fw_factor_handles = [] #self.bw_factor_handles = [] self._register_modules(model) self.fw_merged_comm = MergedCommAllReduce(self.module_names, prefix='forward', merge=True, single_layer=False) self.bw_merged_comm = MergedCommAllReduce(self.module_names, prefix='backward', merge=True, single_layer=False) self.steps = 0 # Dictionaries keyed by `module` to storing the factors and # eigendecompositions self.m_a, self.m_g = {}, {} self.m_A, self.m_G = {}, {} self.m_QA, self.m_QG = {}, {} self.m_dA, self.m_dG = {}, {} self.m_dA_ranks = {} self.m_dG_ranks = {} self.module_ranks = None self.sparse = sparse self.sparse_ratio = sparse_ratio self.residualsA, self.residualsG = {}, {} self.factor_decay = factor_decay self.kl_clip = kl_clip self.fac_update_freq = fac_update_freq self.kfac_update_freq = kfac_update_freq self.diag_blocks = diag_blocks self.diag_warmup = diag_warmup self.batch_averaged = batch_averaged # Compute ideal value for `distribute_layer_factors` based on # registered module count if distribute_layer_factors is None: self.distribute_layer_factors = True \ if hvd.size() > len(self.modules) else False else: self.distribute_layer_factors = distribute_layer_factors self.have_cleared_Q = True if self.diag_warmup == 0 else False self.eps = 1e-10 # for numerical stability self.rank_iter = cycle(list(range(hvd.size()))) def _save_input(self, module, input): """Hook for saving layer input""" if torch.is_grad_enabled() and self.steps % self.fac_update_freq == 0: self.m_a[module] = input[0].data def _compute_forward_factor(self, module, input): if torch.is_grad_enabled() and self.steps % self.fac_update_freq == 0: self.m_a[module] = input[0].data self._update_module_A(module) if hvd.size() > 1: name = self.module_name_map[module] self.fw_merged_comm.allreduce_async_(name, self.m_A[module].data) def _save_grad_output(self, module, grad_input, grad_output): """Hook for saving gradient w.r.t output""" if self.steps % self.fac_update_freq == 0: self.m_g[module] = grad_output[0].data def _compute_backward_factor(self, module, grad_input, grad_output): if self.steps % self.fac_update_freq == 0: self.m_g[module] = grad_output[0].data self._update_module_G(module) if hvd.size() > 1: name = self.module_name_map[module] self.bw_merged_comm.allreduce_async_(name, self.m_G[module].data) def _register_modules(self, model): """Register hooks to all supported layers in the model""" name_idx = 0 for module in model.modules(): classname = module.__class__.__name__ if classname in self.known_modules: self.modules.append(module) module.register_forward_pre_hook(self._compute_forward_factor) module.register_backward_hook(self._compute_backward_factor) module_name = 'module_name_%s_%d' % (classname, name_idx) self.module_names.append(module_name) self.name_module_map[module_name] = module self.module_name_map[module] = module_name name_idx += 1 def _init_A(self, factor, module): """Initialize memory for factor A and its eigendecomp""" self.m_A[module] = torch.diag(factor.new(factor.shape[0]).fill_(1)) self.m_dA[module] = factor.new_zeros(factor.shape[0]) self.m_QA[module] = factor.new_zeros(factor.shape) def _init_G(self, factor, module): """Initialize memory for factor G and its eigendecomp""" self.m_G[module] = torch.diag(factor.new(factor.shape[0]).fill_(1)) self.m_dG[module] = factor.new_zeros(factor.shape[0]) self.m_QG[module] = factor.new_zeros(factor.shape) def _clear_eigen(self): """Clear eigendecompositions Useful for when switching between `diag_blocks=1` and `diag-blocks>1` because eigendecompositions saved in place and the off-diagonals must be cleared. """ for module in self.modules: self.m_QA[module].fill_(0) self.m_QG[module].fill_(0) self.m_dA[module].fill_(0) self.m_dG[module].fill_(0) def _update_module_A(self, module): a = self.computeA(self.m_a[module], module) if self.steps == 0: self._init_A(a, module) update_running_avg(a, self.m_A[module], self.factor_decay) if self.sparse: sparsification(self.m_A[module], module, ratio=self.sparse_ratio, residuals=self.residualsA) def _update_A(self): """Compute and update factor A for all modules""" for module in self.modules: self._update_module_A(module) #if hvd.rank() == 0: # data = self.m_A[module] #ComputeA.get_data(self.m_a[module], module) # d = data.view(-1) # indexes = d.nonzero().data.squeeze().view(-1) # nnz = indexes.numel() # numel = d.numel() # sparsity = (numel-nnz)*1.0/numel # logger.info('vector A name: %s, shape: %s, sparsity: %f', module, data.shape, sparsity) def _update_module_G(self, module): g = self.computeG(self.m_g[module], module, self.batch_averaged) #logger.info('G Name: %s, shape: %s', module, g.shape) if self.steps == 0: self._init_G(g, module) update_running_avg(g, self.m_G[module], self.factor_decay) if self.sparse: sparsification(self.m_G[module], module, ratio=self.sparse_ratio, residuals=self.residualsG) def _update_G(self): """Compute and update factor G for all modules""" for module in self.modules: self._update_eigen_G(module) #if hvd.rank() == 0: # data = self.m_G[module] #ComputeG.get_data(self.m_g[module], module, self.batch_averaged) # d = data.view(-1) # indexes = d.nonzero().data.squeeze().view(-1) # nnz = indexes.numel() # numel = d.numel() # sparsity = (numel-nnz)*1.0/numel # logger.info('vector G name: %s, shape: %s, sparsity: %f', module, data.shape, sparsity) def _update_eigen_A(self, module, ranks): """Compute eigendecomposition of A for module on specified workers Note: all ranks will enter this function but only the ranks specified in `ranks` will continue to actually compute the eigendecomposition. All other ranks will simply zero out their buffer for the eigendecomposition for the current module. This is done so we can sum the eigendecompositions across all ranks to communicate the results of locally computed eigendecompositions. Args: module: module to compute eigendecomposition of A on ranks: list of horovod ranks (i.e. workers) to use when computing the eigendecomposition. """ if hvd.rank() in ranks: #if self.m_A[module].shape[1] >= 2304: # return self._distributed_compute_eigen(self.m_A[module], self.m_QA[module], self.m_dA[module], ranks) else: self.m_QA[module].fill_(0) self.m_dA[module].fill_(0) def _update_eigen_G(self, module, ranks): """Compute eigendecomposition of A for module on specified workers See `_update_eigen_A` for more info` """ if hvd.rank() in ranks: #if self.m_G[module].shape[1] >= 2304: # return self._distributed_compute_eigen(self.m_G[module], self.m_QG[module], self.m_dG[module], ranks) else: self.m_QG[module].fill_(0) self.m_dG[module].fill_(0) def _distributed_compute_eigen(self, factor, evectors, evalues, ranks): """Computes the eigendecomposition of a factor across ranks Assigns each rank in `ranks` to enter this function to compute a diagonal block of `factor`. Results are written to `evectors` and `evalues`. If `len(ranks)==1`, then that rank computes the eigendecomposition of the entire `factor`. Args: factor (tensor): tensor to eigendecompose evectors (tensor): tensor to save eigenvectors of `factor` to evalues (tensor): tensor to save eigenvalues of `factor` to ranks (list): list of ranks that will enter this function """ i = ranks.index(hvd.rank()) n = len(ranks) if n > min(factor.shape): n = min(factor.shape) if i < n: start, end = get_block_boundary(i, n, factor.shape) block = factor[start[0]:end[0], start[1]:end[1]] #d, Q = torch.symeig(block, eigenvectors=True) d, Q = tcmm.f_symeig(block) Q = Q.transpose(-2, -1) d = torch.mul(d, (d > self.eps).float()) evalues.data[start[0]:end[0]].copy_(d) evectors.data[start[0]:end[0], start[1]:end[1]].copy_(Q) def _get_diag_blocks(self, module, diag_blocks): """Helper method for determining number of diag_blocks to use Overrides `diag_blocks` if the `module` does not support `diag_blocks>1`. I.e. for a Linear layer, we do not want to use a `diag_blocks>1`. Args: module: module diag_blocks (int): default number of diag blocks to use """ return diag_blocks if module.__class__.__name__ == 'Conv2d' else 1 def _get_grad(self, module): """Get formated gradient of module Args: module: module/layer to get gradient of Returns: Formatted gradient with shape [output_dim, input_dim] for module """ if module.__class__.__name__ == 'Conv2d': # n_filters * (in_c * kw * kh) grad = module.weight.grad.data.view( module.weight.grad.data.size(0), -1) else: grad = module.weight.grad.data if module.bias is not None: grad = torch.cat([grad, module.bias.grad.data.view(-1, 1)], 1) return grad def _get_preconditioned_grad(self, module, grad): """Precondition gradient of module Args: module: module to compute preconditioned gradient for grad: formatted gradient from `_get_grad()` Returns: preconditioned gradient with same shape as `grad` """ v1 = self.m_QG[module].t() @ grad @ self.m_QA[module] v2 = v1 / (self.m_dG[module].unsqueeze(1) * self.m_dA[module].unsqueeze(0) + self.damping) v = self.m_QG[module] @ v2 @ self.m_QA[module].t() if module.bias is not None: v = [v[:, :-1], v[:, -1:]] v[0] = v[0].view(module.weight.grad.data.size()) # weight v[1] = v[1].view(module.bias.grad.data.size()) # bias else: v = [v.view(module.weight.grad.data.size())] return v def _update_scale_grad(self, updates): """Update the gradients in place and scale Updates the gradients in-place for all modules using the preconditioned gradients and scales the gradients. Args: updates (dict): dict of {module: precon_grad} """ vg_sum = 0 for module in self.modules: v = updates[module] vg_sum += (v[0] * module.weight.grad.data * self.lr**2).sum().item() if module.bias is not None: vg_sum += (v[1] * module.bias.grad.data * self.lr**2).sum().item() nu = min(1.0, math.sqrt(self.kl_clip / abs(vg_sum))) for module in self.modules: v = updates[module] module.weight.grad.data.copy_(v[0]) module.weight.grad.data.mul_(nu) if module.bias is not None: module.bias.grad.data.copy_(v[1]) module.bias.grad.data.mul_(nu) def step(self, closure=None, epoch=None): """Perform one K-FAC step Note: - this function should always be called before `optimizer.step()` - gradients must be averaged across ranks before calling `step()` Args: closure: for compatibility with the base optimizer class. `closure` is ignored by KFAC epoch (int, optional): epoch to use for determining when to end the `diag_warmup` period. `epoch` is not necessary if not using `diag_warmup` """ # Update params, used for compatibilty with `KFACParamScheduler` group = self.param_groups[0] self.lr = group['lr'] self.damping = group['damping'] self.fac_update_freq = group['fac_update_freq'] self.kfac_update_freq = group['kfac_update_freq'] updates = {} handles = [] if epoch is None: if self.diag_warmup > 0: print("WARNING: diag_warmup > 0 but epoch was not passed to " "KFAC.step(). Defaulting to no diag_warmup") diag_blocks = self.diag_blocks else: diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 if hvd.size() > 1 and self.steps % self.fac_update_freq == 0: self.fw_merged_comm.synchronize() self.bw_merged_comm.synchronize() #for handle in self.fw_factor_handles: # hvd.synchronize(handle) #self.fw_factor_handles.clear() #for handle in self.bw_factor_handles: # hvd.synchronize(handle) #self.bw_factor_handles.clear() # if we are switching from no diag approx to approx, we need to clear # off-block-diagonal elements if not self.have_cleared_Q and \ epoch == self.diag_warmup and \ self.steps % self.kfac_update_freq == 0: self._clear_eigen() self.have_cleared_Q = True if self.steps % self.kfac_update_freq == 0: # reset rank iter so device get the same layers # to compute to take advantage of caching self.rank_iter.reset() handles = [] #eigen_ranks = self._generate_eigen_ranks(epoch) eigen_ranks = self._generate_eigen_ranks_uniform(epoch) #eigen_ranks = self._generate_eigen_ranks_naive(epoch) for module in self.modules: ranks_a, ranks_g = eigen_ranks[module] self.m_dA_ranks[module] = ranks_a[0] self.m_dG_ranks[module] = ranks_g[0] rank_a = ranks_a[0] rank_g = ranks_g[0] self._update_eigen_A(module, ranks_a) h1 = hvd.broadcast_async_(self.m_QA[module], rank_a) h2 = hvd.broadcast_async_(self.m_dA[module], rank_a) self._update_eigen_G(module, ranks_g) h3 = hvd.broadcast_async_(self.m_QG[module], rank_g) h4 = hvd.broadcast_async_(self.m_dG[module], rank_g) handles.append((h1, h2, h3, h4)) if hvd.size() > 1: #for handle in handles: # hvd.synchronize(handle) #self._allreduce_eigendecomp() #self._broadcast_eigendecomp() pass for i, module in enumerate(self.modules): if hvd.size() > 1 and len(handles) > 0: h1, h2, h3, h4 = handles[i] hvd.synchronize(h1) hvd.synchronize(h2) hvd.synchronize(h3) hvd.synchronize(h4) grad = self._get_grad(module) precon_grad = self._get_preconditioned_grad(module, grad) updates[module] = precon_grad #self._update_scale_grad(updates) self.steps += 1 def _generate_eigen_ranks_naive(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 buckets = [0] * hvd.size() for module in self.modules: # Get ranks to compute this layer on n = self._get_diag_blocks(module, diag_blocks) ranks_a = self.rank_iter.next(n) ranks_g = self.rank_iter.next(n) if self.distribute_layer_factors \ else ranks_a module_ranks[module] = (ranks_a, ranks_g) buckets[ranks_a[0]] += self.m_A[module].shape[1] buckets[ranks_g[0]] += self.m_G[module].shape[1] self.module_ranks = module_ranks if hvd.rank() == 0: logger.info('buckets: %s', buckets) logger.info('module_ranks: %s', module_ranks.values()) return module_ranks def _generate_eigen_ranks_uniform(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 buckets = [0] * hvd.size() dimensions = [] module_factors = [] for i, m in enumerate(self.modules): name = self.module_names[i] a_dimension = self.m_A[m].shape[1] g_dimension = self.m_G[m].shape[1] dimensions.append(a_dimension) module_factors.append(name + '-A') dimensions.append(g_dimension) module_factors.append(name + '-G') descending_sorted_idx = np.argsort(dimensions)[::-1] A_ranks = {} G_ranks = {} for i in descending_sorted_idx: factor = module_factors[i] dimension = dimensions[i] m_i = self.module_names.index(factor[0:-2]) m = self.modules[m_i] bi = np.argmin(buckets) buckets[bi] += dimension if factor[-1] == 'A': A_ranks[m] = (bi, ) else: G_ranks[m] = (bi, ) for m in self.modules: module_ranks[m] = (A_ranks[m], G_ranks[m]) self.module_ranks = module_ranks if hvd.rank() == 0: logger.info('buckets: %s', buckets) logger.info('module_ranks: %s', module_ranks.values()) return module_ranks def _generate_eigen_ranks(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 buckets = [0] * hvd.size() for module in self.modules: i = np.argmin(buckets) if hvd.rank() == 0: logger.info('A Name: %s, shape: %s', module, self.m_A[module].shape) logger.info('G Name: %s, shape: %s', module, self.m_G[module].shape) a_dimension = self.m_A[module].shape[1] g_dimension = self.m_G[module].shape[1] #buckets[i] += (a_dimension) + g_dimension) buckets[i] += a_dimension ranks_a = (i, ) i = np.argmin(buckets) ranks_g = (i, ) buckets[i] += g_dimension module_ranks[module] = (ranks_a, ranks_g) self.module_ranks = module_ranks if hvd.rank() == 0: logger.info('buckets: %s', buckets) logger.info('module_ranks: %s', module_ranks.values()) return module_ranks def _allreduce_factors(self): """Allreduce the factors for all layers""" handles = [] for m in self.modules: handles.append( hvd.allreduce_async_(self.m_A[m].data, op=hvd.Average)) handles.append( hvd.allreduce_async_(self.m_G[m].data, op=hvd.Average)) for handle in handles: hvd.synchronize(handle) def _allgather_factors(self): """Allgather the factors for all layers""" handles = [] def _get_value_and_idx(sparse_tensor): tensor = sparse_tensor.data.view(-1) one_indexes = tensor != 0 indexes = one_indexes.nonzero().data.squeeze().view(-1) values = tensor.data[indexes] return values, indexes.int() for i, m in enumerate(self.modules): module_name = self.module_names[i] A_values, A_indexes = _get_value_and_idx(self.m_A[m].data) A_value_name = module_name + '_A_value' A_idx_name = module_name + '_A_idx' h_value = allgather_async(A_values, A_value_name) h_idx = allgather_async(A_indexes, A_idx_name) G_values, G_indexes = _get_value_and_idx(self.m_G[m].data) G_value_name = module_name + '_G_value' G_idx_name = module_name + '_G_idx' h_value_G = allgather_async(G_values, G_value_name) h_idx_G = allgather_async(G_indexes, G_idx_name) handles.append((h_value, h_idx, h_value_G, h_idx_G)) for i, handle in enumerate(handles): module_name = self.module_names[i] module = self.modules[i] m_A = self.m_A[module].view(-1) m_A.fill_(0.0) m_G = self.m_G[module].view(-1) m_G.fill_(0.0) h_value_A, h_idx_A, h_value_G, h_idx_G = handle A_values = hvd.synchronize(h_value_A) A_indexes = hvd.synchronize(h_idx_A).long() m_A.scatter_add_(0, A_indexes, A_values) m_A.div_(hvd.size()) G_values = hvd.synchronize(h_value_G) G_indexes = hvd.synchronize(h_idx_G).long() m_G.scatter_add_(0, G_indexes, G_values) m_G.div_(hvd.size()) def _allreduce_eigendecomp(self): """Allreduce the eigendecompositions for all layers Note: we use `op=hvd.Sum` to simulate an allgather`. Each rank will either compute the eigendecomposition for a factor or just return zeros so we sum instead of averaging. """ handles = [] for m in self.modules: handles.append(hvd.allreduce_async_(self.m_QA[m].data, op=hvd.Sum)) handles.append(hvd.allreduce_async_(self.m_QG[m].data, op=hvd.Sum)) handles.append(hvd.allreduce_async_(self.m_dA[m].data, op=hvd.Sum)) handles.append(hvd.allreduce_async_(self.m_dG[m].data, op=hvd.Sum)) for handle in handles: hvd.synchronize(handle) def _broadcast_eigendecomp(self): """Broadcasts the eigendecompositions for all layers Note: we use `op=hvd.Sum` to simulate an allgather`. Each rank will either compute the eigendecomposition for a factor or just return zeros so we sum instead of averaging. """ handles = [] rank = hvd.rank() for i, m in enumerate(self.modules): rank_a = self.m_dA_ranks[m] rank_g = self.m_dG_ranks[m] name = self.module_names[i] h = hvd.broadcast_async_(self.m_QA[m], rank_a, name=name + 'mQA') handles.append(h) h = hvd.broadcast_async_(self.m_dA[m], rank_a, name=name + 'mdA') handles.append(h) h = hvd.broadcast_async_(self.m_QG[m], rank_g, name=name + 'mQG') handles.append(h) h = hvd.broadcast_async_(self.m_dG[m], rank_g, name=name + 'mdG') handles.append(h) for handle in handles: hvd.synchronize(handle)
class KFAC(optim.Optimizer): """KFAC Distributed Gradient Preconditioner Computes the natural gradient of a model in place with a layer-wise FIM approximation. Layer computations are distributed across workers using Horovod. Usage: optimizer = optim.SGD(model.parameters(), ...) optimizer = hvd.DistributedOptimizer(optimizer, ...) preconditioner = KFAC(model, ...) ... for i, (data, target) in enumerate(train_loader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.synchronize() preconditioner.step() with optimizer.skip_synchronize(): optimizer.step() Args: model (nn): Torch model to precondition lr (float, optional): learning rate (default: 0.1) factor_decay (float, optional): running average coefficient for Kronecker factors (default: 0.95) damping (float, optional): Tikhonov damping parameter (default: 0.001) kl_clip (float, optional): clipping parameter for gradient scaling (default: 0.001) fac_update_freq (int, optional): iterations between calculating and updating the running average of the Kronecker factors (default: 10) kfac_update_freq (int, optional): iterations between applying gradient preconditioning (default: 100) batch_averaged (bool, optional): boolean representing if the gradient is alrady averaged across the batches (default: True) diag_blocks (int, optional): Experimental: number of diagonal blocks to approximate the Kronecker factor eigendecomposition with. `diag_blocks=1` computes the eigendecomposition of the entire factor (default: 1) diag_warmup (int, optional): number of epochs to wait before starting the block diagonal factor approximation (default: 0) distribute_layer_factors (bool, optional): if `True`, computes factors A and G on different workers else computes A and G for a single layer on the same worker. If `None`, determines best value based on layer count (default: None) """ def __init__(self, model, lr=0.1, factor_decay=0.95, damping=0.001, kl_clip=0.001, fac_update_freq=10, kfac_update_freq=100, batch_averaged=True, diag_blocks=1, diag_warmup=0, distribute_layer_factors=None, sparse=False, sparse_ratio=0.01, exclude_parts=''): #exclude_parts='CommunicateInverse,ComputeInverse,CommunicateFactor,ComputeFactor'): if not 0.0 <= lr: raise ValueError("Invalid learning rate: {}".format(lr)) if not 0.0 < factor_decay <= 1: raise ValueError( "Invalid factor decay rate: {}".format(factor_decay)) if not 0.0 < damping: raise ValueError("Invalid damping: {}".format(damping)) if not 0.0 < kl_clip: raise ValueError("Invalid clipping value: {}".format(kl_clip)) if not 0 < fac_update_freq: raise ValueError( "Invalid factor update frequency: {}".format(fac_update_freq)) if not 0 < kfac_update_freq: raise ValueError( "Invalid K-FAC update frequency: {}".format(kfac_update_freq)) if not 0 == kfac_update_freq % fac_update_freq: print( "WARNING: it is suggested that kfac_update_freq be a multiple of fac_update_freq" ) if not 0 < diag_blocks: raise ValueError( "Invalid diagonal block approx count: {}".format(diag_blocks)) if not 0 <= diag_blocks: raise ValueError( "Invalid diagonal block approx count: {}".format(diag_blocks)) if not 1 == diag_blocks: print( "WARNING: diag_blocks > 1 is experimental and may give poor results." ) # For compatibility with `KFACParamScheduler` defaults = dict(lr=lr, damping=damping, fac_update_freq=fac_update_freq, kfac_update_freq=kfac_update_freq) super(KFAC, self).__init__(model.parameters(), defaults) self.computeA = ComputeA() self.computeG = ComputeG() self.known_modules = {'Linear', 'Conv2d'} self.modules = [] self.module_names = [] self.fw_factor_handles = [] self.bw_factor_handles = [] self.module_name_map = {} self.reduce_module_names = [] self._register_modules(model) self.steps = 0 merge = NUM_NEARBY_LAYERS > 1 self.fw_merged_comm = MergedCommReduce(tensor_names=None, prefix='forward', merge=True, single_layer=False, symmetric=True, fp16=False) self.bw_merged_comm = MergedCommReduce(tensor_names=None, prefix='backward', merge=False, single_layer=False, symmetric=True, fp16=False) self.fw_allreduce_comm = MergedCommAllReduce(self.module_names, prefix='forward', merge=False, single_layer=False, symmetric=True, fp16=False) self.bw_allreduce_comm = MergedCommAllReduce(self.module_names, prefix='backward', merge=False, single_layer=False, symmetric=True, fp16=False) self.multi_comm = MultiTensorComm(symmetric=True, fp16=False) self.fusion_groups_A = None self.fusion_groups_G = None self.dynamic_merge = True self.profiling = False #if self.profiling: self.fw_profiler = LayerwiseProfiler(self.module_names) self.bw_profiler = LayerwiseProfiler(self.module_names[::-1]) self.fw_factor_sizes = {} self.bw_factor_sizes = {} # Dictionaries keyed by `module` to storing the factors and # eigendecompositions self.m_a, self.m_g = {}, {} self.m_A, self.m_G = {}, {} self.m_QA, self.m_QG = {}, {} self.m_dA, self.m_dG = {}, {} self.m_dA_ranks = {} self.m_dG_ranks = {} self.module_ranks = None self.eigen_ranks = None self.sparse = sparse self.sparse_ratio = sparse_ratio self.residualsA, self.residualsG = {}, {} self.factor_decay = factor_decay self.kl_clip = kl_clip self.fac_update_freq = fac_update_freq self.kfac_update_freq = kfac_update_freq self.diag_blocks = diag_blocks self.diag_warmup = diag_warmup self.batch_averaged = batch_averaged self.exclude_communicate_inverse = True if exclude_parts.find( 'CommunicateInverse') >= 0 else False self.exclude_compute_inverse = True if exclude_parts.find( 'ComputeInverse') >= 0 else False self.exclude_communicate_factor = True if exclude_parts.find( 'CommunicateFactor') >= 0 else False self.exclude_compute_factor = True if exclude_parts.find( 'ComputeFactor') >= 0 else False # Compute ideal value for `distribute_layer_factors` based on # registered module count if distribute_layer_factors is None: self.distribute_layer_factors = True \ if hvd.size() > len(self.modules) else False else: self.distribute_layer_factors = distribute_layer_factors self.have_cleared_Q = True if self.diag_warmup == 0 else False self.eps = 1e-10 # for numerical stability self.rank_iter = cycle(list(range(hvd.size()))) def _compute_forward_factor(self, module, input): if torch.is_grad_enabled() and self.steps % self.fac_update_freq == 0: if self.profiling: name = self.module_name_map[module] if name == self.module_names[0]: self.fw_profiler.start() self.fw_profiler.layer_begin(name) if module in self.m_A: d = self.m_A[module].data.numel() self.fw_factor_sizes[name] = d self.m_a[module] = input[0].data if not self.exclude_compute_factor: self._update_module_A(module) if not self.exclude_communicate_factor: if hvd.size() > 1: name = self.module_name_map[module] if self.eigen_ranks is not None: ranks_a, ranks_g = self.eigen_ranks[module] rank_a = ranks_a[0] if self.m_A[module].shape[1] > REDUCE_THRESHOLD: self.fw_merged_comm.reduce_async_( name, self.m_A[module].data, rank_a) else: self.fw_allreduce_comm.allreduce_async_( name, self.m_A[module].data) def _compute_backward_factor(self, module, grad_input, grad_output): if self.steps % self.fac_update_freq == 0: if self.profiling: name = self.module_name_map[module] if name == self.module_names[-1]: self.bw_profiler.start() self.bw_profiler.layer_end(name) if module in self.m_G: d = self.m_G[module].data.numel() self.bw_factor_sizes[name] = d self.m_g[module] = grad_output[0].data if not self.exclude_compute_factor: self._update_module_G(module) if not self.exclude_communicate_factor: if hvd.size() > 1: name = self.module_name_map[module] if self.eigen_ranks is not None: ranks_a, ranks_g = self.eigen_ranks[module] rank_g = ranks_g[0] if self.m_A[module].shape[1] > REDUCE_THRESHOLD: self.bw_merged_comm.reduce_async_( name, self.m_G[module].data, rank_g) else: self.bw_allreduce_comm.allreduce_async_( name, self.m_G[module].data) #self.bw_allreduce_handlers.append(hvd.allreduce_async_(self.m_G[module].data, op=hvd.Average)) def _register_modules(self, model): """Register hooks to all supported layers in the model""" name_idx = 0 for module in model.modules(): classname = module.__class__.__name__ if classname in self.known_modules: self.modules.append(module) module.register_forward_pre_hook(self._compute_forward_factor) module.register_backward_hook(self._compute_backward_factor) module_name = 'module_name_%s_%d' % (classname, name_idx) self.module_names.append(module_name) self.module_name_map[module] = module_name name_idx += 1 def _init_A(self, factor, module): """Initialize memory for factor A and its eigendecomp""" self.m_A[module] = torch.diag(factor.new(factor.shape[0]).fill_(1)) self.m_dA[module] = factor.new_zeros(factor.shape[0]) self.m_QA[module] = factor.new_zeros(factor.shape) def _init_G(self, factor, module): """Initialize memory for factor G and its eigendecomp""" self.m_G[module] = torch.diag(factor.new(factor.shape[0]).fill_(1)) self.m_dG[module] = factor.new_zeros(factor.shape[0]) self.m_QG[module] = factor.new_zeros(factor.shape) def _clear_eigen(self): """Clear eigendecompositions Useful for when switching between `diag_blocks=1` and `diag-blocks>1` because eigendecompositions saved in place and the off-diagonals must be cleared. """ for module in self.modules: self.m_QA[module].fill_(0) self.m_QG[module].fill_(0) self.m_dA[module].fill_(0) self.m_dG[module].fill_(0) def _update_module_A(self, module): a = self.computeA(self.m_a[module], module) if self.steps == 0: self._init_A(a, module) update_running_avg(a, self.m_A[module], self.factor_decay) if self.sparse: sparsification(self.m_A[module], module, ratio=self.sparse_ratio, residuals=self.residualsA) def _update_A(self): """Compute and update factor A for all modules""" for module in self.modules: self._update_module_A(module) def _update_module_G(self, module): g = self.computeG(self.m_g[module], module, self.batch_averaged) #logger.info('G Name: %s, shape: %s', module, g.shape) if self.steps == 0: self._init_G(g, module) update_running_avg(g, self.m_G[module], self.factor_decay) if self.sparse: sparsification(self.m_G[module], module, ratio=self.sparse_ratio, residuals=self.residualsG) def _update_G(self): """Compute and update factor G for all modules""" for module in self.modules: self._update_module_G(module) def _update_eigen_A(self, module, ranks): """Compute eigendecomposition of A for module on specified workers Note: all ranks will enter this function but only the ranks specified in `ranks` will continue to actually compute the eigendecomposition. All other ranks will simply zero out their buffer for the eigendecomposition for the current module. This is done so we can sum the eigendecompositions across all ranks to communicate the results of locally computed eigendecompositions. Args: module: module to compute eigendecomposition of A on ranks: list of horovod ranks (i.e. workers) to use when computing the eigendecomposition. """ if hvd.rank() in ranks: self._distributed_compute_eigen(self.m_A[module], self.m_QA[module], self.m_dA[module], ranks) else: if ranks[0] == -1: self._local_computer_inverse(self.m_A[module], self.m_QA[module]) else: self.m_QA[module].fill_(0) def _local_computer_inverse(self, factor, inverse): block = factor[0:, 0:] block = add_value_to_diagonal(block, self.damping) inv = torchsso.utils.inv(block) inverse.data[0:, 0:].copy_(inv) def _update_eigen_G(self, module, ranks): """Compute eigendecomposition of A for module on specified workers See `_update_eigen_A` for more info` """ if hvd.rank() in ranks: self._distributed_compute_eigen(self.m_G[module], self.m_QG[module], self.m_dG[module], ranks) else: if ranks[0] == -1: self._local_computer_inverse(self.m_G[module], self.m_QG[module]) else: self.m_QG[module].fill_(0) def _distributed_compute_eigen(self, factor, evectors, evalues, ranks): """Computes the eigendecomposition of a factor across ranks Assigns each rank in `ranks` to enter this function to compute a diagonal block of `factor`. Results are written to `evectors` and `evalues`. If `len(ranks)==1`, then that rank computes the eigendecomposition of the entire `factor`. Args: factor (tensor): tensor to eigendecompose evectors (tensor): tensor to save eigenvectors of `factor` to evalues (tensor): tensor to save eigenvalues of `factor` to ranks (list): list of ranks that will enter this function """ i = ranks.index(hvd.rank()) n = len(ranks) if n > min(factor.shape): n = min(factor.shape) if i < n: start, end = get_block_boundary(i, n, factor.shape) block = factor[start[0]:end[0], start[1]:end[1]] block = add_value_to_diagonal(block, self.damping) inverse = torchsso.utils.inv(block) evectors.data[start[0]:end[0], start[1]:end[1]].copy_(inverse) def _get_diag_blocks(self, module, diag_blocks): """Helper method for determining number of diag_blocks to use Overrides `diag_blocks` if the `module` does not support `diag_blocks>1`. I.e. for a Linear layer, we do not want to use a `diag_blocks>1`. Args: module: module diag_blocks (int): default number of diag blocks to use """ return diag_blocks if module.__class__.__name__ == 'Conv2d' else 1 def _get_grad(self, module): """Get formated gradient of module Args: module: module/layer to get gradient of Returns: Formatted gradient with shape [output_dim, input_dim] for module """ if module.__class__.__name__ == 'Conv2d': # n_filters * (in_c * kw * kh) grad = module.weight.grad.data.view( module.weight.grad.data.size(0), -1) else: grad = module.weight.grad.data if module.bias is not None: grad = torch.cat([grad, module.bias.grad.data.view(-1, 1)], 1) return grad def _get_preconditioned_grad(self, module, grad): """Precondition gradient of module Args: module: module to compute preconditioned gradient for grad: formatted gradient from `_get_grad()` Returns: preconditioned gradient with same shape as `grad` """ #v = self.m_QG[module].t() @ grad @ self.m_QA[module] v = self.m_QG[module] @ grad @ self.m_QA[module] if module.bias is not None: v = [v[:, :-1], v[:, -1:]] v[0] = v[0].view(module.weight.grad.data.size()) # weight v[1] = v[1].view(module.bias.grad.data.size()) # bias else: v = [v.view(module.weight.grad.data.size())] return v def _update_scale_grad(self, updates): """Update the gradients in place and scale Updates the gradients in-place for all modules using the preconditioned gradients and scales the gradients. Args: updates (dict): dict of {module: precon_grad} """ vg_sum = 0 for module in self.modules: v = updates[module] vg_sum += (v[0] * module.weight.grad.data * self.lr**2).sum().item() if module.bias is not None: vg_sum += (v[1] * module.bias.grad.data * self.lr**2).sum().item() if self.exclude_communicate_inverse: nu = 1 else: nu = min(1.0, math.sqrt(self.kl_clip / abs(vg_sum))) for module in self.modules: v = updates[module] module.weight.grad.data.copy_(v[0]) module.weight.grad.data.mul_(nu) if module.bias is not None: module.bias.grad.data.copy_(v[1]) module.bias.grad.data.mul_(nu) def step(self, closure=None, epoch=None): """Perform one K-FAC step Note: - this function should always be called before `optimizer.step()` - gradients must be averaged across ranks before calling `step()` Args: closure: for compatibility with the base optimizer class. `closure` is ignored by KFAC epoch (int, optional): epoch to use for determining when to end the `diag_warmup` period. `epoch` is not necessary if not using `diag_warmup` """ # Update params, used for compatibilty with `KFACParamScheduler` group = self.param_groups[0] self.lr = group['lr'] self.damping = group['damping'] self.fac_update_freq = group['fac_update_freq'] self.kfac_update_freq = group['kfac_update_freq'] updates = {} handles = [] if epoch is None: if self.diag_warmup > 0: print("WARNING: diag_warmup > 0 but epoch was not passed to " "KFAC.step(). Defaulting to no diag_warmup") diag_blocks = self.diag_blocks else: diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 if self.steps % self.fac_update_freq == 0: if self.eigen_ranks is None: if not self.exclude_compute_factor: self._update_A() self._update_G() self.eigen_ranks, self.fusion_groups_A, self.fusion_groups_G = self._generate_eigen_ranks_blockpartition_opt( epoch) if not self.exclude_communicate_factor: if hvd.size() > 1: self._reduce_factors(self.eigen_ranks) self.fw_merged_comm.init_tensor_group(self.reduce_module_names) self.bw_merged_comm.init_tensor_group( self.reduce_module_names[::-1]) if hvd.rank() == 0: print('module_names: ', self.module_names) print('fusion_groups_A: ', self.fusion_groups_A) self.fw_merged_comm.update_tensor_fusion(self.fusion_groups_A) else: # starting from the 2nd iteration if not self.exclude_communicate_factor: if hvd.size() > 1: self.fw_merged_comm.synchronize() self.bw_merged_comm.synchronize() self.fw_allreduce_comm.synchronize() self.bw_allreduce_comm.synchronize() eigen_ranks = self.eigen_ranks # if we are switching from no diag approx to approx, we need to clear # off-block-diagonal elements if not self.have_cleared_Q and \ epoch == self.diag_warmup and \ self.steps % self.kfac_update_freq == 0: self._clear_eigen() self.have_cleared_Q = True if self.steps % self.kfac_update_freq == 0: # reset rank iter so device get the same layers # to compute to take advantage of caching self.rank_iter.reset() handles = [] merged_name_AGs = [[]] * hvd.size() merged_tensor_AGs = [[]] * hvd.size() for i, module in enumerate(self.modules): ranks_a, ranks_g = eigen_ranks[module] self.m_dA_ranks[module] = ranks_a[0] self.m_dG_ranks[module] = ranks_g[0] rank_a = ranks_a[0] rank_g = ranks_g[0] name = self.module_name_map[module] if not self.exclude_compute_inverse: self._update_eigen_A(module, ranks_a) if not self.exclude_compute_inverse: self._update_eigen_G(module, ranks_g) merged_name_AGs[rank_a].append(name + '-A') merged_name_AGs[rank_g].append(name + '-G') merged_tensor_AGs[rank_a].append(self.m_QA[module]) merged_tensor_AGs[rank_g].append(self.m_QG[module]) if not self.exclude_communicate_inverse: if hvd.size() > 1: #for rank, names in enumerate(merged_name_AGs): # merged_names = merged_name_AGs[rank] # merged_tensors = merged_tensor_AGs[rank] # self.multi_comm.bcast_async_(merged_names, merged_tensors, rank) self._broadcast_eigendecomp() if self.exclude_communicate_inverse and not self.exclude_compute_inverse: # should have a barriar if hvd.size() > 1: barrier() if hvd.size() > 1 and self.steps % self.kfac_update_freq == 0: self.multi_comm.synchronize() for i, module in enumerate(self.modules): grad = self._get_grad(module) precon_grad = self._get_preconditioned_grad(module, grad) updates[module] = precon_grad self._update_scale_grad(updates) if hvd.size() > 1 and self.steps % self.kfac_update_freq == 0: if self.steps == 5: self.profiling = True elif self.steps == 25: fw_layerwise_times = torch.tensor( self.fw_profiler.get_results()) bw_layerwise_times = torch.tensor( self.bw_profiler.get_results()) hvd.broadcast_(fw_layerwise_times, root_rank=0) hvd.broadcast_(bw_layerwise_times, root_rank=0) fw_layerwise_times = fw_layerwise_times.numpy() bw_layerwise_times = bw_layerwise_times.numpy() if hvd.rank() == 0: pass fw_factor_sizes = [ self.fw_factor_sizes[m] for m in self.module_names ] bw_factor_sizes = [ self.bw_factor_sizes[m] for m in self.module_names[::-1] ] self.fw_merged_comm.update_groups(self.fusion_groups_A, fw_factor_sizes, fw_layerwise_times, reverse=False) #self.bw_merged_comm.update_groups(self.fusion_groups_G, bw_factor_sizes, bw_layerwise_times, reverse=False) self.profiling = False self.steps += 1 def _generate_eigen_ranks_naive(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 buckets = [0] * hvd.size() for module in self.modules: # Get ranks to compute this layer on n = self._get_diag_blocks(module, diag_blocks) ranks_a = self.rank_iter.next(n) ranks_g = self.rank_iter.next(n) if self.distribute_layer_factors \ else ranks_a module_ranks[module] = (ranks_a, ranks_g) buckets[ranks_a[0]] += self.m_A[module].shape[1] buckets[ranks_g[0]] += self.m_G[module].shape[1] self.module_ranks = module_ranks if hvd.rank() == 0: #logger.info('buckets: %s', buckets) #logger.info('module_ranks: %s', module_ranks.values()) logger.info( 'buckets mean: %f, std: %f, max: %f, min: %f, max-min=%f', np.mean(buckets), np.std(buckets), np.max(buckets), np.min(buckets), np.max(buckets) - np.min(buckets)) return module_ranks def _generate_eigen_ranks_match_merging(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 assigned_rank = 0 assigned_count = 0 #print('nfactors: ', len(self.modules)) buckets = [0] * hvd.size() BIs = [] for i, module in enumerate(self.modules): # Get ranks to compute this layer on name = self.module_name_map[module] a_dimension = self.m_A[module].shape[1] g_dimension = self.m_G[module].shape[1] if a_dimension > REDUCE_THRESHOLD: self.reduce_module_names.append(name) rank = assigned_rank assigned_count += 1 else: rank = -1 ranks_a = (rank, ) ranks_g = (rank, ) module_ranks[module] = (ranks_a, ranks_g) buckets[rank] += (a_dimension + g_dimension) BIs.append(rank) if assigned_count > 0 and assigned_count % NUM_NEARBY_LAYERS == 0: assigned_rank += 1 assigned_rank %= hvd.size() self.module_ranks = module_ranks if hvd.rank() == 0: logger.info( 'buckets mean: %f, std: %f, max: %f, min: %f, max-min=%f', np.mean(buckets), np.std(buckets), np.max(buckets), np.min(buckets), np.max(buckets) - np.min(buckets)) logger.info('BIs: %s', BIs) return module_ranks def _generate_eigen_ranks_blockpartition_opt(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 target_ranks = np.arange(hvd.size()) buckets = [0] * hvd.size() dimensions_A = [] module_factors_A = [] factors_A = [] dimensions_G = [] module_factors_G = [] factors_G = [] for i, m in enumerate(self.modules): name = self.module_names[i] a_dimension = self.m_A[m].shape[1] a_dimension = a_dimension * (a_dimension + 1) / 2 dimensions_A.append(a_dimension) module_factors_A.append(name + '-A') factors_A.append(m) for i, m in enumerate(self.modules): name = self.module_names[i] g_dimension = self.m_G[m].shape[1] g_dimension = g_dimension * (g_dimension + 1) / 2 dimensions_G.append(g_dimension) module_factors_G.append(name + '-G') factors_G.append(m) dimensions_G = dimensions_G[::-1] factors_G = factors_G[::-1] placement_A = get_optimal_block_partition(len(self.modules), hvd.size(), dimensions_A) placement_G = get_optimal_block_partition(len(self.modules), hvd.size(), dimensions_G) if hvd.rank() == 0: print('dimension_G: ', dimensions_G) print('placement_G: ', placement_G) placement = list(placement_A) + list(placement_G) A_ranks = {} G_ranks = {} buckets = [0] * hvd.size() for i, p in enumerate(placement_A): AorG = factors_A[i] A_ranks[AorG] = (p, ) dimension = self.m_A[AorG].shape[1] buckets[p] += dimension for i, p in enumerate(placement_G): AorG = factors_G[i] G_ranks[AorG] = (p, ) dimension = self.m_G[AorG].shape[1] buckets[p] += dimension for m in self.modules: module_ranks[m] = (A_ranks[m], G_ranks[m]) self.module_ranks = module_ranks tensor_names_As = [] tensor_names_Gs = [] c_rank = A_ranks[factors_A[0]] current_group = [] tensor_names_As.append(current_group) for i, name in enumerate(self.module_names): m = factors_A[i] a_rank = A_ranks[m][0] if a_rank != c_rank: current_group = [] tensor_names_As.append(current_group) current_group.append(name) c_rank = a_rank c_rank = G_ranks[factors_G[0]] current_group = [] tensor_names_Gs.append(current_group) for i, name in enumerate(self.module_names[::-1]): m = factors_G[i] g_rank = G_ranks[m][0] if g_rank != c_rank: current_group = [] tensor_names_Gs.append(current_group) current_group.append(name) c_rank = g_rank if hvd.rank() == 0: logger.info('placement: %s', placement) logger.info('buckets: %s', buckets) logger.info('module_ranks: %s', module_ranks.values()) logger.info( 'buckets mean: %f, std: %f, max: %f, min: %f, max-min=%f', np.mean(buckets), np.std(buckets), np.max(buckets), np.min(buckets), np.max(buckets) - np.min(buckets)) return module_ranks, tensor_names_As, tensor_names_Gs def _generate_eigen_ranks_uniform(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 buckets = [0] * hvd.size() dimensions = [] module_factors = [] for i, m in enumerate(self.modules): name = self.module_names[i] a_dimension = self.m_A[m].shape[1] #g_dimension = self.m_G[m].shape[1] dimensions.append(a_dimension) module_factors.append(name + '-A') #dimensions.append(g_dimension) #module_factors.append(name+'-G') descending_sorted_idx = np.argsort(dimensions)[::-1] A_ranks = {} G_ranks = {} BIs = [] for i in descending_sorted_idx: factor = module_factors[i] dimension = dimensions[i] m_i = self.module_names.index(factor[0:-2]) m = self.modules[m_i] bi = np.argmin(buckets) BIs.append(bi) buckets[bi] += dimension if factor[-1] == 'A': A_ranks[m] = (bi, ) G_ranks[m] = (bi, ) else: G_ranks[m] = (bi, ) for m in self.modules: module_ranks[m] = (A_ranks[m], G_ranks[m]) self.module_ranks = module_ranks if hvd.rank() == 0: logger.info('buckets: %s', buckets) logger.info('module_ranks: %s', module_ranks.values()) logger.info( 'buckets mean: %f, std: %f, max: %f, min: %f, max-min=%f', np.mean(buckets), np.std(buckets), np.max(buckets), np.min(buckets), np.max(buckets) - np.min(buckets)) logger.info('BIs: %s', BIs) return module_ranks def _generate_eigen_ranks(self, epoch): if self.module_ranks is not None: return self.module_ranks module_ranks = {} diag_blocks = self.diag_blocks if epoch >= self.diag_warmup else 1 buckets = [0] * hvd.size() for module in self.modules: i = np.argmin(buckets) if hvd.rank() == 0: logger.info('A Name: %s, shape: %s', module, self.m_A[module].shape) logger.info('G Name: %s, shape: %s', module, self.m_G[module].shape) a_dimension = self.m_A[module].shape[1] g_dimension = self.m_G[module].shape[1] #buckets[i] += (a_dimension) + g_dimension) buckets[i] += a_dimension ranks_a = (i, ) i = np.argmin(buckets) ranks_g = (i, ) buckets[i] += g_dimension module_ranks[module] = (ranks_a, ranks_g) self.module_ranks = module_ranks if hvd.rank() == 0: logger.info('buckets: %s', buckets) logger.info('module_ranks: %s', module_ranks.values()) logger.info( 'buckets mean: %f, std: %f, max: %f, min: %f, max-min=%f', np.mean(buckets), np.std(buckets), np.max(buckets), np.min(buckets), np.max(buckets) - np.min(buckets)) return module_ranks def _reduce_factors(self, eigen_ranks): """Allreduce the factors for all layers""" handles = [] for m in self.modules: name = self.module_name_map[m] ranks_a, ranks_g = eigen_ranks[m] rank_a = ranks_a[0] rank_g = ranks_g[0] handles.append( hvd.allreduce_async_(self.m_A[m].data, op=hvd.Average)) handles.append( hvd.allreduce_async_(self.m_G[m].data, op=hvd.Average)) for handle in handles: hvd.synchronize(handle) def _allgather_factors(self): """Allgather the factors for all layers""" handles = [] def _get_value_and_idx(sparse_tensor): tensor = sparse_tensor.data.view(-1) one_indexes = tensor != 0 indexes = one_indexes.nonzero().data.squeeze().view(-1) values = tensor.data[indexes] return values, indexes.int() for i, m in enumerate(self.modules): module_name = self.module_names[i] A_values, A_indexes = _get_value_and_idx(self.m_A[m].data) A_value_name = module_name + '_A_value' A_idx_name = module_name + '_A_idx' h_value = allgather_async(A_values, A_value_name) h_idx = allgather_async(A_indexes, A_idx_name) G_values, G_indexes = _get_value_and_idx(self.m_G[m].data) G_value_name = module_name + '_G_value' G_idx_name = module_name + '_G_idx' h_value_G = allgather_async(G_values, G_value_name) h_idx_G = allgather_async(G_indexes, G_idx_name) handles.append((h_value, h_idx, h_value_G, h_idx_G)) for i, handle in enumerate(handles): module_name = self.module_names[i] module = self.modules[i] m_A = self.m_A[module].view(-1) m_A.fill_(0.0) m_G = self.m_G[module].view(-1) m_G.fill_(0.0) h_value_A, h_idx_A, h_value_G, h_idx_G = handle A_values = hvd.synchronize(h_value_A) A_indexes = hvd.synchronize(h_idx_A).long() m_A.scatter_add_(0, A_indexes, A_values) m_A.div_(hvd.size()) G_values = hvd.synchronize(h_value_G) G_indexes = hvd.synchronize(h_idx_G).long() m_G.scatter_add_(0, G_indexes, G_values) m_G.div_(hvd.size()) def _allreduce_eigendecomp(self): """Allreduce the eigendecompositions for all layers Note: we use `op=hvd.Sum` to simulate an allgather`. Each rank will either compute the eigendecomposition for a factor or just return zeros so we sum instead of averaging. """ handles = [] for m in self.modules: handles.append(hvd.allreduce_async_(self.m_QA[m].data, op=hvd.Sum)) handles.append(hvd.allreduce_async_(self.m_QG[m].data, op=hvd.Sum)) handles.append(hvd.allreduce_async_(self.m_dA[m].data, op=hvd.Sum)) handles.append(hvd.allreduce_async_(self.m_dG[m].data, op=hvd.Sum)) for handle in handles: hvd.synchronize(handle) def _broadcast_eigendecomp(self): """Broadcasts the eigendecompositions for all layers Note: we use `op=hvd.Sum` to simulate an allgather`. Each rank will either compute the eigendecomposition for a factor or just return zeros so we sum instead of averaging. """ rank = hvd.rank() for i, m in enumerate(self.modules): rank_a = self.m_dA_ranks[m] rank_g = self.m_dG_ranks[m] name = self.module_names[i] if rank_a >= 0: self.multi_comm.bcast_async_([name + 'mQA'], [self.m_QA[m]], rank_a) if rank_g >= 0: self.multi_comm.bcast_async_([name + 'mQG'], [self.m_QG[m]], rank_g) self.multi_comm.synchronize()