FP8 以下の浮動小数点量子化精度は、もはや H100 に限定されません。

Huang 氏は誰もが INT8/INT4 を使用するように望んでいましたが、Microsoft DeepSpeed チームは NVIDIA からの公式サポートなしにA100 上で FP6 を実行することに成功しました。

テスト結果によると、新しい方法 TC-FPx は、A100 上で INT4 に近いか、場合によっては INT4 を上回る FP6 量子化速度を実現し、後者よりも高い精度を実現しています。

これを基に、エンドツーエンドの大規模モデルのサポートも提供されており、これはオープンソース化され、DeepSpeed などのディープラーニング推論フレームワークに統合されています。

この成果は、大規模モデルの高速化にも即座に効果を発揮します。このフレームワークで 1 枚のカードを使用して Llama を実行すると、2 枚のカードを使用した場合よりも 2.65 倍のスループットが達成されます。

これを見たある機械学習研究者は、マイクロソフトの研究は狂気としか言いようがないと語った。

絵文字パックもすぐにリリースされました。次のようになります:

Nvidia: H100 のみが FP8 をサポートします。

Microsoft: 分かりました。自分で対処します。

では、このフレームワークはどのような効果を達成できるのでしょうか。また、その背後にはどのようなテクノロジーが使われているのでしょうか。

Llama を FP6 で実行すると、シングル カードの方がデュアル カードよりも高速になります。

A100 で FP6 精度を使用すると、カーネルレベルのパフォーマンスが向上します。

研究者らは、さまざまなサイズのLlamaおよびOPTモデルから線形レイヤーを選択し、NVIDIA A100-40GB GPUプラットフォームでCUDA 11.8を使用してテストしました。

NVIDIA 公式の cuBLAS (W16A16)や TensorRT-LLM (W8A16)と比較すると、 TC-FPx (W6A16) の方が高速です。最大でそれぞれ2.6倍と1.9倍の改善が見られました。

4 ビットの BitsandBytes (W4A16)方式と比較すると、TC-FPx は最大 8.9 倍の速度向上を実現します。

（WとAはそれぞれ重み付き量子化ビット幅と活性化量子化ビット幅を表す）

△cuBLASの結果を1として正規化されたデータ

同時に、TC-FPx コアは DRAM メモリへのアクセスを削減し、DRAM 帯域幅の使用率、Tensor コアの使用率、ALU および FMA ユニットの使用率を向上させます。

TC-FPx に基づいて設計されたエンドツーエンドの推論フレームワーク FP6-LLMも、大規模モデルに大幅なパフォーマンスの向上をもたらします。

Llama-70B を例にとると、シングル カード上の FP6-LLM のスループットは、デュアル カード セットアップ上の FP16 の 2.65 倍であり、バッチ サイズが 16 未満の場合はレイテンシも FP16 よりも低くなります。

OPT-30B (FP16 も 1 枚のカードを使用) などのパラメータの少ないモデルの場合、FP6-LLM によってスループットの大幅な向上とレイテンシの削減も実現されます。

さらに、このような条件下では、単一の FP16 カードは最大バッチ サイズ 4 しかサポートできませんが、FP6-LLM はバッチ サイズ 16 で正常に動作できます。

では、Microsoft チームはどのようにして A100 で FP16 量子化を実行できたのでしょうか?

カーネルソリューションの再設計

最大 6 ビットの精度をサポートするために、TC-FPx チームは、異なるビット幅の量子化重みをサポートできる統合カーネル スキームを設計しました。

従来のデュアルカーネル アプローチと比較して、TC-FPx は、逆量子化と行列乗算を単一のカーネルに統合することで、メモリ アクセスを削減し、パフォーマンスを向上させます。

低精度の量子化を実現するための基本原理は、逆量子化によって FP6 精度のデータを FP16 として「偽装」し、それを GPU に渡して FP16 形式で計算することです。

チームはまた、ビットレベルのプリパッケージング技術を利用して、GPU メモリ システムが 2 の累乗以外のビット幅 (6 ビットなど) に不向きであるという問題にも対処しました。

具体的には、ビットレベルの事前パッケージ化では、モデル推論の前に重みデータを再編成し、6 ビットの量子化された重みを並べ替えて GPU メモリに適した方法でアクセスできるようにします。

さらに、GPU メモリ システムは通常、32 ビットまたは 64 ビットのブロックでデータにアクセスするため、ビット レベルのプリパッキングでは 6 ビットの重みもパックされ、これらの整列したブロックに格納してアクセスできるようになります。

事前パッケージ化後、研究チームは SIMT コアの並列処理機能を使用して、レジスタ内の FP6 重みに対して並列逆量子化を実行し、FP16 形式の重みを生成しました。

逆量子化されたFP16重みはレジスタ内で再構成され、Tensorコアに送られます。再構成されたFP16重みは行列乗算に使用され、線形層の計算が完了します。

このプロセスでは、チームは SMIT コアのビットレベルの並列性を利用して、逆量子化プロセス全体の効率を向上させました。

重み再構築タスクを並行して実行するために、チームは並列重みスプライシング技術も使用しました。

具体的には、各重みは複数の部分に分割され、各部分のビット幅は 2 の累乗です(例: 6 を 2+4 または 4+2 に分割) 。

逆量子化の前に、重みはまず共有メモリからレジスタにロードされます。各重みは複数の部分に分割されているため、実行時にレジスタレベルで完全な重みを再構築する必要があります。

TC-FPxは、実行時のオーバーヘッドを削減するために、重みの抽出と連結を並列化する手法を提案しています。この手法では、2組のレジスタを用いて32個のFP6重みの断片を保存し、これらの重みを並列に再構築します。

一方、重みの抽出と連結を並列に行うには、初期データレイアウトが特定の順序要件を満たしていることを確認する必要があります。そのため、TC-FPxは実行前に重みの断片を並べ替えます。

さらに、TC-FPx は、逆量子化ステップと Tensor Core の行列乗算演算を統合するソフトウェア パイプラインを設計し、命令レベルの並列処理によって全体的な実行効率を向上させています。

論文リンク: https://arxiv.org/abs/2401.14112