メモリを節約しながら Cross Entropy Loss を計算する

長いコンテキストと大きな語彙を持つ LLM の学習で、LM head + cross entropy がなぜ最大級のメモリ消費源の一つになるのかを掘り下げた記事。128K コンテキストでは、logits テンソル 1 つだけで 40GB に迫り、モデルの重みよりも大きくなる。

16B モデルを 128K コンテキストで学習した際に実際に発生した OOM を出発点として、cross entropy の forward/backward を最初から導出し、sequence 軸を単純に chunk に分割する方式ではなぜ peak memory を下げられないのか（autograd が chunk ごとの graph を backward まで保持しているため）を示したうえで、FLCE が各 chunk の gradient を forward pass 内で即座に計算し、大きなテンソルが graph に残らないようにする方法を説明する。最後には、memory/latency のトレードオフ分析と実際の kernel 実装のウォークスルーまで扱う。