Оптимизация нейросетей: как спрямить запутанный граф вычислений ✛
Контекст
Глубокие нейросети для рекомендательных систем — это монстры вычислительной сложности. Чем точнее модель, тем запутаннее её внутренний ландшафт: тысячи операций переплетаются в графы с причудливой топологией. Каждый лишний виток, каждое повторное расщепление и склейка тензоров — как нерасчищенная тропа, замедляющая путника. На аппаратном уровне это оборачивается лавиной мелких ядерных запусков (kernel launch) на GPU, которые съедают время и память. Инженеры пытаются наводить порядок вручную, но модели меняются еженедельно, а архитектуры становятся всё прихотливее. Возникла нужда в автоматическом «спрямителе» графовых троп, способном без вмешательства человека находить и вычищать неэффективные фрагменты кода.
Методы
Сердце Optimus — коллекция из более чем тридцати шаблонов оптимизации, каждый из которых описывает типовой «мусорный» узор в графе: например, бессмысленную цепочку «разрезать тензор на куски и тут же склеить обратно». Поиск таких узоров ведётся жадным алгоритмом, который стартует от «якорей» — ключевых операций вроде split или общих для рекомендательных архитектур. От якоря алгоритм обходит ближайших соседей, сравнивая их с эталонным шаблоном, и при совпадении запускает замену: вместо громоздкой конструкции вставляется компактный эквивалент. Различают два типа починки: горизонтальную — когда несколько независимых одинаковых операций (скажем, дюжина линейных слоёв) сливаются в одну, и вертикальную — когда вырезаются избыточные переходы, например пресловутые split-then-cat. Вся процедура вписана в компилятор PyTorch 2.x и работает на любом его промежуточном представлении (Torch-IR, Aten-IR), что делает её универсальной для разношёрстных моделей.
Результаты
На пяти промышленных рекомендательных моделях Meta разных архитектурных типов (DHEN-подобные, Wukong-подобные, Interformer-подобные и другие) Optimus стабильно увеличивал пропускную способность — до 63% больше обучающих примеров в секунду. Основной вклад дали вертикальные оптимизации: ликвидация паразитных split-cat циклов срезала сотни лишних ядер Triton. Прирост от горизонтальных слияний зависел от количества одинаковых операций — чем их больше, тем весомее выигрыш. Пиковое потребление памяти на больших моделях с квантованием активаций упало на 2–6%, а время холодной компиляции — просто рухнуло: в одном из тестов с 4000+ секунд до примерно 3600, то есть экономия превысила 400 секунд. На открытом иллюстративном примере из TorchBench число ядер (kernel) в обучении сократилось с 12 772 до 24, а в инференсе — с 6 320 до 6, что дало ускорение обучения в 11 раз и инференса в 364 раза. Сравнение с аналогами (TASO, TENSAT) показало, что те вообще не смогли улучшить граф для этой модели, тогда как Optimus справился благодаря гибкой системе правил.
Значение
Optimus — это не разовая поделка, а модульный каркас, встроенный в открытый код PyTorch. Любой исследователь или инженер может добавить новый шаблон оптимизации, зарегистрировав его через стандартный интерфейс, и тут же применить к своему нейросетевому графу. Такой подход превращает некогда муторную ручную возню с кодом в планомерное наведение чистоты на уровне компилятора. Для индустрии это означает, что дорогие GPU будут использоваться эффективнее без переписывания моделей, а заодно снизится углеродный след от обучения гигантских рекомендательных систем.
Развитие
Сейчас основные правила работают на Torch-IR и Aten-IR, но авторы планируют распространить метод на Inductor-IR — более низкоуровневое представление, где можно выжать ещё больше производительности. Другое направление — автоматический выбор порядка и набора шаблонов под конкретную модель: сегодня последовательность оптимизаций подбирается вручную, а хотелось бы поручить это алгоритму, анализирующему структуру графа.
Влияние
Технология непосредственно ускорит обучение и развёртывание рекомендательных систем — ключевого звена современных интернет-платформ. Её также можно адаптировать для любых других глубоких нейросетей, особенно с разветвлёнными графами вычислений: обработка естественного языка, компьютерное зрение, графовые нейросети.
Следующие шаги
Перенос наработок на уровень Inductor-IR и внедрение автоматического планировщика последовательности шаблонов. Также предстоит расширить библиотеку шаблонов, покрывая новые архитектурные мотивы.
Ключевые проблемы
Проблема эффективного исполнения нейросетей на современном «железе» упирается в разрыв между сложностью архитектур и ограниченной пропускной способностью памяти GPU. Optimus — шаг к автоматическому устранению этого разрыва на уровне компилятора, не требуя от разработчика переделки кода. Это вписывается в более широкий контекст поиска масштабируемых методов численной оптимизации вычислений.
Взгляд машины знаний
Этот раздел адресован автору работы и его коллегам. Соседние работы нашёл смысловой поиск по нашему архиву — не по совпадению слов, а по смыслу. Мы не рецензируем: мы показываем поле вокруг.
Что сделано в работе
В работе представлен Optimus — модульный каркас для автоматической оптимизации графов вычислений в PyTorch 2.x, нацеленный на рекомендательные модели. Он использует коллекцию из более чем тридцати шаблонов для поиска и замены неэффективных паттернов (например, split-then-cat) с помощью жадного алгоритма от якорей. На пяти промышленных моделях Meta достигнуто ускорение обучения до 63%, снижение пикового потребления памяти на 2–6% и значительное сокращение времени компиляции (с 4000+ до ~3600 секунд в одном тесте).
В чём сила
Сила работы — в практической интеграции с компилятором PyTorch и гибкой системе правил, позволяющей автоматически находить и устранять неэффективные паттерны в реальных промышленных графах, где аналоги (TASO, TENSAT) не справляются.
Куда можно двигаться
- Может быть интересно посмотреть на опыт SpecFANN по ускорению вычислений через специализированные архитектуры: там используется быстрый эмулятор для генерации спектров, что перекликается с идеей замены тяжёлых фрагментов лёгкими эквивалентами. Возможно, принципы замены на основе обученных суррогатов можно адаптировать для расширения набора шаблонов Optimus.опора: SpecFANN: Молниеносный набросок звёздного спектра arXiv:2607.17348v1
- Рядом лежит работа про нейросеть, которая чувствует напряжение в деталях любой формы: там применяются методы анализа графов для выявления критических узлов. Это может подсказать, как улучшить эвристику выбора якорей в Optimus, чтобы находить более глобальные оптимизации.опора: Нейросеть, предсказывающая механические напряжения в деталях произвольной геометрии arXiv:2606.08287
Что можно попробовать
- Попробовать использовать обученный суррогат для предсказания выигрыша от применения шаблона, чтобы приоритизировать замены.
- Апробировать анализ «напряжения» в графе (по аналогии с работой про детали) для выбора якорей, которые с большей вероятностью приведут к каскадным оптимизациям.
- Проверить, можно ли автоматически генерировать новые шаблоны, обобщая частые паттерны в реальных графах.
Работы рядом в нашем архиве
- Нейросеть, предсказывающая механические напряжения в деталях произвольной геометрииблизость 0.62
- SpecFANN: Молниеносный набросок звёздного спектраблизость 0.62
- Как цифровой двойник спасает мозг за минутыблизость 0.62
- Нейросеть учится течь сквозь ситоблизость 0.60
- Квантовые метки графов: когда сложность уступает простотеблизость 0.60
- Нейроны раскрывают секреты за минуты, а не часыблизость 0.59
Ближайшая работа архива — 0.62 по смысловой близости. Обычный для нашего корпуса уровень.
Оригинальная аннотация
как её написали авторы, на английском arXiv:2607.02945
In large-scale industrial applications, deep learning models that power recommendation and ranking have complex and diverse model architectures. These models are continuously developed and refined by large teams of machine learning engineers, rendering manual optimization infeasible. Consequently, graph-based optimization techniques have become an industry standard for boosting performance, with PyTorch FX transformations leading the charge. These transformations typically rely on a set of human-engineered module-level rewrite rules which are not scalable to diverse model architectures. To address this limitation, we introduce Optimus, a general-purpose model transformation framework built in the PyTorch 2.x (PT2) machine learning compiler. With a concise set of predefined patterns, Optimus applies an efficient greedy search algorithm for pattern matching and replacement, while preserving model semantic. It is designed and implemented as a highly customizable and extensible framework integrated into the PT2 stack. Our evaluation shows that the framework can achieve up to 63% speedup, 6% peak memory reduction, and over 400 second compile time decrease for our industry-scale recommendation models compared to baselines. Optimus is open-sourced together with PyTorch 2.x as a customizable model transformation layer.
Сокращение времени компиляции на 400 секунд — это как если бы вместо того, чтобы варить кофе, проверять почту и ждать, пока откроется среда разработки, вы сразу приступали к делу. А количество ядерных запусков в инференсе иллюстративной модели упало с 6320 до 6 — такое уплотнение графа можно сравнить с заменой толпы курьеров одним грузовиком.
Ключевые числа
- прирост пропускной способности (QPS): до 63%
- снижение пиковой памяти: до 6%
- уменьшение времени компиляции: более 400 секунд
- ядер в TorchBench обучении до/после: 12 772 -> 24
- ядер в TorchBench инференсе до/после: 6 320 -> 6