class_docstr: 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.
docstr: 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.
docstr_tok: ["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","."]
code: 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():
          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.
              idics = (
                  th.cat(idics, dim=0)
                  if len(idics) != 0
                  else th.zeros(
                      (0,), dtype=th.long, device=th.device("cpu")
                  )
              )
              grads = (
                  th.cat(grads, dim=0)
                  if len(grads) != 0
                  else th.zeros(
                      (0, emb.embedding_dim),
                      dtype=th.float32,
                      device=th.device("cpu"),
                  )
              )
              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=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=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=device,
                      ).chunk(self._world_size)
                  )
                  alltoall(
                      self._rank,
                      self._world_size,
                      gather_list,
                      idx_split_size,
                      device,
                  )
                  idx_gather_list = [
                      th.empty(
                          (int(num_emb),),
                          dtype=idics.dtype,
                          device=device,
                      )
                      for num_emb in gather_list
                  ]
                  alltoallv(
                      self._rank,
                      self._world_size,
                      idx_gather_list,
                      idics_list,
                      device,
                  )
                  local_indics[name] = idx_gather_list
                  grad_gather_list = [
                      th.empty(
                          (int(num_emb), grads.shape[1]),
                          dtype=grads.dtype,
                          device=device,
                      )
                      for num_emb in gather_list
                  ]
                  alltoallv(
                      self._rank,
                      self._world_size,
                      grad_gather_list,
                      grad_list,
                      device,
                  )
                  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(device, non_blocking=True),
                  grad.to(device, non_blocking=True),
                  emb,
              )
  
      # synchronized gradient update
      if self._world_size > 1:
          th.distributed.barrier()
  
imports: import abc
  import warnings
  from abc import abstractmethod
  from os.path import exists
  import torch
  import dgl
  from ....None import backend
  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
label_desc: Use image node_id 8 for calling the DistSparseGradOptimizer obj's underlying member method code with example usage: obj.step() without return types