Эффективное с точки зрения памяти обучение LLM с использованием метода онлайн-спуска в подпространство.

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

Авторы: Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu

cs.AI

Аннотация

В последнее время широкое распространение получили алгоритмы обучения с пониженным потреблением памяти для моделей с ограниченной памятью. Эти методы используют структуру низкого ранга градиентов для проецирования состояний оптимизатора в подпространство с использованием матрицы проекции, найденной с помощью сингулярного разложения (SVD). Однако сходимость этих алгоритмов в значительной степени зависит от правил обновления их матрицы проекции. В данной работе мы предоставляем первое гарантированное условие сходимости для произвольных правил обновления матрицы проекции. Это условие общим образом применимо к оптимизаторам, которые могут быть проанализированы с помощью метода гамильтонова спуска, включая наиболее распространенные, такие как LION, Adam. Вдохновленные нашим теоретическим пониманием, мы предлагаем метод Онлайн спуска в подпространство, новое семейство оптимизаторов спуска в подпространство без использования SVD. Вместо обновления матрицы проекции собственными векторами, Онлайн спуск в подпространство обновляет матрицу проекции с помощью онлайн метода главных компонент. Онлайн спуск в подпространство гибок и вносит минимальные накладные расходы на обучение. Мы показываем, что для задачи предварительного обучения моделей LLaMA с числом параметров от 60 миллионов до 7 миллиардов на наборе данных C4, метод Онлайн спуска в подпространство достигает более низкой перплексии и лучшей производительности на последующих задачах, чем передовые методы обучения с низким рангом в различных настройках и сокращает разрыв с полными базовыми значениями.

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.