超速い！ ByteDanceが配列推論エンジンLightSeqをオープンソース化

これは、Transformer や GPT などの複数のモデルの高速推論を完全にサポートする業界初のオープンソース エンジンになります。

[[348155]]

2017年にGoogleはTransformer[1]モデルを提案しました。それ以来、BERT[2]やGPTシリーズ[13]など、このモデルをベースにした優れた事前学習済み言語モデルや機械翻訳モデルが数多く誕生し、多くの自然言語処理タスクの機能を継続的に刷新してきました。同時に、これらのモデルのパラメータの数もほぼ指数関数的に増加しています (下の図を参照)。例えば、最近熱い議論を巻き起こしているGPT-3 [3]は1750億個のパラメータを持ち、パラメータ数の記録を再び更新しました。

このような膨大な数のパラメータは、モデルの推論と展開にも課題をもたらします。機械翻訳を例にとると、WMT[4]コンペティションにおけるSOTAモデルは50層以上に達しています。主流のディープラーニングフレームワークでは、文章を翻訳するのに数秒かかります。これにより、2 つの問題が発生します。1 つ目は、変換時間が長すぎるため、製品のユーザー エクスペリエンスに影響します。2 つ目は、シングル カードの QPS (1 秒あたりのクエリ レート) が低すぎるため、サービス コストが過度に高くなることです。

したがって、今日は、非常に高速で多くの機能をサポートする高性能なシーケンス推論エンジン、LightSeq をお勧めします。 Transformer ベースのシーケンス特徴抽出器 (Encoder) と自己回帰シーケンス デコーダー (Decoder) に徹底的な最適化を施しました。2019 年 12 月にオープンソース化され、Volcano Translation を含む多くのビジネスやシナリオに適用されています。これは、Transformer や GPT などの複数のモデルの高速推論を完全サポートする業界初のオープンソース エンジンになると思われます。

LightSeq は、機械翻訳、自動質疑応答、インテリジェントライティング、対話応答生成など、多くのテキスト生成シナリオに適用でき、オンラインモデル推論の速度を大幅に向上させ、ユーザーエクスペリエンスを向上させ、企業のサービス運用コストを削減します。

他のオープンソースのシーケンス推論エンジンと比較して、LightSeq には次の利点があります。

1. 高いパフォーマンス

LightSeq の推論速度は非常に高速です。たとえば、翻訳タスクでは、LightSeq は Tensorflow と比較して最大 14 倍の高速化を実現できます。また、これは他の現在のオープンソースのシーケンス推論エンジンよりも優れており、たとえば、Faster Transformer よりも最大 1.4 倍高速です。

2. 複数のモデル機能をサポート

LightSeqはBERT、GPT、Transformer、VAEなどの多くのモデルをサポートし、ビームサーチ、ダイバースビームサーチ[5]、サンプリングなどの複数のデコード方法をサポートしています。次の表は、テキスト生成シナリオにおける3つの推論エンジン、Faster Transformer[7]、Turbo Transformers[6]、LightSeqの機能の違いを詳細に示しています。

3. 使いやすく、TensorflowやPyTorchなどのディープラーニングフレームワークにシームレスに接続します。

モデル プロトコルを定義することにより、LightSeq はさまざまなディープラーニング フレームワークによってトレーニングされたモデルの柔軟なインポートをサポートします。また、すぐに使用できるエンドツーエンドのモデル サービスも含まれているため、コードを 1 行も書かずに高速モデル推論を展開できると同時に、マルチレベルの再利用も柔軟にサポートします。

使い方

LightSeq を使用してオンライン サービスを展開するのは比較的簡単です。 LightSeqは、実用的なサービスミドルウェアを多数含むNvidiaのオープンソースGPU推論サーバーであるTriton Inference Server[8]をサポートしています。 LightSeq は、サーバーのカスタム推論エンジン API をサポートしています。したがって、訓練されたモデルがLightSeq [9]で定義されたモデルプロトコルにエクスポートされる限り、コードを記述することなく、ワンクリックでエンドツーエンドの効率的なモデルサービスを開始できます。モデル構成（レイヤー数や埋め込みサイズなど）の変更も簡単にサポートされます。具体的なプロセスは以下のとおりです。

まず、モデルリポジトリを準備します。以下はディレクトリ構造の例です。transformer.pb はモデルプロトコルに従ってエクスポートされたモデル重み、libtransformer.so は LightSeq のコンパイル済み製品です。

 - model_zoo/- model_repo/- config.pbtxt- transformer.pb- 1 /- libtransformer.so

その後、Triton推論サーバー[8]を起動し、モデルサービスを構築します。

 trtserver --model-store=${model_zoo}

パフォーマンステスト

ディープラーニング フレームワーク Tensorflow v1.13 と、デコード シナリオのサポートがより豊富な Faster Transformer v2.1 実装を比較対象として使用し、NVIDIA Tesla P4 および NVIDIA Tesla T4 グラフィック カード上で LightSeq のパフォーマンスをテストしました。 Turbo Transformers のデコード方法は比較的単純です (Beam Search のみをサポートし、テキスト生成でよく使用されるサンプリング デコードはサポートしていません)。実際のアプリケーションのニーズをまだ満たしていないため、比較されていません。

機械翻訳のパフォーマンス

機械翻訳のシナリオでは、ビームサーチデコードを使用して、Transformer ベースモデル (6 層エンコーダー、6 層デコーダー、隠れ層次元 512) のパフォーマンスをテストしました。実験結果は次のとおりです。

小規模バッチのシナリオでは、Faster Transformer と LightSeq はどちらも Tensorflow と比較して約 10 倍の高速化を達成したことがわかります。バッチ サイズが大きくなるにつれて、行列乗算演算の割合が増加するため、Tensorflow の両方の加速率が低下します。 LightSeq は、特に大規模なバッチ シナリオでは比較的スムーズに減衰し、Faster Transformer よりも最大 1.4 倍高速になります。これは、将来の推論最適化作業に対するガイダンスも提供します。小規模バッチ シナリオでは、計算集約的でない演算子が統合されている限り、高い加速化のメリットが得られます。一方、大規模バッチ シナリオでは、行列乗算などの計算集約的な演算子を継続的に最適化する必要があります。

最後に、WMT14 標準のフランス語 - 英語翻訳タスクで Transformer ビッグモデルのパフォーマンスをテストしました。 Tesla P4 グラフィックス カード上の LightSeq の 1 文あたりの平均翻訳遅延は 167 ミリ秒ですが、Tesla T4 では 82 ミリ秒に短縮されます。比較すると、TensorFlow のレイテンシは 1071 ミリ秒ですが、LightSeq はそれぞれ 6.41 倍と 13.06 倍の加速を実現します。さらに、著者は他の多くのモデル構成を試し、比較的一貫した加速効率を得ました。たとえば、より深いモデル構造（エンコーダーが 16 層に深くなる）では、LightSeq によって達成される加速比はそれぞれ 6.97 倍と 13.85 倍になります。

テキスト生成パフォーマンス

前述の機械翻訳では通常、デコードにビームサーチ方式が使用されますが、テキスト生成のシナリオでは、生成される結果の多様性を向上させるためにサンプリングが必要になることがよくあります。次の図は、top-k/top-p サンプリングを使用した Transformer ベース モデルのパフォーマンス比較を示しています。

サンプリング デコードを必要とするタスクでは、LightSeq はほとんどの構成で Faster Transformer よりも優れており、最大 1.4 倍の追加加速を実現できることがわかります。さらに、TensorFlow 実装と比較して、LightSeq は GPT や VAE などの生成モデルに対して 5 倍以上の高速化効果も実現します。

サービスストレステスト

クラウド サービスでは、実際のアプリケーションで GPT シナリオでモデル サービスが Tensorflow から LightSeq に切り替わったときのレイテンシの変化をテストしました (使用したサービス グラフィック カードは NVIDIA Tesla P4)。 pct99 遅延が 3 ～ 5 倍短縮され、ピーク値が約 360 ミリ秒から約 80 ミリ秒に低下していることがわかります。詳細な結果を次の図に示します。

より詳細な比較実験結果については、LightSeqの性能評価レポート[10]をご覧ください。

技術原理

Transformer を例にとると、機械翻訳/テキスト生成モデルの推論プロセスは、シーケンス エンコーディング モジュールの特徴計算と自己回帰デコード アルゴリズムの 2 つの部分で構成されます。特徴計算部分は、自己注意メカニズムと特徴変換（行列乗算、計算集約型）を中心としており、Elementwise（Reshapeなど）やReduce（Layer Normalizationなど）などの多数のIO集約型操作を伴います。デコードアルゴリズムには、語彙Softmax、ビームスクリーニング、キャッシュリフレッシュなどのプロセスが含まれます。計算は単純であり、より複雑な動的形状を導入します。これにより、モデル推論に多くの課題が生じます。

1. IO 集約型コンピューティングのためのきめ細かいカーネル関数呼び出しにより、大量の冗長なビデオ メモリの読み取りと書き込みが発生し、機能コンピューティング パフォーマンスのボトルネックになります。

2. 複雑な動的形状は計算グラフの最適化に課題をもたらし、モデル推論中に大量の動的メモリ要求が発生し、時間がかかります。

3. 文字のデコードと生成の各ステップのロジックが複雑で、ハードウェアを活用するために計算を並列化することが困難です。

LightSeq は優れた推論加速効果を実現します。これらの課題に対処するために、どのようなターゲットを絞った最適化が行われましたか?著者が分析した結果、コア技術には、複数のコンピューティング操作を統合して IO オーバーヘッドを削減すること、ビデオ メモリを再利用して動的アプリケーションを回避すること、およびデコード アルゴリズムを階層的に書き換えて推論速度を向上させることなどが含まれることがわかりました。以下では、各部分の最適化の課題と LightSeq のソリューションについて詳しく紹介します。

オペレータマルチオペレーション融合

近年、Transformer エンコーダー/デコーダー構造は、その効率的な特徴抽出機能により、大量のラベルなしテキストの事前トレーニングなど、さまざまな NLP タスクで広く使用されています。ほとんどのディープラーニング フレームワーク (Tensorflow、Pytorch など) は通常、基本操作ライブラリのカーネル関数を呼び出して、エンコーダー/デコーダーの計算プロセスを実装します。これらのカーネル関数は細かく分割されていることが多く、通常、コンポーネントはそれを実装するために複数のカーネル関数を呼び出す必要があります。

レイヤー正規化を例にとると、Tensorflow は次のように実装します。

平均 = tf.reduce_mean(x, axis=[- 1 ], keepdims=True) 分散 = tf.reduce_mean(tf.square(x - 平均), axis=[- 1 ], keepdims=True) 結果 = (x - 平均) * tf.rsqrt(分散 + イプシロン) * スケール + バイアス

コンパイル最適化テクノロジ (ブロードキャスト操作と要素ごとの操作を自動的に統合) に基づいても、カーネル関数の呼び出しが 3 回 (reduce_mean が 2 回、最終結果を計算するために 1 回)、中間結果 (平均と分散) のメモリ読み取りと書き込みが 2 回必要であることがわかります。 CUDA に基づいて、レイヤーの正規化専用のカーネル関数をカスタマイズし、2 つの中間結果をレジスタに書き込むことができます。これにより、カーネル関数の呼び出しを 1 回実装するだけで、中間結果をビデオ メモリに読み込んだり書き込んだりする必要がなくなり、計算オーバーヘッドが大幅に削減されます。興味のある学生は、記事の最後にある参考リンクで詳細な実装を確認することができます[11]。

この考えに基づき、LightSeqはCUDA行列演算ライブラリcuBLAS[12]が提供する行列乗算とカスタムカーネル関数を使用してTransformerを実装します。具体的な構造を下図に示します。

青い部分はカスタムカーネル関数、黄色い部分は行列乗算です。行列乗算間のすべての演算はカスタマイズされたカーネル関数を使用して実装されており、これによりカーネル関数の呼び出しとビデオメモリの読み取りと書き込みが大幅に削減され、最終的に計算速度が向上することがわかります。

動的メモリ再利用

計算プロセス中にビデオ メモリが解放されるのを防ぎ、ビデオ メモリの使用量を抑えるために、LightSeq はまずモデル内のすべての動的シェイプの最大値 (最大シーケンス長など) を定義し、すべての動的シェイプを静的シェイプに変換します。そして、サービスが開始すると、計算プロセスにおける各中間結果の最大値に応じてビデオメモリが割り当てられ、依存関係のない中間結果についてはビデオメモリが共有されます。このようにして、モデルは各リクエストに対して推論中にビデオ メモリを要求しなくなり、次のことが実現されます。異なるリクエストの同じテンソルはビデオ メモリを再利用します。同じリクエストの異なるテンソルは、形状と依存関係に応じてビデオ メモリを再利用します。

このビデオ メモリ再利用戦略により、LightSeq は T4 グラフィックス カード上に最大 8 つの Transformer ビッグ モデル (batch_size=8、最大シーケンス長 = 8、beam_size = 4、vocab_size = 30,000) を同時に展開できます。これにより、低頻度またはオフピークのシナリオでのグラフィック カードの使用率が大幅に向上します。

階層的デコード計算

自己回帰シーケンス生成シナリオでは、最も複雑で時間のかかる部分はデコードです。 LightSeq は現在、ビーム検索、ダイバーシティ ビーム検索、トップ k/トップ p サンプリングなどの複数のデコード方法をサポートしており、Transformer および GPT と組み合わせて使用​​することで数倍の高速化を実現できます。ここでは、最も広く使用されているビーム検索を例に、LightSeq のデコード プロセスの最適化を紹介します。

まず、ディープラーニング フレームワークで 1 ステップのデコード計算が従来どのように実行されるかを見てみましょう。

# 1. 各トークンで終わるシーケンスの対​​数確率を計算する

log_token_prob = tf.nn.log_softmax(logit) # [batch_size, beam_size, vocab_size]log_seq_prob += log_token_prob # [batch_size, beam_size, vocab_size]log_seq_prob = tf.reshape(log_seq_prob, [- 1 , beam_size * vocab_size])

# 2. 各シーケンス（バッチ要素）のtopkトークンを見つける

topk_log_probs、topk_indices = tf.nn.top_k(log_seq_prob、k=K)

# 3. ビームIDに従ってデコーダーのセルフアテンションモジュールのキーと値のキャッシュを更新する

refresh_cache(キャッシュ、topk_indices)

確率的に上位k個のトークンを選択するには、サイズ[batch_size、beam_size、vocab_size]のロジット行列に対してソフトマックス計算とビデオメモリの読み書きを実行し、その後batch_sizeソートを実行する必要があることがわかります。通常、vocab_size は数万単位なので計算量は非常に大きくなりますが、これは 1 回のデコード手順の計算コストにすぎません。したがって、実際には、自己回帰シーケンス生成タスクにおけるデコード モジュールの累積遅延が大きな割合 (30% 以上) を占めていることがわかります。

LightSeq の革新性は、GPU の計算特性と、検索推奨でよく使用される粗選択と細分類の 2 段階戦略を組み合わせ、デコード計算を階層型に書き換え、ロジット粗選択カーネル関数を設計して、ソフトマックスの計算と数十万の要素の分類をうまく回避していることにあります。粗カーネル関数はロジット行列を 2 回走査します。

• 最初のトラバーサルでは、各ビームのロジット値がランダムに k グループに分割され、各グループの最大値が求められ、次にこれらの k 個の最大値の最小値が求められます。この最小値は、おおよそのトップ k 値 (実際のトップ k 値以下である必要があります) として機能し、R-top-k として示されます。トラバーサル プロセス中に、ビーム内の logit の log_sum_exp 値を同時に計算できます。

• 2 回目のトラバーサルでは、各ビームについて、R-top-k 以上のすべての logit 値を見つけ、(logit - log_sum_exp + batch_id * offset、beam_id * vocab_size + vocab_id) を候補キューに書き込みます。ここで、offset は logit の下限です。

最初のパスでは、ロジット値は通常、正規分布に従うため、計算された R-top-k 値は実際の top-k 値に非常に近くなります。同時に、このステップではレジスタの読み取りと書き込みのみが行われ、アルゴリズムの複雑さが低いため、迅速に実行できます (12 以上の命令サイクル)。トップ 4 設定では、R-top-k に基づいて数万のトークンからおよそ 12 個の候補のみが選択されることが確認されており、これは非常に効率的です。 2 回目のトラバーサルでは、R-top-k に基づいて候補が大まかに選択され、ロジット値が batch_id に従ってオフセットされ、複数のスレッドが同時にビデオ メモリ内の候補キューに書き込まれます。

大まかな選択が完了したら、候補キューでソートを実行してバッチ全体の各シーケンスの正確なトップ k 値を取得し、キャッシュを更新して、ワンステップのデコード処理を迅速に実行します。

以下は、k=2、語彙サイズ=8の場合の具体的な例です（列は出力される文字数、行は各位置の候補を表します）。本来は 16 個の要素をソートする必要があることがわかりますが、階層的デコードを採用した後は、最終的に 5 個の要素のみをソートすればよくなり、ソートの複雑さが大幅に軽減されます。

計算遅延を視覚化して分析する

上記の最適化手法の実際の効果を確認するために、著者は GPU プロファイル ツールを使用して LightSeq 推論プロセスのレイテンシ分析を実行しました。次の図は、32 ビット浮動小数点数と 16 ビット浮動小数点数の精度における各コンピューティング モジュールのレイテンシ比を示しています。

2 つの計算精度では次のことがわかります。

1. 最適化後、cuBLAS の行列乗算計算の遅延はそれぞれ 82% と 88% を占め、推論高速化の新たな主なボトルネックとなりました。対照的に、Tensorflow モデルをテストしたところ、行列乗算のレイテンシはわずか 25% でした。これは、LightSeq のビーム検索最適化によってレイテンシが非常に低いレベルにまで削減されたことを示しています。

2. キャッシュリフレッシュはそれぞれ10%と6%を占めており、割合も高いのですが、最適化を継続することは困難です。将来的には、キャッシュ量を減らす（デコーダー レイヤーの数を減らす、キャッシュ精度を下げるなど）ことで、レイテンシをさらに削減することができます。

3. その他の操作は、レイヤー正規化、ビーム検索、中間結果のメモリ読み取りと書き込みなど、合計で 8% と 6% を占めます。

視覚化の結果は、LightSeq が非常に最適化され、推論速度が大幅に向上したことを示しています。

ポータル:

GitHub プロジェクト アドレス:

https://github.com/bytedance/lightseq