Bellekten Tasarruf Ederek Cross Entropy Loss Hesaplama

Uzun context ve büyük vocab ile LLM eğitilirken, LM head + cross entropy’nin neden en büyük bellek tüketim noktalarından biri haline geldiğini ele alan bir yazı. 128K context’te tek bir logits tensörü 40GB’a yaklaşıyor; bu da model weight’lerinden bile daha büyük oluyor.

16B’lik bir modeli 128K context ile eğitirken gerçekten yaşanan bir OOM’dan yola çıkarak, cross entropy’nin forward/backward sürecini en baştan türetiyor; ardından sequence eksenini basitçe chunk’lara bölme yönteminin neden peak memory’yi düşüremediğini gösteriyor (çünkü autograd, chunk bazlı graph’ı backward aşamasına kadar elinde tutuyor). Sonrasında FLCE’nin her chunk’ın gradient’ini forward pass içinde hemen hesaplayarak büyük tensörlerin graph üzerinde kalmasını nasıl engellediğini açıklıyor. Son bölümde ise bellek/gecikme trade-off analizi ile gerçek kernel implementasyonunun adım adım incelemesi yer alıyor.