Gradient Checkpointing

El checkpointing de gradientes es una técnica de entrenamiento que ahorra memoria almacenando solo determinadas activaciones intermedias del paso hacia delante y recalculando las demás durante la retropropagación. También se conoce como checkpointing de activaciones y cambia más cómputo por un menor uso máximo de memoria. Pese a su nombre, la técnica guarda activaciones, no gradientes de parámetros ni archivos de modelo, por lo que resulta especialmente útil cuando los tensores de activación impiden que una red neuronal quepa en la memoria de la GPU disponible.

Cómo funciona el checkpointing de gradientes#

Durante un paso hacia delante estándar, una red neuronal calcula tensores intermedios llamados activaciones. El sistema de diferenciación automática conserva las activaciones necesarias para calcular los gradientes más adelante, tal y como se describe en la mecánica de autograd de PyTorch. Las redes profundas, los lotes grandes, las imágenes de alta resolución y las secuencias de entrada largas pueden hacer que estos tensores almacenados ocupen mucha memoria.

El checkpointing de gradientes divide la red en segmentos:

1. El paso hacia delante almacena las entradas o las activaciones límite de los segmentos seleccionados.
2. Se descartan las demás activaciones intermedias dentro de esos segmentos.
3. Durante el paso hacia atrás, se vuelve a ejecutar el paso hacia delante de cada segmento con checkpointing para reconstruir los valores que faltan.
4. Las activaciones reconstruidas se utilizan inmediatamente para calcular los gradientes.

La API de checkpointing de activaciones de PyTorch ofrece este comportamiento mediante torch.utils.checkpoint. Existen conceptos equivalentes en el checkpointing de gradientes y la rematerialización de JAX y el checkpointing de cintas de TensorFlow.

La ubicación de los checkpoints determina el equilibrio. En general, aplicar checkpointing a más regiones ahorra más memoria, pero repite más operaciones. Una ubicación selectiva en bloques que consumen muchas activaciones puede ofrecer un mejor equilibrio que recalcular toda la red.

Compromisos y técnicas relacionadas#

El checkpointing de gradientes suele dejar intactos los parámetros del modelo, los gradientes y los estados del optimizador. Su objetivo principal es la memoria de activaciones, y el entrenamiento se ralentiza porque algunos cálculos del paso hacia delante se ejecutan dos veces. El resultado exacto depende de la arquitectura, los límites de los checkpoints, la forma del lote y el hardware.

Se diferencia de varias técnicas relacionadas:

• La acumulación de gradientes procesa varios microlotes antes de actualizar los pesos del modelo, lo que crea un lote efectivo más grande sin cargar todas las muestras a la vez. En cambio, el checkpointing de gradientes reduce las activaciones que se conservan para cada microlote.
• La precisión mixta utiliza tipos de datos de menor precisión para ciertas operaciones. El flujo de trabajo de precisión mixta automática de PyTorch puede reducir el uso de memoria y acelerar las operaciones compatibles, mientras que el checkpointing añade cómputo deliberadamente.
• La reducción del tamaño del lote reduce el uso de memoria al procesar menos muestras juntas. El checkpointing puede permitir trabajar con un lote o una resolución de entrada mayores.
• Los checkpoints de entrenamiento guardan los pesos y el estado del optimizador para recuperar el entrenamiento o realizar inferencias más adelante. La guía de PyTorch sobre checkpoints de modelos describe este mecanismo de persistencia, que no está relacionado con el recálculo de activaciones.

El código con checkpointing debe ser funcionalmente coherente entre su ejecución original y la ejecución recalculada del paso hacia delante. El estado mutable, las transferencias entre dispositivos o la aleatoriedad no controlada dentro de una región con checkpointing pueden causar errores o gradientes incorrectos. En los Transformer de tipo decodificador, puede que también haya que desactivar una caché de clave-valor durante el entrenamiento con checkpointing, ya que el estado de inferencia almacenado en caché puede interferir en la reconstrucción del grafo del paso hacia delante.

Ejemplo en PyTorch#

El siguiente ejemplo aplica checkpointing a un bloque que consume mucha memoria durante un paso de entrenamiento:

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

Solo se conservan las entradas del bloque y la información límite necesaria; sus activaciones internas se reconstruyen durante loss.backward(). En proyectos reales, compara ejecuciones con y sin checkpointing mediante una medición como el máximo de memoria de GPU asignada en PyTorch.

Aplicaciones en el mundo real#

• Visión artificial de alta resolución: La segmentación médica, la detección aérea y la inspección industrial pueden requerir entrenar con imágenes grandes cuyos mapas de características consumen más memoria que los pesos del modelo. Aplicar checkpointing a etapas seleccionadas del backbone puede mantener la resolución de imagen o permitir más muestras por lote, en lugar de reducir drásticamente la resolución de las entradas.

• Entrenamiento de Transformer con secuencias largas: El almacenamiento de activaciones aumenta rápidamente a medida que crecen la longitud de la secuencia y el número de capas. Recalcular bloques Transformer puede permitir procesar contextos más largos o microlotes más grandes en el mismo acelerador. La guía de NVIDIA sobre el recálculo de activaciones muestra el recálculo completo y selectivo de capas Transformer.

Recomendaciones prácticas#

Usa el checkpointing de gradientes cuando el perfilado muestre que las activaciones, y no los parámetros ni el estado del optimizador, son lo que más memoria consume. Empieza por los bloques repetidos de mayor tamaño y compara la memoria máxima, el tiempo por iteración y el comportamiento en validación antes de ampliar su aplicación.

En las configuraciones documentadas de entrenamiento de Ultralytics YOLO, los primeros controles de memoria que conviene probar incluyen la precisión mixta automática, el tamaño de imagen y el tamaño físico del lote. La referencia de AutoBatch de Ultralytics explica cómo seleccionar automáticamente el tamaño del lote en función de la memoria de GPU disponible. Si el hardware local sigue siendo insuficiente, el entrenamiento en la nube de Ultralytics Platform ofrece GPU en la nube configurables para ejecuciones de entrenamiento gestionadas. El checkpointing de gradientes puede complementar estos controles cuando una arquitectura personalizada de PyTorch requiere una gestión más precisa de la memoria de activaciones.