Рекомендательные нейросети — те, что подбирают для вас ленту новостей, — это продукт машинного обучения. Они работают как гигантские рецепты из тысяч шагов. Каждый шаг — простая операция: сложить, умножить, разрезать массив чисел. Со временем, когда модель дорабатывают, в рецептах накапливаются лишние действия. Представьте повара, который режет морковь на дольки, а потом сразу склеивает их обратно — бессмысленно, но рецепт требует. Точно так же в нейросетях иногда оказывается, что данные сперва разделяют на части, а в следующем шаге снова соединяют. Каждое такое действие запускает отдельную задачу на видеокарте, и вместо одного быстрого расчёта получается очередь из сотен крошечных поручений. Раньше инженеры искали такие огрехи вручную, но модели обновляются каждую неделю — без автоматической оптимизации наводить порядок стало невозможно.
Optimus — встроенный в компилятор PyTorch «чистильщик», который знает более тридцати типичных бесполезных узоров. Он проходит по всей схеме вычислений, которую специалисты называют графом (набором связанных шагов), и сверяет её с шаблонами: увидит знакомый узел из «разрезать-и-склеить» — тут же заменяет его на простую передачу данных. Если в рецепте десять одинаковых операций подряд, он сливает их в одну. Вся эта уборка не меняет финальный результат — вы получите тот же ответ, но путь к нему станет прямым. Благодаря тому, что правила работают на общем внутреннем языке, один набор шаблонов подходит для сотен разных моделей, независимо от их топологии (внутреннего устройства связей).
Эффект заметен. Промышленные модели, которые раньше работали медленно из-за динамически меняющихся схем, стали обучаться до 63% быстрее. В экстремальном тесте численное моделирование показало: количество вычислительных ядер при обучении сжалось почти с 13 тысяч до 24, а при использовании готовой модели — с шести тысяч до шести. Пиковая память снизилась на несколько процентов — критично для больших задач. Код Optimus открыт и встроен в PyTorch, так что любой разработчик может пользоваться им и даже добавлять свои шаблоны — делая нейросети ещё быстрее и экологичнее.
Экономия времени компиляции превысила 400 секунд — этого хватило бы, чтобы заварить чай и проверить почту, пока программа готовит нейросеть к работе.