VeLoRA: Entrenamiento Eficiente en Memoria mediante Proyecciones de Sub-Tokens de Rango 1

VeLoRA: Memory Efficient Training using Rank-1 Sub-Token Projections

May 28, 2024

Autores: Roy Miles, Pradyumna Reddy, Ismail Elezi, Jiankang Deng

cs.AI

Resumen

Los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) han surgido recientemente como herramientas poderosas para abordar muchas tareas de procesamiento de lenguaje. A pesar de su éxito, el entrenamiento y el ajuste fino de estos modelos sigue siendo excesivamente intensivo en términos de cómputo y memoria. En este artículo, identificamos y caracterizamos los componentes importantes necesarios para lograr una convergencia efectiva del modelo utilizando el descenso de gradiente. Al hacerlo, descubrimos que las activaciones intermedias utilizadas para implementar la retropropagación pueden comprimirse en exceso sin incurrir en ninguna degradación del rendimiento. Este resultado nos lleva a un algoritmo económico y eficiente en memoria tanto para el ajuste fino como para el preentrenamiento de LLMs. El algoritmo propuesto simplemente divide los tokens en sub-tokens más pequeños antes de proyectarlos en un subespacio unidimensional fijo durante la pasada hacia adelante. Estas características se reconstruyen de manera aproximada durante la pasada hacia atrás para implementar las reglas de actualización. Confirmamos la efectividad de nuestro algoritmo como complemento a muchos métodos PEFT de vanguardia en el punto de referencia de ajuste fino VTAB-1k. Además, superamos a QLoRA en el ajuste fino de LLaMA y mostramos un rendimiento competitivo frente a otros métodos de preentrenamiento eficientes en memoria en el conjunto de datos a gran escala C4.

English

Large language models (LLMs) have recently emerged as powerful tools for tackling many language-processing tasks. Despite their success, training and fine-tuning these models is still far too computationally and memory intensive. In this paper, we identify and characterise the important components needed for effective model convergence using gradient descent. In doing so we find that the intermediate activations used to implement backpropagation can be excessively compressed without incurring any degradation in performance. This result leads us to a cheap and memory-efficient algorithm for both fine-tuning and pre-training LLMs. The proposed algorithm simply divides the tokens up into smaller sub-tokens before projecting them onto a fixed 1-dimensional subspace during the forward pass. These features are then coarsely reconstructed during the backward pass to implement the update rules. We confirm the effectiveness of our algorithm as being complimentary to many state-of-the-art PEFT methods on the VTAB-1k fine-tuning benchmark. Furthermore, we outperform QLoRA for fine-tuning LLaMA and show competitive performance against other memory-efficient pre-training methods on the large-scale C4 dataset.