Авторы Го Ран, Яо Чи, Чжэн Цзекан, Лю Цзюньчэн
Ранее OneFlow выпустилСовместное использование оптимизации производительности OneFlow: как внедрить эффективное ядро Softmax CUDA?", который вводит Softmax после глубокой оптимизации OneFlow, особенно для типа половин, который не рассматривается многими фреймворками, так что производительность значительно превышает реализацию cuDNN.
Сегодня я поделюсь практикой оптимизации производительности другого важного оператора, LayerNorm.
Кроме того, OneFlow также поставляется с OneFlow Softmax, который можно использовать независимо (подробности см. в описании в конце статьи). Вы можете попробовать его и внести предложения.
Результаты тестирования после оптимизации производительности OneFlow
Производительность оптимизированного LayerNorm OneFlow сравнивается с NVIDIA Apex и PyTorch соответственно.Результаты тестирования показывают, что OneFlow LayerNorm имеет очевидные преимущества в производительности.
Сравнение с NVIDIA Apex
В NVIDIA Apex реализовано эффективное объединенное ядро LayerNorm для расширения оператора PyTorch. Мы сравнили ядро LayerNorm, оптимизированное для OneFlow, с NVIDIA Apex. Результаты тестирования следующие:
Горизонтальная ось — это размер num_cols, а вертикальная ось — это время, необходимое для выполнения ядра (чем меньше, тем лучше):
Мы конвертируем время в пропускную способность доступа к памяти, и результаты следующие, вертикальная ось — это эффективная пропускная способность, достигнутая ядром (чем выше, тем лучше):
Среда тестирования — GPU NVIDIA A100-PCIE-40GB, тип данных — половина, Shape = (49152, num_cols), мы динамически меняем последнее измерение и тестируем LayerNorm Kernels разных размеров от 32 до 32768, вы можете видеть, что во всех В этом случае время выполнения ядра OneFlow и эффективная пропускная способность доступа к памяти лучше, чем у реализации Apex.
Сравнение с PyTorch
LayerNorm PyTorch на данный момент не поддерживает половинный тип, поэтому мы сделали набор сравнений с типом float.Следует отметить, что LayerNorm в PyTorch разделен на два ядра CUDA (RowwiseMomentsCUDAKernel и LayerNormForwardCUDAKernel), поэтому кажется, что производительность относительно беден.
Горизонтальная ось — это размер num_cols, а вертикальная ось — это время, необходимое для выполнения ядра (чем меньше, тем лучше):
Видно, что в каждой группе сравнительных экспериментов производительность OneFlow также лучшая.
Оптимизация производительности LayerNorm
LayerNorm — одна из часто используемых операций в языковых моделях. Эффективность реализации ядра CUDA повлияет на конечную скорость обучения многих сетей. Этот метод оптимизации Softmax также применим к LayerNorm. Данные LayerNorm также могут быть выражены как ( num_rows, num_cols), в процессе расчета выполните операцию сокращения над элементами каждой строки, чтобы получить среднюю дисперсию. Поэтому мы используем тот же метод оптимизации, что и Softmax, для оптимизации операции LayerNorm.В этой статье в качестве примера используется прямой расчет LayerNorm.
Метод расчета LayerNorm
Взяв в качестве примера PyTorch, интерфейс LayerNorm выглядит следующим образом:
torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None)
где входная форма: [*, normalized_shape[0], normalized_shape[1], …,normalized_shape[−1]]
Первый параметр normalized_shape может быть только последними размерами входного x_shape, например, x_shape равен (N, C, H, W), normalized_shape может быть (W), (H, W), (C, H, W) или (N, C, H, W). Введите x, чтобы найти среднее значение и дисперсию по измерениям normalized_shape.
Третий параметр elementwise_affine указывает, следует ли преобразовывать результат нормализации, то есть умножать результат нормализации на гамму и добавлять бета. Если elementwise_affine=True, есть еще два параметра модели gamma и beta, а форма — normalized_shape.
Например, если вход x имеет форму (N, C, H, W), а normalized_shape — (H, W), можно понять, что вход x равен (NC, HВт), в ННа линиях C каждая строка имеет HW элементов, возьмите среднее значение и дисперсию по элементам каждой строки и получите NСреднее значение C и inv_variance, а затем рассчитайте y в соответствии с приведенной ниже формулой расчета LayerNorm для входных данных. Если elementwise_affine=True , то HW гаммы и бета, преобразующие элементы H*W в каждой строке.
Методы поиска дисперсии в LayerNorm
Общие методы расчета дисперсии включают двухпроходный метод, наивный метод и алгоритм Велфорда.В этой статье приведены некоторые ключевые формулы и выводы.Подробное введение и вывод см. в: Wiki: Алгоритмы вычисления дисперсии и GiantPandaCV: Реализация LN с Алгоритм Велфорда Обновление дисперсии
- двухпроходный метод
Используемая формула:
Двухпроходный означает, что этот метод должен пройти данные дважды, первый проход накапливает x, чтобы получить среднее значение, а второй проход использует приведенную выше формулу для вычисления дисперсии. Этот метод по-прежнему является численно стабильным, когда n относительно мало.
- наивный метод
Используемая формула:
Этот метод является методом с одним проходом.При расчете дисперсии вам нужно только пройти данные и накопить квадрат x и накопить x, и, наконец, вычислить дисперсию в соответствии с приведенной выше формулой. Этот метод требует прохождения данных только один раз.По сравнению с двухпроходным алгоритмом, легче достичь хорошей производительности.Однако, как описано в справочной ссылке Wiki выше, поскольку SumSquare и (Sum×Sum)/n могут быть очень близко, это может привести к расчету. В результате потери точности велики, поэтому этот метод не рекомендуется использовать на практике.
- Алгоритм Велфорда
Используемая формула:
Алгоритм Велфорда также является однопроходным методом с хорошей числовой стабильностью, поэтому многие фреймворки теперь используют этот метод. Метод Велфорда также используется в коде этой статьи.
Методы OneFlow для глубокой оптимизации ядра CUDA LayerNorm
Как и Softmax, LayerNorm также использует кусочно-функциональную оптимизацию с различными реализациями для разных диапазонов num_cols для достижения высокой эффективной пропускной способности во всех случаях.
В каждой реализации используется общая оптимизация: векторизованный доступ к памяти. В блоге NVIDIA, посвященном оптимизации производительности, повышение производительности с помощью векторизованного доступа к памяти упоминается, что векторизованные операции с памятью могут повысить производительность ядра CUDA. количество инструкций, уменьшить задержку и улучшить использование полосы пропускания.
Теоретически ввод x необходимо считывать дважды во время вычисления LayerNorm, первый раз для вычисления среднего значения и дисперсии. Второй раз используется для процесса расчета после получения среднего значения и дисперсии. Операция доступа к глобальной памяти является дорогостоящей.Если ввод x можно сохранить первым, а не считывать повторно, производительность может быть улучшена. Регистры или общая память могут использоваться для хранения входных данных x в графическом процессоре, но как ресурсы регистров, так и ресурсы общей памяти ограничены. Если num_cols слишком велико, предел использования ресурсов будет превышен. Поэтому мы используем разные реализации для разных num_cols Соответственно вводятся:
1. Случай num_cols
Для num_cols
32 потока, выполняемые параллельно на оборудовании, называются Warp, и 32 потока одного и того же Warp выполняют одну и ту же инструкцию.Warp — это базовая единица планирования и выполнения GPU. Соответствующая взаимосвязь между блоками потоков и элементами показана на рисунке выше.Потоки каждого Warp обрабатывают строку элементов, каждый блок имеет Warp block_size/warp_size, и каждый блок обрабатывает элементы строки block_size/warp_size.
Конкретный поток обработки, как показано на рисунке ниже, каждая строка имеет элементы num_cols, и каждый warp обрабатывает одну строку, поэтому каждый поток должен обрабатывать элементы num_cols/warp_size, и каждый поток считывает элементы, необходимые для обработки, и сохраняет их. в регистре и после вычисления среднего значения и дисперсии с помощью алгоритма Велфорда все потоки в Warp выполняют WelfordWarpAllReduce один раз, так что каждый поток получает правильное среднее значение и дисперсию для участия в последующем вычислении.
WelfordWarpAllReduce завершается операциями WelfordWarpReduce и Broadcast, WelfordWarpReduce реализуется с помощью примитива синхронизации уровня Warp __shfl_down_sync, а операция Broadcast реализуется с помощью __shfl_sync, код выглядит следующим образом:
template<typename T, int thread_group_width = kWarpSize>
__inline__ __device__ void WelfordWarpReduce(T thread_mean, T thread_m2, T thread_count, T* mean,
T* m2, T* count) {
*mean = thread_mean;
*m2 = thread_m2;
*count = thread_count;
for (int mask = thread_group_width / 2; mask > 0; mask /= 2) {
T b_mean = __shfl_down_sync(0xffffffff, *mean, mask);
T b_m2 = __shfl_down_sync(0xffffffff, *m2, mask);
T b_count = __shfl_down_sync(0xffffffff, *count, mask);
WelfordCombine(b_mean, b_m2, b_count, mean, m2, count);
}
}
template<typename T, int thread_group_width = kWarpSize>
__inline__ __device__ void WelfordWarpAllReduce(T thread_mean, T thread_m2, T thread_count, T* mean,
T* m2, T* count) {
WelfordWarpReduce<T, thread_group_width>(thread_mean, thread_m2, thread_count, mean, m2, count);
*mean = __shfl_sync(0xffffffff, *mean, 0, thread_group_width);
*m2 = __shfl_sync(0xffffffff, *m2, 0, thread_group_width);
*count = __shfl_sync(0xffffffff, *count, 0, thread_group_width);
}
Здесь есть параметр шаблона thread_group_width, когда num_cols > pack_size * WarpSize, thread_group_width равен WarpSize. Когда num_cols слишком мал, то есть num_cols Кроме того, при чтении и записи ввода и вывода мы используем оптимизацию векторизованного доступа к памяти.При выполнении условий элементы pack_size упаковываются в более крупные типы данных и считываются. каждый поток Чтение элементов с большими типами данных позволяет лучше использовать пропускную способность памяти. Упакуйте элементы pack_size в более крупный тип данных и считывайте их, но x также участвует в расчетах. Поэтому мы определяем структуру объединения типа Pack, хранилище используется для чтения и записи из глобальной памяти и использует elem[i] для принятия каждого элемента для участия в расчете.Тип Pack определяется следующим образом: Код ядра LayerNormWarpImpl выглядит следующим образом: Значения параметров шаблона реализации LayerNormWarpImpl следующие: Для num_cols > 1024 обрабатывать одну строку в блоке и использовать общую память для хранения входных данных. Для num_cols > 1024 каждый блок обрабатывает строку элементов, сохраняя ввод x в общей памяти. Конкретный поток обработки, как показано на рисунке ниже, каждая строка имеет элементы num_cols, и каждый блок обрабатывает одну строку, поэтому каждый поток должен обрабатывать элементы num_cols/block_size, и каждый поток считывает элементы, необходимые для обработки, и сохраняет их. в общей памяти После вычисления среднего значения и дисперсии с помощью алгоритма Велфорда все потоки в блоке выполняют WelfordBlockAllReduce один раз, так что каждый поток получает правильное среднее значение и дисперсию для участия в последующих вычислениях. WelfordBlockAllReduce завершается с помощью операции WelfordWarpReduce.Специфическая логика заключается в том, что в Блоке не более 32 варпов.WelfordWarpReduce выполняется один раз для всех варпов.После выполнения первый поток в каждом варпе, то есть поток с lane_id =0 Результат текущего WelfordWarpReduce получается из вышеописанного, а затем результат первого потока каждого Warp копируется в буфер Shared Memory, а затем 32 потока первого Warp используются для однократного выполнения WelfordWarpReduce. время, lane_id в первом Warp равно Поток с =0 получает результат сокращения всеми потоками в блоке. С помощью Shared Memory результат транслируется во все потоки в блоке, что завершает операцию WelfordBlockAllReduce. Стоит отметить, что ресурсы разделяемой памяти на графическом процессоре также ограничены, когда num_cols превышает определенный диапазон, разделяемая память, которую необходимо занять, может превысить максимальный предел, и ядро не может быть запущено. Поэтому мы используем функцию cudaOccupancyMaxActiveBlocksPerMultiprocessor, чтобы определить, можно ли успешно запустить ядро при текущих условиях аппаратных ресурсов, и используем эту схему только тогда, когда возвращаемое значение больше 0. Кроме того, поскольку потоки в блоке должны быть синхронизированы, когда блок, который планируется и выполняется в SM, достигает точки синхронизации, исполняемый Warp в SM постепенно уменьшается, в это время он постепенно уменьшается до 0, что приведет к простою вычислительных ресурсов и потерям.Если в то же время выполняются другие блоки, есть еще другие блоки, которые могут быть выполнены, когда блок достигает точки синхронизации. Когда block_size меньше, SM может планировать больше блоков одновременно, поэтому в этом случае, чем меньше block_size, тем лучше. Однако, когда block_size увеличивается, а количество блоков, которые SM может планировать одновременно, остается неизменным, block_size должен быть как можно больше, и чем больше блок, тем выше степень параллелизма. Поэтому при выборе block_size в коде вычисляется cudaOccupancyMaxActiveBlocksPerMultiprocessor для разных block_size, если результат одинаковый, используется больший block_size. Код ядра LayerNormBlockSMemImpl выглядит следующим образом: Когда num_cols велико и метод использования общей памяти не может успешно запустить ядро в текущих условиях аппаратных ресурсов, используется эта реализация: блок обрабатывает элементы строки, не использует общую память и многократно считывает вход x. Этот метод согласуется с соответствующей связью между потоками и элементами во втором случае выше.Единственное отличие состоит в том, что второй метод сохраняет входные данные x в общей памяти, этот метод не хранит x и должен быть переписан из глобальной памяти. Память в каждом вычислении Чтение x вх. Хотя для этого метода требуется еще одна копия x, при фактическом выполнении часть входных данных может быть закеширована в Cache без реального увеличения времени. Стоит отметить, что в этой реализации, чем больше block_size, тем меньше блоков может выполняться параллельно в SM одновременно, тем меньше потребность в Cache и тем больше шанс попасть в Cache, поэтому мы используем больше размер_блока. Код LayerNormBlockUncachedImpl выглядит следующим образом: После неоднократных итераций интерфейс и реализация OneFlow Softmax созрели и стали стабильными, поэтому команда OneFlow отделила его и предоставила как независимый интерфейс, а оптимизированный код был помещен вGitHub.com/oneflow-Inc…, который можно скомпилировать независимо от кода OneFlow. После включения этого заголовочного файла в свой проект вы можете использовать его напрямую. Например, ядро графического процессора Softmax можно реализовать с помощью следующих строк кода. Также очень просто реализовать ядро LogSoftmax: просто замените DispatchSoftmax в приведенном выше коде на DispatchLogSoftmax. Основные преимущества OneFlow Softmax по сравнению с Softmax, предлагаемыми в других местах:template<typename T, int N>
union Pack {
PackType<T, N> storage;
T elem[N];
};
template<typename LOAD, typename STORE, typename ComputeType, int pack_size, int cols_per_thread,
int thread_group_width, int rows_per_access, bool padding>
__global__ void LayerNormWarpImpl(LOAD load, STORE store, const int64_t rows, const int64_t cols,
const double epsilon, ComputeType* mean,
ComputeType* inv_variance) {
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 int64_t global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y;
const int64_t num_global_thread_group = gridDim.x * blockDim.y;
const int64_t 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_mean[rows_per_access];
ComputeType thread_m2[rows_per_access];
ComputeType thread_count[rows_per_access];
#pragma unroll
for (int row_id = 0; row_id < rows_per_access; ++row_id) {
thread_mean[row_id] = 0;
thread_m2[row_id] = 0;
thread_count[row_id] = 0;
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;
const int pack_offset = pack_id * pack_size;
if (!padding || col < cols) {
load.template load<pack_size>(row_buf + pack_offset, row + row_id, col);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
WelfordCombine(row_buf[pack_offset + i], thread_mean + row_id, thread_m2 + row_id,
thread_count + row_id);
}
} else {
#pragma unroll
for (int i = 0; i < pack_size; ++i) { row_buf[pack_offset + i] = 0; }
}
}
}
ComputeType warp_mean[rows_per_access];
ComputeType warp_m2[rows_per_access];
ComputeType warp_count[rows_per_access];
#pragma unroll
for (int row_id = 0; row_id < rows_per_access; ++row_id) {
int global_row_id = row + row_id;
ComputeType* row_buf = buf[row_id];
WelfordWarpAllReduce<ComputeType, thread_group_width>(
thread_mean[row_id], thread_m2[row_id], thread_count[row_id], warp_mean + row_id,
warp_m2 + row_id, warp_count + row_id);
ComputeType row_mean = warp_mean[row_id];
ComputeType row_variance =
max(Div(warp_m2[row_id], warp_count[row_id]), static_cast<ComputeType>(0.0));
ComputeType row_inv_var = Rsqrt(row_variance + static_cast<ComputeType>(epsilon));
if (lane_id == 0) {
mean[global_row_id] = row_mean;
inv_variance[global_row_id] = row_inv_var;
}
#pragma unroll
for (int i = 0; i < cols_per_thread; ++i) {
row_buf[i] = (row_buf[i] - row_mean) * row_inv_var;
}
#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, global_row_id, col);
}
}
}
}
}
2. Случай num_cols > 1024
template<typename LOAD, typename STORE, typename ComputeType, int pack_size, int block_size>
__global__ void LayerNormBlockSMemImpl(LOAD load, STORE store, const int64_t rows,
const int64_t cols, const double epsilon, ComputeType* mean,
ComputeType* inv_variance) {
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_mean = 0;
ComputeType thread_m2 = 0;
ComputeType thread_count = 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) {
buf[i * num_packs + pack_id] = pack[i];
WelfordCombine(pack[i], &thread_mean, &thread_m2, &thread_count);
}
}
ComputeType row_mean = 0;
ComputeType row_m2 = 0;
ComputeType row_count = 0;
WelfordBlockAllReduce<ComputeType>(thread_mean, thread_m2, thread_count, &row_mean, &row_m2,
&row_count);
ComputeType row_variance = max(Div(row_m2, row_count), static_cast<ComputeType>(0.0));
ComputeType row_inv_var = Rsqrt(row_variance + static_cast<ComputeType>(epsilon));
if (threadIdx.x == 0) {
mean[row] = row_mean;
inv_variance[row] = row_inv_var;
}
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) {
pack[i] = (buf[i * num_packs + pack_id] - row_mean) * row_inv_var;
}
store.template store<pack_size>(pack, row, pack_id * pack_size);
}
}
}
3. Когда num_cols велико, общая память не используется.
template<typename LOAD, typename STORE, typename ComputeType, int pack_size, int block_size>
__global__ void LayerNormBlockUncachedImpl(LOAD load, STORE store, const int64_t rows,
const int64_t cols, const double epsilon,
ComputeType* mean, ComputeType* inv_variance) {
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_mean = 0;
ComputeType thread_m2 = 0;
ComputeType thread_count = 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) {
WelfordCombine(pack[i], &thread_mean, &thread_m2, &thread_count);
}
}
ComputeType row_mean = 0;
ComputeType row_m2 = 0;
ComputeType row_count = 0;
WelfordBlockAllReduce<ComputeType>(thread_mean, thread_m2, thread_count, &row_mean, &row_m2,
&row_count);
ComputeType row_variance = max(Div(row_m2, row_count), static_cast<ComputeType>(0.0));
ComputeType row_inv_var = Rsqrt(row_variance + static_cast<ComputeType>(epsilon));
if (threadIdx.x == 0) {
mean[row] = row_mean;
inv_variance[row] = row_inv_var;
}
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
const int pack_offset = pack_id * pack_size;
load.template load<pack_size>(pack, row, pack_offset);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { pack[i] = (pack[i] - row_mean) * row_inv_var; }
store.template store<pack_size>(pack, row, pack_offset);
}
}
}
Библиотеки OneFlow Softmax
oneflow::cuda::softmax::DirectLoad<half, float> load(in, cols);
oneflow::cuda::softmax::DirectStore<float, half> store(out, cols);
oneflow::cuda::softmax::DispatchSoftmax<decltype(load), decltype(store), float>(
cuda_stream, load, store, rows, cols);
Добро пожаловать, чтобы загрузить и испытать новое поколение среды глубокого обучения OneFlow с открытым исходным кодом:GitHub.com/oneflow-Inc…