Основы программирования GPU, оптимизация и ваше первое ядро Triton
Делиться
В эпоху моделей с миллиардами параметров небольшая оптимизация может многое значить. Обучение таких моделей, как GPT4 , обходится более чем в 100 миллионов долларов , что даёт прирост эффективности на 1% более чем в миллион долларов . Эффективный способ оптимизировать эффективность моделей машинного обучения — это писать некоторые их компоненты непосредственно на графическом процессоре . Если вы хоть немного похожи на меня, то одного упоминания ядер CUDA достаточно, чтобы вас бросило в дрожь, поскольку они, как известно, сложны в написании и отладке.
К счастью, в 2021 году компания OpenAI выпустила Triton — новый язык программирования и компилятор, абстрагирующие значительную часть сложности CUDA и позволяющие менее опытным специалистам писать производительные ядра. Ярким примером является Unsloth , сервис обучения LLM, обещающий в 30 раз более быстрое обучение при 60% меньшем использовании памяти благодаря замене слоёв, написанных в PyTorch, ядрами Triton .
В этой серии уроков мы изучим основы архитектуры графических процессоров и способы реализации высокопроизводительных ядер Triton! Весь код, представленный в этой серии, будет доступен по адресу https://github.com/RPegoud/Triton-Kernels.
Основы архитектуры графического процессора
В этом разделе мы рассмотрим самые основы графических процессоров (Nvidia), которые помогут нам приступить к работе и написать наше первое ядро Triton к концу статьи.
Начиная с наименьшего программного блока, иерархию блоков исполнения можно описать следующим образом:
- Потоки : наименьшая единица работы , они выполняют определенный пользователем код ядра.
- Варпы : наименьшая единица планирования , всегда состоящая из 32 параллельных потоков, каждый из которых имеет собственный счётчик адресов инструкций и состояние регистра. Потоки в варпе начинаются вместе, но могут свободно разветвляться и выполняться независимо .
- Блоки потоков : группа варпов, в которой все потоки могут взаимодействовать через общую память и барьеры синхронизации. Требуется, чтобы блоки потоков могли выполняться независимо и в любом порядке, параллельно или последовательно. Эта независимость позволяет планировать блоки потоков в любом порядке на любом количестве ядер , чтобы программы на графических процессорах эффективно масштабировались в зависимости от количества ядер. При необходимости мы можем синхронизировать потоки внутри блока в определённых точках ядра, например, для синхронизации доступа к памяти.
- Потоковый мультипроцессор (SM) : устройство, отвечающее за параллельное выполнение множества варпов . У него есть общая память и кэш L1 (хранит последние строки глобальной памяти, к которым SM обращался). SM имеет специальный планировщик варпов , который извлекает варпы из блоков потоков, готовых к выполнению.
С точки зрения аппаратного обеспечения наименьшей единицей работы является ядро CUDA , физическое арифметико-логическое устройство (АЛУ), которое выполняет арифметические операции для потока (или его части).
Подводя итог этому разделу с помощью аналогии, мы могли бы рассматривать ядра CUDA как отдельных рабочих , в то время как варп — это отряд из 32 рабочих, которым одновременно дана одна и та же инструкция. Они могут выполнять эту задачу или нет одинаково (ветвление) и потенциально могут завершить ее в разный момент времени (независимость). Блок потока состоит из нескольких отрядов, разделяющих общее рабочее пространство (т. е. имеют общую память), работники из всех отрядов в рабочем пространстве могут ждать друг друга, чтобы пообедать в одно и то же время. Потоковый мультипроцессор — это заводской цех со множеством отрядов, работающих вместе и совместно использующих инструменты и хранилище . Наконец, графический процессор — это целый завод со многими этажами.
Основы оптимизации
При оптимизации моделей глубокого обучения мы имеем дело с тремя основными компонентами:
- Вычисления : время, затрачиваемое графическим процессором на вычисления операций с плавающей запятой (FLOPS).
- Память : время, затрачиваемое на передачу тензоров внутри графического процессора.
- Накладные расходы : все остальные операции (интерпретатор Python, диспетчеризация PyTorch, …).
Учёт этих компонентов помогает найти правильный способ устранения узкого места. Например, увеличение вычислительной мощности (например, использование более мощного графического процессора) не поможет, если большая часть времени тратится на передачу данных в память. Однако в идеале большую часть времени следует тратить на вычисления, а точнее, на умножение матриц — именно для этой операции оптимизированы графические процессоры.
Это подразумевает минимизацию затрат на передачу данных, будь то от центрального процессора к графическому процессору (« стоимость передачи данных »), от одного узла к другому (« стоимость сети ») или из глобальной памяти CUDA ( DRAM , дешёвая, но медленная) в общую память CUDA ( SRAM , дорогая, но самая быстрая память на устройстве). Последнее называется затратами на пропускную способность и будет в центре нашего внимания в данный момент. Распространенные стратегии снижения затрат на пропускную способность включают:
- Повторное использование данных, загруженных в общую память, для нескольких шагов. Ярким примером этого является умножение матриц, о котором мы поговорим в следующей публикации.
- Объединяя несколько операций в одном ядре (поскольку каждый запуск ядра подразумевает перемещение данных из DRAM в SRAM), например, можно объединить умножение матриц с функцией активации. Как правило, объединение операторов может обеспечить значительный прирост производительности, поскольку предотвращает множество глобальных операций чтения/записи памяти, а любые два оператора предоставляют возможность для объединения.
В этом примере мы выполняем умножение матриц x на W и сохраняем результат в промежуточной переменной a. Затем мы применяем relu к a и сохраняем результат в переменной y. Для этого графический процессор должен прочитать данные из x и W в глобальной памяти, записать результат в a, снова прочитать данные из a и, наконец, записать данные в y. Вместо этого, слияние операторов позволило бы нам вдвое сократить количество операций чтения и записи в глобальной памяти, выполняя умножение матриц и применяя ReLU в одном ядре.
Тритон
Теперь мы напишем наше первое ядро Triton — простое сложение векторов. Сначала давайте разберём, как эта операция структурируется и выполняется на GPU.
Предположим, что требуется суммировать элементы двух векторов X и Y, каждый из которых содержит 7 элементов (n_elements=7).
Мы поручим графическому процессору обрабатывать эту задачу блоками по 3 элемента за раз (BLOCK_SIZE=3). Таким образом, чтобы охватить все 7 элементов входных векторов, графический процессор запустит 3 параллельные «программы», независимые экземпляры нашего ядра, каждая из которых будет иметь уникальный идентификатор программы, pid:
- Программе 0 присвоены элементы 0, 1, 2.
- Программе 1 присвоены элементы 3, 4, 5.
- Программе 2 присвоен элемент 6.
Затем эти программы запишут результаты в вектор Z, хранящийся в глобальной памяти.
Важная деталь заключается в том, что ядро получает не весь вектор X, а указатель на адрес памяти первого элемента X[0]. Чтобы получить доступ к фактическим значениям X, нам необходимо вручную загрузить их из глобальной памяти.
Доступ к данным каждого блока можно получить, используя идентификатор программы: block_start = pid * BLOCK_SIZE. Отсюда можно получить адреса оставшихся элементов блока, вычислив offsets = block_start + range(0, BLOCK_SIZE) и загрузив их в память.
Однако помните, что программе 2 назначен только элемент 6, но его смещения равны [6, 7, 8]. Чтобы избежать ошибок индексации, Triton позволяет определить маску для определения допустимых целевых элементов, здесь маска = смещения < n_элементов.
Теперь мы можем безопасно загрузить X и Y и сложить их, прежде чем записать результат обратно в выходную переменную Z в глобальной памяти аналогичным образом.
Давайте подробнее рассмотрим код. Вот ядро Triton:
import triton import triton.language as tl @triton.jit def add_kernel( x_ptr, # указатель на первую запись в памяти x y_ptr, # указатель на первую запись в памяти y output_ptr, # указатель на первую запись в памяти output 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 всегда декорируется @triton.jit.
- Во-вторых, некоторые аргументы необходимо объявить как статические, то есть они должны быть известны во время вычислений. Это требуется для BLOCK_SIZE и достигается добавлением аннотации типа tl.constexpr. Также обратите внимание, что мы не аннотируем другие переменные, поскольку они не являются полноценными переменными Python.
- Мы используем tl.program_id для доступа к идентификатору текущего блока, tl.arange ведет себя аналогично np.arange в Numpy.
- Загрузка и сохранение переменных осуществляется вызовом tl.load и tl.store с массивами указателей. Обратите внимание, что оператор return отсутствует, эта роль делегирована tl.store.
Чтобы использовать наше ядро, нам нужно написать обёртку уровня PyTorch , которая предоставляет указатели на память и определяет сетку ядра . Как правило, сетка ядра представляет собой одномерный, двумерный или трёхмерный кортеж, содержащий количество блоков потоков, выделенных ядру по каждой оси . В нашем предыдущем примере мы использовали одномерную сетку из трёх блоков потоков: grid = (3, ).
Для обработки массивов разных размеров по умолчанию используется grid = (ceil(n_elements / BLOCK_SIZE), ).
def add(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor: «»»PyTorch-обертка для `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` в `meta` # и обновит `output` add_kernel[grid](X, Y, output, n_elements, BLOCK_SIZE=1024) return output
Вот два последних замечания по поводу обертки:
Вы, возможно, заметили, что сетка определена как лямбда-функция. Это позволяет Triton вычислять количество блоков потоков, запускаемых при запуске . Таким образом, размер сетки вычисляется на основе размера блока, хранящегося в meta — словаре констант времени компиляции, доступных ядру.
При вызове ядра значение output будет изменено на месте, поэтому нам не нужно переназначать output = add_kernel[…].
Завершим это руководство, проверив, что наше ядро работает правильно:
x, y = torch.randn((2, 2048), device=»cuda») print(add(x, y)) >> тензор([ 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»Максимальная абсолютная разность: {torch.max(abs_difference)}») >> Максимальная абсолютная разность: 0.0
На этом введение закончено, в следующих публикациях мы научимся реализовывать более интересные ядра, такие как плиточное умножение матриц, и увидим, как интегрировать ядра Triton в модели PyTorch с помощью Autograd.
До следующего раза! 👋
Ссылки и полезные ресурсы
- Стоимость обучения GPT4
- Ядра Unsloth
- Учебник Triton: Сложение векторов
- Заставляем глубокое обучение работать с первых принципов
Источник: towardsdatascience.com
