xfromers vs sage-attention ?…No, it's a problem of torch.compile

xfromers vs sage-attention…という問題ではないのですよ、本記事の趣旨はね。最初、そのタイトルを考えましたが、事の本質はそれじゃないのでね、明らかに。

要するに、torch.compileの課題と、その対策という話です。

さて、まず下図ですが、xformersのみで、Flux1の初回生成を行った時のものです。毎回、大体80秒前後…これは以前200秒前後かかっていた訳で、随分速くはなりました。

ベンチマークに使うFlux1のjsonは、例によって以下記事で公開しているMulti ControlNet付のベースです。

さて、以下がsage-attention付きで動作させた同じFlux1です。…ん? 何かがおかしいですね、180秒近くかかってますよ。場合によっては、250秒を超える時もあります。

これを以て、Flux1においては、xformersの方がsage-attentionより速い…訳ではない、という点が重要です。

例えば以下、「ある条件下」で同じくFlux1をsage-attentionで動作させたものです。80秒未満で完走し、十分速いです。

…ん?　て事は?　この差は何なの ?...て話になりますよね。

その原因を突き止めました。

結論から端的に言います。

torch.compileによって作成されるキャッシュファイルが別のモデルによる生成によって置き換わってしまうと、せっかく構築したtorch.compileの最適化設定が全てリセットされ、別モデル用の最適化データによって上書きされてしまう....のが原因です。

上の遅い例の場合、直前にWAN2.1でtorch.compileを動作させています。その後に、Flux1の生成を行うとこうなってしまう訳です。

xformersだから…sage-attentionだから…と言う問題ではないのです。

3枚目の「高速な生成が出来たsage-attention」というのは、実はその遅い初回生成を経た上で、改めてtorch.compileにFlux1のパターンを覚えさせ(逆に言えばWAN2.1の最適化データが失われる)、更に繰り返してFllux1の生成を行った場合の結果なのです。

…という事は、解決策が見えてきましたね。

先日、環境変数によって、torch.compileやtrtionによって発生するキャッシュファイルや一時ファイルを分離して任意のフォルダに作成する…という手順について解説しました。

これらのファイルが、別モデルのロードによって置き換わってしまう事により、Flux1の初回生成が遅くなる(勿論、逆の場合ではWAN2.1が遅くなる)のであれば、その置換が発生しないようにプログラムの流れを組んでやればいい訳です。

まず、Flux1用の起動バッチファイルが以下です。オプションで何も指定せず、かつxformersがインストールされている状態では、ComfyUIは自動でxformersを適用します。

set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:32,garbage_collection_threshold:0.8
set TORCH_COMPILE_DEBUG=0
set TORCH_COMPILE_OPT_LEVEL=2
set TORCH_LOGS=recompiles,dynamo,inductor
set TORCH_INDUCTOR_VERBOSE=1
set TORCH_DYNAMO_CONFIG=force_parameter_static_shapes=false
set TRITON_CACHE_DIR=D:\USERFILES\xformers\triton
set TORCH_HOME=D:\USERFILES\xformers\torch
set TORCHINDUCTOR_CACHE_DIR=D:\USERFILES\xformers\Inductor
set TEMP=D:\USERFILES\xformers\Temp
set TEMP=D:\USERFILES\xformers\Temp
.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --output-directory C:\Users\ussoe\A1111\stable-diffusion-webui\outputs\ComfyUI

続いて以下がsage-attentionを適用し、WAN2.1を想定した場合のバッチファイルです。

set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:32,garbage_collection_threshold:0.8
set TORCH_COMPILE_DEBUG=0
set TORCH_COMPILE_OPT_LEVEL=2
set TORCH_LOGS=recompiles,dynamo,inductor
set TORCH_INDUCTOR_VERBOSE=1
set TORCH_DYNAMO_CONFIG=force_parameter_static_shapes=false
set TRITON_CACHE_DIR=D:\USERFILES\sageattention\triton
set TORCH_HOME=D:\USERFILES\sageattention\torch
set TORCHINDUCTOR_CACHE_DIR=D:\USERFILES\sageattention\Inductor
set TEMP=D:\USERFILES\sageattention\Temp
set TEMP=D:\USERFILES\sageattention\Temp
.\python_embeded\python.exe -s ComfyUI\main.py --use-sage-attention --windows-standalone-build --output-directory C:\Users\ussoe\A1111\stable-diffusion-webui\outputs\ComfyUI

何が違うか…そう、キャッシュファイルの場所です。

つまり、Flux1とWAN2.1(フォルダ名はxformersとsage-attentionで分けた)で、torch.compileによって発生する最適化データのファイルが格納されるフォルダを分離した訳です。

起動時のバッチファイルを使い分ける事で、一度作成したtorch.compileのモデルごとの最適化データが上書きされないようになります。

この結果、WAN2.1で動画を作成した後でも、ComfyUIを一度終了させてバッチファイルを変えて再起動させる事で、別な場所に保存されているtorch.compileにおけるFlux1の最適化データがロードされ、初回生成速度の低下が発生しないようになりました。
(逆の視点では、当然WAN2.1の最適化データも失われない)

尤もComfyUIの再ロードの時間はかかる訳ですが、2分はかかりませんからね…差し引きしても、＋ではあります。
(特にWAN2.1の場合、torch.compileによる初回高速化が2分処ではなく更に長いので、差し引きは絶対的に＋)

ただ、そこまでするかって視点も勿論ありますよ。初回が遅くても2回目以降は、torch.compileは最適化してきますから。

只これも前回記事で書きましたが、「そういう仕組みになっている」という理屈を知るだけでも意味はあると思っています。

何でもそうですが、バックグラウンドを理解せずに使うのと、理解して使うのでは、何かしらトラブルが発生した時の対応力が違ってきますから。