Checkpointing de activaciones
Leerlo en el Capítulo 4 · checkpoint +50 XP →
Descarta activaciones en el forward y recalcúlalas en el backward: cambia ~33% más de cómputo por mucha menos memoria.
La retropropagación necesita las activaciones del forward. Para una capa transformer con FlashAttention son aproximadamente 34 · s · b · h bytes (secuencia × micro-batch × hidden), y se acumulan rápido. Un modelo de 7B con micro-batch 4 necesita más de 70 GB de activaciones.
El checkpointing (también llamado recomputación) guarda solo la entrada de cada capa y vuelve a ejecutar el forward de la capa durante el backward. Costo: un forward extra, ~33% más de cómputo (8N en lugar de 6N FLOPs por token). Beneficio: la memoria de activaciones baja a unos 2 · s · b · h bytes por capa.
Úsalo cuando las activaciones son lo que te impide hacer caber el modelo o un micro-batch mayor. Apágalo cuando la memoria alcanza. Entonces es puro costo. El checkpointing selectivo (solo el interior de la atención) logra la mayor parte del ahorro por una fracción del costo.
Nota que el MFU deliberadamente no cuenta los FLOPs recalculados, así que el checkpointing baja el MFU aunque la GPU esté igual de ocupada.
Úsalo en una misión
Relacionado
Fuentes y lecturas adicionales
- Reducing Activation Recomputation (Korthikanti et al., 2022)
- Distributed Training (T. Tsugriy) — parallelism strategies, collectives, memory
Escrito para LazyLab con nuestras propias palabras; las fuentes se citan para profundizar y verificar.