NNCFと🤗 Optimumを使ったIntel CPU向けStable Diffusionの最適化
概要
Hugging Faceは、Neural Network Compression Framework (NNCF)の量子化認識トレーニングとToken Mergingを組み合わせることで、Intel CPU向けStable Diffusionを最適化する方法を発表しました。これにより、推論速度が最大5.1倍向上し、モデルサイズが元のPyTorchチェックポイントの0.25倍に削減されます。
Stable Diffusionの最適化
Stable DiffusionパイプラインのUNetコンポーネントは計算コストが最も高い部分であり、これを最適化すると大きな速度向上が得られます。従来のポストトレーニング8ビット量子化は、ピクセルレベルの予測タスクがパラメータ変更に非常に敏感であり、何億ものサンプルでのトレーニングにより冗長性がほとんどないため、このモデルにはうまく機能しません。精度を維持するため、Quantization‑Aware Training (QAT)などのより高度な量子化手法が必要です。
最適化ワークフロー
私たちは、Pokemonデータセット（svjack/Stable-Diffusion-Pokemon-en）でファインチューニングされたStable Diffusionモデルから始めました。Diffusersのテキスト・トゥ・イメージファインチューニングの例を使用し、トレーニングスクリプトにNNCFベースのQATを統合し、元のモデルを教師とする知識蒸留損失を追加し、トレーニングの安定性のために量子化子を除くモデルパラメータに指数移動平均 (EMA) を適用しました。グラディエントチェックポイントとEMAモデルをRAMに保持することで、24 GB VRAMの単一GPUで、4096イテレーションを1日未満で実行できる最適化全体を実行できました。
量子化認識トレーニングを超えて
量子化のみではモデルのフットプリント、ロード時間、メモリ消費、レイテンシが削減されます。私たちは、自己注意ブロックの前に冗長なトークンをマージして計算量を削減するToken Merging (ToME)手法と8ビット量子化を組み合わせました。この組み合わせたワークフローには、前述のように知識蒸留、EMA、グラディエントチェックポイントが含まれます。Pokemonファインチューニング済みモデルから再び開始し、量子化の上にマージ比率0.4のToMEを適用しました。その結果得られたモデルは、クライアントまたはエッジCPUでの推論を想定しています。
結果
PyTorchベースラインをOpenVINO FP32に変換すると、1.9倍の速度向上が得られました。8ビット量子化を追加すると、PyTorchと比較して速度向上が3.9倍になり、モデルのフットプリントが元のチェックポイントの0.25倍に削減されました。量子化の上にToken Mergingをスタックすると、フットプリントを同じ0.25倍レベルに保ちながら、推論速度が5.1倍向上しました。すべての測定は、OpenVINO 2022.3を使用し、Hugging Face SpacesのCPUアップグレードインスタンスで、第3世代Intel® Xeon® ScalableプロセッサーおよびIntel® Deep Learning Boostテクノロジーを搭載し、デフォルトの50推論ステップで行われました。ブログでは、ステップ数を減らすと速度は上がりますが、画質に影響を与える可能性があるとし、ステップ数とスケジューラーの実験を推奨しています。
from optimum.intel import OVStableDiffusionPipeline
# Load and compile the pipeline for performance.
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()
# Generate an image.
prompt = "a drawing of a green pokemon with red eyes"
output = pipe(prompt, num_inference_steps=50, output_type="pil\)).images[0]
output.save("image.png
The training and quantization code are available in the Optimum Intel repository, a demonstration notebook is provided, and optimized models can be found on the Hugging Face Hub under the OpenVINO organization. A live demo runs on Hugging Face Spaces.
汎用Stable Diffusionモデルについては？
Pokemonモデルで示されたワークフローは、 modestなトレーニングリソースでも大規模な最適化が可能であることを示しています。ゼロから汎用Stable Diffusionモデルをトレーニングするのはコストがかかりますが、十分な予算とハードウェアがあれば同じアプローチを適用できます。ただし、Token Mergingはモデル容量を減少させるため、より複雑なデータセットでは最適化時にマージ比率を低く設定すべきです。