モデルのボトルネックを「ルート」から見つけよう！第一原理からディープラーニングを分析する

モデルのパフォーマンスを向上させたい場合、まず検索エンジンに問い合わせるのが本能でしょうか?

通常、表示される提案は、インプレース操作の使用、グラデーションを None に設定する、PyTorch バージョンを 1.10.1 から安定バージョン 1.10.0 に戻すなどの、いくつかの難しい操作だけです。

これらの一時的に見つかったトリックは、現在の問題を一時的に解決できますが、使用後にパフォーマンスが満足できるレベルまで改善されない場合は、少し「盲目的」である可能性があります。

ディープラーニング自体は構成要素のブラックボックスモデルですが、このデバッグ方法により、ディープラーニングは科学というよりは錬金術のようなものになるようです。

たとえば、トレーニングセットでのモデルの損失がテスト中の損失よりもはるかに低い場合、モデルが「過剰適合」していることを意味します。この時点でモデルパラメータの数を盲目的に増やすと、完全に時間の無駄になります。たとえば、モデルのトレーニング損失と検証損失が同じ場合、モデルに正則化を追加するのは時間の無駄になります。

そこで、AI 実践者が問題に遭遇した後に根本的な問題を解決できるようにするために、コーネル大学人工知能 (CUAI) の共同創設者である Horace He 氏は最近ブログ記事を公開し、ディープラーニングモデルの時間損失をコンピューティング、メモリ、その他のオーバーヘッドの 3 つの部分に分類し、ディープラーニングモデルを「第一原理」から理解して改善する方法について説明しました。

コンピューティングは、浮動小数点演算を計算するときに GPU が消費する時間、つまり FLOPS を指します。メモリは、テンソルを GPU に書き込むときに消費される時間を指します。

モデルがほとんどの時間をメモリ転送に費やしている場合、GPU の FLOPS を増やしても意味がありません。あるいは、大量の数学演算を実行することにすべての時間を費やしている場合、オーバーヘッドを削減するためにモデルロジックを C++ で書き直すのは意味がありません。

現在の状況を把握することで、最適化の範囲を絞り込むことができ、時間を節約して、安心して怠けることができます。

計算する

通常、ディープラーニングモデルが十分に高速でない理由は、グラフィックカードのパフォーマンスが十分でないためです。新しいグラフィックカードを追加すると、問題が解決します。

しかし、現実は非常に厳しい。カードの性能が高ければ高いほど、価格も高くなる。したがって、コストパフォーマンスを高めるには、グラフィックカードの動作効率を可能な限り向上させ、グラフィックカードをマトリックス内で継続的に実行する必要があります。

コンピューティングがメモリ帯域幅よりも重要であるもう 1 つの理由は、モデルのトレーニング中に必要なコンピューティングの量は、どのような手段を使用しても基本的に削減されないため、コンピューティング能力を最大化することで効率を向上できることです。

ただし、コンピューティングの量が急速に増加すると、コンピューティングの利用率を最大化することが難しくなります。 CPU FLOPS の倍増時間とメモリ帯域幅の倍増時間の表をご覧ください。

コンピューティングについて考える一つの方法は、CPU が工場であると考えることです。ユーザーは、工場を効率的に稼働させるため (コンピューティング)、指示 (オーバーヘッド) と原材料 (メモリ帯域幅) を工場に送信します。

工場の効率が、利用可能な原材料の供給よりも速く上昇する場合、工場が最高効率に到達することはより困難になります。工場の規模（FLOPS）が 2 倍になっても、同時に帯域幅が増加しなければ、パフォーマンスは 2 倍になりません。

FLOPS についてもう 1 つ付け加えておきたいことがあります。最新の機械学習アクセラレーションハードウェアには、Nvidia の Tensor コアなど、行列乗算専用のハードウェアがあります。

つまり、行列乗算を行わない場合、宣伝されている 312 テラフロップスではなく、19.5 テラフロップスしか得られません。これは GPU に特有の欠陥ではありません。TPU は GPU よりもさらに汎用性が低くなります。

実際、GPU はすべての非行列乗算演算で遅く、一見すると大きな影響があるように見えるかもしれませんが、実際には、ニューラルネットワークモデルは基本的に行列乗算です。

BERT モデルの失敗研究では、BERT の 99.8% が行列乗算 (テンソル収縮) 演算であることが判明したため、非行列乗算は 15 倍遅くなりますが、大きな問題ではありません。

しかし、この場合、正規化とドット単位の演算では、行列乗算演算に比べて実際には 250 倍、700 倍少ない FLOPS しかかかりません。

非行列乗算の理論的なパフォーマンスが現実と大きく異なる理由について、研究者らが出した答えは「メモリ帯域幅」です。

メモリ

帯域幅コストは、基本的に、CPU から GPU へのデータ転送や、あるノードから別のノードへのデータ転送など、ある場所から別の場所へデータを移動するために支払われるコストです。通常、この 2 つは「データ転送コスト」と「ネットワークコスト」と呼ばれます。

ディープラーニングモデルの最適化の帯域幅コストは、主に CUDA グローバルメモリから CUDA 共有メモリに転送されます。

工場の例に戻ると、工場はいくつかのコンピューティングタスクを実行できますが、大量のデータを保存するのに適した場所ではありません。一般的なアプローチは、より安価なハードウェア (DRAM) を使用してデータウェアハウスを構築し、倉庫と工場の間で材料を輸送することです (これがメモリ帯域幅です)。

GPU の DRAM サイズは、nvidia-smi コマンドで取得できます。ストレージ容量の不足も、CUDA メモリ不足エラーの主な原因です。

GPU カーネルが実行されるたびに、データを GPU の DRAM との間で移動する必要があることに注意することが重要です。

これで、torch.cos のような単一の操作を実行する場合、このような単純な操作を実行するたびに、データをメモリから GPU に転送する必要があることが分かりました。転送コストは計算コストよりもはるかに高いため、ほぼすべての時間がメモリ内で費やされます。この状況は、メモリバウンド操作とも呼ばれます。

間違ったアプローチは、毎回計算のためにデータを GPU に送信し、結果を返し、その結果を再び計算のために GPU に送信することです。データ転送に多くの時間が費やされていることがわかります。

少し調整するだけで、命令を事前に計算に組み込むと、同じタスクを実行するためのメモリ転送が 1 回に削減されます。

pyTorch コードを使用すると、2 行のコードを 1 行の x.cos().cos() に変換できるため、効率が 2 倍になります。

ただし、この最適化対策はすべてのシナリオに適用できるわけではありません。 GPU は実行されるすべての命令を事前に把握し、CUDA コードを生成する必要があるため、eager モードでは使用できません。また、すべての演算子の融合が点単位の演算子のように単純なわけではありません。

CUDA カーネルコードを書いたことがあるなら、任意の 2 つの PyTorch カーネルを融合して、グローバルメモリの読み取りと書き込みのコストを節約できることをご存知でしょう。 NVFuser や XLA などの既存のコンパイラは通常、いくつかの単純な融合しか実行できず、AI エンジニアの設計ほど優れているわけではありません。自分でカスタム CUDA カーネルを書いてみたい場合は、Triton から始めるのが良いでしょう。

演算子の融合の効果は、同じ時間コストでより多くの演算が行われることです。そのため、gelu には relu よりも明らかに多くの演算があるにもかかわらず、活性化関数の計算コストはほぼ同じになります。

計算機は、操作がメモリ帯域幅に制限されているかどうかを判断する必要がある場合に非常に役立ちます。

単純な演算子の場合、メモリ帯域幅について推論するのは簡単です。たとえば、A100 は 1.5T バイト/秒のグローバルメモリ帯域幅を持ち、19.5T FLOPS の計算を実行できます。したがって、32 ビット浮動小数点 (つまり 4 バイト) を使用すると、GPU は 20 兆回の演算を実行するのと同じ時間で 4000 億の数値を読み込むことができます。さらに、単純な単一の操作 (テンソルを 2 倍にするなど) を実行するには、実際にはテンソルをグローバルメモリに書き戻す必要があります。したがって、1 回の操作を 100 回程度実行すれば、メモリデータを送信できるようになります。

NVFuser のような融合コンパイラーの助けを借りれば、コストを測定するのは実は非常に簡単です。

PyTorch 関数を例にとり、それを Fusion コンパイラーでベンチマークすると、さまざまな繰り返し値で達成される FLOPS とメモリ帯域幅を計算できます。

繰り返し回数を増やすことは、メモリアクセスを増やすことなく計算量を増やす簡単な方法であり、計算強度の増加としても知られています。

テンソルのサイズは N なので、2 * N のメモリアクセスと N * repeat FLOP が実行されます。したがって、達成されるメモリ帯域幅は byte_per_elem * 2 * N / itrs_per_second となり、達成される FLOPS は N * repeat / itrs_per_second となります。

実行時間、フロップス、および達成されたメモリ帯域幅の対数をプロットすると、64 回の乗算が実行されるまで実行時間に顕著な増加がないことがわかります。これは、以前はメモリ帯域幅が制限されており、計算はほとんどアイドル状態であったことも意味します。

その結果、当初は 0.2 テラフロップスしか達成できませんでした。コンピューティングの強度が 2 倍になると、この数値は直線的に増加し、最終的には「コンピューティングの限界」である 9.75 テラフロップスのピークに近づきます。

メモリ帯域幅はピークに近いところから始まり、計算の強度が増すにつれて低下し始めます。これは、メモリへのアクセスよりも実際の計算に多くの時間が費やされるため、予想されたことです。

この場合、いつ計算にバインドされているのか、いつメモリにバインドされているのかが簡単にわかります。

繰り返し回数が 32 回未満の場合は、メモリ帯域幅が飽和し、計算能力が十分に活用されません。逆に、繰り返し回数が 64 回を超えると、計算ワークロードが飽和し (つまり、ピーク FLOPS に近づき)、メモリ帯域幅の使用率が低下し始めることがわかります。

大規模なシステムの場合、両方の組み合わせである可能性があるため、コンピューティング制約なのかメモリ帯域幅制約なのかを判断するのが難しいことがよくあります。

コンピューティング能力の限界を測定する一般的な方法は、達成した FLOPS をピーク FLOPS のパーセンテージとして測定することです。ピーク FLOPS の 80% が達成された場合、コンピューティングリソースが完全に活用され、残りの時間はメモリ帯域幅に費やされる可能性があります。

その他の費用

Python インタープリターで費やされる時間、PyTorch フレームワークで費やされる時間、CUDA カーネルの起動 (実行ではない) に費やされる時間など、テンソルの転送や計算に費やされないコード内の時間はすべてオーバーヘッドと呼ばれます。

オーバーヘッドが問題となる主な理由は、最新の GPU が非常に高速であるためです。 A100 は 1 秒あたり 312 兆回の浮動小数点演算 (312 テラフロップス) を実行できます。それに比べて、Python は非常に遅く、1 秒あたり 3200 万回の加算しか実行できません。

これは、Python が 1 FLOP を実行できる時間内に、A100 が 975 万 FLOP を実行できることも意味します。

PyTorch のようなフレームワークにも、実際のカーネルに入る前に多くのスケジューリングレイヤーがあります。同じ実験を PyTorch で行うと、1 秒あたり 280,000 回の操作しか実行できません。もちろん、小さなテンソルを実行することは PyTorch が構築された目的ではありませんが、科学計算で小さなテンソルを使用すると、PyTorch が C++ に比べて驚くほど遅いことがわかります。

より直感的なグラフは、PyTorch が加算を実行するときに生成されるプロファイルが、小さな四角形を除いて純粋なオーバーヘッドであることを示しています。

最新のディープラーニングモデルは通常、大規模な計算操作を実行し、PyTorch などのフレームワークはそれを非同期的に実行します。つまり、PyTorch が CUDA カーネルを実行している間も、その背後でさらに CUDA カーネルを実行し続け、キューに入れることができます。したがって、PyTorch が CUDA カーネルを「事前に」実行できる限り、フレームワークのオーバーヘッドのほとんどは完全に隠されます。

通常、オーバーヘッドは問題のサイズに応じて拡大縮小しないため (コンピューティングとメモリは比例して拡大縮小します)、バッチサイズを 2 倍にしても実行時間が 10% しか増加しない場合 (実行時間が 2 倍になることが予想されます)、オーバーヘッドが高すぎる可能性が高いと簡単に判断できます。

別のアプローチは、PyTorch プロファイラーを使用することです。ピンクの線は、CPU コアが GPU コアにどれだけ一致しているかを示しています。 GPU が CPU オーバーヘッドを待機している間、多くのギャップが発生します。

CPU が GPU よりも高速に実行されると、ギャップが少なくなります。

nvidia-smi の GPU-Util は、実際に実行されている GPU コアの割合を測定します。これは、オーバーヘッドを測定するのにも適した方法です。

オーバーヘッドのほとんどは、PyTorchのようなフレームワークの柔軟性から生じており、「何をすべきかを考える」のに多くの時間を必要とします。

たとえば、a+b を実行する場合、次の 3 つの手順が必要です。

1. Pythonは__add__がディスパッチしたものを見つける必要がある

2. PyTorchは、どのカーネルを呼び出すかを決定するために、テンソルの多くのプロパティ（dtype、デバイス、Augradが必要かどうかなど）を決定する必要があります。

3. PyTorchは実際にカーネルを起動する必要がある

各ステップには、さまざまな操作をサポートする柔軟性が必要です。この柔軟性に対処する 1 つの方法は、jit.tract、FX、jax.jit などのトレースを実装するか、CUDA グラフを使用してより低いレベルでトレースを実装することです。

モデルの効率を向上させるために最も重要なことは、モデルのパフォーマンスのボトルネックを理解することです。

もちろん、ニューラルネットワークモデルを作成するには、多くのオーバーヘッドの問題を考慮する必要がありますが、これは、ユーザーにとって透過的である必要があるため、これらのシステムとフレームワークの設計上の失敗とも言えます。