NVIDIAのテクノロジーを使用してBaMを追求するMicron 9400 NVMe SSD

データセットの学習規模は、数十億のパラメーターを超えて拡大し続けています。大型化したモデルは、システムメモリに収まらない場合があります。このような場合、データローダーはさまざまな方法でフラッシュストレージに置かれたモデルにアクセスする必要があります。その方法の1つがSSD上で保存されるメモリマップドファイルであり、これにより、データローダーはファイルがメモリにあるかのごとくアクセスできますが、CPUとソフトウェアスタックのオーバーヘッドが学習システムのパフォーマンスを大幅に低下させます。そこで必要になるのが、Big accelerator Memory（BaM）*とGPU-Initiated Direct Storage（GIDS）*データローダーです。

BaMとGIDSとは何ですか？

BaMは、SSDの低レイテンシー、極めて高いスループット、高密度、および耐久性を活用するシステムアーキテクチャーです。BaMの目標は、GPUスレッドがSSDのデータセットにきめ細かくアクセスできるよう効率的な抽象化を提供し、GPUのストレージ要求においてCPUを使うソリューションよりも圧倒的に優れたパフォーマンスを提供することです。BaMの加速化には、特にストレージデバイスに直接アクセスするというGPUの内在的並行性を可能にするよう設計されたカスタムストレージドライバーを使用します。BaMはGPUからSSDへの通信を作成する際にCPUに依存する必要がないため、NVIDIA Magnum IO™ GPUDirect^®ストレージ（GDS）とは異なります。

マイクロンにはこれまでもNVIDIA GDSと連携した実績があります。

• NVIDIA Magnum IO GPUDirectストレージプラットフォームによるMicron^® 9400 NVMe™ SSDの性能 
• Magnum IO GPUDirect Storageとマイクロンの連携がもたらす、AIと機械学習への産業破壊的イノベーション

GIDSデータローダーはBaMのサブシステム上に構築され、GPU-accelerated Graph Neural Network（GNN）学習用のメモリ容量要件に対処すると同時にストレージのレイテンシーをマスキングします。GIDSはSSD上のグラフの特徴データを保存することでこれを達成します。このデータは通常、大型グラフのグラフデータセット全体の最も大きな部分を占めるからです。特徴データよりも一般的にはるかに小さいグラフ構造データは、システムメモリに固定され、高速GPUグラフサンプリングを可能にします。最後に、GIDSデータローダーは、ストレージアクセスを減少させるため、最近アクセスされたノードに関してGPUメモリにソフトウェア定義のキャッシュを割り当てます。

GIDSを使ったグラフニューラルネットワーク学習

BaMとGIDSのメリットを示すため、私たちはIllinois Graph Benchmark（IGB）のヘテロジニアスフルデータセットを使用して、GNN学習を実行しました。このデータセットのサイズは2.28TBで、大半のプラットフォームのシステムメモリに収まりません。私たちは単一のNVIDIA A100 80GB Tensor Core GPUを使用して学習時間を100イテレーションに設定し、図1および表1に示すように、SSDの数を変化させて幅広い範囲の結果を求めました。

図1：IGB-ヘテロジニアスフルデータセットのGIDS学習時間 - 100イテレーション

表1：IGB-ヘテロジニアスフルデータセットのGIDS学習時間 - 100イテレーション

学習の最初の部分は、GPUによるグラフサンプリングとシステムメモリ内のグラフ構造データへのアクセスによるグラフサンプリングです（青で表示）。システムメモリに保存されている構造はこれらのテスト間で変化しないため、この値は異なるテスト設定間で大きな変化は見られませんでした。

もう1つは実際の学習時間です（右端に緑色で表示）。この部分はGPUに大きく依存し、予想どおり、複数のテスト設定間で大きな変化は見られませんでした。

最も重要な部分は特徴集約で、ここに大きな変化が見られました（金色で表示）。このシステムでは、特徴データはMicron 9400 SSDに保存されており、Micron 9400 SSDを1枚から4枚に増やすと特徴集約の処理時間が大きく向上（削減）したことがわかります。SSDを1枚から4枚に増やすと、特徴集約は3.68倍向上しました。

私たちはまた、ベースライン計算で特徴データにアクセスする際にメモリマップの抽象化とDeep Graph Library（DGL）データローダーを使いました。この特徴データへのアクセス方法は、GPUによる直接アクセスの代わりにCPUソフトウェアスタックの使用が必要となるため、学習中にGPUの飽和状態を維持するにはCPUソフトウェアスタックがいかに非効率的であるかわかります。ベースラインに対する特徴の抽象化の向上は、GIDSを使用した1枚のMicron 9400 NVMe SSDでは35.76倍、4枚のMicron 9400 NVMe SSDでは131.87倍でした。他の視点からこのデータを見たものとして、図2と表2にこれらのテストの実行中における有効帯域幅とIOPSを示しました。

図2：GIDS学習対ベースラインの有効帯域幅とIOPS

表2：GIDS学習対ベースラインの有効帯域幅とIOPS

データセットが拡大するにつれ、合理的な時間内にモデルを訓練し、主要GPUが提供する改善を活用するためには、パラダイムシフトが必要であることがわかります。BaMとGIDSは出発点として優れており、私たちは今後もこのタイプのシステムが増えて活用できるようになることを期待しています。

テストシステム

参考資料のリンク

BaMに関する論文とGitHub
GPU-Initiated On-Demand High-Throughput Storage Access in the BaM System Architecture (arxiv.org)
GitHub - ZaidQureshi/bam

GPU Initiated Direct Storageに関する論文とGitHub
Accelerating Sampling and Aggregation Operations in GNN Frameworks with GPU Initiated Direct Storage Accesses (arxiv.org)
GitHub - jeongminpark417/GIDS
GitHub - IllinoisGraphBenchmark/IGB-Datasets: Largest real-world open-source graph dataset - Worked done under IBM-Illinois Discovery Accelerator Institute and Amazon Research Awards and in collaboration with NVIDIA Research.

*注：NVIDIA Big Accelerator Memory（BaM）とNVIDIA GPU Initiated Direct Storage（GIDS）データローダーはNVIDIA Researchによるプロトタイププロジェクトで、一般向けリリースは予定されていません。