Оставить заявку

Перевод статьи от OpenAI про Trition — упрощённый язык программирования на GPU для нейросетей

Перевод большой статьи про новую разработку от OpenAI

Представляем Triton: язык программирования на GPU для нейросетей с открытым исходным кодом 

Мы выпускаем Triton 1.0. Это язык программирования с открытым исходным кодом, похожий на Python. Он позволяет программировать графические процессоры без использования CUDA, который гораздо сложнее для понимания.

Triton позволяет достичь максимальной производительности оборудования с относительно небольшими усилиями. Например, его можно использовать для написания ядер матричного умножения FP16, которые соответствуют производительности cuBLAS — чего не могут сделать многие программисты на GPU менее чем в 25 строках кода.

Наши исследователи уже использовали его для создания ядер, которые в 2 раза эффективнее эквивалентных реализаций Torch.

Мы рады работать с сообществом, чтобы сделать программирование на GPU более доступным для всех! 

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

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

  • либо слишком многословны;
  • либо недостаточно гибкие;
  • либо генерируют код заметно медленнее, чем наши настроенные вручную базовые программы.

Это побудило нас расширить и улучшить Triton, создатель которого сейчас работает на OpenAI.

Проблемы программирования на GPU

Архитектуру современных графических процессоров можно условно разделить на три основных компонента — DRAM, SRAM и ALU. При написании кода на CUDA необходимо заниматься ручной оптимизацией каждого из них:

  • Передача памяти из DRAM должна быть объединена в крупные транзакции, чтобы использовать большую ширину шины современных интерфейсов памяти.
  • Данные должны быть вручную сохранены в SRAM перед повторным использованием и управляться таким образом, чтобы свести к минимуму конфликты банка общей памяти при извлечении.
  • Вычисления должны быть тщательно разделены и запланированы как между потоковыми мультипроцессорами (SM), так и внутри них. Это способствует параллелизму на уровне команд/потоков и позволяет использовать специальные ALU, например, тензорные ядра.
Базовая архитектура графического процессора


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

Triton стремится к широкому применению и поэтому не планирует автоматически работу между SM, оставляя некоторые важные алгоритмические соображения (например: мозаичное размещение, синхронизацию между SM) на усмотрение разработчиков.


CUDATriton
Объединение памятиВручнуюАвтоматически
Управление разделяемой памятьюВручнуюАвтоматически
Планирование (внутри SM)ВручнуюАвтоматически
Планирование (между SM)ВручнуюВручную

Оптимизация компилятора в CUDA vs Triton

Модель программирования

Из всех доступных доменных языков и JIT-компиляторов, Triton, больше всего похож на Numba: ядра определяются как декорированные функции Python и запускаются одновременно с разными program_id в сетке так называемых инстансов.

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

При этом Triton эффективно абстрагируется от всех проблем, связанных с параллелизмом в блоках потоков CUDA (например, объединение памяти, синхронизация/конфликты общей памяти, планирование ядра тензора).

BLOCK = 512

# Это ядро GPU в Triton.
# Различные инстансы этой функции
# могут работать параллельно.
@jit
def add(X, Y, Z, N):
# В Triton каждый инстанс ядра
# выполняет блочные операции в
# одном потоке: здесь нет конструкции
# аналогичной threadIdx
pid = program_id(0)
# блок индексов
idx = pid * BLOCK + arange(BLOCK)
mask = idx < N
# Triton использует указательную 
# арифметику вместо операторов
# индексации
x = load(X + idx, mask=mask)
y = load(Y + idx, mask=mask)
store(Z + idx, x + y, mask=mask)

grid = (ceil_div(N, BLOCK),)
# нет блока потока
add[grid](x, y, z, x.shape[0]
BLOCK = 512

# Это ядро GPU в Numba.
# Различные инстансы этой функции
# могут работать параллельно.

@jit
def add(X, Y, Z, N):
# В Numba/CUDA, каждый инстанс
# ядра сам использует модель 
# выполнения SIMT, где инструкции
# выполняются параллельно для
# разных значений threadIdx

tid = threadIdx.x
bid = blockIdx.x
# скалярный индекс
idx = bid * BLOCK + tid
if id < N:
# В Numba нет указателя.
# Z,X,Y — плотные тензоры

Z[idx] = X[idx] + Y[idx]

grid = (ceil_div(N, BLOCK),)
block = (BLOCK,)
add[grid, block](x, y, z, x.shape[0])

Сложение векторов в Triton

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

Рассмотрим, например, случай объединенного ядра softmax (ниже), в котором каждый экземпляр нормализует другую строку заданного входного тензора X∈R^(M×N). Стандартные реализации этой стратегии распараллеливания в CUDA могут быть сложными для написания, требуя явной синхронизации между потоками. Поскольку они одновременно сокращают одну и ту же строку X. С Triton это реализуется гораздо проще. С помощью него каждый экземпляр ядра загружает интересующую строку и последовательно нормализует ее с помощью примитивов, подобных NumPy.


import triton
import triton.language as tl
@triton.jit def softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N):
# индекс строки
m = tl.program_id(0)
# индексы столбцов
# это конкретное ядро работает только для матриц,
# которые имеют меньше BLOCK_SIZE столбцов

BLOCK_SIZE = 1024
n = tl.arange(0, BLOCK_SIZE)
# адреса всех элементов,
# которые мы хотим загрузить можно вычислить следующим образом

X = X + m * stride_xm + n * stride_xn
# загружаем входные данные; заполняем out-of-bounds элементы нулями
x = tl.load(X, mask=n < N, other=-float('inf'))
# вычисляем численно стабильный softmax
z = x — tl.max(x, axis=0)
num = tl.exp(z)
denom = tl.sum(num, axis=0)
y = num / denom
# записываем обратно в Y
Y = Y + m * stride_ym + n * stride_yn
tl.store(Y, y, mask=n < N)
import torch
# Распределение тензоров ввода-вывода
X = torch.normal(0, 1, size=(583, 931), device=’cuda’)
Y = torch.empty_like(X)
# Сетка запуска SPMD
grid = (X.shape[0], )
# выделение ядра GPU
softmax[grid](Y, Y.stride(0), Y.stride(1),
X, X.stride(0), X.stride(1),
X.shape[0] , X.shape[1])

Fused softmax в Triton

Обратите внимание, что Triton JIT рассматривает X и Y как указатели, а не как тензоры. Нам показалось, что сохранение низкоуровневого контроля доступа к памяти важно для решения более сложных структур данных (например, блочно-разрозненных тензоров).

Данная реализация softmax сохраняет строки X в SRAM на протяжении всего процесса нормализации, что позволяет максимально использовать данные повторно, когда это возможно (~<32K столбцов). Это отличается от внутреннего CUDA-кода PyTorch, чье использование временной памяти делает его более общим, но значительно медленным (см. ниже).

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

Производительность A100 для fused softmax при M = 4096


Более низкая производительность JIT Torch (v1.9) подчеркивает сложность автоматической генерации кода CUDA из последовательностей высокоуровневых тензорных операций.

@torch.jit.script
def softmax(x):
x_max = x.max(dim=1)[0]
z = x — x_max[:, None]
numerator = torch.exp(x)
denominator = numerator.sum(dim=1)
return numerator / denominator[:, None]

Fused softmax с помощью Torch JIT

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

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

Triton отлично справляется и с ними, достигая пиковой производительности всего с ~25 строками кода на Python. Реализация чего-то подобного на CUDA потребовала бы гораздо больше усилий и даже, скорее всего, привела бы к снижению производительности.

@triton.jit
def matmul(A, B, C, M, N, K, stride_am, stride_ak,
stride_bk, stride_bn, stride_cm, stride_cn, **META):
# извлекаем метапараметры
BLOCK_M, GROUP_M = META[‘BLOCK_M’], META[‘GROUP_M’]
BLOCK_N = META[‘BLOCK_N’]
BLOCK_K = META[‘BLOCK_K’]
# программы сгруппированы вместе, чтобы улучшить процент попаданий L2
_pid_m = tl.program_id(0)
_pid_n = tl.program_id(1)
pid_m = _pid_m // GROUP_M
pid_n = (_pid_n * GROUP_M) + (_pid_m % GROUP_M)
# rm (или rn) обозначает диапазон индексов # для строк (или столбцов) C
rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
# rk обозначает диапазон индексов для столбцов
# (или строк) A (или B)

rk = tl.arange(0, BLOCK_K)
# адреса памяти элементов в первом блоке
# A и B могут быть вычислены с помощью nampy-style трансляции
A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak)
B = B + (rk [:, None] * stride_bk + rn[None, :] * stride_bn)
# инициализация и итеративное обновление аккумулятора
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(K, 0, -BLOCK_K):
a = tl.load(A)
b = tl.load(B)
# умножение матрицы на уровне блока
acc += tl.dot(a, b)
# инкрементировать указатели так, чтобы следующие блоки A и B
# были загружены во время следующей итерации

A += BLOCK_K * stride_ak
B += BLOCK_K * stride_bk
# при желании можно предохранить LReLU
# acc = tl.where(acc >= 0, acc, alpha * acc)
# запишите результат

C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
mask = (rm[:, None] < M) & (rn[None, :] < N)
tl.store(C, acc, mask=mask)

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

Одним из важных преимуществ рукописных ядер матричного умножения является то, что они могут быть настроены по желанию для размещения объединённых преобразований их входов (например, нарезки) и выходов (например, Leaky ReLU).

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

Производительность тензорного ядра V100 при умножении матриц с соответствующими значениями для BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M.


Высокоуровневая архитектура системы

Хорошая производительность Triton обеспечивается модульной архитектурой системы, основанной на Triton-IR, промежуточном представлении на основе LLVM, в котором многомерные блоки значений являются объектами первого класса. 

Высокоуровневая архитектура Triton


Декоратор @triton.jit работает путем обхода абстрактного синтаксического дерева (AST) предоставленной Python-функции, чтобы на лету генерировать Triton-IR с использованием общего алгоритма построения SSA.

Полученный IR-код затем упрощается, оптимизируется и автоматически распараллеливается серверной частью нашего компилятора перед преобразованием в высококачественный LLVM-IR — и в PTX — для выполнения на последних графических процессорах NVIDIA.

В настоящее время CPU и AMD GPU не поддерживаются, но мы приветствуем вклад сообщества, направленный на устранение этого ограничения.

Бэкэнд компилятора

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

Например, данные могут быть автоматически убраны в разделяемую память путем просмотра операндов вычислительно интенсивных операций блочного уровня (например, tl.dot), а также распределены/синхронизированы с помощью стандартных методов анализа жизнеспособности.

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


С другой стороны, программы Triton могут быть эффективно и автоматически распараллелены как (1) между SM, путем одновременного выполнения различных экземпляров ядра, так и (2) внутри SM, путем анализа пространства итераций каждой операции на уровне блоков и адекватного разделения его между различными SIMD-блоками, как показано ниже.

Автоматическое распараллеливание в Triton. Каждая операция на уровне блоков определяет блокированное пространство итераций, которое автоматически распараллеливается для использования ресурсов, доступных на потоковом мультипроцессоре (SM).
  1. Определение Triton-программы P, состоящей из трех утверждений S1, S2, S3.
  2. Итерационное пространство S3
  3. Отображение S3 на потоковый мультипроцессор (SM)
  4. Отображение P на GPU

Заключение

Ознакомиться с Triton можно по ссылке на официальный репозиторий: https://github.com/openai/triton

Источник: https://openai.com/blog/triton/