Cómo Solucionar el Error 'CUDA Out of Memory' en PyTorch al Trabajar con LLMs: Guía Práctica de Optimización

1. Anatomía del Consumo de VRAM en Grandes Modelos de Lenguaje (LLMs)

El error torch.cuda.OutOfMemoryError es la pesadilla más recurrente de cualquier ingeniero de machine learning y desarrollador que despliega modelos en entornos locales o clusters de GPU. Para erradicar este problema de raíz en lugar de recurrir al ensayo y error, es imprescindible comprender con exactitud matemática cómo se distribuye cada gigabyte de memoria VRAM en una tarjeta gráfica durante el ciclo de vida de un modelo de lenguaje.

La memoria requerida no se limita al tamaño de los pesos del modelo en disco. En un pipeline de Deep Learning, la VRAM total consumida se divide en cuatro grandes categorías:

• Pesos del modelo (Model Weights): En precisión de coma flotante tradicional de 32 bits (FP32), cada parámetro requiere 4 bytes de memoria. Por tanto, un modelo de 7 mil millones de parámetros (7B) requiere 28 GB de VRAM solo para cargar sus pesos. En FP16 o BFloat16 (2 bytes por parámetro), el consumo se reduce a 14 GB. Si se cuantiza a 4 bits (0.5 bytes por parámetro), se reduce a aproximadamente 3.5 - 4 GB.
• Estados del Optimizador (Optimizer States): Durante el entrenamiento con optimizadores basados en momentum como AdamW, se almacenan copias adicionales de los gradientes, la media móvil del gradiente (primer momento) y la media móvil de los cuadrados de los gradientes (segundo momento). En FP32, AdamW exige 8 bytes adicionales por parámetro (12 bytes si se suman los gradientes en FP32). Esto significa que entrenar un modelo 7B en FP32 completo puede consumir más de 80 GB de VRAM únicamente en pesos y optimizador.
• Gradientes (Gradients): Requieren el mismo espacio que los pesos del modelo (2 bytes por parámetro en FP16/BF16, o 4 bytes en FP32).
• Activaciones intermedias (Activations & KV Cache): Es la memoria consumida por los tensores intermedios calculados durante el paso hacia adelante (forward pass) que deben retenerse para calcular la regla de la cadena durante el paso hacia atrás (backward pass). El volumen de activaciones escala de forma cuadrática respecto a la longitud de la secuencia de contexto en arquitecturas Transformer estándar debido a la matriz de atención (Seq_Len × Seq_Len). En inferencia, esta categoría se traduce en el KV Cache (Key-Value Cache), que crece linealmente con la cantidad de tokens generados y concurrentes.

2. Diferencia entre Memoria Reservada y Memoria Asignada: El Problema de la Fragmentación

Cuando observas la salida típica del error de PyTorch:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB.
GPU 0 has a total capacity of 23.69 GiB of which 180.12 MiB is free.
Including non-PyTorch memory, this process has 23.51 GiB memory in use.
Of the allocated memory 19.80 GiB is allocated by PyTorch, and 2.80 GiB is reserved by PyTorch but unallocated.

Fíjate en la distinción crucial entre allocated memory (memoria asignada) y reserved memory (memoria reservada). PyTorch no solicita memoria al driver de NVIDIA mediante cudaMalloc en cada pequeña operación de tensores porque esa llamada al sistema operativo tiene una latencia inadmisible. En su lugar, PyTorch utiliza un caching memory allocator.

Este asignador reserva grandes bloques de VRAM de la GPU y los divide internamente para servir a los tensores creados. Cuando un tensor se destruye o sale del alcance en Python, PyTorch marca ese espacio como disponible en su piscina interna, pero no se lo devuelve a la GPU. Si la piscina interna sufre de fragmentación de memoria (es decir, existen muchos huecos pequeños libres pero ningún bloque contiguo del tamaño necesario para el siguiente tensor de atención), el asignador intenta pedir un nuevo bloque continuo a CUDA. Si la VRAM física restante es menor que el bloque requerido, salta el error de OOM, a pesar de que la suma de pequeños huecos libres habría sido suficiente.

La solución definitiva a la fragmentación: expandable_segments

A partir de las versiones recientes de PyTorch, se introdujo una bandera de asignación que resuelve este problema de manera radical aprovechando la memoria virtual de CUDA:

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

Al activar expandable_segments:True, PyTorch permite que los segmentos de memoria reservados crezcan de forma modular en el espacio de direcciones virtuales sin requerir bloques físicos contiguos en el momento de la asignación. Esto previene hasta un 90% de los fallos por fragmentación en modelos con secuencias largas.

3. Estrategias Prácticas de Optimización para Eliminar el Error OOM

A. Gradient Accumulation: Mantén el Batch Size Efectivo sin VRAM Extra

Si tu modelo no cabe con un batch size de 16 o 32, no reduzcas simplemente el tamaño del lote perdiendo estabilidad en el descenso de gradiente. Utiliza acumulación de gradientes para simular lotes grandes dividiéndolos en microlotes de 1 o 2 ejemplos:

# Configuración de Gradient Accumulation en PyTorch
accumulation_steps = 8  # Batch efectivo = micro_batch (2) * 8 = 16
optimizer.zero_grad()

for step, batch in enumerate(dataloader):
    inputs, labels = batch['input_ids'].cuda(), batch['labels'].cuda()
    
    # Forward pass con precisión mixta (autocast)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        outputs = model(inputs, labels=labels)
        loss = outputs.loss / accumulation_steps
    
    # Backward pass acumulativo
    loss.backward()
    
    if (step + 1) % accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()

B. Gradient Checkpointing (Activation Checkpointing)

El almacenamiento de activaciones en secuencias largas (por ejemplo, contextos de 4.096 a 32.768 tokens) domina por completo el uso de VRAM. Gradient Checkpointing descarta las activaciones intermedias de la mayoría de las capas Transformer durante el forward pass y solo guarda las activaciones en los límites de cada bloque. Durante el backward pass, recalcula localmente el forward de ese bloque específico.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# Activar Gradient Checkpointing nativo en Hugging Face
model.gradient_checkpointing_enable()

C. Aceleración con FlashAttention-2

El cálculo tradicional de la función softmax en la capa de atención de los Transformers calcula la matriz completa de atención S = Q @ K.T, la almacena en VRAM global (HBM), calcula el softmax y multiplica por V. Para una secuencia de longitud $N$, esta matriz tiene tamaño $N imes N$, exigiendo una cantidad inmensa de memoria.

FlashAttention-2 reestructura el cálculo dividiendo las matrices en bloques pequeños (tiling) que caben íntegramente en la memoria SRAM ultrarrápida del chip de la GPU, calculando el softmax de forma incremental (online softmax) sin materializar jamás la matriz $N imes N$ en la HBM. El resultado es un consumo de memoria lineal respecto a la secuencia y una velocidad de cómputo hasta 3 veces superior.

# Carga con FlashAttention-2 habilitado en Transformers
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map="auto"
)

D. Limpieza rigurosa del grafo en Inferencia

Uno de los errores más elementales que cometen los desarrolladores en producción es ejecutar inferencia sin desactivar el rastreador de autodiferenciación de PyTorch. El decorador clásico torch.no_grad() es adecuado, pero torch.inference_mode() es aún superior, ya que desactiva por completo el rastreo de versiones de tensores, reduciendo la sobrecarga de memoria interna:

import torch

@torch.inference_mode()
def generar_respuesta(prompt_tokens):
    # En este bloque no se construyen nodos de autograd
    output = model.generate(
        prompt_tokens,
        max_new_tokens=256,
        use_cache=True  # Optimiza el KV-cache para generación autoregresiva
    )
    return output

4. Depuración Avanzada de VRAM con torch.cuda.memory_summary

Cuando un script falle o desees perfilar el punto exacto donde la memoria experimenta su pico más alto, inserta un volcado del estado del asignador de CUDA justo antes y después del paso crítico:

print(torch.cuda.memory_summary(device=None, abbreviated=False))

Este informe detalla:

1. Active Memory: Memoria que actualmente contiene datos de tensores vivos.
2. Peak Active Memory: El punto máximo de memoria activa alcanzado desde el inicio del proceso. Si este valor se aproxima a la capacidad física de la GPU, el OOM es inminente.
3. Alloc Retries: Número de veces que el asignador tuvo que pausar y compactar la piscina de memoria para encontrar un espacio contiguo. Si este número es elevado, sufres de severa fragmentación.

5. Lecturas Recomendadas en The Calata para Optimizar tu Infraestructura de IA

Para profundizar en arquitecturas de alto rendimiento y evitar limitaciones de hardware en tus proyectos de Inteligencia Artificial, te sugerimos explorar nuestras siguientes guías técnicas especializadas:

• Descubre las ventajas comparativas de formatos de compresión en nuestra guía sobre Cuantización GGUF vs EXL2 vs AWQ para Inferencia Local.
• Aprende a entrenar modelos con mínima memoria en el tutorial de Fine-Tuning con LoRA y QLoRA Usando Unsloth.
• Si sirves modelos en producción con alta concurrencia, revisa la comparativa entre vLLM vs TGI para Motores de Inferencia.
• Optimiza tus tiempos de respuesta reduciendo el retraso de salida con la guía de Cómo Reducir la Latencia TTFT en APIs de LLMs.