Stable Diffusionの学習速度をいろいろ比較する

gcem156

　タイトルの通り比較します。

学習設定

  • 使うのは自前の訓練コード

  • 学習用データセットは768×768画像128枚

  • Latentは事前に計算しておく（計測対象にはならない）

  • UNetのみを学習対象とする

  • torch2.0を使う

  • diffusersは0.16.1を使う

  • AMPを使用し、型はbfloat16にする

  • AdamW8bitを使う

  • バッチサイズはなるべく大きい値にする

  • WD1-5を対象にする（SD1.x系はもっと遅くなります）

  • nvidia-smiを適当にうちこんでVRAM使用量を記録します

  • プログレスバーの時間をそのまま記録します

  • ColabはCPUハイメモリを選ぶ

  • xformersの有無やgradient_checkpointingの有無で3パターン程度に分けて計測する

結果

　samples/秒は1秒で何枚学習できるかという指標です。samples/costはsamples/秒からPaperspaceの場合1時間辺りの料金（ドル）、Colabの場合は1時間当たりのUnit数で割ったものです。Paperspaceは普通に借りる分にはちょっと高いので参考程度にしてください。

　A100最強！コスパでもA100が最強のようですね。バッチサイズは学習結果にも影響するため、厳密には学習時間のみで比較することはできませんけどね。

考察

torch1.13 vs torch2.0

　最初ローカルPCではtorch==1.13+xformers==0.0.15のままやるつもりだったんですが、RTX3090と同格だと思っていたA5000に大敗していたので、torch==2.0+xformers==0.0.19にアップグレードしました。計算速度もVRAM使用量もかなり変わりました。やっぱり3090とA5000は同じくらいですね。

xformers vs sdpa

　sdpaはtorch2.0から実装された関数のことです。diffusers==0.16.1ではxformersを使わない場合、自動的にこれが使われるはずです。
　今回の結果ではsdpaはxformersよりわずかにVRAM使用量が減少するものの、学習時間は遅くなるという感じみたいですね。GPU（または何らかの設定？）によってはかなり遅くなってしまいます。

gradient_checkpointingの有無

　V100のみがgradient_checkpointingを使った方が早くなるという結果になりました。
　gradient_checkpointingは逆伝搬に必要な順伝搬の履歴を全て保存せず、一部を残して必要になったら再計算するという手法です。再計算する分基本的に速度は落ちますが、VRAM使用量を減らすことができます。V100の場合はバッチサイズを上げることによる計算効率の上昇分が再計算分を上回るためにこういった結果になるのだと思います。これを使用せずに低バッチサイズで学習すると、T4＞V100になるという結果は驚きでした。
　ちなみに前も似たような計測をしましたが、当時の環境ではVRAM24GBでバッチサイズ4を確保できず、A5000やRTX3090でもgradient_checkpointingを使用した方がよいという結果になっていました。このあたりの認識を変えられたのが収穫です。

反省点

　データセットの数を適当に128枚としましたが、もっといろいろなバッチサイズで割り切れる数にすべきでした。余りは指定したバッチサイズ未満で計算されてしまいます。

まとめ

  • torch==2.0、xformers==0.0.19が強い

  • V100（16GB）はgradient_checkpointingを使った方がよい

  • A100がコスパも最強そう（学習時間のみを考えれば）