От MNIST к Transformer. Часть 2. Основы работы с памятью

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

Это вторая статья из цикла От MNIST к Transformer, цель которого пошагово пройти путь от простого CUDA ядра до создания архитектуры Transformer - фундамента современных LLM моделей. Мы не будем использовать готовые высокоуровневые библиотеки. Мы будем разбирать, как все устроено под капотом, и пересобирать их ключевые механизмы своими руками на самом низком уровне. Только так можно по настоящему понять как работают LLM и что за этим стоит. В этой статье разберем основы работы с памятью и две простые операции с точки зрения математики, но не такие простые с точки зрения CUDA ядер.

Приготовьтесь, будет много кода на C++ и CUDA, работы с памятью и погружения в архитектуру GPU. И конечно же математика что за этим стоит. Поехали!

О цикл "От MNIST к Transformer"

Это будет долгий, сложный, но невероятно увлекательный путь. Мы начнем с самых азов и будем шаг за шагом усложнять задачу. Наша цель не только досконально понять как работают современные модели, но и научится программировать на GPU. Статьи для цикла будут выходить постепенно.

1. Самое начало. Разберемся с основами работы GPU. Напишем наш первый код на C++ и CUDA для простейшей операции сложения векторов, hello world в мире GPGPU и CUDA, чтобы понять, как CPU и видеокарта общаются друг с другом. Статья вот тут.

2. Основы работы с памятью (Мы здесь). Разберемся с основами работы и устройства памяти в GPU. Напишем четыре ядра для двух операций (простой подход и оптимизированный), построение гистрограммы для вектора и транспонирование матрицы. Операция транспонирования нам понадобится в будущем что бы собирать наш Transformer.

Типы памяти на GPU

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

Глобальная память или Global Memory (VRAM) - это основное пространство памяти для хранения данных, доступное всем потокам. Она аналогична оперативной памяти DRAM. Ядра, запущенные на графическом процессоре (GPU), имеют прямой доступ к глобальной памяти точно так же, как код, выполняемый на центральном процессоре (CPU), имеет доступ к системной памяти.

В коде из прошлой статьи A, B и С находятся в глобальной памяти.

__global__ void vectorAdd(const float* A, const float* B, float* C, int numElements) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < numElements) { C[i] = A[i] + B[i]; } }

Разделяемая память или shared memory - это пространство памяти, доступное всем потокам в блоке (thread block). Физически она расположена на стриминговом мультипроцессоре (SM) и разделяет эту область с L1 кэшем, сколько памяти отдавать под кэш, а сколько под shared memory можно настраивать в рантайме. Shared memory сохраняются на протяжении всего времени выполнения ядра. Shared memory можно рассматривать как управляемую временную память для использования ядром. Несмотря на небольшой размер по сравнению с глобальной памятью, благодаря ее расположению на SM, она обладает более высокой пропускной способностью (bandwidth) и меньшей задержкой (latency).

Поскольку разделяемая память доступна всем потокам в блоке, необходимо соблюдать осторожность, чтобы избежать data races между потоками в одном и том же блоке. Синхронизация между потоками в одном блоке может быть достигнута с помощью функции syncthreads(). Эта функция блокирует все потоки в одном блоке до тех пор, пока каждый из них не достигнет вызова syncthreads().

Shared memory может быть алоцирована статически:

__shared__ float sharedArray[1024];

или динамически:

extern __shared__ float sharedData[];

Регистры или Registers - расположены на мультипроцессоре (SM) и имеют локальную область видимости для потока (ядра). Они используются для локального хранения данных потока во время выполнения ядра. Количество регистров на один SM и на один блок потоков можно запросить с помощью свойств устройства GPU regsPerMultiprocessor и regsPerBlock.

NVCC позволяет указать максимальное количество регистров, используемых ядром, с помощью опции -maxrregcount. Но нужно быть аккуратным c использование этой опции с одной стороны при уменьшении количества регистров на одном SM можно запустить больше thread blocks с другой если регистров не хватит для хранения локальных данных это может привет к так называемому register spilling.

Локальная память или Local memory - это хранилище, локальное для каждого потока, как и регистры. Однако физически локальная память располагается в пространстве глобальной памяти. Название локальная память относится к её логической области видимости, а не к физическому расположению. Локальная память используется для хранения данных, локальных для потока, во время выполнения ядра. Локальная память значительно медленнее чем регистры.

К данным, которые компилятор, скорее всего, разместит в локальной памяти, относятся:

• Большие структуры или массивы, которые заняли бы слишком много регистрового пространства;

• Eсли ядро использует больше регистров, чем доступно (то есть в случае вытеснения регистров — register spilling) любые переменные будут перемещены в локальную память

Константная память или Constant Memory - имеет область видимости всей сетки (grid scope) и доступна в течение всего времени работы приложения. Она физически находится на устройстве и доступна ядру только для чтения. В связи с этим, такие переменные должны объявляться и инициализироваться на хосте (CPU) с использованием спецификатора __сonstant__ вне каких-либо функций. Общий объем доступной константной памяти можно запросить через свойство устройства totalConstMem. Константная память имеет выделенный кэш. Когда все потоки в варпе (warp) обращаются к одному и тому же адресу, данные считываются один раз и рассылаются всем потокам.

L1/L2 кэш - Кэш L2 расположен на устройстве и является общим для всех мультипроцессоров (SM). Размер кэша L2 можно узнать с помощью свойства l2CacheSize, в структуре которую возвращает функция cudaGetDeviceProperties. L1 кэш находится непосредственно внутри каждого SM, он намного быстрее, но его объем сильно ограничен.

WARP и SIMT что это такое?

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

В основе CUDA лежит архитектура SIMT (Single Instruction, Multiple Threads — одна инструкция, множество потоков) очень похоже на парадигму SIMD (Single Instruction, Multiple Data). Что же это такое? В отличие от классического подхода CPU, GPU группирует потоки в небольшие блоки (не путать с thread block) которые называются варпами или warps.

WARP — это минимальная единица планирования на GPU, состоящая из 32 потоков. Все потоки внутри одного варпа всегда выполняют одну и ту же инструкцию одновременно. Если вы пишете код для одного потока, то знайте он всегда выполняется вместе еще с 31 соседним ядром.

WARP Divergence

WARP и WARP, SIMT и SIMT, но что нам это дает? А давайте представим вот такой простой пример кода внутри ядра.

if (threadIdx.x % 2 == 0) { // Четные потоки делают это } else { // Нечетные потоки делают вот это }

Так вот, так как потоки организованы в варпы то по факту для каждого потока мы должны пройти оба условия. Что это значит, представим у нас есть два соседних потока 0 и 1, так как у нас есть расхождение в коде для одного варпа, то варп сначала пройдет этот код для первого потока (0), а второй поток (1) будет в состоянии idle и ждать пока первый поток выполнит свою работу, потом все тоже самое но наоборот для второго потока (2). Это одна из самых коварных проблем производительности в CUDA, она возникает напрямую из архитектуры SIMT.

Как это обойти? Надо убедиться что все потоки в одном варпе идут по одному пути. Например:

if (threadIdx.x / 32 >= 1) { // Первые 32 потока идут в if } else { // Слудующие 32 потока идут в else }

Правила работы с памятью

Coalesced Memory Access

Помните мы только что обсудили warp и SIMT? Так вот исходя из такой архитектуры возникает еще одна особенность, но она уже связана с работой с памятью. Доступ к глобальной памяти осуществляется через транзакции размером в 32 байта. Когда поток CUDA запрашивает данные из глобальной памяти, соответствующий варп объединяет запросы памяти от всех потоков этого варпа в минимально необходимое количество транзакций. Это количество зависит от размера данных, запрашиваемых каждым потоком, и распределения адресов памяти между потоками.

Например, если поток запрашивает 4-байта, фактическая транзакция, которую варп сделает в глобальную память, составит 32 байта. Чтобы использовать систему памяти наиболее эффективно, варп должен использовать все данные, полученные за одну транзакцию. То есть, если один поток запрашивает 4 байта, а размер транзакции - 32 байта, будет идеально, если остальные потоки в этом варпе смогут использовать оставшиеся байты из этой же 32-байтной транзакции.

В качестве простого примера: если соседние потоки в варпе запрашивают соседние значения из массива интов (int - 4 байта) в памяти, то суммарно варп запросит 32 потока x 4 байта = 128 байт данных. Эти 128 байт будут получены всего за четыре 32-байтовые транзакции. Это обеспечивает 100% эффективность использования памяти - все передаваемые данные используются в одном варпе.

Пример неправильного чтения данных: Если потоки запрашивают данные из массива интов вперемешку, например: thread_0:A[0], thread_1: A[100] и тд. GPU придется совершить 32-байтовую транзакцию для каждого такого потока. Вместо 4 транзакций на варп мы получим 32 транзакции. А размер считываемой памяти будет равно 32 байта x 32 транзакции = 1024 байта. При этом реально использоваться будет только 4 байта x 32 варпа = 128 байт. То есть всего 12.5% из перемещаемой памяти будет использоваться, остальные 87,5% данных будут гоняться туда сюда зря. Можете представить насколько это ухудшит производительность. В дальнейшем мы в этом убедимся.

Расчет гистограммы для вектора

Давайте теперь перейдем от теории к практике. Для начала давайте попробуем рассчитать гистограмму для вектора из int - овых значений. Гистограмма - преобразует массив данных в распределение частот, подсчитывая количество попаданий значений в заранее определенные диапазоны (бины).

Сначала напишем такой вот код для kernel-а который считает гистограмму:

__global__ void naiveHist( int* bins, const int numBins, const double bin_size, const int min_value, const int max_value, // __restrict__ говорит что данных указатель на область памяти // на которую не будет ссылаться какой либо другой указать // это дает возможность оптимизировать доступ к памяти const int *__restrict__ input, const int arraySize ) { // Получаем индекс ядра (потока) int tid = blockIdx.x * blockDim.x + threadIdx.x; // Каждый поток проходит по элементам входных данных со смещением в целую сетку // если массив содержит меньше элементов чем ядер то цикл выполниться всего раз for (int i = tid; i < arraySize; i += blockDim.x * gridDim.x) { // читаем данные из глобальной памяти int bin = input[i]; // вычисляем бин к которому относится значение bin = static_cast<int>((bin - min_value) / bin_size); if (bin >= 0 && bin < numBins) { // atomicAdd синхронизирует запись данных на уровне всей сетки (grid) atomicAdd(&bins[bin], 1); } } }

Что с этим кодом не так? Как думаете? Давайте посмотрим на замеры производительности по сравнению с CPU.

CPU histogram time: 22.2685 ms Naive GPU histogram time: 1071.78 ms Speedup (CPU vs naive GPU): 0.020777x

Так так так, на CPU опять считается быстрее? Может опять размер данных небольшой как было в прошлой статье? Да нет, в этот раз мы так же считаем вектор в 50млн элементов. Так в чем же проблема? А проблема в одной строчке:

atomicAdd(&bins[bin], 1);

atomicAdd - это инструмент синхронизации, но он синхронизирует запись данных на уровне всей сетки (grid). То есть представьте что каждое ядро должно ждать когда выполнят операцию ядра из всех блоков на GPU. По сути мы GPU превратили в CPU такими действиями, где у нас идет последовательно сложение. Как же это можно поправить? А давайте вспомним, что у нас есть shared memory, можно ли ее как то использовать? Давайте попробуем?

__global__ void hist( int* bins, const int numBins, const int bin_size, const int min_value, const int max_value, const int *__restrict__ input, const int arraySize ) { // создаем область памяти в shared memory extern __shared__ int shared_bins[]; // Инициализируем ее, первое ядро в блоке проходит по следующим блокам for (int i = threadIdx.x; i < numBins; i += blockDim.x) { shared_bins[i] = 0; } // Синхронизируем ядра на уровне thread block // каждый поток в блоке ждет когда другие потоки // в этом же блоке дойдут до этой строки __syncthreads(); // Каждый поток обрабатывает элементы в цикле с шагом по сетке (grid-stride loop) // накапливая результаты в локальной для блока shared memory int tid = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; for (int i = tid; i < arraySize; i += stride) { int bin = input[i]; bin = static_cast<int>((bin - min_value) / bin_size); if (bin >= 0 && bin < numBins) { // atomicAdd используется для безопасного обновления массива bins // в разделяемой памяти из нескольких потоков одновременно // эта атомарная операция влияет только на уровень блока, а не на глобальный уровень atomicAdd(&shared_bins[bin], 1); } } __syncthreads(); // Соеденяем данные из shared memory в глобальную память for (int i = threadIdx.x; i < numBins; i += blockDim.x) { atomicAdd(&bins[i], shared_bins[i]); } }

В обновленном ядре мы сначала сохраняем промежуточные результаты на уровне блока далее сохраняем результат в глобальную память. Тут мы используем atomicAdd на уровне одного блока а не всей сетки. Для запуска такого ядра мы можем указать какой размер shared memory мы хотим использовать:

int blocksPerGrid = cuda::ceil_div(input_size, threadsPerBlock); size_t sharedMemSize = numBins * sizeof(int); hist<<<blocksPerGrid, threadsPerBlock, sharedMemSize>>>( bins, numBins, bin_size, min_value, max_value, input, input_size );

Давайте посмотрим на сравнения в скорости выполнения для нового ядра:

CPU histogram time: 22.2685 ms Naive GPU histogram time: 1071.78 ms Optimized GPU hist time: 1.98758 ms Speedup (CPU vs naive GPU): 0.020777x Speedup (CPU vs optimized GPU): 11.2038x Speedup (naive GPU vs optimized GPU): 539.24x

Так, здесь мы уже видим значительный прирост производительности по сравнению с CPU и тем более с naive GPU версией.

Транспонирование матрицы

Давайте так же разберем операцию транспонирования матрицы, которая нам пригодится для построения нашего трансформера, это простая с точки зрения математики операция тоже имеет особенности реализации на GPU с учетом особенности организации памяти.

Давайте так же попробуем написать наивный подход транспонирования с помощью CUDA.

__global__ void naiveTranspose( int m, int n, float* output, const float* __restrict__ input ) { int col = blockIdx.x * blockDim.x + threadIdx.x; int row = blockIdx.y * blockDim.y + threadIdx.y; if (col < m && row < n) { output[row * m + col] = input[col * n + row]; } }

Так ну тут все достаточно просто: берем и для каждого ядра который обрабатывает конкретный элемент матрицы меняем позиции местами x -> y, y -> x. Так же хочу заметить что мы тут используем двухмерную адрессацию для blockIdx и для threadIdx, вообще в CUDA можно использовать одно-, двух- и трехмерную адресацию для удобства. Для запуска такого ядра вызываем его вот так:

dim3 blockDim(ThreadsPerBlock_x, ThreadsPerBlock_y); dim3 gridDim( cuda::ceil_div(m, ThreadsPerBlock_x), cuda::ceil_div(n, ThreadsPerBlock_y) ); naiveTranspose<<<gridDim, blockDim>>>(m, n, output, input);

Давайте опять проведем бэнчмарк и сравним с временем выполнения на CPU. Для сравнения мы возьмем матрицу размером 4096 x 4096.

CPU transpose time: 103.60 ms Naive GPU transpose time: 1053.61 ms Speedup naive GPU vs CPU: 0.10x

Так опять что то не то? Опять на CPU быстрее, давайте разбирать в чем проблема. Тут же нет ни какой atomic операции, все должно быть просто, или нет? На самом деле если мы вспомним про то что уже обсуждали в этой статье, а если быть точным про Coalesced Memory Access то возможно мы поймем что у нас есть проблемы. Давайте посмотрим на эту строчку:

output[row * m + col] = input[col * n + row];

Что мы тут видим? col использует threadIdx.x который меняется чаще чем threadIdx.y в рамках одного блока (размер блока который мы запускаем 32x32 ядра). То есть например ядро 0 имеет индекс col = 0 и row = 0, для соседнего ядра 1 в этом же варпе col = 1 и row = 0. И если мы посмотрим на операцию записи то что мы видим: ядро 0 записывает элемент с индексом 0, ядро 1 записывает элемент с индексом 1. Тут все хорошо, соседние ядра записывают в соседние ячейки в массиве. (помним про Coalesced Memory Access)
Но! Давайте посмотрим на операцию чтения тут ситуация немного другая. Ядро 0 читает элемент с индексом 0, но соседнее ядро 1 читает элемент с индексом 4096. Вот и проблема: классическая ошибка с Coalesced Memory Access. Эффективность работы с памятью в данном случае всего 12.5%

Мы обсуждаем транспонирование матриц, но в коде работаем с массивами. На самом деле тут нет проблем, мы просто представляем матрицу как массив в памяти. Сохранять матрицу можно в row-major порядке и column-major. Мы сохраняем в row-major порядке. На картинке ниже видна разница.

row-major vs column-major

Транспонирование матрицы с учетом Coalesced Memory Access

Давайте для начала я покажу ядро с учетом правильного доступа к памяти, а потом разберем его.

__global__ void transpose( int m, int n, float* output, const float* input ) { // Будем тоже использовать shared memory для сохранения промежуточных элементов. __shared__ float block_data[ThreadsPerBlock_x][ThreadsPerBlock_y]; // используются для проверки границы матрицы для чтения данных const int myCol = blockDim.x * blockIdx.x + threadIdx.x; const int myRow = blockDim.y * blockIdx.y + threadIdx.y; // позиция блока в сетке const int tileX = blockDim.x * blockIdx.x; const int tileY = blockDim.y * blockIdx.y; if( myRow < m && myCol < n ) { const int thread_y_idx = tileY + threadIdx.y; const int thread_x_idx = tileX + threadIdx.x; // читаем данные в shared memory block_data[threadIdx.x][threadIdx.y] = input[thread_y_idx * n + thread_x_idx]; } // синхронизируем ядра на уровне блока, что бы быть уверенными что shared memory заполенена полностью __syncthreads(); const int myColSave = blockDim.x * blockIdx.x + threadIdx.y; const int myRowSave = blockDim.y * blockIdx.y + threadIdx.x; if( myRowSave < m && myColSave < n ) { const int thread_y_idx = tileY + threadIdx.x; const int thread_x_idx = tileX + threadIdx.y; const int transposed_index = thread_x_idx * m + thread_y_idx; // сохраняем данные из shared memory в глобальную память. output[transposed_index] = block_data[threadIdx.y][threadIdx.x]; } }

Ух что тут происходит? На самом деле на словах понятно объяснить очень сложно легче показать на схеме. Для простоты давайте представим что мы транспонируем небольшую матрицу размер 4x4 с помощью сетки (grid) размером 2x2 блока (thread block) в каждом блоке 2x2 потока (thread). Shared memory для блока так же имеет размер 2x2 ячейки.

Тут мы видим что ядро (0, 0) в блоке с индексом (1, 0) читает элемент матрицы (значение - 2) с индексом 2 и сохраняет в shared memory под индексом (0, 0) которая тоже представляет из себя матрицу. Так как это диагональный элемент то для записи мы так же воспользуемся shared memory под индексом (0, 0). Нам нужно записать элемент матрицы (значение - 2) в новую матрицу на позицию 8 (для этого воспользуемся формулой которая указана на картинке). Так, пока все хорошо мы считали данные и записали их в правильную позицию.

Давайте теперь посмотрим на соседнее ядро в варпе (для упрощения положим что варп тоже имеет размер 2). Что мы видим читает оно элемент матрицы (значение - 3) с индексом 3 и сохраняет в shared memory под индексом (1, 0). С чтением все хорошо, соседи по варпу читают соседние элементы. Так как это не диагональный элемент то для записи мы воспользуемся shared memory под индексом (0, 1). Но сохранили то мы в shared memory под индексом (1, 0). А что читать? А читать мы будем то что записало другое ядро, об этом чуть дальше. А сохранить нам надо элемент матрицы (значение - 2) в новую матрицу на позицию 9 (воспользуемся той же самой формулой). Так что мы видим? С записью тоже все хорошо! Соседи по варпу записывают в соседние элементы в памяти. Все согласно Coalesced Memory Access правилу. Давайте посмотрим на другие ядра в это блоке.

По аналогии с предыдущими ядрам обрабатываем ядро (0, 1). Это ядро находится в одном блоке с предыдущими двумя но уже в другом варпе. И тут мы видим что это ядро как раз сохраняет данные в shared memory для чтения ядра из предыдущего шага (рис. 2 ядро (1,0) в блоке (1,0)) а само оно использует данные которые записало предыдущее ядро. То есть по сути они обменялись данными друг для друга в shared memory.

И наконец можно посмотреть как работает последнее ядро в этом блоке. И опять видим что для этого ядра (1, 1) и предыдущего (0, 1) которые являются соседями по варпу, что чтение, что запись выполняются с правилом Coalesced Memory Access. Вот и отлично. Давайте теперь проведем еще один замер скорости выполнения.

CPU transpose time: 103.60 ms Naive GPU transpose time: 1053.61 ms Shared memory GPU transpose time: 0.32 ms Speedup naive GPU vs CPU: 0.10x Speedup shared memory GPU vs naive GPU: 3243.23x Speedup shared memory GPU vs CPU: 318.91x

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

Bank conflict и Shared Memory Access Pattern

Да что еще? Последняя тема которую мы рассмотрим в данной статье, честно!

Shared Memory физически разделена на 32 равных модуля, называемых банками (bank), которые могут обрабатывать запросы параллельно. Bank Conflict возникает в том случае, если несколько потоков в одном варпе одновременно обращаются к разным адресам памяти, которые попадают в один и тот же bank. Тут на самом деле тот же принцип что и в Coalesced Memory Access соседние ядра у нас должны читать соседние ячейки памяти. В этой ситуации GPU вынуждена выполнять запросы последовательно , что кратно увеличивает задержку и практически сводит на нет преимущество высокой скорости этой памяти.

В примере выше у нас возникает bank конфликт для соседних ядер в рамках одного варпа, так как записывают они в последовательные участки памяти в shared memory, а вот читают со смещением из за транспонирования и попадают в один bank. На схеме ниже примерно видно почему возникает bank conflict.

Для того что бы это поправить нам нужно опять переписать ядро. Шутка! На самом деле поправить в данном случае это легко, достаточно просто сделать вот так:

__shared__ float block_data[ThreadsPerBlock_x][ThreadsPerBlock_y + 1];

На схеме ниже видно как это поможет решить bank conflict.

Заключение

Фух! На этом пока все. Мы разобрали относительно простые математические операции. Но рассмотрели очень важные аспекты работы с памятью. И увидели, что простой реализации алгоритма не достаточно, нужно учитывать особенности работы железа, что очень сильно усложняет написание ядер с одной стороны, с другой стороны делает это намного интересней.

Продолжение следует ...

Подробный код для статьи можно посмотреть на github. Так же если хотите следить и данная тема вам интересна можете подписаться на канал в телеграмме созданный специально для этого цикла статей, там будут анонсы статьей по циклу, обсуждения, ваши запросы и многое другое.