Fixed triton3.2.0 ＆ WaveSpeed on Inductor backend finally (but not better nodes)

Loaded triton 3.2.0 on ComfyUI

先月、散々四苦八苦した挙句に、何とか妥協的なcudagraphs backendで完走に漕ぎつけたWaveSpeedのMedel Compile+ノードですが、勿論機能的には公式推奨のInductor backendの方が性能は良い…事はわかっていたのです。

しかし、ここまでどうにもならなかった訳で、Inductorで完走する事は出来ても、コマンドラインをよく見るとエラーを起こして実際にはtorch.compileが働いていなかった訳です。

それが、先だって公開されたWindows用triton3.2.0を使用すれば、3.1.0以前では不可能だったInductor backendが今度こそ正常動作する…事はわかっていました。

しかし、この3.2.0を使用すると別な問題が発生し、ComfyUIにおけるいくつかの重要なCustom Nodesがエラーを起こす為、torch.compileの為だけにそれらを犠牲にする訳には行かず、今日まで妥協的なcudagraphsでの使用を強いられてきました。

その間、別な角度から、MultiGPUノードによる高速化の道も開けていたので、今は寧ろ、torch.compileを削除してMultiGPUノードを組み合わせる方向での設定を進めていた処なのですが…

今日、遂にtriton3.2.0をComfyUI上でノーエラーでロードすることに成功し、下図のようにModel Compile+ノードを公式推奨設定のInductor backendで動作させることに成功しました。 

Tips1

突破口は、別に今日初めて見た訳ではない、以下記事です。

結論から端的に言えば、記事内にあるinclude_libs.zipファイルを展開し、python_embededフォルダに設置する事から突破口が開けました。

実はこれ、3.1.0の当時に既にテストしていたのですが、3.1.0の場合これを設置してもInductorは使えなかったのです。
(ノードは完走しても、コマンドラインをよく見ると、エラーを起こして最適化がされていない)

で、cudagraphsで使用する場合、これはあってもなくても変わらないので、記憶から消していたのです。

しかし今日、triton3.2.0ロード時に発生するエラーメッセージをCursor先生と解析していく中で、以下のようなメッセージが出てきた為に、「あれ ? 何かincludeてワードはtriton絡みのどこかで見た記憶あんなあ」と思い出して、解決につながりました。

triton.hファイルて何処にあるかわかるか ?

triton.hファイルは通常、Tritonのインストールディレクトリの中の以下のようなパスに存在するはずです：

D:/USERFILES/ComfyUI/python_embeded/Lib/site-packages/triton/backends/nvidia/include/triton/

または：

D:/USERFILES/ComfyUI/python_embeded/Lib/site-packages/triton/include/

確認するには、以下のPythonコードを実行して探すことができます：

そう、上の西川氏の記事内で紹介されている、whlとは別に、追加で設置するフォルダの事です。上のCursor先生の返答は、パス自体は間違っていますが、「include」というフォルダ名を見て閃きました、つか連想が全て繋がりました。

「ああ成程、あの時の西川氏の記事にあった「あれ」は、ここで必要になる訳か」と。

但し、上記事内のリンクは、Python3.11.x用なので、3.12.xで動作するComfyUIには適用できません。

Python3.12.x用は、以下から入手しました。

上のtorch.compile設定記事も、WaveSpeed関連記事と公開している各jsonと併せて、この後全て更新していきます。いやはや、一度公開したWaveSpeed搭載型は全部変更ですね。
(更に、この「更新」の内容に関しては、重要な続きがあるので、本記事を最後まで読んでください)

Speed tests

速度に関しては、MultiGPUノード併用型と比較しつつテストしましたが…２回目以降は相当速いですよ、特にFlux1で。

下記事でも解説している、WaveSpeedからModel compile+を削除し、Dual Clipに対してMultiGPUによるDistorchをかました形の２回目以降です。

下図が、Inductor backendを使用した、今度こそフルスペックのWaveSpeed２回目以降です。複数回テストしましたが、やはりWaveSpeedに分があります。特に２回目以降…これは、torch.compileの特性上、そうなる筈です。

Advanced tips

一方、Detailerノードに対してはWaveSpeedのModel Compile+ノードではInductorは仕様として適用できないようで、今まで通りcudagraphs backendとreduce-overhead modeを使用する事が一つ対策となります。

もう一つの対策は、KJNodesによるtorch.compileを使用する事です。どういう訳か、こちらのノードではInductor backendがdetailer周りに対しても、きちんと動作します。

Fixing torch.compile on KJNodes

そもそも、tritonに拘る理由は一つ…ComfyUIにおける有力な高速化手段であるWaveSpeed…というよりも、その中の一つであるtorch.compileを、その最大性能を発揮するInductor baclkendで動作させる為ですが、

前述した通り、torch.compileという技術そのものは別にWavSpeedの専売特許という訳ではありませんで、KJNodesも同様のノードを公開しています。

About KJNodes

で、ここからが重要ですが、近日来、WaveSpeedとKJNodes双方のtorch.compileノードをテストしつつ思うのは、少なくとも現状においてはKJNodesの方が優れている、という点です。

まず最大の理由は、KJNodesのtorch.compileでは下図のように、現状torch.compileにおける最大性能を発揮するInductor backendでのmax-autotune機能が正常に動作する点です。

これが現状、WaveSpeedのCompile Model+ノードでは、前述した通りInductor backendは何とか正常動作するようになりましたが、モードの設定はまだできません。ここで、max-autotune等を指定すると、エラーを起こして完走しません。

max-autotuneは、初回ロード時にcompileの為の処理時間は要するものの、推論の時間…つまりサンプラー処理の時間が確実に短縮されます。

この特性は、例えば私のようにRTX4070 12GBクラスのGPUでFlux1での生成を行い、専用VRAM12GBを超えてしまう処理の際に威力を発揮します。

先日、基幹部分をDDR4からDDR5にアップグレードしている為に、その恩恵もありますが、Compileに時間をかけてもその後の推論時間を短縮できる場合、専用VRAM12GBを超えてしまうFlux1のような処理の場合、このmax-autotuneは効果的という事になります。

実際、テストを重ねてきた範囲では、KJNodesのノードでInductor backendをmax-autotuneで動作させる方が、WaveSpeedのCompile Model+よりも確実に速いことがわかっています。

これを踏まえて、今回triton3.2.0のノーエラーロードに成功した事で、過去公開してきたWaveSpeed対応のワークフローを全て更新する予定でしたが、やや予定を変更し、torch.compile部分に関してはKJNodesを採用することにします。

Advanced tips

但し、例外的にWaveSpeed版のCompile Model+ノードを使用する場合もあります。原因は不明ですが、一部jsonの構成によってはKJNodesではtorch.compileが効かない場合があります。

KJNodes版でエラーを起こす場合、WaveSpeed版を試してください。