用語集 · コンピューティング
CUDAグラフ
別名:CUDA graph、CUDA graphs、graph capture、graphed、torch.cuda.graph、CUDA Graphs、グラフキャプチャ
一連の GPUカーネルを記録したもので、ひとまとまりとして再生できるため、CPU が各カーネルを毎回起動し直す必要がなくなります。
通常、CPU はGPUに何をするかをカーネル1つずつ指示します。学習ループのように、小さなカーネルの長い列が何度も繰り返し実行される場合、これは無駄な手間です。CPU は毎回同じ指示を繰り返し、GPU はその1つ1つを待つことになります。
CUDAグラフ(CUDA は NVIDIA の GPU プログラミング用プラットフォームです)は、その列を一度だけ、すべてのカーネルとそれが読み書きするメモリとともに記録します。その後は、1回の呼び出しで全体を再生できます。GPU はこれから来るものをすべて知っているので、間を空けずに次々と実行します。
ただし落とし穴があります。再生では毎回同じメモリアドレスが使われます。そのため、再生のたびにデータを同じ固定のメモリブロックへコピーしておく必要があり、中のコードが GPU の結果に応じて CPU 側で判断を下すこともできません。
このプロジェクトでは
レンタルGPUで学習する では、トレーナーの2つの部分を CUDAグラフとして記録しました。ミニバッチ1つに対するPPOの更新1回(ネットワークを実行し、勾配を計算し、クリップし、重みを更新する)は、デスクトップの RTX 3090 で1イテレーションあたり 0.6–0.7 秒から 0.2 秒になりました。より遅いプロセッサのサーバーでは 0.7–1.4 秒かかっていました。MuJoCo Warp の物理計算を含む環境ステップ全体もグラフになっており、こちらは約5%の節約になりました。
2人のファイターと視覚 では、ずっと大きなブロック用環境に対して記録した更新が1イテレーションあたり 0.08 秒で済みました。レンタルの RTX 5090 の多くが載っている EPYC サーバーでカーネルを1つずつ実行すると約 0.5 秒かかります。この機能は、記録した直後の最初のロールアウトがそれらのマシンでクラッシュしたため、オフにされていました。更新を環境自体のグラフの後ではなく前に記録したところ、問題なく動きました。原因はまだわかっていません。