Назад к блогу

Сжатие LLM до тернарных весов: как устроены форматы Intel BITCOS и Bonsai

Сжатие LLM до тернарных весов: как устроены форматы Intel BITCOS и Bonsai

Intel и Bonsai по-разному решают одну и ту же задачу — упаковать троичные веса (−1, 0, +1) так, чтобы модель занимала меньше памяти без потери точности. Теоретические 1.58 бита на вес на практике превращаются в 1.6 и выше, и статья разбирает, откуда берётся разница и при каких условиях схема на битовых картах обгоняет классическую пятитритную упаковку. Полезно тем, кто работает с квантизацией LLM и хочет понимать, что именно стоит за цифрами в описаниях форматов.

Тернарная модель хранит каждый вес всего в трёх состояниях: −1, 0 и +1. Это обещает радикальное сокращение памяти, но за «1.58 бита на вес» скрывается несколько разных механизмов — теоретический предел, реальная упаковка байтов, отдельные потоки знаков и пороги, при которых один формат обгоняет другой. Разберём, как эти механизмы устроены и где расходятся теория и практика.

Откуда берётся 1.58 бита и что такое тернарные веса

Теоретический минимум 1.58 бита — это оценка для трёх равновероятных значений веса −1, 0 и +1. Именно столько нужно, чтобы различить три равновероятных варианта.

Реальная упаковка тензора даёт больше из-за неполного использования байтов. Стандартный подход укладывает пять троичных значений в один восьмибитный байт, что даёт в среднем 1.6 бита на вес. Веса обычно хранят блоками по 128, но тогда последний байт остаётся частично неиспользованным, и фактическая скорость поднимается до 1.625 бита на вес.

Формат BITCOS («BITmap and COmpacted Signs») делит веса модели на два потока. Первый отводит один бит каждому весу, чтобы записать, нулевой он или ненулевой; второй отводит знаковый бит только ненулевым весам. Положительный или отрицательный вес поэтому занимает два бита, а ноль — только бит присутствия, потому что знака у него нет.

Если z — доля нулевых весов, BITCOS использует 2 − z бита на вес:

2 − z бита на вес

Это снижение идёт с 1.6 бита при 40% нулей до 1.485 бита при 51.5%. BITCOS становится меньше пятитритной упаковки, когда более 37.5% весов модели равны нулю, — порог, достигнутый в 26 из 29 проверенных чекпоинтов.

Распаковка восстанавливает исходные значения −1, 0 и +1 без потери точности, поскольку меняется только формат хранения. На AVX-512 ядро использует битовую карту присутствия как маску и pdep, чтобы разложить сжатые знаковые биты по позициям ненулевых весов. На GPU Xe2, где эквивалентной инструкции нет, та же операция реализована через таблицу поиска размером 2KB.

Форматы хранения: I2_S, TL1 и TL2

Три формата кодируют одни и те же троичные веса по-разному. I2_S упаковывает троичные значения по четыре в байт и добавляет 32-байтовый хвост, в первых четырёх байтах которого лежит f32 scale. TL1 и TL2 устроены иначе: они переупорядочивают веса в блочную раскладку и по-разному кодируют знаки.

В TL2 пары обрабатываются отдельной функцией, а для троек первый элемент умножается на 9, второй на 3, третий берётся как есть, и знаки кодируются отдельно.

Отдельные ветки для двух и трёх весов нужны потому, что при трёх значениях на блок (K//3) появляется дополнительный знаковый поток, который упаковывается в 16-битные слова, тогда как при двух значениях (K//2) достаточно одной байтовой упаковки.

Как работает quantize_to_i2_s

Функция принимает матрицу весов формы (M, K) и необязательный override_scale. Если он задан, он напрямую используется как scale, а веса уже тернарные. Иначе scale вычисляется как первое ненулевое абсолютное значение, а при отсутствии ненулевых берётся 1e-5; затем веса округляются с ограничением в диапазон [−1, 1].

Тернарные значения отображаются в коды I2_S. Упаковка идёт блоками по 128 значений, дополненными до кратности 128 значением 1. Блок преобразуется в форму (n_blocks, 4, 32) и упаковывается по четыре значения на байт.

Итоговый размер равен packed_size + 32, где packed_size = n // 4, а scale записывается в первые четыре байта хвоста. Хвост выравнивается на 32 байта, потому что при чтении вычисляется nbytes и округляется вверх до кратности 32:

int nbytes = (k - 256) * m / 3 * 5 / 8 + 256 * m / 2 * 4 / 8;
if (nbytes % 32 != 0) nbytes = 32 - nbytes % 32 + nbytes;

После этого scale читается из позиции qweights + nbytes.

Конвейер конвертации HF → GGUF

Путь от safetensors до GGUF разбит на несколько шагов, и порядок их вызова важен.

get_tensors открывает каждый файл модели и по очереди отдаёт пары (имя, тензор). modify_tensors получает один тензор и его имя и возвращает уже переименованные или преобразованные тензоры: для q_proj и k_proj веса переставляются так, чтобы порядок голов внимания совпал с раскладкой, ожидаемой в GGUF, — данные переставляются по числу голов внимания (для k_proj — по числу голов ключей и значений). Эксперты block_sparse_moe.experts накапливаются и, когда набрано n_experts*3, склеиваются в 3d-тензоры w1/w2/w3. generate_extra_tensors не читает исходные тензоры, а порождает дополнительные: при rope_scaling с типом llama3 вычисляются rope_factors и отдаётся тензор ROPE_FREQS. prepare_tensors в BitnetModel только вызывает родительскую реализацию, а затем проверяет, что все эксперты обработаны, — иначе выбрасывается ValueError.

Порядок задан write_tensors: сначала идёт проход по get_tensors для сбора weight_scale, затем второй проход по get_tensors, внутри которого для каждого тензора вызывается modify_tensors, а generate_extra_tensors и prepare_tensors вызываются отдельно от этого цикла.

Обработка офлайн-квантованных моделей

Ветка «Handle offline-quantized weights (uint8 packed with weight_scale)» срабатывает, когда имя тензора без суффикса «.weight» присутствует в scale_map, то есть для весов, у которых рядом лежит отдельный тензор weight_scale. В этом случае тензор приводится к uint8, затем распаковывается: тензор расширяется по новой оси до четырёх копий, сдвигается вправо на величины, заданные сдвигом, и маскируется & 3, после чего преобразуется в float и вычитается 1, что даёт значения из множества {−1,0,1}.

Если целевой формат не I2_S, TL1 или TL2, распакованные значения делятся на соответствующий weight_scale, чтобы получить полные float-значения. Для I2_S/TL1/TL2 значения остаются тернарными, а масштаб передаётся отдельно: в ветке I2_S из scale_map берётся оригинальный weight_scale, и если он есть, его значение передаётся в quantize_to_i2_s как override_scale. При наличии override_scale этот масштаб используется напрямую как scale, а веса берутся как уже тернарные без повторного округления; при отсутствии override_scale масштаб вычисляется как первое ненулевое абсолютное значение, и веса округляются к {−1,0,1}.

Почему неквантованные 2D-веса сохраняются как F16

При квантовании I2_S/TL1/TL2 неквантованные 2D-веса (например, embedding) сохраняются как F16, потому что для них срабатывает ветка по умолчанию: если тип файла — I2_S, TL1 или TL2, размерность n_dims ≥ 2 и имя не оканчивается на «_norm.weight», данные приводятся к np.float16 и тип устанавливается в F16. Это уменьшает размер файла вдвое по сравнению с F32 для всех таких 2D-тензоров, поскольку каждый элемент занимает 2 байта вместо 4.

Одномерные тензоры и нормы (включая «_norm.weight») в эту ветку не попадают и остаются F32, так как для них условие n_dims ≥ 2 или not new_name.endswith(«_norm.weight») не выполняется. В generate-dummy-bitnet-model.py аналогичная логика: extra_f32 включает n_dims == 1 и new_name.endswith(«_norm.weight»), а extra_f16 включает name.endswith(«.weight») и n_dims ≥ 2, поэтому 2D-веса идут в F16, а 1D и нормы — в F32.

Обработка экспертов (MoE)

Ветка «process the experts separately» в modify_tensors перехватывает тензоры, чьё имя содержит подстроку «block_sparse_moe.experts»: она берёт число экспертов из hparams, требует, чтобы номер слоя был задан, и лениво создаёт self._experts как список словарей по числу слоёв. Каждый пришедший тензор эксперта складывается в self._experts[bid][name], а когда в словаре слоя набирается не меньше n_experts * 3 записей, тензоры объединяются: для каждого из «w1», «w2», «w3» собираются данные всех экспертов, удаляются из словаря и складываются в один 3d-тензор через torch.stack(dim=0). Объединённому тензору даётся новое имя, оно прогоняется через map_tensor_name, и результат возвращается как список пар (имя, тензор); пока порог не достигнут, ветка ничего не отдаёт.

В prepare_tensors остаточные эксперты проверяются: если self._experts не None, список словарей «уплощается» в список строк, и если он непуст, выбрасывается ValueError — это защита от того, что какие-то тензоры экспертов не были слиты и остались необработанными.

Генерация LUT-ядер: codegen_tl1 и codegen_tl2

LUT-ядра генерируются скриптами, которые подставляют в C-шаблоны конкретные размеры блоков и собирают готовый код.

В codegen_tl1.py gen_ctor_code возвращает строку C-кода с вспомогательными функциями: per_tensor_quant вычисляет масштаб как 127 / максимум модуля по k элементам, partial_max_reset обнуляет масштаб, lut_ctor строит таблицу подстановки, is_type_supported разрешает только GGML_TYPE_Q4_0 и GGML_TYPE_TL1. gen_body_core_code(bm, by) генерирует тело внутреннего цикла: для i от 0 до 3 загружает uint8x16_t, делит байты на старшие и младшие полубайты, делает четыре vqtbl1q_s8 по индексам, затем vzipq_s8 и накапливает в vec_c. gen_tbl_impl выдаёт макросы размеров блоков и объявляет реализацию таблицы, а gen_top_api генерирует диспетчеры с ветками по (m, k). gen_transform_code создаёт ggml_bitnet_transform_tensor, который проверяет поддержку типа, backend и extra == nullptr, берёт k и m из размерностей тензора, выбирает bm и bk, вычисляет n_tile_num, выделяет scales и читает i2_scales из qweights + k*m/4.

В codegen_tl2.py gen_tbl_impl генерирует three_tbl_impl с KK = BBK/3 и two_tbl_impl с KK = BK2/2, а также соответствующие qgemm-функции, где границы циклов задаются через k_list.

Параметры BM, BK, bm задаются аргументами командной строки --BM, --BK, --bm как списки через запятую, и assert требует, чтобы их длины совпадали с числом kernel_shapes.

Константы и размеры блоков

GGML_BITNET_MAX_NODES 8192 — это константа, объявляемая в генерируемом коде ядра, она задаёт верхнюю границу числа узлов, под которое резервируется память под структуры bitnet_tensor_extra.

Размеры блоков BM задаются макросами вида #define BM3200_3200 320 и #define BM8640_3200 320, то есть в обоих случаях BM равен 320. Эти макросы используются как шаг внешнего цикла по строкам матрицы. В dummy-модели фигурируют другие значения: BMEMD = 256, BYEMD = 256, bmEMD = 32, byEMD = 8, а также BMGQA = 256, BYGQA = 256, bmGQA = 32, byGQA = 8.

Разбиение матрицы весов выполняется последовательно: сначала по строкам с шагом BM, затем по столбцам с шагом BY, затем снова по строкам с шагом bm, затем по столбцам с шагом by, и далее на минорные блоки 16 и 4. В ядрах с тремя LUT дополнительно задаётся BBK3200_3200 96 и BBK8640_3200 96, а внутренний параметр KK вычисляется как BBK/3, что определяет число итераций по k при обработке блока.

Как устроены two_lut_ctor и three_lut_ctor

Обе функции — шаблоны по act_k, принимающие указатель на квантованную таблицу qlut, массив весов b и масштаб lut_scales, и заполняющие qlut значениями LUT.

three_lut_ctor обрабатывает по три веса за итерацию: цикл идёт с шагом act_k / 24, а vec_bi задан как _mm256_set_epi32(84, 72, 60, 48, 36, 24, 12, 0), то есть собирает три группы по 8 значений. two_lut_ctor обрабатывает по два веса: цикл идёт с шагом act_k / 16, а vec_bi задан как _mm256_set_epi32(56, 48, 40, 32, 24, 16, 8, 0), то есть собирает две группы по 8 значений.

Соответственно, в three_lut_ctor вычисляются три вектора vec_b0i, vec_b1i, vec_b2i и строятся 16 записей vec_lut[0..15] как суммы и разности этих трёх значений, а в two_lut_ctor вычисляются только vec_b0 и vec_b1 и записи строятся из двух значений, причём часть записей (vec_lut[9..15]) обнуляется. Это отражает разное число ненулевых весов в блоке: three_lut_ctor рассчитан на блок с тремя ненулевыми весами (три таблицы), two_lut_ctor — на блок с двумя ненулевыми весами (две таблицы).

Ядра на стороне C

Функции three_tbl_impl_ и two_tbl_impl — это AVX2-ядра, которые накапливают в int32-буфере c частичные суммы произведения активированных значений: они читают упакованные байты активаций a (и знаки sign для three-варианта), выбирают из LUT нужные 16-байтовые фрагменты через _mm256_shuffle_epi8 и складывают результаты в 16-битные аккумуляторы, которые затем расширяются до 32 бит и прибавляются к c.

В three-вариантах дополнительно применяется знак: сдвигами формируются маски, и результат корректируется как _mm256_xor_si256(_mm256_add_epi16(x, sign), sign). Различие между three_tbl_impl_3200_8640, three_tbl_impl_3200_3200 и three_tbl_impl_8640_3200 — только в константах BM/BBK (160/96, 320/96, 320/96 соответственно) и в шаге индексации c по BM, а two_tbl_impl3200_8640 и two_tbl_impl3200_3200 отличаются константой BM3200_8640 против BM3200_3200.

Обёртки three_qgemm_lut_3200_8640 и two_qgemm_lut_3200_8640 обнуляют CBits, вызывают соответствующее ядро по k_outer, а затем переносят результат в C; two_qgemm_lut_3200_8640 дополнительно делит на LUT_Scales[bs] и умножает на Scales[0]. В three_qgemm_lut_3200_8640 кода, который бы применял LUT_Scales или Scales, нет.

Диспетчеры и передача lut_scales_size

ggml_preprocessor для каждого элемента батча выполняет пер-тензорное квантование весов B в масштабы LUT_Scales и строит две таблицы подстановки: three_lut_ctor и two_lut_ctor. Она выбирает ветку по тройке (m, two_k, three_k): например, при m == 3200 && two_k == 0 && three_k == 8640 вызывает three_lut_ctor<8640> и two_lut_ctor<0>, а при m == 3200 && two_k == 32 && three_k == 3168 — three_lut_ctor<3168> и two_lut_ctor<32>.

ggml_qgemm_lut — это диспетчер матричного умножения с таблицами подстановки: по значениям m, k и BK он выбирает конкретную реализацию и разворачивает её по размеру батча bs (1, 8, 32, 128, 256, 512).

ggml_bitnet_transform_tensor при первом вызове для поддерживаемого типа и при tensor->extra == nullptr заполняет структуру bitnet_tensor_extras: вычисляет n_tile_num = m / bm, BK = bk, берёт qweights = tensor->data, читает i2_scales из хвоста буфера и сохраняет в extra поля lut_scales_size, BK, n_tile_num, qweights, scales.

lut_scales_size — это число масштабов таблицы подстановки, которое резервируется в структуре bitnet_tensor_extras. В сгенерированном коде оно объявлено как const int lut_scales_size = 1; и записывается в extra. В конвертере отдельный тензор масштаба добавляется как new_name + «_scale» с типом F32.

AVX2-интринсики и платформы без __AVX2__

Функции _mm256_merge_epi32, _mm256_merge_epi64 и _mm256_merge_si128 определены внутри блока #if defined __AVX2__ и служат вспомогательными шагами для Transpose_8_8, которая переставляет восемь 256-битных векторов.

_mm256_merge_epi32 сначала переставляет 64-битные элементы каждого входа через _mm256_permute4x64_epi64 с маской _MM_SHUFFLE(3, 1, 2, 0), затем чередует младшие и старшие 32-битные половины через _mm256_unpacklo_epi32 и _mm256_unpackhi_epi32. _mm256_merge_epi64 делает то же самое, но чередует 64-битные половины, а _mm256_merge_si128 объединяет 128-битные половины двух векторов через _mm256_permute2x128_si256.

Transpose_8_8 применяет эти три функции последовательно: сначала четыре вызова _mm256_merge_epi32 к парам входных векторов, затем четыре вызова _mm256_merge_epi64, затем четыре вызова _mm256_merge_si128, записывая итог обратно в v0..v7. Она вызывается в three_lut_ctor и two_lut_ctor дважды — для ix[0..7] и ix[8..15] — перед упаковкой через _mm256_packs_epi32 и перестановками.

На платформах без __AVX2__ весь этот код исключён препроцессором: определения функций находятся внутри #if defined __AVX2__, а тела three_lut_ctor и two_lut_ctor также обёрнуты в этот блок, поэтому при отсутствии макроса они не компилируются и возвращают 0.

Квантизация и точность

per_tensor_quant(k, lut_scales_, b_) вычисляет один общий масштаб для тензора: он берёт максимум модулей элементов b (через vabsq_f32/vmaxq_f32 на NEON или _mm256_andnot_ps/_mm256_max_ps на AVX2) и записывает в lut_scales значение 127 / этот максимум. partial_max_reset обнуляет масштабы: в TL1-версии просто lut_scales = 0.0, а в TL2-версии циклом по bs элементам.

В transform_to_tl1/tl2 масштаб считается как scale = np.max(np.abs(x)), после чего вызывается preprocess_weights_tl1(x) или preprocess_weights_tl2(x). В preprocess_weights_tl1 значения сначала приводятся к знаку, затем пары элементов объединяются как hi_weight*3 + lo_weight, после чего weight = weight + 4, что и даёт двухбитные коды: значения −1, 0, 1 после sign и сдвига +4 попадают в диапазон, представимый двумя битами.

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

Стенд и методология

Сквозные бенчмарки Intel охватывали семь моделей при размере батча один. Все пять тестовых систем использовали оборудование Intel. Формат ещё не тестировался на оборудовании Nvidia, AMD или Arm. Статья не прошла рецензирование.

Для запуска бенчмарка в README описан скрипт e2e_benchmark.py. Он принимает обязательный аргумент -m/--model (путь к файлу модели) и необязательные: -n/--n-token (число генерируемых токенов, по умолчанию 128), -p/--n-prompt (число токенов промпта, по умолчанию 512), -t/--threads (число потоков, по умолчанию 2). Пример запуска: генерация 200 токенов из промпта в 256 токенов с использованием 4 потоков.

Результаты

По BITCOS: порог доли нулей 37.5% был достигнут в 26 из 29 проверенных чекпоинтов. При доле нулей 40% формат использует 1.6 бита на вес, при 51.5% — 1.485 бита.

По Bonsai 2 27B:

ПоказательЗначениеУсловия
Пропускная способностьдо 143 tokens/secondNVIDIA GeForce RTX 5090
Пропускная способность46.8 tokens/secondM5 Max
Энергопотребление0.714 mWh/tokenRTX 4090
Энергоэффективностьна 40% лучшеRTX 4090, сравнение с 8B-моделью в полной точности
Сохранение качествас 95% до более чем 98%сравнение с Ternary Bonsai 27B

Bonsai 2 27B работает на NVIDIA GPU через CUDA и на устройствах Apple (Mac, iPhone, iPad) через MLX, через собственные low-bit ядра.

По бенчмаркам Bonsai 2 27B в разделе Vision перечислены CharXiv, A-OKVQA, OmniDocBench v1.6, RealWorldQA, OCRBench v2 и приведены значения 78.59, 81.64, 79.82, без указания, какое число какому бенчмарку соответствует. Агрегированный балл Bonsai 2 27B равен 83.9, у Qwen3.8 27B — 85.4, у Qwen3.6 27B — 83.6, что даёт сохранение 98.2% агрегированной производительности Qwen3.8 27B. Отдельно утверждается, что Bonsai 2 27B сохраняет значительную часть производительности полной модели в областях coding agents, tool-use, multimodal workflows и long-horizon tasks.

Что эти цифры не показывают

Цифры BITCOS не показывают поведение на не-Intel оборудовании и при других размерах батча: все пять тестовых систем использовали оборудование Intel, сквозные бенчмарки шли при размере батча один, а формат ещё не тестировался на Nvidia, AMD или Arm. Статья не прошла рецензирование.

Практика: сборка, запуск и воспроизведение

Сначала создайте окружение командой setup_env.py: она принимает --hf-repo/-hr (модель для инференса), --model-dir/-md (каталог сохранения/загрузки модели), --log-dir/-ld (каталог логов), --quant-type/-q (тип квантования i2_s или tl1), --quant-embd (квантовать эмбеддинги в f16) и --use-pretuned/-p (использовать предварительно настроенные параметры ядра).

Затем сгенерируйте тестовую модель через utils/generate-dummy-bitnet-model.py, указав позиционный путь модели, --outfile, --outtype и --model-size. Для запуска инференса используйте run_inference.py: -m/--model — путь к файлу модели, -n/--n-predict — число предсказываемых токенов, -p/--prompt — обязательный промпт, -t/--threads — число потоков, -c/--ctx-size — размер контекста промпта, -temp/--temperature — температура, -cnv/--conversation — режим чата.

Для бенчмарка e2e_benchmark.py обязателен -m/--model, а -n/--n-token по умолчанию 128, -p/--n-prompt по умолчанию 512, -t/--threads по умолчанию 2. Пример запуска:

python utils/generate-dummy-bitnet-model.py models/bitnet_b1_58-large --outfile models/dummy-bitnet-125m.tl1.gguf --outtype tl1 --model-size 125M
python utils/e2e_benchmark.py -m models/dummy-bitnet-125m.tl1.gguf -p 512 -n 128

Требования к окружению

README предлагает создать новое conda-окружение с именем bitnet-cpp и Python версии 3.10 командами conda create -n bitnet-cpp python=3.10 и conda activate bitnet-cpp, после чего установить зависимости через pip install -r requirements.txt.

Проверка правильности версии clang и доступности инструментов Visual Studio выполняется командой clang -v; при ошибке выводится сообщение. На Windows окружение инициализируется запуском VsDevCmd.bat с параметрами -startdir=none -arch=x64 -host_arch=x64 по пути C:\Program Files\Microsoft Visual Studio\2022\Professional\Common7\Tools\VsDevCmd.bat. Для Windows PowerShell README указывает выполнить соответствующие команды.

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

Теоретический минимум 1.58 бита недостижим при стандартной упаковке: пять тритов в байт дают 1.6 бита, а выравнивание блоков по 128 поднимает фактическую скорость до 1.625 бита. Выигрыш BITCOS целиком зависит от доли нулей в конкретной модели: ниже порога 37.5% формат проигрывает пятитритной упаковке, выше — выигрывает, и этот порог был достигнут в 26 из 29 проверенных чекпоинтов. То есть разреженность тернарных весов — не абстрактное свойство, а измеримая величина, от которой зависит выбор формата.

Для I2_S масштаб можно либо вычислить из весов (первое ненулевое абсолютное значение), либо взять готовый weight_scale офлайн-квантованной модели и передать его как override_scale. Во втором случае веса уже тернарные и повторного округления не происходит — это позволяет сохранить масштаб, использованный при исходной квантизации.

Выбор формата влияет и на неквантованные тензоры: при I2_S/TL1/TL2 двумерные веса вроде embedding сохраняются как F16, что вдвое меньше F32, тогда как одномерные тензоры и нормы остаются F32. Это заметно на итоговом размере файла.

Генерация LUT-ядер привязана к конкретным размерам блоков: BM, BK и bm задаются аргументами командной строки, а assert требует, чтобы их длины совпадали с числом kernel_shapes. Ядра собираются под AVX2: без этого макроса тела конструкторов таблиц не компилируются и возвращают 0, то есть на платформах без AVX2 соответствующие пути просто недоступны.

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

Источники

Похожее