¿Qué es el checkpointing de gradiente?

El checkpointing de gradiente es una técnica de optimización de memoria durante el entrenamiento de redes neuronales profundas. La idea es no guardar todas las activaciones intermedias durante la pasada hacia adelante, sino solo algunos puntos de control (checkpoints). Durante la retropropagación, las activaciones faltantes se reconstruyen mediante el recálculo de parte de la pasada hacia adelante.

Esto permite reducir significativamente el consumo de memoria, lo cual es especialmente importante al entrenar modelos muy profundos o con recursos limitados de GPU. Sin embargo, esto aumenta el tiempo de entrenamiento, ya que algunas computaciones deben repetirse.

Ejemplo de uso en PyTorch:

from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # parte compleja del modelo
    return x * x

output = checkpoint(custom_forward, input_tensor)

Por lo tanto, el checkpointing de gradiente es un compromiso entre tiempo y memoria, que permite entrenar modelos grandes con recursos limitados.