Gradient checkpointing

Qué es

Un paso de entrenamiento tiene dos mitades. En el forward pass, cada capa convierte su entrada en una salida, produciendo un conjunto de resultados intermedios por el camino: las activaciones. Luego, el backward pass recorre las capas a la inversa para calcular los gradientes, y los gradientes de cada capa necesitan las activaciones de esa capa. El enfoque sencillo las mantiene todas en memoria hasta que el backward pass llega a ellas. Con ejemplos largos, las activaciones pueden ocupar más memoria que los propios pesos del modelo.

El gradient checkpointing (puntos de control de gradientes), también llamado puntos de control de activaciones o recálculo, conserva solo algunas activaciones, los puntos de control, y descarta el resto. Cuando la propagación hacia atrás necesita una activación descartada, vuelve a ejecutar la propagación hacia adelante desde el punto de control más cercano para recrearla. Chen et al. (2016) demostraron que esto reduce la memoria de activaciones para una red de n capas al orden de √n, a cambio de aproximadamente una propagación hacia adelante adicional por lote. En sus experimentos, una red de 1000 capas pasó de 48 GB a 7 GB con un aumento de aproximadamente el 30 % en el tiempo de ejecución.

Cómo funciona en un modelo de lenguaje

En las bibliotecas de transformadores, los puntos de control suelen situarse en los límites de las capas. Se conserva la entrada de cada capa; todo lo calculado dentro de la capa se descarta y se reconstruye cuando la propagación hacia atrás llega a ese punto. Por lo tanto, en cualquier momento dado solo existe el conjunto completo de activaciones de una capa, en lugar de todas ellas.

activations per GPT-style layer, 16-bit, efficient attention   (Korthikanti et al., 2022)
  ≈ 34 × batch × tokens × hidden size   bytes

illustrative 2B-class model: hidden size 2,048 · 28 layers · batch 2 × 1,024 tokens
  all kept                 34 × 2 × 1,024 × 2,048 × 28   ≈ 4.0 GB
  checkpointed, inputs      2 × 2 × 1,024 × 2,048 × 28   ≈ 0.23 GB
  + one layer rebuilt at a time                          ≈ 0.14 GB

La cantidad de puntos de control que se conservan es ajustable. Conservar más significa menos que recalcular pero más memoria; conservar menos significa lo contrario. Establecer puntos de control en cada límite de capa es la configuración predeterminada habitual porque es simple y ya elimina la mayor parte de la memoria de activaciones: lo que queda es una pequeña entrada por capa más la capa que se está reconstruyendo.

El costo es computacional. La propagación hacia adelante se ejecuta efectivamente dos veces: una vez de forma normal y otra pieza por pieza durante la propagación hacia atrás. La propagación hacia atrás ya cuesta aproximadamente el doble que la propagación hacia adelante, por lo que el trabajo adicional es, como máximo, alrededor de un tercio de un paso; en la práctica, el entrenamiento se ejecuta entre un 20 y un 30 % más lento. El resultado no cambia: los valores recalculados son los mismos, por lo que la curva de pérdida coincide con la de una ejecución sin puntos de control, salvo por pequeñas diferencias numéricas.

Por qué los vocabularios grandes hacen que importe más

Al final del modelo, cada posición en cada ejemplo obtiene una puntuación para cada token del vocabulario: los logits. Su tamaño no depende de la profundidad del modelo, sino únicamente de cuántas posiciones se entrenan y del tamaño del vocabulario. Los modelos más recientes tienen vocabularios de 150 000 a 250 000 tokens, por lo que los logits por sí solos pueden ocupar gigabytes.

logits = batch × tokens × vocabulary size × bytes per number

2 × 1,024 × 248,000 × 4 bytes (32-bit)   ≈ 2.0 GB   before their gradients

Cómo usarlo bien

• Recurra a él antes de recortar datos. Activar los puntos de control cuesta tiempo; reducir la longitud máxima recorta los extremos de los ejemplos largos. Cuando la memoria es el límite, la ralentización suele ser la mejor alternativa.
• Utilice la memoria que libera. Un lote más grande o una longitud máxima mayor pueden recuperar parte de la velocidad perdida por época.
• Combínelo con QLoRA para los modelos más grandes que pueda alojar una máquina: QLoRA reduce los pesos, mientras que los puntos de control reducen las activaciones.
• Conozca lo que no modifica. Los pesos, los gradientes y el estado del optimizador mantienen el mismo tamaño, y no tiene ningún efecto en la ejecución del modelo terminado.
• Compruebe la configuración predeterminada. Muchas recetas de ajuste fino ya lo activan. Si la memoria es suficiente y la velocidad es importante, desactivarlo hace que cada paso sea más rápido.

Errores comunes

• Reducir primero la longitud máxima, cortando los extremos de los ejemplos largos, cuando el checkpointing por sí solo habría hecho que cupieran.
• Esperar que reduzca los pesos o el estado del optimizador. Solo afecta a las activaciones.
• Interpretar la ralentización del 20–30 % como un fallo de la máquina. Es el precio de la memoria ahorrada.
• Asumir que elimina el coste de los logits. Con un vocabulario enorme, aún puede ser necesario un lote o una longitud máxima menores.

Referencias y lecturas adicionales

Las fuentes externas se abren en una pestaña nueva.

1. 01Training Deep Nets with Sublinear Memory Cost↗Chen et al., 2016 · arXiv · arxiv.orgEl artículo que estandarizó el recálculo: memoria √n a cambio de una propagación hacia adelante adicional.
2. 02Reducing Activation Recomputation in Large Transformer Models↗Korthikanti et al., 2022 · arXiv · arxiv.orgDe dónde proviene la estimación de activación por capa y cómo recalcular solo las partes costosas.
3. 03torch.utils.checkpoint↗PyTorch docs · docs.pytorch.org
4. 04GPU memory usage↗Hugging Face Transformers docs · huggingface.coPesos, gradientes, estado del optimizador y activaciones: a dónde se destina la memoria de entrenamiento.

Conceptos relacionados