# NVIDIA ModelOpt

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

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

```bash
pip install -U "nvidia_modelopt[hf]"
```

`NVIDIAModelOptConfig` を [from_pretrained()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.from_pretrained) に渡すことでモデルを量子化できます（事前量子化済みモデルの読み込みも可能です）。[Accelerate](https://hf.co/docs/accelerate/index) による読み込みをサポートし、`torch.nn.Linear` レイヤーを含むモデルであれば、モダリティを問わず動作します。

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

```python
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）や具体的なハードウェア構成に依存します。
> 
> 詳細については、[こちら](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples)を参照してください。

## NVIDIAModelOptConfig

`NVIDIAModelOptConfig` クラスは、3つのパラメータを受け取ります。
- `quant_type`: 以下の量子化タイプのいずれかを指定する文字列値。
- `modules_to_not_convert`: 量子化を実行しないモジュールの完全名または部分名のリストです。例えば、[SD3Transformer2DModel](/docs/diffusers/main/ja/api/models/sd3_transformer2d#diffusers.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 ドキュメント](https://nvidia.github.io/Model-Optimizer/)を参照してください。

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

指定された dtype で量子化モデルをシリアライズするには、まず希望する量子化 dtype でモデルを読み込み、[save_pretrained()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.save_pretrained) メソッドを使用して保存します。

```python
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()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.from_pretrained)メソッドを使用します。

```python
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")
```