Как устроены 4.6-битные сети: идея и имплементация18.06.2024 11:45

Мы уже писали о том, что предложили новую модель квантования нейронных сетей, позволяющую ускорить их на 40% на центральных процессорах, и теперь пришло время рассказать о ней подробнее.

Итак, квантование — это процесс преобразования вещественных весов сети и самого процесса вычислений таким образом, чтобы они использовали целые числа. В результате снижается объем сети в памяти (а это важно, например, для больших языковых моделей), а также могут ускориться вычисления за счет того, что больше значений помещаются в кэш и SIMD регистры.

Основные компоненты, которые нужны для создания и применения квантованной нейросети:

Схема квантования. Она задаёт то, как вещественные числа будут преобразовываться в целые и обратно, а также как мы будем умножать квантованные матрицы. Обычно вычисления в нейросети на центральных процессорах преобразуются именно в матричные операции, и наиболее вычислительно сложная из них — умножение.
Реализация матричного умножения. Само по себе то, что мы сделали числа в матрицах целыми, не сделает работу умножения быстрее. Необходима полноценная реализация быстрого матричного умножения, которая будет использовать кэши, SIMD и распараллеливание.
Метод обучения. Квантованные коэффициенты дискретны, поэтому обучать их обычным градиентным спуском затея несколько сомнительная.

Сегодня мы рассмотрим первые два компонента.

Про схемы квантования

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

$q = \min \left(q_{\max}, \max \left(q_{\min}, Z + \left[\frac{r}{S} \right]\right) \right),$

где $[\cdot]$ — это округление к ближайшему целому, — коэффициент масштабирования, — нулевая точка (целая), , — максимальное и минимальное значения целочисленных коэффициентов.

Тогда матричное умножение будет иметь вид:

$c_{ij} = \sum_k a_{ik} b_{kj} \approx S_a S_b \sum_k (\hat{a}_{ik} - Z_A) (\hat{b}_{kj} - Z_B)$

Введем обозначение:

$\hat{c}_{ij} = \sum_k (\hat{a}_{ik} - Z_A) (\hat{b}_{kj} - Z_B).$

Эту матрицу можно вычислить с использованием только целочисленной арифметики. В некоторых библиотеках ее вычисляют так:

$\hat{c}_{ij} = \sum_k \hat{a}_{ik} \hat{b}_{kj} - Z_A \sum_k \hat{b}_{kj} - Z_B \sum_k \hat{a}_{ik} + K Z_A Z_B,$

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

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

В случае 8-битных сетей, коэффициенты хранятся и загружаются в типе uint8, для хранения результата умножения используется int32. При этом оценим максимальную глубину матриц, которые можно обрабатывать без переполнения. Глубиной называют размерность, которая «схлопывается» при умножении. В сверточных слоях она будет равна произведению размерностей фильтра: ширина на высоту на число каналов.

Итак:

$(2^{31} - 1) / 255 / 255 \approx 33025.5$

Большинство моделей, которые мы вообще хотели бы исполнять на центральном процессоре, этому условию соответствуют.

Однако при дальнейшем снижении размерности возникает проблема: в случае 4-битных сетей коэффициенты могут принимать 16 значений (0–15) и для вычисления их произведений достаточно 8-битного значения, а для суммирования нужен уже 16-разрядный аккумулятор. Однако с ним максимальная глубина матрицы составит

$(2^{15} - 1) / 15 / 15 \approx 145.6$

Современные модели часто оперируют матрицами большей глубины, поэтому необходимо реализовать второй уровень суммирования с 32-битным аккумулятором. Эта схема работает быстрее, чем 8-битное умножение, но ограниченно применяется на практике из-за заметного снижения качества. Если 8-битные сети можно дообучить так, чтобы они работали не хуже вещественных, то с 4-битными этот трюк проходит не во всех задачах, поэтому массового перехода на них пока не произошло.

Однако мы заметили, что стандартное 4-битное квантование, оперирующее беззнаковыми числами 0…15, использует свой 8-битный регистр для произведений недостаточно эффективно (отметим, что в литературе в основном рассматриваются именно такие схемы квантования). Как можно сделать лучше?

4.6-битное квантование

Попробуем создать такую схему квантования, в которой произведение входов и весов помещается в 8-битный регистр:

$−128 ≤xw ≤ 127,\\ |x| ≤ x_{max}, \\ |w| ≤ w_{max}.$

У нас получилось целых для входа и целых для весов , где

$N_x = 2 x_{max} + 1,\\ N_w = 2w_{max} + 1.$

Таких пар значений ровно 21: (255, 3), (127, 5), (85, 7), (63, 9), (51, 11), (43, 13), (37, 15), (31, 17), (29, 19), (25, 21), (23, 23) и симметричные.

Средняя разрядность, приходящаяся на одной квантованное значение составила

$(\log_2 N_x + \log_2 N_w)/ 2 \approx 4.6.$

То есть, за счет знаковости входов и весов 4.6-битное квантование позволяет увеличить число уровней квантования с 16 до 23 (для схемы (23, 23)). Это в 23/16 = 1.468, т.е. почти в полтора раза уменьшает ошибку квантования. Визуально это различие показано на рис. 1.

Рис. 1. Число уровней квантования в 4- и 4.6-битных схемах квантования

Иллюстрация всех рассуждений выше приведена на рис. 2.

Рис. 2. Иллюстрация схем квантования для разных разрядностей.

Реализация матричного умножения

Поскольку распределение весов в нейронной сети обычно симметричное, мы взяли значение для весов равным 0, и это заметно упростило нам вычисления. Сумму столбцов матрицы весов мы вычисляли заранее, так как после обучения сети веса фиксированы.

$\hat{c}_{ij} = \sum_k \hat{a}_{ik} \hat{b}_{kj} - Z_A \sum_k \hat{b}_{kj}.$

Далее мы просто имплементировали высокопроизводительное матричное умножение для 4.6-битных квантованных чисел с на ARM и x86. Общий принцип реализации такого умножения заключается в том, что левая и правая матрица разбиваются на блоки, данные в которых переупорядочиваются и затем умножаются с помощью специальной оптимизированной функции — микроядра. На практике блоки матриц дополнительно ограничиваются по глубине. Цель этого процесса — снизить число промахов кэша и сделать данные последовательными в памяти, чтобы было удобнее их умножать с помощью интринсиков или ассемблерных команд.

Общая организация умножения для 4.6-битных матриц показана на рис. 3.

Рис. 3. Схема умножения 4.6-битных матриц

Микроядра мы написали на ассемблере, для ARM основное ядро сделали размера 24×8, чтобы использовать все 32 регистра, доступных на ARMv8. Дополнительно имплементировали микроядра 24×4, 1×8, 1×4, 24×1 and 1×1 для быстрой обработки матриц произвольного размера.

Приведем псевдокод основного микроядра 24×8.

Входные данные:
    k – глубина матричного умножения;
    A – 8-битный целый блок левой матрицы размера 24 х k,
      данные переупорядочены следующим образом:
        1) Первые 8 значений из первого столбца;
        2) Первые 8 значений из второго столбца;
        3) Следующие 8 значений из первого и второго столбцов; 
        4) (1-3) для всех оставшихся пар столбцов
           (если их количество нечетное – дополняем значением нуля Z);
    B – 8-битный целый блок правой матриц размера k х n,
      данные переупорядочены следующим образом:
       a) 2 значения из первого столбца, затем 2 значения из второго столбца и
          так до 8-го столбца; 
       b) повторяем пункт (a) для оставшихся пар строк (если их количество
          четное – дополняем значением нуля Z);

Выходные данные: 
    C – 32-битный целый блок результирующей матрицы размером 24 х 8    

  for (int r0 = 0; r0 < k; r0 += 258) {
    int8x16_t a[3], b;
    int8x8_t  t[4];
    int16x8_t c[3][8] = {0};
    int n = min(258, k - r0);

    for (int r1 = 0; r1 < n; r_1 += 2) {
      b = следующие 16 значений из B;
      a[0] = следующие 16 значений из A;
      a[1] = следующие 16 значений из A;
      a[2] = следующие 16 значений из A;

      for (int j = 0; j < 16; j += 4) {
          t[0] = vdup_laneq_s8(b, j + 0);
          t[1] = vdup_laneq_s8(b, j + 1);
          t[2] = vdup_laneq_s8(b, j + 2);
          t[3] = vdup_laneq_s8(b, j + 3);
          for (int i = 0; i < 3; i++) {
              c[i][j / 2 + 0] = vmlal_s8(c[i][j / 2 + 0], vget_low_s8(a[i]), t[0]);                    
              c[i][j / 2 + 0] = vmlal_s8(c[i][j / 2 + 0], vget_high_s8(a[i]), t[1]);
              c[i][j / 2 + 1] = vmlal_s8(c[i][j / 2 + 1], vget_low_s8(a[i]), t[2]);                  
              c[i][j / 2 + 1] = vmlal_s8(c[i][j / 2 + 1], vget_high_s8(a[i]), t[3]);                 
          }
      }
    } 
    Загрузка значений из C;    
    Добавление к ним c;
    Сохранение C;
  }

Поясним: 258 — это не магическая константа, а максимальная глубина без переполнений 16-битного типа данных для 4.6-битных сетей.

На x86 регистров меньше, всего 16, поэтому микроядро у нас получилось меньше, всего 8 на 8. Аналогично имплементировали еще микроядра 1×8, 8×1 и 1×1. Также там отсутствуют интринсики для умножения знаковых int8, так что мы использовали инструкцию pmaddubsw (интринсик _mm_maddubs_epi16) для умножения беззнаковых 8-битных целых на знаковые 8-битные целые и сложения соседних пар в знаковый 16-битный тип, и pshufb (интринсик _mm_shuffle_epi8) чтобы упаковать пары 8-битных чисел в регистры.

Для перехода к знаковым целым мы вычли минимальное возможное значение из правой матрицы:

$\hat{c}_{ij} = \sum_k \hat{a}_{ik} \hat{b}_{kj} - Z_A \sum_k \hat{b}_{kj} = \sum_k (\hat{a}_{ik} - A_m + A_m) \hat{b}_{kj} - Z_A \sum_k \hat{b}_{kj}= \\ = \sum_k(\hat{a}_{ik} - A_m) \hat{b}_{kj} - (Z_A - A_m) \sum_k \hat{b}_{kj},$

где — минимально возможное значение левой матрицы. Вычитание можно дешево добавить в квантование входа слоя или переупорядочивание блоков левой матрицы.

Экспериментальная оценка времени работы

Наконец перейдем к экспериментам. Мы измеряли время работы матричного умножения на x86 (AMD Ryzen 9 5950X) и ARM (ARM Cortex A-73).

Обозначим размеры первой матрицы как, а второй как. Размеры варьировались следующим образом: , и. Вычисления выполнялись в 1 поток, каждый тест повторили 100 раз и усреднили полученное время:

$\overline{t} = \mathbb{E}_{(H, W, D)} \frac{T(H, W, D)}{HWD},$

где обозначает среднее время запуска матричного умножения с параметрами. Времена приведены в таблице 1.

Таблица 1. Время работы матричного умножения.

Тип данных	Время на ARM (на умножение), ns	Время на x86 (на умножение), ns
32-битный float	0.2162 ± 0.0015	0.04355 ± 0.00004
8-битный квант.	0.1813 ± 0.0008	0.02963 ± 0.00004
4-битный квант.	0.1049 ± 0.0006	Реализация отсутствует
4.6-битный квант.	0.1089 ± 0.0008	0.02327 ± 0.00007

На ARM в среднем 4.6-битное умножение меньше, чем на 4% медленнее 4-битного. При этом и 4- и 4.6-битное умножение примерно в 1.7 раз быстрее, чем 8-битное и в 2 раза быстрее вещественного. На x86 результат скромнее: 4.6-битное умножение в 1.3 раза быстрее 8-битного и в 1.9 раза быстрее вещественного.

Дальше мы померяли время работы сети целиком. Для преобразования свертки в матричное умножение использовали алгоритм p-im2col. Первый и последний слои сети не квантовались. Батч нормализации и кусочно-линейные активации (ReLU, ReLU6, HardTanh и т.п.) интегрировались в сверточные слои.

Время мы также усреднили по 100 изображениям, которые были 3-канальными размера 32 на 32. Отметим, что 4-битных сетей в таблице нет, так как в опубликованной реализации нет 2-уровневого суммирования и есть ограничение на глубину свертки, а мы рассмотрели 6-слойную сверточную LeNet-подобную модель с 15.6k параметров и ResNet«ы достаточно большого размера. Результаты представлены в таблице 2.

Таблица 2. Время работы нейронных сетей целиком

Модель	#Параметров	float, ms	8-бит квант., ms	4.6-бит квант., ms
CNN6	15.6k	0.49 ± 0.02	0.362 ± 0.012	0.303 ± 0.008
ResNet-18	11.7M	415 ± 5	361.7 ± 1.7	242.8 ± 1.7
ResNet-34	21.8M	802 ± 7	699.6 ± 1.8	444.2 ± 1.6

Ожидаемо, что 4.6-битные модели ускорились. CNN6 ускорилась в 1.6 раза относительно вещественных моделей и в 1.2 раза относительно 8-битных. ResNet«ы ускорились еще лучше, так как размеры сверточных слоев там больше: они работают в 1.7–1.8 раза быстрее, чем вещественные модели и в 1.5–1.6 раза быстрее, чем 8-битные.

Заключение

Таким образом, мы в Smart Engines предложили 4.6-битную схему квантования, которая работает заметно быстрее 8-битной и позволяет представить больше значений, чем 4-битная, то есть является более точной. Для нас эта разработка открывает возможности по снижению разрядности в парке нейронных сетей без необходимости изобретать новые архитектуры и значительно менять процесс обучения (как было бы с бинарными или тернарными сетями), ведь для обучения 4.6-битных сетей можно использовать любые методы, подходящие для равномерных схем квантования. Как делаем именно мы — расскажем в следующей части.