Calculando Cross Entropy Loss economizando memória

Um artigo que examina por que, no treinamento de LLMs com contextos longos e vocabulários grandes, a LM head + cross entropy se torna uma das maiores fontes de consumo de memória. Com um contexto de 128K, apenas o tensor de logits chega perto de 40 GB, ficando maior até do que os weights do modelo.

Partindo de um OOM encontrado na prática ao treinar um modelo de 16B com contexto de 128K, o texto deriva do zero o forward/backward da cross entropy e mostra por que simplesmente dividir o eixo da sequência em chunks não reduz o pico de memória (o autograd mantém o grafo de cada chunk até o backward). Em seguida, explica como o FLCE calcula o gradient de cada chunk diretamente dentro do forward pass, evitando que tensores grandes permaneçam no grafo. No fim, aborda a análise do tradeoff entre memória/latência e até um walkthrough da implementação real do kernel.