La cuantización de modelos es uno de los caminos más directos para reducir el uso de VRAM y acelerar la inferencia en GPUs de consumo como las NVIDIA GeForce RTX. Bajar la precisión computacional sin sacrificar calidad permite que modelos pesados corran en entornos con memoria acotada, algo cada vez más relevante para equipos que despliegan IA fuera del datacenter.
NVIDIA publicó una guía paso a paso para cuantizar el modelo CLIP-ViT-L-14 en formato FP8 usando NVIDIA Model Optimizer (ModelOpt) con el método post-training quantization (PTQ). La receta usa el checkpoint CLIP-ViT-L-14-laion2B-s32B-b82K como base.
¿Qué es NVIDIA Model Optimizer?
ModelOpt es una librería que combina técnicas state-of-the-art de optimización: cuantización, destilación, pruning, decoding especulativo y sparsity. Acepta modelos en formato Hugging Face, PyTorch u ONNX, y expone APIs en Python para combinar técnicas y producir checkpoints optimizados.
La librería soporta formatos de cuantización FP4, FP8, INT8 e INT4, además de algoritmos avanzados como SmoothQuant, AWQ, SVDQuant y Double Quantization. Funciona tanto en PTQ como en quantization-aware training (QAT).
¿Por qué CLIP es un buen caso de estudio?
CLIP (Contrastive Language-Image Pretraining), introducido por OpenAI en 2021, es un vision language model que aprende un espacio de embeddings compartido para imágenes y texto mediante aprendizaje contrastivo sobre grandes pares imagen-texto. Su encoder de texto se reusa como módulo de condicionamiento en pipelines text-to-image (Stable Diffusion) y text-to-video (AnimateDiff). Su encoder visual sirve como backbone en LLMs multimodales como LLaVA y modelos de percepción de vocabulario abierto como OWL-ViT.
¿Cómo se ejecuta el flujo de PTQ con ModelOpt?
La receta usa cuantización W8A8 (FP8 sobre pesos y activaciones), per-tensor estática, calibrada con el algoritmo AbsMax sobre 8.192 pares imagen-texto de MS-COCO. El código clave es directo:
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_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)
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)
q_model.save_pretrained(ckpt_path)
mtq.print_quant_summary(q_model)Después de mtq.quantize, todas las capas Linear de CLIP cargan quantizers de pesos y activaciones, pero los bloques de atención quedan intocados. Esto ocurre porque la atención multi-cabezal despacha a torch.nn.functional.scaled_dot_product_attention, una API funcional que el module walker de ModelOpt no puede interceptar por sí solo.
Para llevar la atención al alcance de la cuantización, hay que registrar un reemplazo cuantizado para CLIPAttention. Cada instancia se actualiza a _QuantAttention del plugin de diffusers, y dentro de su forward pass intercepta de forma transparente la llamada SDPA e inserta cuatro quantizers alrededor del kernel fusionado: q_bmm_quantizer, k_bmm_quantizer y v_bmm_quantizer envuelven los tensores Q/K/V proyectados antes de entrar al kernel, mientras que bmm2_output_quantizer envuelve el output del kernel (softmax @ V) antes de llegar a out_proj.
Para recuperar parte de la calidad, suele recomendarse desactivar algunos quantizers con mtq.disable_quantizer. La función toma como input el nombre de un módulo y, mediante regex, permite seleccionar qué capas excluir. En el ejemplo de NVIDIA se desactivan los quantizers de la capa patch_embedding, lo que mejora notablemente la calidad final.
Resultados en benchmarks
El checkpoint cuantizado se evaluó contra el baseline FP16 en tres tareas del CLIP_benchmark: clasificación zero-shot en CIFAR-100 e ImageNet-1k, y retrieval zero-shot en MS-COCO Captions. Según los resultados publicados, el modelo CLIP-FP8 muestra calidad comparable al CLIP-FP16. Cuando los quantizers se desactivan en la capa patch_embedding, el impacto sobre la calidad del modelo se vuelve "negligible".
¿Qué hace por dentro el flujo de PTQ?
Es importante entender que esta etapa trabaja con fake quantization: el tipo de dato del modelo no cambia. Los quantizers actúan como observadores que simulan los efectos de la cuantización mientras el modelo sigue corriendo en formato floating-point original.
El proceso opera en dos pasos clave:
- Recolección de estadísticas: los quantizers acumulan estadísticas de los tensores (mínimos, máximos) a medida que pasan los datos. Estas estadísticas se usan para calcular parámetros óptimos de cuantización como los scaling factors.
- Simulación de cuantización: los quantizers ejecutan una operación quantize-then-dequantize (QDQ) sobre los tensores. Solo simula el cómputo en baja precisión; las ganancias reales en velocidad y memoria recién aparecen al exportar el modelo a frameworks de despliegue como NVIDIA TensorRT.
El flujo completo de PTQ en ModelOpt sigue seis etapas: preparar la configuración de cuantización, calibrar con un batch representativo, ejecutar la fake quantization, evaluar contra el baseline, iterar si la brecha es inaceptable, y finalmente exportar y desplegar el checkpoint comprimido a TensorRT u ONNX.
QAT recupera la pérdida de calidad inducida por la cuantización mediante fine-tuning con estados de quantizer congelados. Es más intensivo en cómputo que PTQ pero puede mejorar más la calidad del modelo cuantizado. Para más detalles, NVIDIA referencia los ejemplos de ModelOpt.
