[論文レビュー] ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers

Zhewei Yao, Reza Yazdani Aminabadi+4 more|arXiv (Cornell University)|Jun 4, 2022

Advanced Neural Network Applications被引用数 72

ひとこと要約

ZeroQuant は、細粒度の重みと活性化量子化、軽量なレイヤー単位の知識蒸留（LKD）、および最適化されたバックエンドを備えたエンドツーエンドの PTQ パイプラインを提供し、INT8/INT4-INT8 の混合精度を大規模トランスフォーマーで最小の精度低下と大幅なスピードアップを実現します。

ABSTRACT

How to efficiently serve ever-larger trained natural language models in practice has become exceptionally challenging even for powerful cloud servers due to their prohibitive memory/computation requirements. In this work, we present an efficient and affordable post-training quantization approach to compress large Transformer-based models, termed as ZeroQuant. ZeroQuant is an end-to-end quantization and inference pipeline with three main components: (1) a fine-grained hardware-friendly quantization scheme for both weight and activations; (2) a novel affordable layer-by-layer knowledge distillation algorithm (LKD) even without the access to the original training data; (3) a highly-optimized quantization system backend support to remove the quantization/dequantization overhead. As such, we are able to show that: (1) ZeroQuant can reduce the precision for weights and activations to INT8 in a cost-free way for both BERT and GPT3-style models with minimal accuracy impact, which leads to up to 5.19x/4.16x speedup on those models compared to FP16 inference; (2) ZeroQuant plus LKD affordably quantize the weights in the fully-connected module to INT4 along with INT8 weights in the attention module and INT8 activations, resulting in 3x memory footprint reduction compared to the FP16 model; (3) ZeroQuant can be directly applied to two of the largest open-sourced language models, including GPT-J6B and GPT-NeoX20, for which our INT8 model achieves similar accuracy as the FP16 model but achieves up to 5.2x better efficiency.

研究の動機と目的

• メモリと計算資源の制約下で、より大規模なNLPモデルのデプロイの必要性を動機づける。
• 再訓練を回避し、精度を保つポストトレーニング量子化パイプラインを提案する。
• 細粒度の重み/活性化量子化を備えたハードウェア認識量子化スキームを導入する。
• 訓練データなしで機能する軽量なレイヤー単位の知識蒸留法を導入する。
• 推論時の量子化オーバーヘッドを最小化するシステムレベルの最適化を示す。

提案手法

• 細粒度のハードウェアに優しい量子化を適用する：グループごとの重み量子化とトークンごとの活性化量子化。
• Layer-by-layer Knowledge Distillation (LKD) を導入し、元の非量子化レイヤを教師として、1つのレイヤーずつ量子化する。
• データ移動を減らすために前の演算とトークンごとの量子化を統合する高度に最適化された推論バックエンドを開発する。
• CUTLASS ベースの INT8 GeMM カーネルとカーネル融合を活用して量子化/デ量子化のオーバーヘッドを最小化する。
• 重み/活性化の INT8 量子化を精度損失を最小限に抑えつつ BERT および GPT-3 スタイルモデルで実証し、LKD を用いた INT4/INT8 混合精度を示す。
• GPT-J-6B および GPT-NeoX-20B へのスケーラビリティと、 substantial efficiency gains を示す。

実験結果

リサーチクエスチョン

• RQ1 訓練データへアクセスできない状態でのポストトレーニング量子化は、INT8 または INT4 に大規模トランスフォーマーモデルを量子化し、精度低下を最小限に抑えることができるか。
• RQ2 レイヤーごと蒸留アプローチ（LKD）は、元データへの完全な再訓練やアクセスなしで超低精度量子化を可能にするか。
• RQ3 どのハードウェア認識量子化戦略（グループ単位の重み、トークン単位の活性化）が大規模トランスフォーマーにおける最適な精度/レイテンシのトレードオフをもたらすか。
• RQ4 システムレベルの最適化（カーネル融合、バックエンド）が量子化されたトランスフォーマーで実際のレイテンシ利益をもたらすのにどれほど効果的か。
• RQ5 ZeroQuant は十億パラメータ規模のモデルへ適用可能で、競争力のある精度を維持しつつスループットを大幅に向上させられるか。

主な発見

• INT8 量子化による重みと活性化は、BERT および GPT-3 スタイルモデルで FP16 に比べて大幅なスピードアップをもたらし、精度低下は最小限である（BERT-base で最大 5.19x、GPT-3-350M で 4.16x）。
• LKD は INT4/INT8 の混合精度量子化を実現し、FP16 に比べて約 3x のメモリ使用量の削減と最小限の精度損失、クイックな量子化を達成（例：BERT-base の量子化約 33 秒）。
• ZeroQuant-LKD は GPT-J-6B および GPT-NeoX-20B で強力な結果を達成し、最大で 5.2x の効率向上と GPU 要件/レイテンシの削減を実現（例：GPT-NeoX-20B では 2 GPU から 1 に削減、レイテンシ 65ms から 25ms へ）。
• カーネル融合と CUTLASS ベースの INT8 GeMM バックエンドは、量子化/非量子化のオーバーヘッドを大幅に削減し、INT8 トランスフォーマー推論のレイテンシを改善する。
• GPT-3 スタイルモデルでは、量子化下で精度タスクより生成タスクの精度が堅牢で、ZeroQuant は PTQ（W8A8）との差を詰め、W4/8 系よりも LKD で上回る。
• アブレーション研究は、グループ単位の重み量子化とトークン単位の活性化量子化が意味のある精度向上をもたらし、 LKD によってさらに強化されることを示す。