Fine-tuning LLMs to 1.58bit: extreme quantization made easy
[Блог] [Model card] [PR с добавлением BitNet в transformers]

Некоторое время назад здесь мы разбирали статью от мелкософта про BitNet-1.58, где авторам удалось завести QAT c тернарной квантизацией и добиться вполне сносного качества при обучении с нуля.

Лицехваты 🤗 недавно опубликовали весьма занятный и содержательный длиннопост про дообучение LLM в 1.58 (log2(3)) бит.

Метод и эксперименты

Напомню, что при тернарной квантизации каждый элемент тензора принимает одно из трех значений - {-1, 0, 1} и еще есть floating-point масштаб (вместо absmax берут среднюю абсолютную величину элемента тензора).

Дабы квантовалось лучше, перед каждым линейным слоем ставится LayerNorm, и это накидывает на предобучении.

При попытке наивно затюнить Llama-3-8B в BitNet сходимость по лоссу не сильно лучше, чем у случайно инициализированной модели, и авторы начинают копать глубже.

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

Шаг обучения тоже пришлось основательно подбирать, чтобы выжать приемлемое качество. После подбора удачных гиперпараметров на коротких запусках, лучшую конфигурацию учат уже 100B токенов (0.6% бюджета обучения Llama-3-8B). Тем не менее, итоговая просадка по сравнению с исходной моделью довольно значительная.

Для маленьких LM (SmolLM 125M) разогрев (постепенная квантизация) не дает улучшений.

По бенчам просадки значительные по сравнению с fp16, но лучше чем у прошлых квантизаций в похожую битность. Сам протокол сравнения спорный 🤔, не у всех вариантов есть результаты, сравниваются перплексии разных версий Llama с разными словарями. Да и Llama-3 выбивает 65, а не 49 на MMLU 5-shot.

Авторы реализовали kernel для инференса на Тритоне, который дает некоторое ускорение. А с библиотекой BitBlas ускорение становится еще более заметным. Однако работа с ней требует небыстрой компиляции.

Вывод

Неплохой результат и исследование. Но есть вопросы 🙋‍♂️ к протоколу замеров качества, и пока BitNet теряет слишком много в качестве для того чтобы стать общепринятой практикой. Чтобы пользоваться модельками надо поставить версию transformers из PR.