NVIDIA в 10 раз ускорила обучение DeepSeek-V3 на библиотеке JAX
На ускорителях GB300 скорость выросла со 103 до 1 068 терафлопс на видеокарту. Модель обучали так, что эксперты не отбрасывают ни одного токена. Весь набор доработок вошёл в готовый контейнер MaxText.
NVIDIA в 10,4 раза ускорила обучение моделей с архитектурой «смесь экспертов» (MoE) на библиотеке JAX1, говорится в блоге компании для разработчиков. Замер делали на DeepSeek-V3 с 671 млрд параметров и ускорителях GB3002. Исходная сборка выдавала 103 терафлопса3 на видеокарту, доработанная — 1 068.
В MoE-модели вместо одного большого вычислительного блока стоит много небольших сетей-«экспертов». Маршрутизатор решает, каким экспертам отдать каждый токен. Так устроены DeepSeek, Qwen и Mixtral. Трудность в том, что нагрузка на экспертов неравномерна и меняется от шага к шагу: одному эксперту достаётся много токенов, другому мало.
Обычно эту проблему обходят квотами: каждый эксперт обрабатывает фиксированное число токенов, лишние отбрасывают, а недостающие места заполняют пустышками. Видеокартам так удобнее, но модель либо учится на неполных данных, либо тратит вычисления впустую. NVIDIA оптимизировала вариант без отбрасывания токенов: каждый токен доходит до выбранного эксперта.
В исходной сборке 84% времени работы видеокарт уходило на обмен данными между ними. Библиотека NVIDIA Transformer Engine теперь считает умножения матриц для всех экспертов за один вызов, и у каждого эксперта столько строк, сколько токенов он получил на самом деле. Коммуникационный модуль NCCL EP объединяет рассылку токенов по экспертам и сбор результатов. Если токен нужен нескольким экспертам на одном удалённом узле, по сети он идёт один раз. Кроме того, часть промежуточных данных выгружается из памяти видеокарты в память сервера, а компилятор XLA запускает независимые обмены данными параллельно.
На стойках GB300 NVL72 связка сохраняет 97% эффективности масштабирования при обучении DeepSeek-V3 на 1 024 видеокартах, сообщает компания. Доработки уже есть в контейнере MaxText4 из каталога NGC, начиная с образа от 9 сентября 2026 года. В планах NVIDIA — поддержка формата NVFP4 и дальнейшее слияние операций.