Спрямление графовых троп: как Meta ускорила обучение нейросетей на 63% ✛

‹›

Рекомендательные системы, определяющие, какой пост вы увидите следующим, работают на глубоких нейросетях из тысяч операций — сложных графов вычислений (схемах из операций над данными). Эти системы — живые, постоянно эволюционирующие динамические системы: малейшее изменение архитектуры распространяется каскадом по всему графу. С ростом точности моделей тропы запутываются: избыточные разветвления вроде «разрезать тензор (многомерный массив чисел) на куски и тут же склеить обратно» создают настоящие пробки на графических процессорах. Каждый такой извив запускает отдельное вычислительное ядро, а за день обучения могут набежать триллионы лишних запусков. Инженеры пытаются расчищать эти тропы вручную, но модели меняются каждую неделю, и наводить порядок в энтропийном хаосе без помощи автомата стало невозможно.

Optimus — это модуль внутри компилятора PyTorch 2.x, работающий как жадный картограф. Он вооружён библиотекой более чем из тридцати шаблонов оптимизации — типовых мусорных узоров вроде цепочки split-then-cat. Алгоритм стартует от якорей — критических операций, вокруг которых наматывается путаница, — и обходит соседние узлы, сверяя каждый шаг с эталонным шаблоном. Найдя совпадение, он вырезает громоздкую конструкцию и вставляет компактный эквивалент. Горизонтальная чистка сливает десяток одинаковых слоёв в один, вертикальная — убирает бессмысленные расщепления и склейки. Вся процедура не меняет конечный результат, но радикально упрощает внутренний граф. Именно такой оптимизации поддаются графы любой топологии, ведь правила применяются на универсальном промежуточном представлении Torch-IR.

Эффект на промышленных системах Meta оказался ощутимым. Пять рекомендательных моделей разных архитектур показали прирост пропускной способности до 63% — больше обучающих примеров в секунду без замены оборудования. Вертикальные чистки срезали сотни паразитных ядер Triton, а горизонтальные сливали россыпь одинаковых операций в одну; чем больше клонов, тем заметнее выигрыш. Пиковое потребление памяти снизилось на 2–6%, что критически важно для моделей с квантованием активаций, где каждый мегабайт на счету. В иллюстративном тесте на TorchBench количество ядер в обучении рухнуло с 12 772 до 24, а в инференсе — с 6 320 до 6, дав ускорение в 11 и 364 раза соответственно. Такой численный триумф стал возможен благодаря тому, что Optimus оперирует на уровне компилятора, видя граф целиком.

Этот фреймворк превращает беспорядочную сеть тропинок в автобан для данных, и самое главное — все его инструменты открыты каждому. Код встроен в PyTorch, и любой разработчик может не только использовать готовые шаблоны, но и добавлять собственные, обучая систему новым приёмам спрямления. В перспективе авторы планируют перенести оптимизацию на ещё более низкий уровень Inductor-IR и научить алгоритм самому выбирать наилучшую последовательность правил — без участия человека. Это не просто ускорение обучения; это шаг к самооптимизирующимся нейросетям, которые сами подстраиваются под железо, экономя электричество и время. Запутанные тропы графов больше не будут пожирать ресурсы впустую — их выпрямит автопилот.