Проблема масштабирования: 10 ТБ трафика на шаг
В современных системах Reinforcement Learning (RL) веса модели периодически обновляются на GPU инференса (inference) на основе обновлений, полученных на GPU обучения (training). Традиционный метод gather-and-broadcast требует реконструкции полного тензора на каждом шаге, что создает избыточный сетевой трафик. Для современных моделей с весами в BF16 (2 байта на параметр) логический объем передаваемых данных достигает ~2x от размера модели:
| Модель | Размер весов (BF16) | Оценка трафика Refit (All-gather + Broadcast) |
|---|---|---|
| Kimi K3 | 5.6 TB | ≈ 11.2 TB |
| Qwen 3.8 2.4T A95B | 4.8 TB | ≈ 9.6 TB |
При размере модели Kimi K3 в 2.8 триллиона параметров, такой подход становится узким местом, замедляющим обучение.
Решение: Reshard как DTensor-to-DTensor transfer
Команда NVIDIA переосмыслила задачу Refit как проблему перераспределения распределенных тензоров (DTensor). Вместо сбора полных тензоров, алгоритм NCCL M2N Reshard вычисляет точные пересечения фрагментов (shards) между источниками и получателями, передавая только необходимый минимум данных. Это устраняет избыточные передачи, характерные для gather-and-broadcast.
Результаты бенчмарков: до 5.7x ускорения
Тестирование проводилось на системах NVIDIA DGX H100 и DGX GB200 с использованием фреймворка NVIDIA NeMo RL. Применение Reshard показало значительный прирост скорости по сравнению с традиционным методом:
- Ускорение: от 2.7x до 5.7x в зависимости от конфигурации и модели.
- Тестируемые модели: Qwen 3 235B и DeepSeek v3 671B.
- Гибридный подход: Reshard применяется к весам проекций feed-forward (которые составляют основную часть трафика), в то время как для меньших весов сохраняется gather-and-broadcast для обеспечения функциональной полноты.
Доступность и внедрение
Референсная реализация NCCL M2N Reshard доступна в репозитории NVIDIA/nccl-extensions на GitHub. Разработчикам рекомендуется изучить руководство по NeMo RL refit для настройки бэкендов, лейаутов и режимов точности, а также требования к конфигурации NCCL для интеграции в свои пайплайны обучения.
Источник: NVIDIA dev blog ↗