Прорыв в производительности: от 103 до 1068 TFLOPS/GPU
Команда NVIDIA представила оптимизированный стек для обучения моделей Mixture of Experts (MoE) в фреймворке JAX. Ключевым достижением стало использование NVIDIA Transformer Engine, которое позволило достичь 10.4-кратного ускорения пропускной способности на архитектуре NVIDIA GB200. Производительность модели DeepSeek-V3 выросла с базовых 103 TFLOPS/GPU до 1,068 TFLOPS/GPU. Это критически важно, так как в неоптимизированной версии до 84% времени накапливания ядра уходило на меж-GPU коммуникации, что делало обучение неэффективным.
Почему Dropless MoE сложнее Capacity-based?
Стандартные подходы (Capacity-based MoE) фиксируют бюджет токенов для каждого эксперта, обрезая или дополняя данные, что ведет к потере качества или вычислительным потерям. Dropless MoE обрабатывает каждый токен без потерь, но создает «рваные» (ragged) тензоры, так как количество токенов у экспертов варьируется. Это ломает стандартные оптимизации для прямоугольных матриц. NVIDIA решает эту проблему через Grouped GEMM — ядра, которые обрабатывают переменное количество токенов за один вызов, используя cuBLAS и cuBLASLt для полной загрузки Tensor Cores.
Технические детали оптимизаций
Для достижения высоких результатов внедрены следующие компоненты:
- Grouped GEMM: Обработка всех экспертных умножений матриц в одном вызове ядра, избегая циклов и копирования данных Device-to-Host, что сохраняет CUDA-графики.
- NCCL EP (Expert Parallelism): Объединение этапов dispatch и combine, дедупликация токенов для снижения трафика в сети.
- MXFP8 Quantization: Использование 8-битной квантования с блочным масштабированием для экспертных матричных умножений на GPU Blackwell.
- JAX Host Offloading & XLA Multistreaming: Снижение узких мест памяти и перекрытие коммуникаций через NVLink и InfiniBand.
Результаты масштабирования
Полный стек демонстрирует стабильное масштабирование. При обучении модели DeepSeek-V3 (671B параметров) на 1,024 GPU в кластере NVIDIA GB300 NVL72 достигается эффективность масштабирования на уровне 97%. Это подтверждает готовность архитектуры к промышленному развертыванию сверхбольших моделей.
| Метрика | Базовая конфигурация (Unoptimized) | С NVIDIA Transformer Engine + JAX |
|---|---|---|
| Пропускная способность (TFLOPS/GPU) | 103 | 1,068 |
| Время на коммуникации | 84% от времени ядра | Значительно снижено (за счет NCCL EP и fused kernels) |
| Эффективность масштабирования (1024 GPU) | Информация недостаточна | 97% |
| Тип MoE | Capacity-based (обычно) | Dropless (без потери токенов) |
Как попробовать?
Разработчики рекомендуют использовать контейнер NVIDIA NGC MaxText с включенным Transformer Engine. Подробные инструкции по конфигурации MoE доступны в руководстве MaxText, а документация Transformer Engine описывает все доступные возможности библиотеки для ускорения обучения.
Источник: NVIDIA dev blog ↗