WholeMemory#

WholeMemory 可以被视为 GPU 内存的整体视图。无论底层数据如何分布在多个 GPU 上，WholeMemory 都提供一个内存实例的句柄。WholeMemory 假设每个 GPU 由一个独立的进程控制。

WholeMemory 基础#

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

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

由于 WholeMemory 由一组 GPU 拥有，因此需要指定这组 GPU。这可以通过创建 WholeMemory 通信器 并在创建 WholeMemory 时指定该 WholeMemory 通信器来完成。

2. 指定内存位置#

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

  • 主机内存：将使用锁定（pinned）主机内存作为底层存储。

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

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

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

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

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

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

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

WholeMemory 通信器#

WholeMemory 通信器有两个主要目的

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

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

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

WholeMemory 粒度#

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

WholeMemory 映射#

由于 WholeMemory 提供 GPU 对内存的整体视图，通常需要映射来访问 WholeMemory。不同类型的 WholeMemory 支持不同的映射方法，其名称就说明了这一点。一些支持的映射包括

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

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

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

WholeMemory 操作#

可以对 WholeMemory 执行一些操作。这些操作基于 WholeMemory 的映射。

本地操作#

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

加载 / 存储#

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

收集 / 分散#

WholeMemory 还支持收集 / 分散操作，这些操作通常在 WholeMemory 张量 上进行。

WholeMemory 张量#

与 PyTorch 相比，WholeMemory 类似于 PyTorch 存储 (Storage)，而 WholeMemory 张量类似于 PyTorch 张量 (Tensor)。目前，WholeMemory 仅支持 1D 和 2D 张量，即数组和矩阵。只有第一维被分区。

WholeMemory 嵌入#

WholeMemory 嵌入就像 2D WholeMemory 张量，增加了缓存支持和稀疏优化器支持。

缓存支持#

WholeMemory 嵌入支持缓存。要创建带缓存的 WholeMemory 嵌入，首先需要创建 WholeMemory CachePolicy。WholeMemoryCachePolicy 可以通过以下字段创建

  • WholeMemory 通信器：WholeMemory CachePolicy 也需要 WholeMemory 通信器。此 WholeMemory 通信器定义了缓存所有嵌入的 GPU 集合。它可以与创建 WholeMemory 嵌入时使用的 WholeMemory 通信器相同。

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

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

  • 访问类型：访问类型可以是只读或读写。

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

有两种最常用的缓存。它们是

  • 设备缓存的主机内存：当 WholeMemory CachePolicy 的 WholeMemory 通信器与创建 WholeMemory 嵌入时使用的 WholeMemory 通信器相同时，意味着缓存与 WholeMemory 嵌入具有相同的 GPU 集合。因此，每个 GPU 只缓存其自身部分的原始嵌入。通常，当原始 WholeMemory 嵌入位于主机内存上，而缓存位于设备内存上时，每个 GPU 只缓存其自身部分的主机内存。

  • 本地缓存的全局内存：WholeMemory CachePolicy 的 WholeMemory 通信器也可以是 WholeMemory 嵌入的 WholeMemory 通信器的子集。在这种情况下，这组 GPU 共同缓存所有嵌入。通常，原始 WholeMemory 嵌入分布在不同的机器节点上，我们希望在本地机器或本地 GPU 中缓存一些嵌入，那么这组 GPU 可以是本地机器上的所有 GPU。本地缓存的全局内存仅支持只读。

WholeMemory 嵌入稀疏优化器#

WholeMemory 嵌入的另一个特性是它支持嵌入训练。为了高效地训练大型嵌入表，需要一个稀疏优化器。WholeMemory 嵌入稀疏优化器可以在缓存或非缓存的 WholeMemory 嵌入上运行。目前支持的优化器包括 SGD、Adam、RMSProp 和 AdaGrad。