Asignación eficiente de recursos de GPU

La memoria de la GPU es el cuello de botella con el que me topo más a menudo en las cargas de trabajo de machine learning. Ya sea que estés haciendo el fine-tuning de un modelo de 7B de parámetros o sirviendo un modelo de 70B en producción, entender cómo se consume la VRAM —y cómo exprimirla al máximo— marca la diferencia entre un workflow que funciona y uno que falla con un error de CUDA out of memory a los diez minutos.

Requerimientos de VRAM

La huella de memoria depende del conteo de parámetros y de la precisión numérica con la que se almacenen:

Memoria (bytes) = Parámetros × Bytes por Parámetro

Precisión	Bytes por Parámetro	Modelo 7B	Modelo 13B	Modelo 70B
FP32	4	28 GB	52 GB	280 GB
FP16/BF16	2	14 GB	26 GB	140 GB
INT8	1	7 GB	13 GB	70 GB
INT4	0.5	3.5 GB	6.5 GB	35 GB

Invalid viz:graph block: expected JSON like {"name":"my_graph"}.

Memoria de pesos por precisión y tamaño de modelo. Cada paso de precisión reduce la VRAM aproximadamente a la mitad; un modelo de 70B en INT4 (35 GB) finalmente cabe en una sola GPU de clase consumidor de 40 GB.

Esa tabla solo tiene en cuenta los pesos. Durante el entrenamiento, también necesitas memoria para otras cosas. Los estados del optimizador (Adam guarda dos copias adicionales de cada parámetro: momentum y varianza), aproximadamente el doble del tamaño del modelo en FP32. Los gradientes, que requieren una copia completa de todos los parámetros. Las activaciones, que son valores intermedios almacenados para el backward pass y que escalan con el batch size y la longitud de la secuencia. Y el KV cache durante la inferencia, que crece con la longitud de la secuencia y el batch size.

Una regla general para el entrenamiento con Adam en precisión mixta: la VRAM total termina siendo alrededor de 4 veces el tamaño de los pesos del modelo. Un modelo de 7B en BF16 (14 GB de pesos) necesita unos 56-60 GB para entrenarse con un batch size razonable.

Batch size vs memoria

El batch size tiene una relación casi lineal con la memoria de activaciones. Duplicar el batch duplica las activaciones. Esto genera una tensión: los batches más grandes ofrecen un mejor throughput y gradientes más estables, pero es posible que no tengas suficiente VRAM.

La jugada práctica es encontrar el batch size más grande que quepa y luego usar acumulación de gradientes para alcanzar un batch size efectivo mayor.

# En lugar de batch_size=32 (que podría causar OOM):
# Usar batch_size=8 con gradient_accumulation_steps=4
# Batch size efectivo = 8 × 4 = 32, memoria pico = batch_size de 8

from transformers import TrainingArguments

training_args = TrainingArguments(
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    # batch size efectivo: 8 * 4 = 32
    fp16=True,
    output_dir="./output",
)

La acumulación de gradientes ejecuta múltiples pasadas forward-backward con un batch size más pequeño, acumulando los gradientes antes de realizar un único paso del optimizador. El resultado matemático es el mismo que con un batch grande, pero la memoria pico coincide con el batch por paso más pequeño.

Entrenamiento en precisión mixta (Mixed precision)

La precisión mixta utiliza FP16 o BF16 para las pasadas forward y backward, mientras mantiene una copia maestra de los pesos en FP32 para la actualización del optimizador. Esto reduce la memoria a la mitad aproximadamente y mejora el throughput en GPUs modernas con tensor cores.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()

    with autocast(dtype=torch.bfloat16):
        outputs = model(batch["input_ids"])
        loss = loss_fn(outputs, batch["labels"])

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Sobre FP16 vs BF16: BF16 se ha convertido en el estándar para entrenamiento porque mantiene el mismo rango de exponente que FP32 (8 bits), lo que evita los problemas de overflow y underflow que hacen que FP16 sea tedioso. Las complicaciones del escalado de pérdida (loss-scaling) que necesitas con FP16 simplemente desaparecen con BF16. Si tu GPU lo soporta (A100, H100 o más reciente), usa BF16.

Estrategias multi-GPU

Cuando una sola GPU no es suficiente —ya sea porque el modelo no cabe o porque el entrenamiento es demasiado lento— tienes tres estrategias principales de paralelismo.

Paralelismo de datos (Data parallelism)

Es la más simple. Replica el modelo completo en cada GPU y divide los datos. Cada GPU ejecuta un mini-batch diferente, calcula los gradientes y se sincroniza mediante all-reduce.

# PyTorch DistributedDataParallel
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model.to(local_rank), device_ids=[local_rank])

Úsalo cuando el modelo quepa en una sola GPU y quieras más throughput. Esta es la opción por defecto y tu primera elección.

La limitación es que cada GPU debe contener una copia completa del modelo. Para un modelo de 70B en BF16 (140 GB), el paralelismo de datos en GPUs A100 de 80 GB no es una opción.

Paralelismo de modelo o de tensores (Tensor parallelism)

Divide capas individuales entre varias GPUs. Una multiplicación de matrices grande Y = XW puede particionarse por columnas entre las GPUs; cada una calcula una parte de Y y luego se combinan los resultados.

Úsalo cuando el modelo sea demasiado grande para una sola GPU y quieras minimizar la latencia. El paralelismo de tensores mantiene la profundidad del pipeline en 1 (sin burbujas de pipeline), pero necesita interconexiones de gran ancho de banda como NVLink, porque cada capa implica comunicación.

Paralelismo de pipeline (Pipeline parallelism)

Asigna diferentes capas a diferentes GPUs. La GPU 0 ejecuta las capas 0-15, la GPU 1 las capas 16-31, y así sucesivamente. Los micro-batches fluyen a través del pipeline.

Úsalo cuando el modelo sea demasiado grande para una sola GPU y tengas muchas GPUs. El paralelismo de pipeline tiene menos sobrecarga de comunicación que el paralelismo de tensores (solo las activaciones se mueven entre etapas), pero el costo son las burbujas de pipeline: tiempo de inactividad cuando una GPU espera a la etapa anterior.

Combinación de estrategias

El entrenamiento en producción de modelos grandes suele combinar las tres:

Entrenamiento típico de un modelo 70B en 64 GPUs:
├── 8 nodos × 8 GPUs cada uno
├── Paralelismo de tensores: 8 vías dentro de cada nodo (NVLink)
├── Paralelismo de pipeline: 2 vías entre pares de nodos
└── Paralelismo de datos: 4 vías entre los grupos restantes

DeepSpeed, Megatron-LM y FSDP (Fully Sharded Data Parallel) gestionan la complejidad de combinar estas técnicas. FSDP ha sido la herramienta que más he usado últimamente porque distribuye (shards) los estados del optimizador, los gradientes y los parámetros entre las GPUs simultáneamente; básicamente, es paralelismo de datos más ahorro de memoria al fragmentar el estado del modelo.

Gestión práctica de memoria CUDA

Más allá de las estrategias de alto nivel, un puñado de técnicas ayudan en el día a día.

Activation checkpointing

En lugar de almacenar cada activación intermedia para el backward pass, vuelve a calcularlas sobre la marcha durante la retropropagación (backprop). Esto supone un 30% más de tiempo de cómputo a cambio de un 60-70% menos de memoria de activaciones.

from torch.utils.checkpoint import checkpoint

# En lugar de:
# output = self.expensive_layer(input)

# Usar:
output = checkpoint(self.expensive_layer, input, use_reentrant=False)

La mayoría de los frameworks de entrenamiento exponen esto como un flag simple:

# HuggingFace Transformers
training_args = TrainingArguments(
    gradient_checkpointing=True,
    # ...
)

Monitoreo de VRAM

No puedes diagnosticar problemas de memoria sin visibilidad. nvidia-smi está bien para un vistazo rápido, pero para un análisis detallado:

import torch

# Asignación actual
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Reserved:  {torch.cuda.memory_reserved() / 1e9:.2f} GB")

# Habilitar snapshot de memoria para depuración
torch.cuda.memory._record_memory_history()
# ... ejecutar tu código ...
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
# Visualizar en: pytorch.org/memory_viz

Limpieza de caché

El asignador de CUDA de PyTorch pone en caché la memoria liberada para reutilizarla. Normalmente esto es bueno, pero a veces puede bloquear la disponibilidad de memoria para asignaciones de diferentes tamaños:

torch.cuda.empty_cache()  # Libera la memoria en caché de vuelta a CUDA

empty_cache() no libera tensores que aún están referenciados. Si tienes una fuga de memoria (memory leak), casi siempre se debe a referencias de Python que mantienen vivos los tensores, no al asignador de CUDA.

Configuración de límites de memoria

En entornos de GPU compartidos, puedes limitar el uso de VRAM de un proceso:

# Limitar a 40 GB en una GPU de 80 GB
torch.cuda.set_per_process_memory_fraction(0.5, device=0)

Un flujo de decisión

Cuando te encuentres con un muro de memoria de GPU, sigue esta progresión en orden:

Reduce la precisión. Cambia a BF16/FP16 si no lo has hecho (ahorro de 2x en memoria).
Habilita activation checkpointing. Intercambia cómputo por memoria.
Reduce el batch size e incrementa la acumulación de gradientes. El batch size efectivo se mantiene igual.
Usa FSDP o DeepSpeed ZeRO para fragmentar (shard) el estado del modelo entre las GPUs.
Aplica cuantización para inferencia (INT8 o INT4).
Usa paralelismo de pipeline o de tensores para dividir el modelo en sí.
Renta más GPUs o más grandes. A veces, la respuesta correcta es más hardware.

Cada paso añade complejidad, así que empieza por el más sencillo que resuelva el problema. Muchos proyectos reales de ML se pueden manejar con una sola GPU, precisión mixta y activation checkpointing. No siempre necesitas un clúster multi-nodo.