Компилятор Nova ускоряет ИИ: прирост до 10.6% к PyTorch и экономия памяти 29%
Эффективность моделей глубокого обучения напрямую зависит от того, как математические операции отображаются на физическое оборудование. Популярные фреймворки с eager-исполнением не всегда могут использовать весь потенциал GPU. Новый компилятор Nova, представленный на arXiv, предлагает автоматическое решение этой проблемы.
Nova — это сквозной JIT-компилятор, который обеспечивает «абсолютный контроль» над отображением операций на аппаратуру. Он объединяет операции, выходя за границы отдельных математических функций, оптимизирует сложные иерархии памяти и настраивает выполнение вплоть до регистрового уровня.
Технически Nova захватывает eager-выполнение и объединяет прямой и обратный проходы в единый диалект со value-семантикой, что открывает возможности для агрессивной оптимизации всего графа вычислений. Аналитический конфигуратор определяет оптимальные расписания на основе арифметической интенсивности без затрат времени на поиск, а структурное хеширование позволяет генерировать тонко настроенные ядра непосредственно из структуры вычислений.
В тестах на видеокарте RTX 3060 компилятор показал результаты на уровне или немного лучше библиотек cuBLAS и XLA при выполнении матричных умножений в формате TF32, сохраняя относительную ошибку менее 5e-4. На модели с 42 миллионами параметров Nova достигла пропускной способности на 10.6% выше, чем PyTorch, и на 4.4% выше, чем XLA.
Одним из ключевых преимуществ стало сокращение потребления памяти до 29% по сравнению с PyTorch. Благодаря этому Nova успешно обучила модель на 144 миллиона параметров со скоростью 17 900 токенов в секунду, тогда как PyTorch столкнулся с ошибкой нехватки памяти на том же GPU с 12 ГБ видеопамяти.
Разработчики подчеркивают, что Nova предназначена для автоматизации низкоуровневой оптимизации и может быть полезна как исследователям, так и инженерам, работающим с большими моделями на ограниченном оборудовании. Исходный код и дополнительные материалы доступны в статье на arXiv.
