Cross-Entropy-Loss speichersparend berechnen

Ein Artikel darüber, warum der LM Head plus Cross Entropy beim Training von LLMs mit langem Kontext und großem Vokabular zu einem der größten Speicherverbraucher wird. Bei einem Kontext von 128K nähert sich allein der Logits-Tensor 40 GB und wird damit sogar größer als die Modellgewichte.

Ausgehend von einem tatsächlich aufgetretenen OOM beim Training eines 16B-Modells mit 128K Kontext leitet der Beitrag Forward und Backward von Cross Entropy von Grund auf her, zeigt, warum das einfache Aufteilen der Sequenzachse in Chunks den Peak-Speicher nicht senkt (Autograd hält den Graphen für jeden Chunk bis zum Backward fest), und erklärt, wie FLCE die Gradienten jedes Chunks direkt im Forward Pass berechnet, sodass keine großen Tensoren im Graphen verbleiben. Abschließend behandelt er die Analyse des Memory/Latency-Trade-offs sowie einen Walkthrough der tatsächlichen Kernel-Implementierung.