DUICUO

最大 9 倍高速化! ByteDance が 8 ビット混合精度 Transformer エンジンをオープンソース化。

近年、TransformerはNLPやCV分野で主流のモデルとなっていますが、モデルパラメータの多さが効率的な学習と推論の限界となっています。そこで、ByteDanceは2019年12月と2021年6月にそれぞれ効率的な推論・学習エンジンLightSeqをリリースし、Transformerシリーズモデルの学習と推論を大幅に加速するとともに、学習から推論までのTransformerの全プロセスをオープン化し、ユーザーエクスペリエンスを大幅に最適化しました。最近、LightSeq学習エンジンに関する論文[1]が、採択難易度が極めて高いスーパーコンピューティング分野のトップ国際会議であるSC22に採択され、学術界から広く認められました。

  • SC22 に採択された論文: https://sc22.supercomputing.org/presentation/?id=pap211&sess=sess154
  • コードアドレス: https://github.com/bytedance/lightseq

速度をさらに向上させるにはどうすればよいでしょうか?計算精度を下げることは比較的直接的な方法です。2017年以降、fp16混合精度技術[2]が広く使用されています。モデル性能を損なわないという前提の下、モデルの学習と推論の速度は50%以上向上しました。モデル性能を維持するために、低精度の手法(int8など)では、通常、以下の従来の手法を使用する必要があります。

  1. まず、fp16 混合精度を使用してモデルを収束までトレーニングします。
  2. 次に、量子化を考慮したトレーニングのために、モデル内の計算集約型の演算子の重み、入力、および出力に疑似量子化ノードが挿入されます。
  3. 最後に、疑似量子化ノードを含むモデル計算グラフは、サービス展開とモデル推論専用の int8 推論エンジンに変換されます。

上記のアプローチはほとんどのタスクでロスレス モデル パフォーマンスを実現できますが、次の問題が依然として存在します。

  1. 使用方法は複雑です。例えば、追加の量子化知覚学習プロセス[4]が必要であり、擬似量子化ノードを用いた計算グラフ変換も複雑です。
  2. トレーニング速度が遅い。現在普及しているディープラーニングフレームワークはint8精度をサポートしていないため、量子化を考慮したトレーニングでは、int8量子化をシミュレートするためにfp16疑似量子化ノードを挿入する必要があり、その結果、量子化を考慮したトレーニングはfp16混合精度トレーニングよりも2~3倍遅くなります。
  3. 推論の展開は困難で、高速化も低い。fp32やfp16などの型と比較すると、int8のハードウェアおよび基盤となるソフトウェアライブラリの最適化は比較的遅れている。例えば、NVIDIA GPUでは、int8行列乗算の高速化はハードウェアアーキテクチャと特定の形状によって制限され、実際の高速化は理論値よりもはるかに低くなります。

特に指定がない限り、以下のテキストでの量子化は int8 精度の量子化を指します

これらの問題に対処するため、ByteDanceはLightSeq GPU量子化トレーニングおよび推論エンジンの新バージョンをリリースしました。このエンジンはTransformerモデルの量子化トレーニングと推論をサポートし、すぐに使用できるため、ユーザーフレンドリーなエクスペリエンスを提供します。LightSeqは、高速かつ高精度なint8精度の量子化トレーニングと推論を実現します。

  1. 高速: A100 マルチカード トレーニングは最大 5.2 倍高速化、T4 シングルカード推論は最大 8.9 倍高速化。
  2. 精度: トレーニングと推論の効果は基本的に影響を受けません。
  3. 驚異的: 同じデータ量で、GPU メモリ使用量は最大 68% 削減され、モデルのストレージ容量は 75% 削減されます。

全体的に、LightSeq 量子化トレーニングおよび推論エンジンの新しいバージョンには、次の利点があります。

1. 豊富なサポート

Transformerモジュール全体と複数のデコードアルゴリズムをサポートし、Transformer、BERT、GPT、BART、ViTといった様々なモデル構造をサポートしています。また、量子化学習、モデルのエクスポート、量子化推論のためのFairseq、Hugging Face、NeurSTといった複数の学習フレームワークをサポートし、ユーザーが参照できる豊富なサンプルを提供しています。

2. 優れたパフォーマンス

fp16精度のLightSeq推論エンジンと比較して、int8量子化により最大70%の性能向上を実現し、PyTorch推論エンジンと比較して最大8.9倍の高速化を実現します。同時に、メモリ使用量はfp16推論エンジンと比較して約30%削減され、モデルの保存容量は元の4分の1にまで削減されます。最終的に、複数のタスクで検証した結果、推論性能は実質的にロスレスであることが確認されました。

3. 使いやすい

LightSeqは既に複数のトレーニングライブラリの量子化をサポートしており、ワンクリックで量子化トレーニングを行い、LightSeqがサポートするモデル形式への簡単なエクスポート、そして量子化推論が可能です。さらに、LightSeqはトレーニング後の量子化もサポートしており、追加のトレーニングなしで量子化推論が可能です。

使い方

上図に示すように、量子化による損失を最小限に抑えるには、まず浮動小数点モデルをfp16精度で学習し、最高のパフォーマンスを実現する必要があります。次に、量子化を有効にして微調整を行い、微調整された量子化モデルを作成します。この時点で、モデルのパフォーマンスは実質的に浮動小数点モデルのレベルに戻ります。次に、量子化モデルをLightSeqがサポートするPBまたはHDF5モデル形式に変換し、最後にLightSeqを用いて量子化推論を実行します。

インストール方法

LightSeq はインストールが非常に簡単で、必要なのは 1 つのコマンドだけです。

 pip でlightseqをインストール

定量トレーニング

LightSeqは、Fairseq、Hugging Face、NeurSTなどの学習フレームワークとの量子化統合をサポートしており、量子化を有効にしたカスタムモデルの作成も可能です。エンコーダー層を例に挙げると、浮動小数点モデルを定義して量子化を有効にするだけです。

 lightseq.trainingからLSTransformerEncoderLayerをインポート lightseq.training.ops.pytorch.quantizationからenable_quantインポートします
config = LSTransformerEncoderLayer . get_config ( model = "bert-base" , max_batch_tokens = 4096 , max_seq_len = 512 , fp16 = True , local_rank = 0 ,) layer = LSTransformerEncoderLayer ( config ) #量子化レイヤーを有効にする. apply ( enable_quant )

定量的推論

LightSeq は便利な Python 推論インターフェースを提供し、わずか 3 行のコードで高速な量子化推論を可能にします。

 lightseq.inferencelsiとしてインポートします
model = lsi.QuantTransformer ( pb_path , batch_size ) result = model.infer (入力)

さらに、LightSeq は BERT、GPT、ViT などのモデル用の Python インターフェイスを提供しており、それぞれ QuantBert、QuantGpt、QuanVit を呼び出すことで体験できます。

勾配通信量子化

LightSeqはTransformerモデル[5]の勾配通信量子化をサポートしています。分散量子化の学習はFairseqまたはHugging Faceを用いて簡単に開始でき、浮動小数点モデルと量子化モデルの両方をサポートしています。モデルを構築した後は、モデルに勾配通信量子化を開始するための通信フックを登録し、学習プロセスを開始するだけで済みます。

 lightseq.training.gradient_comm_quantizationからencode_and_decodeインポートしGCQStateをtorch.nn.parallelからDistributedDataParallelインポートします
#モデルはFairseqまたはHugging Faceから取得できDDPラップされます。model = DistributedDataParallel ( model ) state = GCQState ( process_group ) #フックを登録します。model.register_comm_hook ( state = statehook = encode_and_decode )

パフォーマンステスト

LightSeq は、複数のタスクで量子化トレーニング、量子化推論、勾配通信量子化の速度をテストし、メモリ使用量と量子化モデルのパフォーマンスを分析しました。

トレーニング速度を定量化する

LightSeq は 8 枚の A100 グラフィック カードでトレーニング実験を実施し、主な比較対象として Fairseq の Transformer、Hugging Face の BERT、GPT2、ViT を使用しました。

4つのモデル構造の加速傾向は似通っており、データ量が増加するにつれて高速化率が低下することがわかります。これには3つの理由があります。

  1. データ量が増加するにつれて、行列乗算GEMMの割合が大幅に増加します。そのため、PyTorch QATによって追加される擬似量子化ノードの時間は徐々に減少し、最終的にはPyTorch fp16の速度に限りなく近づくでしょう。
  2. 一方、GEMMの割合が増えるにつれて、LightSeq fp16カスタム演算子の高速化効果は徐々に減少するため、時間はPyTorch fp16に非常に近くなります。
  3. Ampere アーキテクチャ グラフィック カード上の int8 GEMM は、シェイプが小さい場合は fp16 GEMM よりもさらに遅く、シェイプが大きい場合はわずかに速くなるだけなので、データ量が増加すると、LightSeq int8 は LightSeq fp16 の速度に近づきます。

定量的推論速度

LightSeq は単一の T4 グラフィック カードで推論実験を実施し、主な比較対象として Hugging Face の Transformer、BERT、GPT2、ViT を使用しました。

ご覧のとおり、入力データ量の増加に伴い、LightSeqとPyTorchの差は徐々に縮まっています。これはGEMMの割合が増加したためです。LightSeq fp16とLightSeq int8を比較すると、データ量の増加に伴いLightSeq int8が高速化していることがわかります。これは、T4グラフィックカードにおいて、int8 GEMMの加速は形状の増加に伴い大幅に向上するためです。したがって、T4グラフィックカードで量子化推論を実行する場合、入力データ量が大きいほど加速効果は高くなります。

LightSeqは、複数の言語方向とテストセットにわたる機械翻訳において、異なるバッチサイズにおけるLightSeq int8推論とLightSeq fp16推論の速度向上を比較するテストも実施しました。実験は単一のT4グラフィックカード上で実施され、使用されたモデルはすべて標準的なTransformer-Bigでした。

上記と同じ結論に至ります。バッチサイズが大きくなるにつれて、量子化推論の速度向上は徐々に加速します。LightSeq fp16と比較すると、さらに約70%の速度向上を達成し、オンライン翻訳モデルの推論レイテンシを大幅に削減できます。

最後に、上の画像に示すように、自動GEMMチューニングの効果を示すために、LightSeqテストでA100グラフィックスカード上のTransformerモデルとBERTモデルのレイテンシを、fp16とint8でチューニング前後で比較しました。チューニング前は、一部の形状でint8のGEMM速度がfp16よりもさらに遅かったのに対し、チューニング後は全般的にfp16を上回っていることがわかります。

ビデオメモリ使用量

LightSeqは、異なるバッチサイズにおけるメモリ使用量の観点から、量子化モデルと浮動小数点モデルの比較分析を行いました。バッチサイズが大きくなるにつれて、量子化モデルのメモリ使用量の優位性がより顕著になり、メモリ使用量を最大約30%削減できることがわかります。さらに、LightSeq fp16エンジンはPyTorchモデルと比較してメモリ使用量を大幅に削減します。そのため、LightSeq int8エンジンは最終的にメモリ使用量を最大約68%削減できます。

量子化モデル効果

複数の言語方向と複数のテスト セットにわたる機械翻訳の場合、LightSeq は標準の Transformer-Big モデルを使用して、浮動小数点モデル BLEU と比較した量子化モデル推論の損失をテストしました。

en2zhタスクの大規模データセットでは、LightSeq int8は浮動小数点モデルと比較して比較的大きなBLEU損失を示し、最大で-0.4に達しました。しかし、en2esタスクの小規模データセットでは、LightSeq int8はパフォーマンス損失を示さないだけでなく、浮動小数点モデルよりも優れたパフォーマンスを示しました。全体として、int8量子化モデルの平均BLEUは、浮動小数点モデルと比較して実質的にロスレスです。LightSeqは、GLUEやSQuADを含む複数のタスクにおいて、量子化モデルの有効性も検証しました。

勾配通信量子化

通信ボトルネックは複数マシン、複数GPUのシナリオでより顕著になるため、勾配通信量子化は主に分散学習シナリオで使用されます。そのため、LightSeqは2台のマシンと8つのGPUを搭載したA100プロセッサ上で分散学習の速度テストを実施しました。

勾配通信量子化による学習加速効果は、入力データが増加するにつれて一般的に弱まることがわかります。これは主に、入力データが増加すると計算時間の割合が増加する一方で、勾配通信時間の割合が減少し、勾配量子化のメリットも減少するためです。

LightSeqには、ネットワークインターフェースカード(NIC)の数を変えた追加のトレーニング速度テストも含まれています。勾配通信量子化を用いた分散トレーニング速度は、オリジナルのLightSeq fp16と比較して大幅に向上していることがわかります。

定量技術

int8 定量取引の加速されたメリットは、主に次の側面から生まれます。

  1. GEMM の精度を fp16 から int8 に下げると計算時間が短縮されます。
  2. 入力と出力に int8 を使用した後、カスタム演算子を使用することでデータの読み取りと書き込みの時間が短縮されます。
  3. int8 を使用して勾配を保存することで、複数のマシン間の通信時間が短縮されます。

Transformerモデルを例に挙げると、LightSeq fp16エンジンによる高速化により、カスタム演算子の処理時間が大幅に短縮される一方で、GEMM処理時間が全体の約90%を占めています。そのため、最適化の焦点はGEMM高速化へと移行します。fp16 GEMMをint8 GEMMに置き換えることで、GEMM処理時間が短縮されるだけでなく、前後の演算子の入出力ビット幅も削減され、データの読み書き時間も短縮されます。最後に、マルチマシントレーニングにおけるボトルネックは主に勾配通信であり、勾配をint8精度に量子化することで、分散トレーニングを大幅に高速化できます。

定量原則

量子化による精度の低下を補うため、量子化を考慮したトレーニングは通常、量子化プロセスをシミュレートするために使用されます。上図に示すように、量子化を考慮したトレーニングでは、float GEMMの2つのfloat入力(擬似量子化ノードと呼ばれる)を量子化および逆量子化し、それらをセグメント化された浮動小数点入力に離散化し、float GEMM演算を実行します。その後、結果は再び量子化および逆量子化され、最終的な浮動小数点結果が得られます。量子化プロセスは微分不可能であるため、量子化パラメータの勾配を推定するにはSTE法が必要です。量子化を考慮したトレーニングに擬似量子化ノードが挿入され、float GEMMを使用して量子化プロセスをシミュレートする理由は、TensorFlowやPyTorchなどのトレーニングフレームワークがint8 GEMMをサポートしていないためです。

LightSeqの量子化トレーニングでは、int8 GEMMを直接使用することで量子化プロセスを正確に再現し、従来の実装よりも高速でメモリ効率に優れています。推論時には、int8 GEMM演算に離散化された整数を使用し、最終的に結果を浮動小数点数に逆量子化します。量子化推論プロセスは量子化トレーニングと完全に一致しており、従来の量子化を考慮したトレーニングと完全に同等です。

定量的ポジション

量子化Transformer全体のネットワーク構造は上図のようになっています。赤い矢印は、量子化ノードと逆量子化ノードを追加する必要がある場所を示しています。

まず、int8 GEMMのすべての入力と出力を量子化する必要があります。int8 GEMMの形状制限により、一部のGEMM(アテンションスコアの計算など)では依然としてfloat GEMMが使用されています。さらに、第2層FFNのGEMMは、GEMM入力がReLU活性化関数の出力であり、正数のみを含み非対称であるため、int32出力を使用しています。そのため、int8出力のGEMMを使用した場合、正しい浮動小数点結果に逆量子化することは不可能です。

全てのモデル重みはint8型で保存する必要があるため、重みの量子化が必要になります。しかし、バイアス重みはパラメータ数が少ないため量子化は不要です。そのため、float型の精度を維持することで、モデルのパフォーマンスを向上させることができます。

最後に、デコーダーのキャッシュを量子化する必要があります。デコーダーのキャッシュは推論中に頻繁に読み書きされる必要があるため、キャッシュをint8に量子化することでデコード速度が大幅に向上します。

定量戦略

浮動小数点行列をint8整数行列に量子化する方法はたくさんあります。LightSeqは対称量子化を採用しており、浮動小数点の正負の範囲を整数範囲[-127, 127]に対称的にマッピングします。

現実には、浮動小数点行列の数値範囲は通常対称ではなく、外れ値も非常に少ないため、外れ値の範囲に基づいて行列を直接量子化すると、量子化の精度に影響を及ぼします。そのため、まず行列を切り捨てる必要があります。

LightSeqは切り捨てにPACT法[6]を用い、切り捨て範囲をモデルの学習可能なパラメータとして扱い、STEアルゴリズムを用いてパラメータの勾配を推定し、バックプロパゲーションによる最適化を行う。実践的な経験に基づくと、重みの初期切り捨て範囲を[-1, 1]、中間結果の初期切り捨て範囲を[-16, 16]に設定すると、ほとんどのタスクで最良の結果が得られる。最終的に、切り捨て範囲と他のモデルパラメータの共同最適化により、量子化モデルは実質的にロスレスな性能を実現できる。

勾配通信量子化

分散学習シナリオ向けに、LightSeqは勾配量子化圧縮技術を導入しました。これは、浮動小数点精度の勾配をint8で量子化することで勾配通信にかかる時間を削減し、学習を高速化する技術です。これは勾配通信量子化(GCQ)と呼ばれます。

上の図に示すように、勾配通信量子化の主なプロセスは次のとおりです。

  1. 各カードの勾配のカットオフ範囲を計算します。
  2. カットオフ範囲に対して全削減最大演算を実行します。
  3. 各カードは均一なカットオフ範囲を使用して、int8 を使用して勾配を量子化します。
  4. int8 勾配に対して全削減合計演算を実行します。
  5. 各カードは、all-reduce 後に勾配を逆量子化し、浮動小数点勾配に復元してから、パラメータを更新します。

全削減処理中の int8 勾配のオーバーフロー問題に対処するため、LightSeq はまず各 GPU 上の浮動小数点勾配を GPU 数で除算し、次に除算前の切り捨て範囲を使用して量子化し、最後に全削減演算を実行します。これにより、全削減処理後に各 GPU 上で量子化された int8 整数がオーバーフローするのを防ぎます。ただし、これにより GPU あたりの量子化に実際に使用されるビット数も削減されます。そのため、現在のソリューションは 8 GPU を搭載した 2 台のマシンではほぼロスレスですが、GPU の数が増えるにつれてトレーニングのパフォーマンスが低下します。たとえば、en2de および en2fr 翻訳タスクでは、8 GPU を搭載した 4 台のマシンで分散量子化トレーニングを行うと、BLEU スコアがそれぞれ約 0.4 と 1.5 低下します。LightSeq は今後もこの問題を解決するためのより良い方法を模索していきます。

一般技術

前の章で説明した量子化手法に加えて、LightSeq のこのアップデートでは、量子化モデルだけでなく、他のすべての精度モデルのトレーニングと推論にも適用できるいくつかの一般的な最適化手法も導入されています。

演算子融合

上図は、エンコーダモジュールの量子化トレーニングの計算グラフです。LightSeqは、2つのGEMM演算間のすべての演算を1つの演算子に統合することで[7]、カーネル呼び出し回数を減らし、全体の計算時間を短縮します。

図中の黄色の四角形はint8 GEMM、緑の四角形はfloat GEMMを表しています。ここでは形状の制約によりfloat GEMMが使用されており、int8 GEMMは高速化に適していません。赤い矢印は流れるデータ型がint8であることを示しており、緑の矢印は第2層FFNのGEMM出力がint32データ型であることを示しています。int8 GEMMの入出力に対する量子化および逆量子化演算はフロントカーネルとリアカーネルに統合されており、データの移動量を削減するだけでなく、GPUメモリの使用量も削減します。

LightSeqは推論中にデコーダーの最適化も行います。上図に示すように、自己注意を計算する際、注意スコアの次元は(バッチサイズ, 1, シーケンス長)です。そのため、値積を計算する際にGEMM演算を回避し、重み付き加算演算子を直接記述することで、図の破線枠内の計算を単一のカーネルに統合できます。

自動ビデオメモリ管理

モデル量子化により、より複雑なテンソル型と依存関係が導入され、メモリ管理に新たな課題が生じます。これに対処するため、LightSeqは新たなメモリ管理メカニズムを設計しました。上図に示すように、このメカニズムは主に以下のプロセスで構成されます。

  1. 学習開始前に、各テンソルの有効期間とメモリサイズは、各演算子の位相依存性に基づいて自動的に計算されます。動的な次元を含むテンソルは、機械翻訳タスクにおける最大文長や最大バッチサイズなど、その次元の最大値に基づいて計算されます。これらの最大値は学習前に指定されます。
  2. テンソルの寿命とサイズを決定した後、メモリ再利用の関係を分析します。寿命が重複しないテンソルは、単一のメモリ空間を共有できます。すべてのメモリ空間はデータ型を持たず、任意のデータ型のテンソルに割り当てることができます。
  3. テンソルのメモリ再利用関係に基づいて、複数のメモリ セグメントが割り当てられ、メモリの実際の開始アドレスと終了アドレスが各テンソルに割り当てられます。

テンソルメモリの再利用を分析する際に、LightSeqは論文[3]で提案されたオフセット計算のための貪欲なサイズ法を借用し、3つの改善を加えました。

  1. トレーニング プロセス全体を通じてメモリの再利用 (前方/後方) をサポートします。
  2. 異なるデータ型でメモリの再利用を実現できます (int8/fp16/fp32)。
  3. 単一の非常に大きなメモリ空間を使用する代わりに、すべてのテンソルが複数のメモリ セグメントにまたがって収容されるため、メモリ使用率が効果的に向上します。

自動GEMMチューニング

LightSeqのint8 GEMMは、NVIDIAのcuBLASLtライブラリを使用しています。これは現在、NVIDIAグラフィックスカード上で最も効率的な行列演算ライブラリです。しかし、GEMMが使用する最適な構成(データ配置、GEMMアルゴリズムなど)は、入力データの形状やグラフィックスカードによって異なる場合があり、自動選択が必要です。LightSeqの自動チューニングスキームは以下の通りです。

  1. システムは、さまざまなグラフィック カード (T4 や A100 など) 上のさまざまな形状の最適な GEMM 構成を検索し、その結果を構成ファイルに保存します。この構成ファイルはユーザーがダウンロードできます。
  2. モデルの初期化中に、対応するグラフィックカードモデルの設定ファイルが読み込まれ、解析され、キーと値のペア(形状、最適な構成)の辞書に保存されます。対応するグラフィックカードモデルの設定ファイルがない場合、または必要なGEMM形状が利用できない場合は、ユーザーは手動で検索して保存するか、デフォルトの構成を直接使用するかを選択できます。
  3. モデルの順方向計算または逆方向計算では、入力形状に基づいて辞書内で最適な構成が検索され、その後GEMM計算が実行されます。対応する形状が見つからない場合は、デフォルトの構成が直接使用されます。

今後の仕事

今後、LightSeq は、モバイル デバイスでの低精度の量子化、バックプロパゲーションにおける勾配量子化、大規模モデルの量子化などの分野を継続的に探求していきます。

引用

[1] Wang, Xiaohui, et al. 「LightSeq2: GPU上でのTransformerベースモデルの高速トレーニング」arXivプレプリントarXiv:2110.05722 (2021)。

[2] Micikevicius, Paulius, et al. 「混合精度トレーニング」arXivプレプリントarXiv:1710.03740 (2017).

[3] Pisarchyk, Yury, Juhyun Lee. 「ディープニューラルネット推論のための効率的なメモリ管理」arXivプレプリントarXiv:2001.03288 (2020).

[4] Jacob, Benoit, et al. 「効率的な整数演算のみの推論のためのニューラルネットワークの量子化とトレーニング。」IEEEコンピュータビジョンとパターン認識会議の議事録。2018年。

[5] Alistarh, Dan, et al. 「QSGD: 勾配量子化と符号化による通信効率の高いSGD」神経情報処理システムの進歩 30 (2017)。

[6] Choi, Jungwook, et al. 「Pact: 量子化ニューラルネットワークのためのパラメータ化されたクリッピング活性化」arXivプレプリントarXiv:1805.06085 (2018)。

[7] Wang, Xiaohui, et al. 「LightSeq: Transformersのための高性能推論ライブラリ」arXivプレプリントarXiv:2010.13887 (2020)。