От MNIST к Transformer. Часть 3. Умножение тензоров. Пишем Linear Layer / Хабр

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

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

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

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

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

Самое начало. Разберемся с основами работы GPU. Напишем наш первый код на C++ и CUDA для простейшей операции сложения векторов, hello world в мире GPGPU и CUDA, чтобы понять, как CPU и видеокарта общаются друг с другом. Статья вот тут.
Основы работы с памятью. Разберемся с основами работы и устройства памяти в GPU. Напишем четыре ядра для двух операций (простой подход и оптимизированный), построение гистрограммы для вектора и транспонирование матрицы. Операция транспонирования нам понадобится в будущем что бы собирать наш Transformer. Статья вот тут .
Умножение тензоров. Пишем Linear Layer (Мы здесь). Перейдем от матриц к тензорам, напишем умножение простое умножение тензоров и потом сделаем tiled вариант c broadcasting. Создадим полносвязный слой нейронной сети и напишем прямой проход для нашей модели которая будет распозновать mnist дата сет. Так же не забываем про CUDA: немного углубимся в способы синхронизации CPU и GPU, в реализации mnist модели мы их использовать не будем для простоты, но они пригодятся в дальнейшем.

От матриц к тензорам

До этого момента мы работали с векторами (1D) или матрицами (2D). В простых задачах этого достаточно, но реальные модели, например, Transformer - оперируют данными гораздо большей размерности. Здесь на сцену выходят тензоры.

Тензор - это объект, который обобщает понятие многомерного массива данных. Скаляр (Число) - это тензор ранга 0, вектор - ранга 1, матрица - ранга 2. Все, что имеет три измерения и более, обычно называют просто тензорами произвольного ранга (или размерности).

Физическая память всегда одномерна

Важно понимать одну деталь: как бы мы ни представляли тензор в голове (в виде куба или 4D-гиперкуба), память GPU и CPU - это всегда линейный массив. Физически у нас есть только цепочка ячеек с последовательными адресами.

Многомерность - это абстракция, которую мы создаем программно, договариваясь о том, как именно интерпретировать индексы. Существует два основных способа расположения (layout) многомерных данных:

Row-major order (используется в C/C++): элементы строк идут в памяти друг за другом.
Column-major order: элементы столбцов лежат подряд.

Row-major vs Column-major мы разбирали во второй статье. Но давайте повторим: Так как мы пишем на C++, мы будем придерживаться Row-major. Это критично для производительности CUDA: если обращаться к памяти в неправильном порядке, механизмы объединения запросов (memory coalescing) в GPU не смогут работать эффективно, и скорость вычислений упадет в разы.

Чтобы найти нужный элемент в N-мерном тензоре, нам нужно вычислить его смещение (offset) относительно начала массива.

Для двумерной матрицы размером формула проста:

$index = i \cdot N+ j$

Если у нас есть тензор с размерами , то смещение для индекса вычисляется так:

$offset = i \cdot Stride_0 + j \cdot Stride_1 + k \cdot Stride_2$

Ниже представлен пример расчета stride и плоского индекса для n мерного тензора.

void Tensor::calculate_strides() {
    _strides = std::vector<int>(_shape.size(), 1);
    for (int i = _strides.size() - 2; i >= 0; i--) {
        _strides[i] = _strides[i + 1] * _shape[i + 1];
    }
}

int Tensor::calculate_flat_index(std::vector<int> indices) const {
    int index = 0;
    for (std::vector<int>::size_type i = 0; i < indices.size() - 1; i++) {
        index += indices[i] * _strides[i];
    }
    index += indices[indices.size() - 1];
    return index;
}

Разберем например тензор размеренностью 2x3x2, по сути это можно представить как две матрицы размером 3x2, ниже представлено визуальное представление тензора, как он расположен в памяти и как считается плоский индекс.

Умнож��ние тензоров

Умножение матриц

Умножение матриц - это операция, на которую тратится более 90% вычислительного времени при обучении и инференсе нейросетей. В линейном слое (Linear Layer) это выглядит как , где - входные данные, а - веса. При умножении матриц и умножаются строки матрицы A и столбцы матрицы B между собой. Умножение строк и столбцов матриц происходит по формуле:

$C_{i,j}=\sum_{k=K}^{K}A_{i,k}⋅B_{k,j}$

где - это общая размеренность, то если у нас есть матрица размерностью и матрица то должно совпадать. Например мы можем перемножить две матрицы 2x3 и 3x2, но не можем это сделать для матриц 2x3 и 2x3, в тако случае мы должны сначала транспонировать одну из матриц, соответственно если размер матриц совсем не совпадают например 2x3 и 5x7 то тут уже даже транспонирование не поможет, такие матрицы мы перемножить не сможем. Визуально перемножение матрицы изображено ниже:

Умножение тензоров

Само по себе умножение тензоров ни чем не отличается от умножение матриц, просто тензоры это по сути матрица матриц и в данном случае мы просто перемножаем вложенные матрицы. Например для тензора 2x3x2 и 2x2x3 мы по сути перемножаем четыре матрицы 3x2 и 2x3, а первую размеренность данного тензора можно назвать batch size. Так же мы можем перемножать тензоры например 3x1x3x2 и 1x3x2x3 то есть последние два измерения должны следовать правилу перемножения матриц где внутренняя размеренность должна совпадать (K), а остальные размеренности должны или совпадать или быть равны единице. Для того что бы реализовать такое умножение существует так называемое правило broadcasting для которого мы просто должно правильно считать смещение (stride) по батчам.

Реализация CUDA ядра для умножения тензоров

Для начала давайте напишем просто ядро которое не учитывает broadcasting не учитывает проблемы с памятью которые обсуждали в предыдущих статьях такие какие memory coalescing.

__global__ void naive_matmul_nd_kernel(
    const float* __restrict__ A, 
    const float* __restrict__ B, 
    float* __restrict__ C, 
    int M,
    int K,
    int N
) {
    // Identify which batch (matrix)
    int batch_idx = blockIdx.z;

    // Identify the row and column within that matrix
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    // Boundary check for the matrix
    if (row < M && col < N) {
        // Calculate the linear offset to the start of the current slice
        int offsetA = batch_idx * (M * K);
        int offsetB = batch_idx * (K * N);
        int offsetC = batch_idx * (M * N);

        float sum = 0.0f;
        for (int k = 0; k < K; ++k) {
            sum += A[offsetA + row * K + k] * B[offsetB + k * N + col];
        }
        C[offsetC + row * N + col] = sum;
    }
}

Приведенный выше код это наивный подход где не учитываются особенности железа, не учитывается любой размер тензора, а просто перемножаются две матрицы. Для начала давайте уберем memory coalescing и для этого как и в прошлых статьях мы должны вспомнить про shared memory для блока потоков (thread block). Такой подход умножения матриц называется тайлинг или Tiled Matrix Multiplication. Вот основной код данного алгоритма:

for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; ++t) {
        
    // Load tiles into shared memory (with boundary checks)
    if (row < M && (t * TILE_SIZE + tx) < K)
        tileA[ty][tx] = A[offsetA + row * K + t * TILE_SIZE + tx];
    else
        tileA[ty][tx] = 0.0f;

    if (col < N && (t * TILE_SIZE + ty) < K)
        tileB[ty][tx] = B[offsetB + (t * TILE_SIZE + ty) * N + col];
    else
        tileB[ty][tx] = 0.0f;

    // Wait for all threads to finish loading
    __syncthreads();

    // Compute the product for this tile
    for (int k = 0; k < TILE_SIZE; ++k) {
        sum += tileA[ty][k] * tileB[k][tx];
    }

    // Synchronize before loading the next tile
    __syncthreads();
}

if (row < M && col < N) {
    C[offsetC + row * N + col] = sum;
}

Опять же на словах объяснить алгоритм довольно сложно поэтому ниже представлена схема как примерно работает такая логика.

Так же для упрощения возьмем thread block размером 2x2 (размер желательно должен совпадать с размером выходной матрицы) и посмотрим как он умножает матрицы. Для начала он считывает данные из матриц и и загружает в тайлы которые расположены в shared memory. После этого потоки на уровне блока синхронизируются и записывают результат в локальную переменную sum которая потом потом сохраняется в результирующию матрицу. Так же в данном алгоритме если посмотреть внимательно нет bank конфликта потому что потоки в одно варпе или читают одну и ту же область памяти или из разных bank-ов.

Полный код ядра для умножения тензоров с учетом бродкастинга представлен ниже. Для броадкастинга мы должны передать смещения или stride для внешних размерностей тензора (без последних двух на самом деле, например для тензора 3x3x2x2, внешние размеренности это 3x3)

__global__ void tiled_matmul_broadcast_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M,
    int K,
    int N,
    int num_batch_dims,
    const int* __restrict__ out_batch_shape,
    const int* __restrict__ a_batch_strides, // Pre-computed (0 if broadcasted)
    const int* __restrict__ b_batch_strides, // Pre-computed (0 if broadcasted)
    const int* __restrict__ c_batch_strides
) {
    // Shared memory for tiles
    __shared__ float tileA[TILE_SIZE][TILE_SIZE];
    __shared__ float tileB[TILE_SIZE][TILE_SIZE];

    int bx = blockIdx.x; 
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    int temp_idx = blockIdx.z;
    int offsetA = 0, offsetB = 0, offsetC = 0;

    for (int i = num_batch_dims - 1; i >= 0; --i) {
         int coord = temp_idx % out_batch_shape[i];
         temp_idx /= out_batch_shape[i];
         offsetA += coord * a_batch_strides[i];
         offsetB += coord * b_batch_strides[i];
         offsetC += coord * c_batch_strides[i];
    }

    // Row/Col in the output matrix C
    int row = by * TILE_SIZE + ty;
    int col = bx * TILE_SIZE + tx;

    float sum = 0.0f;

    // Loop over the tiles of the input matrices
    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; ++t) {
        
        // Load tiles into shared memory (with boundary checks)
        if (row < M && (t * TILE_SIZE + tx) < K)
            tileA[ty][tx] = A[offsetA + row * K + t * TILE_SIZE + tx];
        else
            tileA[ty][tx] = 0.0f;

        if (col < N && (t * TILE_SIZE + ty) < K)
            tileB[ty][tx] = B[offsetB + (t * TILE_SIZE + ty) * N + col];
        else
            tileB[ty][tx] = 0.0f;

        // Wait for all threads to finish loading
        __syncthreads();

        // Compute the product for this tile
        for (int k = 0; k < TILE_SIZE; ++k) {
            sum += tileA[ty][k] * tileB[k][tx];
        }

        // Synchronize before loading the next tile
        __syncthreads();
    }

    if (row < M && col < N) {
        C[offsetC + row * N + col] = sum;
    }
}

Запуск такого ядра происходит таким образом:

void launchTiledMatmul(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M,
    int K,
    int N,
    int batch_size,
    int num_batch_dims,
    const int* __restrict__ out_batch_shape,
    const int* __restrict__ a_batch_strides,
    const int* __restrict__ b_batch_strides,
    const int* __restrict__ c_batch_strides
) {
    dim3 blockDim(ThreadsPerBlock_x, ThreadsPerBlock_y);
    dim3 gridDim(
        cuda::ceil_div(M, ThreadsPerBlock_x),
        cuda::ceil_div(N, ThreadsPerBlock_y),
        batch_size
    );
    tiled_matmul_broadcast_kernel<<<gridDim, blockDim>>>(
        A, B, C,
        M, K, N,
        num_batch_dims, out_batch_shape, a_batch_strides, b_batch_strides, c_batch_strides
    );
}

batch_size - в данном случае это общая "внешняя" размеренность тензора, например для тензора 3x3x2x2 batch_size будет равен 9. Так же для упрощения просчета смещения (offset) при броадкастинге если одна из размерностей тензора равна единице, то stride для этой размеренности будет равен 0. Для этого мы должны заранее просчитать смещения:

for (int i = 0; i < matmul_info.num_batch_dims; i++) {
    a_batch_strides[i] = A.shape()[i] == 1 ? 0 : A.strides()[i];
    b_batch_strides[i] = B.shape()[i] == 1 ? 0 : B.strides()[i];
    c_batch_strides[i] = C.strides()[i];
}

На этом с умножением тенозоров пока все, давайте пойдем дальше.

Способы синхронизации CUDA

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

Когда мы запускаем ядро kernel_name<<<...>>>, управление мгновенно возвращается в поток CPU. Это асинхронность из коробки. Однако для построения сложных цепочек вычислений (например, когда следующий слой нейросети ждет результатов предыдущего) нам нужны механизмы управления очередностью.

CUDA Streams: Основа параллелизма

Stream - это потоки выполнения ядер. По умолчанию все операции выполняются в так называемом default stream последовательно. Но чтобы занять GPU на 100% мы можем использовать несколько стримов. Это позволяет одновременно копировать данные из RAM в VRAM и выполнять вычисления например.

cudaStream_t stream1, stream2;
cudaStreamCreate(&stream1);
cudaStreamCreate(&stream2);

// Асинхронное копирование в первом стриме
cudaMemcpyAsync(d_A, h_A, size, cudaMemcpyHostToDevice, stream1);

// Запуск вычислений во втором стриме (если они не зависят от d_A)
kernel<<<grid, block, 0, stream2>>>(d_B);

// Барьер для CPU: ждем завершения всех задач в конкретном стриме
cudaStreamSynchronize(stream1);

Programmatic Dependent Kernel Launch или PDL

В CUDA 12 появился механизм Programmatic Dependency. Раньше, чтобы запустить второе ядро строго после первого, мы либо использовали один стрим, либо синхронизировались через CPU. Теперь мы можем настроить запуск так, чтобы можно было запустить два ядра одновременно если есть на это ресурсы и если одно ядро зависит от другого то мы можем во втором ядре подождать зависимых данных из первого.

cudaLaunchConfig_t config = {0};
config.stream = my_stream;

cudaLaunchAttribute attr[1];
attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attr[0].val.programmaticStreamSerializationAllowed = 1;
config.attrs = attr;
config.numAttrs = 1;

// запуск основного ядра
cudaLaunchKernelEx(&config, primary_kernel);

cudaLaunchKernelEx(&config, secondary_kernel);

Примерный код внутри самих ядер:

__global__ void primary_kernel() {
  // Начальная работа которая нужна для работы второго ядра
  // Можно запускать второе ядро
  cudaTriggerProgrammaticLaunchCompletion();
  // Work that can coincide with the secondary kernel
}

{
__global__ void secondary_kernel()
  // Делаем работу для которой не нужны данные из первого ядра
  // Ждем выполенния первого ядра
  cudaGridDependencySynchronize();
  // прододолжаем работу с зависимыми данными
}

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

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

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

Первый Linear Layer

Линейный слой выполняет простую аффинную трансформацию входных данных. Если отбросить терминологию ML, это просто умножение матрицы входных данных на матрицу весов с последующим прибавлением вектора смещений. Для входного тензора и параметров слоя (весов и смещений ) операция выглядит так:

Ниже представлен код для линейного слоя.

class LinearLayer {
public:
    std::unique_ptr<Tensor> weights; // Shape: (in_features, out_features)
    std::unique_ptr<Tensor> bias; // Shape: (1, out_features)
    std::unique_ptr<Tensor> output; // Shape: (batch_size, out_features)

    LinearLayer(int in_features, int out_features, int batch_size) {
        // in PyTorch, it's (out_features, in_features),
        // but for now we use (in_features, out_features) for simplicity
        float scale = std::sqrt(2.0f / in_features);
        weights = Tensor::create_random({in_features, out_features}, -scale, scale);
        bias = Tensor::create_zeros({1, out_features});
        output = Tensor::create_zeros({batch_size, out_features});
    }

    void forward(Tensor& input) {
        MatmulOperation::matmul_inplace(input, *weights, *output);
        TensorOperation::add_inplace(*output, *bias);
    }

    void prepare_for_gpu_work() {
        weights->prepare_for_gpu_work();
        bias->prepare_for_gpu_work();
        output->prepare_for_gpu_work();
    }
};

Хи (He) инициализация (или Kaiming инициализация)

Вы возможно заметили такую строчку выше

float scale = std::sqrt(2.0f / in_features);

В данном случае мы не просто случайно инициализируем веса в диапазоне от 0 до 1. А используется специальна формула для определения границ. Почему так? Потому что если мы просто заполним веса слишком маленькими значениями, активации на выходе превратятся в нули. Слишком большими - уйдут в бесконечность. Для сетей с функцией активации ReLU стандартом является Инициализация Хи (He Initialization), также известная как Kaiming Initialization.

ReLU

На выходе линейного слоя обычно используется функция активации раньше были популярны функции такие как сигмоида или tanh, но они довольно сложные в вычислительном плане. Ко всему прочему для сигмоиды производная в области больших значений близка к нулю. Это значит, что при обратном распространении ошибки градиент затухает (Vanishing Gradient), и веса первых слоев перестают обновляться. В данном случае ReLU имеет ряд преимуществ, плюс ReLU как выключает часть нейронов (зануляет их). Это создает разреженную активацию, что часто помогает сети лучше обобщать данные и делает её более устойчивой к шуму.

Формула ReLU очень простая:

То есть мы просто зануляем отрицательные значения, все просто.

Код CUDA ядра для ReLU выглядит довольно просто

__global__ void relu_kernel(float* __restrict__ A, int size) {
    // we don't use grid stride loop for simplicity
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        A[idx] = max(A[idx], 0.0f);
    }
}

Neural Network для MNIST

MNIST (Modified National Institute of Standards and Technology) - это классический набор данных, состоящий из рукописных цифр. Он стал эталоном для проверки алгоритмов распознавания, потому что он достаточно мал, чтобы обучаться быстро, но достаточно сложен, чтобы на нем нельзя было достичь высокой точности простой линейной регрессией. Внутри датасета черно-белая картинка размером 28×28 пикселей и еще есть label число от 0 до 9, соответствующее тому, что нарисовано на картинке.

Для распознавания цифр MNIST нам понадобится простая архитектура:

Превращаем картинку 28×28 в вектор из 784 чисел.
Первый линейный слой: [784→128].
функция активации ReLU
Второй и последний слой: [128→10]. 10 выходов - по одному на каждую цифру.

Код для модели для распознования MNIST:

class MNISTModel {
public:
    LinearLayer fc1; // 784 -> 128
    LinearLayer fc2; // 128 -> 10

    MNISTModel(int batch_size) : fc1(784, 128, batch_size), fc2(128, 10, batch_size) {}

    void prepare_for_gpu_work() {
        fc1.prepare_for_gpu_work();
        fc2.prepare_for_gpu_work();
    }

    void forward(Tensor& input) {
        fc1.forward(input);
        TensorOperation::relu_inplace(*fc1.output);
        fc2.forward(*fc1.output);
        cudaDeviceSynchronize();
        fc2.output->copy_to_host();
        TensorOperation::softmax_inplace_cpu(*fc2.output);
    }

    std::vector<int> predict(Tensor& input) {
        forward(input);
        print_vector("Predictions", fc2.output->shape());
        fc2.output->print();
        std::vector<int> predictions;
        for (int i = 0; i < fc2.output->shape()[0]; i++) {
            int max_val_index = std::max_element(
                fc2.output->data() + i * fc2.output->shape()[1],
                fc2.output->data() + (i + 1) * fc2.output->shape()[1]
            ) - fc2.output->data();
            predictions.push_back(max_val_index - fc2.output->shape()[1] * i);
        }
        return predictions;
    }
};

Функция predict использует softmax для получения распределения финальной вероятности.

Softmax

Когда наша нейросеть прогнала данные через все линейные слои и ReLU, на выходе мы получаем 10 чисел (обычно это называется logits). Эти числа могут быть любыми. Но нам нужно понять, какова вероятность того, что на картинке изображена «пятерка». Для этого используется Softmax.

Softmax - это финальный этап классификации. Он берет вектор произвольных чисел и преобразует его в вектор вероятностей, где сумма всех элементов строго равна 1.

Для каждого элемента входного вектора Softmax вычисляется следующим образом:

$S(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{K} e^{x_j}} $

Так как у нас выходной вектор довльно маленький softmax мы вычислим на CPU:

void softmax_inplace_cpu(Tensor& A) {
    float max_val = *std::max_element(A.data(), A.data() + A.total_size());

    float sum = 0.0f;
    for (int i = 0; i < A.total_size(); i++) {
        float val = A.data()[i];
        val = std::exp(val - max_val); // Subtract max_val for stability
        sum += val;
        A.data()[i] = val;
    }

    for (int i = 0; i < A.total_size(); i++) {
        A.data()[i] = A.data()[i] / sum;
    }
}

Заключение

На этом пока все. Мы разобрали тензора и умножение тензоров и матриц. Рассмотрели способ синхронизации CUDA такие как стримы и PDL. Так же написали наш первый линейный слой и собрали сеть для распознавания mnist датасета. В следующей статье мы рассмотрим самое интересное - обучение. Напишем функцию ошибки и оптимизатор.

Надеюсь было понятно и самое главное интересно. Продолжение следует ...

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