RightNow AI запускает AutoKernel для оптимизации GPU-кода

Оптимизация кода для GPU является одной из самых сложных задач в области машинного обучения. Исследователи компании RightNow AI стремятся автоматизировать этот процесс полностью. Они представили AutoKernel — открытый фреймворк, который использует автономный агент на основе языковой модели для оптимизации GPU-ядра для произвольных моделей PyTorch. Суть подхода проста: загрузите модель перед сном, а утром получите более быстрые ядра Triton — без необходимости в экспертизе по GPU.

Ядро GPU — это функция, которая выполняется параллельно на тысячах ядер GPU. При запуске таких трансформерных моделей, как LLaMA или GPT-2, основное время вычислений уходит на операции внутри ядер, такие как матричное умножение, softmax, нормализация слоя и внимание. Эти ядра находятся в библиотеках, таких как cuBLAS и cuDNN, или автоматически генерируются компиляционным конвейером PyTorch. Проблема заключается в том, что для достижения максимальной производительности этих ядер требуется одновременное понимание множества факторов, таких как интенсивность арифметических операций, кэширование памяти и выбор инструкций для тензорных ядер — навыки, которые требуют многих лет практики.

AutoKernel был создан в ответ на недостаток экспертизы в этой области. Основная идея заключается в том, что процесс работы опытного инженера по оптимизации ядер представляет собой простой цикл: написать кандидат, провести бенчмаркинг, сохранить улучшения, отбросить регрессии и повторить. Фреймворк механизирует этот цикл. Агент на основе языковой модели изменяет один файл — kernel.py — фиксированный набор тестов проверяет корректность и измеряет производительность, а результат определяет, сохраняется ли изменение.

Каждый эксперимент сопоставляется с коммитом в системе git. Сохраненные эксперименты продвигают ветку, а отмененные эксперименты удаляются с помощью git reset. Вся история доступна для просмотра с помощью стандартных инструментов git, а результаты экспериментов записываются в простой файл results.tsv, который легко читается и может быть обработан агентом. Каждая итерация занимает примерно 90 секунд: 30 секунд на проверку корректности, 30 секунд на бенчмаркинг производительности и 30 секунд на рассуждения агента и модификацию кода.

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