TT Lab
ブログ

한국어 English 日本語 中文

はじめる
ブログ

ブログ

グラフコンパイラとランタイム:演算融合、torch.compile、ONNX Runtime・TensorRT

ステータス:ドラフト(2026-10-16)、一部実測。 torch.compile・メモリレイアウト・cuDNNの設定は、ジブラルタルの8GBクラスのノートPC用GPU1枚で使い捨てPodを使って測定しました(PyTorch 2.9、ウォームアップ20回の後、50回×5セットの中央値)。融合の原理はノートPCのCPUの例で確認しました(code/03_fusion.py、数秒)。ONNX RuntimeとTensorRTは実験イメージに含まれていないため 未実測です(ダウンロードしていません)。 シリーズ「ビジョンモデルを読む3」第3回。精度削減については 第2回にあります。

目次

1. 直観:料理人と冷蔵庫

GPUの計算装置は高速ですが、データはチップ外のメモリ(冷蔵庫)にあります。eager 実行は、「畳み込み → バッチ正規化 → ReLU → 加算」をそれぞれ別々に行います。畳み込みの結果を冷蔵庫に入れ、バッチ正規化が取り出して修正してまた入れ、ReLUがさらに取り出して入れます。大きな特徴マップであれば、計算よりも出し入れの時間の方が長くなります。 演算融合とは、「取り出したついでに4つすべてを行い、1度だけ書き戻す」ということです。

연산 융합과 컴파일러·런타임 융합 전: 연산마다 커널 하나, 메모리를 오간다 합성곱 BN ReLU 더하기 GPU 메모리 — 읽고 쓰기 4번, 커널 부르기 4번 융합 뒤: 한 커널이 값을 칩 안에 둔 채 끝까지 합성곱 + BN(가중치에 접음) + ReLU + 더하기 GPU 메모리 — 읽고 쓰기 1번, 커널 부르기 1번 같은 모델이 GPU 커널이 되기까지 두 길 PyTorch 모델 torch.compile 그래프 잡기 Inductor Triton 커널 생성 ONNX 파일 표준 그래프 ONNX Runtime TensorRT 엔진
図:融合前は演算ごとに1つのカーネルがメモリを行き来し、融合後は1つのカーネルが値をチップ内に置いたまま最後まで処理する。下図はPyTorchモデルが torch.compile またはONNXを経由してGPUカーネルになるまでの2つの経路

融合には2つの種類があります。

BN 접기:  y = γ·(W∗x + b − μ)/√(σ²+ε) + β
        = (γ/√(σ²+ε))·W ∗ x + (γ·(b − μ)/√(σ²+ε) + β)
        = W' ∗ x + b'           ← 합성곱 하나로

2. 実際にやってみる(ノートPCのCPU)

code/03_fusion.py をノートPCのCPUスレッド1つで実行しました(7回の中央値)。

実験 結果
畳み込み(32→64、3×3、入力144×256)+BNの畳み込み 最大差異4.77e-06(FP32の丸め誤差レベル)。時間5.80 → 5.80ms(差異なし)
要素ごとのチェーン relu(x·1.5 + 0.5)·0.9、要素数1,677万個 eager実行9.7ms → torch.compile 2.4ms(4.1倍)、結果の差異0
同じチェーンを手動でインプレース演算 7.4ms
要素数102万個の加算を1,000回に分けて呼び出す 0.13 → 0.88ms、呼び出しごとに約0.8µs

3. グラフコンパイラは何をするのか

TVM(Chenほか、2018) はモデルをグラフとして受け取り、演算を4つに分類します。要素ごとの処理(injective)、まとめ上げる処理(reduction)、出力に要素ごとの処理を付加できる複雑な処理(畳み込みのような complex-out-fusable)、そしてまとめられない処理です。ルールに従い、最初の3つを統合します。演算融合だけで1.2〜2倍高速化し(2015年デスクトップ級GPU、畳み込み+BN+ReLUなどの4つのタスク、3節・図4)、他のフレームワークと比較してエンドツーエンドで1.6〜3.8倍でした(6.1節)。

PyTorch 2(Anselほか、2024) の torch.compile は2つの部分から構成されます。

データセンターGPUにおいて、TIMM画像モデル62個の推論がeager実行よりもFP32で2.48倍、FP16で2.77倍高速になりました(表3)。融合とインライン化の両方を無効にすると、むしろ0.80倍となり、eager実行よりも遅くなりました(表4、HuggingFace FP16推論)。 速度の大部分は融合に由来することを意味します。冒頭の「推論2.27倍」がどの組み合わせの平均であるか本文で確認できなかったため、ここでは表3・4の値のみを使用します。

4. ONNXとランタイム:モデルを持ち運ぶフォーマット

torch.compile はPyTorch内部で動作します。サーバーにPythonとPyTorchを置きたくない場合や、ブラウザ・スマートフォンで実行する必要がある場合は、モデルを 標準グラフファイルとしてエクスポートする必要があります。

5. 自社での実測:ボールモデルへの torch.compile

TT Scoreのボールモデル(BlurBall、3フレームウィンドウ、512×288、バッチ1)を実験用GPUで測定しました。

実行方式 FP32 FP16 初回呼び出し(コンパイル)
eager実行(現在の運用方式) 15.09ms 9.77ms —
eager実行、cuDNN自動チューニング無効 15.08ms 9.75ms —
eager実行、channels_last(NHWC) 18.62ms 9.02ms —
torch.compile デフォルト 12.91ms(1.17倍) 4.61ms(2.12倍) 10〜11秒
torch.compile reduce-overhead(CUDAグラフ) 12.55ms 4.34ms(2.25倍) 7〜8秒
torch.compile max-autotune(CUDAグラフなし) 12.79ms 4.48ms 41〜46秒

(FP16はautocastを使用。倍率は同じ精度のeager実行と比較したもの。)

同じ実行をプロファイラーで確認したところ、eager実行のFP32の1回の実行で CUDAカーネル478個が呼び出され、GPUの稼働時間14.88msのうち畳み込みが8.09ms(72回)、バッチ正規化が3.27ms(71回、22%)でした。カーネル時間の合計がウォールクロック時間(15.09ms)とほぼ同じであるため、GPUは休止していませんでした。つまり、呼び出しコストではなく メモリを行き来する要素ごとの演算が問題であり、コンパイラがこれをまとめて排除したと解釈しています(コンパイルされたバージョンのカーネル数はカウントしていません)。

読み解き方のポイントです。

6. 制限事項

7. 自プロジェクトへの適用

T.T.の一言。 同じ料理を冷蔵庫の開閉を減らして作ったら、2倍速くなりました。材料(重み)もレシピ(モデル)もそのままなのに!

参考文献

ログインするといいねできます

コメント

まだコメントはありません。

ログインするとコメントできます