Cómo calcular Cross Entropy Loss ahorrando memoria

Un artículo que examina por qué el LM head + cross entropy se convierte en una de las mayores fuentes de consumo de memoria al entrenar LLM con contextos largos y vocabularios grandes. Con un contexto de 128K, un solo tensor de logits puede acercarse a los 40 GB, superando incluso el tamaño de los pesos del modelo.

Partiendo de un OOM ocurrido en un entrenamiento real de un modelo de 16B con contexto de 128K, el texto deriva desde cero el forward/backward de cross entropy y muestra por qué dividir simplemente el eje de la secuencia en chunks no reduce la memoria pico (porque autograd sigue reteniendo el grafo de cada chunk hasta el backward). Luego explica cómo FLCE calcula el gradiente de cada chunk inmediatamente dentro del forward pass para que los tensores grandes no permanezcan en el grafo. Al final también cubre el análisis del tradeoff entre memoria y latencia, además de un recorrido por la implementación real del kernel.