paint-brush
TheStage AI による拡散モデルの高速化: Recraft の 20B および Red Panda モデルのケース スタディ@thestage
2,832 測定値
2,832 測定値

TheStage AI による拡散モデルの高速化: Recraft の 20B および Red Panda モデルのケース スタディ

TheStage AI7m2024/11/26
Read on Terminal Reader

長すぎる; 読むには

この記事では、TheStage AI の Python ツールを使用して、Recraft のテキストから画像への拡散モデルを高速化する方法について説明します。独自のモデルを最適化し、PyTorch 推論の最適化のための効率的なパイプラインを開発する際の課題を検討します。
featured image - TheStage AI による拡散モデルの高速化:
Recraft の 20B および Red Panda モデルのケース スタディ
TheStage AI HackerNoon profile picture
0-item
1-item


Recraft AI は、独自の拡散モデルを搭載した設計ツールです。 同社の新しいモデル Red Panda は、Midjourney、DALL-E 3、FLUX など、既存のテキストから画像への変換モデルをすべて上回っています。Recraft は、デザイナー向けのユニークなユーザー エクスペリエンスと最先端の AI ツールを組み合わせています。拡散モデルでは、何百万人ものユーザーをサポートするために、強力なハードウェアと効率的なソフトウェアを融合した堅牢な推論インフラストラクチャが必要です。この記事では、TheStage AI アクセラレーション ツールが、Recraft の AI エンジニアと研究者が直感的な Python インターフェイスを通じて Nvidia GPU で最大 2 倍のパフォーマンスを実現するのにどのように役立ったかについて説明します。

導入

近年、拡散モデルは、画像、音楽、ビデオ、3D メッシュなどのコンテンツ生成において驚くべき成果を上げています。これらのモデルは、推論時間の計算コストを使用して、各推論ステップで出力をわずかに更新することで、生成結果を反復的に改善します。現在、マルチステップ推論による推論を使用して高品質の回答を提供する LLM に類似点を見ることができます。


TheStage AI では、複雑なモデル加速フローを完全に自動的に処理するために、任意の PyTorch モデル用の一般的な数学的フレームワークを構築しています。当社のシステムは、ハードウェア上の既存の最適化 (量子化、スパース化) を自動的に検出し、各レイヤーに適切なアルゴリズムを選択して、必要なモデル サイズとレイテンシ制約で最高の品質を実現するか、制限された品質制約で最高の加速を見つけます。これは難しい数学的問題ですが、非常に効率的に解決できます。この記事では、Recraft AI とのパートナーシップを通じてこれらのツールをどのように適用するかについて説明します。


ツールを設計する際には、次の原則を尊重することにしました。


  • ハードウェアのカスタマイズ。高品質なAI製品には、すでに優先インフラストラクチャが備わっている。
  • 品質の維持。高品質のAI製品は品質の低下を許容できない
  • プライバシー。高品質のAI製品は、自社の技術を機密に保ち、自社のインフラストラクチャ上のツールで動作することを望んでいます。
  • 任意の DNN。高品質の AI 製品では、オープンソース モデル用の公開アクセラレーション ツールでは複雑な DNN アーキテクチャを処理して正しい出力を生成できない、社内で開発されたアーキテクチャが使用される場合があります。
  • Pytorch。多くの AI エンジニアにとって最も人気があり便利なフレームワークです。


これらの初期条件を考慮して、次の機能を備えたツールを作成することを目的としました。


  • 制御可能な加速。推論の最適化をビジネス最適化の問題として捉え、お客様が希望するモデル サイズ、レイテンシ、データの品質を指定できるようにします。
  • シンプルなコンパイル。生成されたモデルを効率的なハードウェア実行用にコンパイルするには、1 行のコードだけが必要です。グラフのブレークを処理するためのシンプルなインターフェイスも提供しています。
  • 高速コールド スタート。可能な限り高速なコールド スタートを実現するために、コンパイルされたモデルの保存を有効にしています。これが、JIT コンパイラを使用しない理由です。
  • シンプルな展開。最適化されたモデルの展開は、元のモデルの展開と同じくらい簡単です。


テキストから画像への拡散モデル

ノイズ除去拡散プロセスの簡単な視覚化。


拡散プロセスの各反復では、ニューラル ネットワークが変分オートエンコーダの潜在空間で画像のノイズを除去します。次に、新しく取得した画像が再びノイズと混合されますが、その重みは徐々に小さくなります。最初の反復では、拡散モデルが主要なシーンをスケッチし、大きなノイズの重みを活用して大幅なアップグレードを行います。後の反復では、高周波の詳細を洗練します。この観察により、反復から反復へとネットワーク容量をレイヤー全体に戦略的に割り当てて品質を維持し、特定の加速パイプラインを設計できます。ただし、このような割り当てには、数学的洞察とサウンド エンジニアリングを組み合わせた特殊なツールが必要です。ここで、TheStage AI が大いに役立ちます。


雲を見つめて特定の形に焦点を合わせると、私たちの心は自分の考えに似たランダムな形を識別できます。この現象は、脳がノイズ内のパターンを識別し、精神的状況に一致する要素を見つける方法を示しています。同様に、拡散モデルは最初の反復でこの概念を採用し、ノイズからパターンを選択して、目的のイメージの予備スケッチを作成します。


拡散モデルの加速と圧縮

拡散モデルの高速化は、任意の DNN の高速化と見なすことができますが、特定の課題を考慮する必要があります。たとえば、静的量子化は通常大幅な高速化をもたらしますが、反復ごとに活性化分布が変化するため、拡散モデルでは課題が生じます。これに対処するには、すべての反復で最適な値を適切に推定するか、反復ごとに異なる量子化設定を使用する必要があります。


拡散モデルは、トレーニングして高いパフォーマンスを達成するのが困難です。それでも、 Recraft チームが実証した結果は、すべての最新のテキストから画像への変換モデルを上回っています。このようなモデルの劣化を検証するのは難しい場合があり、元の特徴のセマンティクスを維持する加速技術を使用することが重要です。量子化アルゴリズムは、さまざまな活性化分布の課題に対処できる場合、良い選択肢になります。次のセクションで説明する自動パイプラインを見てみましょう。

プロファイリング

特定のデータを使用して特定のモデルをプロファイルすると、次のことが可能になります。


  • 各パラメータのサイズを決定する
  • 各基本ブロックに適用可能な量子化、スパース化、プルーニングアルゴリズムを特定する
  • 異なるメモリレイアウトを持つ個々のブロックのレイテンシを推定する
  • 収集したすべての情報をANNA(自動NNアクセラレータ)用にコンパイルします。

自動圧縮と加速

プロファイラーが必要なデータをすべて収集したら、ANNA ボードを起動してスライダーを動かし、さまざまな最適化されたモデル バージョンを作成できます。その後、ユーザーは品質と推論コストのトレードオフに基づいて最適な候補を選択できます。当社のツールは、これらの主観的な品質の決定を簡単な方法で処理します。


TheStage AI ANNA。スライダーを動かして、品質の低下を最小限に抑えながらモデルのサイズやレイテンシを調整します。

オペレーションの融合とコンパイル

前述のように、JIT コンパイルは新しいノードのコールド スタート時間を増加させるため使用しません。また、既製のコンパイラも使用しません。代わりに、さまざまなアルゴリズムを混在させることができる独自の複雑なアクセラレーション構成をコンパイルします。これを実現するために、ハードウェアに依存しない方法で DNN アクセラレーション レイヤーを定義する独自の内部プロトコルを開発しました。TheStage AI アクセラレーション フレームワークの主な利点の 1 つは、単一のチェックポイントを幅広いハードウェア用にコンパイルできるため、AI ソフトウェアのクロスプラットフォーム互換性が解決されることです。この機能は、アプリケーション開発におけるエッジ デバイスの展開に特に重要になります。


DNN コンパイラの目標は次のとおりです。


  • グラフの簡素化。数学的演算の融合により実行グラフを簡素化し、推論時間を短縮します。
  • メモリ管理。各操作に必要なメモリを計算し、効率的なメモリ再利用で割り当てスケジュールを管理します。
  • 最適な実装。各基本操作の最適な実装をプロファイルします。最適な実装には特定のメモリレイアウトが必要になる場合があり、層間接続の分析につながるため、これは困難な作業です。
  • オペレーションスケジュール。最適化された実行グラフのオペレーションスケジュールを作成します。
  • シリアル化。この情報をすべて保存して、後続の実行でモデルを再コンパイルしないようにします。


連続した操作は、融合カーネルに組み合わせることができます。融合カーネルは、最初の操作の出力をグローバル メモリにコピーする代わりに、2 番目の操作をローカル メモリのレジスタで直接評価します。メモリ転送は実際の計算よりも時間がかかることが多いため、これにより推論が大幅に高速化されます。ただし、すべての操作シーケンスを融合できるわけではありません。融合と完全に互換性がないものもあります。要素単位の操作の場合、融合カーネルを自動的に生成できます。たとえば、Nvidia の NVFuser ツールは、要素単位の操作の任意のシーケンスのカーネルを生成できます。

展開とサービス

推論サーバーと自動スケーリング パイプラインは、受信リクエストをコスト効率よく効率的に処理する上で重要な役割を果たします。また、自動スケーラーの予測スケーリングを設定するために、特定のリクエストのグループ化と統計収集を含めることもできます。今後の記事では、効率的な推論サーバーについてさらに詳しく説明します。

結果

すべてのパイプラインを適用すると、Pytorch コンパイラ ( torch.compile ) よりも優れたパフォーマンスを実現できます。もちろん、float16 の積極的な PyTorch 実行よりも大幅に優れています。さらに、PyTorch コンパイラは各モデルの初期化で JIT コンパイル アプローチを使用するため、多くの入力サイズの再コンパイルが必要になり、レイテンシが非常に重要な実際のアプリケーションではコールド スタートが十分に長くなります。


ビジネス上のメリット

Recraft の製品が高速化されることで得られる主なビジネス上のメリットは次のとおりです。


  • 同じハードウェアで2倍のユーザーにサービスを提供することで、インフラストラクチャのコストを削減します。
  • 画像生成の応答時間が短縮され、ユーザーエクスペリエンスが向上しました。
  • ピーク負荷時により多くの同時ユーザーにサービスを提供できる能力
  • より迅速なサービス提供による競争上の優位性


TheStage AI 最適化ツールを使用すると、品質を低下させることなくテキストから画像へのモデルを高速化できるため、顧客のユーザー エクスペリエンスが向上します。


Recraft CEO、アンナ・ヴェロニカ・ドログシュ

謝辞

これらの結果は、当社のツールと大規模ワークロード製品に関する研究の優れた検証となります。TheStage AI チームは、さらに優れたパフォーマンスの提供に向けて引き続き取り組んでいます。これを実現するために、優れたパートナーと協力しています。以下に深く感謝いたします。


  • 実りある協力をしてくれたRecraftのCEO、Anna Veronikaに感謝します。最高のデザインツールを提供するという彼らの偉大な旅に、少しでも関われることを嬉しく思います。
  • Recraft の AI 責任者である Pavel Ostyakov 氏には、DNN に関する専門知識、ツールに関する強力なフィードバック、そして協力プロジェクトに対する挑戦的な目標の設定に対して感謝の意を表します。
  • この素晴らしい製品を構築してくれたRecraft AIチームに感謝します。この記事の画像は Recraft で生成されました。
  • 私たちの研究のために優れた GPU インフラストラクチャを一貫してサポートしてくれたNebiusチームに感謝します。

連絡先 / リソース

ご質問がございましたら、お気軽にお問い合わせください。推論インフラストラクチャのコスト削減をお手伝いします。

メールアドレス: [email protected]

TheStage AI メインページ: thestage.ai

TheStage AI推論最適化プラットフォーム: app.thestage.ai