ウェブサイト検索

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行列の掛け算など、特定のハードウェア向けの最適化を可能にします。

Quantoの紹介です。これは、Optimumのために設計された新しい量子化バックエンドで、柔軟で簡単な量子化プロセスを提供します。Quantoは、さまざまな機能に対する包括的なサポートを備えており、多様なモデル構成やデバイスとの互換性を確保しています:

Quantoは、ハードウェア最適化のためのツールセットであるHugging Face Optimumの一部でもあります。

  • イagerモードの互換性: トレース不可能なモデルとシームレスに動作します。
  • デバイスの柔軟性: 量子化モデルは、CUDAやMPSを含む任意のデバイスに展開できます。
  • 自動統合：量子化/逆量子化スタブ、機能操作、及び量子化されたモジュールを自動的に挿入します。
  • 効率的なワークフロー：フロートモデルから動的および静的量子化モデルへのスムーズな移行を提供します。
  • シリアライズサポート：PyTorchのweight_onlyおよび🤗 Safetensors形式と互換性があります。
  • 加速された行列乗算：CUDAデバイス上でさまざまな量子化フォーマット（int8-int8、fp16-int4、bf16-int8、bf16-int4）をサポートします。
  • 幅広いサポート：int2、int4、int8、およびfloat8の重みと活性化を処理します。

多くのツールが大規模なAIモデルを小さくすることに焦点を当てる一方で、Quantoはすべての種類のモデルに対してシンプルかつ便利に設計されています。

クワントワークフロー

Quantoをpipを使ってインストールするには、以下のコードを使用してください:-

!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: H100 GPUでの計算を実行するために重要なCUDAのバージョン。
  • PyTorch 2.4.0: モデルのトレーニングと推論に使用される深層学習フレームワーク。
  • ディフューザー：特定のコミットからインストールされた、一貫性を確保するために拡散モデルを構築および実行するために使用されるライブラリ。
  • Quanto: 特定のコミットからインストールされた量子化ツール。
  • FP16の計算: デフォルトでは、すべての計算はFP16で行われ、メモリ使用量を減らし、速度を向上させます。ただし、VAE（変分オートエンコーダ）は潜在的な数値的不安定性を避けるために量子化されていません。主要なパイプラインに焦点を当てる: この研究では、Diffusersから3つのトランスフォーマーベースの拡散パイプラインを選択しました。これらのパイプラインは、テキストから画像への生成に関して最も高度なモデルのいくつかを表しており、モデルの品質を維持しながらメモリ使用量を削減するQuantoの有効性をテストするのに理想的です。
  • PixArt-Sigma
  • ステーブルディフュージョン3
  • オーラフロー

H100 GPUを使用し、これらの主要なモデルに焦点を当てることによって、この研究は、量子化においてLLMと拡散モデルのギャップを埋めるQuantoの可能性を示すことを目指しており、パフォーマンスへの影響を最小限に抑えつつ、重要なメモリ節約を提供します。

拡散パイプラインに量子化を適用するためのQuanto

quantoを使用してモデルを量子化するプロセスは簡単です:-

まず、必要なパッケージをインストールします。ここではtransformers、quanto、torch、sentencepieceをインストールしています。ただし、必要に応じて他のパッケージもインストールする必要があるかもしれませんのでご注意ください。

!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のような小さなデータ型に適応させて変更します。これにより、モデルはメモリ効率が高まり、実行が速くなります。
  • 変換プロセス：浮動小数点数の場合、単純にPyTorchの組み込み変換を使用します。
  • 整数の場合、それらの数値を小さいデータ型に収めるために丸めます。
  • 精度：Quantoは、モデルのパフォーマンスを歪める可能性のある極端な値（あまりにも大きいまたは小さい値）を避けることによって、変換を正確に保とうとします。
  • 最適化された操作：これらの小さいデータ型（低ビット幅）を使用することで、これらの型の操作がより迅速になるため、モデルはより速く実行されます。

Quantoのモジュール

  • Torchモジュールの置き換え：Quantoは、これらの小さなテンソルで動作する特別なQuantoバージョンで通常のPyTorchコンポーネントを入れ替えることができます。
  • 重み変換：モデルがまだ微調整中のとき、重み（モデルの調整可能なパラメータ）は動的に変換されるため、少し遅くなることがあります。しかし、モデルが最終化（フリーズ）されると、この変換は停止します。
  • バイアス: バイアス（別のモデルパラメータのセット）は、変換すると精度が低下するか、非常に高い精度が必要になるため、効率が悪くなります。
  • アクティベーション: モデルのアクティベーション（中間出力）もより小さなデータ型に変換されますが、これはモデルが最適な変換設定を見つけるために調整された後のみ行われます。

量子化できるモジュール

  • 線形層：これは、モデル内の完全に接続された層であり、重みは常に量子化されますが、バイアスはそのままです。
  • Conv2D層：これらは画像処理に頻繁に使用される畳み込み層です。ここでは、重みも常に量子化され、バイアスは量子化されません。
  • レイヤーノルム: これはモデル内のデータを正規化します。Quantoでは、出力は量子化できますが、重みとバイアスはそのまま残されています。

このセットアップは、モデルがより小型で高速になることを可能にしながら、なおかつ良好な精度を維持します。

研究からの最終観察

拡散モデル、特にStable Diffusionで使用されるテキストエンコーダの量子化は、性能とメモリ効率に大きく影響する可能性があります。3つの異なるテキストエンコーダを使用するStable Diffusion 3において、これらのエンコーダの量子化に関する観察は以下の通りです:

  1. 第二のテキストエンコーダの量子化: このアプローチはうまく機能しない可能性が高く、第二のエンコーダの特定の特性が原因です。
  2. 最初のテキストエンコーダー（CLIPTextModelWithProjection）の量子化：このオプションは推奨されており、一般的に効果的です。最初のエンコーダーの量子化は、メモリの節約とパフォーマンスの良いバランスを提供します。
  3. 第三のテキストエンコーダ（T5EncoderModel）の量子化：このアプローチも推奨されています。第三のエンコーダを量子化することで、モデルの性能に大きな影響を与えることなく、かなりのメモリ節約が期待できます。
  4. 最初と第三のテキストエンコーダの両方を量子化すること：最初と第三のテキストエンコーダの両方を量子化することは、許容できるパフォーマンスレベルを維持しながら、メモリ効率を最大化するための良い戦略となる可能性があります。

メモリの節約に関して、この表は、特に異なるテキストエンコーダの組み合わせを量子化することによる相対的な利点とトレードオフに焦点を当てて、これらの量子化戦略を適用することによってどれだけのメモリを節約できるかの詳細を提供しています。

  • すべての場合において拡散トランスフォーマーを量子化することで、観察されたメモリ節約が主にテキストエンコーダの量子化によるものであることが保証されます。
  • H100や4090のような強力なGPUを考慮した場合、bfloat16を使用することでより高速になることがあります。
  • qint8は、効率的な整数演算とハードウェアの最適化により、一般的に推論においてより速くなります。
  • QKVプロジェクションを融合することで、int8カーネルが厚くなり、演算の回数を減らし、効率的なデータ処理を活用することで計算の最適化がさらに進みます。
  • qint4をbfloat16と共にH100 GPUで使用すると、qint4が値ごとにわずか4ビットしか使用しないため、メモリ使用量が改善されます。これにより、重みを保存するために必要なメモリ量が減少します。しかし、これは推論のレイテンシが増加するという代償を伴います。これは、H100 GPUが依然として4ビット整数（int4）での計算をサポートしていないためです。重みが圧縮された4ビット形式で保存されていても、実際の計算は依然としてbfloat16（16ビット浮動小数点形式）で行われるため、ハードウェアはより複雑な演算を処理しなければならず、処理速度が遅くなります。

結論

QuantoはPyTorchのための強力な量子化バックエンドを提供し、重みを低精度フォーマットに変換してモデルのパフォーマンスを最適化します。qint8やqint4のような技術をサポートすることで、Quantoはメモリ消費を削減し、推論を高速化します。さらに、Quantoは異なるデバイス（CPU、GPU、MPS）で機能し、さまざまなセットアップと互換性があります。ただし、MPSデバイスでfloat8を使用するとエラーが発生します。全体として、Quantoはメモリの節約とパフォーマンスのトレードオフをバランス良く取り入れ、深層学習モデルのより効率的なデプロイを可能にします。

参考文献

  • メモリ効率の良い拡散トランスフォーマーとQuantoおよびDiffusers
  • Quanto: OptimumのためのPyTorch量子化バックエンド
  • quanto 0.2.0
  • 最適量 GitHub

関連記事: