chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,30 @@
+"""DGL distributed module"""
+
+from . import optim
+from .dist_context import exit_client, initialize
+from .dist_dataloader import (
+    DistDataLoader,
+    DistEdgeDataLoader,
+    DistNodeDataLoader,
+    EdgeCollator,
+    NodeCollator,
+)
+from .dist_graph import DistGraph, DistGraphServer, edge_split, node_split
+from .dist_tensor import DistTensor
+from .graph_partition_book import GraphPartitionBook, PartitionPolicy
+from .graph_services import *
+from .kvstore import KVClient, KVServer
+from .nn import *
+from .partition import (
+    dgl_partition_to_graphbolt,
+    gb_convert_single_dgl_partition,
+    load_partition,
+    load_partition_book,
+    load_partition_feats,
+    partition_graph,
+)
+from .rpc import *
+from .rpc_client import connect_to_server
+from .rpc_server import start_server
+from .server_state import ServerState
+from .constants import *
@@ -0,0 +1,12 @@
+"""Define all the constants used by DGL rpc"""
+
+# Maximum size of message queue in bytes
+MAX_QUEUE_SIZE = 20 * 1024 * 1024 * 1024
+
+SERVER_EXIT = "server_exit"
+
+DEFAULT_NTYPE = "_N"
+DEFAULT_ETYPE = (DEFAULT_NTYPE, "_E", DEFAULT_NTYPE)
+
+DGL2GB_EID = "_dgl2gb_eid"
+GB_DST_ID = "_gb_dst_id"
@@ -0,0 +1,390 @@
+"""Initialize the distributed services"""
+# pylint: disable=line-too-long
+
+import atexit
+import gc
+import multiprocessing as mp
+import os
+import queue
+import sys
+import time
+import traceback
+from enum import Enum
+
+from .. import utils
+from ..base import dgl_warning, DGLError
+from . import rpc
+from .constants import MAX_QUEUE_SIZE
+from .kvstore import close_kvstore, init_kvstore
+from .role import init_role
+from .rpc_client import connect_to_server
+
+SAMPLER_POOL = None
+NUM_SAMPLER_WORKERS = 0
+INITIALIZED = False
+
+
+def set_initialized(value=True):
+    """Set the initialized state of rpc"""
+    global INITIALIZED
+    INITIALIZED = value
+
+
+def get_sampler_pool():
+    """Return the sampler pool and num_workers"""
+    return SAMPLER_POOL, NUM_SAMPLER_WORKERS
+
+
+def _init_rpc(
+    ip_config,
+    num_servers,
+    max_queue_size,
+    role,
+    num_threads,
+    group_id,
+):
+    """This init function is called in the worker processes."""
+    try:
+        utils.set_num_threads(num_threads)
+        if os.environ.get("DGL_DIST_MODE", "standalone") != "standalone":
+            connect_to_server(ip_config, num_servers, max_queue_size, group_id)
+        init_role(role)
+        init_kvstore(ip_config, num_servers, role)
+    except Exception as e:
+        print(e, flush=True)
+        traceback.print_exc()
+        raise e
+
+
+class MpCommand(Enum):
+    """Enum class for multiprocessing command"""
+
+    INIT_RPC = 0  # Not used in the task queue
+    SET_COLLATE_FN = 1
+    CALL_BARRIER = 2
+    DELETE_COLLATE_FN = 3
+    CALL_COLLATE_FN = 4
+    CALL_FN_ALL_WORKERS = 5
+    FINALIZE_POOL = 6
+
+
+def init_process(rpc_config, mp_contexts):
+    """Work loop in the worker"""
+    try:
+        _init_rpc(*rpc_config)
+        keep_polling = True
+        data_queue, task_queue, barrier = mp_contexts
+        collate_fn_dict = {}
+
+        while keep_polling:
+            try:
+                # Follow https://github.com/pytorch/pytorch/blob/d57ce8cf8989c0b737e636d8d7abe16c1f08f70b/torch/utils/data/_utils/worker.py#L260
+                command, args = task_queue.get(timeout=5)
+            except queue.Empty:
+                continue
+            if command == MpCommand.SET_COLLATE_FN:
+                dataloader_name, func = args
+                collate_fn_dict[dataloader_name] = func
+            elif command == MpCommand.CALL_BARRIER:
+                barrier.wait()
+            elif command == MpCommand.DELETE_COLLATE_FN:
+                (dataloader_name,) = args
+                del collate_fn_dict[dataloader_name]
+            elif command == MpCommand.CALL_COLLATE_FN:
+                dataloader_name, collate_args = args
+                data_queue.put(
+                    (
+                        dataloader_name,
+                        collate_fn_dict[dataloader_name](collate_args),
+                    )
+                )
+            elif command == MpCommand.CALL_FN_ALL_WORKERS:
+                func, func_args = args
+                func(func_args)
+            elif command == MpCommand.FINALIZE_POOL:
+                _exit()
+                keep_polling = False
+            else:
+                raise Exception("Unknown command")
+    except Exception as e:
+        traceback.print_exc()
+        raise e
+
+
+class CustomPool:
+    """Customized worker pool"""
+
+    def __init__(self, num_workers, rpc_config):
+        """
+        Customized worker pool init function
+        """
+        ctx = mp.get_context("spawn")
+        self.num_workers = num_workers
+        # As pool could be used by any number of dataloaders, queues
+        # should be able to take infinite elements to avoid dead lock.
+        self.queue_size = 0
+        self.result_queue = ctx.Queue(self.queue_size)
+        self.results = {}  # key is dataloader name, value is fetched batch.
+        self.task_queues = []
+        self.process_list = []
+        self.current_proc_id = 0
+        self.cache_result_dict = {}
+        self.barrier = ctx.Barrier(num_workers)
+        for _ in range(num_workers):
+            task_queue = ctx.Queue(self.queue_size)
+            self.task_queues.append(task_queue)
+            proc = ctx.Process(
+                target=init_process,
+                args=(
+                    rpc_config,
+                    (self.result_queue, task_queue, self.barrier),
+                ),
+            )
+            proc.daemon = True
+            proc.start()
+            self.process_list.append(proc)
+
+    def set_collate_fn(self, func, dataloader_name):
+        """Set collate function in subprocess"""
+        for i in range(self.num_workers):
+            self.task_queues[i].put(
+                (MpCommand.SET_COLLATE_FN, (dataloader_name, func))
+            )
+        self.results[dataloader_name] = []
+
+    def submit_task(self, dataloader_name, args):
+        """Submit task to workers"""
+        # Round robin
+        self.task_queues[self.current_proc_id].put(
+            (MpCommand.CALL_COLLATE_FN, (dataloader_name, args))
+        )
+        self.current_proc_id = (self.current_proc_id + 1) % self.num_workers
+
+    def submit_task_to_all_workers(self, func, args):
+        """Submit task to all workers"""
+        for i in range(self.num_workers):
+            self.task_queues[i].put(
+                (MpCommand.CALL_FN_ALL_WORKERS, (func, args))
+            )
+
+    def get_result(self, dataloader_name, timeout=1800):
+        """Get result from result queue"""
+        if dataloader_name not in self.results:
+            raise DGLError(
+                f"Got result from an unknown dataloader {dataloader_name}."
+            )
+        while len(self.results[dataloader_name]) == 0:
+            dl_name, data = self.result_queue.get(timeout=timeout)
+            self.results[dl_name].append(data)
+        return self.results[dataloader_name].pop(0)
+
+    def delete_collate_fn(self, dataloader_name):
+        """Delete collate function"""
+        for i in range(self.num_workers):
+            self.task_queues[i].put(
+                (MpCommand.DELETE_COLLATE_FN, (dataloader_name,))
+            )
+        del self.results[dataloader_name]
+
+    def call_barrier(self):
+        """Call barrier at all workers"""
+        for i in range(self.num_workers):
+            self.task_queues[i].put((MpCommand.CALL_BARRIER, tuple()))
+
+    def close(self):
+        """Close worker pool"""
+        for i in range(self.num_workers):
+            self.task_queues[i].put(
+                (MpCommand.FINALIZE_POOL, tuple()), block=False
+            )
+            time.sleep(0.5)  # Fix for early python version
+
+    def join(self):
+        """Join the close process of worker pool"""
+        for i in range(self.num_workers):
+            self.process_list[i].join()
+
+
+def initialize(
+    ip_config,
+    max_queue_size=MAX_QUEUE_SIZE,
+    net_type=None,
+    num_worker_threads=1,
+    use_graphbolt=False,
+):
+    """Initialize DGL's distributed module
+
+    This function initializes DGL's distributed module. It acts differently in server
+    or client modes. In the server mode, it runs the server code and never returns.
+    In the client mode, it builds connections with servers for communication and
+    creates worker processes for distributed sampling.
+
+    Parameters
+    ----------
+    ip_config: str
+        File path of ip_config file
+    max_queue_size : int
+        Maximal size (bytes) of client queue buffer (~20 GB on default).
+
+        Note that the 20 GB is just an upper-bound and DGL uses zero-copy and
+        it will not allocate 20GB memory at once.
+    net_type : str, optional
+        [Deprecated] Networking type, can be 'socket' only.
+    num_worker_threads: int
+        The number of OMP threads in each sampler process.
+    use_graphbolt: bool, optional
+        Whether to use GraphBolt for distributed train.
+
+    Note
+    ----
+    Users have to invoke this API before any DGL's distributed API and framework-specific
+    distributed API. For example, when used with Pytorch, users have to invoke this function
+    before Pytorch's `pytorch.distributed.init_process_group`.
+    """
+    print(
+        f"Initialize the distributed services with graphbolt: {use_graphbolt}"
+    )
+    if net_type is not None:
+        dgl_warning(
+            "net_type is deprecated and will be removed in future release."
+        )
+    if os.environ.get("DGL_ROLE", "client") == "server":
+        from .dist_graph import DistGraphServer
+
+        assert (
+            os.environ.get("DGL_SERVER_ID") is not None
+        ), "Please define DGL_SERVER_ID to run DistGraph server"
+        assert (
+            os.environ.get("DGL_IP_CONFIG") is not None
+        ), "Please define DGL_IP_CONFIG to run DistGraph server"
+        assert (
+            os.environ.get("DGL_NUM_SERVER") is not None
+        ), "Please define DGL_NUM_SERVER to run DistGraph server"
+        assert (
+            os.environ.get("DGL_NUM_CLIENT") is not None
+        ), "Please define DGL_NUM_CLIENT to run DistGraph server"
+        assert (
+            os.environ.get("DGL_CONF_PATH") is not None
+        ), "Please define DGL_CONF_PATH to run DistGraph server"
+        formats = os.environ.get("DGL_GRAPH_FORMAT", "csc").split(",")
+        formats = [f.strip() for f in formats]
+        rpc.reset()
+        serv = DistGraphServer(
+            int(os.environ.get("DGL_SERVER_ID")),
+            os.environ.get("DGL_IP_CONFIG"),
+            int(os.environ.get("DGL_NUM_SERVER")),
+            int(os.environ.get("DGL_NUM_CLIENT")),
+            os.environ.get("DGL_CONF_PATH"),
+            graph_format=formats,
+            use_graphbolt=use_graphbolt,
+        )
+        serv.start()
+        sys.exit()
+    else:
+        num_workers = int(os.environ.get("DGL_NUM_SAMPLER", 0))
+        num_servers = int(os.environ.get("DGL_NUM_SERVER", 1))
+        group_id = int(os.environ.get("DGL_GROUP_ID", 0))
+        rpc.reset()
+        global SAMPLER_POOL
+        global NUM_SAMPLER_WORKERS
+        is_standalone = (
+            os.environ.get("DGL_DIST_MODE", "standalone") == "standalone"
+        )
+        if num_workers > 0 and not is_standalone:
+            SAMPLER_POOL = CustomPool(
+                num_workers,
+                (
+                    ip_config,
+                    num_servers,
+                    max_queue_size,
+                    "sampler",
+                    num_worker_threads,
+                    group_id,
+                ),
+            )
+        else:
+            SAMPLER_POOL = None
+        NUM_SAMPLER_WORKERS = num_workers
+        if not is_standalone:
+            assert (
+                num_servers is not None and num_servers > 0
+            ), "The number of servers per machine must be specified with a positive number."
+            connect_to_server(
+                ip_config,
+                num_servers,
+                max_queue_size,
+                group_id=group_id,
+            )
+        init_role("default")
+        init_kvstore(ip_config, num_servers, "default")
+
+
+def finalize_client():
+    """Release resources of this client."""
+    if os.environ.get("DGL_DIST_MODE", "standalone") != "standalone":
+        rpc.finalize_sender()
+        rpc.finalize_receiver()
+
+
+def _exit():
+    exit_client()
+    time.sleep(1)
+
+
+def finalize_worker():
+    """Finalize workers
+    Python's multiprocessing pool will not call atexit function when close
+    """
+    global SAMPLER_POOL
+    if SAMPLER_POOL is not None:
+        SAMPLER_POOL.close()
+
+
+def join_finalize_worker():
+    """join the worker close process"""
+    global SAMPLER_POOL
+    if SAMPLER_POOL is not None:
+        SAMPLER_POOL.join()
+    SAMPLER_POOL = None
+
+
+def is_initialized():
+    """Is RPC initialized?"""
+    return INITIALIZED
+
+
+def _shutdown_servers():
+    set_initialized(False)
+    # send ShutDownRequest to servers
+    if rpc.get_rank() == 0:  # Only client_0 issue this command
+        req = rpc.ShutDownRequest(rpc.get_rank())
+        for server_id in range(rpc.get_num_server()):
+            rpc.send_request(server_id, req)
+
+
+def exit_client():
+    """Trainer exits
+
+    This function is called automatically when a Python process exits. Normally,
+    the training script does not need to invoke this function at the end.
+
+    In the case that the training script needs to initialize the distributed module
+    multiple times (so far, this is needed in the unit tests), the training script
+    needs to call `exit_client` before calling `initialize` again.
+    """
+    # Only client with rank_0 will send shutdown request to servers.
+    print(
+        "Client[{}] in group[{}] is exiting...".format(
+            rpc.get_rank(), rpc.get_group_id()
+        )
+    )
+    finalize_worker()  # finalize workers should be earilier than barrier, and non-blocking
+    # collect data such as DistTensor before exit
+    gc.collect()
+    if os.environ.get("DGL_DIST_MODE", "standalone") != "standalone":
+        rpc.client_barrier()
+        _shutdown_servers()
+    finalize_client()
+    join_finalize_worker()
+    close_kvstore()
+    atexit.unregister(exit_client)
@@ -0,0 +1,894 @@
+# pylint: disable=global-variable-undefined, invalid-name
+"""Multiprocess dataloader for distributed training"""
+import inspect
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+
+from .. import backend as F, transforms, utils
+from ..base import EID, NID
+from ..convert import heterograph
+from .dist_context import get_sampler_pool
+
+__all__ = [
+    "NodeCollator",
+    "EdgeCollator",
+    "DistDataLoader",
+    "DistNodeDataLoader",
+    "DistEdgeDataLoader",
+]
+
+DATALOADER_ID = 0
+
+
+class DistDataLoader:
+    """DGL customized multiprocessing dataloader.
+
+    DistDataLoader provides a similar interface to Pytorch's DataLoader to generate mini-batches
+    with multiprocessing. It utilizes the worker processes created by
+    :func:`dgl.distributed.initialize` to parallelize sampling.
+
+    Parameters
+    ----------
+    dataset: a tensor
+        Tensors of node IDs or edge IDs.
+    batch_size: int
+        The number of samples per batch to load.
+    shuffle: bool, optional
+        Set to ``True`` to have the data reshuffled at every epoch (default: ``False``).
+    collate_fn: callable, optional
+        The function is typically used to sample neighbors of the nodes in a batch
+        or the endpoint nodes of the edges in a batch.
+    drop_last: bool, optional
+        Set to ``True`` to drop the last incomplete batch, if the dataset size is not
+        divisible by the batch size. If ``False`` and the size of dataset is not divisible
+        by the batch size, then the last batch will be smaller. (default: ``False``)
+    queue_size: int, optional
+        Size of multiprocessing queue
+
+    Examples
+    --------
+    >>> g = dgl.distributed.DistGraph('graph-name')
+    >>> def sample(seeds):
+    ...     seeds = th.LongTensor(np.asarray(seeds))
+    ...     frontier = dgl.distributed.sample_neighbors(g, seeds, 10)
+    ...     return dgl.to_block(frontier, seeds)
+    >>> dataloader = dgl.distributed.DistDataLoader(dataset=nodes, batch_size=1000,
+                                                    collate_fn=sample, shuffle=True)
+    >>> for block in dataloader:
+    ...     feat = g.ndata['features'][block.srcdata[dgl.NID]]
+    ...     labels = g.ndata['labels'][block.dstdata[dgl.NID]]
+    ...     pred = model(block, feat)
+
+    Note
+    ----
+    When performing DGL's distributed sampling with multiprocessing, users have to use this class
+    instead of Pytorch's DataLoader because DGL's RPC requires that all processes establish
+    connections with servers before invoking any DGL's distributed API. Therefore, this dataloader
+    uses the worker processes created in :func:`dgl.distributed.initialize`.
+
+    Note
+    ----
+    This dataloader does not guarantee the iteration order. For example,
+    if dataset = [1, 2, 3, 4], batch_size = 2 and shuffle = False, the order of [1, 2]
+    and [3, 4] is not guaranteed.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        shuffle=False,
+        collate_fn=None,
+        drop_last=False,
+        queue_size=None,
+    ):
+        self.pool, self.num_workers = get_sampler_pool()
+        if queue_size is None:
+            queue_size = self.num_workers * 4 if self.num_workers > 0 else 4
+        self.queue_size = queue_size  # prefetch size
+        self.batch_size = batch_size
+        self.num_pending = 0
+        self.collate_fn = collate_fn
+        self.current_pos = 0
+        self.queue = []  # Only used when pool is None
+        self.drop_last = drop_last
+        self.recv_idxs = 0
+        self.shuffle = shuffle
+        self.is_closed = False
+
+        self.dataset = dataset
+        self.data_idx = F.arange(0, len(dataset))
+        self.expected_idxs = len(dataset) // self.batch_size
+        if not self.drop_last and len(dataset) % self.batch_size != 0:
+            self.expected_idxs += 1
+
+        # We need to have a unique ID for each data loader to identify itself
+        # in the sampler processes.
+        global DATALOADER_ID
+        self.name = "dataloader-" + str(DATALOADER_ID)
+        DATALOADER_ID += 1
+
+        if self.pool is not None:
+            self.pool.set_collate_fn(self.collate_fn, self.name)
+
+    def __del__(self):
+        # When the process exits, the process pool may have been closed. We should try
+        # and get the process pool again and see if we need to clean up the process pool.
+        self.pool, self.num_workers = get_sampler_pool()
+        if self.pool is not None:
+            self.pool.delete_collate_fn(self.name)
+
+    def __next__(self):
+        if self.pool is None:
+            num_reqs = 1
+        else:
+            num_reqs = self.queue_size - self.num_pending
+        for _ in range(num_reqs):
+            self._request_next_batch()
+        if self.recv_idxs < self.expected_idxs:
+            result = self._get_data_from_result_queue()
+            self.recv_idxs += 1
+            self.num_pending -= 1
+            return result
+        else:
+            assert self.num_pending == 0
+            raise StopIteration
+
+    def _get_data_from_result_queue(self, timeout=1800):
+        if self.pool is None:
+            ret = self.queue.pop(0)
+        else:
+            ret = self.pool.get_result(self.name, timeout=timeout)
+        return ret
+
+    def __iter__(self):
+        if self.shuffle:
+            self.data_idx = F.rand_shuffle(self.data_idx)
+        self.recv_idxs = 0
+        self.current_pos = 0
+        self.num_pending = 0
+        return self
+
+    def _request_next_batch(self):
+        next_data = self._next_data()
+        if next_data is None:
+            return
+        elif self.pool is not None:
+            self.pool.submit_task(self.name, next_data)
+        else:
+            result = self.collate_fn(next_data)
+            self.queue.append(result)
+        self.num_pending += 1
+
+    def _next_data(self):
+        if self.current_pos == len(self.dataset):
+            return None
+
+        end_pos = 0
+        if self.current_pos + self.batch_size > len(self.dataset):
+            if self.drop_last:
+                return None
+            else:
+                end_pos = len(self.dataset)
+        else:
+            end_pos = self.current_pos + self.batch_size
+        idx = self.data_idx[self.current_pos : end_pos].tolist()
+        ret = [self.dataset[i] for i in idx]
+        # Sharing large number of tensors between processes will consume too many
+        # file descriptors, so let's convert each tensor to scalar value beforehand.
+        if isinstance(ret[0], tuple):
+            ret = [(type, F.as_scalar(id)) for (type, id) in ret]
+        else:
+            ret = [F.as_scalar(id) for id in ret]
+        self.current_pos = end_pos
+        return ret
+
+
+# [Note] As implementation of ``dgl.distributed.DistDataLoader`` is independent
+# of ``dgl.dataloading.DataLoader`` currently, dedicated collators are defined
+# here instead of using ``dgl.dataloading.CollateWrapper``.
+
+
+def _find_exclude_eids_with_reverse_id(g, eids, reverse_eid_map):
+    if isinstance(eids, Mapping):
+        eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
+        exclude_eids = {
+            k: F.cat([v, F.gather_row(reverse_eid_map[k], v)], 0)
+            for k, v in eids.items()
+        }
+    else:
+        exclude_eids = F.cat([eids, F.gather_row(reverse_eid_map, eids)], 0)
+    return exclude_eids
+
+
+def _find_exclude_eids_with_reverse_types(g, eids, reverse_etype_map):
+    exclude_eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
+    reverse_etype_map = {
+        g.to_canonical_etype(k): g.to_canonical_etype(v)
+        for k, v in reverse_etype_map.items()
+    }
+    exclude_eids.update(
+        {reverse_etype_map[k]: v for k, v in exclude_eids.items()}
+    )
+    return exclude_eids
+
+
+def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
+    """Find all edge IDs to exclude according to :attr:`exclude_mode`.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    exclude_mode : str, optional
+        Can be either of the following,
+
+        None (default)
+            Does not exclude any edge.
+
+        'self'
+            Exclude the given edges themselves but nothing else.
+
+        'reverse_id'
+            Exclude all edges specified in ``eids``, as well as their reverse edges
+            of the same edge type.
+
+            The mapping from each edge ID to its reverse edge ID is specified in
+            the keyword argument ``reverse_eid_map``.
+
+            This mode assumes that the reverse of an edge with ID ``e`` and type
+            ``etype`` will have ID ``reverse_eid_map[e]`` and type ``etype``.
+
+        'reverse_types'
+            Exclude all edges specified in ``eids``, as well as their reverse
+            edges of the corresponding edge types.
+
+            The mapping from each edge type to its reverse edge type is specified
+            in the keyword argument ``reverse_etype_map``.
+
+            This mode assumes that the reverse of an edge with ID ``e`` and type ``etype``
+            will have ID ``e`` and type ``reverse_etype_map[etype]``.
+    eids : Tensor or dict[etype, Tensor]
+        The edge IDs.
+    reverse_eid_map : Tensor or dict[etype, Tensor]
+        The mapping from edge ID to its reverse edge ID.
+    reverse_etype_map : dict[etype, etype]
+        The mapping from edge etype to its reverse edge type.
+    """
+    if exclude_mode is None:
+        return None
+    elif exclude_mode == "self":
+        if isinstance(eids, Mapping):
+            eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
+        return eids
+    elif exclude_mode == "reverse_id":
+        return _find_exclude_eids_with_reverse_id(
+            g, eids, kwargs["reverse_eid_map"]
+        )
+    elif exclude_mode == "reverse_types":
+        return _find_exclude_eids_with_reverse_types(
+            g, eids, kwargs["reverse_etype_map"]
+        )
+    else:
+        raise ValueError("unsupported mode {}".format(exclude_mode))
+
+
+class Collator(ABC):
+    """Abstract DGL collator for training GNNs on downstream tasks stochastically.
+
+    Provides a :attr:`dataset` object containing the collection of all nodes or edges,
+    as well as a :attr:`collate` method that combines a set of items from
+    :attr:`dataset` and obtains the message flow graphs (MFGs).
+
+    Notes
+    -----
+    For the concept of MFGs, please refer to
+    :ref:`User Guide Section 6 <guide-minibatch>` and
+    :doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
+    """
+
+    @property
+    @abstractmethod
+    def dataset(self):
+        """Returns the dataset object of the collator."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def collate(self, items):
+        """Combines the items from the dataset object and obtains the list of MFGs.
+
+        Parameters
+        ----------
+        items : list[str, int]
+            The list of node or edge IDs or type-ID pairs.
+
+        Notes
+        -----
+        For the concept of MFGs, please refer to
+        :ref:`User Guide Section 6 <guide-minibatch>` and
+        :doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def add_edge_attribute_to_graph(g, data_name, gb_padding):
+        """Add data into the graph as an edge attribute.
+
+        For some cases such as prob/mask-based sampling on GraphBolt partitions,
+        we need to prepare such data beforehand. This is because data are
+        usually saved in DistGraph.ndata/edata, but such data is not in the
+        format that GraphBolt partitions require. And in GraphBolt, such data
+        are saved as edge attributes. So we need to add such data into the graph
+        before any sampling is kicked off.
+
+        Parameters
+        ----------
+        g : DistGraph
+            The graph.
+        data_name : str
+            The name of data that's stored in DistGraph.ndata/edata.
+        gb_padding : int, optional
+            The padding value for GraphBolt partitions' new edge_attributes.
+        """
+        if g._use_graphbolt and data_name:
+            g.add_edge_attribute(data_name, gb_padding)
+
+
+class NodeCollator(Collator):
+    """DGL collator to combine nodes and their computation dependencies within a minibatch for
+    training node classification or regression on a single graph with neighborhood sampling.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    nids : Tensor or dict[ntype, Tensor]
+        The node set to compute outputs.
+    graph_sampler : dgl.dataloading.BlockSampler
+        The neighborhood sampler.
+    gb_padding : int, optional
+        The padding value for GraphBolt partitions' new edge_attributes if the attributes in DistGraph are None.
+        e.g. prob/mask-based sampling.
+        Only when the mask of one edge is set as 1, an edge will be sampled in dgl.graphbolt.FusedCSCSamplingGraph.sample_neighbors.
+        The argument will be used in add_edge_attribute_to_graph to add new edge_attributes in graphbolt.
+
+    Examples
+    --------
+    To train a 3-layer GNN for node classification on a set of nodes ``train_nid`` on
+    a homogeneous graph where each node takes messages from all neighbors (assume
+    the backend is PyTorch):
+
+    >>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
+    >>> collator = dgl.dataloading.NodeCollator(g, train_nid, sampler)
+    >>> dataloader = torch.utils.data.DataLoader(
+    ...     collator.dataset, collate_fn=collator.collate,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     train_on(input_nodes, output_nodes, blocks)
+
+    Notes
+    -----
+    For the concept of MFGs, please refer to
+    :ref:`User Guide Section 6 <guide-minibatch>` and
+    :doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
+    """
+
+    def __init__(self, g, nids, graph_sampler, gb_padding=1):
+        self.g = g
+        if not isinstance(nids, Mapping):
+            assert (
+                len(g.ntypes) == 1
+            ), "nids should be a dict of node type and ids for graph with multiple node types"
+        self.graph_sampler = graph_sampler
+
+        self.nids = utils.prepare_tensor_or_dict(g, nids, "nids")
+        self._dataset = utils.maybe_flatten_dict(self.nids)
+
+        # Add prob/mask into graphbolt partition's edge attributes if needed.
+        if hasattr(self.graph_sampler, "prob"):
+            Collator.add_edge_attribute_to_graph(
+                self.g, self.graph_sampler.prob, gb_padding
+            )
+
+    @property
+    def dataset(self):
+        return self._dataset
+
+    def collate(self, items):
+        """Find the list of MFGs necessary for computing the representation of given
+        nodes for a node classification/regression task.
+
+        Parameters
+        ----------
+        items : list[int] or list[tuple[str, int]]
+            Either a list of node IDs (for homogeneous graphs), or a list of node type-ID
+            pairs (for heterogeneous graphs).
+
+        Returns
+        -------
+        input_nodes : Tensor or dict[ntype, Tensor]
+            The input nodes necessary for computation in this minibatch.
+
+            If the original graph has multiple node types, return a dictionary of
+            node type names and node ID tensors.  Otherwise, return a single tensor.
+        output_nodes : Tensor or dict[ntype, Tensor]
+            The nodes whose representations are to be computed in this minibatch.
+
+            If the original graph has multiple node types, return a dictionary of
+            node type names and node ID tensors.  Otherwise, return a single tensor.
+        MFGs : list[DGLGraph]
+            The list of MFGs necessary for computing the representation.
+        """
+        if isinstance(items[0], tuple):
+            # returns a list of pairs: group them by node types into a dict
+            items = utils.group_as_dict(items)
+        items = utils.prepare_tensor_or_dict(self.g, items, "items")
+
+        input_nodes, output_nodes, blocks = self.graph_sampler.sample_blocks(
+            self.g, items
+        )
+
+        return input_nodes, output_nodes, blocks
+
+
+class EdgeCollator(Collator):
+    """DGL collator to combine edges and their computation dependencies within a minibatch for
+    training edge classification, edge regression, or link prediction on a single graph
+    with neighborhood sampling.
+
+    Given a set of edges, the collate function will yield
+
+    * A tensor of input nodes necessary for computing the representation on edges, or
+      a dictionary of node type names and such tensors.
+
+    * A subgraph that contains only the edges in the minibatch and their incident nodes.
+      Note that the graph has an identical metagraph with the original graph.
+
+    * If a negative sampler is given, another graph that contains the "negative edges",
+      connecting the source and destination nodes yielded from the given negative sampler.
+
+    * A list of MFGs necessary for computing the representation of the incident nodes
+      of the edges in the minibatch.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph from which the edges are iterated in minibatches and the subgraphs
+        are generated.
+    eids : Tensor or dict[etype, Tensor]
+        The edge set in graph :attr:`g` to compute outputs.
+    graph_sampler : dgl.dataloading.BlockSampler
+        The neighborhood sampler.
+    g_sampling : DGLGraph, optional
+        The graph where neighborhood sampling and message passing is performed.
+
+        Note that this is not necessarily the same as :attr:`g`.
+
+        If None, assume to be the same as :attr:`g`.
+    exclude : str, optional
+        Whether and how to exclude dependencies related to the sampled edges in the
+        minibatch.  Possible values are
+
+        * None, which excludes nothing.
+
+        * ``'self'``, which excludes the sampled edges themselves but nothing else.
+
+        * ``'reverse_id'``, which excludes the reverse edges of the sampled edges.  The said
+          reverse edges have the same edge type as the sampled edges.  Only works
+          on edge types whose source node type is the same as its destination node type.
+
+        * ``'reverse_types'``, which excludes the reverse edges of the sampled edges.  The
+          said reverse edges have different edge types from the sampled edges.
+
+        If ``g_sampling`` is given, ``exclude`` is ignored and will be always ``None``.
+    reverse_eids : Tensor or dict[etype, Tensor], optional
+        A tensor of reverse edge ID mapping.  The i-th element indicates the ID of
+        the i-th edge's reverse edge.
+
+        If the graph is heterogeneous, this argument requires a dictionary of edge
+        types and the reverse edge ID mapping tensors.
+
+        Required and only used when ``exclude`` is set to ``reverse_id``.
+
+        For heterogeneous graph this will be a dict of edge type and edge IDs.  Note that
+        only the edge types whose source node type is the same as destination node type
+        are needed.
+    reverse_etypes : dict[etype, etype], optional
+        The mapping from the edge type to its reverse edge type.
+
+        Required and only used when ``exclude`` is set to ``reverse_types``.
+    negative_sampler : callable, optional
+        The negative sampler.  Can be omitted if no negative sampling is needed.
+
+        The negative sampler must be a callable that takes in the following arguments:
+
+        * The original (heterogeneous) graph.
+
+        * The ID array of sampled edges in the minibatch, or the dictionary of edge
+          types and ID array of sampled edges in the minibatch if the graph is
+          heterogeneous.
+
+        It should return
+
+        * A pair of source and destination node ID arrays as negative samples,
+          or a dictionary of edge types and such pairs if the graph is heterogenenous.
+
+        A set of builtin negative samplers are provided in
+        :ref:`the negative sampling module <api-dataloading-negative-sampling>`.
+    gb_padding : int, optional
+        The padding value for GraphBolt partitions' new edge_attributes if the attributes in DistGraph are None.
+        e.g. prob/mask-based sampling.
+        Only when the mask of one edge is set as 1, an edge will be sampled in dgl.graphbolt.FusedCSCSamplingGraph.sample_neighbors.
+        The argument will be used in add_edge_attribute_to_graph to add new edge_attributes in graphbolt.
+    --------
+    The following example shows how to train a 3-layer GNN for edge classification on a
+    set of edges ``train_eid`` on a homogeneous undirected graph. Each node takes
+    messages from all neighbors.
+
+    Say that you have an array of source node IDs ``src`` and another array of destination
+    node IDs ``dst``.  One can make it bidirectional by adding another set of edges
+    that connects from ``dst`` to ``src``:
+
+    >>> g = dgl.graph((torch.cat([src, dst]), torch.cat([dst, src])))
+
+    One can then know that the ID difference of an edge and its reverse edge is ``|E|``,
+    where ``|E|`` is the length of your source/destination array.  The reverse edge
+    mapping can be obtained by
+
+    >>> E = len(src)
+    >>> reverse_eids = torch.cat([torch.arange(E, 2 * E), torch.arange(0, E)])
+
+    Note that the sampled edges as well as their reverse edges are removed from
+    computation dependencies of the incident nodes.  This is a common trick to avoid
+    information leakage.
+
+    >>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
+    >>> collator = dgl.dataloading.EdgeCollator(
+    ...     g, train_eid, sampler, exclude='reverse_id',
+    ...     reverse_eids=reverse_eids)
+    >>> dataloader = torch.utils.data.DataLoader(
+    ...     collator.dataset, collate_fn=collator.collate,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pair_graph, blocks in dataloader:
+    ...     train_on(input_nodes, pair_graph, blocks)
+
+    To train a 3-layer GNN for link prediction on a set of edges ``train_eid`` on a
+    homogeneous graph where each node takes messages from all neighbors (assume the
+    backend is PyTorch), with 5 uniformly chosen negative samples per edge:
+
+    >>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
+    >>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
+    >>> collator = dgl.dataloading.EdgeCollator(
+    ...     g, train_eid, sampler, exclude='reverse_id',
+    ...     reverse_eids=reverse_eids, negative_sampler=neg_sampler)
+    >>> dataloader = torch.utils.data.DataLoader(
+    ...     collator.dataset, collate_fn=collator.collate,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
+    ...     train_on(input_nodse, pair_graph, neg_pair_graph, blocks)
+
+    For heterogeneous graphs, the reverse of an edge may have a different edge type
+    from the original edge.  For instance, consider that you have an array of
+    user-item clicks, representated by a user array ``user`` and an item array ``item``.
+    You may want to build a heterogeneous graph with a user-click-item relation and an
+    item-clicked-by-user relation.
+
+    >>> g = dgl.heterograph({
+    ...     ('user', 'click', 'item'): (user, item),
+    ...     ('item', 'clicked-by', 'user'): (item, user)})
+
+    To train a 3-layer GNN for edge classification on a set of edges ``train_eid`` with
+    type ``click``, you can write
+
+    >>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
+    >>> collator = dgl.dataloading.EdgeCollator(
+    ...     g, {'click': train_eid}, sampler, exclude='reverse_types',
+    ...     reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'})
+    >>> dataloader = torch.utils.data.DataLoader(
+    ...     collator.dataset, collate_fn=collator.collate,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pair_graph, blocks in dataloader:
+    ...     train_on(input_nodes, pair_graph, blocks)
+
+    To train a 3-layer GNN for link prediction on a set of edges ``train_eid`` with type
+    ``click``, you can write
+
+    >>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
+    >>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
+    >>> collator = dgl.dataloading.EdgeCollator(
+    ...     g, train_eid, sampler, exclude='reverse_types',
+    ...     reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'},
+    ...     negative_sampler=neg_sampler)
+    >>> dataloader = torch.utils.data.DataLoader(
+    ...     collator.dataset, collate_fn=collator.collate,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
+    ...     train_on(input_nodes, pair_graph, neg_pair_graph, blocks)
+
+    Notes
+    -----
+    For the concept of MFGs, please refer to
+    :ref:`User Guide Section 6 <guide-minibatch>` and
+    :doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
+    """
+
+    def __init__(
+        self,
+        g,
+        eids,
+        graph_sampler,
+        g_sampling=None,
+        exclude=None,
+        reverse_eids=None,
+        reverse_etypes=None,
+        negative_sampler=None,
+        gb_padding=1,
+    ):
+        self.g = g
+        if not isinstance(eids, Mapping):
+            assert (
+                len(g.etypes) == 1
+            ), "eids should be a dict of etype and ids for graph with multiple etypes"
+        self.graph_sampler = graph_sampler
+
+        # One may wish to iterate over the edges in one graph while perform sampling in
+        # another graph.  This may be the case for iterating over validation and test
+        # edge set while perform neighborhood sampling on the graph formed by only
+        # the training edge set.
+        # See GCMC for an example usage.
+        if g_sampling is not None:
+            self.g_sampling = g_sampling
+            self.exclude = None
+        else:
+            self.g_sampling = self.g
+            self.exclude = exclude
+
+        self.reverse_eids = reverse_eids
+        self.reverse_etypes = reverse_etypes
+        self.negative_sampler = negative_sampler
+
+        self.eids = utils.prepare_tensor_or_dict(g, eids, "eids")
+        self._dataset = utils.maybe_flatten_dict(self.eids)
+
+        # Add prob/mask into graphbolt partition's edge attributes if needed.
+        if hasattr(self.graph_sampler, "prob"):
+            Collator.add_edge_attribute_to_graph(
+                self.g, self.graph_sampler.prob, gb_padding
+            )
+
+    @property
+    def dataset(self):
+        return self._dataset
+
+    def _collate(self, items):
+        if isinstance(items[0], tuple):
+            # returns a list of pairs: group them by node types into a dict
+            items = utils.group_as_dict(items)
+        items = utils.prepare_tensor_or_dict(self.g_sampling, items, "items")
+
+        pair_graph = self.g.edge_subgraph(items)
+        seed_nodes = pair_graph.ndata[NID]
+
+        exclude_eids = _find_exclude_eids(
+            self.g_sampling,
+            self.exclude,
+            items,
+            reverse_eid_map=self.reverse_eids,
+            reverse_etype_map=self.reverse_etypes,
+        )
+
+        input_nodes, _, blocks = self.graph_sampler.sample_blocks(
+            self.g_sampling, seed_nodes, exclude_eids=exclude_eids
+        )
+
+        return input_nodes, pair_graph, blocks
+
+    def _collate_with_negative_sampling(self, items):
+        if isinstance(items[0], tuple):
+            # returns a list of pairs: group them by node types into a dict
+            items = utils.group_as_dict(items)
+        items = utils.prepare_tensor_or_dict(self.g_sampling, items, "items")
+
+        pair_graph = self.g.edge_subgraph(items, relabel_nodes=False)
+        induced_edges = pair_graph.edata[EID]
+
+        neg_srcdst = self.negative_sampler(self.g, items)
+        if not isinstance(neg_srcdst, Mapping):
+            assert len(self.g.etypes) == 1, (
+                "graph has multiple or no edge types; "
+                "please return a dict in negative sampler."
+            )
+            neg_srcdst = {self.g.canonical_etypes[0]: neg_srcdst}
+        # Get dtype from a tuple of tensors
+        dtype = F.dtype(list(neg_srcdst.values())[0][0])
+        ctx = F.context(pair_graph)
+        neg_edges = {
+            etype: neg_srcdst.get(
+                etype,
+                (
+                    F.copy_to(F.tensor([], dtype), ctx),
+                    F.copy_to(F.tensor([], dtype), ctx),
+                ),
+            )
+            for etype in self.g.canonical_etypes
+        }
+        neg_pair_graph = heterograph(
+            neg_edges,
+            {ntype: self.g.num_nodes(ntype) for ntype in self.g.ntypes},
+        )
+
+        pair_graph, neg_pair_graph = transforms.compact_graphs(
+            [pair_graph, neg_pair_graph]
+        )
+        pair_graph.edata[EID] = induced_edges
+
+        seed_nodes = pair_graph.ndata[NID]
+
+        exclude_eids = _find_exclude_eids(
+            self.g_sampling,
+            self.exclude,
+            items,
+            reverse_eid_map=self.reverse_eids,
+            reverse_etype_map=self.reverse_etypes,
+        )
+
+        input_nodes, _, blocks = self.graph_sampler.sample_blocks(
+            self.g_sampling, seed_nodes, exclude_eids=exclude_eids
+        )
+
+        return input_nodes, pair_graph, neg_pair_graph, blocks
+
+    def collate(self, items):
+        """Combines the sampled edges into a minibatch for edge classification, edge
+        regression, and link prediction tasks.
+
+        Parameters
+        ----------
+        items : list[int] or list[tuple[str, int]]
+            Either a list of edge IDs (for homogeneous graphs), or a list of edge type-ID
+            pairs (for heterogeneous graphs).
+
+        Returns
+        -------
+        Either ``(input_nodes, pair_graph, blocks)``, or
+        ``(input_nodes, pair_graph, negative_pair_graph, blocks)`` if negative sampling is
+        enabled.
+
+        input_nodes : Tensor or dict[ntype, Tensor]
+            The input nodes necessary for computation in this minibatch.
+
+            If the original graph has multiple node types, return a dictionary of
+            node type names and node ID tensors.  Otherwise, return a single tensor.
+        pair_graph : DGLGraph
+            The graph that contains only the edges in the minibatch as well as their incident
+            nodes.
+
+            Note that the metagraph of this graph will be identical to that of the original
+            graph.
+        negative_pair_graph : DGLGraph
+            The graph that contains only the edges connecting the source and destination nodes
+            yielded from the given negative sampler, if negative sampling is enabled.
+
+            Note that the metagraph of this graph will be identical to that of the original
+            graph.
+        blocks : list[DGLGraph]
+            The list of MFGs necessary for computing the representation of the edges.
+        """
+        if self.negative_sampler is None:
+            return self._collate(items)
+        else:
+            return self._collate_with_negative_sampling(items)
+
+
+def _remove_kwargs_dist(kwargs):
+    if "num_workers" in kwargs:
+        del kwargs["num_workers"]
+    if "pin_memory" in kwargs:
+        del kwargs["pin_memory"]
+        print("Distributed DataLoaders do not support pin_memory.")
+    return kwargs
+
+
+class DistNodeDataLoader(DistDataLoader):
+    """Sampled graph data loader over nodes for distributed graph storage.
+
+    It wraps an iterable over a set of nodes, generating the list
+    of message flow graphs (MFGs) as computation dependency of the said minibatch, on
+    a distributed graph.
+
+    All the arguments have the same meaning as the single-machine counterpart
+    :class:`dgl.dataloading.DataLoader` except the first argument
+    :attr:`g` which must be a :class:`dgl.distributed.DistGraph`.
+
+    Parameters
+    ----------
+    g : DistGraph
+        The distributed graph.
+
+    nids, graph_sampler, device, kwargs :
+        See :class:`dgl.dataloading.DataLoader`.
+
+    See also
+    --------
+    dgl.dataloading.DataLoader
+    """
+
+    def __init__(self, g, nids, graph_sampler, device=None, **kwargs):
+        collator_kwargs = {}
+        dataloader_kwargs = {}
+        _collator_arglist = inspect.getfullargspec(NodeCollator).args
+        for k, v in kwargs.items():
+            if k in _collator_arglist:
+                collator_kwargs[k] = v
+            else:
+                dataloader_kwargs[k] = v
+        if device is None:
+            # for the distributed case default to the CPU
+            device = "cpu"
+        assert (
+            device == "cpu"
+        ), "Only cpu is supported in the case of a DistGraph."
+        # Distributed DataLoader currently does not support heterogeneous graphs
+        # and does not copy features.  Fallback to normal solution
+        self.collator = NodeCollator(g, nids, graph_sampler, **collator_kwargs)
+        _remove_kwargs_dist(dataloader_kwargs)
+        super().__init__(
+            self.collator.dataset,
+            collate_fn=self.collator.collate,
+            **dataloader_kwargs
+        )
+        self.device = device
+
+
+class DistEdgeDataLoader(DistDataLoader):
+    """Sampled graph data loader over edges for distributed graph storage.
+
+    It wraps an iterable over a set of edges, generating the list
+    of message flow graphs (MFGs) as computation dependency of the said minibatch for
+    edge classification, edge regression, and link prediction, on a distributed
+    graph.
+
+    All the arguments have the same meaning as the single-machine counterpart
+    :class:`dgl.dataloading.DataLoader` except the first argument
+    :attr:`g` which must be a :class:`dgl.distributed.DistGraph`.
+
+    Parameters
+    ----------
+    g : DistGraph
+        The distributed graph.
+
+    eids, graph_sampler, device, kwargs :
+        See :class:`dgl.dataloading.DataLoader`.
+
+    See also
+    --------
+    dgl.dataloading.DataLoader
+    """
+
+    def __init__(self, g, eids, graph_sampler, device=None, **kwargs):
+        collator_kwargs = {}
+        dataloader_kwargs = {}
+        _collator_arglist = inspect.getfullargspec(EdgeCollator).args
+        for k, v in kwargs.items():
+            if k in _collator_arglist:
+                collator_kwargs[k] = v
+            else:
+                dataloader_kwargs[k] = v
+
+        if device is None:
+            # for the distributed case default to the CPU
+            device = "cpu"
+        assert (
+            device == "cpu"
+        ), "Only cpu is supported in the case of a DistGraph."
+        # Distributed DataLoader currently does not support heterogeneous graphs
+        # and does not copy features.  Fallback to normal solution
+        self.collator = EdgeCollator(g, eids, graph_sampler, **collator_kwargs)
+        _remove_kwargs_dist(dataloader_kwargs)
+        super().__init__(
+            self.collator.dataset,
+            collate_fn=self.collator.collate,
+            **dataloader_kwargs
+        )
+
+        self.device = device
@@ -0,0 +1,327 @@
+"""Define distributed tensor."""
+
+import os
+
+from .. import backend as F, utils
+
+from .dist_context import is_initialized
+from .kvstore import get_kvstore
+from .role import get_role
+from .rpc import get_group_id
+
+
+def _default_init_data(shape, dtype):
+    return F.zeros(shape, dtype, F.cpu())
+
+
+# These IDs can identify the anonymous distributed tensors.
+DIST_TENSOR_ID = 0
+
+
+class DistTensor:
+    """Distributed tensor.
+
+    ``DistTensor`` references to a distributed tensor sharded and stored in a cluster of machines.
+    It has the same interface as Pytorch Tensor to access its metadata (e.g., shape and data type).
+    To access data in a distributed tensor, it supports slicing rows and writing data to rows.
+    It does not support any operators of a deep learning framework, such as addition and
+    multiplication.
+
+    Currently, distributed tensors are designed to store node data and edge data of a distributed
+    graph. Therefore, their first dimensions have to be the number of nodes or edges in the graph.
+    The tensors are sharded in the first dimension based on the partition policy of nodes
+    or edges. When a distributed tensor is created, the partition policy is automatically
+    determined based on the first dimension if the partition policy is not provided. If the first
+    dimension matches the number of nodes of a node type, ``DistTensor`` will use the partition
+    policy for this particular node type; if the first dimension matches the number of edges of
+    an edge type, ``DistTensor`` will use the partition policy for this particular edge type.
+    If DGL cannot determine the partition policy automatically (e.g., multiple node types or
+    edge types have the same number of nodes or edges), users have to explicity provide
+    the partition policy.
+
+    A distributed tensor can be ether named or anonymous.
+    When a distributed tensor has a name, the tensor can be persistent if ``persistent=True``.
+    Normally, DGL destroys the distributed tensor in the system when the ``DistTensor`` object
+    goes away. However, a persistent tensor lives in the system even if
+    the ``DistTenor`` object disappears in the trainer process. The persistent tensor has
+    the same life span as the DGL servers. DGL does not allow an anonymous tensor to be persistent.
+
+    When a ``DistTensor`` object is created, it may reference to an existing distributed tensor or
+    create a new one. A distributed tensor is identified by the name passed to the constructor.
+    If the name exists, ``DistTensor`` will reference the existing one.
+    In this case, the shape and the data type must match the existing tensor.
+    If the name doesn't exist, a new tensor will be created in the kvstore.
+
+    When a distributed tensor is created, its values are initialized to zero. Users
+    can define an initialization function to control how the values are initialized.
+    The init function has two input arguments: shape and data type and returns a tensor.
+    Below shows an example of an init function:
+
+    .. highlight:: python
+    .. code-block:: python
+
+        def init_func(shape, dtype):
+            return torch.ones(shape=shape, dtype=dtype)
+
+    Parameters
+    ----------
+    shape : tuple
+        The shape of the tensor. The first dimension has to be the number of nodes or
+        the number of edges of a distributed graph.
+    dtype : dtype
+        The dtype of the tensor. The data type has to be the one in the deep learning framework.
+    name : string, optional
+        The name of the embeddings. The name can uniquely identify embeddings in a system
+        so that another ``DistTensor`` object can referent to the distributed tensor.
+    init_func : callable, optional
+        The function to initialize data in the tensor. If the init function is not provided,
+        the values of the embeddings are initialized to zero.
+    part_policy : PartitionPolicy, optional
+        The partition policy of the rows of the tensor to different machines in the cluster.
+        Currently, it only supports node partition policy or edge partition policy.
+        The system determines the right partition policy automatically.
+    persistent : bool
+        Whether the created tensor lives after the ``DistTensor`` object is destroyed.
+    is_gdata : bool
+        Whether the created tensor is a ndata/edata or not.
+    attach : bool
+        Whether to attach group ID into name to be globally unique.
+
+    Examples
+    --------
+    >>> init = lambda shape, dtype: th.ones(shape, dtype=dtype)
+    >>> arr = dgl.distributed.DistTensor((g.num_nodes(), 2), th.int32, init_func=init)
+    >>> print(arr[0:3])
+    tensor([[1, 1],
+            [1, 1],
+            [1, 1]], dtype=torch.int32)
+    >>> arr[0:3] = th.ones((3, 2), dtype=th.int32) * 2
+    >>> print(arr[0:3])
+    tensor([[2, 2],
+            [2, 2],
+            [2, 2]], dtype=torch.int32)
+
+    Note
+    ----
+    The creation of ``DistTensor`` is a synchronized operation. When a trainer process tries to
+    create a ``DistTensor`` object, the creation succeeds only when all trainer processes
+    do the same.
+    """
+
+    def __init__(
+        self,
+        shape,
+        dtype,
+        name=None,
+        init_func=None,
+        part_policy=None,
+        persistent=False,
+        is_gdata=True,
+        attach=True,
+    ):
+        self.kvstore = get_kvstore()
+        assert (
+            self.kvstore is not None
+        ), "Distributed module is not initialized. Please call dgl.distributed.initialize."
+        self._shape = shape
+        self._dtype = dtype
+        self._attach = attach
+        self._is_gdata = is_gdata
+
+        part_policies = self.kvstore.all_possible_part_policy
+        # If a user doesn't provide a partition policy, we should find one based on
+        # the input shape.
+        if part_policy is None:
+            for policy_name in part_policies:
+                policy = part_policies[policy_name]
+                if policy.get_size() == shape[0]:
+                    # If multiple partition policies match the input shape, we cannot
+                    # decide which is the right one automatically. We should ask users
+                    # to provide one.
+                    assert part_policy is None, (
+                        "Multiple partition policies match the input shape. "
+                        + "Please provide a partition policy explicitly."
+                    )
+                    part_policy = policy
+            assert part_policy is not None, (
+                "Cannot find a right partition policy. It is either because "
+                + "its first dimension does not match the number of nodes or edges "
+                + "of a distributed graph or there does not exist a distributed graph."
+            )
+
+        self._part_policy = part_policy
+        assert (
+            part_policy.get_size() == shape[0]
+        ), "The partition policy does not match the input shape."
+
+        if init_func is None:
+            init_func = _default_init_data
+        exist_names = self.kvstore.data_name_list()
+        # If a user doesn't provide a name, we generate a name ourselves.
+        # We need to generate the name in a deterministic way.
+        if name is None:
+            assert (
+                not persistent
+            ), "We cannot generate anonymous persistent distributed tensors"
+            global DIST_TENSOR_ID
+            # All processes of the same role should create DistTensor synchronously.
+            # Thus, all of them should have the same IDs.
+            name = "anonymous-" + get_role() + "-" + str(DIST_TENSOR_ID)
+            DIST_TENSOR_ID += 1
+        assert isinstance(name, str), "name {} is type {}".format(
+            name, type(name)
+        )
+        name = self._attach_group_id(name)
+        self._tensor_name = name
+        data_name = part_policy.get_data_name(name)
+        self._name = str(data_name)
+        self._persistent = persistent
+        if self._name not in exist_names:
+            self._owner = True
+            self.kvstore.init_data(
+                self._name, shape, dtype, part_policy, init_func, is_gdata
+            )
+        else:
+            self._owner = False
+            dtype1, shape1, _ = self.kvstore.get_data_meta(self._name)
+            assert (
+                dtype == dtype1
+            ), "The dtype does not match with the existing tensor"
+            assert (
+                shape == shape1
+            ), "The shape does not match with the existing tensor"
+
+    def __del__(self):
+        initialized = (
+            os.environ.get("DGL_DIST_MODE", "standalone") == "standalone"
+            or is_initialized()
+        )
+        if not self._persistent and self._owner and initialized:
+            self.kvstore.delete_data(self._name)
+
+    def __getitem__(self, idx):
+        idx = utils.toindex(idx)
+        idx = idx.tousertensor()
+        return self.kvstore.pull(name=self._name, id_tensor=idx)
+
+    def __setitem__(self, idx, val):
+        idx = utils.toindex(idx)
+        idx = idx.tousertensor()
+        # TODO(zhengda) how do we want to support broadcast (e.g., G.ndata['h'][idx] = 1).
+        self.kvstore.push(name=self._name, id_tensor=idx, data_tensor=val)
+
+    @property
+    def kvstore_key(self):
+        """Return the key string of this DistTensor in the associated KVStore."""
+        return self._name
+
+    @property
+    def local_partition(self):
+        """Return the local partition of this DistTensor."""
+        return self.kvstore.data_store[self._name]
+
+    def __or__(self, other):
+        new_dist_tensor = DistTensor(
+            self._shape,
+            self._dtype,
+            part_policy=self._part_policy,
+            persistent=self._persistent,
+            is_gdata=self._is_gdata,
+            attach=self._attach,
+        )
+        kvstore = self.kvstore
+        kvstore.union(self._name, other._name, new_dist_tensor._name)
+        return new_dist_tensor
+
+    def __len__(self):
+        return self._shape[0]
+
+    @property
+    def part_policy(self):
+        """Return the partition policy
+
+        Returns
+        -------
+        PartitionPolicy
+            The partition policy of the distributed tensor.
+        """
+        return self._part_policy
+
+    @property
+    def shape(self):
+        """Return the shape of the distributed tensor.
+
+        Returns
+        -------
+        tuple
+            The shape of the distributed tensor.
+        """
+        return self._shape
+
+    @property
+    def dtype(self):
+        """Return the data type of the distributed tensor.
+
+        Returns
+        ------
+        dtype
+            The data type of the tensor.
+        """
+        return self._dtype
+
+    @property
+    def name(self):
+        """Return the name of the distributed tensor
+
+        Returns
+        -------
+        str
+            The name of the tensor.
+        """
+        return self._detach_group_id(self._name)
+
+    @property
+    def tensor_name(self):
+        """Return the tensor name
+
+        Returns
+        -------
+        str
+            The name of the tensor.
+        """
+        return self._detach_group_id(self._tensor_name)
+
+    def count_nonzero(self):
+        """Count and return the number of nonzero value
+
+        Returns
+        -------
+        int
+            the number of nonzero value
+        """
+        return self.kvstore.count_nonzero(name=self._name)
+
+    def _attach_group_id(self, name):
+        """Attach group ID if needed
+
+        Returns
+        -------
+        str
+            new name with group ID attached
+        """
+        if not self._attach:
+            return name
+        return "{}_{}".format(name, get_group_id())
+
+    def _detach_group_id(self, name):
+        """Detach group ID if needed
+
+        Returns
+        -------
+        str
+            original name without group ID
+        """
+        if not self._attach:
+            return name
+        suffix = "_{}".format(get_group_id())
+        return name[: -len(suffix)]
@@ -0,0 +1,178 @@
+"""Module for mapping between node/edge IDs and node/edge types."""
+
+import numpy as np
+import torch
+
+from .. import backend as F, utils
+
+from .._ffi.function import _init_api
+
+
+__all__ = ["IdMap"]
+
+
+class IdMap:
+    """A map for converting node/edge IDs to their type IDs and type-wise IDs.
+
+    For a heterogeneous graph, DGL assigns an integer ID to each node/edge type;
+    node and edge of different types have independent IDs starting from zero.
+    Therefore, a node/edge can be uniquely identified by an ID pair,
+    ``(type_id, type_wise_id)``. To make it convenient for distributed processing,
+    DGL further encodes the ID pair into one integer ID, which we refer to
+    as *homogeneous ID*.
+
+    DGL arranges nodes and edges so that all nodes of the same type have contiguous
+    homogeneous IDs. If the graph is partitioned, the nodes/edges of the same type
+    within a partition have contiguous homogeneous IDs.
+
+    Below is an example adjancency matrix of an unpartitioned heterogeneous graph
+    stored using the above ID assignment. Here, the graph has two types of nodes
+    (``T0`` and ``T1``), and four types of edges (``R0``, ``R1``, ``R2``, ``R3``).
+    There are a total of 400 nodes in the graph and each type has 200 nodes. Nodes
+    of type 0 have IDs in [0,200), while nodes of type 1 have IDs in [200, 400).
+
+    ```
+        0 <- T0 -> 200 <- T1 -> 400
+     0  +-----------+------------+
+        |           |            |
+     ^  |    R0     |     R1     |
+     T0 |           |            |
+     v  |           |            |
+    200 +-----------+------------+
+        |           |            |
+     ^  |    R2     |     R3     |
+     T1 |           |            |
+     v  |           |            |
+    400 +-----------+------------+
+    ```
+
+    Below shows the adjacency matrix after the graph is partitioned into two.
+    Note that each partition still has two node types and four edge types,
+    and nodes/edges of the same type have contiguous IDs.
+
+    ```
+                partition 0              partition 1
+
+        0 <- T0 -> 100 <- T1 -> 200 <- T0 -> 300 <- T1 -> 400
+     0  +-----------+------------+-----------+------------+
+        |           |            |                        |
+     ^  |    R0     |     R1     |                        |
+     T0 |           |            |                        |
+     v  |           |            |                        |
+    100 +-----------+------------+                        |
+        |           |            |                        |
+     ^  |    R2     |     R3     |                        |
+     T1 |           |            |                        |
+     v  |           |            |                        |
+    200 +-----------+------------+-----------+------------+
+        |                        |           |            |
+     ^  |                        |    R0     |     R1     |
+     T0 |                        |           |            |
+     v  |                        |           |            |
+    100 |                        +-----------+------------+
+        |                        |           |            |
+     ^  |                        |    R2     |     R3     |
+     T1 |                        |           |            |
+     v  |                        |           |            |
+    200 +-----------+------------+-----------+------------+
+    ```
+
+    The following table is an alternative way to represent the above ID assignments.
+    It is easy to see that the homogeneous ID range [0, 100) is used for nodes of type 0
+    in partition 0, [100, 200) is used for nodes of type 1 in partition 0, and so on.
+    ```
+    +---------+------+----------
+      range   | type | partition
+    [0, 100)  |   0  |    0
+    [100,200) |   1  |    0
+    [200,300) |   0  |    1
+    [300,400) |   1  |    1
+    ```
+
+    The goal of this class is to, given a node's homogenous ID, convert it into the
+    ID pair ``(type_id, type_wise_id)``. For example, homogeneous node ID 90 is mapped
+    to (0, 90); homogeneous node ID 201 is mapped to (0, 101).
+
+    Parameters
+    ----------
+    id_ranges : dict[str, Tensor].
+        Node ID ranges within partitions for each node type. The key is the node type
+        name in string. The value is a tensor of shape :math:`(K, 2)`, where :math:`K` is
+        the number of partitions. Each row has two integers: the starting and the ending IDs
+        for a particular node type in a partition. For example, all nodes of type ``"T"`` in
+        partition ``i`` has ID range ``id_ranges["T"][i][0]`` to ``id_ranges["T"][i][1]``.
+        It is the same as the `node_map` argument in `RangePartitionBook`.
+    """
+
+    def __init__(self, id_ranges):
+        id_ranges_values = list(id_ranges.values())
+        assert isinstance(
+            id_ranges_values[0], np.ndarray
+        ), "id_ranges should be a dict of numpy arrays."
+        self.num_parts = id_ranges_values[0].shape[0]
+        self.dtype = id_ranges_values[0].dtype
+        self.dtype_str = "int32" if self.dtype == np.int32 else "int64"
+        self.num_types = len(id_ranges)
+        ranges = np.zeros(
+            (self.num_parts * self.num_types, 2), dtype=self.dtype
+        )
+        typed_map = []
+        id_ranges = id_ranges_values
+        id_ranges.sort(key=lambda a: a[0, 0])
+        for i, id_range in enumerate(id_ranges):
+            ranges[i :: self.num_types] = id_range
+            map1 = np.cumsum(id_range[:, 1] - id_range[:, 0], dtype=self.dtype)
+            typed_map.append(map1)
+
+        assert np.all(np.diff(ranges[:, 0]) >= 0)
+        assert np.all(np.diff(ranges[:, 1]) >= 0)
+        self.range_start = utils.toindex(
+            np.ascontiguousarray(ranges[:, 0]), dtype=self.dtype_str
+        )
+        self.range_end = utils.toindex(
+            np.ascontiguousarray(ranges[:, 1]) - 1, dtype=self.dtype_str
+        )
+        self.typed_map = utils.toindex(
+            np.concatenate(typed_map), dtype=self.dtype_str
+        )
+
+    def __call__(self, ids):
+        """Convert the homogeneous IDs to (type_id, type_wise_id).
+
+        Parameters
+        ----------
+        ids : 1D tensor
+            The homogeneous ID.
+
+        Returns
+        -------
+        type_ids : Tensor
+            Type IDs
+        per_type_ids : Tensor
+            Type-wise IDs
+        """
+        if self.num_types == 0:
+            return F.zeros((len(ids),), F.dtype(ids), F.cpu()), ids
+        if len(ids) == 0:
+            return ids, ids
+
+        ids = utils.toindex(ids, dtype=self.dtype_str)
+        ret = _CAPI_DGLHeteroMapIds(
+            ids.todgltensor(),
+            self.range_start.todgltensor(),
+            self.range_end.todgltensor(),
+            self.typed_map.todgltensor(),
+            self.num_parts,
+            self.num_types,
+        )
+        ret = utils.toindex(ret, dtype=self.dtype_str).tousertensor()
+        return ret[: len(ids)], ret[len(ids) :]
+
+    @property
+    def torch_dtype(self):
+        """Return the data type of the ID map."""
+        # [TODO][Rui] Use torch instead of numpy.
+        return torch.int32 if self.dtype == np.int32 else torch.int64
+
+
+_init_api("dgl.distributed.id_map")
@@ -0,0 +1,17 @@
+"""dgl distributed.optims."""
+import importlib
+import os
+import sys
+
+from ...backend import backend_name
+from ...utils import expand_as_pair
+
+
+def _load_backend(mod_name):
+    mod = importlib.import_module(".%s" % mod_name, __name__)
+    thismod = sys.modules[__name__]
+    for api, obj in mod.__dict__.items():
+        setattr(thismod, api, obj)
+
+
+_load_backend(backend_name)
@@ -0,0 +1,2 @@
+"""dgl distributed sparse optimizer for pytorch."""
+from .sparse_emb import DistEmbedding
@@ -0,0 +1,209 @@
+"""Define sparse embedding and optimizer."""
+
+import torch as th
+
+from .... import backend as F, utils
+from ...dist_tensor import DistTensor
+
+
+class DistEmbedding:
+    """Distributed node embeddings.
+
+    DGL provides a distributed embedding to support models that require learnable embeddings.
+    DGL's distributed embeddings are mainly used for learning node embeddings of graph models.
+    Because distributed embeddings are part of a model, they are updated by mini-batches.
+    The distributed embeddings have to be updated by DGL's optimizers instead of
+    the optimizers provided by the deep learning frameworks (e.g., Pytorch and MXNet).
+
+    To support efficient training on a graph with many nodes, the embeddings support sparse
+    updates. That is, only the embeddings involved in a mini-batch computation are updated.
+    Please refer to `Distributed Optimizers <https://docs.dgl.ai/api/python/dgl.distributed.html#
+    distributed-embedding-optimizer>`__ for available optimizers in DGL.
+
+    Distributed embeddings are sharded and stored in a cluster of machines in the same way as
+    :class:`dgl.distributed.DistTensor`, except that distributed embeddings are trainable.
+    Because distributed embeddings are sharded
+    in the same way as nodes and edges of a distributed graph, it is usually much more
+    efficient to access than the sparse embeddings provided by the deep learning frameworks.
+
+    Parameters
+    ----------
+    num_embeddings : int
+        The number of embeddings. Currently, the number of embeddings has to be the same as
+        the number of nodes or the number of edges.
+    embedding_dim : int
+        The dimension size of embeddings.
+    name : str, optional
+        The name of the embeddings. The name can uniquely identify embeddings in a system
+        so that another DistEmbedding object can referent to the same embeddings.
+    init_func : callable, optional
+        The function to create the initial data. If the init function is not provided,
+        the values of the embeddings are initialized to zero.
+    part_policy : PartitionPolicy, optional
+        The partition policy that assigns embeddings to different machines in the cluster.
+        Currently, it only supports node partition policy or edge partition policy.
+        The system determines the right partition policy automatically.
+
+    Examples
+    --------
+    >>> def initializer(shape, dtype):
+            arr = th.zeros(shape, dtype=dtype)
+            arr.uniform_(-1, 1)
+            return arr
+    >>> emb = dgl.distributed.DistEmbedding(g.num_nodes(), 10, init_func=initializer)
+    >>> optimizer = dgl.distributed.optim.SparseAdagrad([emb], lr=0.001)
+    >>> for blocks in dataloader:
+    ...     feats = emb(nids)
+    ...     loss = F.sum(feats + 1, 0)
+    ...     loss.backward()
+    ...     optimizer.step()
+
+    Note
+    ----
+    When a ``DistEmbedding``  object is used in the forward computation, users
+    have to invoke
+    :py:meth:`~dgl.distributed.optim.SparseAdagrad.step` afterwards. Otherwise,
+    there will be some memory leak.
+    """
+
+    def __init__(
+        self,
+        num_embeddings,
+        embedding_dim,
+        name=None,
+        init_func=None,
+        part_policy=None,
+    ):
+        self._tensor = DistTensor(
+            (num_embeddings, embedding_dim),
+            F.float32,
+            name,
+            init_func=init_func,
+            part_policy=part_policy,
+        )
+        self._trace = []
+        self._name = name
+        self._num_embeddings = num_embeddings
+        self._embedding_dim = embedding_dim
+
+        # Check whether it is multi-gpu/distributed training or not
+        if th.distributed.is_initialized():
+            self._rank = th.distributed.get_rank()
+            self._world_size = th.distributed.get_world_size()
+        # [TODO] The following code is clearly wrong but changing it to "raise DGLError"
+        # actually fails unit test.  ???
+        # else:
+        #     assert 'th.distributed should be initialized'
+        self._optm_state = None  # track optimizer state
+        self._part_policy = part_policy
+
+    def __call__(self, idx, device=th.device("cpu")):
+        """
+        node_ids : th.tensor
+            Index of the embeddings to collect.
+        device : th.device
+            Target device to put the collected embeddings.
+
+        Returns
+        -------
+        Tensor
+            The requested node embeddings
+        """
+        idx = utils.toindex(idx).tousertensor()
+        emb = self._tensor[idx].to(device, non_blocking=True)
+        if F.is_recording():
+            emb = F.attach_grad(emb)
+            self._trace.append((idx.to(device, non_blocking=True), emb))
+        return emb
+
+    def reset_trace(self):
+        """Reset the traced data."""
+        self._trace = []
+
+    @property
+    def part_policy(self):
+        """Return the partition policy
+
+        Returns
+        -------
+        PartitionPolicy
+            partition policy
+        """
+        return self._part_policy
+
+    @property
+    def name(self):
+        """Return the name of the embeddings
+
+        Returns
+        -------
+        str
+            The name of the embeddings
+        """
+        return self._tensor.tensor_name
+
+    @property
+    def data_name(self):
+        """Return the data name of the embeddings
+
+        Returns
+        -------
+        str
+            The data name of the embeddings
+        """
+        return self._tensor._name
+
+    @property
+    def kvstore(self):
+        """Return the kvstore client
+
+        Returns
+        -------
+        KVClient
+            The kvstore client
+        """
+        return self._tensor.kvstore
+
+    @property
+    def num_embeddings(self):
+        """Return the number of embeddings
+
+        Returns
+        -------
+        int
+            The number of embeddings
+        """
+        return self._num_embeddings
+
+    @property
+    def embedding_dim(self):
+        """Return the dimension of embeddings
+
+        Returns
+        -------
+        int
+            The dimension of embeddings
+        """
+        return self._embedding_dim
+
+    @property
+    def optm_state(self):
+        """Return the optimizer related state tensor.
+
+        Returns
+        -------
+        tuple of torch.Tensor
+            The optimizer related state.
+        """
+        return self._optm_state
+
+    @property
+    def weight(self):
+        """Return the tensor storing the node embeddings
+
+        Returns
+        -------
+        torch.Tensor
+            The tensor storing the node embeddings
+        """
+        return self._tensor
@@ -0,0 +1,17 @@
+"""dgl distributed.optims."""
+import importlib
+import os
+import sys
+
+from ...backend import backend_name
+from ...utils import expand_as_pair
+
+
+def _load_backend(mod_name):
+    mod = importlib.import_module(".%s" % mod_name, __name__)
+    thismod = sys.modules[__name__]
+    for api, obj in mod.__dict__.items():
+        setattr(thismod, api, obj)
+
+
+_load_backend(backend_name)
@@ -0,0 +1,2 @@
+"""dgl distributed sparse optimizer for pytorch."""
+from .sparse_optim import SparseAdagrad, SparseAdam
@@ -0,0 +1,747 @@
+"""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
@@ -0,0 +1,114 @@
+"""Provide utils for distributed sparse optimizers
+"""
+import torch as th
+import torch.distributed as dist
+
+
+def alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list):
+    """Each process scatters list of input tensors to all processes in a cluster
+    and return gathered list of tensors in output list. The tensors should have the same shape.
+
+    Parameters
+    ----------
+    rank : int
+        The rank of current worker
+    world_size : int
+        The size of the entire communicator
+    output_tensor_list : List of tensor
+        The received tensors
+    input_tensor_list : List of tensor
+        The tensors to exchange
+    """
+    input_tensor_list = [
+        tensor.to(th.device("cpu")) for tensor in input_tensor_list
+    ]
+    for i in range(world_size):
+        dist.scatter(
+            output_tensor_list[i], input_tensor_list if i == rank else [], src=i
+        )
+
+
+def alltoallv_cpu(rank, world_size, output_tensor_list, input_tensor_list):
+    """Each process scatters list of input tensors to all processes in a cluster
+    and return gathered list of tensors in output list.
+
+    Parameters
+    ----------
+    rank : int
+        The rank of current worker
+    world_size : int
+        The size of the entire communicator
+    output_tensor_list : List of tensor
+        The received tensors
+    input_tensor_list : List of tensor
+        The tensors to exchange
+    """
+    # send tensor to each target trainer using torch.distributed.isend
+    # isend is async
+    senders = []
+    for i in range(world_size):
+        if i == rank:
+            output_tensor_list[i] = input_tensor_list[i].to(th.device("cpu"))
+        else:
+            sender = dist.isend(
+                input_tensor_list[i].to(th.device("cpu")), dst=i
+            )
+            senders.append(sender)
+
+    for i in range(world_size):
+        if i != rank:
+            dist.recv(output_tensor_list[i], src=i)
+
+    th.distributed.barrier()
+
+
+def alltoall(rank, world_size, output_tensor_list, input_tensor_list):
+    """Each process scatters list of input tensors to all processes in a cluster
+    and return gathered list of tensors in output list. The tensors should have the same shape.
+
+    Parameters
+    ----------
+    rank : int
+        The rank of current worker
+    world_size : int
+        The size of the entire communicator
+    output_tensor_list : List of tensor
+        The received tensors
+    input_tensor_list : List of tensor
+        The tensors to exchange
+    """
+    if th.distributed.get_backend() == "nccl":
+        th.distributed.all_to_all(output_tensor_list, input_tensor_list)
+    else:
+        alltoall_cpu(
+            rank,
+            world_size,
+            output_tensor_list,
+            input_tensor_list,
+        )
+
+
+def alltoallv(rank, world_size, output_tensor_list, input_tensor_list):
+    """Each process scatters list of input tensors to all processes in a cluster
+    and return gathered list of tensors in output list.
+
+    Parameters
+    ----------
+    rank : int
+        The rank of current worker
+    world_size : int
+        The size of the entire communicator
+    output_tensor_list : List of tensor
+        The received tensors
+    input_tensor_list : List of tensor
+        The tensors to exchange
+    """
+    if th.distributed.get_backend() == "nccl":
+        th.distributed.all_to_all(output_tensor_list, input_tensor_list)
+    else:
+        alltoallv_cpu(
+            rank,
+            world_size,
+            output_tensor_list,
+            input_tensor_list,
+        )
@@ -0,0 +1,247 @@
+"""Manage the roles in different clients.
+
+Right now, the clients have different roles. Some clients work as samplers and
+some work as trainers.
+"""
+
+import os
+
+import numpy as np
+
+from . import rpc
+
+REGISTER_ROLE = 700001
+REG_ROLE_MSG = "Register_Role"
+
+
+class RegisterRoleResponse(rpc.Response):
+    """Send a confirmation signal (just a short string message)
+    of RegisterRoleRequest to client.
+    """
+
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __getstate__(self):
+        return self.msg
+
+    def __setstate__(self, state):
+        self.msg = state
+
+
+class RegisterRoleRequest(rpc.Request):
+    """Send client id and role to server
+
+    Parameters
+    ----------
+    client_id : int
+        ID of client
+    role : str
+        role of client
+    """
+
+    def __init__(self, client_id, machine_id, role):
+        self.client_id = client_id
+        self.machine_id = machine_id
+        self.role = role
+        self.group_id = rpc.get_group_id()
+
+    def __getstate__(self):
+        return self.client_id, self.machine_id, self.role, self.group_id
+
+    def __setstate__(self, state):
+        self.client_id, self.machine_id, self.role, self.group_id = state
+
+    def process_request(self, server_state):
+        kv_store = server_state.kv_store
+        role = server_state.roles.setdefault(self.group_id, {})
+        if self.role not in role:
+            role[self.role] = set()
+            if kv_store is not None:
+                barrier_count = kv_store.barrier_count.setdefault(
+                    self.group_id, {}
+                )
+                barrier_count[self.role] = 0
+        role[self.role].add((self.client_id, self.machine_id))
+        total_count = 0
+        for key in role:
+            total_count += len(role[key])
+        # Clients are blocked util all clients register their roles.
+        if total_count == rpc.get_num_client():
+            res_list = []
+            for target_id in range(rpc.get_num_client()):
+                res_list.append((target_id, RegisterRoleResponse(REG_ROLE_MSG)))
+            return res_list
+        return None
+
+
+GET_ROLE = 700002
+GET_ROLE_MSG = "Get_Role"
+
+
+class GetRoleResponse(rpc.Response):
+    """Send the roles of all client processes"""
+
+    def __init__(self, role):
+        self.role = role
+        self.msg = GET_ROLE_MSG
+
+    def __getstate__(self):
+        return self.role, self.msg
+
+    def __setstate__(self, state):
+        self.role, self.msg = state
+
+
+class GetRoleRequest(rpc.Request):
+    """Send a request to get the roles of all client processes."""
+
+    def __init__(self):
+        self.msg = GET_ROLE_MSG
+        self.group_id = rpc.get_group_id()
+
+    def __getstate__(self):
+        return self.msg, self.group_id
+
+    def __setstate__(self, state):
+        self.msg, self.group_id = state
+
+    def process_request(self, server_state):
+        return GetRoleResponse(server_state.roles[self.group_id])
+
+
+# The key is role, the value is a dict of mapping RPC rank to a rank within the role.
+PER_ROLE_RANK = {}
+
+# The global rank of a client process. The client processes of the same role have
+# global ranks that fall in a contiguous range.
+GLOBAL_RANK = {}
+
+# The role of the current process
+CUR_ROLE = None
+
+IS_STANDALONE = False
+
+
+def init_role(role):
+    """Initialize the role of the current process.
+
+    Each process is associated with a role so that we can determine what
+    function can be invoked in a process. For example, we do not allow some
+    functions in sampler processes.
+
+    The initialization includes registeration the role of the current process and
+    get the roles of all client processes. It also computes the rank of all client
+    processes in a deterministic way so that all clients will have the same rank for
+    the same client process.
+    """
+    global CUR_ROLE
+    CUR_ROLE = role
+
+    global PER_ROLE_RANK
+    global GLOBAL_RANK
+    global IS_STANDALONE
+
+    if os.environ.get("DGL_DIST_MODE", "standalone") == "standalone":
+        if role == "default":
+            GLOBAL_RANK[0] = 0
+            PER_ROLE_RANK["default"] = {0: 0}
+        IS_STANDALONE = True
+        return
+
+    PER_ROLE_RANK = {}
+    GLOBAL_RANK = {}
+
+    # Register the current role. This blocks until all clients register themselves.
+    client_id = rpc.get_rank()
+    machine_id = rpc.get_machine_id()
+    request = RegisterRoleRequest(client_id, machine_id, role)
+    rpc.send_request(0, request)
+    response = rpc.recv_response()
+    assert response.msg == REG_ROLE_MSG
+
+    # Get all clients on all machines.
+    request = GetRoleRequest()
+    rpc.send_request(0, request)
+    response = rpc.recv_response()
+    assert response.msg == GET_ROLE_MSG
+
+    # Here we want to compute a new rank for each client.
+    # We compute the per-role rank as well as global rank.
+    # For per-role rank, we ensure that all ranks within a machine is contiguous.
+    # For global rank, we also ensure that all ranks within a machine are contiguous,
+    # and all ranks within a role are contiguous.
+    global_rank = 0
+
+    # We want to ensure that the global rank of the trainer process starts from 0.
+    role_names = ["default"]
+    for role_name in response.role:
+        if role_name not in role_names:
+            role_names.append(role_name)
+
+    for role_name in role_names:
+        # Let's collect the ranks of this role in all machines.
+        machines = {}
+        for client_id, machine_id in response.role[role_name]:
+            if machine_id not in machines:
+                machines[machine_id] = []
+            machines[machine_id].append(client_id)
+
+        num_machines = len(machines)
+        PER_ROLE_RANK[role_name] = {}
+        per_role_rank = 0
+        for i in range(num_machines):
+            clients = machines[i]
+            clients = np.sort(clients)
+            for client_id in clients:
+                GLOBAL_RANK[client_id] = global_rank
+                global_rank += 1
+                PER_ROLE_RANK[role_name][client_id] = per_role_rank
+                per_role_rank += 1
+
+
+def get_global_rank():
+    """Get the global rank
+
+    The rank can globally identify the client process. For the client processes
+    of the same role, their ranks are in a contiguous range.
+    """
+    if IS_STANDALONE:
+        return 0
+    else:
+        return GLOBAL_RANK[rpc.get_rank()]
+
+
+def get_rank(role):
+    """Get the role-specific rank"""
+    if IS_STANDALONE:
+        return 0
+    else:
+        return PER_ROLE_RANK[role][rpc.get_rank()]
+
+
+def get_trainer_rank():
+    """Get the rank of the current trainer process.
+
+    This function can only be called in the trainer process. It will result in
+    an error if it's called in the process of other roles.
+    """
+    assert CUR_ROLE == "default"
+    if IS_STANDALONE:
+        return 0
+    else:
+        return PER_ROLE_RANK["default"][rpc.get_rank()]
+
+
+def get_role():
+    """Get the role of the current process"""
+    return CUR_ROLE
+
+
+def get_num_trainers():
+    """Get the number of trainer processes"""
+    return len(PER_ROLE_RANK["default"])
+
+
+rpc.register_service(REGISTER_ROLE, RegisterRoleRequest, RegisterRoleResponse)
+rpc.register_service(GET_ROLE, GetRoleRequest, GetRoleResponse)
@@ -0,0 +1,228 @@
+"""Functions used by client."""
+
+import atexit
+import logging
+import os
+import socket
+import time
+
+from . import rpc
+from .constants import MAX_QUEUE_SIZE
+
+if os.name != "nt":
+    import fcntl
+    import struct
+
+
+def local_ip4_addr_list():
+    """Return a set of IPv4 address
+
+    You can use
+    `logging.getLogger("dgl-distributed-socket").setLevel(logging.WARNING+1)`
+    to disable the warning here
+    """
+    assert os.name != "nt", "Do not support Windows rpc yet."
+    nic = set()
+    logger = logging.getLogger("dgl-distributed-socket")
+    for if_nidx in socket.if_nameindex():
+        name = if_nidx[1]
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        try:
+            ip_of_ni = fcntl.ioctl(
+                sock.fileno(),
+                0x8915,  # SIOCGIFADDR
+                struct.pack("256s", name[:15].encode("UTF-8")),
+            )
+        except OSError as e:
+            if e.errno == 99:  # EADDRNOTAVAIL
+                logger.warning(
+                    "Warning! Interface: %s \n"
+                    "IP address not available for interface.",
+                    name,
+                )
+                continue
+            raise e
+
+        ip_addr = socket.inet_ntoa(ip_of_ni[20:24])
+        nic.add(ip_addr)
+    return nic
+
+
+def get_local_machine_id(server_namebook):
+    """Given server_namebook, find local machine ID
+
+    Parameters
+    ----------
+    server_namebook: dict
+        IP address namebook of server nodes, where key is the server's ID
+        (start from 0) and value is the server's machine_id, IP address,
+        port, and group_count, e.g.,
+
+          {0:'[0, '172.31.40.143', 30050, 2],
+           1:'[0, '172.31.40.143', 30051, 2],
+           2:'[1, '172.31.36.140', 30050, 2],
+           3:'[1, '172.31.36.140', 30051, 2],
+           4:'[2, '172.31.47.147', 30050, 2],
+           5:'[2, '172.31.47.147', 30051, 2],
+           6:'[3, '172.31.30.180', 30050, 2],
+           7:'[3, '172.31.30.180', 30051, 2]}
+
+    Returns
+    -------
+    int
+        local machine ID
+    """
+    res = 0
+    ip_list = local_ip4_addr_list()
+    for _, data in server_namebook.items():
+        machine_id = data[0]
+        ip_addr = data[1]
+        if ip_addr in ip_list:
+            res = machine_id
+            break
+    return res
+
+
+def get_local_usable_addr(probe_addr):
+    """Get local usable IP and port
+
+    Returns
+    -------
+    str
+        IP address, e.g., '192.168.8.12:50051'
+    """
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        # should get the address on the same subnet as probe_addr's
+        sock.connect((probe_addr, 1))
+        ip_addr = sock.getsockname()[0]
+    except ValueError:
+        ip_addr = "127.0.0.1"
+    finally:
+        sock.close()
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    sock.listen(1)
+    port = sock.getsockname()[1]
+    sock.close()
+
+    return ip_addr + ":" + str(port)
+
+
+def connect_to_server(
+    ip_config,
+    num_servers,
+    max_queue_size=MAX_QUEUE_SIZE,
+    group_id=0,
+):
+    """Connect this client to server.
+
+    Parameters
+    ----------
+    ip_config : str
+        Path of server IP configuration file.
+    num_servers : int
+        server count on each machine.
+    max_queue_size : int
+        Maximal size (bytes) of client queue buffer (~20 GB on default).
+        Note that the 20 GB is just an upper-bound and DGL uses zero-copy and
+        it will not allocate 20GB memory at once.
+    group_id : int
+        Indicates which group this client belongs to. Clients that are
+        booted together in each launch are gathered as a group and should
+        have same unique group_id.
+
+    Raises
+    ------
+    ConnectionError : If anything wrong with the connection.
+    """
+    assert num_servers > 0, (
+        "num_servers (%d) must be a positive number." % num_servers
+    )
+    assert max_queue_size > 0, (
+        "queue_size (%d) cannot be a negative number." % max_queue_size
+    )
+    # Register some basic service
+    rpc.register_service(
+        rpc.CLIENT_REGISTER,
+        rpc.ClientRegisterRequest,
+        rpc.ClientRegisterResponse,
+    )
+    rpc.register_service(rpc.SHUT_DOWN_SERVER, rpc.ShutDownRequest, None)
+    rpc.register_service(
+        rpc.GET_NUM_CLIENT,
+        rpc.GetNumberClientsRequest,
+        rpc.GetNumberClientsResponse,
+    )
+    rpc.register_service(
+        rpc.CLIENT_BARRIER, rpc.ClientBarrierRequest, rpc.ClientBarrierResponse
+    )
+    rpc.register_sig_handler()
+    server_namebook = rpc.read_ip_config(ip_config, num_servers)
+    num_servers = len(server_namebook)
+    rpc.set_num_server(num_servers)
+    # group_count means how many servers
+    # (main_server + bakcup_server) in total inside a machine.
+    group_count = []
+    max_machine_id = 0
+    for server_info in server_namebook.values():
+        group_count.append(server_info[3])
+        if server_info[0] > max_machine_id:
+            max_machine_id = server_info[0]
+    rpc.set_num_server_per_machine(group_count[0])
+    num_machines = max_machine_id + 1
+    rpc.set_num_machines(num_machines)
+    machine_id = get_local_machine_id(server_namebook)
+    rpc.set_machine_id(machine_id)
+    rpc.set_group_id(group_id)
+    rpc.create_sender(max_queue_size)
+    rpc.create_receiver(max_queue_size)
+    # Get connected with all server nodes
+    max_try_times = int(os.environ.get("DGL_DIST_MAX_TRY_TIMES", 1024))
+    for server_id, addr in server_namebook.items():
+        server_ip = addr[1]
+        server_port = addr[2]
+        try_times = 0
+        while not rpc.connect_receiver(server_ip, server_port, server_id):
+            try_times += 1
+            if try_times % 200 == 0:
+                print(
+                    "Client is trying to connect server receiver: {}:{}".format(
+                        server_ip, server_port
+                    )
+                )
+            if try_times >= max_try_times:
+                raise rpc.DistConnectError(
+                    max_try_times, server_ip, server_port
+                )
+            time.sleep(3)
+    if not rpc.connect_receiver_finalize(max_try_times):
+        raise rpc.DistConnectError(max_try_times)
+    # Get local usable IP address and port
+    ip_addr = get_local_usable_addr(server_ip)
+    client_ip, client_port = ip_addr.split(":")
+    # Register client on server
+    register_req = rpc.ClientRegisterRequest(ip_addr)
+    for server_id in range(num_servers):
+        rpc.send_request(server_id, register_req)
+    # wait server connect back
+    rpc.wait_for_senders(client_ip, client_port, num_servers)
+    print(
+        "Client [{}] waits on {}:{}".format(os.getpid(), client_ip, client_port)
+    )
+    # recv client ID from server
+    res = rpc.recv_response()
+    rpc.set_rank(res.client_id)
+    print(
+        "Machine (%d) group (%d) client (%d) connect to server successfuly!"
+        % (machine_id, group_id, rpc.get_rank())
+    )
+    # get total number of client
+    get_client_num_req = rpc.GetNumberClientsRequest(rpc.get_rank())
+    rpc.send_request(0, get_client_num_req)
+    res = rpc.recv_response()
+    rpc.set_num_client(res.num_client)
+    from .dist_context import exit_client, set_initialized
+
+    atexit.register(exit_client)
+    set_initialized(True)
@@ -0,0 +1,150 @@
+"""Functions used by server."""
+
+import os
+import time
+
+from ..base import DGLError
+from . import rpc
+from .constants import MAX_QUEUE_SIZE, SERVER_EXIT
+
+
+def start_server(
+    server_id,
+    ip_config,
+    num_servers,
+    num_clients,
+    server_state,
+    max_queue_size=MAX_QUEUE_SIZE,
+):
+    """Start DGL server, which will be shared with all the rpc services.
+
+    This is a blocking function -- it returns only when the server shutdown.
+
+    Parameters
+    ----------
+    server_id : int
+        Current server ID (starts from 0).
+    ip_config : str
+        Path of IP configuration file.
+    num_servers : int
+        Server count on each machine.
+    num_clients : int
+        Total number of clients that will be connected to the server.
+        Note that, we do not support dynamic connection for now. It means
+        that when all the clients connect to server, no client will can be added
+        to the cluster.
+    server_state : ServerSate object
+        Store in main data used by server.
+    max_queue_size : int
+        Maximal size (bytes) of server queue buffer (~20 GB on default).
+        Note that the 20 GB is just an upper-bound because DGL uses zero-copy and
+        it will not allocate 20GB memory at once.
+    """
+    assert server_id >= 0, (
+        "server_id (%d) cannot be a negative number." % server_id
+    )
+    assert num_servers > 0, (
+        "num_servers (%d) must be a positive number." % num_servers
+    )
+    assert num_clients >= 0, (
+        "num_client (%d) cannot be a negative number." % num_clients
+    )
+    assert max_queue_size > 0, (
+        "queue_size (%d) cannot be a negative number." % max_queue_size
+    )
+    # Register signal handler.
+    rpc.register_sig_handler()
+    # Register some basic services
+    rpc.register_service(
+        rpc.CLIENT_REGISTER,
+        rpc.ClientRegisterRequest,
+        rpc.ClientRegisterResponse,
+    )
+    rpc.register_service(rpc.SHUT_DOWN_SERVER, rpc.ShutDownRequest, None)
+    rpc.register_service(
+        rpc.GET_NUM_CLIENT,
+        rpc.GetNumberClientsRequest,
+        rpc.GetNumberClientsResponse,
+    )
+    rpc.register_service(
+        rpc.CLIENT_BARRIER, rpc.ClientBarrierRequest, rpc.ClientBarrierResponse
+    )
+    rpc.set_rank(server_id)
+    server_namebook = rpc.read_ip_config(ip_config, num_servers)
+    machine_id = server_namebook[server_id][0]
+    rpc.set_machine_id(machine_id)
+    ip_addr = server_namebook[server_id][1]
+    port = server_namebook[server_id][2]
+    rpc.create_sender(max_queue_size)
+    rpc.create_receiver(max_queue_size)
+    # wait all the senders connect to server.
+    # Once all the senders connect to server, server will not
+    # accept new sender's connection
+    print(
+        "Server is waiting for connections on [{}:{}]...".format(ip_addr, port)
+    )
+    rpc.wait_for_senders(ip_addr, port, num_clients)
+    rpc.set_num_client(num_clients)
+    recv_clients = {}
+    while True:
+        # go through if any client group is ready for connection
+        for group_id in list(recv_clients.keys()):
+            ips = recv_clients[group_id]
+            if len(ips) < rpc.get_num_client():
+                continue
+
+            del recv_clients[group_id]
+            # a new client group is ready
+            ips.sort()
+            client_namebook = dict(enumerate(ips))
+            time.sleep(3)  # wait for clients' receivers ready
+            max_try_times = int(os.environ.get("DGL_DIST_MAX_TRY_TIMES", 120))
+            for client_id, addr in client_namebook.items():
+                client_ip, client_port = addr.split(":")
+                try_times = 0
+                while not rpc.connect_receiver(
+                    client_ip, client_port, client_id, group_id
+                ):
+                    try_times += 1
+                    if try_times % 200 == 0:
+                        print(
+                            "Server~{} is trying to connect client receiver: {}:{}".format(
+                                server_id, client_ip, client_port
+                            )
+                        )
+                    if try_times >= max_try_times:
+                        raise rpc.DistConnectError(
+                            max_try_times, client_ip, client_port
+                        )
+                    time.sleep(1)
+            if not rpc.connect_receiver_finalize(max_try_times):
+                raise rpc.DistConnectError(max_try_times)
+            if rpc.get_rank() == 0:  # server_0 send all the IDs
+                for client_id, _ in client_namebook.items():
+                    register_res = rpc.ClientRegisterResponse(client_id)
+                    rpc.send_response(client_id, register_res, group_id)
+        # receive incomming client requests
+        timeout = 60 * 1000  # in milliseconds
+        req, client_id, group_id = rpc.recv_request(timeout)
+        if req is None:
+            continue
+        if isinstance(req, rpc.ClientRegisterRequest):
+            if group_id not in recv_clients:
+                recv_clients[group_id] = []
+            recv_clients[group_id].append(req.ip_addr)
+            continue
+
+        res = req.process_request(server_state)
+        if res is not None:
+            if isinstance(res, list):
+                for response in res:
+                    target_id, res_data = response
+                    rpc.send_response(target_id, res_data, group_id)
+            elif isinstance(res, str):
+                if res == SERVER_EXIT:
+                    print("Server is exiting...")
+                    return
+                else:
+                    raise DGLError("Unexpected response: {}".format(res))
+            else:
+                rpc.send_response(client_id, res, group_id)
@@ -0,0 +1,81 @@
+"""Server data"""
+
+from .._ffi.function import _init_api
+
+# Remove C++ bindings for now, since not used
+
+
+class ServerState:
+    """Data stored in one DGL server.
+
+    In a distributed setting, DGL partitions all data associated with the graph
+    (e.g., node and edge features, graph structure, etc.) to multiple partitions,
+    each handled by one DGL server. Hence, the ServerState class includes all
+    the data associated with a graph partition.
+
+    Under some setup, users may want to deploy servers in a heterogeneous way
+    -- servers are further divided into special groups for fetching/updating
+    node/edge data and for sampling/querying on graph structure respectively.
+    In this case, the ServerState can be configured to include only node/edge
+    data or graph structure.
+
+    Each machine can have multiple server and client processes, but only one
+    server is the *master* server while all the others are backup servers. All
+    clients and backup servers share the state of the master server via shared
+    memory, which means the ServerState class must be serializable and large
+    bulk data (e.g., node/edge features) must be stored in NDArray to leverage
+    shared memory.
+
+    Attributes
+    ----------
+    kv_store : KVServer
+        reference for KVServer
+    graph : DGLGraph
+        Graph structure of one partition
+    total_num_nodes : int
+        Total number of nodes
+    total_num_edges : int
+        Total number of edges
+    partition_book : GraphPartitionBook
+        Graph Partition book
+    use_graphbolt : bool
+        Whether to use graphbolt for dataloading.
+    """
+
+    def __init__(self, kv_store, local_g, partition_book, use_graphbolt=False):
+        self._kv_store = kv_store
+        self._graph = local_g
+        self.partition_book = partition_book
+        self._roles = {}
+        self._use_graphbolt = use_graphbolt
+
+    @property
+    def roles(self):
+        """Roles of the client processes"""
+        return self._roles
+
+    @property
+    def kv_store(self):
+        """Get data store."""
+        return self._kv_store
+
+    @kv_store.setter
+    def kv_store(self, kv_store):
+        self._kv_store = kv_store
+
+    @property
+    def graph(self):
+        """Get graph data."""
+        return self._graph
+
+    @graph.setter
+    def graph(self, graph):
+        self._graph = graph
+
+    @property
+    def use_graphbolt(self):
+        """Whether to use graphbolt for dataloading."""
+        return self._use_graphbolt
+
+
+_init_api("dgl.distributed.server_state")
@@ -0,0 +1,26 @@
+"""Define utility functions for shared memory."""
+
+from .. import backend as F, ndarray as nd
+from .._ffi.ndarray import empty_shared_mem
+
+DTYPE_DICT = F.data_type_dict
+DTYPE_DICT = {DTYPE_DICT[key]: key for key in DTYPE_DICT}
+
+
+def _get_ndata_path(graph_name, ndata_name):
+    return "/" + graph_name + "_node_" + ndata_name
+
+
+def _get_edata_path(graph_name, edata_name):
+    return "/" + graph_name + "_edge_" + edata_name
+
+
+def _to_shared_mem(arr, name):
+    dlpack = F.zerocopy_to_dlpack(arr)
+    dgl_tensor = nd.from_dlpack(dlpack)
+    new_arr = empty_shared_mem(
+        name, True, F.shape(arr), DTYPE_DICT[F.dtype(arr)]
+    )
+    dgl_tensor.copyto(new_arr)
+    dlpack = new_arr.to_dlpack()
+    return F.zerocopy_from_dlpack(dlpack)
@@ -0,0 +1,127 @@
+"""Define a fake kvstore
+
+This kvstore is used when running in the standalone mode
+"""
+
+from .. import backend as F
+
+
+class KVClient(object):
+    """The fake KVStore client.
+
+    This is to mimic the distributed KVStore client. It's used for DistGraph
+    in standalone mode.
+    """
+
+    def __init__(self):
+        self._data = {}
+        self._all_possible_part_policy = {}
+        self._push_handlers = {}
+        self._pull_handlers = {}
+        # Store all graph data name
+        self._gdata_name_list = set()
+
+    @property
+    def all_possible_part_policy(self):
+        """Get all possible partition policies"""
+        return self._all_possible_part_policy
+
+    @property
+    def num_servers(self):
+        """Get the number of servers"""
+        return 1
+
+    def barrier(self):
+        """barrier"""
+
+    def register_push_handler(self, name, func):
+        """register push handler"""
+        self._push_handlers[name] = func
+
+    def register_pull_handler(self, name, func):
+        """register pull handler"""
+        self._pull_handlers[name] = func
+
+    def add_data(self, name, tensor, part_policy):
+        """add data to the client"""
+        self._data[name] = tensor
+        self._gdata_name_list.add(name)
+        if part_policy.policy_str not in self._all_possible_part_policy:
+            self._all_possible_part_policy[part_policy.policy_str] = part_policy
+
+    def init_data(
+        self, name, shape, dtype, part_policy, init_func, is_gdata=True
+    ):
+        """add new data to the client"""
+        self._data[name] = init_func(shape, dtype)
+        if part_policy.policy_str not in self._all_possible_part_policy:
+            self._all_possible_part_policy[part_policy.policy_str] = part_policy
+        if is_gdata:
+            self._gdata_name_list.add(name)
+
+    def delete_data(self, name):
+        """delete the data"""
+        del self._data[name]
+        if name in self._gdata_name_list:
+            self._gdata_name_list.remove(name)
+
+    def data_name_list(self):
+        """get the names of all data"""
+        return list(self._data.keys())
+
+    def gdata_name_list(self):
+        """get the names of graph data"""
+        return list(self._gdata_name_list)
+
+    def get_data_meta(self, name):
+        """get the metadata of data"""
+        return F.dtype(self._data[name]), F.shape(self._data[name]), None
+
+    def push(self, name, id_tensor, data_tensor):
+        """push data to kvstore"""
+        if name in self._push_handlers:
+            self._push_handlers[name](self._data, name, id_tensor, data_tensor)
+        else:
+            F.scatter_row_inplace(self._data[name], id_tensor, data_tensor)
+
+    def pull(self, name, id_tensor):
+        """pull data from kvstore"""
+        if name in self._pull_handlers:
+            return self._pull_handlers[name](self._data, name, id_tensor)
+        else:
+            return F.gather_row(self._data[name], id_tensor)
+
+    def map_shared_data(self, partition_book):
+        """Mapping shared-memory tensor from server to client."""
+
+    def count_nonzero(self, name):
+        """Count nonzero value by pull request from KVServers.
+
+        Parameters
+        ----------
+        name : str
+            data name
+
+        Returns
+        -------
+        int
+            the number of nonzero in this data.
+        """
+        return F.count_nonzero(self._data[name])
+
+    @property
+    def data_store(self):
+        """Return the local partition of the data storage.
+
+        Returns
+        -------
+        dict[str, Tensor]
+            The tensor storages of the local partition.
+        """
+        return self._data
+
+    def union(self, operand1_name, operand2_name, output_name):
+        """Compute the union of two mask arrays in the KVStore."""
+        self._data[output_name][:] = (
+            self._data[operand1_name] | self._data[operand2_name]
+        )