Gradient checkpointing: ¿cómo ahorra memoria de GPU?
Traducción automática
Este artículo se tradujo automáticamente a partir de la versión original en inglés.
El gradient checkpointing guarda menos activaciones intermedias durante la pasada hacia delante y vuelve a calcular las que faltan durante el cálculo hacia atrás. Úsalo cuando las activaciones impidan que el lote de entrenamiento o la longitud de secuencia quepan en la memoria de GPU, y el cálculo adicional sea aceptable.
Reduce la memoria de las activaciones guardadas. No elimina los pesos del modelo, el estado del optimizador ni todos los búferes temporales.
¿Qué se vuelve a calcular?
El cálculo hacia atrás necesita resultados intermedios de la pasada hacia delante para calcular los gradientes. Sin checkpointing, el framework conserva los resultados necesarios. Con checkpointing, guarda determinadas entradas o resultados intermedios y vuelve a ejecutar las operaciones afectadas cuando se necesitan sus valores.
Por ejemplo, un framework puede guardar la entrada de un grupo de capas de un transformer y descartar sus activaciones internas. Durante el cálculo hacia atrás, vuelve a ejecutar ese grupo a partir de la entrada guardada. La entrada guardada es un checkpoint de activaciones, distinto de un checkpoint del modelo escrito en disco para recuperarlo.
El análisis original de Chen et al. presenta una estrategia con segmentos uniformes que requiere O(sqrt(n)) de memoria de activaciones para una red de n capas y aproximadamente una pasada adicional hacia delante por paso de entrenamiento. Los transformers reales tienen activaciones de tamaños distintos, diferentes implementaciones de atención y políticas del framework. Mide la memoria y el tiempo resultantes en lugar de usar la estimación asintótica como previsión de despliegue.
Activa la integración que espera tu entorno
Para los modelos compatibles con Hugging Face Trainer, gradient_checkpointing=True activa el mecanismo general. Con FSDP, revisa en su lugar la integración de checkpointing de activaciones del framework. La documentación actual de Transformers Trainer recomienda activation_checkpointing para FSDP porque el gradient checkpointing genérico puede provocar un all-gather adicional durante el cálculo hacia atrás.
Compara la misma configuración de entrenamiento con y sin checkpointing. Registra el pico de memoria asignada en GPU, los ejemplos o tokens por segundo y si cabe la longitud de secuencia prevista. Comprueba también que los gradientes sigan siendo válidos cuando tu modelo use operaciones personalizadas o entrenamiento de adaptadores.
Si el estado del optimizador ocupa la mayor parte de la memoria, considera en su lugar dividir ese estado entre dispositivos o usar un ajuste eficiente en parámetros; volver a calcular las activaciones no elimina esa asignación de memoria.
La sección sobre gradient checkpointing de la LLM Engineering Guide explica el cálculo por segmentos.