数据集的训练规模不断扩大，当前已超过数十亿个参数。有些模型可以完全载入系统内存，但较大的模型无法全部载入。这种情况下，数据加载器需要通过各种方法访问闪存上的模型数据。其中一种方法是将内存映射文件存储在 SSD 上。通过这种方式，数据加载器能像访问内存中的文件一样访问这些文件，但 CPU 和软件栈的负载会增大，导致训练系统的性能大幅降低。这正是大型加速器内存 (BaM)* 和 GPU 直接启动存储 (GIDS)* 数据加载器的用武之地。
什么是 BaM 和 GIDS？
BaM 是种系统架构，充分利用了 SSD 的低延迟、高吞吐量、高密度和高耐用度等特性。BaM 旨在提供高效抽象，使 GPU 线程能够实现对 SSD 上数据集的细粒度访问。与需要 CPU 为 GPU 提供存储请求的解决方案相比，BaM 可实现更高的性能。BaM 加速使用定制的存储驱动程序，专门支持具备内在并行性的 GPU 直接访问存储设备。与 NVIDIA Magnum IO™ GPUDirect® 存储 (GDS) 相比，BaM 不依赖 CPU 来准备 GPU 到 SSD 的通信。
美光之前曾与 NVIDIA GDS 合作，具体如下：
- 美光® 9400 NVMe™ SSD 借助 NVIDIA Magnum IO GPUDirect 存储平台实现高性能
- 美光与 Magnum IO GPUDirect Storage 精诚合作，为人工智能 (AI) 和机器学习 (ML) 领域带来颠覆性创新。
GIDS 数据加载器建立在 BaM 子系统上，以满足 GPU 加速图神经网络 (GNN) 训练对内存容量的要求，同时缓解存储延迟。GIDS 通过在 SSD 上存储图的特征数据来实现这一点，对于大规模图而言，这些数据通常是整个图数据集的主要部分。与特征数据相比，图结构数据通常要小得多，可常驻在系统内存中，以便 GPU 快速执行图采样。最后，GIDS 数据加载器会在 GPU 内存上为最近访问过的节点分配软件定义的缓存，以减少存储访问。
使用 GIDS 训练图神经网络
为了展示 BaM 和 GIDS 的优势，我们使用 Illinois 图基准 (IGB) 异构全数据集执行了图神经网络训练。这个数据集的规模高达 2.28TB，大多数平台的系统内存都无法容纳。如图 1 和表 1 所示，我们使用单个 NVIDIA A100 80GB Tensor Core GPU 完成了 100 次迭代计时训练，并配置了不同的 SSD 数量，以得出各种结果。
图 1：IGB 异构全数据集的 GIDS 训练时间——100 次迭代
|
GIDS（4 个 SSD） |
GIDS（2 个 SSD） |
GIDS（1 个 SSD） |
DGL 内存映射抽象 |
采样 |
4.75 |
4.93 |
4.08 |
4.65 |
特征聚合 |
8.57 |
15.9 |
31.6 |
1,130 |
训练 |
1.98 |
1.97 |
1.87 |
2.13 |
端到端 |
15.3 |
22.8 |
37.6 |
1,143 |
表 1：IGB 异构全数据集的 GIDS 训练时间——100 次迭代
训练的第一部分是图采样，由 GPU 完成，并访问系统内存中的图结构数据（蓝色部分）。由于存储在系统内存中的结构不会因测试的不同而发生变化，因此这个值在不同测试配置中变化不大。
另一部分是实际训练时间（最右边的绿色部分）。这部分在很大程度上取决于 GPU，我们可以看到，正如预期的那样，不会因测试配置的不同而明显变化。
最重要的部分，也是差异最大的部分是特征聚合（金色部分）。由于该系统的特征数据存储在美光 9400 SSD 上，我们可以看到，从 1 个扩展到 4 个美光 9400 SSD，可以显著减少特征聚合处理时间。当我们从 1 个 SSD 扩展到 4 个 SSD 时，特征聚合能力提高了 3.68 倍。
我们还在测试中包含了基线计算，使用内存映射抽象和 Deep Graph Library (DGL) 数据加载器来访问特征数据。由于这种访问特征数据的方法需要使用 CPU 软件栈，而不是由 GPU 直接访问，因此我们可以看到，要在训练过程中保持 GPU 饱和，CPU 软件栈的效率非常低。与基线相比，使用 GIDS 的单块美光 9400 NVMe SSD 的特征抽象能力提高了 35.76 倍，4 块美光 9400 NVMe SSD 的特征抽象能力提高了 131.87 倍。图 2 和表 2 展示了这些数据的另一视图，其中显示了这些测试期间的有效带宽和 IOP。
图 2：GIDS 训练的有效带宽和 IOPS 与基线对比
|
DGL 内存映射 |
GIDS（1 个 SSD） |
GIDS（2 个 SSD） |
GIDS（4 个 SSD） |
有效带宽 (GB/s) |
0.194 |
6.9 |
13.8 |
25.6 |
实现的 IOP (M/s) |
0.049 |
1.7 |
3.4 |
6.3 |
表 2：GIDS 训练的有效带宽和 IOPS 与基线对比
随着数据集不断扩大，有必要转变模式，以便在合理的时间内训练这些模型，并利用先进的 GPU 所带来的性能提升。BaM 和 GIDS 是很好的起点，我们期待今后与更多此类系统合作。
测试系统
组件 |
详细信息 |
服务器 |
Supermicro® AS 4124GS-TNR |
CPU |
|
内存 |
1 TB 美光 DDR4-3200 |
GPU |
内存时钟：1512 MHz SM 时钟：1410 MHz |
SSD |
4 个美光 9400 MAX 6.4TB |
操作系统 |
Ubuntu 22.04 LTS，内核 5.15.0.86 |
NVIDIA 驱动程序 |
535.113.01 |
软件堆栈 |
在 NVIDIA Docker 容器中运行的 CUDA 12.2、DDL 1.1.2、Pytorch 2.1 |
参考链接
大加速器内存 (Big Accelerator Memory) 论文和 GitHub
BaM 系统架构中由 GPU 启动的按需高吞吐量存储访问 (arxiv.org)
GitHub - ZaidQureshi/bam
GPU 直接启动存储论文和 GitHub
利用 GPU 直接启动的存储访问加速 GNN 框架中的采样和聚合操作 (arxiv.org)
GitHub - jeongminpark417/GIDS
GitHub - IllinoisGraphBenchmark/IGB-数据集：大型真实开源图数据集，在 IBM-Illinois 发现加速器研究所和 Amazon 研究奖的支持下，与 NVIDIA 研究中心合作完成。
*请注意：NVIDIA 大加速器内存（Big Accelerator Memory，BaM）和 NVIDIA GPU 直接启动存储 (GIDS) 数据加载器是 NVIDIA 研究中心的原型项目，并不打算正式发布。
