Эффективное развертывание низкобитных больших языковых моделей на CPU
Ключевая инновация T-MAC заключается в использовании вычислительной парадигмы на основе таблиц поиска (LUT) вместо традиционной парадигмы умножения с накоплением (MAC). T-MAC использует таблицы поиска для прямой поддержки низкобитных вычислений, тем самым устраняя необходимость в операциях деквантизации, обязательных в других системах, и значительно сокращая количество операций умножения и сложения.
Эксперименты показали превосходную производительность T-MAC: на ПК Surface AI с новейшим чипсетом Qualcomm Snapdragon X E lite скорость генерации для 3B BitNet-b1.58 модели достигает 48 токенов в секунду, для 2-битной 7B llama модели - 30 токенов в секунду, а для 4-битной 7B llama модели - 20 токенов в секунду.
Это даже превосходит производительность NPU!
При развертывании модели llama-2-7b-4bit, хотя NPU может генерировать 10.4 токена в секунду, CPU с помощью T-MAC достигает 12.6 токенов в секунду, используя только два ядра, и может достичь максимума в 22 токена в секунду.
Все это значительно превышает среднюю скорость чтения человека и в 4-5 раз быстрее, чем оригинальный фреймворк llama.cpp.
Даже на устройствах более низкого уровня, таких как Raspberry Pi 5, T-MAC для 3B BitNet-b1.58 может достичь скорости генерации 11 токенов в секунду. T-MAC также имеет значительные преимущества в энергопотреблении: для достижения той же скорости генерации T-MAC требует только 1/4-1/6 количества ядер по сравнению с оригинальным llama.cpp, снижая энергопотребление и оставляя вычислительные ресурсы для других приложений.
Стоит отметить, что вычислительная производительность T-MAC линейно увеличивается с уменьшением количества битов, что трудно наблюдать в GPU и NPU, основанных на деквантизации. Однако T-MAC может достичь 10 токенов в секунду на одном ядре и 28 токенов в секунду на четырех ядрах при 2 битах, значительно превосходя производительность NPU.
Умножение матриц без умножения, только поиск по таблице (LUT)
Для низкобитных параметров (весов) T-MAC группирует каждый бит отдельно (например, группа из 4 битов), эти биты умножаются на вектор активации, все возможные частичные суммы предварительно вычисляются и затем сохраняются с использованием LUT.
Затем T-MAC использует операции сдвига и накопления для поддержки масштабируемого количества битов от 1 до 4. Этим методом T-MAC отказывается от неэффективных на CPU инструкций FMA (умножение-сложение) в пользу более энергоэффективных и производительных инструкций TBL/PSHUF (поиск по таблице).
Вычисления на основе битов вместо вычислений на основе типов данных
Традиционные вычисления, основанные на деквантизации, фактически являются вычислениями, основанными на типах данных, что требует индивидуальной настройки для каждого различного типа данных.
Каждая комбинация битовой ширины активации и весов, такая как W4A16 (веса int4, активация float16) и W2A8, требует специфической компоновки весов и вычислительного ядра.
Например, компоновка W3 требует отдельной упаковки 2 битов и 1 бита, используя различные методы чередования или перемешивания для выравнивания памяти или быстрого декодирования. Затем соответствующее вычислительное ядро должно распаковать эту специфическую компоновку в типы данных, поддерживаемые оборудованием, для выполнения.
T-MAC, рассматривая низкобитное умножение матриц с точки зрения битов, требует оптимальной структуры данных только для одного бита, а затем расширяет ее до более высоких 2/3/4 битов путем наложения.
В то же время, для векторов активации различной точности (float16/float32/int8) изменяется только процесс построения таблицы, при поиске по таблице больше не нужно учитывать различные структуры данных.
Кроме того, в традиционных методах, основанных на деквантизации, при снижении с 4 битов до 3/2/1 битов, хотя использование памяти уменьшается, объем вычислений не уменьшается, и из-за увеличения накладных расходов на деквантизацию производительность может даже ухудшиться.
Но объем вычислений T-MAC линейно уменьшается с уменьшением количества битов, что приводит к лучшему ускорению при меньшем количестве битов, обеспечивая эффективное решение для развертывания 1-битных/2-битных моделей, выпущенных в последних работах, таких как BitNet, EfficientQAT и др.
Высокооптимизированная реализация операторов
Вычисления на основе битов имеют много преимуществ, но их реализация на CPU все еще представляет значительные проблемы:
(1) По сравнению с последовательным доступом к данным активации и весов, доступ к таблице является случайным. Размещение таблицы в быстрой встроенной памяти особенно важно для конечной производительности вывода;
(2) Однако встроенная память ограничена, и метод поиска по таблице (LUT) увеличивает использование встроенной памяти по сравнению с традиционным mpGEMV. Это потому, что таблица поиска должна сохранять результаты умножения вектора активации со всеми возможными битовыми паттернами. Это намного больше, чем сама активация.
Для решения этих проблем исследователи из Microsoft Research Asia глубоко изучили поток данных вычислений на основе таблиц и разработали эффективные структуры данных и вычислительные процессы для этой вычислительной парадигмы, включая:
-
Хранение LUT во встроенной памяти для использования векторных инструкций поиска по таблице (TBL/PSHUF) на CPU для повышения производительности случайного доступа к памяти.
-
Изменение порядка вычислений осей матрицы для максимального повышения коэффициента повторного использования данных ограниченных LUT, помещенных во встроенную память.
-
Разработка оптимального метода разбиения матрицы (Tiling) специально для поиска по таблице, в сочетании с поиском оптимальных параметров разбиения с помощью autotvm.
-
Оптимизация компоновки параметров весов
a) Перестановка весов для максимально последовательного доступа и повышения попаданий в кэш
b) Чередование весов для повышения эффективности декодирования
- Целевая оптимизация для CPU Intel/ARM, включая
a) Перестановку регистров для быстрого создания таблицы поиска
b) Быстрое 8-битное накопление с использованием инструкции усреднения
Исследователи применяли различные оптимизации шаг за шагом на базовой реализации, в итоге достигнув значительного ускорения по сравнению с современными низкобитными операторами: