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

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

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

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

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

展開とサービス

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

結果

すべてのパイプラインを適用すると、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チームに感謝します。

連絡先 / リソース