Назад к блогу

Тернарные LLM и барьер 1.58 бита: откуда берётся число и как его обходят

Тернарные LLM и барьер 1.58 бита: откуда берётся число и как его обходят

Тернарные языковые модели обещают радикальное сжатие: каждый вес принимает всего три значения, а теоретический минимум — около 1.58 бита. На практике реальные форматы хранения расходуют больше, и статья разбирает, откуда берётся этот разрыв и как его обойти за счёт неравномерного распределения нулей и табличных подстановок вместо умножения.

тернарная модель. Отсюда и берётся «1.58 бита»: три равновероятных состояния требуют log2(3) ≈ 1.585 бита. Но это информационно-теоретическая оценка, а не то, сколько модель занимает на диске. Реальное развёртывание упаковывает пять тернарных весов в один байт, и из-за группировки по степеням двойки стоимость округляется вверх до 1.625 бита на вес. Разрыв между 1.585 и 1.625 — это и есть «барьер», который пытаются пробить: он возникает не из-за энтропии данных, а из-за формата хранения, который трактует три символа как равновероятные. А они не равновероятны.

Откуда берётся 1.58 и почему формат его не достигает

Три равновероятных состояния дают log2(3) ≈ 1.585 бита — это нижняя граница для хранения одного веса, если бы символы встречались одинаково часто. Распространённый формат развёртывания округляет эту стоимость до 1.625 бита на вес. То есть формат платит за удобство упаковки, а не за реальную энтропию данных.

Измерение фактического распределения символов в 29 тернарных LLM-моделях показало, что нули составляют до 51.5% всех весов. Если больше половины весов — нули, трактовать три символа как равновероятные расточительно. На этом наблюдении и строится раскладка, в которой отдельно хранится битовая карта присутствия — по одному биту на вес, отмечающему, равен ли вес нулю, — и компактный вектор знаков только для ненулевых весов; стоимость такой раскладки равна 2 − z бит на элемент веса при плотности нулей z. Чем больше нулей, тем дешевле хранение: на самой разреженной из протестированных моделей такая раскладка достигает 1.485 бита на вес — ниже и теоретического log2(3), и стоимости распространённого формата.

Контекст, в котором это появилось: BitNet b1.58 — вариант 1-битной LLM, в которой каждый параметр (или вес) тернарный {-1, 0, 1}. Именно на этой архитектуре и проверяются форматы хранения.

Как устроено ядро: таблица подстановки вместо умножения

Ключевая механика тернарного инференса — замена умножения-накопления на таблицу подстановки. Идея проста: если активация — это 4-битное значение, у неё всего 16 возможных состояний, и для каждого можно заранее посчитать, что оно даст в комбинации с весами. Поэтому таблица состоит ровно из 16 слотов: каждому из 16 возможных 4-битных шаблонов активаций соответствует один заранее вычисленный элемент.

При построении таблицы float-веса умножаются на масштаб, округляются до целых и упаковываются в int16-векторы. Из двух половин весов собираются 16 векторов таблицы как суммы и разности этих половин. Каждый элемент таблицы — заранее вычисленная линейная комбинация двух весов, соответствующая одному из 16 возможных 4-битных шаблонов активаций.

При вычислении индекс формируется из самих активаций: байты загружаются, из них выделяются старшие полубайты и младшие, и эти 4-битные значения используются как индексы в таблице. Результат выборки сразу накапливается сложением — вместо умножения каждого веса на активацию и последующего накопления выполняется одна табличная выборка и одно сложение. В конце целочисленные накопления делятся на масштаб таблицы и умножаются на масштаб весов, возвращая результат в float.

Построение таблицы: lut_ctor и его варианты

В ARM-версии (tl1) построение таблицы начинается с чтения масштаба, затем для каждого блока загружаются веса, умножаются на масштаб, округляются до int32 и сужаются до int16. Из пар векторов строятся 16 векторов таблицы как суммы и разности. Далее выполняется перестановка Transpose_8_8, а табличная выборка байтов сохраняет результат в буфер по вычисленным смещениям.

В x86-версии (TL2) есть два варианта построения таблицы, и различаются они числом весовых значений, которые обрабатываются за одну итерацию, — от этого зависит, сколько комбинаций знаков попадает в 16 слотов таблицы. Вариант для трёх весов за итерацию даёт 8 комбинаций знаков на 16 слотов: цикл идёт с шагом act_k / 24, а индексы для сбора задаются с шагом 12. Вариант для двух весов за итерацию использует только два вектора, а старшие слоты таблицы обнуляются: цикл идёт с шагом act_k / 16.

Оба варианта нужны, потому что соответствуют разным типам данных. В одном наборе ядер проверка допускает Q4_0 или TL2, в другом — Q4_0 или TL1, и для каждого типа требуется своя форма таблицы.

Квантизация и предобработка

Масштаб для тензора считается как 127, делённое на максимум абсолютных значений по всем элементам. В AVX2-варианте абсолютное значение получается через маску знака, максимум накапливается векторно, финальный масштаб — 127 / max. В NEON-варианте то же самое через горизонтальный максимум.

Отдельная деталь — сброс накопленного масштаба между блоками. Функция сброса обнуляет массив масштабов: в TL1-версии один элемент, в TL2-версии — цикл по bs элементам. Это нужно потому, что построитель таблицы читает масштаб из общего места и в конце записывает его обратно — значение переиспользуется между вызовами, и без сброса накопленный ранее масштаб попал бы в следующий блок.

Переупаковка весов тензора начинается с проверки: тип должен поддерживаться, тензор должен быть на CPU, и поле extra ещё не заполнено — иначе функция сразу возвращает управление. Затем берутся размеры, выбираются блочные константы по паре (m, k), вычисляется число тайлов, а масштаб выделяется отдельно и заполняется из хвоста данных весов. В разных вариантах ядер смещение до масштабов вычисляется по-разному: в одном случае просто k·m/4, в другом — через более сложную формулу с выравниванием до 32.

Перестановка весов при переупаковке отключается только для типа TL1, для всех остальных типов она выполняется. Transpose_8_8 — inline-функция для ARM NEON, выполняющая серию vzipq_s16 над восемью векторами int16, то есть переупорядочивающая элементы внутри блока 8×8.

Квантизация активаций

Активации проходят отдельный путь. Сначала находится максимум модуля по всем входным float-значениям, он сохраняется как масштаб. Затем каждый элемент превращается в один из трёх кодов: если модуль меньше 1e-6, пишется 1 (нуль); иначе пишется 2 для положительных и 0 для отрицательных — то есть знак кодируется как 2/0, а ноль как 1.

Далее байты упаковываются. В параллельной ветке для каждого блока из QK_I2_S элементов берётся группа и позиция внутри группы, и байт сдвигается на (6 − 2·group_idx) и OR-ится в нужную позицию. В ветке 1x4 (требуется nrow % 4 == 0) четыре строки упаковываются в один байт сдвигами 6, 4, 2, 0. Масштаб записывается как float сразу после упакованных данных, а возвращаемый размер включает дополнительные 32 байта под выравнивание.

Горизонтальное суммирование восьми int32 в одном векторе сворачивает 128-битные половины, затем 64-битные, затем перестановкой два оставшихся int32. В ядре скалярного произведения результат записывается как float.

Формы матриц и блочные размеры

Конкретное ядро выбирается по точному совпадению размеров m и k с заранее заданными константами. В tl1 это четыре формы: 14336×4096, 4096×14336, 1024×4096, 4096×4096. В tl2 добавляется разветвление по BK, и внутри каждой ветви — по размеру батча bs.

Формы 3200×8640, 3200×3200, 8640×3200 в диспетчеризации не обрабатываются; для 3200×8640 есть только шаблоны с соответствующими блочными константами.

Размер блока по строкам для конкретного ядра задаёт, сколько строк обрабатывается за один проход, а размер блока по K — шаг внешнего цикла по K. Например, в ядре 14336×4096 блок строк равен 256 и обрабатывается по 64 за шаг, а внешний цикл по K идёт шагами блока по K. Для 4096×14336 в tl1 блок строк 256, в tl2 — 128, что меняет и шаг по строкам, и вычисление KK. Директива #pragma unroll стоит перед циклами, границы которых выражены через эти константы, то есть разворачивание привязано к вычисленным из них границам.

Как накапливается результат

Функции qgemm_lut_* сначала обнуляют буфер целочисленных накоплений, затем в цикле по внешнему измерению вызывают соответствующую реализацию, передавая ей указатели на буфер, на срез таблицы и на срез активаций. Внутри таблица загружается в векторы, выборка выполняется по индексам из активаций, а накопление идёт в int16-векторы: результаты попарно складываются, затем расширяются до int32 и прибавляются к выходному буферу.

Масштаб таблицы применяется в самом конце: результат делится на LUT_Scales[0] и умножается на Scales[0]. Выбор конкретной реализации определяется размерами m и k, а в вариантах с батчем — ещё и BK.

Диспетчеризация и интеграция с ggml

Выбор ядра основан на точном совпадении размеров с константами; никаких проверок типов, форм или выравнивания внутри диспетчера нет. При несовпадении не выполняется ничего. В варианте с батчем выбор дополнительно зависит от BK и bs.

Проверка типа и выравнивания выполняется отдельно, при переупаковке тензора: если тип не поддерживается, backend не CPU или extra уже задан, функция возвращает управление без изменений. Поддерживаются только Q4_0 и TL1 (в одной версии) либо Q4_0 и TL2 (в другой).

Допустимость умножения проверяется отдельно: тип первого операнда должен поддерживаться, тип второго и результата — F32. В ARM-ветке дополнительно требуется, чтобы число столбцов второго операнда не превышало 1. Размер рабочего буфера вычисляется по формуле, зависящей от размеров и типа float, и округляется вверх до кратного 64. В ARM-ветке используется один множитель на int8, в x86-ветке — другой, плюс место под масштабы и смещения.

Барьер вызывается только при числе потоков больше 1 и служит для синхронизации: поток с номером 0 выполняет инициализацию общих данных, остальные ждут её завершения перед вычислением.

Жизненный цикл LUT-структур устроен так: инициализация выделяет массив под метаданные тензоров и сбрасывает индекс, освобождение удаляет массив и обнуляет указатель. Сами таблицы в этих функциях не создаются — они формируются в рабочем буфере во время умножения. Битность выбирается отдельно: для TL1 или TL2 возвращается 2, для Q4_0 — 4, иначе 0. Полученное значение передаётся в инициализацию и вычисление, определяя вариант обработки.

Векторные пути: ARM и x86

Общее скалярное произведение выбирает реализацию по числу строк: если оно делится на размер параллелизма, вызывается один вариант, иначе другой; во всех остальных случаях — базовый.

В базовом варианте на AVX2 данные читаются как байты, из каждого байта сдвигами вправо на 2, 4, 6 и маской извлекаются четыре 2-битных значения, а второй операнд читается четырьмя 32-байтовыми блоками. Произведения накапливаются через умножение с насыщением и сложение пар, затем горизонтальное суммирование сворачивает 8 int32 в одно число.

На ARM NEON те же четыре значения получаются сдвигами и маской. При наличии аппаратной поддержки скалярного произведения произведение сразу аккумулируется в int32; без неё — через умножение с накоплением в int16 с последующим расширением до int32.

В варианте на 4 строки на AVX2 обрабатываются сразу 4 строки, для каждой свой аккумулятор, а второй операнд загружается один раз. На ARM NEON этот путь оставлен пустым. В варианте на N строк на AVX2 обрабатывается сразу несколько строк, второй операнд загружается один раз, а первый читается для каждой строки; на ARM NEON этот путь также пуст.

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

Вспомогательные примитивы транспонирования

Функции слияния переставляют 64-битные блоки в двух входных векторах, а затем объединяют их — для 32-битных элементов, для 64-битных или для 128-битных половин. Транспонирование применяет их каскадом: сначала слияние по 32-битным элементам к парам входов, затем по 64-битным к промежуточным векторам, затем по 128-битным, записывая итог обратно.

В построителях таблицы эти транспонирования применяются к массиву из 16 векторов до упаковки и перестановок, после чего результат сохраняется. В ARM-версии аналогичное транспонирование строится из vzipq_s16 и применяется к двум половинам таблицы перед табличной выборкой и записью. Эти примитивы нужны, чтобы переупорядочить заранее вычисленные значения таблицы в раскладку, пригодную для последующей упаковки и табличных выборок.

Пограничные случаи и корректность

Обработки неполного блока нет: все циклы по строкам идут с шагом, равным блочной константе, и предполагают, что размер кратен ей. Внешние циклы по K также делят размер ровно, без отдельной ветви для хвоста. Постобработка идёт по всем элементам блока без проверки на неполноту. Диспетчер выбирает функцию только по точным парам (m, k), иных вариантов не предусмотрено.

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

Как сравнивали и что получилось

Для запуска бенчмарка используется скрипт, принимающий путь к модели, число генерируемых токенов, длину промпта и число потоков. Пример запуска задаёт 200 генерируемых токенов, промпт 256 и 4 потока. Отдельно упоминается команда проверки версии clang.

Параметры нагрузки: число генерируемых токенов по умолчанию 128, длина промпта по умолчанию 512, число потоков по умолчанию 2.

Известны только относительные величины:

  • ускорение декодирования до 1.18× на CPU и 1.27× на GPU;
  • ускорение до 1.28× относительно production-ядер тернарного матрично-векторного умножения при плотностях нулей, которые показывают реальные тернарные модели;
  • достижение 1.485 бита на вес на самой разреженной модели;
  • BITCOS хранит веса компактнее, чем five-trit packing, в 26 из 29 протестированных моделей.

Что из этого следует на практике

Барьер 1.58 бита — это не физический предел, а следствие формата. Теоретическая граница log2(3) ≈ 1.585 бита предполагает равновероятные символы, но реальные модели имеют до 51.5% нулей, и адаптивный формат это использует: стоимость 2 − z бит на вес падает с ростом плотности нулей, достигая 1.485 бита на самой разреженной модели. То есть чем разреженнее модель, тем выгоднее адаптивная раскладка по сравнению с упаковкой по пять тритов.

Практический выигрыш зависит от того, насколько распределение символов в конкретной модели отклоняется от равномерного. Формат, который трактует три символа как равновероятные, платит 1.625 бита независимо от того, сколько в модели нулей. Адаптивный формат платит меньше там, где нулей больше.

Ядро на таблице подстановки заменяет умножение-накопление на выборку и сложение: если активация 4-битная, все возможные произведения весов на неё можно посчитать заранее и разложить по 16 слотам. Это работает только при малом числе состояний активации — именно поэтому тернарные веса и квантованные активации образуют пару, для которой табличный подход имеет смысл.

Диспетчеризация жёсткая: ядро выбирается по точному совпадению размеров с константами, и никаких проверок форм или выравнивания внутри диспетчера нет. Неподдерживаемая форма просто не обрабатывается — ни ошибки, ни фолбэка. Это значит, что набор форм матриц, для которых есть ядра, определяет, какие модели можно запустить без изменения кода.

Где смотреть в коде

Источники

Похожее