diffusers/main/ja/quantization/modelopt.md

NVIDIA ModelOpt

NVIDIA-ModelOptは、量子化、プルーニング、蒸留、投機的デコーディングなどの最先端のモデル最適化技術を含む統合ライブラリです。TensorRT-LLMやTensorRTなどの下流デプロイメントフレームワーク向けに深層学習モデルを圧縮し、推論速度の最適化を実現します。

開始する前に、nvidia_modeloptがインストールされていることを確認してください。

pip install -U "nvidia_modelopt[hf]"

NVIDIAModelOptConfig を from_pretrained() に渡すことでモデルを量子化できます（事前量子化済みモデルの読み込みも可能です）。Accelerate による読み込みをサポートし、torch.nn.Linear レイヤーを含むモデルであれば、モダリティを問わず動作します。

以下の例では、重みをFP8に量子化します。

import torch
from diffusers import AutoModel, SanaPipeline, NVIDIAModelOptConfig

model_id = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
dtype = torch.bfloat16

quantization_config = NVIDIAModelOptConfig(quant_type="FP8", quant_method="modelopt")
transformer = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=dtype,
)
pipe = SanaPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=dtype,
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB")

prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")

注:

NVIDIA-ModelOpt の量子化手法は、モデルの性能を維持しつつ、さまざまな QAT（Quantization-Aware Training）および PTQ（Post-Training Quantization）技術を用いてモデル重みのメモリフットプリントを削減するように設計されています。ただし、推論時の実際の性能向上は、デプロイメントフレームワーク（例：TRT-LLM、TensorRT）や具体的なハードウェア構成に依存します。

詳細については、こちらを参照してください。

NVIDIAModelOptConfig

NVIDIAModelOptConfig クラスは、3つのパラメータを受け取ります。

• quant_type: 以下の量子化タイプのいずれかを指定する文字列値。
• modules_to_not_convert: 量子化を実行しないモジュールの完全名または部分名のリストです。例えば、SD3Transformer2DModelのpos_embedプロジェクションブロックに対して量子化を一切行わない場合は、modules_to_not_convert=["pos_embed.proj.weight"]と指定します。
• disable_conv_quantization: 真偽値。これをTrueに設定すると、モデル内のすべての畳み込み層に対する量子化が無効になります。これは、チャネル量子化やブロック量子化が一般的に畳み込み層（INT4、NF4、NVFP4で使用）と相性が良くないため、有用です。特定の畳み込み層の量子化のみを無効にしたい場合は、代わりにmodules_to_not_convertを使用してください。
• algorithm: スケールを決定するために使用するアルゴリズムです。デフォルト値は"max"です。その他のアルゴリズムや詳細については、modeloptのドキュメントを参照してください。
• forward_loop: 量子化時のアクティベーションのキャリブレーションに使用するフォワードループ関数。指定しない場合、重みのみを使用して計算された静的なスケール値に依存します。
• kwargs: quant_typeに基づいて呼び出される、基盤となる量子化メソッドに渡すキーワード引数の辞書。

サポートされている量子化タイプ

ModelOptは、weight-only、channel、blockの量子化に対応しており、int8、fp8、int4、nf4、nvfp4をサポートしています。これらの量子化手法は、推論時のモデルの性能を維持しながら、モデルの重みのメモリフットプリントを削減することを目的に設計されています。

Weight-only quantization は、モデルの重みを特定の低ビットのデータ型で保存し、bfloat16 などの高精度なデータ型で計算を実行します。これにより、モデルの重みによるメモリ要件は低減されますが、アクティベーションの計算に伴うメモリピークは維持されます。

サポートされている量子化手法は以下の通りです:

量子化タイプ	サポートされているスキーム	必須のKwargs	補足
INT8	int8 weight only, int8 channel quantization, int8 block quantization	quant_type, quant_type + channel_quantize, quant_type + channel_quantize + block_quantize
FP8	fp8 weight only, fp8 channel quantization, fp8 block quantization	quant_type, quant_type + channel_quantize, quant_type + channel_quantize + block_quantize
INT4	int4 weight only, int4 block quantization	quant_type, quant_type + channel_quantize + block_quantize	channel_quantize = -1 is only supported for now
NF4	nf4 weight only, nf4 double block quantization	quant_type, quant_type + channel_quantize + block_quantize + scale_channel_quantize + scale_block_quantize	channel_quantize = -1 and scale_channel_quantize = -1 are only supported for now
NVFP4	nvfp4 weight only, nvfp4 block quantization	quant_type, quant_type + channel_quantize + block_quantize	channel_quantize = -1 is only supported for now

利用可能な量子化手法と、利用可能な設定オプションの包括的なリストについて詳しく理解するには、公式の modelopt ドキュメントを参照してください。

量子化モデルのシリアライズとデシリアライズ

指定された dtype で量子化モデルをシリアライズするには、まず希望する量子化 dtype でモデルを読み込み、save_pretrained() メソッドを使用して保存します。

import torch
from diffusers import AutoModel, NVIDIAModelOptConfig
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()

model_id = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
quant_config_fp8 = {"quant_type": "FP8", "quant_method": "modelopt"}
quant_config_fp8 = NVIDIAModelOptConfig(**quant_config_fp8)
model = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quant_config_fp8,
    dtype=torch.bfloat16,
)
model.save_pretrained('path/to/sana_fp8', safe_serialization=False)

シリアライズされた量子化モデルを読み込むには、from_pretrained()メソッドを使用します。

import torch
from diffusers import AutoModel, NVIDIAModelOptConfig, SanaPipeline
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()

quantization_config = NVIDIAModelOptConfig(quant_type="FP8", quant_method="modelopt")
transformer = AutoModel.from_pretrained(
    "path/to/sana_fp8",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
pipe = SanaPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_600M_1024px_diffusers",
    transformer=transformer,
    dtype=torch.bfloat16,
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"
prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")