NVIDIA ModelOpt

NVIDIA-ModelOpt es una biblioteca unificada de técnicas de optimización de modelos de última generación, como cuantización, poda, destilación, decodificación especulativa, etc. Comprime modelos de aprendizaje profundo para frameworks de despliegue posteriores como TensorRT-LLM o TensorRT, con el fin de optimizar la velocidad de inferencia.

Antes de comenzar, asegúrate de tener instalado nvidia_modelopt.

pip install -U "nvidia_modelopt[hf]"

Cuantice un modelo pasando NVIDIAModelOptConfig a from_pretrained() (también puede cargar modelos pre-cuantizados). Esto funciona para cualquier modelo de cualquier modalidad, siempre que admita la carga con Accelerate y contenga capas torch.nn.Linear.

El ejemplo siguiente solo cuantiza los pesos a 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")

Nota:

Los métodos de cuantización en NVIDIA-ModelOpt están diseñados para reducir la huella de memoria de los pesos del modelo mediante diversas técnicas de QAT (Quantization-Aware Training) y PTQ (Post-Training Quantization), manteniendo al mismo tiempo el rendimiento del modelo. Sin embargo, la mejora real del rendimiento durante la inferencia depende del marco de despliegue (p. ej., TRT-LLM, TensorRT) y de la configuración de hardware específica.

Más detalles se pueden encontrar aquí.

NVIDIAModelOptConfig

La clase NVIDIAModelOptConfig acepta tres parámetros:

• quant_type: un valor de cadena que menciona uno de los tipos de cuantización siguientes.
• modules_to_not_convert: Una lista de nombres completos o parciales de módulos para los que no se debe realizar la cuantización. Por ejemplo, para no cuantizar los bloques de proyección pos_embed del SD3Transformer2DModel, se especificaría: modules_to_not_convert=["pos_embed.proj.weight"].
• disable_conv_quantization: un valor booleano que, si se establece en True, deshabilita la cuantización para todas las capas convolucionales del modelo. Esto es útil porque la cuantización por canal y por bloque generalmente no funciona bien con las capas convolucionales (se usa con INT4, NF4, NVFP4). Si deseas deshabilitar la cuantización para capas convolucionales específicas, usa modules_to_not_convert en su lugar.
• algorithm: El algoritmo que se utilizará para determinar la escala, por defecto es "max". Puede consultar la documentación de modelopt para conocer más algoritmos y detalles.
• forward_loop: La función del buque de propagación hacia adelante que se utilizará para calibrar la activación durante la cuantización. Si no se proporciona, se basa en valores de escala estáticos calculados utilizando solo los pesos.
• kwargs: Un diccionario de argumentos por nombre para pasar al método de cuantización subyacente que se invocará en función de quant_type.

Tipos de cuantización admitidos

ModelOpt admite cuantización int8, fp8, int4, nf4 y nvfp4, ya sea solo de pesos, por canal o por bloque. Los métodos de cuantización están diseñados para reducir la huella de memoria de los pesos del modelo manteniendo su rendimiento durante la inferencia.

La cuantización solo de pesos almacena los pesos del modelo en un tipo de datos de bajo número de bits, pero realiza los cálculos con un tipo de datos de mayor precisión, como bfloat16. Esto reduce los requisitos de memoria derivados de los pesos del modelo, pero mantiene los picos de memoria para el cálculo de las activaciones.

Los métodos de cuantización soportados son los siguientes:

Consulte la documentación oficial de modelopt para comprender mejor los métodos de cuantización disponibles y la lista exhaustiva de opciones de configuración disponibles.

Serialización y deserialización de modelos cuantizados

Para serializar un modelo cuantizado en un dtype dado, primero cargue el modelo con el dtype de cuantización deseado y luego guárdelo usando el método 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)

Para cargar un modelo cuantizado serializado, use el método 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")