Triton GPU: OpenAI дала писать ядра без CUDA, Unsloth ускорил обучение LLM в 30 раз

Triton GPU-программирование появилось как ответ на конкретную проблему: ядра CUDA сложно писать и отлаживать, а потребность в оптимизации GPU-вычислений растёт вместе с размером моделей. Компания OpenAI создала Triton как компилятор, который абстрагирует сложность CUDA и открывает разработчикам без глубокого опыта в аппаратном программировании возможность писать производительные ядра. Весь код из этого разбора доступен в репозитории на GitHub: https://github.com/RPegoud/Triton-Kernels.

Как устроен GPU изнутри: аналогия с заводом

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

• Поток — мельчайшая рабочая единица. Каждый поток выполняет код ядра, заданный программистом.
• Варп (warp) — бригада из 32 потоков, которые стартуют одновременно, но могут ветвиться и завершаться независимо. Это наименьшая единица планирования.
• Блок потоков — группа варпов, работающих с общей разделяемой памятью. Блоки выполняются независимо друг от друга и в любом порядке. Именно поэтому GPU-программы хорошо масштабируются при увеличении числа ядер.
• Потоковый мультипроцессор (SM) — аппаратный блок, в котором параллельно работают несколько варпов. Он владеет разделяемой памятью, кэшем L1 и собственным планировщиком варпов.
• Ядро CUDA — физическое арифметическое логическое устройство (АЛУ), которое выполняет арифметические операции для потока.

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

Три узких места, которые надо знать до оптимизации

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

• Вычисления (FLOPS) — время, которое GPU тратит на арифметику с числами с плавающей точкой. GPU заточены под перемножение матриц, и в идеале именно туда должна уходить основная доля времени.
• Память (пропускная способность) — время на перемещение тензоров (многомерных массивов данных) внутри GPU: из глобальной памяти DRAM (дешёвой, но медленной) в разделяемую SRAM (дорогую, но быструю).
• Накладные расходы — всё остальное: интерпретатор Python, диспетчеризация PyTorch и прочее.

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

Две стратегии снижения затрат на пропускную способность

Переиспользование данных в разделяемой памяти. Классический пример — перемножение матриц с разбиением на блоки (tiled matrix multiplication). Данные загружаются в быструю SRAM один раз и используются повторно на нескольких этапах вычислений.

Слияние операций в одном ядре (kernel fusion). Каждый запуск ядра требует переноса данных из DRAM в SRAM. Если объединить, например, умножение матриц и функцию активации в одно ядро, глобальных операций чтения и записи становится меньше, и производительность вырастает. Любые два оператора потенциально подходят для слияния.

Что понадобится

• Python 3.8 или новее
• PyTorch (актуальная версия с поддержкой CUDA)
• Библиотека Triton (pip install triton)
• GPU Nvidia с поддержкой CUDA (от архитектуры Volta и выше)
• Репозиторий с примерами: https://github.com/RPegoud/Triton-Kernels
• Время: от 30 до 60 минут на первое рабочее ядро

Пошаговая инструкция: первое ядро на Triton

1. Установите Triton и убедитесь, что GPU доступен.

pip install triton
python -c "import torch; print(torch.cuda.is_available())"

Если вывод True, среда готова.

1. Клонируйте репозиторий с примерами.

git clone https://github.com/RPegoud/Triton-Kernels.git
cd Triton-Kernels

1. Изучите структуру Triton-ядра. Ядро Triton — это функция Python с декоратором @triton.jit, внутри которой вы работаете не с отдельными потоками (как в CUDA), а с блоками данных. Triton сам решает, как разложить блоки по потокам и варпам.

2. Напишите простое ядро сложения двух векторов. Это стандартная отправная точка.

import triton
import triton.language as tl
import torch

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + 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)
    tl.store(output_ptr + offsets, x + y, mask=mask)

1. Вызовите ядро из обычного кода PyTorch.

def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    return output

x = torch.rand(1000000, device='cuda')
y = torch.rand(1000000, device='cuda')
result = add(x, y)

1. Проверьте результат, сравнив с PyTorch.

assert torch.allclose(result, x + y), "Результаты не совпадают"
print("Ядро работает корректно")

1. Переходите к более сложным ядрам. В репозитории есть примеры слияния операций (умножение матриц плюс активация) и работы с разделяемой памятью. Именно на этом этапе Triton GPU-оптимизация начинает давать реальный выигрыш: вы заменяете стандартные слои PyTorch на собственные ядра, как это делает Unsloth.

Что делать с этим прямо сейчас, по ролям

Разработчику, работающему с LLM. Начните с замены одного узкого слоя в вашей модели PyTorch на Triton-ядро. Функция активации или нормализация — хорошие кандидаты для слияния операций. Проект Unsloth показал, что именно такая замена даёт кратное ускорение обучения.

Автору Дзена, пишущему про технологии. Тема Triton GPU-оптимизации сейчас на подъёме: разработчики ищут способы снизить стоимость обучения моделей, и пошаговые разборы с кодом собирают техническую аудиторию. Разберите конкретный кейс из репозитория и покажите замер «до и после».

Предпринимателю и тимлиду в РФ. Triton работает на любом GPU Nvidia с поддержкой CUDA, локализация не нужна. Если ваша команда тратит деньги на облачные GPU для обучения или инференса моделей, даже небольшая оптимизация ядер может ощутимо сократить счёт. Из доступных в РФ аналогов для ускорения обучения LLM стоит смотреть на Unsloth (открытый проект, работает поверх Triton).

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