WholeGraph 介绍#

WholeGraph 帮助训练大规模图神经网络 (GNN)。WholeGraph 提供了称为 WholeMemory 的底层存储结构。WholeMemory 是一种类似 Tensor 的存储，并提供多 GPU 支持。它针对像 DGX A100 服务器这样的 NVLink 系统进行了优化。通过与 cuGraph、cuGraph-DGL、cuGraph-PyG 以及上游的 DGL 和 PyG 协同工作，构建 GNN 应用程序将变得容易。

WholeMemory#

WholeMemory 可以被视为 GPU 内存的整体视图。无论底层数据如何跨多个 GPU 存储，WholeMemory 都暴露一个内存实例的句柄。WholeMemory 假定使用单独的进程来控制每个 GPU。

WholeMemory 基础知识#

要定义 WholeMemory，我们需要指定以下内容

1. 指定处理内存的 GPU 集#

由于 WholeMemory 归一组 GPU 所有，因此必须指定这组 GPU。这可以通过创建 WholeMemory Communicator 并在创建 WholeMemory 时指定 WholeMemory Communicator 来完成。

2. 指定内存位置#

虽然 WholeMemory 归一组 GPU 所有，但内存本身可以位于主机内存或设备内存中。需要指定内存位置，可以指定两种类型的位置。

  • 主机内存：将使用固定主机内存作为底层存储。

  • 设备内存：将使用 GPU 设备内存作为底层存储。

3. 指定内存的地址映射模式#

由于 WholeMemory 归多个 GPU 所有，每个 GPU 都将访问整个内存空间，因此我们需要地址映射。有三种类型的地址映射模式（也称为 WholeMemory 类型），它们是

  • 连续模式 (Continuous)：每个 GPU 的所有内存都将映射到每个 GPU 的单个连续内存地址空间中。在此模式下，每个 GPU 都可以使用单个指针和偏移量直接访问整个内存，就像使用普通设备内存一样。软件不会看到差异。硬件对等访问将处理底层通信。

  • 分块模式 (Chunked)：每个 GPU 的内存将映射到不同的内存块中，每个 GPU 一个块。在此模式下，也支持直接访问，但不是使用单个指针。软件将看到分块内存。然而，一个抽象层可以帮助隐藏这一点。

  • 分布式模式 (Distributed)：来自其他 GPU 的内存不会映射到当前 GPU 中，因此不支持直接访问。要访问其他 GPU 的内存，需要显式通信。

要了解 WholeMemory 位置和 WholeMemory 类型的更多详细信息，请参阅 WholeMemory 实现细节

WholeMemory Communicator#

WholeMemory Communicator 有两个主要目的

  • 定义一组在 WholeMemory 上协同工作的 GPU。WholeMemory Communicator 由所有希望协同工作的 GPU 创建。只要所需的 GPU 集相同，WholeMemory Communicator 就可以重复使用。

  • 提供 WholeMemory 所需的底层通信通道。WholeMemory 在创建期间以及对某些类型的 WholeMemory 进行某些操作时，可能需要 GPU 之间的通信器。

要创建 WholeMemory Communicator，首先需要创建一个 WholeMemory Unique ID，这通常由 GPU 集中的第一个 GPU 创建，然后广播给所有希望协同工作的 GPU。然后此通信器中的所有 GPU 将使用此 WholeMemory Unique ID、当前 GPU 的 rank 以及所有 GPU 计数来调用 WholeMemory Communicator 创建函数。

WholeMemory 粒度#

由于底层存储可能会在物理上分割到多个 GPU 中，这通常不是单个用户数据块内所希望的。为解决此问题，在创建 WholeMemory 时可以指定数据的粒度。然后 WholeMemory 被视为具有相同粒度的多个块，并且在粒度内部不会被分割。

WholeMemory 映射#

由于 WholeMemory 向 GPU 提供内存的整体视图，因此访问 WholeMemory 通常需要映射。不同类型的 WholeMemory 支持不同的映射方法，与其名称相符。支持的一些映射包括

  • 所有 WholeMemory 类型都支持映射本地 GPU 负责的内存范围。也就是说，每个 rank 可以在所有 WholeMemory 类型中直接访问“本地”内存。这里的“本地”内存不一定非得在当前 GPU 的内存上，它可以位于主机内存甚至可能在其他 GPU 上，但它保证可以被当前 GPU 直接访问。

  • 分块模式 (Chunked) 和连续模式 (Continuous) 的 WholeMemory 也支持分块映射。也就是说，所有 GPU 的内存都可以映射到当前 GPU，每个 GPU 一个连续块。当前 GPU 可以直接访问每个块。但不同块的内存不保证连续。

  • 连续模式 (Continuous) 的 WholeMemory 可以映射到连续的内存空间中。也就是说，所有 GPU 的内存都映射到单个虚拟内存范围，访问此内存的不同位置将物理上访问不同的 GPU。此映射由硬件（CPU 页表或 GPU 页表）处理。

WholeMemory 操作#

WholeMemory 上可以执行一些操作。它们基于 WholeMemory 的映射。

本地操作#

由于所有 WholeMemory 都支持本地内存映射，因此支持本地内存操作。操作可以是读或写。像使用当前设备的 GPU 内存一样使用即可。

加载和存储#

为了方便文件操作，支持从文件加载 WholeMemory 或将 WholeMemory 存储到文件。WholeMemory 对磁盘操作使用原始二进制文件格式。对于加载，输入文件可以是单个文件或文件列表，如果是一个列表，它们将被逻辑连接在一起然后加载。对于存储，每个 GPU 将其本地内存存储到文件中，生成文件列表。

收集和分散 (Gather and Scatter)#

WholeMemory 还支持收集/分散操作，它们通常在 WholeMemory Tensor 上运行。

WholeMemory Tensor#

与 PyTorch 相比，WholeMemory 类似于 PyTorch Storage，而 WholeMemory Tensor 类似于 PyTorch Tensor。目前，WholeMemory 仅支持 1D 和 2D tensors，即数组和矩阵。仅第一个维度被分区。

WholeMemory Embedding#

WholeMemory Embedding 就像一个 2D WholeMemory Tensor，增加了两个功能。它们支持缓存和稀疏优化器。

缓存支持#

要创建带缓存的 WholeMemory Embedding，首先需要创建 WholeMemory CachePolicy。WholeMemory CachePolicy 可以通过以下字段创建

  • WholeMemory Communicator：WholeMemory CachePolicy 也需要 WholeMemory Communicator。WholeMemory Communicator 定义了缓存所有 Embedding 的 GPU 集。它可以与创建 WholeMemory Embedding 时使用的 WholeMemory Communicator 相同。

  • WholeMemory 类型：WholeMemory CachePolicy 使用 WholeMemory 类型来指定缓存的 WholeMemory 类型。

  • WholeMemory 位置：WholeMemory CachePolicy 使用 WholeMemory 位置来指定缓存的位置。

  • 访问类型：访问类型可以是只读 (readonly) 或读写 (readwrite)。

  • 缓存比例：指定缓存将使用多少内存。此比例针对缓存整个 embedding 的每个 GPU 集计算。

两种最常用的缓存是

  • 设备缓存的主机内存 (Device cached host memory)：当 Cache Policy 的 WholeMemory Communicator 与创建 WholeMemory Embedding 时使用的 WholeMemory Communicator 相同时，这意味着缓存与 WholeMemory Embedding 具有相同的 GPU 集。因此，每个 GPU 只缓存其自己部分的原始 Embedding。最常见的情况是，当原始 WholeMemory Embedding 位于主机内存上，而缓存位于设备内存上时，每个 GPU 只缓存其自己部分的主机内存。

  • 本地缓存的全局内存 (Local cached global memory)：WholeMemory CachePolicy 的 WholeMemory Communicator 也可以是 WholeMemory Embedding 的 WholeMemory Communicator 的一个子集。在这种情况下，GPU 的子集共同缓存所有 embedding。通常，当原始 WholeMemory Embedding 分区在不同的机器节点上，并且我们想在本地机器或本地 GPU 中缓存一些 embedding 时，GPU 的子集可以是本地机器中的所有 GPU。对于本地缓存的全局内存，仅支持只读。

WholeMemory Embedding 稀疏优化器#

WholeMemory Embedding 的另一个特点是它支持 embedding 训练。为了高效训练大型 embedding 表，需要稀疏优化器。WholeMemory Embedding 稀疏优化器可以在带缓存或不带缓存的 WholeMemory Embedding 上运行。当前支持的优化器包括 SGD、Adam、RMSProp 和 AdaGrad。

图结构#

WholeGraph 中的图结构也基于 WholeMemory。在 WholeGraph 中，图以 CSR 格式存储。ROW_INDEX (记为 csr_row_ptr) 和 COL_INDEX (记为 csr_col_ind) 都存储在 WholeMemory Tensor 中。因此，加载图结构可以使用 WholeMemory Tensor 加载机制。