Memory बचाते हुए Cross Entropy Loss की गणना करना

लंबे context और बड़े vocab वाले LLM training में LM head + cross entropy सबसे बड़े memory consumers में से एक क्यों बन जाता है, इस पर चर्चा करने वाला लेख। 128K context में सिर्फ logits tensor ही 40GB के करीब पहुंच जाता है, यानी model weights से भी बड़ा हो जाता है।

16B model को 128K context के साथ train करते समय वास्तव में आए OOM से शुरू करते हुए, यह लेख cross entropy के forward/backward को शुरू से derive करता है, फिर दिखाता है कि sequence axis को साधारण chunks में बांटने से peak memory क्यों कम नहीं होती (autograd हर chunk के graph को backward तक पकड़े रखता है)। इसके बाद यह समझाता है कि FLCE कैसे हर chunk के gradient को forward pass के अंदर ही तुरंत calculate करता है, ताकि बड़े tensors graph में बचे न रहें। अंत में memory/latency tradeoff analysis और वास्तविक kernel implementation walkthrough तक कवर किया गया है।