Acumulación de gradientes: lotes efectivos sin perder la normalización correcta

Qué es

La acumulación de gradientes procesa varios microlotes antes de actualizar los parámetros del modelo. Sus gradientes contribuyen a un único paso del optimizador, lo que permite un lote efectivo mayor cuando un lote físico no cabe en la memoria del dispositivo. Las llamadas a backward de PyTorch suman los gradientes a los que ya tienen los parámetros hasta que estos se borran. [1][2]

Cómo funciona

Si los microlotes son del mismo tamaño y la pérdida es la media sobre ejemplos con el mismo peso, hay que dividir cada pérdida entre el número de microlotes de la ventana de acumulación. La suma de esos gradientes da entonces el gradiente de la pérdida media del conjunto. Si los tamaños son desiguales, hay que ponderar por el número real de ejemplos que contribuyen. En las pérdidas por token, se cuentan los tokens válidos, es decir, los no enmascarados; los pesos por clase u otras reducciones pueden requerir un denominador distinto. Promediar sin más las medias de los microlotes cambia el objetivo. Estas reglas de normalización se deducen de derivar la media ponderada que se pretende optimizar.

Flujo de trabajo práctico

Hay que borrar los gradientes al inicio de cada ventana, ejecutar forward y backward para cada microlote normalizado y, después, realizar un único paso del optimizador. A continuación se borran los gradientes para la ventana siguiente. Los pases forward independientes normalmente no requieren conservar los grafos de cómputo anteriores. [2] La última ventana, que queda incompleta, debe tratarse de forma explícita. Con precisión mixta, se mantiene constante la escala del gradiente durante toda la acumulación, se aplica unscale una sola vez tras el último backward y después se hace clipping del gradiente combinado antes del paso del optimizador. El scaler se actualiza en ese límite. [1] Un scheduler configurado por actualización del optimizador solo debe avanzar cuando se produce una actualización.

Ejemplo ilustrativo

Supongamos que un dispositivo procesa cuatro microlotes de ocho ejemplos. El lote efectivo contiene 32 ejemplos, así que la media de cada microlote recibe un peso de 1/4. Una última ventana con 8, 8, 8 y 2 ejemplos contiene 26 ejemplos, y los pesos correspondientes son 8/26, 8/26, 8/26 y 2/26. Se trata de aritmética, no del resultado medido de un entrenamiento.

Límites y verificación

La normalización por lotes (batch normalization) usa las estadísticas de cada microlote físico, por lo que la acumulación no reproduce automáticamente un único lote grande. [3] Las operaciones aleatorias y el orden de la suma en coma flotante también pueden alterar los resultados. Aplicar clipping a cada microlote por separado suele cambiar la actualización combinada. El método reduce el almacenamiento simultáneo de activaciones; no elimina los pesos del modelo ni los estados del optimizador, y los pases adicionales pueden aumentar el tiempo de ejecución.

Antes de un entrenamiento largo, conviene comparar una actualización acumulada con una actualización de lote combinado en un modelo pequeño y determinista, sin capas que dependan del lote. Hay que revisar la normalización de la pérdida, las diferencias entre gradientes dentro de una tolerancia adecuada, el número de actualizaciones y el comportamiento de la ventana parcial. Son comprobaciones recomendadas, no pruebas realizadas para este artículo.

Fuentes

[1] PyTorch: Automatic Mixed Precision examples [2] PyTorch: torch.Tensor.backward [3] PyTorch: BatchNorm1d