NVIDIA Model Optimizer mejora el rendimiento de modelos AI mediante cuantización post-entrenamiento

La cuantización de modelos es un método eficaz que reduce el uso de VRAM y mejora el rendimiento de inferencia en dispositivos como las GPU NVIDIA GeForce RTX. Este proceso disminuye los requisitos computacionales y de memoria, manteniendo la calidad del modelo, lo que permite que los modelos de IA funcionen de manera más eficiente en entornos con recursos limitados.

Este artículo describe cómo utilizar el NVIDIA Model Optimizer para cuantizar un modelo CLIP en formato FP8, empleando el método de cuantización posterior al entrenamiento (PTQ). Para una introducción general a la cuantización de modelos, consulte Cuantización de Modelos: Conceptos, Métodos y Por Qué Es Importante.

¿Qué es NVIDIA Model Optimizer?

La biblioteca NVIDIA Model Optimizer (ModelOpt) incluye técnicas avanzadas de optimización de modelos para comprimir y acelerar modelos de IA. Estas técnicas abarcan cuantización, destilación, poda, decodificación especulativa y escasez. ModelOpt acepta modelos en formatos Hugging Face, PyTorch o ONNX como entrada, y proporciona APIs de Python para facilitar la combinación de diferentes técnicas de optimización.

ModelOpt admite formatos de cuantización altamente eficientes como FP4, FP8, INT8 e INT4, así como algoritmos avanzados como SmoothQuant, AWQ, SVDQuant y Double Quantization. Además, soporta tanto PTQ como entrenamiento consciente de la cuantización (QAT).

¿Qué es CLIP?

CLIP (Contrastive Language-Image Pretraining), presentado por OpenAI en 2021, es un modelo de lenguaje-visual que aprende un espacio de incrustación compartido para imágenes y texto a través del aprendizaje contrastivo en grandes pares de imagen-texto. Su capacidad para generar representaciones semánticamente alineadas lo ha convertido en un componente fundamental en sistemas multimodales modernos.

El codificador de texto de CLIP se utiliza ampliamente como módulo de acondicionamiento para la síntesis de texto a imagen (por ejemplo, Stable Diffusion) y texto a video (por ejemplo, AnimateDiff). Su codificador visual actúa como el respaldo visual en modelos de LLM multimodales, como LLaVA, y modelos de percepción de vocabulario abierto, como OWL-ViT. Sucesores como OpenCLIP y SigLIP escalan los datos y refinan el objetivo, pero preservan el paradigma de contraste de doble codificador.

Receta de cuantización

La siguiente receta de cuantización se utiliza como una guía paso a paso para ejecutar la cuantización del modelo CLIP con ModelOpt y entender cómo funciona el proceso.

Primero, prepare los modelos y conjuntos de datos correspondientes como se muestra a continuación:

• Modelo base CLIP: CLIP-ViT-L-14-laion2B-s32B-b82K
• Conjunto de datos de calibración para cuantización: subconjunto de 10K de MS-COCO
• Tareas de evaluación de precisión del modelo enfocadas en tres del CLIP_benchmark
  • cifar100 (clasificación cero-shot)
  • imagenet1k (clasificación cero-shot)
  • mscoco_captions (recuperación cero-shot)

Cómo ejecutar PTQ con ModelOpt

El siguiente ejemplo de código muestra cómo ejecutar PTQ para el modelo CLIP en FP8 utilizando ModelOpt:

 import torch from torch.utils.data import DataLoader, Subset from transformers import CLIPModel, CLIPTokenizer, CLIPImageProcessor from transformers.models.clip.modeling_clip import CLIPAttention import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.plugins.diffusion.diffusers import _QuantAttention # FP8 (E4M3) cuantización estática por tensor FP8_CFG = { "quant_cfg": { "*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, "*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, "*[qkv]_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, "*bmm2_output_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, "default": {"enable": False}, }, "algorithm": "max", } mto.enable_huggingface_checkpointing() mtq.QuantModuleRegistry.register({CLIPAttention: "CLIPAttention"})(_QuantAttention) model = CLIPModel.from_pretrained(args.model_ckpt, attn_implementation="sdpa").half().eval().cuda() tokenizer = CLIPTokenizer.from_pretrained(args.model_ckpt) processor = CLIPImageProcessor.from_pretrained(args.model_ckpt) calib_set = Subset(CLIP_COCO_dataset(ANN, IMG_DIR, tokenizer, processor), range(8192)) loader = DataLoader(calib_set, batch_size=512, num_workers=4) # Calibración: 8k pares de imagen-texto de MS-COCO def calibrate(m): for img, txt in loader: m.get_text_features(input_ids=txt.cuda()) m.get_image_features(pixel_values=img.cuda()) q_model = mtq.quantize(model, FP8_CFG, forward_loop=calibrate) # Guardar el checkpoint del modelo cuantizado q_model.save_pretrained(ckpt_path) mtq.print_quant_summary(q_model) 

FP8_CFG es solo una de las recetas: W8A8 (FP8 en pesos y activaciones), cuantización estática por tensor, calibrada con el simple algoritmo AbsMax. ModelOpt admite muchas más dimensiones de elección (cuantización por canal / granularidad por bloque, cuantización dinámica de activaciones, algoritmos de calibración avanzados como AWQ / GPTQ, y muchos más).

Para obtener el esquema de configuración detallado, consulte la guía de cuantización de ModelOpt. Los hiperparámetros en la configuración de cuantización siempre se pueden ajustar según sea necesario, y encontrar los valores óptimos suele requerir algunas iteraciones.

Después de que mtq.quantize se complete, todas las capas Linear de CLIP llevarán cuantizadores de peso y activación, pero los bloques de atención permanecerán intactos. Esto se debe a que la atención multi-cabeza despacha a torch.nn.functional.scaled_dot_product_attention, una API funcional que el módulo de ModelOpt no puede interceptar por sí solo.

Para incluir la atención en el ámbito de cuantización, registre un reemplazo cuantizado para CLIPAttention:

 mtq.QuantModuleRegistry.register({CLIPAttention: "CLIPAttention"})(_QuantAttention) 

Cada instancia de CLIPAttention se actualiza ahora a _QuantAttention del complemento de ModelOpt. Dentro de su paso hacia adelante, _QuantAttention intercepta de manera transparente la llamada SDPA e inserta cuatro cuantizadores alrededor del núcleo fusionado:

• q_bmm_quantizer, k_bmm_quantizer, v_bmm_quantizer envuelven los tensores proyectados Q / K / V antes de que ingresen al núcleo
• bmm2_output_quantizer envuelve la salida del núcleo (softmax @ V) antes de que fluya hacia out_proj

Esto asegura una cuantización adecuada a lo largo del mecanismo de atención.

Para restaurar precisión, a menudo se recomienda desactivar algunos de los cuantizadores utilizando mtq.disable_quantizer. Esto toma una función como entrada, donde la función misma recibe como entrada el nombre de un módulo. Usando expresiones regulares o coincidencia de cadenas, puede seleccionar las capas a desactivar. En el siguiente ejemplo, se desactivan los cuantizadores en la capa patch_embedding del modelo CLIP.

 import re def filter_func(name): pattern = re.compile( r".*(patch_embedding).*" ) return pattern.match(name) is not None mtq.disable_quantizer(q_model, filter) 

Evaluación del benchmark CLIP

El checkpoint guardado de ModelOpt puede restaurarse en cualquier script de evaluación posterior. Para más detalles, consulte Restaurando ModelOpt Models. El checkpoint cuantizado de CLIP fue evaluado en tres benchmarks: clasificación cero-shot (CIFAR-100, ImageNet-1k) y recuperación cero-shot (MS-COCO Captions). El modelo CLIP FP16 sirve como base.

Basado en los resultados de evaluación, el modelo CLIP-FP8 cuantizado demuestra una calidad comparable al modelo CLIP-FP16. Notablemente, cuando se desactivan los cuantizadores en la capa de embedding de parche, el impacto de la cuantización en la calidad del modelo se vuelve insignificante.

Dentro del flujo PTQ de ModelOpt

Es importante entender que esta etapa implica trabajar con “cuantización falsa” porque el tipo de dato real del modelo no ha cambiado. En cambio, estos cuantizadores insertados actúan como observadores que simulan los efectos de la cuantización mientras mantienen el modelo en su formato de punto flotante original.

El proceso de cuantización falsa funciona de dos maneras clave:

• Recolección de estadísticas: Los cuantizadores recogen estadísticas de tensor (valores mínimos y máximos, por ejemplo) a medida que los datos fluyen a través de ellos. Estas estadísticas se utilizan para calcular parámetros de cuantización óptimos, como factores de escalado.
• Simulación de cuantización: Los cuantizadores realizan una operación de cuantización-descuantización (QDQ) en tensores que fluyen a través de la red. Solo simulan el cálculo de baja precisión y el verdadero aumento de velocidad y ahorro de memoria deben lograrse exportando el modelo a marcos de despliegue como NVIDIA TensorRT.

Esta simulación es crucial porque permite evaluar la precisión del modelo antes de comprometerse a la cuantización real. Los cuantizadores aplican las mismas limitaciones de redondeo y precisión que ocurrirían en el modelo cuantizado desplegado con marcos de inferencia posteriores, por lo que se puede:

• Medir el impacto en la precisión antes del despliegue
• Experimentar con diferentes configuraciones de cuantización
• Identificar capas problemáticas que puedan requerir un manejo especial

En general, el flujo PTQ de ModelOpt sigue seis etapas:

1. Preparar: Establecer la configuración de cuantización para insertar módulos cuantizadores alrededor de los pesos y/o activaciones del modelo.
2. Calibrar: Pasar un pequeño lote de datos representativos a través del modelo para que cada cuantizador pueda recoger estadísticas (por ejemplo, amax de activación) y derivar su factor de escalado.
3. Cuantización falsa: Los cuantizadores ahora aplican un recorrido Q → DQ en punto flotante, simulando fielmente la pérdida de precisión del formato objetivo, mientras el modelo sigue ejecutándose en FP16/BF16.
4. Evaluar: Medir la precisión en un conjunto de evaluación reservado y compararla con la línea base no cuantizada.
5. Iterar: Si la diferencia es inaceptable, ajustar la configuración de cuantización (granularidad, algoritmo, capas cuantizadas), desactivar la cuantización para capas sensibles y recalibrar.
6. Exportar y desplegar: Una vez que la precisión es aceptable, los pesos cuantizados falsamente se comprimen en su verdadera forma de baja precisión y se exportan como un checkpoint para motores posteriores. En nuestro caso, exportamos el checkpoint de PyTorch a ONNX y ejecutamos inferencia con TensorRT. Las mejoras de velocidad y ahorro de memoria ocurrirán allí.

El QAT recupera la pérdida de calidad inducida por la cuantización al ajustar los pesos del modelo con estados de cuantizador congelados. Es más intensivo en cómputo que el PTQ, pero puede mejorar mejor la calidad del modelo cuantizado. Para más detalles, consulte los ejemplos de ModelOpt.

Comenzar con NVIDIA Model Optimizer

Este artículo introdujo el NVIDIA Model Optimizer y demostró un flujo de trabajo típico de cuantización posterior al entrenamiento al cuantizar el modelo CLIP a FP8 con un ejemplo de código práctico. Los resultados en tres conjuntos de datos de evaluación muestran que la cuantización FP8 puede preservar la calidad del modelo mientras habilita un camino de despliegue más eficiente.

¿Listo para comenzar a usar ModelOpt con sus propios modelos? Siga este flujo de trabajo: prepare el modelo y los datos de calibración, establezca la configuración de cuantización, calibre, valide el modelo cuantizado contra métricas de calidad específicas de la tarea, guarde y restaure los checkpoints de ModelOpt.

Para explorar flujos de trabajo adicionales y adaptar ModelOpt a sus propios casos de uso, consulte la documentación de ModelOpt.