"""Node embedding optimizers for distributed training""" import abc import warnings from abc import abstractmethod from os.path import exists import torch as th import dgl from .... import backend as F from ...dist_tensor import DistTensor from ...graph_partition_book import EDGE_PART_POLICY, NODE_PART_POLICY from ...nn.pytorch import DistEmbedding from .utils import alltoall, alltoallv EMB_STATES = "emb_states" WORLD_SIZE = "world_size" IDS = "ids" PARAMS = "params" STATES = "states" class DistSparseGradOptimizer(abc.ABC): r"""The abstract dist sparse optimizer. Note: dgl dist sparse optimizer only work with dgl.distributed.DistEmbedding Parameters ---------- params : list of DistEmbedding The list of DistEmbedding. lr : float The learning rate. """ def __init__(self, params, lr): self._params = params self._lr = lr self._rank = None self._world_size = None self._shared_cache = {} self._clean_grad = False self._opt_meta = {} self._state = {} ## collect all hyper parameters for save self._defaults = {} if th.distributed.is_initialized(): self._rank = th.distributed.get_rank() self._world_size = th.distributed.get_world_size() else: self._rank = 0 self._world_size = 1 def local_state_dict(self): """Return the state pertaining to current rank of the optimizer. Returns ------- dict Local state dict Example Dict of Adagrad Optimizer: .. code-block:: json { "params": { "_lr": 0.01, "_eps": "1e-8", "world_size": 2 }, "emb_states": { "emb_name1": { "ids": [0, 2, 4, 6 ,8 ,10], ## tensor, "emb_name1_sum": [0.1 , 0.2, 0.5, 0.1, 0.2] ## tensor, }, "emb_name2": { "ids": [0, 2, 4, 6 ,8 ,10], ## tensor, "emb_name2_sum": [0.3 , 0.2, 0.4, 0.5, 0.2] ## tensor, } } } :param json: json object See Also -------- load_local_state_dict """ local_state_dict = {} local_state_dict[EMB_STATES] = {} local_state_dict[PARAMS] = {WORLD_SIZE: self._world_size} for emb in self._params: trainers_per_machine = self._world_size // max( 1, dgl.distributed.get_num_machines() ) emb_state_dict = {} part_policy = ( emb.part_policy if emb.part_policy else emb.weight.part_policy ) idx = self._get_local_ids(part_policy) if trainers_per_machine > 1: kv_idx_split = (idx % trainers_per_machine).long() local_rank = self._rank % trainers_per_machine mask = kv_idx_split == local_rank idx = F.boolean_mask(idx, mask) emb_state_dict.update({IDS: idx}) emb_state = {} states = ( list(self._state[emb.name]) if isinstance(self._state[emb.name], tuple) else [self._state[emb.name]] ) emb_state = {state.name: state[idx] for state in states} emb_state_dict.update({STATES: emb_state}) local_state_dict[EMB_STATES].update({emb.name: emb_state_dict}) local_state_dict[PARAMS].update(self._defaults) return local_state_dict def load_local_state_dict(self, local_state_dict): """Load the local state from the input state_dict, updating the optimizer as needed. Parameters ---------- local_state_dict : dict Optimizer state; should be an object returned from a call to local_state_dict(). See Also -------- local_state_dict """ for emb_name, emb_state in local_state_dict[EMB_STATES].items(): idx = emb_state[IDS] # As state of an embedding of different optimizers can be a single # DistTensor(Adagrad) or a tuple(Adam) of that, converting it to list for # consistency. The list contains reference(s) to original DistTensor(s). states = ( list(self._state[emb_name]) if isinstance(self._state[emb_name], tuple) else [self._state[emb_name]] ) if len(emb_state[STATES]) != len(states): raise ValueError( f"loaded state dict has a different number of states" f" of embedding {emb_name}" ) name_to_index = { state.name: index for index, state in enumerate(states) } for name, state in emb_state[STATES].items(): if name not in name_to_index: raise ValueError( "loaded state dict contains a state {name}" "that can't be found in the optimizer states" ) state_idx = name_to_index[name] state = state.to( th.device("cpu"), states[name_to_index[name]].dtype ) states[state_idx][idx] = state self._defaults.update(local_state_dict[PARAMS]) self.__dict__.update(local_state_dict[PARAMS]) def save(self, f): """Save the local state_dict to disk on per rank. Saved dict contains 2 parts: * 'params': hyper parameters of the optimizer. * 'emb_states': partial optimizer states, each embedding contains 2 items: 1. ```ids```: global id of the nodes/edges stored in this rank. 2. ```states```: state data corrseponding to ```ids```. NOTE: This needs to be called on all ranks. Parameters ---------- f : Union[str, os.PathLike] The path of the file to save to. See Also -------- load """ if self._world_size > 1: th.distributed.barrier() f = f if isinstance(f, str) else str(f, "UTF-8") f = f"{f}_{self._rank}" th.save(self.local_state_dict(), f) if self._world_size > 1: th.distributed.barrier() def load(self, f): """Load the local state of the optimizer from the file on per rank. NOTE: This needs to be called on all ranks. Parameters ---------- f : Union[str, os.PathLike] The path of the file to load from. See Also -------- save """ if self._world_size > 1: th.distributed.barrier() f = f if isinstance(f, str) else str(f, "UTF-8") f_attach_rank = f"{f}_{self._rank}" # Don't throw error here to support device number scale-out # after reloading, but make sure your hyper parameter is same # as before because new added local optimizers will be filled # in nothing if not exists(f_attach_rank): warnings.warn(f"File {f_attach_rank} can't be found, load nothing.") else: old_world_size = self._load_state_from(f_attach_rank) # Device number scale-in if self._world_size < old_world_size: for rank in range( self._rank + self._world_size, old_world_size, self._world_size, ): self._load_state_from(f"{f}_{rank}") if self._world_size > 1: th.distributed.barrier() def _load_state_from(self, f): local_state_dict = th.load(f) world_size = local_state_dict[PARAMS].pop(WORLD_SIZE) self.load_local_state_dict(local_state_dict) return world_size def _get_local_ids(self, part_policy): if EDGE_PART_POLICY in part_policy.policy_str: return part_policy.partition_book.partid2eids( part_policy.part_id, part_policy.type_name ) elif NODE_PART_POLICY in part_policy.policy_str: return part_policy._partition_book.partid2nids( part_policy.part_id, part_policy.type_name ) else: raise RuntimeError( "Cannot support policy: %s " % part_policy.policy_str ) def step(self): """The step function. The step function is invoked at the end of every batch to push the gradients of the embeddings involved in a mini-batch to DGL's servers and update the embeddings. """ with th.no_grad(): # [Rui] # As `gloo` supports CPU tensors only while `nccl` supports GPU # tensors only, we firstly create tensors on the corresponding # devices and then copy the data to target device if needed. # Please note that the target device can be different from the # preferred device. target_device = None preferred_device = ( th.device(f"cuda:{self._rank}") if th.distributed.get_backend() == "nccl" else th.device("cpu") ) local_indics = {emb.name: [] for emb in self._params} local_grads = {emb.name: [] for emb in self._params} for emb in self._params: name = emb.weight.name kvstore = emb.weight.kvstore trainers_per_server = self._world_size // kvstore.num_servers idics = [] grads = [] for trace in emb._trace: if trace[1].grad is not None: idics.append(trace[0]) grads.append(trace[1].grad.data) else: assert len(trace[0]) == 0 # If the sparse embedding is not used in the previous forward step # The idx and grad will be empty, initialize them as empty tensors to # avoid crashing the optimizer step logic. # # Note: we cannot skip the gradient exchange and update steps as other # working processes may send gradient update requests corresponding # to certain embedding to this process. # # [WARNING][TODO][Rui] # For empty idx and grad, we blindly create data on the # preferred device, which may not be the device where the # embedding is stored. idics = ( th.cat(idics, dim=0) if len(idics) != 0 else th.zeros((0,), dtype=th.int64, device=preferred_device) ) grads = ( th.cat(grads, dim=0) if len(grads) != 0 else th.zeros( (0, emb.embedding_dim), dtype=th.float32, device=preferred_device, ) ) target_device = grads.device # will send grad to each corresponding trainer if self._world_size > 1: # get idx split from kvstore idx_split = kvstore.get_partid(emb.data_name, idics) idx_split_size = [] idics_list = [] grad_list = [] # split idx and grad first for i in range(kvstore.num_servers): mask = idx_split == i idx_i = idics[mask] grad_i = grads[mask] if trainers_per_server <= 1: idx_split_size.append( th.tensor( [idx_i.shape[0]], dtype=th.int64, device=preferred_device, ) ) idics_list.append(idx_i) grad_list.append(grad_i) else: kv_idx_split = th.remainder( idx_i, trainers_per_server ).long() for j in range(trainers_per_server): mask = kv_idx_split == j idx_j = idx_i[mask] grad_j = grad_i[mask] idx_split_size.append( th.tensor( [idx_j.shape[0]], dtype=th.int64, device=preferred_device, ) ) idics_list.append(idx_j) grad_list.append(grad_j) # if one machine launch multiple KVServer, they share the same storage. # For each machine, the pytorch rank is num_trainers * # machine_id + i # use scatter to sync across trainers about the p2p tensor size # Note: If we have GPU nccl support, we can use all_to_all to # sync information here gather_list = list( th.empty( [self._world_size], dtype=th.int64, device=preferred_device, ).chunk(self._world_size) ) alltoall( self._rank, self._world_size, gather_list, idx_split_size, ) idx_gather_list = [ th.empty( (int(num_emb),), dtype=idics.dtype, device=preferred_device, ) for num_emb in gather_list ] alltoallv( self._rank, self._world_size, idx_gather_list, idics_list, ) local_indics[name] = idx_gather_list grad_gather_list = [ th.empty( (int(num_emb), grads.shape[1]), dtype=grads.dtype, device=preferred_device, ) for num_emb in gather_list ] alltoallv( self._rank, self._world_size, grad_gather_list, grad_list, ) local_grads[name] = grad_gather_list else: local_indics[name] = [idics] local_grads[name] = [grads] if self._clean_grad: # clean gradient track for emb in self._params: emb.reset_trace() self._clean_grad = False # do local update for emb in self._params: name = emb.weight.name idx = th.cat(local_indics[name], dim=0) grad = th.cat(local_grads[name], dim=0) self.update( idx.to(target_device, non_blocking=True), grad.to(target_device, non_blocking=True), emb, ) # synchronized gradient update if self._world_size > 1: th.distributed.barrier() @abstractmethod def update(self, idx, grad, emb): """Update embeddings in a sparse manner Sparse embeddings are updated in mini batches. We maintain gradient states for each embedding so they can be updated separately. Parameters ---------- idx : tensor Index of the embeddings to be updated. grad : tensor Gradient of each embedding. emb : dgl.distributed.DistEmbedding Sparse node embedding to update. """ def zero_grad(self): """clean grad cache""" self._clean_grad = True def initializer(shape, dtype): """Sparse optimizer state initializer Parameters ---------- shape : tuple of ints The shape of the state tensor dtype : torch dtype The data type of the state tensor """ arr = th.zeros(shape, dtype=dtype) return arr class SparseAdagrad(DistSparseGradOptimizer): r"""Distributed Node embedding optimizer using the Adagrad algorithm. This optimizer implements a distributed sparse version of Adagrad algorithm for optimizing :class:`dgl.distributed.DistEmbedding`. Being sparse means it only updates the embeddings whose gradients have updates, which are usually a very small portion of the total embeddings. Adagrad maintains a :math:`G_{t,i,j}` for every parameter in the embeddings, where :math:`G_{t,i,j}=G_{t-1,i,j} + g_{t,i,j}^2` and :math:`g_{t,i,j}` is the gradient of the dimension :math:`j` of embedding :math:`i` at step :math:`t`. NOTE: The support of sparse Adagrad optimizer is experimental. Parameters ---------- params : list[dgl.distributed.DistEmbedding] The list of dgl.distributed.DistEmbedding. lr : float The learning rate. eps : float, Optional The term added to the denominator to improve numerical stability Default: 1e-10 """ def __init__(self, params, lr, eps=1e-10): super(SparseAdagrad, self).__init__(params, lr) self._eps = eps self._defaults = {"_lr": lr, "_eps": eps} # We need to register a state sum for each embedding in the kvstore. for emb in params: assert isinstance( emb, DistEmbedding ), "SparseAdagrad only supports dgl.distributed.DistEmbedding" name = emb.name + "_sum" state = DistTensor( (emb.num_embeddings, emb.embedding_dim), th.float32, name, init_func=initializer, part_policy=emb.part_policy, is_gdata=False, ) assert ( emb.name not in self._state ), "{} already registered in the optimizer".format(emb.name) self._state[emb.name] = state def update(self, idx, grad, emb): """Update embeddings in a sparse manner Sparse embeddings are updated in mini batches. We maintain gradient states for each embedding so they can be updated separately. Parameters ---------- idx : tensor Index of the embeddings to be updated. grad : tensor Gradient of each embedding. emb : dgl.distributed.DistEmbedding Sparse embedding to update. """ eps = self._eps clr = self._lr state_dev = th.device("cpu") exec_dev = grad.device # only perform async copies cpu -> gpu, or gpu-> gpu, but block # when copying to the cpu, so as to ensure the copy is finished # before operating on the data on the cpu state_block = state_dev == th.device("cpu") and exec_dev != state_dev # the update is non-linear so indices must be unique grad_indices, inverse, cnt = th.unique( idx, return_inverse=True, return_counts=True ) grad_values = th.zeros( (grad_indices.shape[0], grad.shape[1]), device=exec_dev ) grad_values.index_add_(0, inverse, grad) grad_values = grad_values / cnt.unsqueeze(1) grad_sum = grad_values * grad_values # update grad state grad_state = self._state[emb.name][grad_indices].to(exec_dev) grad_state += grad_sum grad_state_dst = grad_state.to(state_dev, non_blocking=True) if state_block: # use events to try and overlap CPU and GPU as much as possible update_event = th.cuda.Event() update_event.record() # update emb std_values = grad_state.sqrt_().add_(eps) tmp = clr * grad_values / std_values tmp_dst = tmp.to(state_dev, non_blocking=True) if state_block: std_event = th.cuda.Event() std_event.record() # wait for our transfers from exec_dev to state_dev to finish # before we can use them update_event.wait() self._state[emb.name][grad_indices] = grad_state_dst if state_block: # wait for the transfer of std_values to finish before we # can use it std_event.wait() emb._tensor[grad_indices] -= tmp_dst class SparseAdam(DistSparseGradOptimizer): r"""Distributed Node embedding optimizer using the Adam algorithm. This optimizer implements a distributed sparse version of Adam algorithm for optimizing :class:`dgl.distributed.DistEmbedding`. Being sparse means it only updates the embeddings whose gradients have updates, which are usually a very small portion of the total embeddings. Adam maintains a :math:`Gm_{t,i,j}` and `Gp_{t,i,j}` for every parameter in the embeddings, where :math:`Gm_{t,i,j}=beta1 * Gm_{t-1,i,j} + (1-beta1) * g_{t,i,j}`, :math:`Gp_{t,i,j}=beta2 * Gp_{t-1,i,j} + (1-beta2) * g_{t,i,j}^2`, :math:`g_{t,i,j} = lr * Gm_{t,i,j} / (1 - beta1^t) / \sqrt{Gp_{t,i,j} / (1 - beta2^t)}` and :math:`g_{t,i,j}` is the gradient of the dimension :math:`j` of embedding :math:`i` at step :math:`t`. NOTE: The support of sparse Adam optimizer is experimental. Parameters ---------- params : list[dgl.distributed.DistEmbedding] The list of dgl.distributed.DistEmbedding. lr : float The learning rate. betas : tuple[float, float], Optional Coefficients used for computing running averages of gradient and its square. Default: (0.9, 0.999) eps : float, Optional The term added to the denominator to improve numerical stability Default: 1e-8 """ def __init__(self, params, lr, betas=(0.9, 0.999), eps=1e-08): super(SparseAdam, self).__init__(params, lr) self._eps = eps # We need to register a state sum for each embedding in the kvstore. self._beta1 = betas[0] self._beta2 = betas[1] self._defaults = { "_lr": lr, "_eps": eps, "_beta1": betas[0], "_beta2": betas[1], } for emb in params: assert isinstance( emb, DistEmbedding ), "SparseAdam only supports dgl.distributed.DistEmbedding" state_step = DistTensor( (emb.num_embeddings,), th.float32, emb.name + "_step", init_func=initializer, part_policy=emb.part_policy, is_gdata=False, ) state_mem = DistTensor( (emb.num_embeddings, emb.embedding_dim), th.float32, emb.name + "_mem", init_func=initializer, part_policy=emb.part_policy, is_gdata=False, ) state_power = DistTensor( (emb.num_embeddings, emb.embedding_dim), th.float32, emb.name + "_power", init_func=initializer, part_policy=emb.part_policy, is_gdata=False, ) state = (state_step, state_mem, state_power) assert ( emb.name not in self._state ), "{} already registered in the optimizer".format(emb.name) self._state[emb.name] = state def update(self, idx, grad, emb): """Update embeddings in a sparse manner Sparse embeddings are updated in mini batches. We maintain gradient states for each embedding so they can be updated separately. Parameters ---------- idx : tensor Index of the embeddings to be updated. grad : tensor Gradient of each embedding. emb : dgl.distributed.DistEmbedding Sparse embedding to update. """ beta1 = self._beta1 beta2 = self._beta2 eps = self._eps clr = self._lr state_step, state_mem, state_power = self._state[emb.name] state_dev = th.device("cpu") exec_dev = grad.device # only perform async copies cpu -> gpu, or gpu-> gpu, but block # when copying to the cpu, so as to ensure the copy is finished # before operating on the data on the cpu state_block = state_dev == th.device("cpu") and exec_dev != state_dev # the update is non-linear so indices must be unique grad_indices, inverse, cnt = th.unique( idx, return_inverse=True, return_counts=True ) # update grad state state_idx = grad_indices.to(state_dev) # The original implementation will cause read/write contension. # state_step[state_idx] += 1 # state_step = state_step[state_idx].to(exec_dev, non_blocking=True) # In a distributed environment, the first line of code will send write requests to # kvstore servers to update the state_step which is asynchronous and the second line # of code will also send read requests to kvstore servers. The write and read requests # may be handled by different kvstore servers managing the same portion of the # state_step dist tensor in the same node. So that, the read request may read an old # value (i.e., 0 in the first iteration) which will cause # update_power_corr to be NaN state_val = state_step[state_idx] + 1 state_step[state_idx] = state_val state_step = state_val.to(exec_dev) orig_mem = state_mem[state_idx].to(exec_dev) orig_power = state_power[state_idx].to(exec_dev) grad_values = th.zeros( (grad_indices.shape[0], grad.shape[1]), device=exec_dev ) grad_values.index_add_(0, inverse, grad) grad_values = grad_values / cnt.unsqueeze(1) grad_mem = grad_values grad_power = grad_values * grad_values update_mem = beta1 * orig_mem + (1.0 - beta1) * grad_mem update_power = beta2 * orig_power + (1.0 - beta2) * grad_power update_mem_dst = update_mem.to(state_dev, non_blocking=True) update_power_dst = update_power.to(state_dev, non_blocking=True) if state_block: # use events to try and overlap CPU and GPU as much as possible update_event = th.cuda.Event() update_event.record() update_mem_corr = update_mem / ( 1.0 - th.pow(th.tensor(beta1, device=exec_dev), state_step) ).unsqueeze(1) update_power_corr = update_power / ( 1.0 - th.pow(th.tensor(beta2, device=exec_dev), state_step) ).unsqueeze(1) std_values = clr * update_mem_corr / (th.sqrt(update_power_corr) + eps) std_values_dst = std_values.to(state_dev, non_blocking=True) if state_block: std_event = th.cuda.Event() std_event.record() # wait for our transfers from exec_dev to state_dev to finish # before we can use them update_event.wait() state_mem[state_idx] = update_mem_dst state_power[state_idx] = update_power_dst if state_block: # wait for the transfer of std_values to finish before we # can use it std_event.wait() emb._tensor[state_idx] -= std_values_dst