VeLoRA: Geheugenefficiënte training met rang-1 sub-token projecties

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

May 28, 2024

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

cs.AI

Samenvatting

Grote taalmmodellen (LLMs) zijn recent naar voren gekomen als krachtige tools voor het aanpakken van veel taalverwerkingstaken. Ondanks hun succes zijn het trainen en finetunen van deze modellen nog steeds veel te reken- en geheugenintensief. In dit artikel identificeren en karakteriseren we de belangrijke componenten die nodig zijn voor effectieve modelconvergentie met behulp van gradient descent. Hierbij ontdekken we dat de tussenliggende activaties die worden gebruikt voor backpropagation overmatig kunnen worden gecomprimeerd zonder dat dit ten koste gaat van de prestaties. Dit resultaat leidt ons naar een goedkope en geheugenefficiënte algoritme voor zowel het finetunen als het pretrainen van LLMs. Het voorgestelde algoritme verdeelt eenvoudigweg de tokens in kleinere sub-tokens voordat ze tijdens de forward pass worden geprojecteerd op een vaste 1-dimensionale deelruimte. Deze kenmerken worden vervolgens grof gereconstrueerd tijdens de backward pass om de update-regels te implementeren. We bevestigen de effectiviteit van ons algoritme als aanvulling op veel state-of-the-art PEFT-methoden op de VTAB-1k finetuning benchmark. Bovendien overtreffen we QLoRA voor het finetunen van LLaMA en laten we competitieve prestaties zien tegenover andere geheugenefficiënte pretrainingsmethoden op de grootschalige C4-dataset.

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.