WholeMemory 实现细节#

WholeMemory 简介 所述,WholeMemory 有两种位置和三种类型。因此,总共有六种 WholeMemory。

类型 连续 连续 分块 分块 分布式 分布式
位置 设备 主机 设备 主机 设备 主机
分配者 每个 第一个 每个 第一个 每个 每个
分配 API 驱动程序 主机 运行时 主机 运行时 运行时
IPC 映射 Unix 文件描述符 mmap cudaIpc mmap 无 IPC 映射 无 IPC 映射

对于“连续”和“分块”类型的 WholeMemory,所有内存都映射到每个 GPU,因此这两种类型都是“映射”WholeMemory,与“分布式”WholeMemory(所有内存均未映射)形成对比。

WholeMemory 布局#

由于单个 WholeMemory 对象的底层内存可能位于多个 GPU 设备上,WholeGraph 库会将数据分区到这些 GPU 设备中。分区方法保证每个 GPU 可以访问整个内存的连续一部分。这里的“可以访问”意味着可以直接从 CUDA 内核访问,但内存不必物理地位于该 GPU 上。例如,它可以位于主机内存或可以通过 P2P 访问的其他 GPU 设备内存上。在这种情况下,存储的数据有其自身的粒度,不应被拆分。数据粒度可以在创建 WholeMemory 时指定。然后每个数据粒度可以被视为一个数据块。

下图显示了 15 个数据块在 4 个 GPU 上的布局。WholeMemory Layout

对于 WholeMemory 张量,它们可以是 1D 或 2D 张量。对于 1D 张量,数据粒度是一个元素。对于 2D 张量,数据粒度是其 1D 张量。布局将如下所示:WholeMemory Tensor Layout

WholeMemory 分配#

由于有六种类型的 WholeMemory,每种类型的分配过程如下

设备连续 WholeMemory#

对于设备连续 WholeMemory,首先在每个 GPU 中预留一个虚拟地址空间范围,该范围覆盖整个内存范围。然后在每个 GPU 中分配一部分物理内存,如下图所示。Device Continuous WholeMemory Allocation Step 1 之后,每个 GPU 收集所有 GPU 的内存句柄,并将它们映射到预留的地址空间。Device Continuous WholeMemory Allocation Step 2

设备分块 WholeMemory#

对于设备分块 WholeMemory,首先每个 GPU 使用 CUDA 运行时 API 分配其自己的内存部分,这将为其自己的内存创建虚拟地址空间和物理内存。Device Chunked WholeMemory Allocation Step 1 每个 GPU 收集所有其他 GPU 的内存的 Ipc 句柄,并将其映射到其自己的虚拟地址空间。Device Chunked WholeMemory Allocation Step 2

主机映射 WholeMemory#

对于主机,连续和分块使用相同的方法。首先,排名并分配主机物理内存,然后将其共享给所有排名。Host Mapped WholeMemory Allocation Step 1 然后每个排名将该主机内存注册到 GPU 地址空间。Host Mapped WholeMemory Allocation Step 2

分布式 WholeMemory#

对于分布式 WholeMemory,每个 GPU 只需 malloc 其自己的内存部分,无需与其他 GPU 共享。Distributed WholeMemory Allocation