PAI-Blade を使用した安定した拡散推論プロセスの最適化

バックグラウンド

AIGC は、人工知能コンピューティングの分野で急速に成長している重要なビジネスです。Stable Diffusionはその中でも最も人気のあるオープンソースモデルであり、大きな注目を集めています。しかし、アプリケーションシナリオが継続的に拡大するにつれて、安定拡散が直面する推論遅延と計算コストの問題がますます顕著になってきました。

序章

PAI-Blade は、PAI が提供する汎用の推論最適化ツールであり、モデルシステムの共同最適化により最適な推論パフォーマンスを実現できます。PAI-Blade は、完全に動的サイズの AI コンパイラー BladeDISCと、画像生成モデル Stable Diffsuion、大規模言語モデル LLM、大規模スパース推奨モデル CTR、音声認識モデルを含む深層学習自動スケジューリングに基づく高性能コンピューティング ライブラリ BlaDNNに依存しています。 ASR を含む多くのモデルは、自動的な高性能推論の最適化を提供します。

BladeDISC は完全に動的ディメンションをサポートする AI コンパイラーであり、フロントエンドは Pytorch モデルと Tensorflow モデルをサポートします。Pytorch モデルは、TorchScript と TorchDynamo の 2 つの入力モードをサポートでき、バックエンドは AStitch 大規模オペレーター融合テクノロジーと効率的なコード生成ロジックを使用して、モデルのメモリを大量に使用するオペレーターの実行効率を向上させます。BladeDISC は現在、github でオープンソースです (プロジェクト アドレス: https://github.com/alibaba/BladeDISC ) 。

BlaDNN は、深層学習の自動スケジューリングに基づいたハイパフォーマンス コンピューティング ライブラリです。Ansor のアップグレード バージョンとして、BlaDNN は Ansor よりも優れたカーネル パフォーマンスを生成するだけでなく、Tuning チューニングを使用せずに DNN 自動スケジューリングに完全に依存できるため、Dynamic Shape ビジネス シナリオのオンライン自動スケジューリングが可能になります。DNN 自動スケジューリングに基づいて生成された GPU は、演算集約 オペレータの平均パフォーマンスは、最終的なチューニング パフォーマンスの 99.39% に達します モデル システムの共同最適化により、DNN 推論遅延は 2us と低く、1 つの CPU コアのみが使用されるため、 GPU モデル自体のパフォーマンスに対するジッター。

PAI-Blade 高速推論最適化技術の採用により、メモリ集約型オペレータと最適化されたコード生成の大規模な融合、および計算集約型オペレータの自動スケジューリングにより、安定拡散の推論遅延とメモリ使用量を大幅に削減し、計算コストを削減できます。 。PAI-Blade を使用して安定拡散を最適化すると、次の 3 つの利点があります。

1. 高性能の Blade を使用すると、Text2Img や Img2Img などの推論プロセスのエンドツーエンド遅延を 2.42 ～ 3.05 倍削減でき、同時にメモリ使用量を最大 5.27 倍削減でき、業界の SOTA 最適化を上回ります。 TensorRT-8.5 などのメソッド。
2. 完全に動的な形状のサポート。1 回の最適化の後、あらゆる形状とバッチ サイズの入力をサポートできます。
3. 使いやすさとスケーラビリティ: 複数の種類のパイプラインでブレードの最適化を有効にするために必要なコードはわずか数行であり、同時に LoRA などの推論スキームの最適化をサポートできます。

使用例

次に、この記事では、人気のある「runwayml/stable-diffusion-v1-5」 Text2Img パイプラインを例として、さまざまな使用シナリオでの PAI-Blade の使用方法を詳しく紹介します。

環境のインストール

次の例の完全な実行スクリプトと関連環境は、registry.cn-beijing.aliyuncs.com/blade_demo/blade_diffusiondocker に統合されています。この Docker では、 を介して推論例を直接実行python /blade/blade_diffusion.pyできます。

公式モデルの最適化

PAI-Blade を使用して安定拡散モデルを最適化する手順は、次の手順に分けることができます。

まず、事前トレーニングされたモデルを読み込みます。

from diffusers import StableDiffusionPipeline

device = torch.device("cuda:0")
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", revision="fp16", torch_dtype=torch.float16).to(device)

2 番目のステップでは、PAI-Blade を使用して最適化します。PAI-Blade は完全に動的な形状最適化ツールであるため、最適化後の推論には任意の形状を使用できることに注意してください。

import torch_blade

opt_cfg = torch_blade.Config()
opt_cfg.enable_fp16 = True
with opt_cfg, torch.no_grad():
    encoder = blade_optimize(pipe.text_encoder, model_inputs=encoder_inputs, allow_tracing=True)
    unet = blade_optimize(pipe.unet, model_inputs=unet_inputs, allow_tracing=True)
    decoder = blade_optimize(pipe.vae.decoder, model_inputs=decoder_inputs, allow_tracing=True)

最後に、元のモデルを最適化されたモデルに置き換えて、元のパイプラインと同じ方法で推論を実行します。

@dataclass
class UNet2DConditionOutput:
    sample: torch.FloatTensor

class TracedUNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.config = pipe.unet.config
        self.in_channels = pipe.unet.in_channels
        self.device = pipe.unet.device

    def forward(self, latent_model_input, t, encoder_hidden_states, **kwargs):
        sample = unet(latent_model_input.half(), t.half(), encoder_hidden_states.half())["sample"]
        return UNet2DConditionOutput(sample=sample)

class TracedEncoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.config = pipe.text_encoder.config
        self.device = pipe.text_encoder.device
        self.dtype = torch.half

    def forward(self, input_ids, **kwargs):
        embeddings = encoder(input_ids.long())
        return [embeddings["last_hidden_state"]]

class TracedDecoder(torch.nn.Module):
    def forward(self, input):
        return decoder(input.half())

pipe.text_encoder = TracedEncoder()
pipe.unet = TracedUNet()
pipe.vae.decoder = TracedDecoder()

A100の性能比較

画像サイズ	サンプルステップ	ピトーチの時代	PAIブレードの時間	スピードアップ	Pytorch のメモリ使用量 (GB)	PAI-Blade のメモリ使用量 (GB)
1024x1024	50	13.26	4.34	3.06倍	32.91	6.25
768x768	50	5.65	2.00	2.83倍	14.99	5.91
512x512	50	2.24	0.84	2.67倍	6.60	5.42

A10の性能比較

画像サイズ	サンプルステップ	ピトーチの時代	PAIブレードの時間	スピードアップ	Pytorch のメモリ使用量 (GB)	PAI-Blade のメモリ使用量 (GB)
1024x1024	50	OOM	13.86	-	OOM	6.89
768x768	50	13.13	5.61	2.34倍	12.60	6.22
512x512	50	4.53	2.11	2.15倍	6.28	5.47

推論結果の検証

PAI-Blade を使用して最適化した後、生成された画像を Pytorch の元の出力と比較して、最適化の結果が正しいかどうかを確認します。左の図は Pytorch の Eager モードの出力、右の図は PAI-Blade の最適化モデルの出力です。

認証されたパイプラインの種類

1. 安定した拡散パイプライン
2. StableDiffusionImg2ImgPipeline
3. 安定した拡散ペイントパイプライン
4. AltDiffusionパイプライン

LoRAの最適化

LoRA は、元のモデルに基づいて事前トレーニングされたモデルを微調整するために低ランク行列を追加し、新しく追加された重みのみをトレーニングすることを指します。これにより、微調整のコストが大幅に削減されます。LoRA の重みは、ディフューザーの公式トレーニング コードを微調整することで取得できます。LoRA を使用してディフューザーが読み込まれた後、モデルは元のモデルとは若干異なる動作をするため、追加の計算オーバーヘッドが発生します。

PAI-Blade は現在、ハグフェイス/ディフューザーの LoRA 最適化手法に適応しています。同様に、Blade は同じパイプラインに対して 1 回最適化するだけで済み、任意の LoRA 重みを推論に使用できます。次回はPAI-Bladeを使ってLoRAを最適化する方法を紹介しますので、ご期待ください。

展望

現在、Stable Diffusion 関連の技術は進化を続けており、PAI-Blade チームは常にコミュニティの動向に注目し、最適化や各種ツールへの適応を行っています。現在のチームは主に次のことに重点を置いています。

1. 関連する最適化をstable-diffusion-webuiに統合します。
2. トレーニング速度の微調整を最適化します。

{{名前}}

{{名前}}