安定拡散~Raspberry Pi

Raspberry Piでも安定拡散が動作します！260MB の RAM が 10 億のパラメータを持つ大規模モデルを「保持」します

Stable Diffusion は 11 か月前に誕生し、消費者向けの GPU で実行できるというニュースは多くの研究者を勇気づけました。それだけでなく、Apple は正式に正式に Stable Diffusion を iPhone、iPad、Mac に「詰め込んで」実行することにしました。これにより、Stable Diffusion のハードウェア デバイスの要件が大幅に軽減され、徐々に誰もが使用できる「ブラック テクノロジ」になっていきます。

現在では、Raspberry Pi Zero 2 でも動作します。

Raspberry Pi Zero 2 「とにかく小さい。５倍の速さだよ」

これはどのようなコンセプトですか？Stable Diffusion の実行は簡単ではありません。これには 10 億パラメータの大規模な Transformer モデルが含まれており、推奨される最小 RAM/VRAM は通常 8GB です。RPI Zero 2 は、512MB のメモリを搭載した単なるマイクロコンピュータです。

これは、RPI Zero 2 で Stable Diffusion を実行することが大きな課題であることを意味します。さらに、作成者は、実行中に記憶域スペースを増やしたり、中間結果をディスクにオフロードしたりしませんでした。

一般に、主要な機械学習フレームワークとライブラリは、推論レイテンシの最小化やスループットの最大化に焦点を当てていますが、そのすべてがメモリ使用量を犠牲にしています。したがって、著者は、メモリ消費を最小限に抑えることに特化した、超小型でハッキング可能な推論ライブラリを作成することにしました。

OnnxStream がそれを行います。

プロジェクトアドレス https://github.com/vitoplantamura/OnnxStream

OnnxStream は、WeightsProvider から派生したクラスであるモデルの重みを提供するコンポーネントから推論エンジンを切り離すという考えに基づいています。WeightsProvider 特殊化では、あらゆるタイプのモデル パラメーターの読み込み、キャッシュ、プリフェッチを実装できます。たとえば、カスタム WeightsProvider は、ディスクに何もロードしたり書き込んだりせずに、HTTP サーバーからデータを直接ダウンロードすることを決定できます (これが、OnnxStream の名前に Stream が含まれている理由です)。使用可能なデフォルトの WeightsProvider には、DiskNoCache と DiskPrefetch の 2 つがあります。

Microsoft の推論フレームワーク OnnxStream と比較すると、OnnxStream は同じ効果を達成するために消費するメモリの量が 1/55 だけですが、速度 (CPU 上) は前者より 0.5 ～ 2 倍遅いだけです。

次に、RPI Zero 2 で実行される Stable Diffusion の効果とその背後にある方法を確認します。実行速度は遅くなりますが、より小型で制限されたデバイスで大規模なモデルを実行するという新しい試みであることに注意してください。

ネチズンはこのプロジェクトが素晴らしいと考えています

Raspberry Pi Zero 2 で安定した拡散を実行する

VAE デコーダは、単精度または半精度で RPI Zero 2 RAM に適合しない Stable Diffusion の唯一のモデルです。これは、モデル内の残留接続、非常に大きなテンソル、畳み込みが原因です。唯一の解決策は静的量子化 (8 ビット) です。

以下のこれらのイメージは、著者のリポジトリに含まれる Stable Diffusion サンプル実装から、VAE デコーダーのさまざまな精度で OnnxStream を使用して生成されました。

最初の画像は、RPI Zero 2 によって生成されたものと同じ潜在を使用して、作成者の PC 上で生成されました。

W16A16精度のVAEデコーダの生成効果 

 W8A32精度のVAEデコーダの生成効果

3 番目のグラフは、RPI Zero 2 によって約 3 時間で生成されました。

凡例: W8A8 精度の VAE デコーダの生成効果

OnnxStreamの特徴

• 推論エンジンは WeightsProvider から切り離されています

• WeightsProvider は DiskNoCache、DiskPrefetch、またはカスタムにすることができます

• 注意スライス

• 動的量子化 (8 ビット符号なし、非対称、パーセンタイル)

• 静的量子化 (W8A8 符号なし、非対称、パーセンタイル)

• 定量的モデルを簡単に校正

• FP16 サポート (FP16 操作の有無にかかわらず)

• 24 個の ONNX 演算子 (最も一般的に使用される演算子) を実装

• 操作は順番に実行されますが、すべてのオペレーターはマルチスレッドです

• 単一の実装ファイル + ヘッダー ファイル

• XNNPACK 呼び出しは XnnPack クラスにラップされます (将来の置き換えのために)

また、OnnxStream は、MatMul、Convolution、要素ごとの Add/Sub/Mul/Div、Sigmoid、Softmax などの特定のプリミティブを高速化するために XNNPACK に依存していることに注意してください。

性能比較

安定拡散は、テキスト エンコーダ (672 の演算と 1 億 2,300 万のパラメータ)、UNET モデル (2050 の演算と 8 億 5,400 万のパラメータ)、および VAE デコーダ (276 の演算と 4,900 万のパラメータ) の 3 つのモデルで構成されます。

バッチ サイズが 1 であると仮定すると、完全なイメージを生成するには 10 ステップかかります。良好な結果を得るには、テキスト エンコーダーを 2 回実行し、UNET モデルを 20 (つまり 2*10) 実行し、VAE デコーダーを 1 回実行する必要があります。 (Euler Ancestral スケジューラを使用)。

この表は、安定拡散の 3 つのモデルのさまざまな推論時間とメモリ消費量 (つまり、Windows のピーク ワーキング セット サイズまたは Linux の最大常駐セット サイズ) を示しています。

UNET モデル (FP16 精度で実行時に OnnxStream で FP16 演算が有効になっている) では、OnnxStream のメモリ消費量は OnnxRuntime のメモリ消費量の 1/55 にすぎませんが、速度は 0.5 ～ 2 倍遅いだけであることがわかります。

このテストについて注意すべき点がいくつかあります。

• OnnxRuntime の最初の実行は、その InferenceSession が最初の実行前に作成され、後続のすべての実行で再利用されるため、ウォームアップ推論です。また、OnnxStream には、その設計が純粋に「熱心」であるため、ウォームアップ推論がありません (ただし、その後の実行では、オペレーティング システムの重みファイルのキャッシュの恩恵を受けることができます)。

• 現在、OnnxStream はバッチ サイズ ! = 1 の入力をサポートしていません。これは、UNET モデルの実行時にバッチ サイズ = 2 を使用することで拡散プロセス全体を大幅に高速化できる OnnxRuntime とは異なります。

• テストでは、OnnxRuntime の SessionOptions (EnableCpuMemArena や ExecutionMode など) を変更しても、結果に目立った影響はありませんでした。

• OnnxRuntime のパフォーマンスは、メモリ消費量と推論時間の点で NCNN (別のフレームワーク) と非常に似ています。

• テスト済みの実行条件: Windows Server 2019、16GB RAM、8750H CPU (AVX2)、970 EVO Plus SSD、VMWare 上の 8 仮想コア。

スライスと量子化に注意

UNET モデルの実行時に「アテンション スライシング」を採用し、VAE デコーダーに W8A8 量子化を使用することは、モデルのメモリ消費を RPI Zero 2 での実行に適したレベルまで削減するために重要でした。

量子化ニューラル ネットワークに関する情報はインターネット上にたくさんありますが、「アテンション スライス」に関する情報はほとんどありません。

ここでの考え方は単純です。目標は、UNET モデルのさまざまなマルチヘッド アテンションのスケーリングされたドット積アテンションを計算するときに、完全な Q@K^T 行列の生成を回避することです。UNET モデルでは、注目ヘッドの数が 8 の場合、Q の形状は (8,4096,40) となり、K^T は (8,40,4096) になります。したがって、最初の MatMul の最終的な形状は (8,4096,4096) で、これは 512MB テンソル (FP32 精度) です。ワオソフト アイオット http://143ai.com

解決策は、Q を垂直に分割し、各 Q ブロックに対して通常どおりアテンション操作を実行することです。Q_sliced 形状は (1,x,40) です。x は 4096 (この場合)、onnxstream::Model::m_attention_fused_ops_parts で除算されます (デフォルトは 2 ですが、カスタマイズ可能です)。

この簡単なトリックにより、FP32 精度で実行する場合、UNET モデルの全体的なメモリ消費量が 1.1 GB から 300 MB に削減されます。より効率的な代替方法は、FlashAttendant を使用することですが、FlashAttendant では、サポートされているアーキテクチャ (AVX、NEON) などごとにカスタム カーネルを作成する必要があり、作成者が示した例では XnnPack をバイパスします。

参考リンク：

https://www.reddit.com/r/MachineLearning/comments/152ago3/p_onnxstream_running_stable_diffusion_in_260mb_of/

https://github.com/vitoplantamura/OnnxStream