Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Training

要約

GPU メモリの不足は、大規模な DNN トレーニングの主なボトルネックになっています。
トレーニング中にメモリ使用量を削減する方法は、集中的な研究の注目を集めています。
以前の勾配の蓄積は活性化メモリを削減しますが、勾配の保存と勾配の解放の間に矛盾があるため、勾配メモリの削減と互換性がないことがわかりました。
この問題に対処するために、アクティベーション メモリと勾配メモリの両方を削減できる、Adam Accumulation (AdamA) という名前の Adam 用の新しいオプティマイザ蓄積方法を提案します。
具体的には、AdamA は勾配をオプティマイザー状態に直接統合し、マイクロバッチにわたってオプティマイザー状態を蓄積するため、勾配は使用後すぐに解放できます。
私たちは、AdamA が Adam と同じ収束特性をもたらすことを数学的および実験的に実証します。
トランスフォーマーベースのモデルで評価すると、AdamA は、トレーニング スループットの低下が 2% 未満で、勾配累積と比較して最大 23% のメモリ削減を達成します。
特に、AdamA は、オプティマイザー状態のメモリ削減メソッドと連携して、メモリ容量が異なる GPU 上の PyTorch および DeepSpeed ベースラインの 1.26 倍から 3.14 倍の大きなモデルに適合させることができます。

要約(オリジナル)

Running out of GPU memory has become a main bottleneck for large-scale DNN training. How to reduce the memory footprint during training has received intensive research attention. We find that previous gradient accumulation reduces activation memory but fails to be compatible with gradient memory reduction due to a contradiction between preserving gradients and releasing gradients. To address this issue, we propose a novel optimizer accumulation method for Adam, named Adam Accumulation (AdamA), which enables reducing both activation and gradient memory. Specifically, AdamA directly integrates gradients into optimizer states and accumulates optimizer states over micro-batches, so that gradients can be released immediately after use. We mathematically and experimentally demonstrate AdamA yields the same convergence properties as Adam. Evaluated on transformer-based models, AdamA achieves up to 23% memory reduction compared to gradient accumulation with less than 2% degradation in training throughput. Notably, AdamA can work together with memory reduction methods for optimizer states to fit 1.26x~3.14x larger models over PyTorch and DeepSpeed baseline on GPUs with different memory capacities.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google