見出し画像

flash-attention2.7.0+xformersの統合


 前回ソースからビルドした、pytorch2.5.1(cuda12.6.3+cudnn9.5.1)+flash-attention2.7.0.post2に、xformersを統合できないのかと、色々やってみて結論的には、現状xformersを統合できるけど生成速度、vram消費量的に統合させない方がよいことが分かりました。つまり、xformersはインストールしない方がよい。(あくまで現時点での話です)

 python -m xformers.info このコマンドで、仮想環境にインストールしてあるxformersの有効化している機能や統合したpytorch+cudaバージョンなどを確認することができます。その項目の中に下記の表記があります。

memory_efficient_attention.fa2F@v2.7.0 post2-pt: available memory_efficient_attention.fa2B@v2.7.0 post2-pt: available

 私の場合、上記の項目が最初に@v2.5.7と表示されていたのですが、pytorchをソースからビルドした際にデフォルトでflash-attention v2.5.7がビルドと統合されているためです。上記のように表示させて、仮想環境にインストールしてあるflash-attention2.7.0をxformersと統合させる方法をgithubのxformers 公式に相談したところ、metaのContributorから回答があり、現状flash-attention2.6.3までしか対応しておらず、スクリプトを編集すれば上記のように2.7.0pt2が出るが動作的に要テストが必要とのことでした。
 
 力技みたいな方法でスクリプトを編集して、flash-attention v2.7.0がxformersと統合させることができましたがやはり、"value error"が発生してComfyUIでは生成に失敗しました。
 
 一応、xformers/xformers/ops/fmha/flash.py の66行目 LAST 2.6.3 という部分を2.7.0に編集すればfa2Fと、fa2Bが v2.7.0post2に変わって、xformersを仮想環境に開発版でインストールする際にflash-attentionのモジュールが追加されて統合可能です。

 最終的に私の環境において、最も効率的に画像生成できるライブラリの組み合わせを調べてみて、pytorch(+デフォルトで組み込まれているflash-attention v2.5.7)のみ、pytorch+xformers+flash-attention2.6.3、pytorch+flash-attention2.7.0pt2という組み合わせで、生成速度やvram消費量、生成時のGPU温度などを計測しました。
 個人的な画像生成の範疇での意見となりますが、pytorch2.5.1(cuda12.6.3+cudnn9.5.1)+flash-attention2.7.0post2の組み合わせが最も速度(平均的に0.15秒ほど早くて)、vram消費量面(どちらもほぼ差がない)となり、少しでも効率的に画像生成できる組み合わせを選択しました。
 画像生成で同一のプロンプトで連続画像生成をしているうちにAIの学習が進むにつれて、pytorch+flash-attention2.7.0pt2の場合、vram消費量がxformersと同程度だったものが10枚ほど生成すると1.5GBほど減ったかな?という感じです。次は同条件でxformers統合させた組み合わせでも学習を進めればvram消費量が減るかも確かめてみます。
 


いいなと思ったら応援しよう!