[論文レビュー] Training Large Neural Networks with Constant Memory using a New Execution Algorithm

Bharadwaj Pudipeddi, Maral Mesmakhosroshahi+2 more|arXiv (Cornell University)|Feb 13, 2020

Ferroelectric and Negative Capacitance Devices参考文献 15被引用数 24

ひとこと要約

本稿では、L2L（レイヤー・ツー・レイヤー）と呼ばれる新しい実行アルゴリズムを提案する。この手法は、フルモデルをCPUベースのイージー・パラメータ・サーバー（EPS）にオフロードすることで、GPUメモリに常に現在のレイヤーのパラメータと活性化値のみを保持するため、定常的なメモリ使用量で大規模ニューラルネットワークの学習を可能にする。この方法により、最先端のベースラインと比較して45%低いメモリ使用量と40%高いスループットを達成し、1つの16GB V100 GPU（512GB CPUメモリ）でモデル分割やミニバッチサイズの制限なしに500億パラメータのモデルを学習可能にする。

ABSTRACT

Widely popular transformer-based NLP models such as BERT and Turing-NLG have enormous capacity trending to billions of parameters. Current execution methods demand brute-force resources such as HBM devices and high speed interconnectivity for data parallelism. In this paper, we introduce a new relay-style execution technique called L2L (layer-to-layer) where at any given moment, the device memory is primarily populated only with the executing layer(s)'s footprint. The model resides in the DRAM memory attached to either a CPU or an FPGA as an entity we call eager param-server (EPS). To overcome the bandwidth issues of shuttling parameters to and from EPS, the model is executed a layer at a time across many micro-batches instead of the conventional method of minibatches over whole model. L2L is implemented using 16GB V100 devices for BERT-Large running it with a device batch size of up to 256. Our results show 45% reduction in memory and 40% increase in the throughput compared to the state-of-the-art baseline. L2L is also able to fit models up to 50 Billion parameters on a machine with a single 16GB V100 and 512GB CPU memory and without requiring any model partitioning. L2L scales to arbitrary depth allowing researchers to develop on affordable devices which is a big step toward democratizing AI. By running the optimizer in the host EPS, we show a new form of mixed precision for faster throughput and convergence. In addition, the EPS enables dynamic neural architecture approaches by varying layers across iterations. Finally, we also propose and demonstrate a constant memory variation of L2L and we propose future enhancements. This work has been performed on GPUs first, but also targeted towards all high TFLOPS/Watt accelerators.

研究の動機と目的

• BERT や GPT-3 のような大規模トランスフォーマー・モデルの増大するメモリおよび計算要件に対処すること。これらのモデルは、標準的なGPUの容量を超えている。
• 高帯域幅メモリ（HBM）デバイスやモデル分割を必要とせず、安価なハードウェアで数十億パラメータのモデルを学習可能にすること。
• 任意のモデルの深さにスケーリング可能な定常的メモリ実行手法を開発すること。
• モデル重みと最適化器状態をCPUベースのイージー・パラメータ・サーバー（EPS）に移動させ、レイヤーを逐次実行することで、メモリ圧力を軽減し、スループットを向上させること。
• GPUとCPU間の新しい低オーバーヘッドのパラメータ転送メカニズムにより、混合精度学習と効率的なデータ並列処理を実現すること。

提案手法

• L2Lは、リレーオンスタイルの実行を採用し、GPUメモリには常に現在のレイヤーのパラメータと活性化値のみを保持し、フルモデルはCPUまたはFPGAのDRAM上に配置されたイージー・パラメータ・サーバー（EPS）として保持する。
• EPSは、実行直前に次のレイヤーのパラメータをGPUに事前にロードし、内側のループによる最適化により、アイドル時間を最小限に抑え、転送頻度を低減する。
• 本手法は、ミニバッチではなくマイクロバッチをレイヤーごとに逐次処理することで、メモリフットプリントを削減し、モデルの深さに関係なく定常的メモリ使用量を維持する。
• EPSはGPU計算と並列して勾配の縮約と重み更新を処理し、より高速な収束を実現する新しい形の混合精度学習を可能にする。
• 将来の拡張であるL2Lpでは、EPS内で完全に並列化された縮約と重み更新が実装され、NVLinkは次のレイヤーのロードにのみ使用されるため、帯域幅依存性が最小限に抑えられる。
• レイヤーが独立して実行可能であり、イテレーションごとに構造を変更できるため、動的ニューラルアーキテクチャサーチをサポートする。

実験結果

リサーチクエスチョン

• RQ1 標準的なGPU上で、モデルをCPUベースのパラメータサーバーにオフロードすることで、大規模トランスフォーマー・モデルを定常的メモリで学習可能か？
• RQ2 マイクロバッチを用いたレイヤー単位の実行は、従来のミニバッチ学習と比較して、メモリ使用量とスループットにおいてどのように異なるか？
• RQ3 L2L手法は、1つの16GB V100でも、極めて深いモデル（例：384レイヤー）をメモリ不足を起こさずにスケーリング可能か？
• RQ4 EPSベースのアーキテクチャは、最適化されたパラメータ転送と混合精度学習により、どれほど収束を高速化し、スループットを向上できるか？
• RQ5 L2Lは、再コンパイルや再構成を必要とせず、イテレーションごとにレイヤー構造を変更可能であるため、動的ニューラルアーキテクチャサーチをサポートできるか？

主な発見

• L2Lは、1つの16GB V100 GPUでBERT-Largeを学習する際、最先端のベースラインと比較してGPUメモリ使用量を45%削減した。
• 本手法は、メモリ圧力を軽減しながらも、ベースラインと比較して40%高いトレーニングスループットを達成した。
• L2Lは、1つの16GB V100で最大256のデバイスバッチサイズでBERT-Largeを学習可能であり、これはベースラインがバッチサイズ2でさえも苦戦するのと比べて顕著な改善である。
• 本アプローチは、モデル分割やメモリオーバーフローのエラーなしに、1つの16GB V100と512GB CPUメモリで最大500億パラメータのモデルをサポートする。
• L2Lは、モデルの深さに関係なく定常的メモリ使用量を維持し、384レイヤーのモデルでさえも、メモリオーバーフローを起こさずに学習可能である。
• 検証曲線から、FP32モードおよび混合精度モードの両方で、L2Lがベースラインよりも収束が速いことが示され、トレーニング効率の向上が裏付けられた。