NeuZip : Entraînement et inférence efficaces en mémoire avec compression dynamique des réseaux neuronaux

NeuZip: Memory-Efficient Training and Inference with Dynamic Compression of Neural Networks

October 28, 2024

Auteurs: Yongchang Hao, Yanshuai Cao, Lili Mou

cs.AI

Résumé

Les performances des réseaux neuronaux s'améliorent lorsque davantage de paramètres sont utilisés. Cependant, les tailles des modèles sont limitées par la mémoire disponible sur l'appareil pendant l'entraînement et l'inférence. Bien que l'application de techniques telles que la quantification puisse atténuer cette contrainte, elles souffrent d'une dégradation des performances. Dans ce travail, nous présentons NeuZip, un nouveau schéma de compression des poids basé sur l'entropie des nombres en virgule flottante dans les réseaux neuronaux. Avec NeuZip, nous parvenons à obtenir un entraînement et une inférence efficaces en termes de mémoire sans sacrifier les performances. Notamment, nous réduisons de manière significative l'empreinte mémoire de l'entraînement d'un modèle Llama-3 8B de 31 Go à moins de 16 Go, tout en conservant pleinement la dynamique d'entraînement inchangée. En inférence, notre méthode peut réduire l'utilisation de la mémoire de plus de la moitié tout en maintenant des performances quasi sans perte. Notre code est disponible publiquement.

English

The performance of neural networks improves when more parameters are used. However, the model sizes are constrained by the available on-device memory during training and inference. Although applying techniques like quantization can alleviate the constraint, they suffer from performance degradation. In this work, we introduce NeuZip, a new weight compression scheme based on the entropy of floating-point numbers in neural networks. With NeuZip, we are able to achieve memory-efficient training and inference without sacrificing performance. Notably, we significantly reduce the memory footprint of training a Llama-3 8B model from 31GB to less than 16GB, while keeping the training dynamics fully unchanged. In inference, our method can reduce memory usage by more than half while maintaining near-lossless performance. Our code is publicly available.