Проблема сложности и решение NVIDIA
Разработка агентного обучения с подкреплением (Agentic RL) часто упирается в сложность инфраструктуры: каждое изменение алгоритма требует правок в слоях трейнера, распределенного бэкенда и логики руллов. NVIDIA NeMo выпустила Molt — фреймворк на базе PyTorch, спроектированный так, чтобы его код был компактным и понятным как для человека, так и для AI-ассистентов. Это позволяет исследователям быстрее экспериментировать, не тратя время на дебаггинг огромных кодовых баз.
Компактность против конкурентов
Ключевая метрика Molt — размер кодовой базы. Фреймворк насчитывает всего около 8.6K строк RL-кода (измеряется через граф импорта от точки входа). Для сравнения, другие популярные решения значительно объемнее, что усложняет их модификацию:
| Фреймворк | Размер кода (строк) | Примечание |
|---|---|---|
| NVIDIA Molt | ~8.6K | PyTorch-native, агентный RL |
| OpenRLHF | 7.2K | Ближайший аналог по размеру |
| Slime | 25K | — |
| Verl | ~62K | Значительно сложнее для модификации |
Архитектура: Композиция, а не форки
Molt не форкает существующие библиотеки, а компону их через контейнерные зависимости. Это обеспечивает автоматическое получение апстрим-апдейтов. Архитектура состоит из трех основных компонентов:
- Ray: для размещения задач и асинхронных очередей.
- vLLM: для генерации руллов (rollout).
- NVIDIA AutoModel с FSDP2: для обучения модели.
Рантайм работает как пул агентов: vLLM-движения за роутером запросов и один обучаемый актор. Используется потоковый пул, который держит группы промптов в активном состоянии, чтобы двигатели не простаивали. При обновлении модели происходит частичная пауза руллов, прямая трансляция шардов актора через NCCL и возобновление запросов без их отбрасывания.
Гибкость агентов и корректность
Агент в Molt — это обычный Python-модуль с экспортом AgentRunner. Поддерживаются два режима:
- Env: фреймворк управляет циклом LLM в стиле Gymnasium (
step()). - ChatAgent: пользователь управляет циклом через стандартные SDK OpenAI или Anthropic. Molt запускает сервер-петлю, который транслирует протоколы, обеспечивая точное накопление токенов.
Фреймворк строго соблюдает инварианты корректности: «никогда не обучайся на токенах, которые вы не генерировали». Для MoE-моделей реализован rollout routing replay: vLLM возвращает ID экспертов для каждого токена, а тренировочный форвард воспроизводит их, избегая расхождений из-за численных ошибок.
Масштабируемость и доступность
Molt выпущен под лицензией Apache 2.0 и включает слурм-скрипты и готовые контейнеры. Рекомендуемая конфигурация для запуска — 2 узла с 8 GPU H100 (8 на обучение, 8 на руллы). Фреймворк поддерживает масштабирование от плотных моделей (4B) до MoE-моделей (700B) с помощью флага --fsdp.ep_size 256. Производительность сопоставима с стеком на базе Megatron, что делает Molt привлекательным для лабораторий с доступом к H100/H200 и стартапов, занимающихся пост-обучением.
Источник: MarkTechPost ↗