Huawei представляет Pangu Ultra MoE: разреженная языковая модель с 718 млрд параметров, оптимизированная для Ascend NPU

Эффективные разреженные языковые модели на основе Mixture of Experts

Разреженные большие языковые модели (LLM), использующие архитектуру Mixture of Experts (MoE), становятся популярными благодаря способности масштабироваться, активируя только часть параметров для каждого токена. Такой подход сохраняет высокую емкость модели при уменьшении вычислительной нагрузки на один токен. Однако с ростом моделей до триллионов параметров их эффективное обучение требует инновационных алгоритмов и тесной интеграции программного и аппаратного обеспечения.

Проблемы обучения разреженных LLM на специализированном оборудовании

Основная проблема при обучении разреженных LLM на нестандартных AI-ускорителях, таких как Ascend NPU, — неэффективное использование ресурсов оборудования. Поскольку для каждого токена активна лишь часть параметров, нагрузка между устройствами распределяется неравномерно, что приводит к задержкам синхронизации и недозагрузке процессоров. Память также используется неравномерно: некоторые эксперты обрабатывают больше токенов, чем позволяют их ресурсы. Эти проблемы усугубляются на больших масштабах, когда тысячи чипов сталкиваются с узкими местами в коммуникациях и управлении памятью, что ограничивает пропускную способность и мешает полноценному использованию преимуществ разреженности.

Существующие методы и их ограничения

Для решения этих задач предлагаются разные методы: вспомогательные функции потерь для балансировки распределения токенов, стратегии drop-and-pad для ограничения нагрузки экспертов. Однако они могут снижать качество модели или приводить к дополнительным неэффективностям. Эвристическое размещение экспертов и традиционные схемы коммуникации, такие как All-to-All, часто плохо масштабируются и не обеспечивают высокую производительность. Классические методы экономии памяти, например, рекомпьютинг, обычно выполняются на уровне целых слоев, что увеличивает время работы без адекватной экономии памяти.

Pangu Ultra MoE от Huawei: специализированное решение

Команда Huawei Cloud Pangu разработала Pangu Ultra MoE — разреженную LLM с 718 миллиардами параметров, оптимизированную для Ascend NPU. Их подход начинается с симуляционного конфигурирования модели, в ходе которого оцениваются тысячи архитектурных вариантов с использованием метрик, основанных на реальном поведении оборудования. Это позволяет сэкономить вычислительные ресурсы и точно подобрать гиперпараметры до начала обучения.

Продвинутые методы параллелизма и коммуникаций

Симуляция анализирует комбинации количества слоев, размера скрытого пространства и числа экспертов с использованием пятимерного параллелизма: Pipeline, Tensor, Expert, Data и Context Parallelism. Итоговая модель включает 256 экспертов, скрытый размер 7680 и 61 трансформерный слой. Для повышения производительности применены Adaptive Pipe Overlap для маскировки затрат на коммуникацию и иерархическая All-to-All коммуникация для снижения межузлового трафика. Рекомпьютинг выполнен на уровне ключевых и значений в модулях внимания, а динамический обмен тензорами позволяет выгружать активации в память хоста.

Результаты и показатели

Pangu Ultra MoE достигла Model Flops Utilization (MFU) в 30,0% и обрабатывала 1,46 миллиона токенов в секунду на 6000 Ascend NPU, что значительно превосходит базовый показатель MFU 18,9% и 0,61 миллиона токенов в секунду на 4000 NPU. Динамическое размещение экспертов улучшило баланс нагрузки и увеличило MFU на 10%. Модель показала высокие результаты на бенчмарках: 81,3% на AIME2024, 97,4% на MATH500, 94,8% на CLUEWSC и 91,5% на MMLU. В медицинской области она превзошла DeepSeek R1, набрав 87,1% на MedQA и 80,8% на MedMCQA, подтвердив свою силу в специализированных доменах.

Значение для масштабируемого обучения AI

Данное исследование демонстрирует, как системный поиск архитектуры, оптимизация коммуникаций и управление памятью позволяют раскрыть потенциал масштабных разреженных моделей на специализированном оборудовании. Pangu Ultra MoE от Huawei задает ориентир для будущих AI-систем, гармонично сочетающих алгоритмы и аппаратные возможности для эффективного обучения.