StreamBP : Rétropropagation exacte à faible consommation de mémoire pour l'entraînement de séquences longes dans les LLM

Résumé

L'entraînement de modèles de langage sur des données de séquences longues constitue une exigence exigeante pour améliorer la capacité du modèle sur des tâches complexes, telles que le raisonnement à longue chaîne. Cependant, à mesure que la longueur des séquences augmente, le coût mémoire pour stocker les valeurs d'activation devient considérable pendant le processus de rétropropagation (BP), même avec l'application de la technique de vérification des gradients. Pour relever ce défi, nous proposons une méthode de BP économe en mémoire et exacte, appelée StreamBP, qui effectue une décomposition linéaire de la règle de la chaîne le long de la dimension de séquence de manière couche par couche, réduisant ainsi significativement le coût mémoire des valeurs d'activation et des logits. La méthode proposée est applicable à des objectifs courants tels que SFT, GRPO et DPO. D'un point de vue implémentation, StreamBP réalise moins d'opérations FLOPs et une vitesse de BP plus rapide en exploitant la structure causale du modèle de langage. Par rapport à la vérification des gradients, StreamBP augmente la longueur maximale des séquences de BP de 2,8 à 5,5 fois, tout en utilisant un temps de BP comparable ou même inférieur. Il est à noter que la capacité de StreamBP à augmenter la longueur des séquences peut être directement transférée à l'augmentation de la taille des lots pour accélérer l'entraînement. Nous développons également une version distribuée et économe en communication de StreamBP pour soutenir efficacement l'entraînement multi-GPU et élargir son applicabilité. Notre code peut être facilement intégré dans le pipeline d'entraînement de tout modèle de transformateur et est disponible à l'adresse https://github.com/Ledzy/StreamBP.

English

Training language models on long sequence data is a demanding requirement for enhancing the model's capability on complex tasks, e.g., long-chain reasoning. However, as the sequence length scales up, the memory cost for storing activation values becomes huge during the Backpropagation (BP) process, even with the application of gradient checkpointing technique. To tackle this challenge, we propose a memory-efficient and exact BP method called StreamBP, which performs a linear decomposition of the chain rule along the sequence dimension in a layer-wise manner, significantly reducing the memory cost of activation values and logits. The proposed method is applicable to common objectives such as SFT, GRPO, and DPO. From an implementation perspective, StreamBP achieves less computational FLOPs and faster BP speed by leveraging the causal structure of the language model. Compared to gradient checkpointing, StreamBP scales up the maximum sequence length of BP by 2.8-5.5 times larger, while using comparable or even less BP time. Note that StreamBP's sequence length scaling ability can be directly transferred to batch size scaling for accelerating training. We further develop a communication-efficient distributed StreamBP to effectively support multi-GPU training and broaden its applicability. Our code can be easily integrated into the training pipeline of any transformer models and is available at https://github.com/Ledzy/StreamBP.

StreamBP : Rétropropagation exacte à faible consommation de mémoire pour l'entraînement de séquences longes dans les LLM

StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs

Résumé

Support