Source code for dgl.optim.pytorch.sparse_optim

"""Node embedding optimizers"""
import abc
from abc import abstractmethod

import torch as th

from ...cuda import nccl
from ...nn.pytorch import NodeEmbedding
from ...partition import NDArrayPartition
from ...utils import (
    create_shared_mem_array,
    gather_pinned_tensor_rows,
    get_shared_mem_array,
    pin_memory_inplace,
    scatter_pinned_tensor_rows,
)


class SparseGradOptimizer(abc.ABC):
    r"""The abstract sparse optimizer.

    Note: dgl sparse optimizer only work with dgl.NodeEmbedding

    Parameters
    ----------
    params : list of NodeEmbedding
        The list of NodeEmbeddings.
    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._comm = None
        self._first_step = True
        self._device = None
        # hold released shared memory to let other process to munmap it first
        # otherwise it will crash the training
        self.shmem_buffer_holder = []

        assert len(params) > 0, "Empty parameters"
        # if we are using shared memory for communication
        for emb in params:
            assert isinstance(
                emb, NodeEmbedding
            ), "DGL SparseOptimizer only supports dgl.nn.NodeEmbedding"

            if self._rank is None:
                self._rank = emb.rank
                self._world_size = emb.world_size
            else:
                assert (
                    self._rank == emb.rank
                ), "MultiGPU rank for each embedding should be same."
                assert (
                    self._world_size == emb.world_size
                ), "MultiGPU world_size for each embedding should be same."
        assert not self._rank is None
        assert not self._world_size is None

    def step(self):
        """The step function.

        The step function is invoked at the end of every batch to update embeddings
        """
        # on the first step, check to see if the grads are on the GPU
        if self._first_step:
            for emb in self._params:
                for _, data in emb._trace:
                    if data.grad.device.type == "cuda":
                        # create a communicator
                        if self._device:
                            assert (
                                self._device == data.grad.device
                            ), "All gradients must be on the same device"
                        else:
                            self._device = data.grad.device
                    else:
                        assert (
                            not self._device
                        ), "All gradients must be on the same device"

            # distributed backend use nccl
            if self._device and (
                not th.distributed.is_initialized()
                or th.distributed.get_backend() == "nccl"
            ):
                # device is only set if the grads are on a GPU
                self._comm_setup()
            else:
                self._shared_setup()
            self._first_step = False

        if self._comm:
            self._comm_step()
        else:
            self._shared_step()

    @abstractmethod
    def setup(self, params):
        """This is function where subclasses can perform any setup they need
        to. It will be called during the first step, and communicators or
        shared memory will have been setup before this call.

        Parameters
        ----------
        params : list of NodeEmbedding
            The list of NodeEmbeddings.
        """

    def _comm_setup(self):
        self._comm = True

    def _shared_setup(self):
        for emb in self._params:
            emb_name = emb.name
            if self._rank == 0:  # the master gpu process
                opt_meta = create_shared_mem_array(
                    emb_name + "_opt_meta",
                    (self._world_size, self._world_size),
                    th.int32,
                ).zero_()

            if self._rank == 0:
                emb.store.set(emb_name + "_opt_meta", emb_name)
                self._opt_meta[emb_name] = opt_meta
            elif self._rank > 0:
                # receive
                emb.store.wait([emb_name + "_opt_meta"])
                opt_meta = get_shared_mem_array(
                    emb_name + "_opt_meta",
                    (self._world_size, self._world_size),
                    th.int32,
                )
                self._opt_meta[emb_name] = opt_meta

    def _comm_step(self):
        with th.no_grad():
            idx_in = {}
            grad_in = {}
            for emb in self._params:  # pylint: disable=too-many-nested-blocks
                emb_name = emb.name
                partition = emb.partition

                if not partition:
                    # use default partitioning
                    partition = NDArrayPartition(
                        emb.num_embeddings,
                        self._world_size if self._world_size > 0 else 1,
                        mode="remainder",
                    )

                # we need to combine gradients from multiple forward paths
                if len(emb._trace) == 0:
                    idx = th.zeros((0,), dtype=th.long, device=self._device)
                    grad = th.zeros(
                        (0, emb.embedding_dim),
                        dtype=th.float32,
                        device=self._device,
                    )
                elif len(emb._trace) == 1:
                    # the special case where we can use the tensors as is
                    # without any memcpy's
                    idx, grad = emb._trace[0]
                    grad = grad.grad.data
                else:
                    idx = []
                    grad = []
                    for i, data in emb._trace:
                        idx.append(i)
                        grad.append(data.grad.data)
                    idx = th.cat(idx, dim=0)
                    grad = th.cat(grad, dim=0)

                (
                    idx_in[emb_name],
                    grad_in[emb_name],
                ) = nccl.sparse_all_to_all_push(idx, grad, partition=partition)
                if emb.partition:
                    # if the embedding is partitioned, map back to indexes
                    # into the local tensor
                    idx_in[emb_name] = partition.map_to_local(idx_in[emb_name])

            if self._clean_grad:
                # clean gradient track
                for emb in self._params:
                    emb.reset_trace()
                self._clean_grad = False

            for emb in self._params:
                emb_name = emb.name
                idx = idx_in[emb_name]
                grad = grad_in[emb_name]
                self.update(idx, grad, emb)

    def _shared_step(self):
        with th.no_grad():
            # Frequently alloc and free shared memory to hold intermediate tensor is expensive
            # We cache shared memory buffers in shared_emb.
            shared_emb = {emb.name: ([], []) for emb in self._params}

            # Go through all sparse embeddings
            for emb in self._params:  # pylint: disable=too-many-nested-blocks
                emb_name = emb.name

                # we need to combine gradients from multiple forward paths
                idx = []
                grad = []
                for i, data in emb._trace:
                    idx.append(i)
                    grad.append(data.grad.data)
                # 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.
                idx = (
                    th.cat(idx, dim=0)
                    if len(idx) != 0
                    else th.zeros((0,), dtype=th.long, device=th.device("cpu"))
                )
                grad = (
                    th.cat(grad, dim=0)
                    if len(grad) != 0
                    else th.zeros(
                        (0, emb.embedding_dim),
                        dtype=th.float32,
                        device=th.device("cpu"),
                    )
                )

                device = grad.device
                idx_dtype = idx.dtype
                grad_dtype = grad.dtype
                grad_dim = grad.shape[1]
                if self._world_size > 1:
                    if emb_name not in self._shared_cache:
                        self._shared_cache[emb_name] = {}

                    # Each training process takes the resposibility of updating a range
                    # of node embeddings, thus we can parallel the gradient update.
                    # The overall progress includes:
                    #   1. In each training process:
                    #     1.a Deciding which process a node embedding belongs to according
                    #         to the formula: process_id = node_idx mod num_of_process(N)
                    #     1.b Split the node index tensor and gradient tensor into N parts
                    #         according to step 1.
                    #     1.c Write each node index sub-tensor and gradient sub-tensor into
                    #         different DGL shared memory buffers.
                    #   2. Cross training process synchronization
                    #   3. In each traning process:
                    #     3.a Collect node index sub-tensors and gradient sub-tensors
                    #     3.b Do gradient update
                    #   4. Done
                    idx_split = th.remainder(idx, self._world_size).long()
                    for i in range(self._world_size):
                        mask = idx_split == i
                        idx_i = idx[mask]
                        grad_i = grad[mask]

                        if i == self._rank:
                            shared_emb[emb_name][0].append(idx_i)
                            shared_emb[emb_name][1].append(grad_i)
                        else:
                            # currently nccl does not support Alltoallv operation
                            # we need to use CPU shared memory to share gradient
                            # across processes
                            idx_i = idx_i.to(th.device("cpu"))
                            grad_i = grad_i.to(th.device("cpu"))
                            idx_shmem_name = "idx_{}_{}_{}".format(
                                emb_name, self._rank, i
                            )
                            grad_shmem_name = "grad_{}_{}_{}".format(
                                emb_name, self._rank, i
                            )

                            # Create shared memory to hold temporary index and gradient tensor for
                            # cross-process send and recv.
                            if (
                                idx_shmem_name
                                not in self._shared_cache[emb_name]
                                or self._shared_cache[emb_name][
                                    idx_shmem_name
                                ].shape[0]
                                < idx_i.shape[0]
                            ):

                                if (
                                    idx_shmem_name
                                    in self._shared_cache[emb_name]
                                ):
                                    self.shmem_buffer_holder.append(
                                        self._shared_cache[emb_name][
                                            idx_shmem_name
                                        ]
                                    )
                                    self.shmem_buffer_holder.append(
                                        self._shared_cache[emb_name][
                                            grad_shmem_name
                                        ]
                                    )

                                # The total number of buffers is the number of NodeEmbeddings *
                                # world_size * (world_size - 1). The minimun buffer size is 128.
                                #
                                # We extend the buffer by idx_i.shape[0] * 2 to avoid
                                # frequent shared memory allocation.
                                # The overall buffer cost will be smaller than three times
                                # the maximum memory requirement for sharing gradients.
                                buffer_size = (
                                    128
                                    if idx_i.shape[0] < 128
                                    else idx_i.shape[0] * 2
                                )
                                idx_shmem = create_shared_mem_array(
                                    "{}_{}".format(idx_shmem_name, buffer_size),
                                    (buffer_size,),
                                    idx_dtype,
                                )
                                grad_shmem = create_shared_mem_array(
                                    "{}_{}".format(
                                        grad_shmem_name, buffer_size
                                    ),
                                    (buffer_size, grad_dim),
                                    grad_dtype,
                                )
                                self._shared_cache[emb_name][
                                    idx_shmem_name
                                ] = idx_shmem
                                self._shared_cache[emb_name][
                                    grad_shmem_name
                                ] = grad_shmem

                            # Fill shared memory with temporal index tensor and gradient tensor
                            self._shared_cache[emb_name][idx_shmem_name][
                                : idx_i.shape[0]
                            ] = idx_i
                            self._shared_cache[emb_name][grad_shmem_name][
                                : idx_i.shape[0]
                            ] = grad_i
                            self._opt_meta[emb_name][self._rank][
                                i
                            ] = idx_i.shape[0]
                else:
                    shared_emb[emb_name][0].append(idx)
                    shared_emb[emb_name][1].append(grad)

            # make sure the idx shape is passed to each process through opt_meta
            if self._world_size > 1:
                th.distributed.barrier()
            for emb in self._params:  # pylint: disable=too-many-nested-blocks
                emb_name = emb.name
                if self._world_size > 1:
                    # The first element in shared_emb[emb_name][0] is the local idx
                    device = shared_emb[emb_name][0][0].device
                    # gather gradients from all other processes
                    for i in range(self._world_size):
                        if i != self._rank:
                            idx_shmem_name = "idx_{}_{}_{}".format(
                                emb_name, i, self._rank
                            )
                            grad_shmem_name = "grad_{}_{}_{}".format(
                                emb_name, i, self._rank
                            )
                            size = self._opt_meta[emb_name][i][self._rank]

                            # Retrive shared memory holding the temporal index and gradient
                            # tensor that is sent to current training process
                            if (
                                idx_shmem_name
                                not in self._shared_cache[emb_name]
                                or self._shared_cache[emb_name][
                                    idx_shmem_name
                                ].shape[0]
                                < size
                            ):
                                buffer_size = 128 if size < 128 else size * 2
                                idx_shmem = get_shared_mem_array(
                                    "{}_{}".format(idx_shmem_name, buffer_size),
                                    (buffer_size,),
                                    idx_dtype,
                                )
                                grad_shmem = get_shared_mem_array(
                                    "{}_{}".format(
                                        grad_shmem_name, buffer_size
                                    ),
                                    (buffer_size, grad_dim),
                                    grad_dtype,
                                )
                                self._shared_cache[emb_name][
                                    idx_shmem_name
                                ] = idx_shmem
                                self._shared_cache[emb_name][
                                    grad_shmem_name
                                ] = grad_shmem

                            idx_i = self._shared_cache[emb_name][
                                idx_shmem_name
                            ][:size]
                            grad_i = self._shared_cache[emb_name][
                                grad_shmem_name
                            ][:size]
                            shared_emb[emb_name][0].append(
                                idx_i.to(device, non_blocking=True)
                            )
                            shared_emb[emb_name][1].append(
                                grad_i.to(device, non_blocking=True)
                            )

            if self._clean_grad:
                # clean gradient track
                for emb in self._params:
                    emb.reset_trace()
                self._clean_grad = False

            for emb in self._params:
                emb_name = emb.name

                idx = th.cat(shared_emb[emb_name][0], dim=0)
                grad = th.cat(shared_emb[emb_name][1], dim=0)
                self.update(idx, grad, 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.nn.NodeEmbedding
            Sparse node embedding to update.
        """

    def zero_grad(self):
        """clean grad cache"""
        self._clean_grad = True

    def state_dict(self, **kwargs):  # pylint: disable=unused-argument
        """Return a copy of the whole optimizer states stored in CPU memory.
        If this is a multi-processing instance, the states will be returned in
        shared memory. If the underlying embedding is currently stored on
        multiple GPUs, all processes must call this method in the same order.

        NOTE: This method must be called by all processes sharing the
        underlying embedding, or it may result in a deadlock.

        Returns
        -------
        dictionary of optimizer states
            The optimizer states stored in CPU memory.
        """
        return {
            "state": {
                emb.name: emb._all_get_optm_state() for emb in self._params
            },
            "param_groups": self.param_groups,
        }

    def load_state_dict(
        self, state_dict, **kwargs
    ):  # pylint: disable=unused-argument
        """Load the optimizer states. This method must be called by all
        processes sharing the underlying embedding with identical
        :attr:`state_dict`.

        NOTE: This method must be called by all processes sharing the
        underlying embedding, or it may result in a deadlock.

        Parameters
        ----------
        state_dict : dictionary of optimizer states
            The global states to pull values from.
        """
        for emb in self._params:
            emb._all_set_optm_state(state_dict["state"][emb.name])
        self._set_param_groups(state_dict["param_groups"])

    @property
    @abstractmethod
    def param_groups(self):
        """Emulate 'param_groups' of torch.optim.Optimizer.
        Different from that, the returned 'param_groups' doesn't contain
        parameters because getting the whole embedding is very expensive.
        It contains other attributes, e.g., lr, eps, for debugging.
        """

    @abstractmethod
    def _set_param_groups(self, groups):
        """A helper method to load param_groups from saved state_dict."""




class SparseAdagrad(SparseGradOptimizer):
    r"""Node embedding optimizer using the Adagrad algorithm.

    This optimizer implements a sparse version of Adagrad algorithm for
    optimizing :class:`dgl.nn.NodeEmbedding`. 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.nn.NodeEmbedding]
        The list of dgl.nn.NodeEmbedding.
    lr : float
        The learning rate.
    eps : float, Optional
        The term added to the denominator to improve numerical stability
        Default: 1e-10

    Examples
    --------
    >>> def initializer(emb):
            th.nn.init.xavier_uniform_(emb)
            return emb
    >>> emb = dgl.nn.NodeEmbedding(g.num_nodes(), 10, 'emb', init_func=initializer)
    >>> optimizer = dgl.optim.SparseAdagrad([emb], lr=0.001)
    >>> for blocks in dataloader:
    ...     ...
    ...     feats = emb(nids, gpu_0)
    ...     loss = F.sum(feats + 1, 0)
    ...     loss.backward()
    ...     optimizer.step()
    """

    def __init__(self, params, lr, eps=1e-10):
        super(SparseAdagrad, self).__init__(params, lr)
        self._eps = eps

        # setup tensors for optimizer states
        self.setup(self._params)

    def setup(self, params):
        # We need to register a state sum for each embedding in the kvstore.
        for emb in params:
            assert isinstance(
                emb, NodeEmbedding
            ), "SparseAdagrad only supports dgl.nn.NodeEmbedding"

            emb_name = emb.name
            if th.device(emb.weight.device) == th.device("cpu"):
                # if our embedding is on the CPU, our state also has to be
                if self._rank < 0:
                    state = th.empty(
                        emb.weight.shape,
                        dtype=th.float32,
                        device=th.device("cpu"),
                    ).zero_()
                elif self._rank == 0:
                    state = create_shared_mem_array(
                        emb_name + "_state", emb.weight.shape, th.float32
                    ).zero_()

                    if self._world_size > 1:
                        emb.store.set(emb_name + "_opt", emb_name)
                elif self._rank > 0:
                    # receive
                    emb.store.wait([emb_name + "_opt"])
                    state = get_shared_mem_array(
                        emb_name + "_state", emb.weight.shape, th.float32
                    )
            else:
                # distributed state on on gpu
                state = th.empty(
                    emb.weight.shape,
                    dtype=th.float32,
                    device=emb.weight.device,
                ).zero_()
            emb.set_optm_state((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.nn.NodeEmbedding
            Sparse embedding to update.
        """
        eps = self._eps
        clr = self._lr

        # 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=grad.device
        )
        grad_values.index_add_(0, inverse, grad)
        grad_values = grad_values / cnt.unsqueeze(1)

        grad_sum = grad_values * grad_values
        (state,) = emb.optm_state
        state_dev = state.device
        state_idx = grad_indices.to(state_dev)
        grad_state = state[state_idx].to(grad.device)
        grad_state += grad_sum
        state[state_idx] = grad_state.to(state_dev)

        std_values = grad_state.add_(eps).sqrt_()
        tmp = clr * grad_values / std_values
        emb.weight[state_idx] -= tmp.to(state_dev)

    @property
    def param_groups(self):
        """Emulate 'param_groups' of torch.optim.Optimizer.
        Different from that, the returned 'param_groups' doesn't contain
        parameters because getting the whole embedding is very expensive.
        It contains other attributes, e.g., lr, eps, for debugging.
        """
        return [{"lr": self._lr, "eps": self._eps}]

    def _set_param_groups(self, groups):
        """A helper method to load param_groups from saved state_dict."""
        self._lr = groups[0]["lr"]
        self._eps = groups[0]["eps"]





class SparseAdam(SparseGradOptimizer):
    r"""Node embedding optimizer using the Adam algorithm.

    This optimizer implements a sparse version of Adagrad algorithm for
    optimizing :class:`dgl.nn.NodeEmbedding`. 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.nn.NodeEmbedding]
        The list of dgl.nn.NodeEmbeddings.
    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
    use_uva : bool, Optional
        Whether to use pinned memory for storing 'mem' and 'power' parameters,
        when the embedding is stored on the CPU. This will improve training
        speed, but will require locking a large number of virtual memory pages.
        For embeddings which are stored in GPU memory, this setting will have
        no effect.
        Default: True if the gradients are generated on the GPU, and False
        if the gradients are on the CPU.
    dtype : torch.dtype, Optional
        The type to store optimizer state with. Default: th.float32.

    Examples
    --------
    >>> def initializer(emb):
            th.nn.init.xavier_uniform_(emb)
            return emb
    >>> emb = dgl.nn.NodeEmbedding(g.num_nodes(), 10, 'emb', init_func=initializer)
    >>> optimizer = dgl.optim.SparseAdam([emb], lr=0.001)
    >>> for blocks in dataloader:
    ...     ...
    ...     feats = emb(nids, gpu_0)
    ...     loss = F.sum(feats + 1, 0)
    ...     loss.backward()
    ...     optimizer.step()
    """

    def __init__(
        self,
        params,
        lr,
        betas=(0.9, 0.999),
        eps=1e-08,
        use_uva=None,
        dtype=th.float32,
    ):
        super(SparseAdam, self).__init__(params, lr)
        self._lr = lr
        self._beta1 = betas[0]
        self._beta2 = betas[1]
        self._eps = eps
        self._use_uva = use_uva
        self._nd_handle = {}
        self._is_using_uva = {}
        assert dtype in [th.float16, th.float32], (
            "Unsupported dtype {}. Valid choices are th.float32 "
            "and th.float32".format(dtype)
        )
        self._dtype = dtype

        # setup tensors for optimizer states
        self.setup(self._params)

    def _setup_uva(self, name, mem, power):
        self._is_using_uva[name] = True
        mem_nd = pin_memory_inplace(mem)
        power_nd = pin_memory_inplace(power)
        self._nd_handle[name] = [mem_nd, power_nd]

    def setup(self, params):
        # We need to register a state sum for each embedding in the kvstore.
        for emb in params:
            assert isinstance(
                emb, NodeEmbedding
            ), "SparseAdam only supports dgl.nn.NodeEmbedding"
            emb_name = emb.name
            self._is_using_uva[emb_name] = self._use_uva
            if th.device(emb.weight.device) == th.device("cpu"):
                # if our embedding is on the CPU, our state also has to be
                if self._rank < 0:
                    state_step = th.empty(
                        (emb.weight.shape[0],),
                        dtype=th.int32,
                        device=th.device("cpu"),
                    ).zero_()
                    state_mem = th.empty(
                        emb.weight.shape,
                        dtype=self._dtype,
                        device=th.device("cpu"),
                    ).zero_()
                    state_power = th.empty(
                        emb.weight.shape,
                        dtype=self._dtype,
                        device=th.device("cpu"),
                    ).zero_()
                elif self._rank == 0:
                    state_step = create_shared_mem_array(
                        emb_name + "_step", (emb.weight.shape[0],), th.int32
                    ).zero_()
                    state_mem = create_shared_mem_array(
                        emb_name + "_mem", emb.weight.shape, self._dtype
                    ).zero_()
                    state_power = create_shared_mem_array(
                        emb_name + "_power", emb.weight.shape, self._dtype
                    ).zero_()

                    if self._world_size > 1:
                        emb.store.set(emb_name + "_opt", emb_name)
                elif self._rank > 0:
                    # receive
                    emb.store.wait([emb_name + "_opt"])
                    state_step = get_shared_mem_array(
                        emb_name + "_step", (emb.weight.shape[0],), th.int32
                    )
                    state_mem = get_shared_mem_array(
                        emb_name + "_mem", emb.weight.shape, self._dtype
                    )
                    state_power = get_shared_mem_array(
                        emb_name + "_power", emb.weight.shape, self._dtype
                    )

                if self._is_using_uva[emb_name]:
                    # if use_uva has been explicitly set to true, otherwise
                    # wait until first step to decide
                    self._setup_uva(emb_name, state_mem, state_power)
            else:
                # make sure we don't use UVA when data is on the GPU
                self._is_using_uva[emb_name] = False

                # distributed state on on gpu
                state_step = th.empty(
                    [emb.weight.shape[0]],
                    dtype=th.int32,
                    device=emb.weight.device,
                ).zero_()
                state_mem = th.empty(
                    emb.weight.shape,
                    dtype=self._dtype,
                    device=emb.weight.device,
                ).zero_()
                state_power = th.empty(
                    emb.weight.shape,
                    dtype=self._dtype,
                    device=emb.weight.device,
                ).zero_()
            state = (state_step, state_mem, state_power)
            emb.set_optm_state(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.nn.NodeEmbedding
            Sparse embedding to update.
        """
        with th.no_grad():
            state_step, state_mem, state_power = emb.optm_state
            exec_dtype = grad.dtype
            exec_dev = grad.device
            state_dev = state_step.device

            # whether or not we need to transfer data from the GPU to the CPU
            # while updating the weights
            is_d2h = state_dev.type == "cpu" and exec_dev.type == "cuda"

            # 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 = is_d2h

            if self._is_using_uva[emb.name] is None and is_d2h:
                # we should use UVA going forward
                self._setup_uva(emb.name, state_mem, state_power)
            elif self._is_using_uva[emb.name] is None:
                # we shouldn't use UVA going forward
                self._is_using_uva[emb.name] = False

            use_uva = self._is_using_uva[emb.name]

            beta1 = self._beta1
            beta2 = self._beta2
            eps = self._eps

            clr = self._lr
            # There can be duplicated indices due to sampling.
            # Thus unique them here and average the gradient here.
            grad_indices, inverse, cnt = th.unique(
                idx, return_inverse=True, return_counts=True
            )
            state_idx = grad_indices.to(state_dev)
            state_step[state_idx] += 1
            state_step = state_step[state_idx].to(exec_dev)

            if use_uva:
                orig_mem = gather_pinned_tensor_rows(state_mem, grad_indices)
                orig_power = gather_pinned_tensor_rows(
                    state_power, grad_indices
                )
            else:
                orig_mem = state_mem[state_idx].to(exec_dev)
                orig_power = state_power[state_idx].to(exec_dev)
            # convert to exec dtype
            orig_mem = orig_mem.to(dtype=exec_dtype)
            orig_power = orig_power.to(dtype=exec_dtype)

            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

            if use_uva:
                scatter_pinned_tensor_rows(
                    state_mem, grad_indices, update_mem.to(dtype=self._dtype)
                )
                scatter_pinned_tensor_rows(
                    state_power,
                    grad_indices,
                    update_power.to(dtype=self._dtype),
                )
            else:
                update_mem_dst = update_mem.to(dtype=self._dtype).to(
                    state_dev, non_blocking=True
                )
                update_power_dst = update_power.to(dtype=self._dtype).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()

            if not use_uva:
                if state_block:
                    # 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.weight[state_idx] -= std_values_dst

    @property
    def param_groups(self):
        """Emulate 'param_groups' of torch.optim.Optimizer.
        Different from that, the returned 'param_groups' doesn't contain
        parameters because getting the whole embedding is very expensive.
        It contains other attributes, e.g., lr, betas, eps, for debugging.
        """
        return [
            {
                "lr": self._lr,
                "betas": (self._beta1, self._beta2),
                "eps": self._eps,
            }
        ]

    def _set_param_groups(self, groups):
        """A helper method to load param_groups from saved state_dict."""
        self._lr = groups[0]["lr"]
        self._beta1, self._beta2 = groups[0]["betas"]
        self._eps = groups[0]["eps"]