ZAYA1: прорывная модель ИИ, доказавшая эффективность обучения на GPU AMD

Zyphra работала совместно с AMD и IBM в течение года, тестируя графические процессоры и платформу AMD, чтобы проверить возможность обучения крупномасштабных моделей ИИ. Результатом этой работы стала модель ZAYA1.

В рамках сотрудничества компании обучили ZAYA1, которую они называют первой крупной базовой моделью Mixture-of-Experts, полностью построенной на графических процессорах и сетевых технологиях AMD. По мнению участников проекта, этот результат показывает, что рынок ИИ может масштабироваться без обязательной зависимости от решений NVIDIA.

Модель обучалась на GPU AMD Instinct MI300X, сетевых технологиях Pensando и программном обеспечении ROCm, работающем в инфраструктуре IBM Cloud. Примечательно, что сама конфигурация выглядит вполне традиционно. Вместо необычного оборудования или экспериментальных схем Zyphra использовала архитектуру, похожую на корпоративные кластеры, но без компонентов NVIDIA.

По утверждению Zyphra, модель ZAYA1 не уступает признанным открытым моделям в задачах логического мышления, математики и программирования, а в отдельных областях показывает более высокие результаты. Для компаний, которые сталкиваются с ограничениями в поставках или ростом стоимости GPU, это открывает дополнительный вариант без необходимости уменьшать производительность.

Как Zyphra использовала GPU AMD, чтобы снизить расходы без потери качества обучения?

Большинство организаций при планировании вычислительных ресурсов обычно ориентируется на стабильность времени итераций, скорость передачи данных и доступный объем памяти. Это важнее, чем максимальные теоретические показатели пропускной способности.

Zyphra использовала узлы с восемью GPU MI300X, объединенными через InfinityFabric, и подключила каждый GPU к отдельному сетевому адаптеру Pollara. Отдельная сеть обслуживала задачи чтения данных и создание контрольных точек. Архитектура получилась простой, и в этом заключалась ее цель: чем меньше сложных сетевых элементов, тем ниже стоимость инфраструктуры и тем стабильнее обрабатываются итерации.

ZAYA1 как модель искусственного интеллекта, рассчитанная на высокую эффективность

ZAYA1-base активирует ограниченное количество параметров и проходит обучение на большом объеме данных в несколько этапов. Архитектура использует сжатое внимание, улучшенную маршрутизацию токенов к нужным экспертам и упрощенное остаточное масштабирование, что улучшает устойчивость при обучении глубоких слоев.

Для оптимизации процессов Zyphra использовала комбинацию оптимизаторов Muon и AdamW. Чтобы повысить эффективность Muon на оборудовании AMD, команда изменила конфигурации ядер и уменьшила лишний трафик памяти. Размер пакетов увеличивался постепенно, в зависимости от характеристик систем хранения, которые должны были обеспечивать достаточную скорость доставки данных.

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

Совместимость ROCm с GPU AMD

Zyphra отмечает, что перенос рабочего процесса, изначально ориентированного на NVIDIA, на ROCm потребовал значительных усилий. Команда изучала особенности оборудования, адаптировала размеры моделей, вычислительные шаблоны и параметры микропакетов к предпочтительным режимам работы MI300X.

InfinityFabric наиболее эффективно работает, когда все восемь GPU в узле выполняют коллективные операции. Pollara лучше справляется с передачей данных в крупноблочном режиме, поэтому буферы были настроены соответствующим образом. Длинноконтекстное обучение на основе сегментированных последовательностей и деревообразного внимания помогало избегать узких мест при обработке данных.

Вопросы хранения данных потребовали не менее внимательного подхода. Меньшие модели создают больше операций ввода-вывода, а крупные требуют стабильной пропускной способности. Zyphra объединяла сегменты наборов данных, чтобы сократить количество разрозненных обращений, и увеличивала объем кэша страниц для ускорения восстановления контрольных точек, что важно для длительных обучений.

Обеспечение надежной работы кластеров

Длительные задачи обучения редко проходят без сбоев. Сервис Aegis от Zyphra анализирует системные логи и метрики, выявляет неполадки, включая сетевые ошибки или проблемы с корректностью памяти, и автоматически выполняет корректирующие действия. Команда также увеличила время ожидания RCCL, чтобы кратковременные сетевые задержки не останавливали выполнение всей задачи.

Контрольные точки распределяются между всеми GPU, а не сохраняются в одной точке, что ускоряет процесс сохранения. По оценкам Zyphra, такая схема уменьшает затраты времени более чем в несколько раз и снижает нагрузку на операторов.

Что означает обучение ZAYA1 на AMD для процессов закупки в области ИИ?

В отчете проводится сравнение технологий NVIDIA и AMD: NVLINK и InfinityFabric, NCCL и RCCL, cuBLASLt и hipBLASLt. Авторы считают, что стек решений AMD уже достаточно развит для обучения крупномасштабных моделей.

Речь не идет о полном отказе предприятий от инфраструктуры NVIDIA. Более реалистичная стратегия заключается в использовании NVIDIA для продуктивных систем и AMD для этапов обучения, где важны открытость ROCm и возможности работы с крупными объемами данных. Такой подход распределяет риски и увеличивает возможности масштабирования.

Авторы делают несколько выводов: необходимо рассматривать форму модели как гибкую, а не фиксированную; проектировать сети с учетом коллективных операций; строить инфраструктуру, которая обеспечивает бесперебойную работу GPU; и оптимизировать работу с контрольными точками, чтобы они не прерывали процесс обучения.

Эти наблюдения основаны на опыте Zyphra, AMD и IBM при обучении крупной модели искусственного интеллекта MoE на GPU AMD. Для организаций, которые стремятся увеличить свои возможности в области ИИ и не зависеть от одного поставщика, это может стать полезным ориентиром.

Источник: https://www.artificialintelligence-news.com/