Speichereffizientes Training von LLM mit Online-Subspace-Abstieg

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

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

cs.AI

Zusammenfassung

In letzter Zeit haben eine Vielzahl von speicher-effizienten LLM-Trainingsalgorithmen erheblich an Popularität gewonnen. Diese Methoden nutzen die Niedrigrangstruktur der Gradienten, um die Optimiererzustände mithilfe einer Projektionsmatrix in einen Unterraum zu projizieren, die durch Singulärwertzerlegung (SVD) gefunden wird. Die Konvergenz dieser Algorithmen hängt jedoch stark von den Aktualisierungsregeln ihrer Projektionsmatrix ab. In dieser Arbeit bieten wir die erste Konvergenzgarantie für beliebige Aktualisierungsregeln der Projektionsmatrix. Diese Garantie ist im Allgemeinen auf Optimierer anwendbar, die mit Hamiltonian Descent analysiert werden können, einschließlich der häufigsten wie LION und Adam. Inspiriert von unserem theoretischen Verständnis schlagen wir Online Subspace Descent vor, eine neue Familie von Subspace-Descent-Optimierern ohne SVD. Anstatt die Projektionsmatrix mit Eigenvektoren zu aktualisieren, aktualisiert Online Subspace Descent die Projektionsmatrix mit Online-PCA. Online Subspace Descent ist flexibel und führt nur minimale zusätzliche Belastung beim Training ein. Wir zeigen, dass Online Subspace Descent für die Aufgabe des Pretrainings von LLaMA-Modellen mit 60M bis 7B Parametern auf dem C4-Datensatz eine niedrigere Perplexität und eine bessere Leistung bei nachgelagerten Aufgaben erzielt als modernste Niedrigrang-Trainingsmethoden in verschiedenen Einstellungen und den Abstand zu vollrangigen Baselines verringert.

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.