NeuZip: Entrenamiento e Inferencia Eficientes en Memoria con Compresión Dinámica de Redes Neuronales

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

October 28, 2024

Autores: Yongchang Hao, Yanshuai Cao, Lili Mou

cs.AI

Resumen

El rendimiento de las redes neuronales mejora cuando se utilizan más parámetros. Sin embargo, los tamaños de los modelos están limitados por la memoria disponible en el dispositivo durante el entrenamiento y la inferencia. Aunque la aplicación de técnicas como la cuantificación puede aliviar la restricción, sufren de degradación del rendimiento. En este trabajo, presentamos NeuZip, un nuevo esquema de compresión de pesos basado en la entropía de números de punto flotante en redes neuronales. Con NeuZip, logramos alcanzar un entrenamiento e inferencia eficientes en memoria sin sacrificar el rendimiento. Es destacable que reducimos significativamente la huella de memoria para el entrenamiento de un modelo Llama-3 8B de 31GB a menos de 16GB, manteniendo completamente inalterada la dinámica de entrenamiento. En inferencia, nuestro método puede reducir el uso de memoria en más de la mitad manteniendo un rendimiento casi sin pérdidas. Nuestro código está disponible públicamente.

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.