Cómo reducir el tamaño de un modelo LLM usando quantization

Qué es la quantization y por qué te importa

La quantization consiste en representar los pesos y activaciones del modelo con menos bits (8-bit, 4-bit, etc.) en lugar de usar formatos flotantes de 16/32 bits. Al hacerlo se consiguen:

• Modelos 2-8× más pequeños (caben en GPUs de menor VRAM o incluso en CPU).
• Menor consumo de memoria y energía.
• Inferencia más rápida gracias a operaciones de menor precisión.

El precio a pagar es una posible (y normalmente pequeña) pérdida de precisión, que puede mitigarse con las técnicas correctas.

Requisitos previos

• Conocimientos básicos de Python y PyTorch.
• Una GPU con soporte CUDA (opcional pero recomendado).
• Python ≥ 3.8.
• Paquetes: transformers, bitsandbytes, accelerate.

pip install -U transformers bitsandbytes accelerate

3 caminos de quantization

En esta guía usaremos PTQ de 4-bit con bitsandbytes porque ofrece el mejor balance entre simplicidad y calidad.

Paso a paso: cuantizar en 4-bit con bitsandbytes

# quantize_llm.py
# Cuantiza un modelo existente a 4 bits de forma inmediata.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 1) Identificador del modelo en Hugging Face
model_id = 'facebook/opt-1.3b'  # ~2.7 GB en FP16

# 2) Cargamos el modelo directamente en 4-bit
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_4bit=True,      # Activa la cuantización
    device_map='auto',      # Aprovecha todas las GPUs/CPU disponibles
    quantization_config={
        'bnb_4bit_compute_dtype': torch.bfloat16,  # Precisión de cálculo
        'bnb_4bit_use_double_quant': True,         # Segunda cuantización para más calidad
        'bnb_4bit_quant_type': 'nf4'               # Normal-Float4, mejor que int4
    }
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# 3) Inferencia de prueba
prompt = 'Dame 3 consejos para aprender machine learning:'
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Resultado típico:

• Tamaño en disco: de 2.7 GB a ~700 MB.
• VRAM utilizada: ≈4 GB al inferir.
• Latencia: 20-40 % más rápida (dependiendo del hardware).

Verificar la degradación de calidad

Cuantizar sin medir es apostar a ciegas. Dos métricas rápidas:

1. Perplejidad (modelos de lenguaje)
2. Exactitud en tareas de benchmark (ej. MMLU, GSM8K)

Ejemplo con evaluate:

from evaluate import load
import torch

perplexity = load('perplexity', module_type='metric')
score = perplexity.compute(
    model_id='facebook/opt-1.3b',
    add_start_token=True,
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
)
print('Perplejidad original:', score['perplexities'][0])

score_q = perplexity.compute(
    model=model,  # modelo cuantizado
    tokenizer=tokenizer,
    add_start_token=True,
)
print('Perplejidad 4-bit:', score_q['perplexities'][0])

Regla práctica: si la métrica empeora menos de 5 %, la cuantización es “gratis” para producción.

Ajustes finos con Quantization Aware Training

Si el modelo pierde demasiado rendimiento:

# Idea general (no runnable): inserta capas fake-quant durante el fine-tuning
from torch.ao.quantization import get_default_qat_qconfig, prepare_qat, convert

model.train()
model.qconfig = get_default_qat_qconfig('fbgemm')
model_prepared = prepare_qat(model)
# ... continúa tu entrenamiento ...
model_int8 = convert(model_prepared)  # Modelo listo en 8-bit

QAT requiere más tiempo de entrenamiento pero recupera casi toda la precisión original.

Otros formatos y herramientas populares

• GPTQ / Auto-GPTQ – apropiado para despliegue en CPU.
• GGML / GGUF – ideal para dispositivos móviles o Raspberry Pi.
• TensorRT-LLM – quantization + optimización para GPU NVIDIA.
• Intel Neural Compressor – simplifica PTQ y QAT en hardware Intel.

Errores comunes y cómo evitarlos

1. Mezclar capas de precisión distinta sin fijar torch.autocast: produce overflows.
2. Olvidar el to(model.device) en la tokenización: el modelo funciona pero más lento.
3. Cuantizar embeddings de entrada en tareas de clasificación pequeñas: suelen perder precisión desproporcionada; usa formato mixto (embeddings en FP16).

Checklist rápida antes de subir a producción

• [ ] La métrica clave baja < 5 %.
• [ ] No aparecen mensajes de overflow/underflow en logs.
• [ ] El modelo cabe en la memoria estimada con margen del 10 %.
• [ ] Latencia medida con torch.cuda.synchronize() es aceptable.
• [ ] Se guarda copia del modelo original por si hay que revertir.

Recursos recomendados

• Paper “SmoothQuant” (OpenAI & Microsoft) – mejora PTQ en Transformers.
• Guía oficial de bitsandbytes: https://github.com/TimDettmers/bitsandbytes
• Cheat-sheet de Hugging Face para quantization: https://huggingface.co/docs/transformers/perf_quantization

¡Listo! Ahora tienes un LLM más ligero que cabe en tu hardware y responde más rápido, sin sacrificar calidad de manera significativa.