rwkv-quant: одна матрица в bf16 снимает 39% потерь при квантовании RWKV-7 2.9B

Квантование моделей семейства RWKV с линейной рекуррентной структурой остаётся менее изученным, чем для трансформеров. Автор инструмента rwkv-quant, который сжимает RWKV-7 двумя пресетами (около 2.1x и около 3.1x), заметил аномалию. Агрессивный COMPRESSION на 2.9B даёт +4.04% по перплексии, а на 1.5B только +2.85%. Разбор причин показал, что значительная часть потерь сидит в одной матрице 2560×2560. Ее хранение в bf16 стоит 0.5% файла и снимает 39% KL.

Два пресета и шесть размеров

Исследуется линейка RWKV-7 из шести размеров. Три пары позволяют разделить влияние ширины и глубины: 0.4B/1.5B (одна глубина, ширина вдвое больше), 2.9B/7.2B (глубина 32, ширина 2560 против 4096) и 7.2B/13.3B (ширина 4096, глубина 32 против 61).

• 0.1B: 12 слоёв, ширина 768
• 0.4B: 24 слоя, ширина 1024
• 1.5B: 24 слоя, ширина 2048
• 2.9B: 32 слоя, ширина 2560
• 7.2B: 32 слоя, ширина 4096
• 13.3B: 61 слой, ширина 4096

Пресет REDUCTION почти без потерь: проекции в 8 бит с симметричной раскладкой и группой 16 весов, cmix (FFN) в 6 бит, эмбеддинг и голова в 8 бит. COMPRESSION агрессивнее: проекции в 4 бита с асимметричной суперблочной раскладкой и группой 32, cmix в 4 бита (ffn.key в 5), голова в 5 бит, эмбеддинг в 6 бит. LoRA-ветки в обоих пресетах идут в 6 бит (g в 8). В обоих масштаб и сдвиг групп подбираются поиском, взвешенным по среднему квадрату входной активации канала E[x²].

Где ломается 2.9B

Разложение KL по группам матриц методом leave-one-out показало, что на 2.9B o_proj даёт 46% всего KL, а эмбеддинг 17%. На 1.5B эти доли составляют 16% и 3%, и главный вклад вносит cmix.value (37%). Внутри o_proj 85–86% вреда приходится на одну матрицу, выходную проекцию нулевого блока.

Прямой эксперимент на 2.9B (реальный путь, MLX) подтвердил это:

• базовый пресет: KL 0.052048, +4.04% ppl;
• o_proj слоя 0 в bf16: KL 0.031966, +2.43% ppl, файл больше на 9.4 МБ;
• o_proj слоя 31 в bf16: эффекта нет (+4.07%);
• весь o_proj в 5 бит: +3.54%, +26.2 МБ;
• весь o_proj в 6 бит: +2.89%, +52.4 МБ.

Одна матрица в bf16 работает лучше, чем перевод всего o_proj в 6 бит, и стоит в пять с лишним раз меньше. На других размерах эффект слабее: KL падает на 7.6% на 0.1B, на 1.5% на 0.4B, на 0.4% на 1.5B и на 3.9% на 7.2B. Значимы только 0.1B и 2.9B.

Вход этой матрицы несёт «массивный канал». У 2.9B на восемь каналов приходится 98% энергии, номер канала 517 совпадает у двух серий чекпоинтов, у 1.5B это канал 1202. Значит, он определяется шириной и архитектурой, а не конкретным обучением. Но сила канала вред не предсказывает: у 7.2B концентрация выше всех (99%), а базовый KL ниже всех (0.023 против 0.078 у 2.9B). Без AW-взвешивания 4 бита в o_proj разрушают 2.9B: KL 0.340, +37.2% ppl.

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

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

Первый оценочный набор (38 окон по 512 токенов, en/ru/sr) целиком лежал внутри калибровочного корпуса: общих 16-грамм 100%. Масштаб утечки измерен на GPTQ. На 1.5B выигрыш по KL составил −61.8% на старом наборе и −27.1% на чистом, то есть утечка завысила результат в 2.3 раза. Автоподбор битности на ней не выиграл: −16.1% против −17.6%.

Поэтому разделы 7–10 работы используют только чистые наборы: Википедия en/ru/sr (36 окон), закрытый набор кода Python/Swift (24 окна) и открытый (48 окон). Интервалы строятся парным бутстрэпом по окнам, 20 000 повторов, полуширина 95%-интервала для KL составляет 7–12% от значения.

Поверх этого в работе описаны автоподбор битности по правилу «вред на байт», который работает в обе стороны, и GPTQ на сетке пресета без смены формата, ядер и размера файла. Эти методы складываются почти мультипликативно. Отдельно приведено сравнение с GGUF и MLX при равном размере файла.

Практически вывод простой. Если вы квантуете RWKV-7 для локального запуска, особенно на Mac через MLX, держите o_proj первого слоя в bf16: это 6–10 МБ на моделях 1.5B–2.9B. Эмбеддинг стоит держать минимум в 6 битах. Прежде чем доверять цифрам калиброванного метода, проверьте нулевое пересечение 16-грамм между калибровочным и оценочным текстом. Лучше мерить вклад отдельных групп матриц через leave-one-out, чем предполагать, что на большой модели потери распределятся равномерно.