PyTorch チームは「すべてを分割」モデルを書き直しました。これは元の実装より 8 倍高速です

Meta の「すべてを分割」モデルを最適化するにはどうすればよいでしょうか? PyTorch チームが書いたこのブログは、単純なものから深いものまで、この問題に答えるのに役立ちます。

今年の初めから現在まで、生成 AI は急速に発展しました。しかし、多くの場合、特に PyTorch を使用する場合、生成 AI のトレーニングや推論などをどのように高速化するかという難しい問題に直面しなければなりません。

この記事では、PyTorch チームの研究者がソリューションを提供します。この記事では、純粋なネイティブ PyTorch を使用して生成 AI モデルを高速化する方法に焦点を当てており、PyTorch の新しい機能とそれらを組み合わせる方法の実践例も紹介しています。

#結果はどうなりましたか? PyTorch チームは、Meta の「Split Everything」(SAM) モデルを書き直した結果、精度を損なうことなく元の実装よりも 8 倍高速なコードが得られ、すべてネイティブ PyTorch を使用して最適化されたと述べました。

ブログ アドレス: https://pytorch.org/blog/accelerated-generative-ai/

参照この記事を最後まで読むと、次のことがわかるでしょう:

• Torch.compile: PyTorch モデル コンパイラーである PyTorch 2.0 には、torch と呼ばれる新しい関数が追加されました。コンパイル ()、1 行のコードで既存のモデルを高速化できます;
• GPU 量子化: 計算精度を下げることでモデルを高速化します;
• SDPA (スケーリング ドット プロダクト アテンション): メモリ効率の高いアテンション実装;
• 半構造化 (2:4) スパースネス: GPU 用に最適化されたスパース メモリ フォーマット;
• ネストされたテンソル: ネストされたテンソルは、{tensor, Mask} を一緒にパックして、異なるサイズの画像など、不均一なサイズのデータ​​を単一のテンソルにバッチ処理します。
• Triton カスタム オペレーション: Triton Python DSL を使用して GPU オペレーションを記述し、カスタム オペレーター登録を通じてそれらを PyTorch のさまざまなコンポーネントに簡単に統合します。

# 「PyTorch によってスループットが向上し、メモリ オーバーヘッドが削減される」のネイティブ機能。

SAM は Meta によって提案されました。この研究の詳細については、「

CV はもう存在しません? リリースされました」を参照してください。 by Meta「すべてを分割する」AI モデル、CV が GPT-3 の瞬間を迎える可能性がある

」

次に、この記事では、パフォーマンス分析、ボトルネックの特定、およびこれらの新機能を PyTorch に統合して SAM が直面する問題を解決する方法など、SAM の最適化プロセスを紹介します。さらに、この記事では、PyTorch のいくつかの新機能 (torch.compile、SDPA、Triton カーネル、Nested Tensor、および半構造化スパース性) についても紹介します。

この記事の内容はステップバイステップで詳しく説明されています。記事の最後に、SAM の高速バージョンが紹介されます。興味のある友人は、次のリンクからダウンロードできます。 GitHub. また、この記事では Perfetto UI を使用して、各 PyTorch 機能のアプリケーション価値を示すためにデータが視覚化されています。

GitHub アドレス: https://github.com/pytorch-labs/segment-anything-fast

すべてのモデルを分割する SAM の書き換え

調査によると、この記事で使用されている SAM ベースライン データ型は float32 dtype、バッチ サイズは1、および PyTorch Profiler でカーネル トレースを表示した結果は次のとおりです:

この記事では、SAM には最適化できる箇所が 2 つあることがわかりました:

1 つ目は、テンソル インデックスによって引き起こされる aten::index への長い呼び出しです。基礎となる呼び出しによって引き起こされる操作 ([] など)。ただし、GPU が aten::index に費やす実際の時間は比較的短く、その理由は、2 つのコアを起動するプロセス中に、aten::index がコア間の cudaStreamSynchronize をブロックするためです。これは、2 番目のコアが起動されるまで、CPU が GPU の処理が完了するのを待つことを意味します。したがって、SAM を最適化するには、アイドル時間を引き起こす GPU 同期のブロックを排除するよう努めるべきであるとこの文書では考えています。

2 つ目は、SAM が行列の乗算 (上の画像の濃い緑色) に多くの GPU 時間を費やすことです。これは、Transformers では一般的です。 SAM モデルが行列の乗算に費やす GPU 時間を削減できれば、SAM を大幅に高速化できる可能性があります。

次に、この記事では、SAM のスループット (img/s) とメモリ オーバーヘッド (GiB) を使用してベースラインを確立します。その後、最適化プロセスが始まります。

Bfloat16 半精度 (および GPU 同期とバッチ処理)

解決するには上記の問題を解決する、つまり行列の乗算にかかる時間を短縮するには、この記事では bfloat16 を取り上げます。 Bfloat16 は一般的に使用される半精度タイプで、各パラメーターとアクティベーションの精度を下げることで、計算時間とメモリを大幅に節約できます。

bfloat16 を使用してパディング タイプを置き換えます

削除GPU 同期については、この記事で最適化できる場所が 2 つあることがわかりました。

具体的には (理解しやすいように上の図を参照してください。表示される変数名はすべてコード内にあります)、研究では次のことがわかりました。画像エンコーダには、座標スケーラーとして機能する変数 q_coords と k_coords があり、これらの変数は CPU 上で割り当てられ、処理されます。ただし、これらの変数が rel_pos_resize でのインデックス付けに使用されると、これらのインデックス付け操作によってこれらの変数が GPU に自動的に移動され、このコピーによって GPU 同期が行われます。上記の問題を解決するために、研究では、この部分を上記のように torch.where を使用して書き換えることで解決できることに注目しました。

カーネル トレース

これらの変更を適用した後、この記事では、単一のカーネルが特に小さなバッチ (ここでは 1) の場合、呼び出しの間にはかなりの時間間隔があります。この現象をより深く理解するために、この記事ではバッチ サイズ 8 での SAM 推論のパフォーマンス分析から始めます。コア、この記事 SAM の GPU 時間のほとんどが要素ごとのカーネルとソフトマックス操作に費やされていることが観察されています。

これで、行列乗算の相対コストがはるかに小さいことがわかります。

GPU 同期と bfloat16 最適化を組み合わせると、SAM パフォーマンスが 3 倍向上します。

Torch.compile（グラフ ブレークと CUDA グラフ）

この記事では、詳細な処理の過程で多くの小さな操作があることがわかりました。 SAM の研究. コンパイラを使用して操作を融合することには大きな利点があると考えられているため、PyTorch は torch.compile に対して次の最適化を行いました:

• Change nn.LayerNorm または nn.GELU などの一連の操作は単一の GPU カーネルに融合されます;
• #行列乗算カーネルの直後に操作を融合して GPU の数を削減しますカーネル呼び出し。

#これらの最適化により、研究では GPU グローバル メモリの往復回数が削減され、推論が高速化されました。 SAM の画像エンコーダで torch.compile を試すことができるようになりました。パフォーマンスを最大化するために、この記事ではいくつかの高度なコンパイル手法を使用します:

カーネル トレース

結果は、torch.compile が非常にうまく機能することを示しています。

#softmax が時間の大部分を占め、その後にさまざまな GEMM バリアントが続いていることがわかります。次の測定値は、バッチ サイズ 8 以上のものです。

SDPA:scaled_dot_product_attention

次に、この記事では SDPA (scaled_dot_product_attention) についてレビューします。 ) 実験では、研究の焦点は注意メカニズムです。一般に、ネイティブ アテンション メカニズムは、時間とメモリのシーケンスの長さに応じて二次関数的にスケールします。 PyTorch の SDPA 操作は、Flash Attendant、FlashAttendantV2、および xFormer のメモリ効率の高いアテンション原則に基づいて構築されており、GPU アテンションを大幅に高速化できます。この操作を torch.compile と組み合わせると、MultiheadAttendant のバリアントでの共通パターンの表現と融合が可能になります。小さな変更を加えた後、モデルはscaled_dot_product_attentionを使用できるようになりました。

カーネル トレース

メモリ効率の高いアテンション、カーネルの占有率、膨大な計算時間を確認できるようになりました。 GPU の場合:

PyTorch のネイティブのscaled_dot_product_attentionを使用すると、バッチ サイズを大幅に増やすことができます。以下のグラフは、バッチ サイズ 32 以上の変化を示しています。

その後、研究では Triton、NestedTensor、バッチ処理 Predict_torch、int8 量子化、半構造化 (2:4) スパース性、その他の操作も実験しました。

たとえば、この記事ではカスタム位置 Triton カーネルを使用し、バッチ サイズ 32 での測定結果を観察します。

Nested Tensor を使用し、バッチ サイズ 32 以上のバリエーション。

#量子化を追加した後のバッチ サイズ 32 以上の測定。

記事の最後は半構造化スパーシティです。この研究は、行列の乗算が依然として直面する必要のあるボトルネックであることを示しています。解決策は、スパース化を使用して行列乗算を近似することです。スパース行列 (つまり、値をゼロにする) により、重みとアクティベーション テンソルを格納するために使用できるビットが少なくなります。テンソル内のどの重みがゼロに設定されるかを設定するプロセスは、枝刈りと呼ばれます。より小さい重みを取り除くと、精度を大幅に損なうことなくモデルのサイズを削減できる可能性があります。

枝刈りには、完全に構造化されていないものから高度に構造化されたものまで、さまざまな方法があります。非構造化枝刈りは理論的には精度に最小限の影響を与えますが、GPU は大規模な密行列の乗算を行う際には非常に効率的ですが、疎な場合には大幅なパフォーマンス低下に見舞われる可能性があります。最近 PyTorch でサポートされたプルーニング手法は、半構造化 (または 2:4) スパース性と呼ばれるバランスをとることを目的としています。この疎なストレージにより、元のテンソルが 50% 削減され、同時に高密度のテンソル出力が生成されます。以下の図を参照してください。

#このスパース ストレージ形式と関連する高速カーネルを使用するには、次に行うべきことは重みを取り除くことです。この記事では、スパース度 2:4 でプルーニングするための最小の 2 つの重みを選択します。重みをデフォルトの PyTorch (「ストライド」) レイアウトからこの新しい半構造化スパース レイアウトに変更するのは簡単です。 apply_sparse (モデル) を実装するには、32 行の Python コードのみが必要です。

この記事では、スパース度 2:4 で、バッチ サイズが 2:4 のときの vit_b と SAM を観察します。 32 ピーク パフォーマンス:

最後に、この記事を一文で要約します: この記事では、一連の新しい機能を利用して、これまでのところ PyTorch で最速の Segment Anything 実装を紹介します。正式にリリースされた関数。この記事では、精度を損なうことなく、純粋な PyTorch で元の SAM を書き直しています。

興味のある読者は、元のブログで詳細を確認してください。

^{参考リンク: https://pytorch.org/blog/accelerated-generative-ai/}

以上がPyTorch チームは「すべてを分割」モデルを書き直しました。これは元の実装より 8 倍高速ですの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。