[論文レビュー] AutoFreeze: Automatically Freezing Model Blocks to Accelerate Fine-tuning

Yuhan Liu, Saurabh Agarwal+1 more|arXiv (Cornell University)|Feb 2, 2021

Topic Modeling参考文献 72被引用数 16

ひとこと要約

AutoFreeze は、勾配ノルムに基づく基準を用いて、収束が早い層を動的に凍結することで、BERT の微調整を高速化する。活性化キャッシュと分散実行を組み合わせることで、1 GPU で最大 2.55× の高速化、64 GPU クラスタでは 5.03× のコスト削減を達成し、モデルの精度を損なわず、最先端の精度を維持する。

ABSTRACT

With the rapid adoption of machine learning (ML), a number of domains now use the approach of fine tuning models which were pre-trained on a large corpus of data. However, our experiments show that even fine-tuning on models like BERT can take many hours even when using modern accelerators like GPUs. While prior work proposes limiting the number of layers that are fine-tuned, e.g., freezing all layers but the last layer, we find that such static approaches lead to reduced accuracy. We propose, AutoFreeze, a system that uses an adaptive approach to choose which layers are trained and show how this can accelerate model fine-tuning while preserving accuracy. We also develop mechanisms to enable efficient caching of intermediate activations which can reduce the forward computation time when performing fine-tuning. We extend AutoFreeze to perform distributed fine-tuning and design two execution modes that minimize cost and running time respectively. Our evaluation on ten NLP tasks shows that AutoFreeze, with caching enabled, can improve fine-tuning on a single GPU by up to 2.55x. On a 64 GPU cluster, for fine-tuning on the AG's news dataset, AutoFreeze is able to achieve up to 4.38x speedup when optimizing for end-to-end training time and 5.03x reduction in total cost when optimizing for efficiency, without affecting model accuracy.

研究の動機と目的

• BERT などの大規模事前学習モデルの微調整にかかる高い計算コスト（4 GPU クラスタで最大 27 時間）を軽減すること。
• 固定された層を凍結する静的微調整戦略やロットリート・プルーニングの限界を克服し、モデル精度を低下させることがある。
• 微調整中に動的に、実行時において決定される方法でモデルブロックを凍結する、適応的で精度を保つ高速化手法を開発すること。
• 効率的な活性化キャッシュと分散実行モードを統合し、マルチ GPU 環境におけるトレーニング時間とコストをさらに削減すること。

提案手法

• AutoFreeze は、トレーニング中に層の変化の速さを勾配ノルムに基づくテストで順位付けし、早期に収束する層を特定する。
• 勾配ノルムの減少が最も遅い層が凍結対象とされ、逆伝播計算が回避され、計算と通信のオーバーヘッドが削減される。
• 動的活性化キャッシュを採用することで、順伝播時に中間出力を再計算する必要がなくなり、順伝播の計算時間が短縮される。
• AutoFreeze は、エンドツーエンドのトレーニング時間を最適化するモードと、総コストを最小化するモードの2つの分散実行モードをサポートしており、クラスタ全体での効率的なスケーリングを可能にする。
• この手法は PyTorch に統合されており、最小限のランタイムオーバーヘッドで、1 GPU およびマルチ GPU の微調整をサポートする。
• 凍結の意思決定はトレーニング中に定期的に行われ、速度向上と精度のバランスを取るために評価間隔が最適化されている。

実験結果

リサーチクエスチョン

• RQ1 微調整中に適応的層の凍結を実施することで、モデル精度を損なわずトレーニング時間を顕著に短縮できるか？
• RQ2 勾配ノルム基準は、静的凍結やロットリートプルーニングと比較して、速度向上と精度保持の点でどのように優れているか？
• RQ3 活性化キャッシュと分散実行を組み合わせることで、マルチ GPU 環境における微調整の高速化はどの程度達成できるか？
• RQ4 層の収束を評価する最適な頻度は何か？速度向上とモデル性能のバランスを取るには？

主な発見

• 1 GPU では、キャッシュを活用した AutoFreeze が、10 の NLP タスクにおいて標準的な微調整と比較して最大 2.55× の高速化を達成した。
• 64 GPU クラスタでは、AutoFreeze がエンドツーエンドのトレーニング時間を最大 4.38× 短縮し、総コストを最大 5.03× 削減したが、精度の低下は認められなかった。
• 勾配ノルムテストによるオーバーヘッドは非常に小さく、区間時間の 1% 未満、1 エポックあたり 453MB のメモリ使用量にとどまり、実世界での利用に適している。
• 変化が遅い層のみを凍結することで、モデル精度が保持される。一方、静的アプローチでは初期層を凍結するため、MRPC タスクで精度が 87% から 76.5% に低下する。
• この手法は、顕著なトレーニング時間短縮を実現しながらも、最先端の精度を維持しており、精度の低下を伴う先行手法（例：EarlyBERT）を上回っている。
• 最適な凍結頻度は、1 エポックあたり 2 ～ 5 回の評価間隔に位置し、速度向上とモデル安定性のバランスが取れている。