Автор | Го Ран
Операция Softmax — одна из наиболее часто используемых операций в моделях глубокого обучения. В задаче классификации глубокого обучения окончательный классификатор сети часто представляет собой комбинацию Softmax + CrossEntropy:
Хотя математический вывод может быть уменьшен при совместном использовании Softmax и CrossEntropy, все еще существует множество сценариев, в которых Softmax Op используется по отдельности. Например, на уровне внимания каждого уровня кодировщика BERT Softmax используется отдельно для решения вероятностного распределения внимания; многоголовая часть внимания GPT-2 также использует только Softmax и так далее.
Все вычислительные операторы в среде глубокого обучения преобразуются в функции ядра CUDA на графическом процессоре, и операция Softmax не является исключением. Softmax — это оператор, широко используемый в большинстве сетей, и эффективность реализации его ядра CUDA повлияет на конечную скорость обучения многих сетей. Так как же реализовать эффективное ядро Softmax CUDA? В этой статье будут представлены методы оптимизированного ядра CUDA Softmax в OneFlow и сравнение его с операцией Softmax в cuDNN.Результаты показывают, что использование пропускной способности памяти Softmax после глубокой оптимизации OneFlow может быть близко к теоретическому верхний предел, который намного выше, чем у cuDNN.
\
Основы GPU и принципы оптимизации производительности CUDA
\
Введение в основы графического процессора, а также принципы и цели оптимизации производительности CUDA см. в предыдущих статьях:
zhuanlan.zhihu.com/p/271740706
\
Он кратко представляет аппаратную структуру и принцип работы графического процессора:
- Ядро: Функция ядра CUDA — это основная единица описания вычислительных задач графического процессора. Каждое ядро будет выполняться параллельно многими потоками на графическом процессоре в соответствии с параметрами конфигурации.Вычисления на графическом процессоре эффективны, поскольку одновременно могут выполняться тысячи ядер (потоков), а эффективность вычислений намного превышает эффективность ЦП.
- Потоки графического процессора логически разделены на три уровня: поток, блок и сетка, а аппаратное обеспечение разделено на ядро,warpдва уровня;
- Память графического процессора разделена на три уровня: глобальная память, общая память и локальная память.
- GPU в основном предоставляет два вида ресурсов:вычислительные ресурсы и ресурсы пропускной способности памяти. Если мы сможем в полной мере использовать эти два ресурса, а потребность в ресурсах уменьшить нельзя, то производительность будет оптимизирована до предела, а время выполнения будет кратчайшим. В большинстве случаев вычислительные ресурсы графического процессора при обучении глубокому обучению используются полностью, а ядро графического процессора CUDAоптимизировать цельЭто должно максимально использовать ресурсы пропускной способности памяти.
Как оценить, полностью ли ядро CUDA использует ресурсы пропускной способности памяти?
Для ресурсов пропускной способности видеопамяти «полное использование» означает, что эффективная пропускная способность видеопамяти для чтения и записи ядра достигла пропускной способности видеопамяти устройства.верхний предел, где пропускную способность памяти устройства можно получить, выполнив пропускную способность в cuda. Эффективная пропускная способность памяти ядра оценивается объемом данных, прочитанных и записанных ядром, и временем выполнения ядра:
Реализация Naive Softmax:
Прежде чем представить методы оптимизации, давайте посмотрим, какова максимальная теоретическая пропускная способность неоптимизированного ядра Softmax. Как показано на рисунке ниже, в простейшей реализации расчета Softmax для выполнения общего расчета вызывается несколько базовых функций ядра CUDA:
Предполагая, что размер входных данных равен D, shape = (num_rows, num_cols), то есть D = num_rows * num_cols, большинство операций Navi будут обращаться к глобальной памяти несколько раз, где:
- ReduceMax = D + num_rows (D для чтения, num_rows для записи)
- BroadcaseSub = 2 * D + num_rows (чтение — это D + num_rows, запись — это D)
- Exp = 2 * D (чтение и запись равны D)
- ReduceSum = D + num_rows (D для чтения, num_rows для записи)
- BroadcastDive = 2 * D + num_rows (чтение равно D + num_rows, запись равно D)
Всего требуется 8 * D + 4 * num_rows накладных расходов на выборку. Поскольку num_rows можно игнорировать по сравнению с D, версия ядра CUDA Softmax для Navie должна иметь доступ как минимум к 8-кратному объему видеопамяти данных, а именно:Для видеокарт GeForce RTX™ 3090 теоретический предел пропускной способности составляет936ГБ/с, то верхняя граница пропускной способности памяти, используемой версией ядра CUDA Softmax для военно-морского флота, составляет 936/8 = 117 ГБ/с.
в статье
zhuanlan.zhihu.com/p/271740706Здесь мы находимся в методе:Слияние ядер с сокращением вычислений с помощью общей памятиУпоминается, что доступ к памяти в Softmax Kernel оптимизирован до 2*D, но это все еще не оптимизировано до предела. После оптимизации, описанной в этой статье, использование пропускной способности памяти ядром OneFlow Softmax CUDA может приблизиться к теоретической пропускной способности в большинстве сценариев.
Сравнение OneFlow и cuDNN
Мы сравнили пропускную способность доступа к памяти между ядром Softmax CUDA после глубокой оптимизации OneFlow и ядром Softmax в cuDNN.Результаты тестирования следующие:
Пропускная способность ядра Softmax:
Пропускная способность ядра Softmax Grad:
Тестовая среда — графический процессор GeForce RTX™ 3090, тип данных — половина, Softmax’s Shape = (49152, num_cols), где 49152 = 32 * 12 * 128 — передняя трехмерность тензора внимания в сети BERT-base, мы исправили В первых трех измерениях последнее измерение динамически изменяется, и тестируется эффективная пропускная способность памяти ядра Softmax прямого и обратного ядра разных размеров от 32 до 32768. Из приведенных выше двух рисунков видно, что OneFlow в большинстве случаев может приблизиться к теоретической пропускной способности (около 800 ГБ/с, что эквивалентно пропускной способности доступа к памяти cudaMemcpy. Официально заявленная теоретическая пропускная способность составляет 936 ГБ/с). И во всех случаях эффективная пропускная способность выборки памяти ядра CUDA OneFlow лучше, чем у реализации cuDNN.
Методы OneFlow для глубокой оптимизации ядра Softmax CUDA
Входная форма функции Softmax: (num_rows, num_cols), изменение num_cols повлияет на эффективную пропускную способность, потому что нет никогоУниверсальныйМетод оптимизации может обеспечить оптимальную передачу во всех случаях num_cols. Поэтому для оптимизации SoftmaxKernel в OneFlow используется кусочная функция: для разных диапазонов num_cols выбираются разные реализации для достижения высокой эффективной пропускной способности во всех случаях. См. Оптимизация ядра softmax cuda.
В OneFlow есть три реализации, и softmax оптимизируется поэтапно:
(1) аWarpОбработка расчета одной строки, подходит для случая num_cols
32 потока, выполняющихся параллельно на оборудовании, называются варпом, и 32 потока одного и того же варпа выполняют одну и ту же инструкцию. warp — это базовая единица выполнения планирования графического процессора.
(2) Один блок обрабатывает вычисление одной строки и использует общую память для сохранения промежуточных данных результата.Это подходит для случая, когда требуемые ресурсы общей памяти соответствуют условиям запуска запуска ядра.В этой тестовой среде это 1024
(3) Блок обрабатывает расчет одной строки, не использует разделяемую память и многократно считывает ввод x, что подходит для случаев, когда (1) и (2) не поддерживаются.
Три реализации описаны ниже как пример прямых вычислений:
Реализация 1: Каждый Warp обрабатывает одну или две строки элементов.
\
Каждый Warp обрабатывает одну или две строки элементов, и операция Reduce каждой строки должна выполнять операцию Reduce внутри Warp. Мы реализуем WarpAllReduce для выполнения операций Global Max и Global Sum между потоками в Warp. WarpAllReduce использует уровень Warp. примитив__shfl_xor_sync Реализован, код такой.
template<template<typename> typename ReductionOp, typename T>
__inline__ __device__ T WarpAllReduce(T val) {
for (int mask = kWarpSize / 2; mask > 0; mask /= 2) {
val = ReductionOp<T>()(val, __shfl_xor_sync(0xffffffff, val, mask));
}
return val;
}
Реализация SoftmaxWarpImpl имеет следующие параметры шаблона:
LOAD и STORE представляют вход и выход соответственно, используйтеload.template load<pack_size>(ptr, row_id, col_id);
иstore.template store<pack_size>(ptr, row_id, col_id);
Прочитайте и напишите. Использование LOAD и STORE имеет два преимущества: 1. Вы можете заботиться только о типе вычисления ComputeType в ядре CUDA, но не о конкретном типе данных T. 2. Нужно всего лишь добавить несколько строк кода, чтобы быстро поддерживать Softmax и другие Kernel Fuse, снизить требования к пропускной способности и повысить общую производительность. Обычный SoftmaxKernel использует DirectLoad и DirectStore напрямую, FusedSoftmaxKernel, например
FusedScaleSoftmaxDropoutKernel необходимо только определить структуру ScaleLoad и структуру DropoutStore для предварительной обработки Scale для входных данных x и постобработки Dropout для выходных данных y.
ComputeType представляет тип расчета. pack_size представляет количество элементов пакета векторизованной операции доступа к памяти Мы упаковываем несколько элементов для чтения и записи, чтобы улучшить использование полосы пропускания. cols_per_thread представляет количество элементов, обработанных каждым потоком.
thread_group_width представляет собой ширину группы потоков, которая обрабатывает элемент. Когда cols > pack_size * warp_size, thread_group_width равен warp_size, что равно 32. Когда cols
rows_per_access представляет количество строк, обрабатываемых каждой группой thread_group за раз.Когда cols мал, а thread_group_width не равен warp_size 32, если строки могут делиться на 2, мы позволяем каждому потоку обрабатывать 2 строки, чтобы увеличить параллелизм инструкций и повысить производительность.
padding показывает, выполняется ли заполнение в данный момент.Если cols не является целым числом, кратным warp_size, мы дополним его до ближайшего целого числа, кратного.
алгоритм представляет используемый алгоритм, а параметры — Algorithm::kSoftmax или Algorithm::kLogSoftmax.
Логика основного цикла, выполняемая ядром CUDA, выглядит следующим образом: во-первых, в соответствии с информацией num_cols вычисляется cols_per_thread, который будет обрабатываться каждым потоком, и каждому потоку выделяетсяrows_per_access * cols_per_thread
Размер регистра, ввод x считывается в регистр, а последующие вычисления считываются из регистра.
Теоретически обработка строки элементов в единицах Warp является самой быстрой, но так как входные данные x необходимо кэшировать с помощью регистров, а ресурсы регистров ограничены, то в случае num_cols>1024 следует использовать (2) разделяемый метод памяти работает достаточно быстро, поэтому используйте реализацию Warp только при num_cols
template<typename LOAD, typename STORE, typename ComputeType, int pack_size, int cols_per_thread,
int thread_group_width, int rows_per_access, bool padding, Algorithm algorithm>
__global__ void SoftmaxWarpImpl(LOAD load, STORE store, const int64_t rows, const int64_t cols) {
static_assert(cols_per_thread % pack_size == 0, "");
static_assert(thread_group_width <= kWarpSize, "");
static_assert(kWarpSize % thread_group_width == 0, "");
constexpr int num_packs = cols_per_thread / pack_size;
assert(cols <= cols_per_thread * thread_group_width);
ComputeType buf[rows_per_access][cols_per_thread];
const int global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y;
const int num_global_thread_group = gridDim.x * blockDim.y;
const int lane_id = threadIdx.x;
for (int64_t row = global_thread_group_id * rows_per_access; row < rows;
row += num_global_thread_group * rows_per_access) {
ComputeType thread_max[rows_per_access];
#pragma unroll
for (int row_id = 0; row_id < rows_per_access; ++row_id) {
thread_max[row_id] = -Inf<ComputeType>();
ComputeType* row_buf = buf[row_id];
#pragma unroll
for (int pack_id = 0; pack_id < num_packs; ++pack_id) {
const int col = (pack_id * thread_group_width + lane_id) * pack_size;
if (!padding || col < cols) {
load.template load<pack_size>(row_buf + pack_id * pack_size, row + row_id, col);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
thread_max[row_id] = max(thread_max[row_id], row_buf[pack_id * pack_size + i]);
}
} else {
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
row_buf[pack_id * pack_size + i] = -Inf<ComputeType>();
}
}
}
}
ComputeType warp_max[rows_per_access];
#pragma unroll
for (int row_id = 0; row_id < rows_per_access; ++row_id) {
warp_max[row_id] = WarpAllReduce<MaxOp, ComputeType, thread_group_width>(thread_max[row_id]);
}
ComputeType thread_sum[rows_per_access];
#pragma unroll
for (int row_id = 0; row_id < rows_per_access; ++row_id) {
thread_sum[row_id] = 0;
ComputeType* row_buf = buf[row_id];
#pragma unroll
for (int i = 0; i < cols_per_thread; ++i) {
if (algorithm == Algorithm::kSoftmax) {
row_buf[i] = Exp(row_buf[i] - warp_max[row_id]);
thread_sum[row_id] += row_buf[i];
} else if (algorithm == Algorithm::kLogSoftmax) {
row_buf[i] -= warp_max[row_id];
thread_sum[row_id] += Exp(row_buf[i]);
} else {
__trap();
}
}
}
ComputeType warp_sum[rows_per_access];
#pragma unroll
for (int row_id = 0; row_id < rows_per_access; ++row_id) {
warp_sum[row_id] = WarpAllReduce<SumOp, ComputeType, thread_group_width>(thread_sum[row_id]);
}
#pragma unroll
for (int row_id = 0; row_id < rows_per_access; ++row_id) {
ComputeType* row_buf = buf[row_id];
#pragma unroll
for (int i = 0; i < cols_per_thread; ++i) {
if (algorithm == Algorithm::kSoftmax) {
row_buf[i] = Div(row_buf[i], warp_sum[row_id]);
} else if (algorithm == Algorithm::kLogSoftmax) {
row_buf[i] -= Log(warp_sum[row_id]);
} else {
__trap();
}
}
#pragma unroll
for (int i = 0; i < num_packs; ++i) {
const int col = (i * thread_group_width + lane_id) * pack_size;
if (!padding || col < cols) {
store.template store<pack_size>(row_buf + i * pack_size, row + row_id, col);
}
}
}
}
}
Реализация 2: Блок обрабатывает ряд элементов
Блок обрабатывает ряд элементов, Row Reduce должен выполнять операцию Reduce внутри блока, выполнять синхронизацию потоков внутри блока и использовать BlockAllReduce для завершения операций Global Max и Global Sum между потоками в Warp. BlockAllReduce реализуется с помощью метода Cub BlockReduce, код выглядит следующим образом:
template<template<typename> typename ReductionOp, typename T, int block_size>
__inline__ __device__ T BlockAllReduce(T val) {
typedef cub::BlockReduce<T, block_size> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
__shared__ T result_broadcast;
T result = BlockReduce(temp_storage).Reduce(val, ReductionOp<T>());
if (threadIdx.x == 0) { result_broadcast = result; }
__syncthreads();
return result_broadcast;
}
Различные параметры шаблона ядра CUDA подробно описаны в (1).Логика выполнения основного цикла следующая.Согласно num_cols, требуемый размер общей памяти рассчитывается как параметр запуска ядра, а ввод сохраняется с Общая память и последующие вычисления считываются непосредственно из общей памяти.
Так как ресурсы Shared Memory в SM тоже ограничены, то когда num_cols превысит определенный диапазон и приложение для Shared Memory превысит максимальный лимит при старте ядра, оно не запустится Поэтому только при вызове
cudaOccupancyMaxActiveBlocksPerMultiprocessor использует схему общей памяти, когда возвращаемое значение больше 0.
Кроме того, следует отметить, что, поскольку потоки в блоке должны быть синхронизированы, когда блок, который планируется и выполняется в SM, достигает точки синхронизации, исполняемый Warp в SM постепенно уменьшается. в это время уменьшится до 0, что приведет к простою вычислительных ресурсов и потерям.Если в то же время выполняются другие блоки, есть еще другие блоки, которые могут выполняться, когда блок достигает точки синхронизации. Когда block_size меньше, SM может планировать больше блоков одновременно, поэтому в этом случае, чем меньше block_size, тем лучше. Однако, когда block_size увеличивается, а количество блоков, которые SM может планировать одновременно, остается неизменным, block_size должен быть как можно больше, и чем больше блок, тем выше степень параллелизма. Поэтому при выборе block_size в коде вычисляются разные block_size
cudaOccupancyMaxActiveBlocksPerMultiprocessor, если результат тот же, используйте больший размер блока.
template<typename LOAD, typename STORE, typename ComputeType, int pack_size, int block_size,
Algorithm algorithm>
__global__ void SoftmaxBlockSMemImpl(LOAD load, STORE store, const int64_t rows,
const int64_t cols) {
extern __shared__ __align__(sizeof(double)) unsigned char shared_buf[];
auto* buf = reinterpret_cast<ComputeType*>(shared_buf);
const int tid = threadIdx.x;
assert(cols % pack_size == 0);
const int num_packs = cols / pack_size;
for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
ComputeType thread_max = -Inf<ComputeType>();
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
load.template load<pack_size>(pack, row, pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
buf[i * num_packs + pack_id] = pack[i];
thread_max = max(thread_max, pack[i]);
}
}
const ComputeType row_max = BlockAllReduce<MaxOp, ComputeType, block_size>(thread_max);
ComputeType thread_sum = 0;
for (int col = tid; col < cols; col += block_size) {
if (algorithm == Algorithm::kSoftmax) {
const ComputeType exp_x = Exp(buf[col] - row_max);
buf[col] = exp_x;
thread_sum += exp_x;
} else {
const ComputeType x = buf[col] - row_max;
buf[col] = x;
thread_sum += Exp(x);
}
}
const ComputeType row_sum = BlockAllReduce<SumOp, ComputeType, block_size>(thread_sum);
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
if (algorithm == Algorithm::kSoftmax) {
pack[i] = Div(buf[i * num_packs + pack_id], row_sum);
} else if (algorithm == Algorithm::kLogSoftmax) {
pack[i] = buf[i * num_packs + pack_id] - Log(row_sum);
} else {
__trap();
}
thread_max = max(thread_max, pack[i]);
}
store.template store<pack_size>(pack, row, pack_id * pack_size);
}
}
}
Реализация 3: блок обрабатывает элементы строки, не использует разделяемую память и многократно считывает вход x
Как и в реализации 2, это по-прежнему блок для обработки строки элементов. Разница в том, что ввод x больше не кэшируется в общей памяти, а ввод x перечитывается каждый раз при его вычислении. Эта реализация не имеет ограничений на максимальное число столбцов и может поддерживать любой размер.
Кроме того, следует отметить, что в этой реализации block_size должен быть установлен как можно большим, чем больше block_size, тем меньше количество блоков, которые могут выполняться параллельно в SM, тем меньше потребность в кэше, и чем больше возможностей Нажмите на кеш, прочитайте x несколько раз, не будет обращаться к глобальной памяти несколько раз, поэтому в реальном тесте, когда можно использовать кеш, эффективная пропускная способность не будет уменьшена в несколько раз при трехкратном чтении x.
template<typename LOAD, typename STORE, typename ComputeType, int pack_size, int block_size,
Algorithm algorithm>
__global__ void SoftmaxBlockUncachedImpl(LOAD load, STORE store, const int64_t rows,
const int64_t cols) {
const int tid = threadIdx.x;
assert(cols % pack_size == 0);
const int num_packs = cols / pack_size;
for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
ComputeType thread_max = -Inf<ComputeType>();
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
load.template load<pack_size>(pack, row, pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { thread_max = max(thread_max, pack[i]); }
}
const ComputeType row_max = BlockAllReduce<MaxOp, ComputeType, block_size>(thread_max);
ComputeType thread_sum = 0;
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
load.template load<pack_size>(pack, row, pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { thread_sum += Exp(pack[i] - row_max); }
}
const ComputeType row_sum = BlockAllReduce<SumOp, ComputeType, block_size>(thread_sum);
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
load.template load<pack_size>(pack, row, pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
if (algorithm == Algorithm::kSoftmax) {
pack[i] = Div(Exp(pack[i] - row_max), row_sum);
} else if (algorithm == Algorithm::kLogSoftmax) {
pack[i] = (pack[i] - row_max) - Log(row_sum);
} else {
__trap();
}
}
store.template store<pack_size>(pack, row, pack_id * pack_size);
}
}
}
Общие советы по оптимизации
В дополнение к методам оптимизации сегментации для softmax, представленным выше, OneFlow также использует некоторые общедоступные методы оптимизации в реализации softmax, которые используются не только в softmax, но и в других реализациях ядра. Вот два:
1. Упакуйте тип Half в Half2 для доступа, улучшите пропускную способность инструкций без изменения задержки, аналогично оптимизации шаблона CUDA для поэлементных ядер.
2. Конфликты банков в общей памяти
CUDA делит Shared Memory на 32 банка в соответствии с размером 4 байта или 8 байт.Для архитектуры Volta это 4 байта.Здесь для примера используются 4 байта.Как показано на рисунке ниже, адреса 0 -128 байт соответственно находятся в банке 0. -31, 128-256 также находятся в банке 0-31 соответственно.
Примечание: это изображение и следующие изображения банковских конфликтов взяты из архитектуры VOLTA и оптимизации производительности.
Конфликты банков возникают, когда разные потоки внутри варпа обращаются к разным адресам одного и того же банка. Когда возникают конфликты банков, потоки могут обращаться только последовательно, увеличивая задержку.На следующем рисунке показана задержка, когда n потоков в Warp одновременно обращаются к разным адресам в одном и том же банке.
Примечание: Рисунок из анализа архитектуры графического процессора NVIDIA Volta с помощью микробенчмаркинга.
Несколько ситуаций банковских конфликтов:
Если каждый поток в Warp читает 4 байта и обращается к ним последовательно, конфликтов банков не будет.Если каждый поток в Warp читает 8 байтов, адрес, к которому обращается поток 0 в Warp, находится между 0-м и 1-м банками адрес, к которому обращается поток 1, находится во 2-м и 3-м банках и т. д., адрес, к которому обращается поток 16, находится в 0-м и 1-м банках, что отличается от адреса, когда поток 0 обращается к тому же банку. генерируется в это время.
В предыдущей реализации (2) в процессе выделения разделяемой памяти, если используется следующий метод, когда размер пакета = 2, каждый поток записывает два последовательных 4-байтовых адреса, и будут генерироваться конфликты банков.
#pragma unroll
for (int j = 0; j < pack_size; ++j) {
buf[pack_id * pack_size * j] = pack[j];
thread_max = max(thread_max, pack[j]);
}
Следовательно, в реализации (2) для общей памяти принята новая структура памяти, которая позволяет избежать доступа одного и того же варпа к разным адресам одного и того же банка и позволяет избежать конфликтов банков.
#pragma unroll
for (int j = 0; j < pack_size; ++j) {
buf[num_packs * j + pack_id] = pack[j];
thread_max = max(thread_max, pack[j]);
}
Использованная литература:
Using CUDA Warp-Level Primitives | NVIDIA Developer Blog
CUDA Pro Tip: Increase Performance with Vectorized Memory Access
Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking
VOLTA Architecture and performance optimization