DisagMoE: ускорение обучения MoE-моделей до 1,8 раза за счёт разделения вычислений
Обучение больших языковых моделей на архитектуре Mixture-of-Experts (MoE) сталкивается с серьёзными узкими местами из-за необходимости интенсивного обмена данными между GPU. Экспертный параллелизм (expert parallelism), широко применяемый для MoE, требует all-to-all коммуникации, что особенно критично при распределении экспертов по разным узлам.
Существующие подходы пытались решить эту проблему, перекрывая all-to-all обмен с вычислениями в сетях прямой связи (FFN) и самовнимания (attention). Однако из-за разного соотношения вычислений и коммуникации в этих слоях часто остаются простои сети.
Новый метод DisagMoE, описанный в препринте на arXiv, предлагает кардинально другой подход. Система разделяет слои внимания и FFN на отдельные, несвязанные GPU-группы. Вводится многостадийный конвейер с однонаправленной связью типа «многие ко многим», что позволяет эффективнее совмещать вычисления и передачу данных.
Для оптимального распределения пропускной способности GPU и сети между группами используется roofline-модель, которая сбалансирует нагрузку и минимизирует простои.
DisagMoE реализован на базе популярного фреймворка Megatron-LM. Тестирование проводилось на кластере из 16 узлов, каждый с 8 ускорителями NVIDIA H800. Эксперименты с различными MoE-моделями показали ускорение обучения до 1,8 раза по сравнению с традиционным экспертным параллелизмом.
Разработка представляет практический интерес для компаний, занимающихся обучением сверхбольших разрежённых моделей, так как позволяет более эффективно использовать имеющееся оборудование и сокращать время тренировки.
