Gradient Accumulation

La acumulación de gradientes es una técnica de entrenamiento que ejecuta varios lotes pequeños, llamados microlotes, en el pase hacia delante y la retropropagación, suma sus gradientes y solo entonces actualiza los pesos del modelo. La actualización aproxima un paso sobre un lote del tamaño combinado de todos los microlotes, mientras que la GPU solo almacena las activaciones de un microlote cada vez. Esto permite entrenar con un dispositivo limitado por memoria y usar el tamaño de lote que requiere la receta, sin reducir la resolución de imagen ni el tamaño del modelo. El tamaño efectivo del lote es el tamaño del microlote multiplicado por el número de pasos de acumulación y el número de dispositivos; así, microlotes de 4 imágenes acumulados durante 8 pasos en 4 GPU equivalen a un lote de 128.

Cómo funciona la acumulación de gradientes#

En PyTorch y la mayoría de los demás marcos, cada llamada a backward() suma nuevos gradientes a los que ya están almacenados en los parámetros, en lugar de sustituirlos. La acumulación aprovecha este comportamiento directamente: ejecuta los pases hacia delante y hacia atrás para cada microlote, llama al optimizador cuando se alcanza el número establecido de microlotes, borra los gradientes y empieza la siguiente ventana:

La mayoría de las funciones de pérdida calcula la media de las muestras de un lote, así que se divide la pérdida de cada microlote por el número de pasos para que la suma de gradientes se convierta en una media. Con este ajuste, cuatro microlotes de 8 producen el mismo gradiente que un lote de 32:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(16, 1)
x, y = torch.randn(32, 16), torch.randn(32, 1)
loss_fn = nn.MSELoss()

# Reference: one batch of 32 samples
loss_fn(model(x), y).backward()
batch_grad = model.weight.grad.clone()
model.zero_grad()

# Four micro-batches of 8: backward() adds into .grad, and dividing each loss by 4 turns the sum into a mean
steps = 4
for x_micro, y_micro in zip(x.chunk(steps), y.chunk(steps)):
    (loss_fn(model(x_micro), y_micro) / steps).backward()

print("Matches the full-batch gradient:", torch.allclose(model.weight.grad, batch_grad, atol=1e-6))
print(f"Largest difference: {(model.weight.grad - batch_grad).abs().max():.1e}")

Imprime Matches the full-batch gradient: True con una diferencia máxima de 1.6e-07, debida al redondeo de coma flotante al sumar los mismos números en un orden distinto.

En qué se diferencia la acumulación de un lote grande#

La equivalencia se cumple para pérdidas que calculan la media de muestras independientes, pero se rompe en tres casos habituales:

• Normalización por lotes: Calcula la media y la varianza de cada microlote, por lo que las estadísticas de 4 imágenes son más ruidosas que las de 128.
• Secuencias de longitud variable: Cuando un modelo de lenguaje calcula la media de la pérdida sobre los tokens de cada microlote, los microlotes con pocos tokens tienen el mismo peso que los más largos.
• Recorte de gradientes: Debe aplicarse una sola vez al gradiente acumulado, justo antes del paso del optimizador; recortar cada microlote cambia el resultado.

Hugging Face y Unsloth solucionaron en octubre de 2024 el desajuste en el recuento de tokens para el ajuste fino de un modelo de lenguaje de gran tamaño (LLM), dividiendo la suma de las pérdidas por el número total de tokens de toda la ventana de acumulación. La acumulación tampoco acelera el proceso: el mismo número de muestras sigue pasando por la red, un microlote cada vez, y los microlotes pequeños aprovechan la GPU con menos eficiencia. En el entrenamiento distribuido, DistributedDataParallel de PyTorch ofrece un contexto no_sync() que omite la operación all-reduce de gradientes en todos los microlotes salvo el último, lo que reduce el volumen de comunicación. Un lote efectivo de mayor tamaño suele requerir una tasa de aprendizaje más alta: investigadores de Facebook entrenaron ResNet-50 en ImageNet en una hora con lotes de 8.192 imágenes en 256 GPU, escalando la tasa de aprendizaje linealmente con el tamaño del lote y aumentándola gradualmente al principio.

Técnicas relacionadas#

• Puntos de control de gradiente: Recalcula las activaciones durante la retropropagación en lugar de almacenarlas, lo que ahorra memoria en cada microlote a costa de más cómputo.
• Precisión mixta: Almacena las activaciones en formatos de 16 bits, lo que deja espacio para microlotes más grandes. El escalador de pérdidas desescala los gradientes una sola vez, después del último microlote.
• Paralelismo de canalización: Divide cada lote en microlotes que pasan por etapas del modelo en distintas GPU y acumula sus gradientes antes de ejecutar un único paso del optimizador.

El entrenamiento de Ultralytics YOLO26 acumula gradientes de forma predeterminada: en modo de entrenamiento, el entrenador ejecuta el optimizador cada round(nbs / batch) lotes, por lo que los valores predeterminados batch=16 y el tamaño de lote nominal nbs=64 dan cuatro lotes por actualización, y batch=8 da ocho. El número aumenta gradualmente desde uno durante el calentamiento, y la decaída de pesos se escala mediante batch × accumulation / nbs, así que los hiperparámetros predeterminados para un lote de 64 imágenes siguen siendo válidos en una GPU que solo puede almacenar unas pocas imágenes cada vez.