Hugging Face's logo

Diffusers documentation

NVIDIA ModelOpt

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

NVIDIA ModelOpt

NVIDIA-ModelOpt 是一个集成了量化、剪枝、蒸馏、投机解码等前沿模型优化技术的统一库。它针对 TensorRT-LLM 或 TensorRT 等下游部署框架压缩深度学习模型，以优化推理速度。

在开始之前，请确保已安装 nvidia_modelopt。

pip install -U "nvidia_modelopt[hf]"

通过向 from_pretrained() 传入 NVIDIAModelOptConfig 来量化模型（也可以加载预量化模型）。只要模型支持使用 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（量化感知训练）和 PTQ（训练后量化）技术来降低模型权重的内存占用，同时保持模型性能。然而，推理期间的实际性能提升取决于部署框架（例如 TRT-LLM、TensorRT）以及具体的硬件配置。

更多详细信息请参见此处。

NVIDIAModelOptConfig

NVIDIAModelOptConfig 类接受三个参数：

  • 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：用于确定 scale 的算法，默认为 "max"。有关更多算法和详细信息，请参阅 modelopt 文档。
  • forward_loop：用于在量化过程中校准激活值的前向循环函数。如果未提供，则依赖仅使用权重计算得出的静态缩放值。
  • kwargs：一个字典，包含要传递给底层量化方法的关键字参数，该方法将根据 quant_type 被调用。

支持的量化类型

ModelOpt 支持权重、通道和块量化的 int8、fp8、int4、nf4 和 nvfp4。这些量化方法旨在降低模型权重的内存占用，同时保持模型在推理期间的性能。

仅权重量化以特定的低比特数据类型存储模型权重，但使用更高精度的数据类型（例如 bfloat16）执行计算。这降低了模型权重带来的内存需求，但保留了激活值计算时的内存峰值。

支持的量化方法如下：

量化类型支持的方案必需的 Kwargs补充说明
INT8int8 weight only, int8 channel quantization, int8 block quantizationquant_type, quant_type + channel_quantize, quant_type + channel_quantize + block_quantize
FP8fp8 weight only, fp8 channel quantization, fp8 block quantizationquant_type, quant_type + channel_quantize, quant_type + channel_quantize + block_quantize
INT4int4 weight only, int4 block quantizationquant_type, quant_type + channel_quantize + block_quantizechannel_quantize = -1 is only supported for now
NF4nf4 weight only, nf4 double block quantizationquant_type, quant_type + channel_quantize + block_quantize + scale_channel_quantize + scale_block_quantizechannel_quantize = -1 and scale_channel_quantize = -1 are only supported for now
NVFP4nvfp4 weight only, nvfp4 block quantizationquant_type, quant_type + channel_quantize + block_quantizechannel_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")
Update on GitHub

←quanto AutoRound→