美光 9400 NVMe SSD 探索使用 NVIDIA 技術的大型加速器記憶體

資料集訓練規模持續擴大，參數數量已超過數十億。雖然有些模型可以完整地裝入系統記憶體，但較大的模型則不行。在此情況下，資料載入器需要透過各種方法存取位於快閃儲存裝置上的模型。應對方法之一就是儲存在 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 次迭代

表 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 與基準對比

表 2：GIDS 訓練的有效頻寬和 IOPS 與基準對比

隨著資料集的不斷擴大，我們需要轉變模式，以便在合理的時間內訓練這些模型，並充分利用領先 GPU 所提供的增強能力。BaM 和 GIDS 是一個很好的起點，我們期待著未來更多地利用此類系統。

測試系統

參考連結

大型加速器記憶體文章與 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 的原型專案，不打算全面發佈。