PyTorchでOOMを克服：通常のGPUで巨大グラフをトレーニング

本稿では、PyTorch向けの新しいアウトオブコア最適化アルゴリズムであるDisk Sparse Adam（DSA）を紹介します。これは、数百万ノードを持つナレッジグラフ埋め込みなど、大規模なスパースモデルをコンシューマー向けGPUで学習するために設計されています。標準のtorch.optim.SparseAdamは、一次モーメントと二次モーメントの状態をメモリに保持しますが、1000万エンティティ、128次元ベクトルのモデルでは約10.24GBの追加メモリを消費し、しばしばCUDAのアウトオブメモリエラーを引き起こします。DSAは、これらのモーメント行列をバイナリファイルとしてディスクに保存し、OSのmmapメカニズムを使用してメモリにマッピングし、各最適化ステップでアクティブなバッチの状態のみを読み書きします。この最適化アルゴリズムは、標準のPyTorchパイプラインにそのまま置き換えて使用できるように設計されており、ユークリッド空間と双曲空間（ポアンカレ球）の両方をサポートしています。著者によると、最適化アルゴリズムの状態のメモリ使用量はギガバイトからメガバイトに削減され、100万エンティティを使用したKaggleでのベンチマークでは、VRAMオーバーヘッドが0.00MB、スループットが毎秒134,212サンプルでした。DSAはMITライセンスでGitHub上で公開されており、制限として、最高のパフォーマンスを得るにはNVMe SSDが必要であること、また、torch.nn.EmbeddingやEmbeddingBagからのようなスパース勾配のみに対応していることが挙げられます。