Menghitung Cross Entropy Loss Sambil Menghemat Memori

Artikel yang membahas mengapa LM head + cross entropy menjadi salah satu konsumen memori terbesar dalam pelatihan LLM dengan context panjang dan vocab besar. Pada context 128K, satu tensor logits saja mendekati 40GB, bahkan lebih besar daripada weight model.

Berangkat dari OOM yang benar-benar terjadi saat melatih model 16B dengan context 128K, artikel ini menurunkan forward/backward cross entropy dari awal, lalu menunjukkan mengapa sekadar membagi sumbu sequence menjadi chunk tidak dapat menurunkan peak memory (autograd tetap menahan graph tiap chunk sampai backward). Setelah itu, artikel ini menjelaskan bagaimana FLCE menghitung gradient tiap chunk langsung di dalam forward pass agar tensor besar tidak tersisa di graph. Di bagian akhir, artikel ini juga membahas analisis tradeoff memory/latency serta walkthrough implementasi kernel nyata.