Success Story - rus

Оптимизация вычислений графа нейросетевой модели

Описание проекта:

Перед командой MIL Team стояла R&D-задача: разработать эффективный и адаптивный подход к оптимизации вычислительных графов нейросетей, ориентированный на ускорение процессов на аппаратной платформе Huawei Atlas 800.

Заказчик столкнулся с рядом технологических ограничений: существующие решения, такие как TASO, основаны на фиксированных алгоритмах и плохо адаптируются к современным нейросетевым архитектурам. Это мешает эффективно реализовать:
  • device placement — распределение операций по вычислительным устройствам (CPU/NPU);
  • operation fusion — агрегацию последовательных операций для снижения накладных расходов;
  • scheduling – выбор оптимального порядка выполнения операций;
  • data layout — эффективную организацию данных в памяти.

Кроме того, традиционные методы не позволяют решать сразу несколько задач оптимизации согласованно, что приводит к снижению производительности и нерациональному использованию вычислительных ресурсов.

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

Решение:

Наша команда MIL Team разработала систему оптимизации графов на основе обучения с подкреплением (Reinforcement Learning, RL). В отличие от классических подходов (TASO), RL-агент итеративно принимает решения об оптимизации графа, обучаясь на собственном опыте. Такой пошаговый подход позволяет учитывать глобальные зависимости в графе, адаптироваться к его топологии и аккумулировать знания о долгосрочных эффектах локальных изменений, что особенно важно при одновременной оптимизации нескольких аспектов.

В рамках проекта были реализованы и сравнительно исследованы различные пространства действий:
  • DAS1 / DAS2 (Discrete Action Spaces) — дискретное присваивание устройств узлам/сегментам;
  • CAS1 / CAS2 / CAS3 (Continuous Action Spaces) — предсказание границ разбиения в непрерывном пространстве при фиксированном числе подграфов;
  • Iterative action space v3 — гибридный подход с последовательным уточнением стратегии и выбором переменного числа подграфов.

Каждая стратегия имела свои особенности: дискретные подходы были проще в реализации, но ограничены по разнообразию решений из-за маскирования; непрерывные подходы требовали точного контроля границ разбиений и балансировки нагрузки.
Для повышения стабильности и ускорения обучения были интегрированы дополнительные механизмы мотивации агента:
  • ICM (Intrinsic Curiosity Module) — усиливал изучение новых траекторий за счёт моделирования переходов между состояниями, повышая разнообразие действий и ускоряя обучение;
  • HER (Hindsight Experience Replay) — переосмысливал неудачные эпизоды как потенциально успешные при других целях, эффективно работая на ранних стадиях, но ограниченно применимый в условиях высокой динамичности графа.

RL-агент обучался на графах с различной топологией, получая в качестве входа: типы операций, связи, размещение по устройствам, объёмы трафика между узлами и метаинформацию о ресурсоёмкости. Вознаграждение строилось на измерении ускорения исполнения графа по сравнению с базовой реализацией.

Итоги:
В результате предложенное решение ускорило исполнение на аппаратной платформе Huawei Atlas 800. Предложенный RL-подход продемонстрировал преимущество над TASO за счёт способности учитывать глобальные зависимости и совместно решать несколько задач оптимизации.

Разработанное решение успешно прошло тестирование на различных нейросетевых архитектурах и продемонстрировало высокую степень масштабируемости — до графов с миллионами узлов и миллиардами рёбер. В среднем была достигнута прибавка в производительности в х1.39, а на отдельных топологиях — до х2 ускорения относительно стандартного плейсмента.

Решение было построено на open-source фреймворках, что обеспечивает лёгкую интеграцию, гибкость и расширяемость под прикладные задачи. Переданные заказчику результаты закладывают основу в создания универсальных, самообучающихся решений для оптимизации вычислений, способных эффективно работать на различных аппаратных платформах и в условиях реальных производственных ограничений.
Research Compression