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()
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.
        
python/dgl/distributed/optim/pytorch/sparse_optim.py