資料集訓練規模持續擴大，參數數量已超過數十億。雖然有些模型可以完整地裝入系統記憶體，但較大的模型則不行。在此情況下，資料載入器需要透過各種方法存取位於快閃儲存裝置上的模型。應對方法之一就是儲存在 SSD 上的記憶體對應檔案。這樣，資料載入器就能快速存取檔案，如同存取記憶體中的檔案一樣，但 CPU 和軟體棧的負擔會大大降低訓練系統的效能。這便是大型加速器記憶體（BaM）* 以及 GPU 主動直接儲存（GIDS）* 資料載入器的獨到之處。
什麼是 BaM 和 GIDS？
BaM 是一種利用 SSD 的低延遲、極高吞吐量、大容量和耐用性的系統架構。BaM 的目標是提供高效的抽象，使 GPU 執行緒能夠對 SSD 上的資料集進行細粒度存取，並實現比要求 CPU 提供儲存請求以服務 GPU 的解決方案更高的效能。BaM 加速使用自訂的儲存驅動程式，該驅動程式專為實現 GPU 的固有平行性以直接存取儲存裝置而設計。BaM 與 NVIDIA Magnum IO™ GPUDirect® Storage（GDS）不同，因為 BaM 不依賴 CPU 來準備從 GPU 到 SSD 的通訊。
如下文所述，美光以前的產品也曾採用 NVIDIA GDS：
- Micron® 9400 NVMe™ SSD 透過 NVIDIA Magnum IO GPUDirect 儲存平台實現高效能
- 美光與 Magnum IO GPUDirect 儲存技術合作為 AI 和 ML 帶來顛覆業界的創新
GIDS 資料載入器基於 BaM 子系統構建，可滿足 GPU 加速的圖神經網路（GNN）訓練對記憶體容量的要求，同時還能屏蔽儲存裝置延遲。GIDS 將圖的特徵資料儲存在 SSD 上，因為對於大規模圖形而言，該資料通常是整個圖形資料集中最大的部分。與特徵資料相比，圖形結構資料通常要小得多，本身被固定在系統記憶體中，以實現快速 GPU 圖形取樣。最後，GIDS 資料載入器會在 GPU 記憶體中為最近存取的節點分配軟體定義的快取，以減少儲存裝置存取。
使用 GIDS 進行圖神經網路訓練
為了展示 BaM 和 GIDS 的優勢，我們使用 Illinois Graph Benchmark（IGB）異構全資料集進行 GNN 訓練。該資料集有 2.28TB 之大，大多數平台的系統記憶體都裝不下。我們使用單個 NVIDIA A100 80GB Tensor Core GPU 進行了 100 次迭代計時訓練，並變更 SSD 的數量，以提供廣泛的結果，如圖 1 和表 1 所示。
圖 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 倍。
我們還進行了基線計算，使用記憶體映射抽象和深度圖庫（DGL）資料載入器存取特徵資料。由於這種存取特徵資料的方法需要使用 CPU 軟體棧，而不是 GPU 直接存取，從中可以看出 CPU 軟體棧在訓練期間保持 GPU 飽和的效率低下。與基線相比，使用 GIDS 的 1 個美光 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 |
2 個 AMD EPYC™ 7702（64 核心） |
記憶體 |
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，DGL 1.1.2，Pytorch 2.1 |
參考連結
大型加速器記憶體文章與 GitHub
BaM 系統架構中的 GPU 主動隨需高傳輸量儲存裝置存取（arxiv.org）
GitHub - ZaidQureshi/bam
GPU 主動直接儲存文章和 GitHub
利用 GPU 主動直接儲存存取加速 GNN 框架中的取樣和聚合操作（arxiv.org）
GitHub - jeongminpark417/GIDS
GitHub - IllinoisGraphBenchmark/IGB-Datasets：最大的真實開源圖形資料集——在 IBM - 伊利諾伊發現加速器研究所（IBM-Illinois Discovery Accelerator Institute）和 Amazon 研究獎的支持下，與 NVIDIA Research 合作完成。
*註：NVIDIA 大型加速器記憶體（BaM）和 NVIDIA GPU 主動直接儲存（GIDS）資料載入器是 NVIDIA Research 的原型專案，不打算全面發佈。
