Meta представляет KernelLLM: 8-миллиардная модель, преобразующая PyTorch модули в эффективные Triton GPU ядра

Meta представила KernelLLM, языковую модель с 8 миллиардами параметров, дообученную на базе Llama 3.1 Instruct, предназначенную для автоматизации преобразования модулей PyTorch в оптимизированные Triton GPU ядра. Эта инновация направлена на упрощение программирования GPU и облегчение процесса разработки ядер.

Технические детали

KernelLLM обучалась на датасете KernelBook, содержащем около 25 000 пар примеров модулей PyTorch и соответствующих им Triton ядер. Датасет сформирован на основе кода из The Stack и дополнен синтетически сгенерированными образцами с помощью torch.compile() и различных методов подсказок.

Обучение проводилось с использованием контролируемой настройки по инструкциям и шаблонам подсказок, включающим примеры формата как на этапе тренировки, так и на оценке. Модель обучалась в течение 10 эпох с размером батча 32, используя 16 GPU примерно 12 часов, что составляет около 192 часов GPU.

Результаты производительности

Эффективность модели оценивалась с помощью KernelBench-Triton — бенчмарка для генерации Triton ядер из модулей PyTorch. KernelLLM достигла Pass@1 результата 20.2, превзойдя более крупные модели, такие как GPT-4o (~200 млрд параметров) и DeepSeek V3 (671 млрд параметров), которые набрали 15 и 16 соответственно.

При нескольких попытках генерации, Pass@10 и Pass@20 модели составили 51.8 и 57.1, что свидетельствует о высокой надежности в создании корректных ядер.

Влияние на программирование GPU

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

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

Модель доступна на Hugging Face. Все заслуги в исследовании принадлежат команде проекта. Следите за обновлениями в Twitter, присоединяйтесь к ML SubReddit с более чем 95 тысячами участников и подписывайтесь на нашу рассылку для получения дополнительных материалов.