Ultra-Scale Playbook, часть 4: внутри GPU — ядра, FlashAttention, смешанная точность

Четвёртая, заключительная часть перевода книги The Ultra-Scale Playbook. Предыдущие части: 1 — память на одной GPU, 2 — DP, ZeRO, TP и SP, 3 — CP, PP, EP и 5D.

Как читать вместе с курсом ChinaAI

Эта глава почти точно ложится на модуль 5.7 «GPU изнутри»: 5.7.1 «Архитектура GPU и профилирование», 5.7.2 «FlashAttention», 5.7.3 «Ядра на Triton». Раздел о смешанной точности дополняет урок 5.4.1. Числа в оригинале — для H100; у вашей RTX 4080 (архитектура Ada) другие объёмы, но иерархия памяти та же. Код CUDA в оригинале сплющен в одну строку — мы восстановили форматирование и исправили явные опечатки.

Погружение в GPU: слияние, потоки и смешивание точностей

Заметка

В оригинале и здесь можно послушать подкаст NotebookLM по следующим разделам.

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

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

В этом разделе мы углубимся в детали архитектуры GPU. Сосредоточимся на архитектуре NVIDIA, но общие идеи, как это часто бывает, применимы и к похожим ускорителям.

Сначала кратко объясним, как устроены GPU, затем перейдём к революции FlashAttention, к эффективному планированию работы на GPU и к эффективному использованию разных точностей.

Краткое введение в GPU

В целом GPU устроены очень иерархично. С вычислительной стороны GPU — это массив вычислительных блоков, которые называются потоковыми мультипроцессорами (streaming multiprocessors, SM). Каждый SM содержит набор потоковых процессоров, они же ядра (cores), и управляет ими. Например, у NVIDIA H100 132 SM по 128 ядер, всего 16 896 ядер (подробнее — в документации по тензорным ядрам), и каждое может одновременно обслуживать несколько потоков.

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

Со стороны памяти тоже сильная иерархия с несколькими уровнями кэша и памяти. Регистры — самые маленькие единицы, они принадлежат потокам во время выполнения. Разделяемая память (shared memory) и кэш L1 общие для потоков, работающих на одном SM. Выше — кэш L2, общий для всех SM, и, наконец, глобальная память (global memory) — самая большая память GPU (например, заявленные 80 ГБ у H100), но и самая медленная для обращений.

flowchart TB
R["Регистры: свои у каждого потока, быстрее всего"] --> S["Shared memory и L1: общие для потоков одного SM"]
S --> L2["Кэш L2: общий для всех SM"]
L2 --> G["Глобальная память HBM: десятки ГБ, медленнее всего"]

Схема переводчика: иерархия памяти GPU (вниз — больше объём, ниже скорость)

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

Код, выполняющийся на ядре GPU, называется ядром (kernel). Его можно написать на высоком уровне, например на CUDA или Triton, а затем он компилируется в Parallel Thread Execution (PTX) — низкоуровневый ассемблер GPU NVIDIA.

Чтобы запустить ядро, нужен ещё код хоста (host code): он выполняется на CPU (хосте) и выделяет память, загружает данные и код.

// Код хоста
void vecAdd(float* h_A, float* h_B, float* h_C, int n) {
    // Выделить векторы в памяти устройства (GPU)
    int size = n * sizeof(float);
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, size);
    cudaMalloc(&d_B, size);
    cudaMalloc(&d_C, size);

    // Скопировать векторы из памяти хоста в память устройства
    cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);

    // Запустить ядро
    int threadsPerBlock = 256;
    int blocksPerGrid = (n + threadsPerBlock - 1) / threadsPerBlock;
    VecAdd<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, n);

    // Скопировать результат из памяти устройства в память хоста
    // h_C содержит результат в памяти хоста
    cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);

    // Освободить память устройства
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
}

Код хоста для ядра CUDA, складывающего два вектора; адаптирован по CUDA C Programming Guide и блогу Coding Confessions.

// Код устройства
__global__ void VecAdd(float* A, float* B, float* C, int N) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;  // глобальный номер потока
    if (i < N)                                      // лишние потоки ничего не делают
        C[i] = A[i] + B[i];
}

Код устройства с определением ядра сложения векторов (из тех же источников).

Ядра обычно планируются так:

• потоки группируются в варпы (warps) по 32 потока. Все потоки варпа синхронно выполняют одни и те же инструкции, но над разными частями данных;
• варпы объединяются в более крупные блоки (blocks) гибкого размера (например, 512 или 1024 потока в блоке), и каждый блок назначается одному SM. SM может выполнять несколько блоков параллельно. Но в зависимости от доступных ресурсов не все блоки сразу получают место; некоторые ждут в очереди, пока освободятся ресурсы.

Главное, что стоит запомнить: есть разные ограничения на размеры и распределение ресурсов (размеры разных видов памяти, число одновременно работающих блоков и потоков в варпах), и их нужно учитывать, чтобы использовать архитектуру GPU максимально эффективно.

Чаще всего спускаться до такого уровня точности не нужно: можно использовать ядра и код, подготовленные другими людьми из сообщества. Но несколько советов для начала работы с ядрами мы всё же дадим!

Как ускорить работу с помощью ядер

Если вы хотите добавить новую операцию, для которой нет оптимизированного ядра, или ускорить существующую функцию PyTorch, написать ядро с нуля может показаться самым прямым путём. Но создание высокопроизводительных ядер CUDA с нуля требует большого опыта, и кривая обучения крутая. Обычно лучше начать с torch.compile: он динамически оптимизирует код PyTorch, перехватывая ваши операции и генерируя низкоуровневые высокопроизводительные ядра на Triton.

Пусть вы хотите написать ядро для функции активации ELU (Exponential Linear Unit):

$$\text{ELU}(x) = \begin{cases} e^x - 1 & \text{если } x < 0 \\ x & \text{если } x \geq 0 \end{cases}$$

Можно начать с простой реализации на PyTorch и просто добавить сверху декоратор @torch.compile:

import torch

@torch.compile
def elu(x, alpha=1.0):
    # там, где x < 0, берём alpha * (e^x - 1), иначе оставляем x
    return torch.where(x < 0, alpha * (torch.exp(x) - 1), x)

На графике оригинала видна огромная разница в производительности между скомпилированной и нескомпилированной версиями — и это при том, что мы добавили только декоратор (N на графике — число столбцов).

Если такого прироста не хватает, можно реализовать ядра на Triton. Для начала можно посмотреть на ядро Triton, сгенерированное @torch.compile. Для этого достаточно установить переменную окружения TORCH_LOGS в "output_code":

export TORCH_LOGS="output_code"

Если запустить скрипт Python с декоратором @torch.compile, он сгенерирует и выведет соответствующее ядро Triton — в нашем случае такое:

@triton.jit
def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 100000000
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = 0.0
    tmp2 = tmp0 < tmp1
    tmp3 = tl_math.exp(tmp0)
    tmp4 = 1.0
    tmp5 = tmp3 - tmp4
    tmp6 = tl.where(tmp2, tmp5, tmp0)
    tl.store(out_ptr0 + (x0), tmp6, xmask)

Для читаемости можно переименовать переменные, добавить комментарии и немного поправить код (или попросить LLM сделать это за нас):

@triton.jit
def elu_kernel(input_ptr, output_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    # Начальный индекс для этого блока
    block_start = tl.program_id(0) * BLOCK_SIZE
    # Массив индексов для этого блока
    block_indices = block_start + tl.arange(0, BLOCK_SIZE)[:]
    # Маска, чтобы обрабатывать только допустимые индексы
    valid_mask = block_indices < num_elements
    # Загрузить входные значения по допустимым индексам
    input_values = tl.load(input_ptr + block_indices, valid_mask)
    # Параметры ELU
    zero_value = 0.0  # порог активации ELU
    negative_mask = input_values < zero_value
    exp_values = tl.math.exp(input_values)
    # Сдвиг выхода ELU
    one_value = 1.0
    shifted_exp_values = exp_values - one_value
    output_values = tl.where(negative_mask, shifted_exp_values, input_values)
    # Записать вычисленные значения обратно по указателю выхода
    tl.store(output_ptr + block_indices, output_values, valid_mask)

Здесь tl.program_id(0) даёт уникальный номер блока, по которому мы определяем, какой участок данных обрабатывает этот блок. По номеру блока block_start вычисляет начальный индекс участка, а block_indices задаёт диапазон индексов внутри него. valid_mask гарантирует, что обрабатываются только индексы меньше num_elements, и данные безопасно загружаются через tl.load. Затем применяется функция ELU (значения меняются в зависимости от знака), и результаты записываются обратно в память через tl.store.

На бенчмарке сгенерированного ядра через triton.testing.Benchmark (график в оригинале) это отдельное ядро даже быстрее @torch.compile на малых размерах, но, скорее всего, это лишь артефакт времени компиляции torch.compile. В любом случае не обязательно начинать с нуля: можно взять такие сгенерированные ядра за основу и сосредоточиться на оптимизации их производительности, сэкономив много времени.

Даже на Triton иногда не удаётся достичь пиковой производительности устройства из-за ограничений языка в работе с низкоуровневыми деталями вроде разделяемой памяти и планирования внутри потоковых мультипроцессоров. Возможности Triton ограничены блоками и их распределением по SM. Для ещё более глубокого контроля придётся писать ядра прямо на CUDA, где доступны все низкоуровневые детали.

В CUDA для повышения эффективности ядер есть разные приёмы. Мы разберём лишь несколько: оптимизация шаблонов обращения к памяти для снижения задержки, использование разделяемой памяти для часто используемых данных и распределение работы между потоками для сокращения простоев.

Прежде чем перейти к примерам на CUDA, подытожим инструменты, позволяющие писать код ядер для GPU:

1. PyTorch: просто, но медленно;
2. @torch.compile: просто, быстро, но не гибко;
3. Triton: сложнее, быстрее и гибче;
4. CUDA: сложнее всего, быстрее всего и гибче всего (если сделать правильно).

Начнём с одного из самых частых применений CUDA — оптимизации обращений к памяти. Глобальная память GPU (самая большая, как мы видели) по сравнению с кэшем имеет большую задержку и низкую пропускную способность, и для многих приложений это главное узкое место. Эффективный доступ к глобальной памяти может сильно ускорить работу.

Объединение обращений к памяти

Чтобы эффективно использовать пропускную способность глобальной памяти, нужно понимать её устройство. В устройствах CUDA глобальная память реализована на DRAM.

Объединение обращений к памяти (memory coalescing) использует то, что DRAM выдаёт данные пачками при каждом обращении к адресу. При каждом обращении к ячейке DRAM несколько датчиков чипа параллельно считывают последовательность соседних ячеек (включая запрошенную). Прочитанные данные затем быстро передаются процессору одной пачкой. В CUDA объединение использует это поведение, чтобы максимально эффективно обращаться к памяти: потоки варпа (32 потока, выполняющих одну и ту же инструкцию в ногу) должны обращаться к последовательным ячейкам памяти. Например, если поток 0 обращается к ячейке M, поток 1 — к M + 1, поток 2 — к M + 2 и так далее, железо GPU объединяет эти запросы в один большой эффективный запрос на пачку DRAM, а не обрабатывает каждое обращение по отдельности.

Возьмём умножение матриц. В простой прямолинейной реализации каждый поток вычисляет один элемент выходной матрицы:

__global__ void matmul_naive(int M, int N, int K, const float *A, const float *B, float *C) {
    const uint x = blockIdx.x * blockDim.x + threadIdx.x;  // строка результата
    const uint y = blockIdx.y * blockDim.y + threadIdx.y;  // столбец результата

    if (x < M && y < N) {
        float tmp = 0.0;
        for (int i = 0; i < K; ++i) {
            tmp += A[x * K + i] * B[i * N + y];
        }
        C[x * N + y] = tmp;
    }
}

Отличная визуализация этого ядра — в фантастическом посте Саймона Бёма.

Однако профилирование этого ядра инструментом вроде ncu показывает проблемы, в том числе низкую пропускную способность памяти и необъединённые обращения к памяти.

Причина в том, что в этом ядре два потока одного блока с номерами (0, 0) и (1, 0) (которые окажутся в одном варпе) загружают элементы одного и того же столбца матрицы B, но разных строк матрицы A. Поскольку элементы матрицы хранятся построчно (row-major: элементы строки лежат по последовательным адресам), на первой итерации $i = 0$ поток (0, 0) загрузит $A_{0,0}$, а поток (1, 0) — $A_{1,0}$. Эти элементы лежат в памяти не рядом, и такое рассогласование повторяется на каждой итерации, не давая объединять обращения к памяти.

Пример переводчика

Пусть $K = 4$. Элементы $A$ лежат в памяти так: $A_{0,0}, A_{0,1}, A_{0,2}, A_{0,3}, A_{1,0}, \ldots$ Поток (0, 0) читает адрес 0, поток (1, 0) — адрес 4, поток (2, 0) — адрес 8: между соседними потоками «дыры», и одна пачка DRAM их не покрывает.

Чтобы ускорить ядро, можно изменить вычисление координат x и y так:

const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);
const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);

if (x < M && y < N) {
    float tmp = 0.0;
    for (int i = 0; i < K; ++i) {
        tmp += A[x * K + i] * B[i * N + y];
    }
    C[x * N + y] = tmp;
}

Вместо двумерного блока мы переходим к одномерному и по-новому определяем x и y. Теперь потоки одного варпа (у них близкие threadIdx.x) имеют одинаковый x, но разные y. Значит, они загружают одну и ту же строку матрицы A, но разные столбцы матрицы B. В результате обращения к памяти для построчно хранимой матрицы можно объединить.

При профилировании нового ядра предупреждение о необъединённых обращениях исчезает, а пропускная способность памяти GPU вырастает примерно в 10 раз.

Заодно время выполнения ядра уменьшается в 10 раз. Потрясающе!

Теперь разберём ещё один приём, который часто упоминается в литературе, — тайлинг (tiling).

Тайлинг

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

Например, при умножении матриц каждому потоку блока могут понадобиться элементы двух матриц, A и B. Если каждый поток независимо загружает нужные ему строку и столбец из глобальной памяти, получится много лишних загрузок: несколько потоков блока обращаются к пересекающимся данным. Вместо этого с помощью тайлинга можно один раз загрузить блок («плитку», tile) A и B в разделяемую память, и все потоки блока будут использовать одни и те же общие данные.

При тайлинге на каждой итерации все потоки блока сообща загружают две плитки — одну из матрицы A, другую из матрицы B — в разделяемую память. Точнее, потоки загружают плитку матрицы A (размера BLOCK_SIZE_M на BLOCK_SIZE_K) и плитку матрицы B (размера BLOCK_SIZE_K на BLOCK_SIZE_N). Когда плитки в разделяемой памяти, потоки перемножают их — это эффективно, потому что все нужные данные быстро доступны. Результаты умножения плиток записываются в накопительную матрицу промежуточных результатов. После каждой итерации результат умножения текущих плиток прибавляется к ней, и так до тех пор, пока не будут обработаны все плитки обеих матриц. (Рисунок в оригинале — из учебника Седрика Ньюгерена.)

Посмотрим на важные части реализации:

// Установить указатели на начальные элементы
A += blockRow * TILE_SIZE * K;                        // начать со строки blockRow, столбца 0
B += blockCol * TILE_SIZE;                            // начать со строки 0, столбца blockCol
C += blockRow * TILE_SIZE * N + blockCol * TILE_SIZE; // строка blockRow, столбец blockCol
float sum = 0.0;
// Внешний цикл идёт по плиткам A (вдоль столбцов) и B (вниз по строкам)
for (int tileIdx = 0; tileIdx < K; tileIdx += TILE_SIZE) {
    sharedA[localRow * TILE_SIZE + localCol] = A[localRow * K + localCol];
    sharedB[localRow * TILE_SIZE + localCol] = B[localRow * N + localCol];

    // Дождаться, пока все потоки блока закончат загрузку
    __syncthreads();

    // Сдвинуть указатели на следующую плитку
    A += TILE_SIZE;
    B += TILE_SIZE * N;

    // Частичное скалярное произведение для этой плитки
    for (int i = 0; i < TILE_SIZE; ++i) {
        sum += sharedA[localRow * TILE_SIZE + i] * sharedB[i * TILE_SIZE + localCol];
    }
    // Снова синхронизироваться, чтобы ни один поток не начал загружать новые данные
    // в разделяемую память, пока остальные не закончили вычисления
    __syncthreads();
}
C[localRow * N + localCol] = sum;

Для простоты берём квадратную плитку.

Каждый поток сначала загружает по одному элементу из матрицы A и из матрицы B в разделяемую память. Здесь легко добиться объединённых обращений: если взять threadIdx.x в качестве локального номера столбца (localCol), потоки одного варпа будут обращаться к соседним элементам обеих матриц. Когда все потоки блока загрузили свои элементы в разделяемую память (это гарантирует вызов __syncthreads()), они вычисляют скалярное произведение двух плиток. Пройдя все плитки — по горизонтали для A и по вертикали для B, — итоговая сумма записывается в соответствующую ячейку матрицы C.

На бенчмарке этого ядра через ncu пропускная способность памяти выросла до 410 Гбит/с, а время выполнения ядра сократилось примерно на 43 %, что дало производительность около 6,6 TFLOPS.

Укрупнение потоков

Тайлинг заметно ускорил ядро. Однако анализ состояний варпов — сколько тактов проведено в каждом состоянии — показывает следующую картину (график в оригинале).

Значение этих загадочных названий состояний можно найти в разделе «Warp Stall Reasons» руководства NVIDIA по профилированию ядер. Там сказано, что smsp__pcsamp_warps_issue_stalled_mio_throttle означает: «Варп простаивал в ожидании, пока очередь инструкций MIO (memory input/output) перестанет быть заполненной. Эта причина простоя велика при экстремальной загрузке конвейеров MIO, куда входят специальные математические инструкции, динамические ветвления и инструкции разделяемой памяти. Если причина в обращениях к разделяемой памяти, уменьшить нагрузку на конвейер может использование меньшего числа, но более широких загрузок».

Похоже, варпы простаивают в ожидании ответа от разделяемой памяти! Решить это можно приёмом укрупнения потоков (thread coarsening): несколько потоков объединяются в один «укрупнённый». Это сильно сокращает обращения к разделяемой памяти, потому что каждый укрупнённый поток может обрабатывать несколько выходных элементов.

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

Минимизация расхождения потока управления

Потоковый мультипроцессор выполняет все потоки варпа по модели SIMD (Single Instruction, Multiple Data — «одна инструкция, много данных»). В каждый момент одна инструкция выбирается и выполняется одновременно для всех потоков варпа. Потоки варпа работают над разными частями данных, но выполняют одну и ту же инструкцию (отсюда и название). Главное преимущество SIMD — эффективность: управляющее железо, отвечающее за выборку и рассылку инструкций, общее для многих исполнительных блоков. Это сокращает аппаратные затраты на управление, и бо́льшая часть железа может работать на арифметическую пропускную способность.

Расхождение потока управления (control divergence) возникает, когда потоки одного варпа идут разными путями выполнения. Например, если условный оператор (вроде if) приводит к тому, что одни потоки выполняют один блок кода, а другие — другой, варпу приходится выполнять их последовательно, и часть потоков простаивает в ожидании остальных. Чтобы минимизировать расхождение, ядра нужно проектировать так, чтобы потоки одного варпа шли одним путём. Этого можно добиться, перестроив код с меньшим числом ветвлений, используя структуры данных, при которых все потоки идут похожими путями, или применяя приёмы вроде предикации (predication).

Мы разобрали основные соображения при написании своих ядер и улучшении производительности и потребления памяти операций на GPU. Но перед реальным примером нужно рассмотреть ещё одно важное понятие — слияние ядер (kernel fusion).

Слитые ядра

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

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

Эту идею прекрасно иллюстрируют схемы Горация Хе (Horace He). На первой — последовательность ядер, каждое из которых требует хождения туда и обратно между глобальной памятью и вычислительными блоками. На второй — вместо того чтобы отправлять наш «треугольник» обратно в глобальную память только для того, чтобы снова его прочитать, мы выполняем все операции за один заход.

flowchart LR
subgraph NF["Без слияния: 3 ядра"]
  direction LR
  M1["HBM"] --> K1["ядро: x*2"] --> M2["HBM"] --> K2["ядро: +1"] --> M3["HBM"] --> K3["ядро: relu"] --> M4["HBM"]
end
subgraph FU["Слитое ядро"]
  direction LR
  N1["HBM"] --> KF["одно ядро: relu(x*2+1)"] --> N2["HBM"]
end

Схема переводчика: слияние поэлементных операций убирает лишние чтения и записи в глобальную память

Как избежать такого хождения туда и обратно? Лучший способ — сделать GPU как можно более самостоятельной. Для этого как можно больше последовательных вычислительных операций упаковывают в одно ядро, которое GPU выполняет целиком, — «слитое ядро» (fused kernel).

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

В модели-трансформере такое «слияние» можно применять всякий раз, когда идёт цепочка поэлементных операций, например в вычислениях слоёв LayerNorm.

Теперь у нас есть всё, чтобы оценить настоящий шедевр инженерии ядер — FlashAttention.

FlashAttention

FlashAttention предложил Три Дао (Tri Dao): он оптимизировал вычисление внимания, написав свои ядра CUDA, которые сделали его гораздо быстрее и экономнее по памяти. Идея FlashAttention — эффективно использовать разные виды памяти GPU и не слишком полагаться на самую медленную — глобальную.

Глобальная память современных GPU часто построена по технологии High Bandwidth Memory (HBM) («память с высокой пропускной способностью»), которая, несмотря на название, в иерархии памяти GPU медленнее SRAM. Термин HBM понадобится, когда будем обсуждать детали реализации FlashAttention.

Простая реализация механизма внимания требует много передач между памятью и исполнителями. Нужно материализовать в HBM матрицу S ($S = QK^T$, оценки внимания) и матрицу P ($P = \text{softmax}(S)$, нормированные веса внимания), то есть результаты отправляются в HBM, а затем обратно в SRAM для следующих вычислений.

flowchart LR
Q["Q, K из HBM"] --> S["S = QKᵀ → записать в HBM"]
S --> P["прочитать S, P = softmax(S) → записать в HBM"]
P --> O["прочитать P и V, O = PV → записать в HBM"]

Схема переводчика: наивное внимание трижды гоняет большие матрицы через медленную HBM

Пропускная способность HBM гораздо ниже, поэтому здесь возникает серьёзное узкое место. Можно ли лучше? Три Дао говорит: да!

Ключевой элемент — вычислять матрицу S маленькими кусками, помещающимися в небольшую разделяемую память SM. Но можно сделать ещё лучше и вообще не материализовать огромную матрицу S, храня только статистики, нужные для вычисления нормирующего множителя softmax. Тогда часть O можно вычислить сразу, за одно вычисление в SRAM, не гоняя промежуточные результаты туда и обратно. При этом мы не только используем разделяемую память, но и снимаем узкое место по памяти, вызванное материализацией одной из самых больших матриц активаций в модели (при длинном контексте) — матрицы внимания. (Рисунок в оригинале — из статьи о FlashAttention.)

Пример переводчика: что такое «онлайн-softmax»

Softmax строки можно считать по кускам. Храним текущий максимум $m$ и сумму $\ell = \sum e^{s_j - m}$. Пришёл новый кусок оценок с максимумом $m'$: новый максимум $m_{new} = \max(m, m')$, а старую сумму пересчитываем как $\ell \cdot e^{m - m_{new}}$ и добавляем вклад нового куска. Так же масштабируется уже накопленная сумма $\sum e^{s_j - m} v_j$. В конце делим на $\ell$ — и получаем точный результат, не храня всю строку $S$. Подробно — в уроке 5.7.2.

Идея FlashAttention устраняет столько узких мест при обучении моделей, что быстро стала стандартным способом вычислять внимание во всех трансформерах. В частности:

• не материализуя матрицу S, мы снижаем нагрузку внимания на память;
• мы также убираем бо́льшую часть прямого влияния квадратичной стоимости $O(S^2)$ внимания.

Все варианты линейного внимания и субквадратичные подходы к приближению внимания (появившиеся вскоре после изобретения архитектуры трансформера) в основном отошли на второй план в пользу этого точного и быстрого механизма FlashAttention.

После FlashAttention-1 та же лаборатория выпустила две улучшенные версии: FlashAttention-2 и -3. По сравнению с FlashAttention-1 улучшения в версиях 2 и 3 касаются не столько общего механизма внимания, сколько более точной подгонки низкоуровневой реализации под GPU: (1) как можно меньше операций, не являющихся умножением матриц, (2) аккуратное распределение работы между варпами и блоками потоков (FlashAttention-2) и (3) тщательная оптимизация под FP8 и тензорные ядра новейшей архитектуры Hopper (H100) (FlashAttention-3).

FlashAttention ограничивает, какие шаблоны внимания можно ускорить. Посмотрите на FlexAttention — быстрый и гибкий вариант.

FlashAttention — мастерская демонстрация прорывных улучшений, которых можно добиться, учитывая внутреннее устройство памяти и вычислений современных GPU-ускорителей.

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

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

Обучение со смешанной точностью

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

Обучение со смешанной точностью (mixed precision training), как следует из названия, смешивает разные точности при обучении. По умолчанию тензоры PyTorch хранятся в формате с плавающей точкой одинарной точности, он же FP32 или float32: каждое число занимает 32 бита, то есть 4 байта. Биты, отведённые под число, делятся на три части:

• знак (sign): первый бит определяет, положительное число или отрицательное;
• порядок (exponent): задаёт масштаб числа;
• мантисса (mantissa): задаёт значащие цифры числа.

Принцип чисел с плавающей точкой легко пояснить через научную запись чисел, например $-5{,}734 \times 10^{7}$: сначала знак, затем мантисса и порядок. Так можно представлять числа очень разного масштаба с адаптивной точностью. Хотя float32 — формат по умолчанию, в PyTorch доступен ряд форматов с плавающей точкой:

Формат	Всего бит	Знак	Порядок	Мантисса
float32	32	1	8	23
float16	16	1	5	10
bfloat16	16	1	8	7
float8 (e4m3)	8	1	4	3
float8 (e5m2)	8	1	5	2

Примечание: возможно, вам интересно, откуда «b» в bfloat16. Формат разработали в Google Brain, так что «b» — от «brain».

Уменьшение общего числа бит имеет цену (и тут без бесплатных обедов), но мы можем выбирать, чем платить: жертвовать битами мантиссы или порядка. Поэтому есть и два формата float8 (FP8), названные по числу бит порядка и мантиссы, чтобы можно было гибко выбрать подходящий. На графике оригинала — диапазоны чисел каждого формата.

Видно, что float32 охватывает 80 порядков величины, float16 сильно жертвует диапазоном, а bfloat16 сохраняет полный диапазон. Два формата float8 сужают диапазон ещё больше: e5m2 сохраняет диапазон float16, а у e4m3 он ещё меньше.

Как одни форматы сохраняют полный диапазон, а другие нет? Исследуем их разрешение: возьмём 10 000 точек между 1 и 2 и округлим каждую до ближайшего представимого числа в каждом формате (график в оригинале).

Видно, что bfloat16 сохраняет диапазон float32 (в отличие от float16), но ценой большей потери точности. Для float8 ситуация ещё печальнее: e4m3 может представить в интервале $[1, 2]$ только 7 чисел, а e5m2 — только 3.

Распространённая мера разрешения формата — эпсилон (epsilon): первое представимое число после 1,00. Для float32 верхняя граница — $10^{-4}$ (на самом деле $1{,}19 \times 10^{-7}$). Для float16 это около $10^{-3}$, а для bfloat16 — ещё в 10 раз больше.

Примечание переводчика

В оригинале опечатка: «1.19^{-7}» вместо $1{,}19 \times 10^{-7}$. Точные значения эпсилон: float32 — $2^{-23} \approx 1{,}19 \cdot 10^{-7}$, float16 — $2^{-10} \approx 9{,}8 \cdot 10^{-4}$, bfloat16 — $2^{-7} \approx 7{,}8 \cdot 10^{-3}$. Проверить можно так: torch.finfo(torch.bfloat16).eps.

Идея обучения со смешанной точностью — использовать эти форматы пониженной точности для части вычислений, сохраняя качество обучения в полной точности. Оказывается, полностью отказаться от float32 нельзя: обычно часть вычислений приходится выполнять в полной точности.

Посмотрим теперь на обучение моделей в 16 битах, а затем — можно ли пойти дальше, вплоть до 8 бит.

Обучение в FP16 и BF16

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

1. Копия весов в FP32. С весами в FP16 возможны две проблемы. Во время обучения некоторые веса могут стать очень маленькими и округлиться до 0. Но даже если сами веса не близки к нулю, при очень маленьких обновлениях разница масштабов может привести к потере значимости (underflow) при сложении. Если вес стал нулём, он останется нулём до конца обучения: сигнал градиента через него больше не проходит.
2. Масштабирование функции потерь (loss scaling). Похожая проблема есть и у градиентов: они обычно намного меньше 1 и рискуют потерять значимость. Простая, но действенная стратегия — умножить функцию потерь на масштаб перед обратным проходом и разделить градиенты на него после. Так при обратном проходе не будет потери значимости, а на обучение масштабирование не влияет, потому что мы снимаем масштаб до дальнейшей обработки градиентов (например, клиппинга) и шага оптимизации.
3. Накопление (accumulation). Наконец, при некоторых арифметических операциях в 16-битной точности (например, средних или суммах) тоже возможны потеря значимости или переполнение. Решение — накапливать промежуточные результаты в FP32 во время операции и лишь итоговый результат переводить обратно в 16 бит.

С этими приёмами обучение стабильно, и при этом мы выигрываем в пропускной способности за счёт более быстрых операций пониженной точности.

Естественно, как любопытный читатель (и уже слегка зависимый от максимизации пропускной способности), вы можете спросить: можно ли пойти дальше и быстрее, чем 16 бит?

Возможно!

Пример переводчика: BF16 в PyTorch

На ваших RTX 4080 (Ada) BF16 поддерживается аппаратно, и масштабирование потерь для него обычно не нужно — достаточно автоматической смешанной точности:

import torch

model = torch.nn.Linear(1024, 1024).cuda()        # веса хранятся в FP32 («мастер-веса»)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
x = torch.randn(8, 1024, device='cuda')

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    y = model(x)                 # умножение матриц выполняется в BF16
    loss = y.float().pow(2).mean()   # функцию потерь считаем в FP32
loss.backward()                  # градиенты весов получаются в FP32
opt.step()

Для FP16 понадобился бы ещё torch.amp.GradScaler — это и есть «масштабирование функции потерь» из пункта 2.

Предобучение в FP8

Даже если идеально перекрыть коммуникацию вычислениями, рано или поздно мы упрёмся в низкоуровневый теоретический предел FLOPS самого железа — эффективность каждой отдельной операции. Здесь решающей становится числовая точность. Например, на H100 умножения матриц в FP8 (операции GEMM) дают вдвое больше теоретических FLOPS, чем в BF16, что делает обучение в пониженной точности привлекательным путём для дальнейшей оптимизации.

Недавние исследования — включая FP8-LM, torchao и DeepSeek-V3 — показали потенциал обучения в FP8 для крупных моделей. Но предобучение в FP8 ставит серьёзную проблему — стабильность. В пониженной точности числовая нестабильность часто приводит к расхождению функции потерь, и трудно сравняться по точности с обучением в более высокой точности.

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

Первое успешное обучение в очень большом масштабе со смешанной точностью FP8 публично описано в техотчёте DeepSeek-V3. Авторы тщательно проанализировали каждую операцию прямого прохода (Fprop), а также обратных проходов по активациям (Dgrad) и по весам (Wgrad). Как и при смешанной точности BF16, часть агрегаций и мастер-веса хранятся в более высокой точности, а сами операции выполняются в FP8.

Чтобы перейти от высокой точности (например, FP32 или BF16) к пониженной (например, FP16 или FP8) с меньшим диапазоном, нужно нормировать диапазон значений активаций, например вычисляя их абсолютный максимум. DeepSeek-V3 ввёл особую схему квантования, где диапазоны нормируются поплиточно: 1×128 для входов и активаций и 128×128 для весов и масштабных множителей. Так на нормировку меньше влияют выбросы в активациях. Авторы предложили и ряд дополнительных приёмов для сокращения памяти и коммуникации — о них можно прочитать в разделе 3.3 техотчёта.

Сводка нескольких известных подходов к обучению в FP8:

	Точность GEMM	Мастер-веса	Накопленные градиенты	Веса модели	Градиенты	Состояния оптимизатора	Всего памяти
Базовый вариант: BF16 со смешанной точностью FP32	BF16	FP32	FP32	BF16	BF16	FP32 + FP32	4 + 4 + 2 + 2 + 4 + 4 = 20 байт
То же без накопления градиентов в FP32	BF16	FP32	—	BF16	BF16	FP32 + FP32	4 + 2 + 2 + 4 + 4 = 16 байт (−20 %)
Transformer Engine	FP8	—	—	FP32	FP32	FP32 + FP32	4 + 4 + 4 + 4 = 16 байт (−20 %)
FP8-LM, уровень O3	FP8	FP16	FP16	FP8	FP8	FP8 + FP16	2 + 2 + 1 + 1 + 1 + 2 = 9 байт (−55 %)
DeepSeek-V3	FP8	FP32	FP32	FP8	BF16	BF16 + BF16	4 + 4 + 1 + 2 + 2 + 2 = 15 байт (−25 %)
FP8 в Nanotron	FP8	BF16	FP32	FP8	FP8	FP8 + FP8	2 + 4 + 1 + 1 + 1 + 1 = 10 байт (−50 %)

В целом FP8 (на начало 2025 года) остаётся экспериментальной техникой, и методы ещё развиваются. Учитывая очевидные преимущества, он, скорее всего, станет стандартом и скоро заменит смешанную точность BF16. Открытую реализацию приёмов обучения в FP8 можно посмотреть в этом PR в Nanotron.

Если заглянуть ещё дальше: для следующего поколения чипов NVIDIA, Blackwell, объявлена поддержка обучения в FP4 — это ещё ускорит обучение, но, несомненно, принесёт и новые проблемы стабильности.

Этот раздел завершает наше долгое путешествие в страну быстрого обучения больших моделей на десятках и тысячах GPU. Пора потихоньку остановить наш кластер GPU и оглянуться на всё, что мы узнали по пути!

Заключение

Поздравляем, дорогой читатель, вы дошли до конца! Путь был немалый: мы начали с того, как обучить простую модель на одной GPU, и дошли до освоения разных хитроумных приёмов, которыми эффективно обучают огромные языковые модели вроде Llama-405B и DeepSeek-V3 на тысячах GPU. Теперь вы можете (относительно) легко прочитать схему вроде 4D-параллелизма Llama-3 (рисунок в оригинале).

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

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

Это был долгий путь обучения — и не только для вас! Прогнать тысячи бенчмарков на кластере GPU оказалось сложнее, чем мы ожидали, и мы хотели поделиться и своим опытом.

Что дальше?

Теперь у вас хороший обзор основных понятий распределённого обучения, но мы лишь слегка коснулись многих инструментов и приёмов. Углубиться в тему можно по-разному; вот шаги, которые мы рекомендуем:

• внимательно прочитать несколько этапных или совсем свежих статей — обширный список самых влиятельных статей, постов и книг, которые мы знаем, есть в разделе References оригинала;
• начать с нуля и реализовать алгоритм самостоятельно — часто метод по-настоящему «щёлкает», только когда сам его реализовал;
• погрузиться в один из широко используемых фреймворков и начать вносить вклад: исправлять ошибки, отвечать на issues или реализовывать новые возможности. Это лучший способ войти в любую область машинного обучения!

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

Благодарности

Мы благодарим Эли за тщательные рецензии и создание аудио с помощью NotebookLM. Особая благодарность Гинеку за оптимизацию производительности фронтенда. Также благодарим Симона за решение нескольких проблем на хабе.

Страница обсуждения

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

Что сделать руками

Это не часть оригинала — короткий план от переводчика для вашей машины (4× RTX 4080).

1. Запустите примеры коллективных операций из приложения A0 через torchrun --nproc_per_node=4 и замерьте время all-reduce тензора на 1 ГБ.
2. Разберите файлы picotron: data_parallel.py, tensor_parallel.py, pipeline_parallel.py — каждый умещается на нескольких экранах.
3. Сравните скорость elu без @torch.compile и с ним на тензоре из $10^8$ элементов и посмотрите сгенерированное ядро Triton через TORCH_LOGS="output_code".

Словарик

EN	RU
streaming multiprocessor (SM)	потоковый мультипроцессор
registers, shared memory, L1/L2 cache, global memory (HBM)	регистры, разделяемая память, кэш L1/L2, глобальная память
kernel, host code	ядро (программа для GPU), код хоста
thread, warp, block	поток, варп (32 потока), блок потоков
memory coalescing	объединение обращений к памяти
tiling	тайлинг (разбиение на плитки)
thread coarsening	укрупнение потоков
control divergence	расхождение потока управления
kernel fusion, fused kernel	слияние ядер, слитое ядро
online softmax	онлайн-softmax (softmax по кускам)
sign, exponent, mantissa	знак, порядок, мантисса
loss scaling	масштабирование функции потерь
underflow / overflow	потеря значимости / переполнение