BitSparse stocke les activations feed-forward parcimonieuses en utilisant moins de mémoire
En termes simples
Pour entraîner un modèle de langage, il faut conserver certaines valeurs calculées lors de la passe avant afin de pouvoir ensuite calculer les mises à jour. Les couches feed-forward de la famille ReLU produisent souvent des zéros exacts, qui occupent quand même de la place lorsqu’on les stocke de façon habituelle.
BitSparse conserve les valeurs non nulles et une carte indiquant leurs positions, puis reconstitue les valeurs enregistrées au moment voulu. Dans les charges de travail testées, cette méthode a réduit l’utilisation de mémoire, y compris pour deux modèles de langage. Les mesures ne couvrent pas toutes les étapes d’un entraînement complet.
La recherche
Les systèmes d’entraînement doivent conserver les activations feed-forward pour la rétropropagation, et la mémoire nécessaire augmente avec le nombre de jetons traités ensemble.
Peut-on exploiter les zéros exacts produits par les activations de la famille ReLU pour réduire cet espace de stockage sans recalculer ni modifier les valeurs ?
BitSparse stocke les activations feed-forward non nulles avec leurs positions et les reconstitue pour la rétropropagation, ce qui réduit la mémoire occupée par leur cache.
L’article examine si les zéros exacts des activations de la famille ReLU conservées peuvent réduire la mémoire d’entraînement sans recalcul ni approximation. La quantité d’activations conservées augmente avec le nombre de jetons traités simultanément.
BitSparse stocke les valeurs non nulles et leurs positions, puis les reconstitue pour la rétropropagation. Dans les exécutions testées de modèles de langage avec passes avant et arrière, la mémoire ajustée par jeton a diminué de 31 % pour Modded-NanoGPT et de 30 % pour Nemotron-H, par rapport au stockage dense des activations.
Ces exécutions ne mettaient pas à jour les paramètres et ne conservaient pas l’état de l’optimiseur.
Les mesures de mémoire portent sur certaines couches feed-forward de la famille ReLU et sur des charges de travail précises. La comparaison d’entraînement complète utilisait aussi TensorBuffer, et la source ne précise pas si les deux exécutions sont parties du même état.
Analyse
Détails
- Auteurs de l’article original
- Max Zhu, Bohang Sun, Pietro Lio, Jialin Yu
- Soumis à arXiv
- Licence
- CC BY 4.0