Calculer la Cross Entropy Loss en économisant de la mémoire

Un article qui explique pourquoi, lors de l’entraînement de LLM avec un contexte long et un grand vocabulaire, la tête LM + la cross entropy devient l’un des plus gros postes de consommation mémoire. Avec un contexte de 128K, le seul tenseur des logits approche les 40 Go, devenant même plus volumineux que les poids du modèle.

En partant d’un OOM réellement rencontré lors de l’entraînement d’un modèle 16B avec un contexte de 128K, l’article dérive depuis le début le forward/backward de la cross entropy, puis montre pourquoi découper simplement l’axe de séquence en chunks ne réduit pas le pic mémoire (autograd conserve le graphe de chaque chunk jusqu’au backward). Il explique ensuite comment FLCE calcule le gradient de chaque chunk directement dans la passe forward, afin qu’aucun gros tenseur ne reste dans le graphe. Enfin, il couvre l’analyse du compromis mémoire/latence ainsi qu’un walkthrough de l’implémentation réelle du kernel.