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 张量,它们可以是 1D 或 2D 张量。对于 1D 张量,数据粒度是一个元素。对于 2D 张量,数据粒度是其 1D 张量。布局将如下所示:
WholeMemory 分配#
由于有六种类型的 WholeMemory,每种类型的分配过程如下
设备连续 WholeMemory#
对于设备连续 WholeMemory,首先在每个 GPU 中预留一个虚拟地址空间范围,该范围覆盖整个内存范围。然后在每个 GPU 中分配一部分物理内存,如下图所示。 之后,每个 GPU 收集所有 GPU 的内存句柄,并将它们映射到预留的地址空间。
设备分块 WholeMemory#
对于设备分块 WholeMemory,首先每个 GPU 使用 CUDA 运行时 API 分配其自己的内存部分,这将为其自己的内存创建虚拟地址空间和物理内存。 每个 GPU 收集所有其他 GPU 的内存的 Ipc 句柄,并将其映射到其自己的虚拟地址空间。
主机映射 WholeMemory#
对于主机,连续和分块使用相同的方法。首先,排名并分配主机物理内存,然后将其共享给所有排名。 然后每个排名将该主机内存注册到 GPU 地址空间。
分布式 WholeMemory#
对于分布式 WholeMemory,每个 GPU 只需 malloc 其自己的内存部分,无需与其他 GPU 共享。