В эпоху, когда размер моделей искусственного интеллекта перевалил за триллионы параметров, традиционные подходы к распределенным вычислениям начинают показывать свою уязвимость. Одним из самых критичных узких мест в процессе обучения с подкреплением (Reinforcement Learning, RL) для больших языковых моделей (LLM) стала операция, известная как refit (подгонка весов). Представьте себе сценарий, где модель обучается на одном кластере GPU, а генерирует траектории для обучения — на другом. Чтобы обучение было эффективным, веса, обновленные на обучающих GPU, должны быть мгновенно перенесены на инференсные GPU. Для современных гигантов, таких как Qwen 3 или DeepSeek v3, объем этих данных исчисляется десятками терабайт. Старые методы передачи, основанные на сборе и широковещательной рассылке (gather-and-broadcast), превращаются в узкое горлышко, тормозящее весь процесс обучения.
Команда NVIDIA представила революционное решение — Reshard collective. Это не просто оптимизация, а фундаментальная переформулировка проблемы передачи данных. Вместо того чтобы отправлять полные копии тензоров, Reshard перемещает только теоретически необходимый минимум данных, напрямую сопоставляя фрагменты (shards) обучающих и инференсных устройств. В этой статье мы подробно разберем, как работает эта технология, почему она ускоряет процесс в 2,7–5,7 раза и как разработчики могут интегрировать её в свои рабочие процессы с помощью фреймворка NeMo RL.
01Проблема масштабирования: почему старые методы больше не работают
Чтобы понять масштаб проблемы, нужно заглянуть в архитектуру современных систем обучения с подкреплением. В типичном сценарии инференс и обучение разделены (disaggregated setup). Обучающие GPU обновляют веса модели, а инференсные GPU используют текущую версию модели для генерации данных. Поскольку веса обновляются на каждом шаге, система должна периодически синхронизировать их. Эта операция называется refit.
Существует две ключевые характеристики refit, которые делают его сложной задачей:
- Различные схемы параллелизма. Обучающие и инференсные кластеры часто имеют разную конфигурацию GPU и используют разные методы параллелизма (data, tensor, pipeline, expert parallelism). Это означает, что веса распределены по-разному: на обучающих GPU тензор может быть разрезан по строкам, а на инференсных — по столбцам.
- Огромный объем передаваемых данных. Поскольку refit должен обновить все веса, объем передаваемых данных сопоставим с размером всей модели. Для моделей уровня Kimi K3 (2,8 трлн параметров) или Qwen 3.8 (2,4 трлн параметров) размер весов в формате BF16 составляет около 5,6 ТБ и 4,8 ТБ соответственно.
Традиционный подход использует двухэтапный процесс: сначала все обучающие GPU собирают (all-gather) свои локальные фрагменты, чтобы воссоздать полный тензор, а затем этот полный тензор рассылается (broadcast) всем инференсным GPU. Проблема в том, что инференсный GPU получает полный тензор, хотя ему нужен лишь его небольшой фрагмент. Это создает избыточную коммуникацию на обоих этапах. Логический объем передаваемых данных при таком подходе удваивается: для модели в 5 ТБ объем передачи составляет около 10 ТБ. При выполнении этой операции на каждом шаге обучения такие затраты становятся неприемлемыми.
02Переосмысление задачи: от Gather-and-Broadcast к DTensor-to-DTensor
Решение NVIDIA заключается в отказе от идеи передачи полных тензоров. Вместо этого команда переформулировала задачу refit как проблему перераспределения (resharding) тензоров. Ключевым инструментом здесь стала абстракция PyTorch Distributed Tensor (DTensor). DTensor позволяет описывать, как тензор разрезан или размножен (replicated) между GPU, учитывая глобальную форму тензора, сетку устройств (device mesh) и размещение (placements).
Вместо того чтобы собирать полный тензор, Reshard вычисляет точные пересечения между исходными фрагментами (на обучающих GPU) и целевыми фрагментами (на инференсных GPU). Рассмотрим простой пример: матрица 8x8, разрезанная по строкам на двух обучающих GPU и по столбцам на четырех инференсных GPU. Каждый обучающий GPU хранит фрагмент 4x8, а каждый инференсный GPU нуждается в фрагменте 8x2. Чтобы получить нужный фрагмент, инференсный GPU должен получить область 4x2 от каждого из двух обучающих GPU. Эти области напрямую вытекают из пересечения исходных и целевых фрагментов.
Этот подход позволяет передавать только те данные, которые действительно нужны получателю. Никаких лишних копий, никаких полных тензоров в промежуточной стадии. Это не только снижает объем передаваемых данных, но и делает процесс более предсказуемым и эффективным с точки зрения использования сетевых ресурсов.
03Как работает NCCL M2N Reshard
Реализация этой концепции доступна в виде открытого исходного кода под названием NCCL M2N Reshard (доступно в репозитории NVIDIA/nccl-extensions). Эта библиотека, построенная на базе NCCL (NVIDIA Collective Communications Library), принимает описания исходной и целевой раскладок тензоров и локальные буферы для выполнения операции.
Главная инновация NCCL M2N Reshard заключается в том, как он использует физическую топологию сети GPU. Вместо того чтобы отправлять каждый исходный фрагмент каждому инференсному GPU независимо, Reshard группирует реплики по доменам NVIDIA NVLink. Внутри каждого домена он распределяет входящие данные максимально равномерно между рангами (GPU), так что каждый получающий узел получает уникальный набор данных по сети.
Затем эти ранги обмениваются недостающими частями данных через высокоскоростную шину NVLink. Это позволяет каждому инференсному GPU собрать полный фрагмент без дублирования сетевого трафика. Если реплики находятся в разных доменах NVLink, одна копия каждого фрагмента пересылается между доменами, а затем происходит тот же сбалансированный локальный обмен. Сетевые передачи и обмен через NVLink конвейеризируются (pipelined), что позволяет скрыть задержки одного процесса за работой другого.
04Гибридный подход в NeMo RL
Хотя Reshard предлагает теоретическую эффективность, на практике реализация идеального решения для всех типов слоев трансформера может быть сложной из-за разнообразия конфигураций параллелизма. Команда NVIDIA разработала умный гибридный подход, интегрированный в фреймворк NVIDIA NeMo RL.
Анализ показал, что веса проекции feed-forward (FFN) в блоках трансформера доминируют в объеме данных refit, составляя около 95% всех параметров модели. Поэтому в гибридной схеме именно эти доминирующие веса передаются через эффективный путь Reshard. Остальные, более мелкие веса (например, проекции QKV или веса в архитектуре Mamba) передаются по традиционному пути gather-and-broadcast. Этот путь оптимизирован для широкой функциональной совместимости и простоты реализации.
Такой подход позволяет сочетать максимальную эффективность передачи основного объема данных с надежностью и универсальностью традиционных методов для остальных компонентов. Это обеспечивает как скорость, так и стабильность работы системы.
05Результаты бенчмарков: ускорение до 5,7 раз
Эффективность нового подхода была протестирована на реальных моделях и аппаратных платформах NVIDIA. Бенчмарки проводились на системах NVIDIA DGX H100 (NVL8) и NVIDIA DGX GB200 (NVL72) с использованием фреймворка NeMo RL. Результаты продемонстрировали впечатляющее ускорение.
Для модели Qwen 3 235B на DGX H100 время refit сократилось с 6,4 секунд (традиционный метод) до 1,9 секунд (Reshard), что дает ускорение в 3,3 раза. Для еще более крупной модели DeepSeek v3 671B на той же платформе время сократилось с 14,9 секунд до 2,6 секунд — ускорение почти в 5,7 раза. На новой платформе DGX GB200 результаты также показали значительное улучшение: для DeepSeek v3 время упало с 6,1 до 2,3 секунд, а для Nemotron Ultra v3 — с 8,2 до 2,3 секунд.
Особого внимания заслуживает тест на модели Qwen3.8 2.4T BF16. В микро-бенчмарке, измеряющем только время refit, новый метод справился с передачей данных за 4,3 секунды. Это доказывает, что технология Reshard масштабируется и эффективна даже для моделей с триллионами параметров, где традиционные методы были бы практически неприменимы из-за чрезмерных задержек.
06Что это значит на практике
Для исследователей и инженеров, работающих с большими языковыми моделями, внедрение Reshard collective означает переход от компромиссов к оптимизации. Раньше увеличение размера модели часто приводило к непропорциональному росту времени синхронизации, что ограничивало скорость итераций обучения. Теперь, благодаря возможности передавать только необходимые фрагменты данных, время refit становится предсказуемым и минимальным.
Это открывает новые возможности для:
- Более частых обновлений: Меньшее время refit позволяет обновлять веса чаще, что может улучшить качество обучения с подкреплением.
- Экономии ресурсов: Снижение нагрузки на сеть позволяет использовать пропускную способность для других задач или уменьшать количество необходимых узлов.
- Масштабирования: Как показывают тесты, метод работает эффективно даже для моделей с 2,4+ трлн параметров, что делает его актуальным для моделей следующего поколения.
Для тех, кто хочет попробовать эту технологию, NVIDIA предоставляет готовую реализацию в NeMo RL. Включение опции достаточно просто: в конфигурационном файле нужно установить параметр refit_transport: nccl_reshard и убедиться, что инференс и обучение не совмещены на одних и тех же GPU. Документация NeMo RL и NCCL Reshard подробно описывает поддерживаемые бэкенды, раскладки и режимы точности.
Reshard collective — это не просто еще одна оптимизация, это сдвиг парадигмы в том, как мы думаем о передаче данных в распределенных системах ИИ. Переход от передачи «всего» к передаче «необходимого» становится стандартом для работы с моделями будущего.
Источник: NVIDIA Developer ↗