Quanto を使用した AI モデルの最適化
導入
トランスフォーマーベースの拡散モデルは日々改良されており、テキストから画像への生成モデルに革命をもたらすことが証明されています。トランスフォーマーの機能により、あらゆるモデルのスケーラビリティとパフォーマンスが向上し、その結果、モデルの複雑さが増大します。
「大いなる力には大いなる責任が伴う」
この場合モデルが非常に複雑になるため、 電力とメモリの消費量も膨大になります。
たとえば、Stable Diffusion 3 のようなモデルで推論を実行するには、テキスト エンコーダ、拡散バックボーン、画像デコーダなどのコンポーネントが関与するため、膨大な GPU メモリが必要になります。この高いメモリ要件は、コンシューマ グレードの GPU を使用する場合に妨げとなり、アクセシビリティと実験の両方を妨げます。
モデルの量子化を入力します。リソースを大量に消費するモデルを、その有効性を犠牲にすることなく、より管理しやすいサイズにスケールダウンできることを想像してみてください。量子化は、高解像度の画像をよりコンパクトな形式に圧縮するようなもので、モデルのパラメータを低精度の表現に変換します。これにより、メモリ使用量が削減されるだけでなく、計算が高速化され、複雑なモデルがよりアクセスしやすくなり、操作が容易になります。
この投稿では、Quanto の量子化ツールが Transformer ベースの拡散パイプラインのメモリ効率を大幅に向上させる方法を検討します。
前提条件
- AI モデルの基本的な理解: 深層学習フレームワーク、特に PyTorch に関する知識。
- 量子化の知識: モデルの量子化の基本概念と AI モデルの最適化におけるその目的。
- セットアップ環境: Python がインストールされ、PyTorch や Optimum などの必要な依存関係がインストールされたマシンにアクセスします。
- コンピューティング リソース: パフォーマンスのベンチマーク用の GPU 対応環境 (NVIDIA H100 など)。
- サンプル モデル/データセット: Quanto を使用したテストおよび実験用の事前トレーニングされたモデルまたはデータセット。
Quanto の紹介: 多用途の PyTorch 量子化バックエンド
量子化は、深層学習モデルの計算量とメモリの需要を軽減し、消費者向けデバイスへの展開により適したものにするための重要な技術です。 32 ビット浮動小数点 (float32) の代わりに 8 ビット整数 (int8) などの低精度データ型を使用することで、量子化によりメモリ ストレージ要件が軽減されるだけでなく、CUDA での int8 または float8 行列乗算などの特定のハードウェアの最適化も可能になります。デバイス。
Optimum の新しい量子化バックエンドである Quanto は、多用途で簡単な量子化プロセスを提供するように設計されています。 Quanto は、さまざまな機能の包括的なサポートで際立っており、さまざまなモデル構成やデバイスとの互換性を保証します。
Quanto は、ハードウェア最適化のためのツールセットである Hugging Face Optimum の一部でもあります
- Eager モードの互換性: トレース不可能なモデルでもシームレスに動作します。
- デバイスの柔軟性: 量子化モデルは、CUDA や MPS を含むあらゆるデバイスに展開できます。
- 自動統合: 量子化/量子化解除スタブ、関数演算、および量子化モジュールを自動的に挿入します。
- 合理化されたワークフロー: 浮動小数点モデルから動的および静的な量子化モデルへの簡単な移行を実現します。
- シリアル化のサポート: PyTorch
weight_only
および 🤗 Safetensors 形式と互換性があります。 - 高速行列乗算: CUDA デバイスでさまざまな量子化形式 (int8-int8、fp16-int4、bf16-int8、bf16-int4) をサポートします。
- 幅広いサポート: int2、int4、int8、float8 の重みとアクティベーションを処理します。
多くのツールは大規模な AI モデルを小さくすることに重点を置いていますが、Quanto はシンプルであらゆる種類のモデルに役立つように設計されています。
クアントのワークフロー
pip を使用して Quanto をインストールするには、以下のコードを使用してください:-
!pip install optimum-quanto
モデルの量子化
以下のコードは、標準モデルを量子化モデルに変換するのに役立ちます。
from optimum.quanto import quantize, qint8
quantize(model, weights=qint8, activations=qint8)
校正する
Quanto のキャリブレーション モードでは、量子化パラメータがモデル内の実際のデータ分布に合わせて調整され、量子化モデルの精度と効率が向上します。
from optimum.quanto import Calibration
with Calibration(momentum=0.9):
model(samples)
量子化を意識したトレーニング
モデルのパフォーマンスが影響を受ける場合は、数エポックの間モデルを調整してモデルのパフォーマンスを向上させることができます。
import torch
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data).dequantize()
loss = torch.nn.functional.nll_loss(output, target)
loss.backward()
optimizer.step()
整数の重みを固定する
モデルをフリーズしている間、浮動小数点数の重みは量子化された重みに変換されます。
from optimum.quanto import freeze
freeze(model)
H100 ベンチマーク調査
H100 GPU は、変圧器や拡散モデルなどの大規模モデルのトレーニングや推論など、要求の厳しい AI タスク向けに特別に設計された高性能グラフィックス カードです。このベンチマークにそれが選ばれた理由は次のとおりです。
- 最高級のパフォーマンス: H100 は並外れた速度とパワーを提供し、テキストから画像への生成パイプラインやテキストからビデオへの生成パイプラインなど、大規模モデルで必要とされる複雑な操作の処理に最適です。
- FP16 のサポート: この GPU は FP16 (半精度浮動小数点) での計算を効率的に処理するため、精度を大幅に犠牲にすることなくメモリ使用量を削減し、計算を高速化します。
- 高度なハードウェア機能: H100 は、混合精度のトレーニングと推論の最適化された操作をサポートしており、パフォーマンスを維持しながらモデル サイズを削減することを目的とした量子化手法に最適です。
ベンチマーク調査では、 新しい量子化ツールであるQuantoを拡散モデルに適用することに主な焦点が当てられています。量子化は大規模言語モデル (LLM) の実践者の間ではよく知られていますが、拡散モデルではあまり一般的には使用されません。 Quanto は、品質をほとんどまたはまったく低下させずに、これらのモデルでメモリを節約できるかどうかを調査するために使用されます。
研究内容は次のとおりです。
環境設定
- CUDA 12.2: 使用される CUDA のバージョン。これは、H100 GPU で計算を実行するために重要です。
- PyTorch 2.4.0: モデルのトレーニングと推論に使用される深層学習フレームワーク。
- ディフューザー: 拡散モデルの構築と実行に使用されるライブラリ。一貫性を確保するために特定のコミットからインストールされます。
- Quanto: 量子化ツール。これも特定のコミットからインストールされます。
- FP16 での計算: デフォルトでは、メモリ使用量を削減し、速度を向上させるために、すべての計算が FP16 で実行されました。ただし、VAE (変分オートエンコーダー) は、潜在的な数値的不安定性を回避するために量子化されません。 主要なパイプラインに焦点を当てる: この研究では、ディフューザーから 3 つの変圧器ベースの拡散パイプラインが選択されています。これらのパイプラインが選択された理由は、これらのパイプラインがテキストから画像への生成の最も高度なモデルの一部を表しており、モデルの品質を維持しながらメモリ使用量を削減する Quanto の有効性をテストするのに最適であるためです。
- PixArt-シグマ
- 安定拡散3
- オーラの流れ
この研究では、H100 GPU を使用し、これらの主要モデルに焦点を当てることで、量子化の観点から LLM と拡散モデルの間のギャップを埋める Quanto の可能性を実証し、パフォーマンスへの影響を最小限に抑えながら大幅なメモリ節約を実現することを目的としています。
Quanto を適用して DiffusionPipeline を量子化する
quanto を使用したモデルの量子化は、次のような簡単なプロセスに従います。
まず必要なパッケージをインストールします。ここでは変圧器、クアント、トーチ、センテンスピースをインストールします。ただし、要件に応じてさらにいくつかのパッケージをインストールする必要がある場合があることに注意してください。
!pip install transformers==4.35.0
!pip install quanto==0.0.11
!pip install torch==2.1.1
!pip install sentencepiece==0.2.0
!pip install optimum-quanto
quantize()
関数は、量子化されるモジュールで呼び出され、量子化の対象となるコンポーネントを指定します。このコンテキストでは、パラメータのみが量子化され、アクティベーションは変更されません。パラメーターは FP8 データ型に量子化されます。最後に、freeze()
関数が呼び出され、元のパラメータが新しく量子化されたパラメータに置き換えられます。
from optimum.quanto import freeze, qfloat8, quantize
from diffusers import PixArtSigmaPipeline
import torch
pipeline = PixArtSigmaPipeline.from_pretrained(
"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
).to("cuda")
quantize(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)
完了したら、パイプラインを使用します。
image = pipeline("ghibli style, a fantasy landscape with castles").images[0]
FP8 の拡散トランスを使用して生成された画像 (画像ソース)
以下のコードは、テキスト エンコーダーを量子化するために使用できます。
quantize(pipeline.text_encoder, weights=qfloat8)
freeze(pipeline.text_encoder)
テキスト エンコーダーはトランスフォーマー モデルでもあり、量子化することもできます。テキスト エンコーダと拡散バックボーンの両方を量子化することにより、大幅にメモリを節約できます。
LLM パイプライン
optimum-quanto
は、Hugging Face 量子化モデルを量子化、保存、リロードするためのヘルパー クラスを提供します。
以下のコードは、Transformers ライブラリを使用して、事前トレーニングされた言語モデル (Meta-Llama-3-8B) を読み込みます。次に、Optimum Quanto の QuantizedModelForCausalLM
クラスを使用してモデルに量子化を適用し、具体的にはモデルの重みを qint4 データ型に変換します。 lm_head
(出力レイヤー) は、精度を維持するために量子化から除外され、最終的な出力品質が高いままであることが保証されます。
from transformers import AutoModelForCausalLM
from optimum.quanto import QuantizedModelForCausalLM, qint4
model = AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B')
qmodel = QuantizedModelForCausalLM.quantize(model, weights=qint4, exclude='lm_head')
# quantized model can be saved using save_pretrained
qmodel.save_pretrained('./Llama-3-8B-quantized')
# reload the model using from_pretrained
from optimum.quanto import QuantizedModelForCausalLM
qmodel = QuantizedModelForCausalLM.from_pretrained('Llama-3-8B-quantized')
トランスフォーマーとの統合
Quanto は、Hugging Face トランスフォーマー ライブラリとシームレスに統合します。 QuantoConfig
にモデルを渡すだけで、どのモデルでも量子化できます。
from transformers import AutoModelForCausalLM, AutoTokenizer, QuantoConfig
model_id = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(model_id)
quantization_config = QuantoConfig(weights="int8")
quantized_model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config= quantization_config
)
Quanto を使用すると、CPU/GPU/MPS (Apple Silicon) を使用しているかどうかに関係なく、モデルを量子化して実行できます。
from transformers import AutoModelForSpeechSeq2Seq
model_id = "openai/whisper-large-v3"
quanto_config = QuantoConfig(weights="int8")
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="cuda",
quantization_config=quanto_config
)
クアントのテンソル
- Quanto の機能: Quanto は、int8 や float8 などのより小さいデータ型に適合するように変更する特別な Tensor (データ構造) を処理します。これにより、モデルのメモリ効率が向上し、実行速度が向上します。
- 変換プロセス: 浮動小数点数の場合は、PyTorch の組み込み変換を使用するだけです。
- 整数の場合は、より小さいデータ型に収まるように数値を丸めます。
- 精度: Quanto は、モデルのパフォーマンスを歪める可能性のある極端な値 (大きすぎるか小さすぎる) を回避することで、変換の正確さを維持しようとします。
- 最適化された操作: これらの小さいデータ型 (低ビット幅) を使用すると、これらの型の操作が高速になるため、モデルをより高速に実行できます。
クアントのモジュール
- Torch モジュールの置き換え: Quanto は、通常の PyTorch コンポーネントを、これらの小さな Tensor で動作する特別な Quanto バージョンに置き換えることができます。
- ウェイト変換: モデルがまだ微調整されているときは、ウェイト (モデルの調整可能なパラメータ) が動的に変換されるため、処理が少し遅くなる可能性があります。ただし、モデルが完成 (フリーズ) されると、この変換は停止します。
- バイアス: バイアス (モデル パラメーターの別のセット) は、変換すると精度が低下するか、非常に高い精度が必要となり非効率になるため、変換されません。
- アクティベーション: モデルのアクティベーション (中間出力) もより小さいデータ型に変換されますが、これは最適な変換設定を見つけるためにモデルが調整された後にのみ変換されます。
量子化できるモジュール
- 線形層: これらはモデル内の全結合層であり、重みは常に量子化されますが、バイアスはそのまま残されます。
- Conv2D レイヤー: 画像処理でよく使用される畳み込みレイヤーです。ここで、重みも常に量子化されますが、バイアスは量子化されません。
- LayerNorm: モデル内のデータを正規化します。 Quanto では、出力を量子化できますが、重みとバイアスはそのまま残されます。
この設定により、良好なレベルの精度を維持しながら、モデルをより小さく、より高速にすることができます。
研究からの最終所見
安定拡散で使用されるような拡散モデルでテキスト エンコーダを量子化すると、パフォーマンスとメモリ効率に大きな影響を与える可能性があります。 3 つの異なるテキスト エンコーダを使用する Stable Diffusion 3 の場合、これらのエンコーダの量子化に関連する観察結果は次のとおりです。
- 2 番目のテキスト エンコーダの量子化: このアプローチは、おそらく 2 番目のエンコーダの固有の特性が原因で、うまく機能しません。
- 最初のテキスト エンコーダーの量子化 (CLIPTextModelWithProjection): このオプションは推奨されており、最初のエンコーダーの量子化によりメモリの節約とパフォーマンスのバランスが適切に保たれるため、一般に効果的です。
- 3 番目のテキスト エンコーダー (T5EncoderModel) の量子化: このアプローチも推奨されます。 3 番目のエンコーダーを量子化すると、モデルのパフォーマンスに大きな影響を与えることなく、メモリを大幅に節約できます。
- 1 番目と 3 番目のテキスト エンコーダの両方の量子化: 1 番目と 3 番目のテキスト エンコーダの両方の量子化を組み合わせるのは、許容可能なパフォーマンス レベルを維持しながらメモリ効率を最大化するための優れた戦略となります。
メモリの節約に関して、この表は、これらの量子化戦略のそれぞれを適用することでどれだけのメモリを節約できるかについての詳細を示し、特にテキスト エンコーダのさまざまな組み合わせの量子化間の相対的な利点とトレードオフに焦点を当てています (画像ソース)。
- すべての場合において拡散トランスを量子化すると、観察されたメモリの節約が主にテキスト エンコーダの量子化によるものであることが保証されます。
- H100 や 4090 などの強力な GPU を考慮すると、bfloat16 を使用すると高速になる可能性があります。
qint8
は、効率的な整数演算とハードウェアの最適化により、一般に推論が高速になります。- QKV プロジェクションを融合することで int8 カーネルが強化され、演算数が減り効率的なデータ処理が活用されることで計算がさらに最適化されます。
- H100 GPU で
bfloat16
とともにqint4
を使用すると、qint4
は値ごとに 4 ビットのみを使用するため、メモリ使用量が改善され、メモリ使用量が削減されます。重みを保存するために必要なメモリ。ただし、これには推論のレイテンシーが増加するという代償が伴います。これは、H100 GPU がまだ 4 ビット整数 (int4
) での計算をサポートしていないためです。重みは圧縮された 4 ビット形式で保存されますが、実際の計算は依然としてbfloat16
(16 ビット浮動小数点形式) で実行されます。つまり、ハードウェアはより複雑な演算を処理する必要があります。処理時間の低下につながります。
結論
Quanto は、PyTorch に強力な量子化バックエンドを提供し、重みを低精度の形式に変換することでモデルのパフォーマンスを最適化します。 qint8
や qint4
などの手法をサポートすることで、Quanto はメモリ消費を削減し、推論を高速化します。さらに、Quanto はさまざまなデバイス (CPU、GPU、MPS) で動作し、さまざまなセットアップと互換性があります。ただし、MPS デバイスで float8
を使用するとエラーが発生します。 全体として、Quanto はディープ ラーニング モデルのより効率的な展開を可能にし、メモリの節約とパフォーマンスのトレードオフのバランスをとります。
参考文献
- Quanto および Diffuser を備えたメモリ効率の高い拡散トランス
- Quanto: Optimum の PyTorch 量子化バックエンド
- クアント0.2.0
- 最適なクアント Github