[論文レビュー] Pufferfish: Communication-efficient Models At No Extra Cost

Wang, Hongyi, Saurabh Agarwal+1 more|arXiv (Cornell University)|Mar 5, 2021

Advanced Neural Network Applications参考文献 89被引用数 5

ひとこと要約

Pufferfish は、フルランクモデルから導出された低ランクで事前に因子分解された深層ネットワークを用いることで、通信および計算効率に優れた分散学習フレームワークである。最初に学習の一部（例：10%）をフルランクネットワークで初期化し、その後SVDを適用して低ランクに変換することで、Pufferfishは圧縮のオーバーヘッドを排除し、精度を損なわずに PyTorch の DDP より最大 1.64× のエンドツーエンドの高速化を達成しており、ロットリート・チケット仮説および構造的プルーニング手法を上回っている。

ABSTRACT

To mitigate communication overheads in distributed model training, several studies propose the use of compressed stochastic gradients, usually achieved by sparsification or quantization. Such techniques achieve high compression ratios, but in many cases incur either significant computational overheads or some accuracy loss. In this work, we present Pufferfish, a communication and computation efficient distributed training framework that incorporates the gradient compression into the model training process via training low-rank, pre-factorized deep networks. Pufferfish not only reduces communication, but also completely bypasses any computation overheads related to compression, and achieves the same accuracy as state-of-the-art, off-the-shelf deep models. Pufferfish can be directly integrated into current deep learning frameworks with minimum implementation modification. Our extensive experiments over real distributed setups, across a variety of large-scale machine learning tasks, indicate that Pufferfish achieves up to 1.64x end-to-end speedup over the latest distributed training API in PyTorch without accuracy loss. Compared to the Lottery Ticket Hypothesis models, Pufferfish leads to equally accurate, small-parameter models while avoiding the burden of "winning the lottery". Pufferfish also leads to more accurate and smaller models than SOTA structured model pruning methods.

研究の動機と目的

• 頻繁な勾配転送に起因する分散ディープラーニング学習における通信ボトル neck 問題を解消すること。
• スパarsification や量子化といった勾配圧縮技術に通常伴う計算オーバーヘッドを排除すること。
• モデルの精度を損なわず、複雑な実装変更を要せず通信効率を達成すること。
• PyTorch などの既存のディープラーニングフレームワークと互換性のある即挿し可能なソリューションを提供すること。
• ロットリート・チケット仮説や構造的プルーニングといった既存手法を、精度およびモデルサイズの面で上回ること。

提案手法

• 全ランクの深層ニューラルネットワークを、全学習時間のわずかな割合（例：10%）にわたり訓練して重みを安定化させる。
• 各層に対して特異値分解（SVD）を適用し、低ランクの成分に因子分解することで、事前因子分解された低ランクモデルを作成する。
• 得られた低ランクモデルを、残りの学習スケジュールにわたり微調整する。
• 明示的な勾配圧縮ステップを回避するため、低ランクモデルを直接学習パイプラインに統合する。
• 初期学習段階でフルランクと低ランクのコンポーネントを組み合わせたハイブリッドアーキテクチャを採用し、精度の低下を軽減する。
• PyTorch DDP などの既存の分散学習APIを最小限の変更で活用し、エンドツーエンドの効率性を実現する。

実験結果

リサーチクエスチョン

• RQ1 勾配圧縮をモデルアーキテクチャ自体に統合することで、計算オーバーヘッドを排除しつつ精度を維持できるか？
• RQ2 部分的なフルランク学習後にSVDを用いてモデルを事前に因子分解することで、精度の劣化を伴わず通信効率を達成できるか？
• RQ3 このアプローチは、カスタム通信プリミティブを必要とせず、最先端の分散学習APIを上回る高速化を達成できるか？
• RQ4 ロットリート・チケット仮説と比較して、本手法は精度、モデルサイズ、学習効率の面でどのように差をつけるか？
• RQ5 低ランク因子分解は、ハイパーパramータの大幅なチューニングを要せず、さまざまなアーキテクチャやタスクに適用可能か？

主な発見

• Pufferfish は、実際の分散環境において、精度を損なわず PyTorch の DistributedDataParallel (DDP) より最大 1.64× のエンドツーエンドの高速化を達成した。
• CIFAR-10 における ResNet-18 では、速度最適化設定で 1.16× のエポック単位の高速化を達成し、MAC 演算数が 22% 減少した。
• 初期段階でのフルランクウォームアップを組み込んだハイブリッドトレーニング戦略により、VGG-19-BN (CIFAR-10) では 0.36%、ResNet-50 (ImageNet) では 4.8% の精度向上が得られた。
• Pufferfish モデルはロットリート・チケット仮説のモデルと同等またはそれ以上の精度を達成しており、高コストなロットリート・ウィンニングプロセスを回避している。
• Pufferfish は、最先端の構造的プルーニング手法よりも小さく、より正確なモデルを生成し、パラメータ効率の優位性を示している。