[論文レビュー] GACT: Activation Compressed Training for Generic Network Architectures

Xiaoxuan Liu, Lianmin Zheng+10 more|arXiv (Cornell University)|Jun 22, 2022

Advanced Neural Network Applications被引用数 5

ひとこと要約

GACTは、収束保証付きで動的に入力活性化を圧縮することで、畳み込みネットワーク、トランスフォーマー、グラフネットワークを含む多様なニューラルネットワークアーキテクチャのメモリ効率の良い学習を可能にする一般化されたアクティベーション圧縮学習（ACT）フレームワークです。最大8.1倍の活性化メモリ削減を実現し、精度損失を最小限に抑えながら最大24.7倍の大きなバッチサイズでの学習が可能になります。また、既存のモデルやメモリ節約技術とシームレスに統合可能なPyTorchライブラリとして実装されています。

ABSTRACT

Training large neural network (NN) models requires extensive memory resources, and Activation Compressed Training (ACT) is a promising approach to reduce training memory footprint. This paper presents GACT, an ACT framework to support a broad range of machine learning tasks for generic NN architectures with limited domain knowledge. By analyzing a linearized version of ACT's approximate gradient, we prove the convergence of GACT without prior knowledge on operator type or model architecture. To make training stable, we propose an algorithm that decides the compression ratio for each tensor by estimating its impact on the gradient at run time. We implement GACT as a PyTorch library that readily applies to any NN architecture. GACT reduces the activation memory for convolutional NNs, transformers, and graph NNs by up to 8.1x, enabling training with a 4.2x to 24.7x larger batch size, with negligible accuracy loss. We implement GACT as a PyTorch library at https://github.com/LiuXiaoxuanPKU/GACT-ICML.

研究の動機と目的

• アーキテクチャや演算子固有の仮定を一切行わず、多様なニューラルネットワークアーキテクチャの学習に一般化可能なアクティベーション圧縮フレームワークを開発すること。
• 演算子の種別やモデル構造の事前知識なしに、アクティベーション圧縮学習の収束を証明すること。
• 勾配への影響推定に基づき、実行時における各テンソルごとに最適な圧縮率を選択する適応的アルゴリズムを設計すること。
• GACTをPyTorchライブラリとして実装し、既存の学習パイプラインや勾配チェックポイントやスワッピングなどの他のメモリ節約技術とシームレスに統合できること。

提案手法

• アーキテクチャに依存しない一般化を図るため、ACTを演算子の計算グラフとして定式化し、レイヤー固有の仮定を回避すること。
• 線形化近似を用いてACTの確率的勾配を分析し、バイアスなしで構造的分散を持つことを証明し、数値的推定が可能であることを示すこと。
• 異なる圧縮戦略における勾配分散を予測する数値的アルゴリズムを開発し、最適な戦略選択を可能にすること。
• 各テンソルの勾配更新への感受性に基づき、最適な圧縮率を自動的に決定するための整数計画法の近似を用いること。
• 複数の最適化レベルを備えたPyTorchライブラリとしてGACTを実装し、ユーザー定義演算子のサポートと、既存のメモリ効率技術との統合を可能にすること。
• 学習中に各テンソルが勾配分散に与える影響を推定することで、実行時における圧縮率の適応的変更を可能にすること。

実験結果

リサーチクエスチョン

• RQ1 アーキテクチャに関する仮定なしに、一般化されたニューラルネットワークアーキテクチャにわたるアクティベーション圧縮学習を、収束保証付きで実現できるか？
• RQ2 各テンソルごとに圧縮率を適応的に選択することで、学習の安定性と精度を維持できるか？
• RQ3 カスタム実装なしに、任意の演算子（ユーザー定義を含む）をサポートする汎用ACTフレームワークを構築できるか？
• RQ4 一般化されたディープラーニングモデルにおいて、メモリ削減、学習速度、精度損失のトレードオフはどのように変化するか？
• RQ5 GACTは、勾配チェックポイントやスワッピングなどの既存のメモリ効率技術とどのように統合され、その効果が向上するか？

主な発見

• GACTは、畳み込みネットワーク、トランスフォーマー、グラフニューラルネットワークの全分野で最大8.1倍の活性化メモリ削減を実現し、精度損失は最小限に抑えられます。
• GACTにより、同じGPUメモリ予算下で最大24.7倍の大きなバッチサイズでの学習が可能となり、学習スループットが顕著に向上します。
• ResNet-152では、16GBのGPUメモリ制約下で、7.0倍の深さ、3.6倍の幅、または3.0倍の解像度にスケーリング可能になります。
• BERT-largeでは、同じメモリ制約下で2.0倍の深さ、1.6倍の幅のモデルが利用可能になります。
• GACTと勾配チェックポイントを組み合わせることで、バッチサイズ288の条件下でピーク活性化メモリが5.4倍削減され、精度に影響なしに実現されました。
• GACTとスワッピングを組み合わせることで、圧縮されたテンソルのCPU-GPU間通信が高速化され、最大2.3倍の学習速度向上が達成され、プリフェッチによる追加の約7％の高速化が、最小限のメモリオーバーヘッドで実現されました。