OpenAI представил новый язык Triton для GPU-вычислений

Компания OpenAI выпустила Triton, Python-подобный язык программирования с открытым исходным кодом, который позволяет исследователям писать высокоэффективный GPU код для рабочих задач, связанных с ИИ. OpenAI утверждает, что Triton позволяет достичь максимальной производительности оборудования с относительно небольшими усилиями, создавая код, равный тому, что мог бы сделать эксперт, всего за 25 строк.

Глубокие нейронные сети стали важным элементом искусственного интеллекта, они способны обеспечить высочайшую производительность при обработке естественного языка, в компьютерном зрении и в других областях. Сила этих моделей заключается в их иерархической структуре, которая основывается на большом объеме параллелизируемой работы, хорошо подходящей для многоядерных процессоров, таких как GPU. Фреймворки для работы с GPU, такие как CUDA и OpenCL, в последние годы упростили разработку высокопроизводительных программ. Тем не менее, программы для графических процессоров по-прежнему сложно оптимизировать, отчасти потому, что их архитектуры быстро развиваются.

Для решения этой проблемы появились предметно-ориентированные языки и компиляторы, но эти системы, как правило, менее гибкие и более медленные, чем низкоуровневые ядра библиотек, таких как cuBLAS, cuDNN или TensorRT. Управление всеми этими факторами может оказаться сложной задачей даже для опытных программистов. Таким образом, цель Triton — автоматизировать эти оптимизации, чтобы разработчики могли сосредоточиться на высокоуровневой логике своего кода.

«Новые исследовательские идеи в области глубокого обучения обычно реализуются с использованием комбинации операторов нативных фреймворков… Использование специализированных ядер графического процессора может повысить производительность, но часто возникают сложности из-за многих вопросов программирования графических процессоров. И хотя в последнее время появилось множество систем, облегчающих этот процесс, мы обнаружили, что они либо слишком подробны, либо не обладают гибкостью, либо генерируют код заметно медленнее, чем наши настроенные вручную операторы», — говорит Филипп Тилле, создатель Triton, который сейчас работает в OpenAI. «Наши исследователи уже использовали Triton для создания ядер, которые в 2 раза более эффективны, чем эквивалентные реализации Torch, и мы рады работать с сообществом, чтобы сделать программирование на GPU более доступным для всех».

Согласно OpenAI, Triton, берущий свое начало в статье 2019 года, представленной на Международном семинаре по машинному обучению и языкам программирования, упрощает разработку специализированных ядер, которые могут быть намного быстрее, чем ядра в библиотеках общего назначения. Его компилятор упрощает код и автоматически оптимизирует и распараллеливает его, преобразовывая в код для выполнения на новых графических процессорах NVIDIA (CPU, графические процессоры AMD и другие платформы, кроме Linux, в настоящее время не поддерживаются).

“Основная проблема — это планирование работы, то есть то, как работа, выполняемая каждым экземпляром программы, должна быть разделена для эффективного выполнения на современных графических процессорах”, — объясняет Тилле на веб-сайте Triton. «Чтобы решить эту проблему, компилятор Triton интенсивно использует анализ потока данных на уровне блоков, метод статического планирования итерационных блоков на основе структуры управления и потока данных целевой программы. Полученная система действительно работает на удивление хорошо: нашему компилятору удается автоматически применять широкий спектр интересных оптимизаций».

Первая стабильная версия Triton вместе с руководствами доступна в репозитории проекта на GitHub.