298MBのRAMで動作するStable Diffusion XL 1.0

 (github.com/vitoplantamura)
4 ポイント 投稿者 GN⁺ 2023-10-04 | 1件のコメント | WhatsAppで共有
  • OnnxStreamは、512MB RAMのRaspberry Pi Zero 2でStable Diffusion 1.5とSDXL 1.0 Baseを、スワップ追加やディスクへの中間保存なしで実行することを目指す超小型推論ライブラリ
  • 推論エンジンと**WeightsProvider**を分離し、DiskNoCache、DiskPrefetch、カスタムのロード・キャッシュ・プリフェッチをサポート。HTTPから重みを直接受け取り、ディスク書き込みなしで処理することも可能
  • SDXL 1.0 BaseはHugging Face Diffusers 0.19.3実装からエクスポートしたONNXファイルを使用し、Raspberry Pi Zero 2では10ステップの画像生成に約11時間かかる
  • SDXLのVAEデコーダはFP32で4.4GB RAMを使い、FP16・UINT8だけでは解決が難しいため、重なり合うタイルデコードとブレンディングでメモリを298MBまで削減
  • CPU基準でOnnxRuntimeより最大55倍少ないメモリで動作する代わりに0.5〜2倍遅くなる場合があり、バッチサイズ1と一部のONNX演算子中心という制約がある

OnnxStreamの目標と設計

  • OnnxStreamは、Stable Diffusion 1.5をRaspberry Pi Zero 2で実行する問題を目標に始まった
    • Raspberry Pi Zero 2は512MB RAMを搭載したマイクロコンピュータ
    • Stable Diffusionの推奨最小RAM/VRAMは一般に8GB
    • 目標条件は、スワップ領域の追加なし、ディスクへの中間結果のオフロードなし
  • 一般的な機械学習フレームワークは推論レイテンシやスループットを優先するため、RAM使用量が大きくなりやすい
  • OnnxStreamはメモリ使用量の最小化に焦点を当てた、小さく修正しやすい推論ライブラリ
  • 中核構造は、推論エンジンとモデル重み提供コンポーネントである**WeightsProvider**の分離
    • 派生クラスがモデルパラメータのロード、キャッシュ、プリフェッチ方式を実装できる
    • カスタムWeightsProviderはHTTPサーバーからデータを直接ダウンロードし、ディスクの読み書きなしで処理可能
    • 基本実装はDiskNoCacheとDiskPrefetch
  • CPU基準でOnnxRuntimeより最大55倍少ないメモリを使いながら、0.5〜2倍遅い程度

Stable Diffusion 1.5の実行

  • Stable Diffusion 1.5の例では、OnnxStreamを使って異なるVAEデコーダ精度で画像を生成する
  • Raspberry Pi Zero 2のRAMに収まらなかったコンポーネントはVAEデコーダだけだった
    • 原因はモデル内部のresidual connection、大きなテンソル、大きなconvolution
    • 単精度や半精度だけではRaspberry Pi Zero 2のRAMに収まらなかった
    • 解決策は8ビット静的量子化だった
  • W8A8精度のVAEデコーダ画像はRaspberry Pi Zero 2で生成され、MAX_SPEEDオプション使用時で約1.5時間かかる
  • 比較用のW16A16画像は同じlatentを使ってPCで生成された

Stable Diffusion XL 1.0 Base対応

  • OnnxStreamのStable Diffusion例はSDXL 1.0 Baseをサポートし、Refinerは含まない
  • ONNXファイルはHugging Face DiffusersライブラリのSDXL 1.0実装からエクスポートしたもので、使用バージョンは0.19.3
  • SDXL 1.0はSD 1.5より計算コストがはるかに大きい
    • 最大の違いは512x512ではなく1024x1024画像を生成できる点
    • 12コアPCと32GB RAMでHugging Face Diffusersを使い10ステップ画像を生成するのに26分かかる
    • SDXLの推奨最小VRAMは一般に12GB
  • OnnxStreamはSDXL 1.0を300MB未満のRAMで実行できるため、Raspberry Pi Zero 2でも動作する
    • スワップ追加なし
    • 推論中のディスク書き込みなし
    • Raspberry Pi Zero 2での10ステップ画像生成は約11時間かかる

SDXL専用メモリ最適化

  • SDXL 1.0にはSD 1.5と同じ最適化セットを適用するが、一部に違いがある
  • UNETモデルはRaspberry Pi Zero 2で300MB未満のRAMで実行するため、UINT8動的量子化を使用する
    • 量子化対象は大きな中間テンソルの特定のサブセットに限定される
  • SDXL 1.0のVAEデコーダはSD 1.5より扱いが難しい
    • SDXL 1.0 VAEデコーダはSD 1.5 VAEデコーダより4倍大きい
    • OnnxStreamでFP32精度で実行すると4.4GB RAMを使用する
    • SD 1.5ではVAEデコーダをUINT8静的量子化し、RAM使用量を260MBまで減らせた
  • SDXL 1.0 VAEデコーダはFP16演算でoverflowが発生し、activationの数値範囲が大きいため、UINT8量子化では良質な画像を得にくい
  • sdxl-vae-fp16-fixのようなFP16の解決策もあるが、メモリを半分にしても2.2GBで、Raspberry Pi Zero 2にはなお大きい
  • 最終的な解決策は、Hugging Face DiffusersのVAEデコーダ実装から着想を得たタイルデコード
    • diffusion結果テンソルのshapeは(1,4,128,128)
    • これを(1,4,32,32) shapeの重なり合うテンソル5x5、合計25個に分割
    • 各タイルは左側および上側のタイルと25%重なる
    • 各タイルのデコード結果は(1,3,256,256)テンソルで、最終画像にブレンドされる
    • ブレンディングをオフにするとタイル境界が見え、オンにすると最終結果で境界が見えない
  • この方式により、SDXL VAEデコーダのRAM使用量は4.4GBから298MBに削減される

サポート機能と依存関係

  • OnnxStreamは、メモリ節約型推論に必要な機能を小さくまとめている
    • 推論エンジンとWeightsProviderの分離
    • DiskNoCache、DiskPrefetch、カスタムWeightsProvider
    • attention slicing
    • 8ビットunsigned asymmetric percentile動的量子化
    • W8A8 unsigned asymmetric percentile静的量子化
    • 量子化モデルのキャリブレーション
    • FP16サポート
    • よく使われる25個のONNX演算子の実装
  • 演算は逐次実行されるが、各operatorはマルチスレッドで動作する
  • 実装は単一の実装ファイルとヘッダーファイル構成で、XnnPackクラスでXNNPACK呼び出しをラップする
  • 一部の高速化primitiveはXNNPACKに依存する
    • MatMul
    • Convolution
    • element-wise Add/Sub/Mul/Div
    • Sigmoid
    • Softmax

性能比較と制約

  • Stable Diffusionは3つのモデルで構成される
    • text encoder: 672 operations、1億2,300万parameters
    • UNET: 2050 operations、8億5,400万parameters
    • VAE decoder: 276 operations、4,900万parameters
  • バッチサイズ1基準の10ステップ画像生成には、次の実行が必要
    • text encoder 2回
    • UNET 20回
    • VAE decoder 1回
  • FP16 UNET基準で、OnnxStreamとOnnxRuntimeのメモリ・時間差は大きい
    • OnnxStream: 約0.133GB、18.2〜19.8秒
    • OnnxRuntime: 5.085〜7.353GB、7.28〜12.8秒
    • OnnxStreamはメモリを最大55倍少なく使うが、0.5〜2倍遅い
  • FP32 text encoderはOnnxStreamが0.147GB、OnnxRuntimeが0.641GBを使用
  • FP32 VAE decoderはOnnxStreamが1.004GB、OnnxRuntimeが1.330〜2.026GBを使用
  • 比較結果にはいくつかの条件が付く
    • OnnxRuntimeの初回実行はInferenceSession再利用前のwarm-up推論
    • OnnxStreamはeager設計のためwarm-upの概念はないが、以降の実行ではOSのweightsファイルキャッシュの利点を受けられる
    • OnnxStreamは現在、batch size 1ではない入力をサポートしていない
    • OnnxRuntimeはUNET実行でbatch size 2を使用し、diffusion全体を大きく高速化できる
    • OnnxRuntimeのSessionOptionsのうちEnableCpuMemArena、ExecutionModeの変更は、テストで有意な差を生まなかった
    • NCNNはメモリ使用量と推論時間の面でOnnxRuntimeと非常によく似ていた
    • テスト環境はWindows Server 2019、16GB RAM、8750H CPU AVX2、970 EVO Plus SSD、VMWare 8 virtual cores

Attention slicingと量子化

  • UNET実行時のattention slicingとVAEデコーダのW8A8量子化は、Raspberry Pi Zero 2で実行可能な水準までメモリを下げるうえで重要だった
  • attention slicingはmulti-head attentionでscaled dot-product attentionを計算する際、Q @ K^T行列全体をmaterializeしないようにする
  • UNETモデルのattention head countが8のとき、テンソルshapeは次のとおり
    • Q: (8,4096,40)
    • K^T: (8,40,4096)
    • 最初のMatMul結果: (8,4096,4096)
    • FP32精度では512MBテンソルになる
  • 解決方法は、Qを縦方向に分割し、各chunkに対してattention演算を実行すること
    • Q_sliced shapeは(1,x,40)
    • xは4096をonnxstream::Model::m_attention_fused_ops_partsで割った値
    • デフォルト値は2で、カスタマイズ可能
  • この方式はFP32 UNETモデル全体のメモリ使用量を1.1GBから300MBに下げる
  • FlashAttentionはより効率的な代替になり得るが、AVX・NEONなど対応アーキテクチャ別のカスタムカーネルを書く必要があり、その場合XNNPACKを迂回する必要がある

モデル変換と実行方式

  • OnnxStreamはpath_to_model_folder/model.txtで定義されたモデルを実行する
    • すべてのモデル演算はASCII形式のmodel.txtに入っている
    • weightsファイルは同じフォルダの.binファイル群として存在する必要がある
  • Modelオブジェクトには複数の選択パラメータを設定できる
    • 別のweights providerの指定
    • 量子化モデルのactivation clipping rangeファイルの読み書き
    • モデルキャリブレーションモード
    • FP16演算の使用
    • UINT8演算の使用
    • UINT8動的量子化の使用
    • attention slicingの有効化
    • attention分割数の設定
  • ONNXファイルをOnnxStreamで使うには、onnx2txt.ipynbノートブックでmodel.txtと.bin weightsファイルをエクスポートする
  • PyTorch nn.ModuleをONNXにエクスポートする際はいくつかの条件がある
    • torch.onnx.export呼び出し時、dynamic_axesは空にする必要がある
    • OnnxStreamは動的shape入力をサポートしない
    • 変換前にONNX Simplifierを実行することが強く推奨される

ビルドと実行準備

  • Linux、Mac、Windows、TermuxでStable Diffusion例をビルドできる
    • WindowsではVisual Studio Toolsのx64 Native Tools Command Promptを使用
    • Macではbrew install cmakeでcmakeのインストールが必要
  • まずXNNPACKをビルドする必要がある
    • XNNPACK関数prototypeはいつでも変わり得るため、特定時点に合うcommitへcheckoutする手順が含まれる
    • 基準時点は2023-06-27 00:00以前のmaster commit
  • その後OnnxStreamリポジトリをcloneし、src/buildでcmakeによりビルドする
    • MAX_SPEED=ON
    • XNNPACK_DIRでXNNPACK cloneパスを指定
  • MAX_SPEEDオプションは性能を高められるが、ビルド時により多くのメモリを使用する
    • Windowsでは約**10%**の性能向上
    • Raspberry Piでは50%以上の性能向上
    • 生成された実行ファイルが動作しない場合があり、Termuxテストで問題があった
    • 問題があれば、まずMAX_SPEED=OFFに変えることが推奨される
  • Stable Diffusion 1.5例のweightsはリポジトリのReleasesから取得でき、サイズは約2GB
  • Stable Diffusion XL 1.0 Base weightsはHugging Faceから取得でき、サイズは約8GB

Stable Diffusion例のオプション

  • 例の実行ファイルは、モデル選択、入出力、プロンプト、デコード方式をコマンドラインオプションで制御する
    • --xl: Stable Diffusion 1.5の代わりにStable Diffusion XL 1.0を実行
    • --models-path: Stable Diffusionモデルフォルダを指定
    • --ops-printf: 推論中の現在のoperationをstdoutに出力
    • --output: 出力PNGファイルを指定
    • --decode-latents: diffusionをスキップし、指定したlatentsファイルをデコード
    • --prompt: positive promptを指定
    • --neg-prompt: negative promptを指定
    • --steps: diffusion step数を指定
    • --save-latents: diffusion後のlatentsを指定ファイルに保存
  • Raspberry Piとデコーダ関連のオプションも別途提供される
    • --decoder-calibrate: SD 1.5でのみquantized VAE decoderをキャリブレーション
    • --decoder-fp16: SD 1.5でのみFP16 VAE decoderを使用
    • --not-tiled: SDXL 1.0でのみtiled VAE decoderを使用しない
    • --rpi: Raspberry Piで実行するようにモデルを構成
    • --rpi-lowmem: SDXL 1.0でのみRaspberry Pi Zero 2向け低メモリ構成を適用

関連記事

1件のコメント

 
GN⁺ 2023-10-04
Hacker News のコメント

  • 興味深い。核心はこの一文だと思う: 「OnnxStream は OnnxRuntime よりもメモリ使用量が最大55分の1でありながら、速度低下は0.5〜2倍にとどまる場合がある」
    ビデオメモリ/メモリ使用量と推論時間のトレードオフは、Raspberry Pi のようにメモリが限られる場合だけでなく、ほかの状況でも有利に働きそうに見える
    実際にこの重みのアンロード方式で、同じメモリ内でより大きなバッチサイズを処理できるなら、レイテンシは増えてもスループットは大きく伸ばせるのではないかと気になる

    • これを LLM に使いたい。メモリ占有がそれだけ減るなら、1枚の GPU により多くのモデルを同時に載せられるし、クロックが追いつくという前提なら、個々のモデルの推論速度低下を十分に相殺できるかもしれない
    • 「0.5〜2倍遅い」はタイプミスでは？ 0.5倍遅いなら、むしろ2倍速いという意味になってしまう
      おそらく「1.5〜2倍遅い」と書きたかったのだと思う
    • かなり単純に読んだ限りでは、必要なメモリ帯域幅を減らしたのではなく、ワーキングセットのサイズだけを小さくしたように見える
      推論はたいてい「このモデルがこのシステムに収まるか」という段階を超えるとメモリ帯域幅に縛られるので、より大きなバッチサイズでスループットを伸ばすうえでは、この手法はあまり役に立たない気がする。インスタンス1つだけでも、すでにメモリコントローラを飽和させている可能性が高い
      ただし学習のほうでは役に立つかもしれない

  • 11時間とは、昔 Amiga 500 でレイトレーシングしていた頃を思い出す。「最終」レンダーは間違いなく一晩中回す作業だった

    • 今でも時々そうする。現代的な双方向レイトレーサーでは面白い技法が使えるし、水たまりにできる明るい線のようなコースティクスを見たかった
      ただしコースティクスは明るく見えても統計的にはかなりまれな現象なので、きちんと得るにはレンダーエンジンの制限を外して、一晩そのまま回す必要がある
      出来上がるのは、見事なコースティクスを備えた、腕の足りない作者による平凡なシーン画像1枚だ。本業は続けたほうがよさそうだ
    • 先に低品質レンダーを走らせるのは、間違ったまま一晩を無駄にするより、1時間だけ無駄にして正しいか確認するほうがましだったからだ
      その頃、別の趣味が必要だと思った。ちょうどその直前に、ある優秀な人が OpenGL でシーンを先に見られるようにレンダリングするツールを出していた。Amiga では無理だっただろうが、自分のマシンではぎりぎり動いた
    • 自分の 286 でも同じだった。povray を設定して寝て、朝、学校へ行く前に画像を確認していた
    • もっと後の話だが、386 上で怪しげなコピーの 3DSMAX を使って似たようなことをしていた
    • C64 で Mandelbrot フラクタルを回していた記憶がある。コードのデバッグが本当に大変だった

  • MBP で invoke.ai を使って Stable Diffusion を動かしているのだが、SD のパラメータをもっとうまく詰めるおすすめはあるだろうか？ 同じプロンプトと見た目には同じ設定、たとえば Euler A のような同じモデルを使っても、ネットで見かける画像品質にどうしても届かない

    • これまで試した中ではこれが一番よかったが、Mac 対応はなさそうだ。元の ControlNet 開発者が作った Fooocus の多機能フォークで、短いプロンプトから出る品質が驚くほど高い: https://github.com/MoonRide303/Fooocus-MRE
      標準の SD 1.5 には、速いという理由で Volta を使っている: https://github.com/VoltaML/voltaML-fast-stable-diffusion/com...
      本当に良い SD 1.5 の画像品質は、ファインチューニング済みモデル、LoRA、ControlNet、その他の補強機能を惜しみなく使わないと出ない。たとえば元画像を構造としてなぞらせたり、画像内の特定領域ごとにプロンプトを指定したりするようなものだ。InvokeAI も実際には機能が多く、ノード UI の中にこうした補強機能がたくさん隠れているが、Volta のような別の UI はもっと直接的に露出している
    • カスタム重みを使っている？ たぶん使っているとは思うが、標準の RunwayML 1.5 の重みと、特定目的向けにファインチューニングされたモデルとの間には大きな差がある
      一般に、印象的なファインチューニング済みモデルは標準の重みより汎用性がかなり低いが、実用上は大きな問題にならず、結果はかなり良くなることがある
    • MBP M1 で Invoke.ai や MochiDiffusion を使ったとき、同じ経験をした。ほかの画像品質に合わせられたのは Automatic1111（https://github.com/AUTOMATIC1111/stable-diffusion-webui）だけだった
      Invoke や Nvidia のグラフィックカードより時間とメモリは多く必要だが、そこまで悪くはない。標準的な 512x768px 品質の画像は 1〜2 s/it、Hires Fix を使った高品質な 1024x1536px 画像は 14〜20 s/it 程度だ
    • その画像が生成器から直接出てきたものだと明記されている？ 自分が見た制作過程の動画では、「緑の野原に立つ少女」から始めて、手やポーズなどを直すためにインペインティングを1時間以上やっていた
    • Draw Things に CUDA 互換シードモードが追加されたので、Mac でも NVDIA カードで生成した画像に合わせられる

  • これをデジタルフォトフレームや壁掛け絵に組み込んで動かしたら、本当に格好よさそうだ

    • 以前、Raspberry Pi Zero 2 W で Stable Diffusion を動かしていた旧バージョンで、これを作ってみた: https://hackaday.com/2023/09/19/e-paper-news-feed-illustrate...
      https://github.com/rvdveen/epaper-slow-generative-art/
    • 今まさにそれを e-ink ディスプレイで作っている。残念ながら、リポジトリ内の XNNPACK 関連部分を Pi Zero 2W でビルドできていない
    • 良いアイデアだ。10時間ごとくらいに、ユーザーが与えたテーマのようなものをもとに、自分で新しい画像を作って更新すればよさそうだ
    • ただし環境にやさしくはない

  • 驚くべき成果だが、画像を生成するのには当然ものすごく時間がかかる。README には 11時間と書かれている

    • その通り。この実装を自分が必要としたり使ったりすることはないだろうが、ここで使われている手法は他のツールへ移っていくだろうし、それは素晴らしいことだ
    • より高度な方式と比べたとき、コストと電力の面でどの程度同等なのかを見ると面白そうだ
      例えば Pi Zero 2 を100台、別に W である必要もないが、それで画像100枚を作るのにかかる時間・コスト・必要な全ハードウェア・電力を、平均的なミドルレンジ PC のようなものと比較する、といった具合だ
      おそらく PC が依然として勝つ気がする
      Zero 2 は実用性というより挑戦課題として使われた感じなので、Pi 4 や 5 のほうがより良い基準点かもしれない
    • Raspberry Pi、それも Zero 2 で動くというのが肝心な点だ

  • 印象的だ
    まことに、電灯やトースターでさえ優れた知性を宿す時代が近づいているのだ
    この分野を何年も見てきたが、この10年は驚くべきものだった
    「だった」と言うのは、最近の 6〜18か月の加速はまたまったく別の次元だからだ
    2年後に何ができるかではなく、発展の速度がどれほどさらに速くなるのかが心配だ。そしてまた、さらに、またさらに速くなるのだろう

    • プロンプトを受け取ってその画像を トーストに描いてくれるトースターとはいいね。GPU の熱を実際にトーストを焼くのに活用することもできそうだ
      スタートアップを作ろう

  • ここまで来ると、Stable Diffusion 系の技術を規制しようとする試みは終わったのではないか？モデルと推論インフラを縮小して PS2 でも動かせるレベルになるなら、全体主義的な監視国家なしにこの技術を止めるのは不可能に見えるし、そういう国家でもかろうじて可能かどうかだ

    • 汎用コンピューティングとの戦争は続いているが、人々が汎用コンピューティング機器を所有できないようにするほどには、まだ踏み込めていない
    • その理屈なら窃盗の規制も終わったのでは？体さえあれば道具なしでも窓を開けられるのだから、全体主義的な監視国家なしには窃盗を防げない、という話になる
      メディアの「海賊版」やランサムウェアも同じだ
      国家は純粋に技術的には執行できないものも、ずっと昔から規制してきた
    • この種のモデルの本来の要件は RAM 16GBで、これは20ドル未満で買える。GPU でははるかに高速に動き、そうした GPU も200ドル未満で買える。普通の人々が何百万人も、すでにその両方を持っている
    • PS2 の RAM は 32MBしかなかった。PS3 も 256MB にすぎなかった
      面白い誇張例だというのは分かるが、PS2 で動かすにはこれよりはるかにさらに小さくする必要がある
    • 規制の試みの大半は、既存モデルの生成利用というより、ますます大規模化する 学習の実行に焦点を合わせているのだと思っていた。モデルの利用そのものをめぐる規制が別にあるのだろうか？

  • Windows 3.1 向けの Stable Diffusion が出る日が待ちきれない

  • やばい。11時間かかろうが何だろうが、Pi Zero のようなハードウェアで SD が動かせるとはまったく予想していなかった