Трава: Вычисление эффективного обучения низкопамятичной модели языкового моделирования с структурированными разреженными градиентами

Grass: Compute Efficient Low-Memory LLM Training with Structured Sparse Gradients

June 25, 2024

Авторы: Aashiq Muhamed, Oscar Li, David Woodruff, Mona Diab, Virginia Smith

cs.AI

Аннотация

Обучение и донастройка больших языковых моделей (LLM) часто затрудняются ограниченной памятью GPU. Существующие методы оптимизации на основе проекций решают эту проблему путем проецирования градиентов в низкоразмерное подпространство для уменьшения памяти состояния оптимизатора, однако они обычно используют плотные матрицы проекций, что может привести к избыточным вычислительным и памятьным нагрузкам. В данной работе мы предлагаем Grass (GRAdient Stuctured Sparsification), новый подход, который использует разреженные проекции для преобразования градиентов в структурированные разреженные обновления. Этот дизайн не только значительно снижает использование памяти для состояний оптимизатора, но также минимизирует память для градиентов, вычислительные и коммуникационные затраты, что приводит к существенному увеличению пропускной способности. Обширные эксперименты на задачах предварительного обучения и донастройки показывают, что Grass достигает конкурентоспособных результатов по сравнению с обучением полного ранга и существующими методами на основе проекций. Заметно, что Grass позволяет предварительное обучение модели LLaMA с 13 миллиардами параметров с использованием половины точности на одном GPU A100 с памятью 40 ГБ - подвиг, недостижимый для предыдущих методов, и обеспечивает увеличение пропускной способности до 2 раз на системе с 8 GPU. Код можно найти по ссылке https://github.com/aashiqmuhamed/GRASS.

English

Large language model (LLM) training and finetuning are often bottlenecked by limited GPU memory. While existing projection-based optimization methods address this by projecting gradients into a lower-dimensional subspace to reduce optimizer state memory, they typically rely on dense projection matrices, which can introduce computational and memory overheads. In this work, we propose Grass (GRAdient Stuctured Sparsification), a novel approach that leverages sparse projections to transform gradients into structured sparse updates. This design not only significantly reduces memory usage for optimizer states but also minimizes gradient memory footprint, computation, and communication costs, leading to substantial throughput improvements. Extensive experiments on pretraining and finetuning tasks demonstrate that Grass achieves competitive performance to full-rank training and existing projection-based methods. Notably, Grass enables half-precision pretraining of a 13B parameter LLaMA model on a single 40GB A100 GPU--a feat infeasible for previous methods--and yields up to a 2times throughput improvement on an 8-GPU system. Code can be found at https://github.com/aashiqmuhamed/GRASS .