Квантование LLM. Part 1

Введение

Ранее мы говорили о квантовании CV-моделей. Однако оно отличается от квантования LLM. В основном это связано с тем, что для CV-моделей есть возможность применить Quantization Aware Training (далее — QAT). А для LLM дообучение очень затратно, поэтому его могут позволить только очень большие компании.

В этой статье мы сделаем акцент именно на тех методах, где дообучение не применяется вообще или применяется крайне ограниченно. Разберём два алгоритма, которые, на наш взгляд, являются ключевыми для понимания различных современных работ.

Итак, план статьи следующий:

1. Введение
2. AWQ — история о том, как 99% весов не делают это…
3. QuIP — мощный baseline
4. Заключение

Давайте начинать 🙂

AWQ, или Activation-aware Weight Quantization

Статья очень похожа на своего предшественника от той же команды.

Идея следующая: квантование портит выбросы. Выбросы — это значения, отстающие от основной массы. При квантовании они делают больший шаг дискретизации, что, в свою очередь, размывает значения вблизи нуля. Это приводит к снижению точности. В классическом CV есть способы решения такой проблемы через дообучение, но существуют и альтернативные методы. Они позволяют избегать выбросы при помощи перемасштабирования. Визуализация метода представлена ниже:

Авторы статьи AWQ выяснили: в LLM всего лишь 1% весов владеет 90% богатств мира сильно влияет на ошибку после квантования. Найти эти веса можно в тех каналах, куда приходят большие значения активаций. Таким образом, даже малое изменение подобных весов влечёт за собой значительные трансформации следующих слоёв. Как нам быть?

Авторы предлагают перемасштабировать эти веса, а также активации, которые им соответствуют. Более того, они показывают, что такой подход может привести к снижению ошибки квантования.

Следовательно, они предлагают оптимизировать следующее выражение:

Здесь \(W\)— веса квантуемого слоя, \(X\) — активации, которые приходят на вход, \(s\) — тот самый перемасштабирующий фактор, а \(Q\) — функция fake-квантования, то есть квантизации и деквантизации одновременно.

Есть небольшая проблема — такое выражение недифференцируемое. По мнению авторов, мы можем не обращать на это внимание, а просто подбирать \(s\) перебором по сетке. Это становится возможным благодаря ограничениям, которые они вводят. В частности, \(s = s_{X}^{\alpha}\), где \(\alpha\) лежит в диапазоне от 0 до 1 и, в свою очередь, перебирается по сетке, а \(s_{X}\) — среднее значение конкретного канала активации. Также авторы значительную часть статьи посвящают особенностям инференса LLM и имплементации предложенной процедуры для эффективного запуска.

QuIP, или Quantization with Incoherence Processing

Очень хорошая работа с мощнейшей математикой. По мнению авторов, различные эмпирические процедуры (как та, что выше) — это, конечно, хорошо, но вот работают они на низких разрядностях слабовато. Нам хочется получить такое квантование, чтобы в 4, 3 или даже 2 бита всё работало гораздо лучше. В статье представлено много теории, но основных идей две: адаптивная процедура округления и, опять же, борьба с выбросами (ну, в каком-то смысле).

Давайте сначала посмотрим на адаптивную процедуру округления, которую авторы называют LDLQ. Здесь они оптимизируют прокси-функцию «послойного» отличия между квантованными \(\hat{W}\) и оригинальными \(W\) весами, которые перевели в диапазон квантованных (посредством сдвигов и масштабирования, но без округления). Как можно заметить, механизм похож на тот, что был выше.

\(Е_{x}\) здесь — математическое ожидание, которое можно заменить на «его» эмпирическую версию, усреднение. Тогда получим следующее (\(W, \hat{W} \in R^{[m, n]}, x \in R^{[n, 1]}\)):

\(E_{x}[||(\hat{W} -W)x||^{2}] \approx \frac{1}{N} \sum_{i=0}^{N}[(\hat{W} -W)x_{i}]^2 =\\ \frac{1}{N} \sum_{i=0}^{N}x_{i}^T(\hat{W} -W)^T(\hat{W} -W)x_{i} =\\ \frac{1}{N} \sum_{i=0}^{N}tr(x_{i}^T(\hat{W} -W)^T(\hat{W} -W)x_{i}) =\\ \frac{1}{N} \sum_{i=0}^{N}tr((\hat{W} -W)x_{i}x_{i}^T(\hat{W} -W)^T) =\\ tr(\frac{1}{N} \sum_{i=0}^{N}(\hat{W} -W)x_{i}x_{i}^T(\hat{W} -W)^T) =\\ tr((\hat{W} -W)[\frac{1}{N} \sum_{i=0}^{N}x_{i}x_{i}^T](\hat{W} -W)^T) =\\ tr((\hat{W} -W)H(\hat{W} -W)^T)\)

В данном случае H — гессиан MSE-функции.

Затем авторы обновляют столбец весов за столбцом на основе LDL — разложения этого гессиана по формулам ниже (поэтому LDLQ и называется так, LDL + Quantization):

Основная идея заключается в том, чтобы линейным образом учесть ошибку от квантования предыдущих весов в следующие. В статье исследователи утверждают: если применить к гессиану LDL-разложение (L — нижнетреугольная, D — диагональная), то можно в качестве линейной добавки взять матрицу \(\’U\).

Авторы показывают, что эта процедура оптимальна в том смысле, что она либо лучше (ошибка от квантования меньше), либо совпадает с любой процедурой округления (например, к ближайшему целому или «стохастическому» округлению).

Затем они вводят понятие «некогерентности». Это свойство матрицы, при котором она не имеет слишком больших значений.

Зачем нам про это думать? Ведь у нас по-прежнему есть эффективная процедура округления, LDLQ. На самом деле если это свойство не соблюдается, то LDLQ ничего не стоит, просто совпадает с обычной процедурой округления. Поэтому нам нужно сначала особым образом обработать матрицу \(W\) и матрицу \(H\) и после этого уже заквантовать.

Отметим, что сопряжение симметричной матрицы со случайно ортогональной делает её собственные вектора случайными величинами, норма которых сконцентрирована около \(n^{-1/2}\).

Действительно, пусть \(H = Q\Lambda Q^{T}\). Тогда матрица \(\hat{H} = VHV^T\), где \(V\) — ортогональная, будет иметь следующие случайные вектора в качестве собственных: \(\hat{H} = VHV^T = VQ\Lambda Q^{T}V^{T} = A \Lambda A^T\), где \(A = VQ\). Затем, применив математику, мы понимаем, что значительно больших элементов у нас не будет. Таким образом, мы можем взять произвольную ортогональную матрицу V и при помощи неё предобусловить нашу матрицу H, чтобы она стала когерентной.

При этом, обработать одну только \(H\) у нас не получится, надо применить подобную предобработку к \(W\) ( \(\tilde{W} = UWV^T\), где матрица \(U\)- тоже случайная ортогональная, но уже другая), так как мы хотим, чтобы наша прокси-метрика не изменилась: \(\tilde{H} = VHV^T\), \(\tilde{W} = UWV^T\)

Таким образом, у нас получается такая стратегия:

1. Генерируем U, V.
2. Модифицируем H, W (incoherence pre-processing).
3. Реализуем LDLQ, значит, получаем пожатые веса, которые легко ложатся в память.
4. В рантайме исполняем деквантование и возвращение матрицы W в исходное пространство (incoherence post-processing). Выполняем умножение в FP16.

Но при слишком длительном умножении получится сильный оверхед на post-procession, следовательно, мы потеряем все плюшки от того, что поджали модельку.

Для решения проблемы авторы предлагают сконструировать матрицы \(U, V\) как Кронекеровские произведения матриц. Так мы имеем меньше операций, необходимых для их умножения, чем если бы мы делали это напрямую. Авторы также доказывают — при такой процедуре матрицы будут некогерентны.

Есть ряд эвристик для улучшения квантования:

1. Авторы применяют метод ближайшего на тот момент конкурента (GPT-Q) для регуляризации Гессиана: \(H = H + \alpha * mean(diag(H))I\).
2. Дополнительно предлагают перемасштабировать \(W\) и \(H\) следующим образом: \(W = WD\), \(H = D^{-1}HD^{-1}\), где матрица \(D = \sqrt[4]{\frac{diag(H)}{diag(W*W^T)}}\). Так мы имеем примерно одинаковый спектр у матриц \(H\) и \(WW^{T}\). Можем в этом убедиться, если возьмём след после перемасштабирования от \(H\) и \(WW^{T}\). Это нужно для того, чтобы, опять же, размазать большие значения спектра гессиана (или весов) между друг другом. Ведь если мы сделали матрицу некогерентной — все выбросы сосредоточатся только в больших собственных значениях матриц, то есть в спектре.
3. После того, как мы сделали нашу матрицу \(W\) некогерентной (см. определение выше), мы можем отнормировать её не на типичный \(max_{i,j}|W_{ij}|\), а на большую величину, \(\rho \frac{|W|_{F}}{\sqrt{mn}}\). Это взято из определения некогерентности, только вместо \(\mu\) у нас \(\rho\), гиперпараметр. Так мы сокращаем диапазон квантования и (в каком-то смысле) боремся с выбросами.
4. Дополнительно авторы предлагают использовать жадное обновление, то есть покоординатный спуск в том же порядке, что и при исходном методе квантования: изменяем каждый вес по очереди (в квантованном диапазоне), чтобы минимизировать прокси-лосс, оставляем при этом другие веса фиксированными.

В итоге процедура QuIP выглядит следующим образом:

Заключение

Мы не сравниваем приведенные выше статьи, поскольку сегодня они уже не являются SOTA-методами. Тем не менее они важны для понимания современных методов AQLM и QUIP#, у которых есть открытая имплементация с возможностью использования.