온라인 부분 공간 강하를 이용한 메모리 효율적인 LLM 훈련

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

저자: Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu

cs.AI

초록

최근에는 메모리를 효율적으로 활용하는 다양한 LLM 훈련 알고리즘이 상당한 인기를 얻고 있습니다. 이러한 방법들은 그래디언트의 저랭크 구조를 활용하여 최적화기 상태를 투영 행렬을 사용하여 부분 공간으로 변환하는데 SVD에 의해 찾아진 투영 행렬을 활용합니다. 그러나 이러한 알고리즘의 수렴은 그들의 투영 행렬의 업데이트 규칙에 매우 의존적입니다. 본 연구에서는 임의의 투영 행렬 업데이트 규칙에 대한 최초의 수렴 보장을 제시합니다. 이 보장은 대부분의 일반적인 최적화기에 대해 해밀토니안 하강을 통해 분석할 수 있는 옵티마이저에 일반적으로 적용됩니다. 이에는 LION, Adam과 같은 가장 흔한 것들이 포함됩니다. 우리의 이론적 이해를 바탕으로, 우리는 SVD 없이 새로운 서브스페이스 하강 옵티마이저인 온라인 서브스페이스 하강을 제안합니다. 온라인 서브스페이스 하강은 고유벡터로 투영 행렬을 업데이트하는 대신 온라인 PCA로 투영 행렬을 업데이트합니다. 온라인 서브스페이스 하강은 유연하며 훈련에 최소한의 오버헤드만 도입합니다. 우리는 C4 데이터셋에서 60M에서 7B 매개변수 범위의 LLaMA 모델 사전 훈련 작업에 대해, 온라인 서브스페이스 하강이 최첨단 저랭크 훈련 방법보다 낮은 퍼플렉서티와 더 나은 하위 작업 성능을 달성하며 다양한 설정에서 전체 랭크 기준과의 격차를 줄입니다.

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.