October 28, 2024

Linear Quantization

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

Source: КПД, Quantization in Depth


Brief Literary Review

Вообще исследователи уже давно стараются оптимизировать модельки, уменьшать их вес и размер, при сохранении качества работы. О прунинге случайного леса я слышал еще в самом начале своей карьеры, а с появлением LLM эта тема и вовсе получила финансовый толчок от разных компаний, ведь инференсить квантизованную модельку дешевле, удобнее и быстрее, чем тратить миллионы у.е. на деплой Llama70B.

Фундаментальной работой в этой области стало исследование GPT3.int8(): 8-bit Matrix Multiplication for Transformers at Scale в которой авторы показали возможность эффективной 8-битной квантизации крупных языковых моделей. Далее ребята из Яндекса сказали, что это не предел и выкатили Extreme Compression of Large Language Models via Additive Quantization с возможностью сжатия до 2-3 битов, при чем их метод аддитивной квантизации является Парето-оптимальным с точки зрения соотношения точности и размера модели при сжатии до менее чем 3 битов на параметр. Этот же метод потом раскатили и на диффузионных моделях Accurate Compression of Text-to-Image Diffusion Models via Vector Quantization, что дало результаты лучше, чем Q-Diffusion: Quantizing Diffusion Models и Post-training Quantization on Diffusion Models. Ну и вскоре китайцы в работе VPTQ: Extreme Low-bit Vector Post-Training Quantization for Large Language Models показали, что векторная квантизация лучше всех и мы можем сжимать модели до 2 бит, при этом значительно снижать перплексию квантованных моделей и улучшать точность на задачах вопросов и ответов. А тут EDEN: Communication-Efficient and Robust Distributed Mean Estimation for Federated Learning ребята показали, что квантовать можно не только веса и активации, но и градиенты.

Из этого короткого литобзора видна довольно высокая актуальность методов квантизации, что еще раз подтолкнуло меня на разбор некоторых методов более подробно. Однако я не хочу, чтобы кто-то подумал, что я профессионал квантизации и разберу сразу все в своих постах. Если вы хотите углубиться в эту тему или развиваться в ней, то рекомендую канал КПД - там вы найдете море постов про квантизацию моделей, а сам автор является тем самым профессионалом в этой сфере.

Моей целью будет являться разбор базовых понятий, которые помогли бы в дальнейшем понимать сложные вещи, которые я буду разбирать чуть позже. Мне показалось, что сделать это будет проще всего по курсу о квантизации от HuggingFace, так что если вы его не смотрели или у вас не хватало времени на него, то у вас появилась отличная возможность прочитать summary этого курса :)

Linear Quantization

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

Параметры s (scale) и z (zero point) имеют тот же тип данных, как у original value и quantized value соответственно. Из формулы выше закономерно вытекает формула квантизация вектора.

Как видно, линейная квантизация не является каким-то rocket science, однако возникает закономерный вопрос - а как определить те самые s (scale) и z (zero point)? Для этого существуют отдельные формулы, из которых, с помощью границ области определения оригинальных значений и квантизованных, определяются s и z:

В случае, если zero point выходит из области значений q, то он приравнивается ближайшему крайнему значению q. Например, если z < q_min, тогда z = q_min.

Linear Quantization бывает двух типов: asymmetric и symmetric. Asymmetric мы разбирали только что, когда r_min не равен по модулю r_max (аналогично с q_min и q_max). В symmetric нам нет необходимости использовать zero point, ведь если q_min = q_max, то это просто середина отрезка, которая сопадает с серединой отрезка [-r_max; r_max] и равняется нулю. Из-за этого формулы немного изменятся, но суть останется той же:

Granularity

Также стоит добавить, что нет необходимости расчитывать в Linear Quantization параметры zero point и scale для всего тензора - это может негативно повлиять на точность. Лучше рассчитывать отдельные zero point и scale для каждой группы (например для каждой строки или даже n элементов тензора) - это и называется Granularity.

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

Activations Linear Quantization

Помимо весов, мы можем квантовать активации. В курсе приводится следующее мнение:

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

2) Когда мы квантуем веса и активации, то часто нам необходимы типы данных в 8 битной памяти, поэтому используются такие как INT8 или INT4. Я не очень понял, что они имели в виду, когда писали, что деквантизация не поддерживается на некоторых устройствах, поэтому если вы поняли этот момент, то напишите в комментариях.

Weights Packing

Мы точно можем квантовать до 8 бит, но в сфере больших моделей это может вызвать проблемы - особенно когда значения будут все чаще выходить из этого диапазона, а точность закономерно будет падать, то мощностей мы сильно не сэкономим. Хочется квантовать в еще менее малое количество битов, чтобы эффективно сохранять память.

А можем ли мы беспрепятственно квантовать до 2 бит? Кажется, мы можем столкнуться с ограничениями системы счисления, ведь довольно странно отображать числа в двух битах. Однако мы можем схитрить и таки перевести 8-битный тензор в 2-битный.

Допустим у нас будет тензор, состоящий из 4 чисел в формате UINT8:

[1, 0, 3, 2]

Их двоичная запись будет выглядеть вот так:

Разумеется, нам совсем не хочется хранить нулевые биты, которые не несут никакой информации:

Поэтому мы уберем их и соединим информативные биты в единое чило формата 8 бит:

Или же в десятичной записи:

[177]

Итак, наша 2-битная квантизация заработала. Она позволяет отображать реальный объем памяти, который занимают 2-битные квантизованные веса. Однако у этого метода есть и недостаток: если мы хотим прогнать инференс модели, нам нужно распаковать параметры, потому что pytorch не поддерживает инференс на 2-битных весах.

В коде это можно иплементировать с помощью операции OR, постепенно сдвигая биты и записывая их в заранее созданный тензор нужного размера:

for i in range(num_values):
    for j in range(num_steps):
        packed_tensor[i] |= uint8tensor[unpacked_idx] << (bits * j)
        unpacked_idx += 1
return packed_tensor

А unpacking:

mask = 2 ** bits - 1

for i in range(uint8tensor.shape[0]):
    for j in range(num_steps):
        unpacked_tensor[unpacked_idx] |= uint8tensor[i] >> (bits * j)
        unpacked_idx += 1

unpacked_tensor &= mask

Beyond Linear Quantization

LLMs становятся все больше и больше, поэтому линейная квантизация, конечно же, уже не обеспечивает должного эффекта экономии памяти и вычислений. Помимо нее существуют и другие методы, такие как LLM.int8, GPTQ, SmoothQuant, AWQ и так далее. Их статьи я буду разбирать в следующих постах.


На этом у меня все :)
Спасибо, что дочитали до конца!