Facebook AI Research(FAIR)の科学者チームは、独自に開発したAIハードウェアプラットフォームのZionEX上でPyTorchを使ってディープラーニングレコメンデーションモデル(DLRM)をトレーニングするシステムを発表した。このシステムを使って同チームは、最大12Tパラメータを持つモデルのトレーニングにおいて、他のシステムよりも1桁近いトレーニング時間のスピードアップを実現している。
研究の内容は、arXivに公開された論文に説明されている。Zionプラットフォームの最新バージョンであるZionEXハードウェアは、独自開発した計算ノード数千台で構成されるクラスタである。各ノードは4ソケットのCPUと8個のGPUの他、CPU用に4つのネットワークインターフェースコントローラ(NIC)、GPU用にRDMA over Convereged Ethernet(ROCE) NICを備える。トレーニングソフトウェアは、FacebookのディープラーニングフレームワークであるPyTorchをベースに、フレームワークの分散トレーニング機能に関して、GPUパイプラインやネットワーク通信のチューニングなど、カスタムコードによるオペレータと拡張を加えたものを使用する。システムのパフォーマンスを評価するため、同チームは、95Bから12Tまで、さまざまなサイズのDLRMをトレーニングした。前世代のトレーニングシステムと比較した場合、ZionEXはスループットで1桁近いスピードアップを達成している。
ディープラーニングレコメンデーションモデルは、Facebookのニュースフィードや検索機能において、結果のランク付けやクリック率の予測などに幅広く使用されている。Facebookの表現を借りれば、DLRMは、"データセンタのインフラストラクチャ需要の点から見て、単一としては最大のAIアプリケーション"である。DLRMのアーキテクチャは、多層パーセプトロン(MLP)と、高次元入力のための多数の埋め込みテーブルで構成される。このためモデルは、OpenAIのGPT-3やGoogleのSwitch Transformerにも匹敵するような、極めて大規模なものになる可能性があるのだ。2019年、Facebookは、自社のDLRM実装をオープンソースとして公開すると同時に、独自開発したAIトレーニングハードウェアのZionを発表した。
しかしながら、同社のZionとそのDLRMトレーニングアルゴリズムは、大規模モデルの需要に伴うスケーリング上の課題をいくつか抱えていた。MLPモジュールをトレーニングするために、Facebookは、単一パラメータサーバと分散型トレーナというデータ並列構成を使用していたが、その一方で組み込みテーブルに関しては、複数パラメータサーバによるモデル並列スキームを用いていた。このZionノードのネットワーク設計によって、分散サーバをコーディネートするための通信要件がボトルネックとなっていたのだ。
特に問題なのは、ノードのNICがCPUにアタッチされていたため、異なるノードのGPU間の通信にCPUサイクルが必要であったことだ。データセンタのネットワークを使用するためにTCP/IPネットワークプロトコルの使用が必要であった点も、分散トレーニングには最適とは言えなかった。再設計されたZionEXのノードは、GPU毎にNICを追加し、効率のよいRDMA/GPUDirectプロトコルを使用した専用のバックエンドネットワークに接続することで、ネットワークの改善をサポートしている。このバックエンドネットワークは、数千ノードの規模までスケールアップが可能だ。
新ハードウェアのメリットをすべて引き出すため、トレーニングソフトウェアも再設計された。従来のシステムでは、組み込みパラメータの非同期更新を使用していたため、大規模なモデルにおいて精度の低下があった。再設計されたトレーニングシステムでは、引き続きハイブリッドなアプローチ -- MLPはデータ並列、組み込みにはモデル並列 -- を引き続き採用しているが、パラメータサーバは廃止し、組み込みパラメータに同期更新を使用するようにした。組み込みトレーニングのパフォーマンスをさらに改善するため、いくつかのPyTorchオペレータについて、トレーニング中のメモリとネットワークの使用量を削減した効率のよい実装を開発した。これらはオープンソースとして公開されている。
システムのパフォーマンスを評価するために、同チームは、それぞれ95B、793B、845B、12Tのパラメータを備えた4種類のDLRMをトレーニングした。モデルのトレーニングは最大16ノード、あるいは128GPUのクラスタ上で行われた。ノード数の増加によってクラスタのトレーニング時間はスピードアップしたが、845Bパラメータモデルで8倍に達した以降は、GPU間のネットワーク通信が、特に組み込みパラメータの更新に必要な"all-to-all"通信中にボトルネックになる傾向があった、と研究者らは記している。このオペレーション中のデータ転送量を削減するため、チームは、このデータを16ビットに量子化する選択をした。これにより、モデルの精度を大幅に低減することなく、ネットワークトラフィックを低減することに成功したのだ。
大規模なディープラーニングモデルの急増により、モデルトレーニングのスケーラビリティを改善する試みがいくつも進められている。巨大なモデルには複数GPUの使用が必須であるため、分散オペレーションを処理するようなトレーニングソフトウェアの変更が必要になる。Googleは、自社のTensorFlowフレームワーク用に、分散レイヤのMesh TensorFlowと、XLAコンパイラを拡張して並列処理をサポートするGShardを開発した。Microsoftは先頃、同社のDeepSpeedライブラリをアップデートして、大規模モデルのメモリ使用量を削減した。これにより、単一GPUで数千億、並列テクニックを使用すれば兆単位のパラメータを持ったモデルのトレーニングが可能になっている。
Facebookのオペレータライブラリのソースコードとベンチマークは、同社のDLRM実装と合わせてGitHubで入手が可能である。