跳到正文
Prime Intellect Blog·· 2026-06-22精选AI 评分64

Prime Intellect 发布 prime-rl 0.6.0,支持万亿参数规模智能体 RL 训练

RL at 1T Scale: prime-rl Performance Deep Dive

AI 导读

Prime Intellect 发布 prime-rl 0.6.0,可在 28 个 H200 节点上以 131k 序列长度、256 rollouts 批量和低于 5 分钟的步时训练 GLM-5 完成 SWE 任务。

推荐理由

原文给出万亿参数 MoE 模型智能体 RL 训练的完整优化路径,可迁移到自建后训练栈的工程取舍。

正文 · AI 翻译

1T 规模下的强化学习:prime-rl 性能深度解析

今天我们发布 prime-rl 0.6.0 版本。该版本使我们(以及你)能够在高负载智能体工作负载下,以最高效率训练万亿参数规模的模型。我们一直在持续优化强化学习基础设施,以最大化大型 MoE 模型上的性能,降低在智能体工作流上对 OSS 模型进行后训练所需的成本、时间和痛苦。我们能够在仅 28 个 H200 节点上,以最高 131k 序列长度、低于 5 分钟的步长时间和 256 次 rollout 的批量大小,在 SWE 任务上训练 GLM-5。

prime-rl step time stays flat at 1T scale on GLM-5 SWE training

在本博客中,我们将介绍带来这些结果的所有优化,从低精度推理和训练,到预填充与解码分离的推理部署。我们将以 zai-org/GLM-5.1 作为模型示例,但我们的优化适用于任何大型混合专家模型,例如 moonshotai/Kimi-K2.7-Code、nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 等。

使用 prime-rl,可以在 Slurm 集群上通过单条命令运行 GLM-5.1 训练:

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

从第一性原理出发的智能体强化学习

prime-rl 从零开始构建,旨在支持高效的智能体后训练,拥抱异步强化学习。智能体任务通常存在长尾离群值;这些 rollout 可能长达数小时,尤其是长时程编码任务。如果等到这些 rollout 完成后再延迟策略更新,会降低 GPU 利用率并损害性能。异步强化学习通过允许推理策略在训练器部署上的优化器步骤完成后立即更新来解决这个问题。在异步强化学习中,训练器和推理是分离的,可以独立优化。

Asynchronous RL: disaggregated trainer and inference

训练器和推理之间存在一个固有的同步点——策略更新。每次优化器步骤后,rollout 策略会用新权重更新。在 prime-rl 中,我们一有新权重集就立即更新。为了不拖慢推理,已分派的 rollout 的活动前缀缓存不会被重置——这些 rollout 将由多个策略生成的 token 组成,KV 缓存也由多个版本产生。然而,新的 rollout 即使与旧 rollout 共享前缀,也会重新填充自己的 KV 缓存;我们使用 KV 缓存盐来强制这一点。最后,如果请求由过旧的策略生成,则会被丢弃;你可以通过 max_off_policy_steps 值来控制这一点。

In-flight weight updates during rollouts

从系统优化的角度来看,这些交互产生了一个有趣的问题:如何优化两个系统——训练器和推理——同时保持它们兼容。

在以下各节中,我们将剖析这两个系统以及我们所做的优化。

推理

推理是强化学习训练生命周期中的关键部分。模型在此与环境交互,产生被评估并分配奖励的 rollout。其中一些功能已存在于推理框架中;对于其他功能,我们与 vLLM 和 Dynamo 等框架紧密合作,目标只有一个:为社区提供最高性能的推理,并附带经过验证、易于使用的配方。

FP8 推理

推理吞吐量通常是你强化学习系统的瓶颈。推理吞吐量在很大程度上受益于预填充和解码部署中的低精度。我们大量利用 FP8 推理,并结合来自 DeepEP 和 DeepGEMM 的优化内核,以实现更低的延迟和更高的吞吐量。

宽专家并行

在关于推理性能的其他文章中,你可能注意到很多关注点都放在最小化延迟上,以便为用户实现最高的交互性。但 RL 并非如此——我们的首要关注点是最大化吞吐量,同时将延迟保持在一定的范围内(稍后会详细说明)。

实现这一目标的最佳配置之一是 Wide EP——大规模专家并行,通常跨越 ≥32 个 GPU。为了最大化吞吐量,我们将这一策略与大规模数据并行 rank 相结合——例如 32——创建一大组 GPU,每个 GPU 持有独立的专家,各自作为独立的端点。同步按层进行,分别在 dispatch 和 combine 操作中完成。

Wide expert parallelism across GPUs

Prefill 与 Decode 分离

Prefill 吞吐量是智能体 rollout 的一大瓶颈——某些模型↔环境组合产生的 prefill:decode token 比例高达 4:1。让同一批推理 worker 同时服务 prefill 和 decode 请求会增加端到端延迟,显著削弱 PipelineRL 的优势。

如前所述,RL 的优先事项是最大化推理吞吐量,而不是最小化延迟。然而,如果由于推理批次被 prefill 请求主导而导致延迟大幅增长,就可以观察到已完成的推理 rollout 出现“分组”现象,从而导致训练器与推理步骤的重叠度降低。

使用 prime-rl,你可以无缝地使用 P/D 分离。当 prefill 和 decode worker 分离后,较长的 prefill 请求(冗长的工具输出等)不会拖慢 decode worker,使其能够以可预测的延迟推进。这会更快地完成模型轮次,工具调用随后更快地到达沙箱执行,如此循环往复,有时会跨越数百个轮次。

Prefill and decode disaggregation

KV cache 管理

最大化吞吐量需要高并发。而这又需要大量的 KV cache 空间。如果空间不足,就可能出现 KV cache 抖动和低前缀缓存命中率,从而降低吞吐量。prime-rl 始终紧跟推理框架的最新特性,并端到端地支持它们——KV cache offloading 就是其中之一。

我们通过原生 vLLM offloading 和 Mooncake 支持将 KV cache 分层卸载到 CPU 和磁盘。有了更多的 KV cache 空间,我们就可以提高并发度,从而摊销更多的训练器成本。

这两种方法之间主要有 2 个区别:

  • vLLM 原生 offloading——一种简单的方法,为每个 worker(DP rank)创建一个单独的 CPU/磁盘池;只有该 worker 才能从这个缓存中加载。
  • Mooncake Store 则作为一个集中式存储运行,将所有客户端(节点)的 RAM/磁盘汇集到一个大池中,然后可以从任何节点的任何推理 worker 访问——这提供了显著的优势,尤其是在使用更复杂的路由策略时。

请求路由

为了将一切串联起来,推理请求需要被高效地路由,以实现高效的前缀复用、负载感知路由等。

prime-rl 中的默认路由选项是我们 fork 的 vllm-router,这是一个极简、轻量级的解决方案,以最小的配置开销提供强大的性能。根据你的需求,你可以选择自己的路由策略,针对负载均衡、KV cache 复用或其他目标进行优化。

我们还支持将 NVIDIA Dynamo 路由器作为即插即用的替代方案。这使我们能够为更大规模的运行开发和部署更复杂的路由策略。这些策略结合了不同因素,例如 KV 缓存复用、队列深度、KV 缓存利用率或当前负载,基于推理工作节点的实时指标,为每个工作节点计算一个分数。然后根据策略和分数选择工作节点。

结合 Mooncake Store(作为集中式 KV 缓存卸载层),这实现了跨副本的前缀缓存命中,同时公平地分配负载并对实时推理指标做出响应。

Request routing across inference workers

我们还在积极与 vLLM 和 Dynamo 团队合作,持续改进路由解决方案,为推理带来更多性能提升。

路由器重放

训练器↔推理不匹配会悄无声息地毁掉你的 RL 训练。为了解决这个问题,你可以使用路由器重放——即 prime-rl 中的 R3。其工作原理是捕获推理期间做出的路由决策,然后直接在训练器上重放这些决策。这有效地将训练器与推理之间的 KL 不匹配降低了一个数量级,从而实现更稳定的训练。

Router replay reduces KL mismatch between trainer and inference

这并非没有代价——在大规模部署中,路由专家数据可达数十 Gbps,给处理带来巨大压力。这让我们头疼不已,但现在 prime-rl 可以在处理这些数据的同时服务数千个并发的智能体 rollout。

路由专家是一个形状为 [num_layers, top_k, seq_len] 的大型负载,可以迅速增长到数百 GB,进而给 Python 处理带来巨大负载——即使是看似简单的操作,例如将响应转换为 Python 字典,也可能导致显著的事件循环延迟和 CPU 瓶颈。为了消除这种开销,prime-rl 将路由专家视为不透明负载,仅由高度优化的 PyTorch 操作进行处理,从而减轻 CPU 压力。

路由器重放与其他推理优化完全兼容,包括 P/D 分离,让你可以轻松部署生产就绪的堆栈。

训练

我们的训练器基于 torchtitan——一个高性能的 PyTorch 原生大规模训练代码库。我们从 torchtitan 改编了大量训练器代码,涵盖 FSDP、EP 等各种抽象,同时加入了自己的特色和改进。

并行

prime-rl 主要依赖三维并行,确切地说是:FSDP、CP 和 EP。每一种都有各自的用例、优点和缺点。要让大规模运行顺利进行,你需要不同程度地组合使用它们。在我们的 GLM-5 案例研究中,我们三种都用上了。

让我们简单回顾一下它们。

FSDP。完全分片数据并行是我们的基线分布策略。参数、梯度和优化器状态在数据并行 rank 之间分片,并在前向和反向传播期间按需收集。对于参数超过 1T 的模型,这是摊销完整优化器状态或参数内存占用的必要条件。我们使用 PyTorch 的 fully_shard(FSDP2)作为我们的 FSDP 实现;这使其可以轻松与其他策略组合,后文将进一步讨论。

Fully Sharded Data Parallel (FSDP)

专家并行(EP)。 即便使用了 FSDP,大型模型的层在 FSDP all-gather 之后仍然太大,无法有效放入单 GPU 的 HBM 中。对于 78 层、800B 参数且主权重为 float32 的模型,单层 all-gather 大约需要 (800B × 4) / 78 ≈ 40GB 的缓冲区。若 FSDP overlap 为 1 层,则仅活跃层权重就需要约 80GB 内存。

这正是 EP 的用武之地:我们不再 all-gather 整个层,而是设置一个单独的内部 EP 度,例如 EP=8,专家不会在该范围内被 gather。取而代之的是,token 会通过 all2all 原语进行 dispatch 和 combine。由于专家是层内存占用的主要来源,这显著降低了活跃内存。

在 prime-rl 中,我们允许两种独立的 EP 配置——torch 原生 all2all 和 DeepEP。根据我们的观察,torch 原生在单节点 EP 跨度(即 EP=8)上提供略好的吞吐量,但当 EP 跨节点时会显著下降。此时使用 DeepEP 会快得多。

Expert parallelism with all2all dispatch and combine

上下文并行(CP)。 在 131k+ 序列长度下,中间激活——而非参数——成为主要的内存开销。上下文并行将序列维度分片到各个 rank 上,以减少每 GPU 的激活内存。

在 prime-rl 中,我们为所有自定义模型支持上下文并行。我们支持两种主要使用上下文并行的方式:

  1. Ring Attention —— 在整个模型前向过程中,batch 按序列分片。当到达核心注意力时,每个 rank 持有其 Q、K 和 V 的分片,同时以环形模式处理来自其他 rank 的 K 和 V。
  2. Ulysses —— 与 ring attention 一样,数据在整个模型前向过程中按序列长度分片。当到达注意力时,all2all 操作将布局从序列分片切换为头分片,然后在头维度上计算注意力。注意力计算完成后,布局再次通过另一个 all2all 换回。这能很好地适用于大多数非标准注意力(线性注意力、Mamba 等),也是我们的默认方案。

然而,也有一些例外——其中之一是 GLM-5 中使用的 DSA。

对于无法同时用 Ulysses 和 Ring Attention 并行的注意力模型,我们编写自定义的上下文并行实现,GLM-5 就是其中之一。

我们的上下文并行实现保持序列分片并计算投影。之后,K 和 V 被 gather——这很廉价,因为它们被投影到潜在空间——以便 indexer 能看到完整序列。indexer 为全局序列计算稀疏索引,并基于这些索引计算核心注意力。由于 DSA 具有固定的 top_k,其成本也是固定的(除了 KV 的内存成本,如前所述可忽略不计)。

该方案允许每个注意力层仅进行一次 all-gather 集合通信,将成本保持在最低。

Custom context parallelism for GLM-5 DSA

GLM-5 DSA

为了高效计算 DSA,我们使用自定义 kernel,大量基于参考实现并根据我们的需求进行调整,同时提供快速的前向和反向。

FP8 训练

如前所述,训练器与推理之间的不匹配会损害你的训练。为了解决这个问题,我们使用 DeepGEMM 内核来执行块缩放 FP8,正如 DeepSeek V3 所提出的。与普遍看法相反,由于量化开销(特定配置除外),这实际上并不会真正提高吞吐量;然而,它大幅降低了训练器与推理之间的 KL 不匹配,因为两者现在使用相同的精度,在某些情况下甚至使用相同的内核。这反过来又带来了更稳定的训练。

FP8 training reduces trainer-inference KL mismatch

未来工作

我们继续探索提升 RL 引擎性能的其他方式,积极与其他框架合作——尤其是 vLLM、Dynamo 和 llm-d 以加速推理侧,或与 PyTorch 合作打造极速训练器——探索诸如推测解码、NVFP4 训练与推理等技术,或诸如容错、弹性伸缩以及超大模型训练器↔推理权重亚秒级传输等基础设施改进。

我们正在招聘!

在我们看来,大规模智能体 RL 是当今 AI 中最令人兴奋的系统挑战之一。构建高效的 RL 栈需要优化众多组件,既要单独优化,也要作为一个整体系统来优化:训练、推理、请求路由、权重广播、飞行中权重更新、环境、代码执行沙箱等等。

在大规模下,每一种开销来源都至关重要。成功来自于理解这些系统如何交互、识别瓶颈,并在整个栈中不懈地追求效率。

如果这听起来让你感兴趣,并且你希望日常工作涉及在大规模下构建和优化这些系统、试验新的分布式训练与推理技术,以及从数千块 GPU 中榨取最后一点性能,我们很乐意听到你的声音。

在此处申请,或在 X 上联系我们。

@article{primeintellect2026rlat1tscale,
author = {Matej Sirovatka and Prime Intellect Team},
title = {RL at 1T Scale: prime-rl Performance Deep Dive},
journal = {Prime Intellect Blog},
year = {2026},
month = {June},
note = {https://www.primeintellect.ai/blog/rl-at-1t-scale}
}

来源:Prime Intellect Blog · primeintellect.ai