Geheugenefficiënte LLM-training met Online Subruimte Afdaling

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

Auteurs: Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu

cs.AI

Samenvatting

Onlangs hebben een breed scala aan geheugenefficiënte LLM-trainingsalgoritmen aanzienlijke populariteit verworven. Deze methoden maken gebruik van de low-rank structuur van gradiënten om optimizer-statussen te projecteren in een deelruimte met behulp van een projectiematrix die wordt gevonden via singuliere waardeontbinding (SVD). De convergentie van deze algoritmen is echter sterk afhankelijk van de update-regels van hun projectiematrix. In dit werk bieden we de eerste convergentiegarantie voor willekeurige update-regels van de projectiematrix. Deze garantie is algemeen toepasbaar op optimizers die kunnen worden geanalyseerd met Hamiltonian Descent, waaronder de meest gebruikelijke, zoals LION en Adam. Geïnspireerd door ons theoretisch inzicht, stellen we Online Subspace Descent voor, een nieuwe familie van subspace descent-optimizers zonder SVD. In plaats van de projectiematrix bij te werken met eigenvectoren, werkt Online Subspace Descent de projectiematrix bij met online PCA. Online Subspace Descent is flexibel en introduceert slechts minimale overhead tijdens het trainen. We laten zien dat voor de taak van het pretrainen van LLaMA-modellen variërend van 60M tot 7B parameters op de C4-dataset, Online Subspace Descent een lagere perplexiteit en betere prestaties op downstream taken bereikt dan state-of-the-art low-rank trainingsmethoden in verschillende instellingen, en de kloof met full-rank baselines verkleint.

English

Recently, a wide range of memory-efficient LLM training algorithms have gained substantial popularity. These methods leverage the low-rank structure of gradients to project optimizer states into a subspace using projection matrix found by singular value decomposition (SVD). However, convergence of these algorithms is highly dependent on the update rules of their projection matrix. In this work, we provide the first convergence guarantee for arbitrary update rules of projection matrix. This guarantee is generally applicable to optimizers that can be analyzed with Hamiltonian Descent, including most common ones, such as LION, Adam. Inspired by our theoretical understanding, we propose Online Subspace Descent, a new family of subspace descent optimizer without SVD. Instead of updating the projection matrix with eigenvectors, Online Subspace Descent updates the projection matrix with online PCA. Online Subspace Descent is flexible and introduces only minimum overhead to training. We show that for the task of pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset, Online Subspace Descent achieves lower perplexity and better downstream tasks performance than state-of-the-art low-rank training methods across different settings and narrows the gap with full-rank baselines.