คำนวณ Cross Entropy Loss โดยประหยัดหน่วยความจำ

บทความที่ชี้ให้เห็นว่าในการฝึก LLM ที่มี context ยาวและ vocab ขนาดใหญ่ ทำไม LM head + cross entropy จึงกลายเป็นหนึ่งในจุดที่ใช้หน่วยความจำมากที่สุด ที่ context 128K เพียง tensor ของ logits ตัวเดียวก็มีขนาดเกือบ 40GB จนใหญ่กว่า weight ของโมเดลเสียอีก

เริ่มจากกรณี OOM ที่พบจริงระหว่างฝึกโมเดล 16B ด้วย context 128K จากนั้นไล่อนุพันธ์ forward/backward ของ cross entropy ตั้งแต่ต้น แสดงให้เห็นว่าการแบ่งแกน sequence ออกเป็น chunk แบบง่าย ๆ ไม่สามารถลด peak memory ได้อย่างไร (เพราะ autograd ยังคงถือ graph ของแต่ละ chunk ไว้จนถึง backward) แล้วอธิบายวิธีที่ FLCE คำนวณ gradient ของแต่ละ chunk ทันทีภายใน forward pass เพื่อไม่ให้ tensor ขนาดใหญ่ค้างอยู่ใน graph ช่วงท้ายยังครอบคลุมการวิเคราะห์ tradeoff ระหว่าง memory/latency และ walkthrough การ implement kernel จริงด้วย