Компания Moore Threads сделала проект Megatron-MUSA-Titan общедоступным, представив инфраструктуру для обучения больших моделей, которую она разработала для использования на собственных графических процессорах. Этот проект предназначен для работы с моделями типа 'смесь экспертов' (mixture-of-experts) с триллионом параметров на кластерах, состоящих из более чем 10 000 GPU от Moore Threads, и охватывает как предварительное обучение, так и обучение с подкреплением.
Код проекта доступен на GitHub. Компания утверждает, что данная система непрерывно использовалась для предварительного обучения MoE с триллионом параметров на кластерах мощностью 10 000 GPU, где длительные сессии обучения обеспечивали высокую утилизацию FLOPs модели и стабильность процесса.
В репозитории представлен пример конфигурации для модели MoE в стиле Kimi-K2 с 1 триллионом параметров. Центральным принципом проектирования стало сохранение близости к существующим решениям. Megatron-MUSA-Titan построен на основе фреймворка Megatron-LM от NVIDIA. Благодаря интеграции PyTorch через torch_musa и совместимости MUSA C с CUDA C, разработчики, имеющие среду MUSA, могут клонировать официальные репозитории Megatron-LM, slime и miles и запускать обучение на GPU Moore Threads без необходимости модификации исходного кода.
Проект указывает, что общие алгоритмы и исправления ошибок следует направлять в вышестоящие проекты, тогда как Moore Threads фокусируется на двух типах изменений: крупномасштабных оптимизациях, адаптированных под оборудование и кластеры Moore Threads, а также на обходных решениях для специфических проблем MUSA.
Функции предварительного обучения включают блочное обучение в формате FP8, которое опирается на MT-TransformerEngine от Moore Threads. Диспетчеризация и объединение токенов MoE ускоряются с помощью DeepEP, работающего поверх Async Copy Engine компании, который также позволяет коммуникации экспертов параллельно с вычислениями. Среди других опций — выгрузка активаций с использованием виртуального параллелизма конвейера и планирование одного прямого прохода и одного обратного прохода, а также параллельная коммуникация конвейера, совмещенная с вычислениями, благодаря Mooncake RDMA.
Инструмент анализа распределенной трассировки считывает трассировки PyTorch с каждого GPU для обнаружения узких мест, таких как пузырьки конвейера и узлы с низкой производительностью. Проект заявляет, что он способен анализировать запуск MoE с 1 триллионом параметров на 10 000 GPU всего за несколько часов.
Для обучения с подкреплением стек поддерживает фреймворки slime и miles; Megatron занимается обучением, а MUSA-сборка SGLang генерирует прогоны, что позволяет одной базовой системе обслуживать как предварительное обучение, так и RL. Оптимизация для автономного инференса не входит в сферу действия данного проекта.
Для разработчиков, рассматривающих отечественные ускорители, практический вопрос заключается в том, какая часть существующего кода обучения останется работоспособной. Ответ Moore Threads состоит в том, что код экосистемы Megatron не требует изменений.
