NeuZip: Treinamento e Inferência Eficientes em Memória com Compressão Dinâmica de Redes Neurais

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

October 28, 2024

Autores: Yongchang Hao, Yanshuai Cao, Lili Mou

cs.AI

Resumo

O desempenho das redes neurais melhora quando mais parâmetros são utilizados. No entanto, os tamanhos dos modelos são limitados pela memória disponível no dispositivo durante o treinamento e inferência. Embora a aplicação de técnicas como quantização possa aliviar a restrição, elas sofrem de degradação de desempenho. Neste trabalho, apresentamos o NeuZip, um novo esquema de compressão de pesos baseado na entropia dos números de ponto flutuante em redes neurais. Com o NeuZip, somos capazes de alcançar treinamento e inferência eficientes em memória sem sacrificar o desempenho. Notavelmente, reduzimos significativamente a pegada de memória do treinamento de um modelo Llama-3 8B de 31GB para menos de 16GB, mantendo a dinâmica de treinamento completamente inalterada. Na inferência, nosso método pode reduzir o uso de memória em mais da metade, mantendo um desempenho quase sem perdas. Nosso código está disponível publicamente.

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.