模型の列車の速度を上げるコツは何でしょうか?まず、この問題の第一原理を理解しましょう。

誰もがモデルをより速くトレーニングしたいと考えていますが、本当に適切なアプローチを探していますか?コーネル大学の学部生であり、PyTorch チームのインターンでもある Horace He 氏の見解では、この問題はいくつかのステップで解決する必要があります。まず、トレーニングが遅い理由、つまりボトルネックがどこにあるかを把握し、それに応じた解決策を見つける必要があります。基本原理（第一原理）を理解せずに何かをしようとするのは時間の無駄です。

この記事では、Horace He がコンピューティング、メモリ帯域幅、オーバーヘッドという 3 つの観点から考えられるボトルネックを分析し、現在どのボトルネックが発生しているかを判断する方法をいくつか紹介します。これにより、より的を絞った方法でシステムを高速化できるようになります。この記事は、Chen Tianqi 氏をはじめとする多くの上級研究者や開発者から賞賛されました。

以下は元のコンテンツです。

ディープラーニングモデルのパフォーマンスを向上させるにはどうすればよいでしょうか?ほとんどの人は、オンライン ブログにまとめられている「システムの組み込み演算子を使用する、勾配を 0 に設定する、PyTorch バージョン 1.10.1 ではなく 1.10.0 を使用する」などのランダムなトリックを選択します。

現代のシステム（特にディープラーニング）が科学というよりも錬金術のように感じられるこの分野では、ユーザーがこのランダムなアプローチに惹かれるのも理解できます。それでも、この分野には従うべきいくつかの基本原則があり、それによって多数の方法を排除し、問題を解決しやすくなります。

たとえば、トレーニング損失がテスト損失よりもはるかに低い場合、過剰適合の問題が発生している可能性があり、モデル容量を増やすために時間を無駄にしている可能性があります。たとえば、トレーニング損失と検証損失が同じである場合、モデルを正規化するのは賢明ではありません。

同様に、効率的なディープラーニングの問題を 3 つの異なる要素に分けることができます。

1. 計算: GPU が実際の浮動小数点演算の計算に費やす時間 (FLOPS)。
2. メモリ: GPU 内でテンソルを転送するのにかかった時間。
3. 追加のオーバーヘッド: 他のコンポーネントに費やされる時間。

機械学習モデルをトレーニングするときにどのような問題に直面しているかを知ることは、モデルを効率的にするという問題と同様に重要です。たとえば、モデルがメモリから GPU への転送に多くの時間を費やす場合 (つまり、メモリ帯域幅が狭い場合)、GPU の FLOPS を増やしても効果はありません。一方、行列乗算を大量に実行している場合 (つまり、計算負荷が高い場合)、オーバーヘッドを軽減するためにプログラムを C++ で書き直しても役に立ちません。

したがって、GPU をスムーズに実行したい場合は、上記の 3 つの側面についての議論と調査が不可欠です。

痛い教訓の背後には、GPU を効率的に稼働させ続ける多数のエンジニアがいます。

注: このブログのコンテンツのほとんどは GPU と PyTorch の例に基づいていますが、原則は基本的にハードウェアとフレームワーク間で共通です。

計算する

ディープラーニング システムを最適化する際の 1 つの側面は、計算に費やす時間を最大化することです。 312 テラフロップスの料金を支払ったので、それをコンピューティングに使用できるようにしたいと考えています。しかし、高価な行列乗算から費用対効果を得るには、他の部分に費やす時間を削減する必要があります。

しかし、なぜここではメモリ帯域幅の最大化ではなく計算の最大化に重点が置かれているのでしょうか?理由は簡単です。オーバーヘッドやメモリ消費量を削減することはできますが、実際の操作を変更しないと計算量を削減することはほとんどできないからです。

メモリ帯域幅と比較した計算速度の増加により、計算利用率を最大化することが難しくなります。次のグラフは、CPU FLOPS を 2 倍にし、メモリ帯域幅を 2 ​​倍にするのにかかる時間を示しています (黄色の列に注目してください)。

コンピューティングについて考える方法の 1 つは、それを工場として考えることです。私たちは工場に指示を伝え（過剰消費）、原材料を供給します（メモリ帯域幅）。すべては工場をより効率的に稼働させるためです（コンピューティング）。

したがって、工場の生産能力が原材料の供給能力を上回る速度で拡大した場合、最高効率に到達するのは困難になるでしょう。

工場の生産能力 (FLOP) を 2 倍にしても、帯域幅が追いつかなければパフォーマンスは 2 倍にはなりません。

FLOPS についてもう 1 つ言及しておくべきことは、Nvidia の「Tensor Cores」など、行列乗算専用のハードウェア構成を備えた機械学習アクセラレータが増えていることです。

したがって、行列乗算を行わない場合、312 テラフロップスではなく、19.5 テラフロップスしか得られません。 GPU だけが特別なものではないことに注意してください。実際、TPU は GPU よりも特化したコンピューティング モジュールです。

行列乗算とは別に、GPU は行列乗算以外の演算の処理が遅く、一見問題があるように見えるかもしれませんが、レイヤー正規化や活性化関数などの他の演算子はどうでしょうか?実際、これらの演算子は、FLOPS での行列乗算の丸め誤差に似ています。たとえば、BERT のさまざまな演算子タイプが占有する FLOP の数については、次の表をご覧ください。ここで、「テンソル収縮」は行列乗算を指します。

ご覧のとおり、非行列乗算演算は全演算の 0.2% しか占めていないため、行列乗算の 1/15 の速度であっても、大きな問題にはなりません。

実際、正規化と点ごとの演算では、行列乗算の FLOPS の 1/250 と 1/700 しか使用されません。では、なぜ非行列乗算演算には必要以上に長い実行時間がかかるのでしょうか?

上記の「工場」の例えに戻ると、原因は多くの場合、原材料が工場との間でどのように輸送されるか、つまり「メモリ帯域幅」にあります。

帯域幅

帯域幅の消費は、本質的には、データをある場所から別の場所に移動するコストです。つまり、CPU から GPU へ、あるノードから別のノードへ、さらには CUDA グローバル メモリから CUDA 共有メモリへデータを移動することを意味します。最後のものがこの記事の焦点であり、一般的に「帯域幅消費」または「メモリ帯域幅消費」と呼ばれます。最初の 2 つは一般に「データ転送消費」または「ネットワーク消費」と呼ばれ、この記事の範囲外です。

「工場」の例えに戻りましょう。実際の作業は工場内で行いますが、大規模な保管には適していません。量で勝つのではなく、ストレージが十分に効率的で、すぐに使用できる (SRAM) ことを確認する必要があります。

では、実際の結果と「原材料」はどこに保存するのでしょうか?通常、土地が十分に安く、十分なスペース (DRAM) がある倉庫を所有しています。その後、工場との間で物品を配送できるようになります (メモリ帯域幅)。

コンピューティング ユニット間でデータを移動する際にかかるコストは、「メモリ帯域幅」コストと呼ばれます。実は、nvidia-smi コマンドで表示される「メモリ」とは DRAM のことで、よく人を困らせる「CUDA メモリ不足」というのはこの DRAM のことを指しています。

GPU コア操作を実行するたびに、データを倉庫 (DRAM) に転送して戻す必要があることに注意してください。

ここで、単項演算 (torch.cos など) を実行するときに、データを倉庫 (DRAM) から工場 (SRAM) に送信し、工場で小さな計算ステップを実行してから、結果を倉庫に送り返す必要があると想像してください。転送には非常に時間がかかり、この場合、実際の計算を行うのではなく、データの転送にほぼすべての時間を費やします。

すべての時間をメモリ帯域幅で費やしているため、この操作はメモリバインド操作とも呼ばれ、計算に多くの時間を費やしていないことを意味します。

明らかに、これは私たちが望んでいることではありません。それで、私たちに何ができるでしょうか?演算子のシーケンスがどのようになるかを見てみましょう。

点単位の演算子のシーケンスは次のようになります。

グローバル メモリとコンピューティング ユニット間でデータを転送することは明らかに最適ではありません。より良いアプローチは、データ ファクトリ内のすべての操作を一度に実行し、データを送信し返すことです。

これは演算子の融合であり、ディープラーニング コンパイラーで最も重要な最適化です。簡単に言えば、このアプローチでは、データをグローバル メモリに書き込んで再度読み取るのではなく、一度に複数の計算を実行することで余分なメモリ アクセスを回避します。

たとえば、x.cos().cos() 操作を実行するには、メモリへの書き込みに 4 回のグローバル読み取りと書き込みが必要です。

 x1 = x . cos () # グローバルメモリのx から読み取り、 x1 に書き込みます
x2 = x1 . cos () # グローバルメモリのx1 から読み取り、 x2 に書き込みます

演算子の融合では、グローバル メモリの読み取りと書き込みが 2 回のみ必要なので、2 倍の高速化が実現します。

 x2 = x.cos (). cos ( ) # グローバルメモリのx から読み取り、 x2 に書き込みます

しかし、このアプローチは簡単ではなく、特定の条件が必要です。まず、GPU は現在の操作を実行した後に次に何が起こるかを知る必要があるため、この最適化は PyTorch の Eager モード (一度に 1 つの演算子を実行する) では実行できません。次に、CUDA コードを記述する必要がありますが、これも簡単な作業ではありません。

すべての演算子の融合がポイントごとの演算子のように単純であるとは限りません。点単位の演算子を縮約または行列乗算に融合できます。行列の乗算自体も、ブロードキャスト乗算と減算の融合と考えることができます。

任意の 2 つの PyTorch 演算子を融合して、グローバル メモリの読み取り/書き込みのメモリ帯域幅コストを節約できます。さらに、多くの既存のコンパイラは、多くの場合「単純な」融合を実行できます (例: NVFuser と XLA)。ただし、より複雑な融合は依然として人が手作業で記述する必要があるため、カスタム CUDA カーネルを自分で記述したい場合は、Triton が適切な出発点となります。

驚くべきことに、融合された x.cos().cos() 操作には、x.cos() への単一の呼び出しとほぼ同じ時間がかかります。そのため、gelu では relu よりも明らかに多くの操作が行われるにもかかわらず、活性化関数のコストはほぼ同じになります。

したがって、チェックポイントを再実装/アクティブ化すると、興味深い結果がいくつか生じます。基本的に、追加の再計算を行うと、メモリ帯域幅が減少し、実行時間が短縮される可能性があります。したがって、AOTAutograd で簡潔な最小カット最適化パスを再実装して構築することで、メモリ使用量と実行時間を削減できます。

推論メモリ帯域幅コスト

単純な操作の場合、メモリ帯域幅を直接推論することが可能です。たとえば、A100 は 1.5 TB/秒のグローバル メモリ帯域幅を持ち、19.5 テラフロップス/秒の計算を実行できます。したがって、32 ビットの浮動小数点数 (4 バイト) を使用すると、GPU が 20 兆回の演算を実行するのと同じ時間で 4000 億個の数値をロードできます。

さらに、単純な単項演算 (テンソルを 2 倍にするなど) を実行するには、実際にはテンソルをグローバル メモリに書き戻す必要があります。

したがって、約 100 個の単項演算が実行されるまで、メモリを実際に計算するよりもメモリにアクセスするのにかかる時間の方が長くなります。

次の PyTorch 関数を実行すると:

 定義f ( x : テンソル[ N ] ):
 _が範囲内にある場合( 繰り返し):
 x = x * 2
 x を返す

フュージョンコンパイラでベンチマークすることで、繰り返し値ごとに FLOPS とメモリ帯域幅を計算できます。繰り返し値を増やすことは、メモリ アクセスを増やすことなく計算量を増やす簡単な方法です。これは、計算強度の増加とも呼ばれます。

具体的には、このコードをベンチマークし、まず 1 秒あたりに実行される反復回数を調べ、次に 2N (N はテンソルのサイズ) のメモリ アクセスと N*repeat FLOP を実行します。したがって、メモリ帯域幅は bytes_per_elem * 2 * N /itrs_per_second となり、FLOPS は N * repeat /itrs_per_second となります。

ここで、計算強度の 3 つの関数 (実行時間、フロップス、メモリ帯域幅) をプロットしてみましょう。

64 回の乗算が実行されるまで、実行時間はまったく大幅に増加しないことに注意してください。つまり、この時点までは、主にメモリ帯域幅によって制限され、計算はほとんどアイドル状態でした。

初期の FLOPS 値は 0.2 テラフロップスでした。計算強度を 2 倍にすると、この数値は直線的に増加し、ピークの 9.75 テラフロップスに近づきます。この時点で「計算能力に限界がある」と見なされます。

最後に、メモリ帯域幅はピーク付近から始まり、コンピューティングの強度が増すにつれて低下し始めることがわかります。これはまさに予想通りのことで、メモリへのアクセスではなく、実際の計算の実行に多くの時間が費やされていることを示しています。

この場合、コンピューティング バウンドとメモリ バウンドがいつ発生するかを簡単に確認できます。 repeat< 32 の場合、メモリ帯域幅が飽和に近くなり、十分な計算が実行されません。repeat> 64 の場合、計算が飽和に近くなり (つまり、ピーク FLOPS に近くなり)、メモリ帯域幅が減少し始めます。

大規模なシステムの場合、通常は両方の組み合わせであるため、コンピューティング制約かメモリ帯域幅制約かを判断するのが難しいことがよくあります。計算能力の限界を測定する一般的な方法は、実際の FLOPS とピーク FLOPS のパーセンテージを計算することです。

ただし、メモリ帯域幅のコストとは別に、GPU がスムーズに動作しない原因となるもう 1 つの要因があります。

追加費用

オーバーヘッドは、コードがテンソルの転送や計算以外の処理に時間を費やすときに発生します。たとえば、Python インタープリターで費やされる時間、PyTorch フレームワークで費やされる時間、CUDA カーネルの起動 (実行ではない) に費やされる時間はすべて間接的なオーバーヘッドです。

追加のオーバーヘッドが重要な理由は、最新の GPU が非常に高速であるためです。 A100 は、1 秒あたり 312 兆回の浮動小数点演算 (312TeraFLOPS) を実行できます。それに比べると、Python はひどく遅く、1 秒あたり約 3,200 万回の加算を実行します。

つまり、Python が 1 FLOP を実行するのにかかる時間で、A100 は 975 万 FLOP を実行できたことになります。

さらに悪いことに、Python インタープリターは間接参照の唯一のソースではなく、PyTorch などのフレームワークにも、実際のカーネルに到達する前に多くの層のスケジューリングがあります。 PyTorch は 1 秒あたり約 280,000 回の演算を実行できます。小さなテンソルを扱う場合（科学計算など）、PyTorch は C++ に比べて非常に遅いと感じるかもしれません。

たとえば、下の図では、PyTorch を使用して単一の加算を実行していますが、グラフのごく一部のみが実際に計算を実行しており、残りは純粋なオーバーヘッドです。

これを踏まえると、PyTorch が主流のフレームワークになったという事実に戸惑うかもしれません。これは、現代のディープラーニング モデルが通常、大規模な操作を実行するためです。さらに、PyTorch などのフレームワークは非同期で実行されます。したがって、フレームワークのオーバーヘッドのほとんどは完全に無視できます。

GPU 演算子が十分に大きい場合、CPU は GPU より先に実行できます (したがって、CPU オーバーヘッドはわずかです)。一方、GPU 演算子が小さすぎる場合、GPU はほとんどの時間を文鎮として費やすことになります。

では、この問題を抱えているかどうかはどうやってわかるのでしょうか?通常、オーバーヘッドは問題のサイズに応じて拡大しません (計算とメモリは拡大します)。そのため、最も簡単な方法は、データのサイズを増やすことです。実行時間が比例して増加しない場合は、オーバーヘッドの制限に達していると言えます。たとえば、バッチ サイズを 2 倍にしても実行時間が 10% しか増加しない場合は、オーバーヘッドによって制限される可能性があります。

別のアプローチは、PyTorch プロファイラーを使用することです。下の図に示すように、ピンク色のブロックは CPU コアと GPU コアの一致を示しています。

CPU は GPU よりもはるかに高速に動作します。

一方、nvidia-smi の「GPU-Util」(「Volatile GPU-Util」ではありません) エントリは、実際に実行されている GPU コアの割合を測定するため、オーバーヘッド制限に達しているかどうかを確認するもう 1 つの良い方法です。このオーバーヘッドは、PyTorch のようなすべての柔軟なフレームワークの特徴であり、本質的に「何をすべきかを判断する」のに多くの時間を必要とします。

これは、Python (属性の検索または適切な関数へのディスパッチ) または PyTorch のコードからのものになります。たとえば、a + b を実行すると、次の手順が実行されます。

1. Python は、__add__ が a にディスパッチする対象を検索する必要があります。
2. PyTorch は、どのカーネルを呼び出すかを決定するために、テンソルの多くのプロパティ (dtype、デバイス、autograd が必要かどうかなど) を決定する必要があります。
3. PyTorch は実際にカーネルを起動する必要があります。

基本的に、このオーバーヘッドは、各ステップで異なる操作を実行できる柔軟性から生じます。この柔軟性が必要ない場合、これを回避する 1 つの方法は、たとえば jit.trace、FX、または jax.jit を使用してトレースすることです。あるいは、代わりに CUDA Graphs などを使用して、この操作をより低いレベルで実行することもできます。

残念ながら、これは柔軟性を犠牲にすることになります。両方の長所を最大限に活用する 1 つの方法は、VM レベルでイントロスペクションを行うことで、「実際の」 JIT に沿ったものを記述することです。詳細については、TorchDynamo を参照してください。

要約する

ディープラーニング システムを高速化したい場合、モデル内のボトルネックが何であるかを理解することが最も重要です。ボトルネックによって、システムを高速化するのに適した方法が決まるからです。

PyTorch コードの高速化に関心のある研究者やその他の人々が、取り組んでいる問題を理解せずに盲目的に試みているのを頻繁に目にします。

もちろん、一方で、ユーザーがこれらのことを考慮する必要がある場合、それはフレームワークの部分的な失敗を反映しているとも言えます。 PyTorch は活発な関心領域ですが、PyTorch のコンパイラやプロファイル API は使いやすいとは言えません。

全体として、システムの基礎を理解することはほとんどの場合役に立つと私は思っており、それが皆さんにとっても役に立つことを願っています。