✅ Обучаем триллионный MoE с prime‑rl 0.6: просто и эффективно

TL;DR

prime-rl 0.6.0 trains trillion‑parameter MoE models on agentic RL workloads.
GLM‑5 trained on SWE at 131 k sequence length, sub‑5‑minute steps, 28 H200 nodes.
Asynchronous RL separates trainer and inference for independent optimisation.
Inference uses FP8, Wide EP, P/D disaggregation, KV offloading and router replay.
Training uses 3‑D parallelism (FSDP, EP, CP) plus block‑scaled FP8.

Что такое prime‑rl 0.6.0?

prime‑rl — открытый фреймворк для асинхронного reinforcement learning. Он позволяет пост‑тренировать крупные открытые модели на агентных задачах. Версия 0.6.0 выводит поддержку триллион‑параметров MoE‑масштаба.

Пример модели из анонса — zai-org/GLM-5.1. Оптимизации работают и с другими крупными MoE‑моделями, например moonshotai/Kimi-K2.7-Code и nvidia/NVIDIA‑Nemotron‑3‑Ultra‑550B‑A55B‑BF16.

Запуск полного эксперимента GLM‑5.1 на кластере Slurm сводится к одной команде:

uv run rl @ examples/glm5_llmd/rl.toml --output-dir /shared/outputs/glm5-llmd

Роль асинхронного RL

В агентных задачах часто встречаются «длинные хвосты»‑выбросы: отдельные роллауты могут занимать часы. Ожидание их завершения перед обновлением политики просто простаивает GPU.

Асинхронный RL решает проблему, разъединяя trainer и inference. Они работают и масштабируются независимо, а политика обновляется сразу после завершения шага оптимизатора.

Синхронизация происходит только в точке обновления весов. Уже запущенные роллауты сохраняют кеш активного префикса, поэтому один роллаут может использовать токены из разных версий политики. Новые роллауты пересоздают кеш KV, даже если префиксы совпадают. Запросы от слишком старой политики отбрасываются; порог задаётся параметром max_off_policy_steps.

Оптимизации инференса

Инференс обычно является «бутылочным горлышком» в RL‑системе. prime‑rl ускоряет инференс, сохраняя контролируемую задержку.

FP8 инференс: низкая точность ускоряет prefill и decode, реализовано через ядра DeepEP и DeepGEMM.
Wide Expert Parallelism: распределяет эксперты по ≥32 GPU, совместно с крупным data‑parallel ранком (например, 32). Синхронизация происходит слой‑за‑слоем через dispatch/combine.
Disaggregation prefill/decoder: отдельные воркеры для prefill и decode снижают влияние длительных выводов инструментов на latency decode‑воркеров.
KV‑кеш: поддерживается tiered offloading на CPU и диск; централизованный пул RAM/диска (Mooncake Store) повышает конкуренцию.
Роутинг запросов: в комплекте идёт форк vllm‑router, также поддерживается NVIDIA Dynamo router.
Router replay (R3): фиксирует решения роутера инференса и воспроизводит их в trainer, уменьшая KL‑расхождение приблизительно в 10 раз.

Оптимизации обучения

Тренер построен на torchtitan и использует 3‑D параллелизм: FSDP, Expert Parallelism (EP) и Context Parallelism (CP).

Стратегия	Что шарит	Основное назначение	Ключевая деталь
FSDP (FSDP2)	Параметры, градиенты, состояния оптимизатора	Базовое уменьшение памяти	Грузит веса по‑запросу через `fully_shard`
Expert Parallelism (EP)	Эксперты внутри слоя	Сокращает активную память слоя	`all2all` dispatch/combine; torch‑native или DeepEP
Context Parallelism (CP)	Последовательность (dim)	Длинные контексты	Ulysses (по‑умолчанию) или Ring Attention

EP необходим, потому что после FSDP слои остаются огромными: у 78‑слойной модели с 800 B параметров в float32 один all‑gather весит ~40 GB. При EP=8 токены распределяются, а не собираются полностью, что экономит память.

CP важен при длине последовательности > 131 k, когда активные памяти превалируют над параметрами. Для GLM‑5 используется кастомный CP‑модуль, совместимый с DSA.

FP8 обучение: block‑scaled FP8 (DeepSeek V3) синхронен с инференсом, снижает KL‑расхождение и стабилизирует обучение.

Примеры применения

Длинные SWE‑агенты: обучение на реальных issue‑репозиториях; роллауты могут растягиваться на сотни ходов и вызовы инструментов. P/D disaggregation сохраняет предсказуемую latency decode‑воркеров.
Триллион‑параметров на небольшом кластере: запуск GLM‑5 помещается на 28 H200‑нодах благодаря Wide EP и KV offloading.
Стабильный агентный RL на масштабе: Router replay и FP8‑training снижают KL‑расхождение, что приводит к более надёжному обучению.