Изучаем Triton ядро за ядром: сложение векторов

Основы программирования для GPU, оптимизация и ваше первое ядро на Triton

Весь код к этому посту выложен на Github по адресу https://github.com/RPegoud/Triton-Kernels ^[4].

Основы архитектуры GPU

В этом разделе будет в самом общем виде разобрано устройство GPU (от Nvidia), чтобы, дочитав эту статью, вы смогли приступить к написанию вашего первого ядра на Triton.

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

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

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

• Блок потоков: группа варпов, в рамках которой все потоки могут кооперироваться, пользуясь разделяемой памятью и барьерами синхронизации. Необходимо, чтобы блоки потоков могли выполняться независимо друг от друга и в любом порядке, параллельно или последовательно. Благодаря такой независимости, можно распланировать блоки потоков в любом порядке на любом количестве вычислительных ядер. Поэтому программы для GPU эффективно масштабируются с увеличением числа ядер процессора. Можно синхронизировать объединённые в блок потоки в конкретных точках работы ядра — например, чтобы синхронизировать доступ к памяти.

• Потоковый мультипроцессор (SM): блок, отвечающий за параллельное выполнение множества варпов. Он владеет разделяемой памятью и кэшем L1 (то есть, содержит самые свежие кэш-линии глобальной памяти, к которым недавно обращался потоковый мультипроцессор). В потоковом мультипроцессоре есть выделенный планировщик варпов, который подтягивает варпы из тех блоков потоков, которые готовы к запуску.

С аппаратной стороны мельчайшая единица выполнения — это ядро CUDA, физическое арифметическое логическое устройство (АЛУ), выполняющее арифметические операции для потока (или его частей).

Чтобы резюмировать этот раздел при помощи аналогии, можно сказать, что ядра CUDA подобны отдельным работникам, тогда как варп — это бригада из 32 работников, которым одновременно поручена одна и та же задача. Они могут выполнять эту задачу как одинаково, так и по-разному (ветвление) и потенциально могут справиться со своей частью задачи в разное время (независимость). Блок потоков состоит из нескольких бригад, работающих в одном цеху (т.e., они пользуются общей памятью), причём, рабочие из всех бригад в данном цеху могут дожидаться, пока другие справятся со своими делами, чтобы всем вместе пойти на обед. Потоковый мультипроцессор — это заводской этаж, в котором работают многие бригады, совместно пользующиеся как инструментами, так и складскими помещениями. Наконец, GPU — это целый многоэтажный завод.

Основы оптимизации

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

• Вычислительные ресурсы: время, которое GPU тратит на операции над числами с плавающей точкой (FLOPS).

• Память: время, затрачиваемое на перенос тензоров внутри GPU.

• Накладные расходы: все прочие операции (интерпретатор Python, диспетчеризация PyTorch, …).

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

Таким образом, нужно минимизировать издержки, связанные с перемещением данных, либо от ЦП к графическому процессору (”стоимость переноса данных”), либо от узла к узлу (”стоимость работы с сетью”), либо из глобальной памяти CUDA (DRAM, дешёвая, но медленная) в совместно используемую память CUDA (SRAM, дорогая, но зато самая быстрая память на устройстве). Последние издержки называются расходами на пропускную способность, и здесь мы уделим им особое внимание ^[5]. Среди обычных стратегий снижения расходов на пропускную способность есть такие:

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

• Слияние множественных операций в одном ядре (поскольку при каждом пуске ядра приходится перемещать данные из  DRAM в SRAM). Например, можно сливать перемножение матриц с функцией активации. Как правило, слияние операторов позволяет сильно нарастить производительность, так как исключает множество глобальных операций чтения и записи в память. При этом, любые два оператора потенциально подходят для слияния.

В этом примере перемножаются матрицы x@W, и их произведение сохраняется в промежуточной переменной a. После этого к a применяется relu, и результат этой операции сохраняется в переменной y. Для этого GPU должен прочитать информацию из x и W, находящихся в глобальной памяти, записать результат в a, снова прочитать из a и, наконец, записать результат в y. Напротив, при слиянии операторов мы могли бы при перемножении матриц вдвое сократить количество операций чтения и записи в глобальную память, так как применили бы ReLU в единственном ядре.

Triton

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

Например, мы хотим просуммировать записи из двух векторов X и Y, каждый из которых содержит по 7 элементов (n_elements=7).

Поручим GPU решать эту задачу фрагмент за фрагментом, по 3 элемента за один раз  (BLOCK_SIZE=3). Следовательно, чтобы охватить все 7 элементов входных векторов, GPU запустит 3 параллельные «программы», каждая из которых представляет собой самостоятельный экземпляр ядра, каждый из них — с собственным уникальным программным ID, pid:

• Программе 0 присваиваются элементы 0, 1, 2.

• Программе 1 присваиваются элементы 3, 4, 5.

• Программе 2 присваивается элемент 6.

Затем эти программы запишут полученные результаты обратно в вектор Z, сохранённый в глобальной памяти.

Здесь есть важная деталь: ядро получает не целый вектор X, а только указатель на адрес в памяти, по которому расположен первый элемент X[0]. Чтобы обратиться к конкретным значениям X, необходимо вручную загрузить их из глобальной памяти.

Можно обратиться к данным каждого блока по ID программы: block_start = pid * BLOCK_SIZE. Начиная отсюда, можно получить адреса оставшихся элементов из данного блока, вычислив offsets = block_start + range(0, BLOCK_SIZE) и загрузив их в память.

Но не забывайте, что программе 2 присвоен только элемент 6, а её смещения равны [6, 7, 8]. Чтобы избежать каких-либо ошибок с индексированием, Triton позволяет определить маску для идентификации валидных целевых элементов. Здесь mask = offsets < n_elements.

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

Давайте подробнее рассмотрим код. Вот ядро Triton:

import triton
import triton.language as tl

@triton.jit
def add_kernel(
	x_ptr, # указатель на первую запись x, расположенную в памяти
	y_ptr, # указатель на первую запись y, расположенную в памяти
	output_ptr, # указатель на первую запись вывода, расположенную в памяти 
	n_elements, # размерность x и y
	BLOCK_SIZE: tl.constexpr, # размер отдельно взятого блока
):
	# --- вычисляем смещения и маску ---
	pid = tl.program_id(axis=0) # индекс блока
	block_start = pid * BLOCK_SIZE # начальный индекс для текущего блока
	offsets = block_start + tl.arange(0, BLOCK_SIZE) # диапазон индекса
	mask = offsets < n_elements # маска элементов, выходящих за границы
	
	# --- Загружаем переменные из глобальной памяти  ---
	x = tl.load(x_ptr + offsets, mask=mask)
	y = tl.load(y_ptr + offsets, mask=mask)

	# --- Операция ---
	output = x + y	
	
	# --- Сохранение результатов в глобальной памяти ---
	tl.store(pointer=output_ptr + offsets, value=output, mask=mask)

Давайте разберём некоторые детали Triton-специфичного синтаксиса:

• Во-первых, ядро Triton всегда декорируется <a href=”http://twitter.com/triton ^[6]” target=”_blank” rel=”noreferrer noopener”>@triton</a>.jit.

• Во-вторых, некоторые аргументы требуется определить как статические — то есть, они известны во время вычислений. Это необходимо для BLOCK_SIZE и достигается путём добавления аннотации типа tl.constexpr. Также обратите внимание, что мы не аннотируем другие переменные, поскольку они не являются полноценными переменными Python.

• При помощи tl.program_id мы получаем доступ к ID актуального блока. tl.arange ведёт себя примерно как np.arange из NumPy.

• Загрузка и сохранение переменных выполняются путём вызова tl.load и tl.store ^[7] с массивами указателей. Обратите внимание: здесь нет оператора return, его роль делегируется tl.store ^[7].

Теперь, чтобы использовать наше ядро, нужно написать для него обёртку на уровне PyTorch, в которой будут предоставляться указатели на память и определяться сетка ядра. Как правило, сетка ядра — это 1D-, 2D- или 3D-кортеж, в котором указано, сколько блоков потоков выделено ядру вдоль каждой из осей. В нашем предыдущем примере мы использовали 1D-грид из 3 блоков потоков: grid = (3, ).

Для работы с массивами переменного размера по умолчанию определяем, что grid = (ceil(n_elements / BLOCK_SIZE), ).

def add(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
	"""PyTorch wrapper for `add_kernel`."""
	output = torch.zeros_like(x) # выделяем память для выходных значений
	n_elements = output.numel()  # измерения осей X и Y
	
	# cdiv = ceil div, вычисляет, сколько блоков будет использоваться
	grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
	# при вызове ядра мы будем автоматически сохранять `BLOCK_SIZE` in `meta`
	# и обновлять `output`
	add_kernel[grid](X, Y, output, n_elements, BLOCK_SIZE=1024)
	
	return output

Вот два последних замечания об обёртке:

Возможно, вы заметили, что grid определяется как лямбда-функция. Благодаря этому Triton может вычислить, сколько блоков потоков должно стартовать во время запуска. Следовательно, мы вычисляем размер сетки, основываясь на том размере блока, который сохранён в meta, словаре предоставляемых ядру констант времени компиляции.

При вызове ядра значение output будем изменять прямо на месте, поэтому нам и не требуется заново присваивать output = add_kernel[…].

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

x, y = torch.randn((2, 2048), device="cuda")

print(add(x, y))
>> tensor([ 1.8022, 0.6780, 2.8261, ..., 1.5445, 0.2563, -0.1846], device='cuda:0')

abs_difference = torch.abs((x + y) - add(x, y))
print(f"Max absolute difference: {torch.max(abs_difference)}")
>> Max absolute difference: 0.0

Ссылки на полезные ресурсы

• Сколько стоит обучение GPT4 ^[8]

• Ядра Unsloth ^[9]

• Руководство по Triton: сложение векторов ^[10]

• Что делать, чтобы глубокое обучение работало — разбор с азов ^[11]