June 12, 2023

Memory-Efficient Fine-Tuning of Compressed Large Language Models via sub-4-bit Integer Quantization

Введение

Значится, есть у нас квантованная модель, скажем, в 4-бита с некоторой просадкой в качестве на целевой задаче. И мы хотим каким-то образом устранить эту просадку. Естественный вариант - дообучить на целевой задаче.

Можно обучать все параметры, но это слишком накладно по памяти, можно применять один из известных вариантов PEFT - (адаптеры, LoRA, prompt tuning).

Метод

В этой статье предложен новый вариант PEFT, специфичный для квантования, под названием Parameter-Efficient and Quantization-aware Adaptation (PEQA), где по существу в качестве обучаемых параметров выступают скейлы.

Напомним, что при квантовании веса нейронной сети представляются в виде целочисленных матриц которые при домножении на скаляр/вектор масштабов дают что-то похожее на исходные веса. Существует множество алгоритмов квантования, в том числе использующих прокси для гессиана, как OPTQ, используемый в сравнении, но в данной работе авторы для простоты используют RTN квантование (округление к ближайшему значению на сетке).

QAT (Quantization-Aware training) заключается в том, чтобы обучать все значения алгоритмом обратного распространия ошибка, но он требует хранения градиентов и состояний оптимизатора для всех значений весов, что слишком дорого. И в данной работе предлалагается держать целочисленные матрицы W0 фиксированными , и обучать только скейлы s0. Если иметь один скейл на один канал, то для типичной LLM, скейлов будет несколько миллионов против нескольких миллиардов параметров модели. И, как утверждается, подобная техника даже более эффективна по параметрам, чем пресловутая LoRA.

Статья про QLoRA появилась на arXiv в тот же день, потому у авторов есть уважительная причина использовать обычную LoRA вместо QLoRA в качестве бейзлайна.

Эксперименты

Для экспериментальной проверки предложенного метода авторы сначала обучают квантованную сеть (вернее только ее скейлы) на задаче CausalLM (предсказания следующего токена по предыдущим) на стандартных бенчмарках Wikitext2 и PTB для вариаций GPT и моделей из семейства LLaMA.

Согласно нижеприведенной таблице, их метод уверенно бьет конкурентые подходы (однако, здесь замечу, что полученные значения для OPTQ с LoRA хуже, чем заявленные в репозитории, что навевает смутные сомнения).

Далее для оценки качества эффективности PEQA на instruction-tuning берется датасет ALPACA, на котором сколько-то шагов обучаются скейлы, а затем замеряется качество модели на бенчмарках по Natural Language Understanding и CommonSense Reasoning. Полученный подход в Zero-Shot выступает наравне с LoRA в fp16.

Далее модель валидируют на 5 бенчмарках из MMLU (без сравнения с конкуретными подходами, небось проиграли заметно LoRA).

Полученный подход имеет совсем небольшое число обучаемых параметров, в несколько раз меньше, чем LoRA. Под размером модели подразумевается квантованная модель в их случае против fp16/bf16 в случае LoRA.

Итог

Авторы предложили дешевый и сердитый способ дообучения квантованной модели. К достоинствам метода следует отнести простоту и дешевизну, однако, в то же время выразительность метода довольно ограничена. В качестве возможных направлений для развития метода можно было бы использовать квантование более совершенным алгоритмом (тем же OPTQ), и квантовать меньшими группами - по 128 весов в каждой например, что с одной с стороны, повысило бы точность метода, и увеличило количество обучаемых параметров.