NVIDIA в 10 раз ускорила обучение моделей-«смесей экспертов» на JAX

NVIDIA ускорила обучение моделей со смесью экспертов во фреймворке JAX примерно в десять раз. Компания проверила это на DeepSeek-V3 с 671 млрд параметров и чипах GB300. Без оптимизаций один GPU выдавал 103 терафлопса, с ними — 1 068.

Главной проблемой была не арифметика, а пересылка данных. В исходной версии обмен данными между видеокартами занимал 84% времени работы вычислительных ядер. Причина в устройстве таких моделей. Каждый токен отправляется лишь к нескольким «экспертам» из многих, и выбирает их обученный маршрутизатор. Одному эксперту может достаться гора токенов, другому — почти ничего, и эта картина меняется от шага к шагу.

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

Чтобы это не тормозило обучение, в библиотеку Transformer Engine добавили несколько вещей. Сгруппированное умножение матриц считает всех экспертов за один вызов, и у каждого — его реальное число токенов, без дополнения до максимума. Коммуникационный модуль NCCL EP объединяет рассылку токенов по GPU и сбор результатов в единую цепочку. Если токен нужен нескольким экспертам на одном узле, он передаётся по сети один раз и копируется уже на месте. Кроме того, часть промежуточных данных выгружается в память хоста, а компилятор XLA гоняет пересылки внутри стойки и между узлами параллельно.

Масштабируется система хорошо: на 1 024 GPU в стойках GB300 NVL72 она сохраняет 97% эффективности. Всё уже собрано в контейнер MaxText в каталоге NGC. Параметры для DeepSeek-V3 компания опубликовала, но предупреждает, что другие модели придётся настраивать отдельно. Дальше NVIDIA обещает поддержку формата NVFP4 и совмещение пересылок с вычислениями.