MoE routing в прямых топологиях упирается в непредсказуемый трафик. Разбор MoX показывает, как статическая маршрутизация снижает перегрузки и приближается к производительности switch.
В классических ML-кластерах сеть проектируется под регулярные коллективные операции. Это работает, пока трафик предсказуем. MoE routing ломает это предположение. Каждый токен выбирает top-K экспертов, и возникает разреженный AllToAll-V с сильной вариативностью. Нагрузка меняется от батча к батчу и распределяется неравномерно. В direct-connect топологиях это быстро приводит к деградации: появляются «дорогие» промежуточные хопы (bandwidth tax) и перегруженные линki. Важно, что время выполнения определяется самым загруженным линком, а не средним значением, поэтому даже умеренный skew становится системной проблемой.
MoX предлагает прагматичный сдвиг: не адаптировать топологию под трафик, а заранее оптимизировать MoE routing под класс нагрузки. Вместо динамической реконфигурации сети или предсказания traffic matrix используется статическая политика. Она опирается на две идеи. Первая — token-aware multicast: один токен распространяется по дереву, а не отправляется независимо каждому эксперту. Это снижает дублирование передачи. Вторая — балансировка через предвычисленные веса линков (per-link weights), которые минимизируют ожидаемую нагрузку на самый перегруженный канал. Это компромисс: система игнорирует конкретный runtime-скью, но выигрывает в стабильности и простоте.
Реализация MoE routing в MoX строится вокруг multicast-дерева для каждого токена. Источник постепенно расширяет множество узлов, которые уже получили embedding, и выбирает пути к новым назначениям. Приоритет — маршруты, где промежуточные узлы сами являются получателями. Это превращает relay-хопы в «полезные» доставки. Если такого пути нет, используется обычный relay с неизбежным bandwidth tax. Для балансировки между эквивалентными путями применяется weighted round robin, где вероятность выбора задаётся статическими весами линков. Эти веса вычисляются оффлайн через итеративную оптимизацию, приближающую задачу multicast tree packing. При этом runtime-состояние остаётся компактным: одно значение на линк, без хранения всех возможных маршрутов.
Отдельный момент — фаза combine. MoX использует обратный обход того же дерева и выполняет частичную агрегацию (partial reduction) на промежуточных узлах. Это снижает обратный трафик, поскольку не все результаты экспертов идут независимо к источнику. Такой подход перекладывает часть работы на GPU, но уменьшает давление на сеть.
Результаты показывают, что такой MoE routing даёт системный эффект. В моделировании MoX ускоряет полный MoE-блок до 1.8× по сравнению с min-hop routing. При этом производительность близка к идеальному packet switch: отклонение составляет единицы процентов в ряде конфигураций. Важно, что эффект достигается без динамической адаптации сети. Более того, статическая стратегия показывает лучшие результаты, чем топология, оптимизированная под конкретный трафик, если та использует наивную маршрутизацию. Это указывает на то, что выбор routing policy может быть важнее, чем сама физическая топология.
На практике также видно снижение hotspot-эффекта. В топологии уровня Boardfly bottleneck-нагрузка уменьшается до 47%. Это прямое следствие выравнивания трафика по линкам, а не сокращения числа хопов. Интересно, что основное улучшение даёт именно multicast-структура, а веса лишь доводят балансировку.
Главный вывод: MoE routing можно сделать предсказуемым без контроля над трафиком. Статическая, load-oblivious стратегия работает, если она учитывает семантику MoE (multicast и reduction) и оптимизирует худший линк. Это хорошо ложится на тренд упрощения сетей: меньше динамики в инфраструктуре, больше интеллекта в алгоритмах маршрутизации.
Источник информации
arXiv — крупнейший открытый репозиторий препринтов (с 1991 года, под эгидой Корнелла), где исследователи оперативно размещают рабочие версии статей; материалы общедоступны, но не проходят полное рецензирование, поэтому результаты следует считать предварительными и, по возможности, сверять с обновленными версиями или рецензируемыми журналами. arxiv.org