Back to Articles

This article is also available in Chinese 简体中文.

モデルサイズが拡大するにつれ、生成AIの実装には大量の推論リソースが必要になります。これにより生成あたりのコストだけでなく、こうしたリクエストを処理するための消費電力も増加します。

テキスト生成向けの推論最適化は、レイテンシ、インフラコスト、消費電力の削減に不可欠です。これにより、ユーザー体験の向上とテキスト生成タスクの効率化につながります。

アシストデコーディングは、テキスト生成を高速化するための一般的な手法です。私たちはこれをIntel Gaudi向けに適応・最適化しました。Intel Gaudiは以前の投稿で示されたようにNvidia H100 GPUと同等の性能を発揮しつつ、価格はNvidia A100 80GB GPUと同水準です。この取り組みは現在Optimum Habanaの一部となっており、TransformersやDiffusersなどさまざまなHugging Faceライブラリを拡張して、AIワークフローをIntel Gaudiプロセッサ向けに完全に最適化できるようにしています。

Speculative Sampling - Assisted Decoding

投機的サンプリングは、テキスト生成を高速化するために用いられる手法です。ドラフトモデルがK個のトークンを生成し、それらをターゲットモデルで評価するという仕組みで動作します。ドラフトモデルが却下された場合はターゲットモデルで次のトークンを生成し、このプロセスを繰り返します。投機的サンプリングを使用することで、テキスト生成の速度を向上させつつ、自己回帰サンプリングと同等のサンプリング品質を実現できます。この手法では、テキスト生成時にドラフトモデルを指定できます。大規模なTransformerベースモデルでは約2倍の高速化が示されています。全体として、これらの手法はIntel Gaudiプロセッサ上でのテキスト生成を加速し、パフォーマンスを向上させることができます。

ただし、ドラフトモデルとターゲットモデルはサイズが異なるためKVキャッシュに異なる形で表現され、別々の最適化戦略を同時に活用することが課題となります。本記事では量子化モデルを想定し、KVキャッシングと投機的サンプリングを併用します。各モデルには独自のKVキャッシュがあり、ドラフトモデルがK個のトークンを生成してターゲットモデルで評価される点に注意してください。ドラフトモデルが却下された場合はターゲットモデルで次のトークンを生成し、ドラフトモデルが次のK個のトークンを生成してプロセスを繰り返します。

著者[2]は、投機的サンプリングを実行するとターゲット分布が回復することを証明しており、これはターゲットモデル自体での自己回帰サンプリングと同等のサンプリング品質を保証します。したがって、投機的サンプリングを活用しないことが不利になるケースは、ドラフトモデルの相対的なサイズによる節約が十分でない場合や、ドラフトモデルの受理率がドラフトモデルのサイズの小ささを活かせないほど低い場合です。

投機的サンプリングと類似の手法として、アシスト生成があります。これはほぼ同時期に独立して開発されたものです[3]。著者はこの手法をHugging Face Transformersに統合し、.generate()呼び出しに任意パラメータassistant_modelを追加してこの手法を有効化できるようにしました。

Usage & Experiments

アシスト生成の使い方は簡単です。使用例はこちらで提供されています。 想定通り、--assistant_modelパラメータでドラフトモデルを指定します。ドラフトモデルがK個のトークンを生成してターゲットモデルで評価され、ドラフトモデルが却下された場合はターゲットモデルで次のトークンを生成します。ドラフトモデルが次のK個のトークンを生成してプロセスを繰り返します。ドラフトモデルの受理率は入力テキストに一部依存します。一般的に、大規模なTransformerベースモデルでは約2倍の高速化を確認しています。

Conclusion

アシスト生成によるGaudiでのテキスト生成の高速化がサポートされ、簡単に利用可能になりました。これによりIntel Gaudiプロセッサ上でのパフォーマンス向上が図れます。この手法は投機的サンプリングに基づいており、大規模なTransformerベースモデルのパフォーマンス向上に有効であることが示されています。

References

[1] N. Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need,” Nov. 2019. arXiv:1911.02150.

[2] C. Chen, S. Borgeaud, G. Irving, J.B. Lespiau, L. Sifre, and J. Jumper, “Accelerating Large Language Model Decoding with Speculative Sampling,” Feb. 2023. arXiv:2302.01318.

[3] J. Gante, “Assisted Generation: a new direction toward low-latency text generation,” May 2023, https://huggingface.co/blog/assisted-generation.