PyTorch は、大規模モデル用の「アクセラレーション パッケージ」を作成し、1,000 行未満のコードで 10 倍の速度を実現します。 NVIDIA Scientist: minGPT 以来の最高のチュートリアル リポジトリの 1 つ

PyTorch チームは、大規模モデルの推論を10 倍高速化しました。

そして、純粋なネイティブ PyTorch コードは 1000 行未満しか使用しません。

このプロジェクトはGPT-fastと呼ばれ、加速効果は次のようになります。

滑らかですよ、本当に滑らかです！

重要な点は、チームがコードと詳細な「チュートリアル」を直接リリースしたことです。棒人間バージョンなので、とてもわかりやすいです。

開発チームのメンバーであるホレス・ヒー氏は次のように述べています。

私たちはこれをライブラリやフレームワークとしてではなく、ニーズに応じて「コピーして貼り付ける」ことができる例として捉えています。

ネットユーザーはすぐに衝撃を受けた。NvidiaのAI科学者ジム・ファン氏は次のようにコメントした。

これは、Andrej Karpathy が minGPT をリリースして以来、最高のチュートリアル スタイルのリポジトリの 1 つです。

オープンソースの世界には、minGPT や GPT-Fast のようなプロジェクトがさらに必要です。

では、GPT-fast はどのようにして大規模モデルを高速化するのでしょうか?

大型模型「アクセラレーターパック」開封

一般的には、以下の方法が使用されます。

• Torch.compile : モデルの実行効率を向上できる、PyTorch モデル専用に設計されたコンパイラー。
• GPU 量子化: 計算の精度を下げることでモデルの計算を高速化します。
• 投機的デコード: より小さなモデルを使用してより大きなモデルの出力を予測し、大規模な言語モデルの操作を高速化します。
• テンソル並列処理: モデルの操作を複数のハードウェア デバイスに分散することで処理を高速化します。

これらについて一つずつ詳しく説明しましょう。

開発チームは当初、シンプルな PyTorch を使用して実装しましたが、効果は良くありませんでした (25.5 tok/s)。

トレースを調べたところ、CPU の使用量が多すぎるために推論パフォーマンスが制限されていたことが原因の 1 つであることがわかりました。

それで、どうやって解決するのでしょうか?

GPU が巨大な工場 (利用可能な計算能力が非常に大きい) であり、CPU が工場に「物資を運ぶ」小さなカートであるというシナリオを想像してみてください。

多くの場合、CPU は GPU に十分な速さでデータを「供給」することができません。

そのため、開発チームは、GPU にさらに多くのワークロードを与えるか、一度に処理するタスクの「チャンク」を大きくすることを推奨しています。

推論中にこれを行うには、 torch.compileをインポートできます。

torch.compile は、モデルのより大きな領域をキャプチャし、それらを単一のコンパイル済み領域にコンパイルすることができます。特に「オーバーヘッド削減」モードで実行する場合、CPU オーバーヘッドの削減に非常に効果的です。

効果はすぐに現れ、パフォーマンスは 25 tok/s から 107 tok/s へと 4 倍向上しました。

次に、開発チームは速度をさらに向上させたいと考えましたが、メモリ帯域幅のボトルネックが発生しました。

開発チームはモデルの帯域幅使用率を計算し、それが 72% に達することを発見しました。

つまり、速度をさらに向上させる余地は限られている可能性があるということです。

上記の式を再検討した結果、チームは、モデル パラメータの数や GPU のメモリ帯域幅を実際に変更することはできないものの (少なくとも、追加費用をかけずには)、各パラメータを格納するために使用されるバイト数を変更できることを発見しました。

つまり、モデルのサイズを変更したりハードウェアをアップグレードしたりしてもパフォーマンスは向上しませんが、モデル パラメータを保存するために必要なデータの量を減らすことで効率を高めることができます。

これは通常、各パラメータを表すために必要なビット数を削減する量子化技術によって実現されます。

その結果、開発チームは次の技術であるint8 量子化を導入しました。

int8 重み量子化を使用すると、メモリ負荷が軽減され、パフォーマンスがさらに向上します (157.4 tok/s)。

量子化を使用した後には別の問題が発生します。100 個のトークンを生成するには、モデルの重みを 100 回読み込む (または呼び出す) 必要があります。モデルの重みを頻繁にロードすると、非効率になる可能性もあります。

一見すると、自己回帰生成モデルには厳密なシリアル依存性があるため、解が存在しないように思えます。

しかし開発チームは、この厳密な順次依存性は投機的デコードを活用することで解消できると指摘しました。

別の例として、技術的な判断は常に正しいが、コードを書くのが比較的遅い、Verity という名前の上級エンジニアを想像してください。

同時に、Drake という名前のジュニア エンジニアがいます。Verity とは異なり、技術的な決定を下すのは得意ではありませんが、より速く、より低コストでコードを書くことができます。

では、さまざまな人々の強みを活用して全体的な効率を向上させるにはどうすればよいでしょうか?

アプローチはシンプルです。まずは Drake にコードを書かせ、その過程で技術的な決定を下すのです。次に、コードはレビューのために Verity に引き渡され、間違いがあれば Drake がやり直します。

Transformer モデルの推論では、大規模な検証モデルは Verity ロールであり、Drake はテキストをより速く生成できる小規模なドラフト モデルです。

開発チームはドラフト モデルを使用して 8 つのトークンを生成し、検証モデルを使用してそれらを並列処理し、一致しないトークンを破棄しました。

これにより、シリアル依存性が解消され、速度が再び向上します。

投機的デコードでは出力の品質は変わらないことに留意してください。このアプローチは、トークンを生成し、ドラフト モデルを使用してそれらのトークンを検証するのに必要な時間が、それらのトークンを個別に生成するのに必要な時間よりも短い限り機能します。

実際、ネイティブ PyTorch を使用してこの手法を実装するのは非常に簡単で、実装プロセス全体で約 50 行のネイティブ PyTorch コードのみが必要です。

AMD は Triton と torch.compile バックエンドもサポートしているため、以前に Nvidia GPU に適用されたすべての最適化を AMD GPU にも再適用できます。

開発チームは、int8 量子化の速度が 22 tok/s から 102 tok/s に向上したことを確認しました。

その後、開発チームは int4 量子化を使用して速度をさらに向上させましたが、モデルの精度は低下しました。

したがって、重みサイズを削減するために、グループ量子化と GPTQ が使用されます。

最終的に、精度を確保しながら速度は 202.1 tok/s に増加しました。

上記のテクノロジーを組み合わせることで、244.7 tok/s というより高速な速度を実現できます。

これまで、R&D チームは単一の GPU 上でプロセスを高速化してきました。しかし、実際には、多くの場合、複数の GPU を使用できます。

複数の GPU を使用するとメモリ帯域幅が増加し、モデルの全体的なパフォーマンスが向上します。

並列処理戦略を選択する場合は、トークンの処理を複数のデバイスに分割する必要があるため、テンソル並列処理を使用する必要があります。

PyTorch は、torch.compile と組み合わせて使用​​できるテンソル並列処理用の低レベル ツールも提供します。

開発チームは、テンソル並列性を表現するためのより高レベルの API も開発中であることを明らかにしました。

ただし、高レベル API がなくても、テンソルの並列処理を追加するのは簡単で、150 行のコードで実現でき、モデルを変更する必要もありません。

前述のすべての最適化は、テンソル並列処理と組み合わせることができます。これらの最適化を組み合わせると、Llama-70B で 55 トークン/秒の int8 量子化が可能になります。

最後に、結果をまとめると、量子化を無視すると、高速推論、投機的デコード、テンソル並列処理を実現するために、わずか 766 行のコード (model.py 244 行のコード、generate.py 371 行のコード、tp.py 151 行のコード) が使用されています。

Llama-7B の場合、コンパイル + int4 量子化 + 投機的デコードを使用すると、速度は 241 tok/s に達します。 Llama-70B では、テンソル並列処理を追加することで 80 tok/s を達成します。

これらのパフォーマンスは、現在の SOTA に近いか、それを上回ります。