您需要了解的所有关于在GLM-5上执行强化学习的基础设施

https://t.co/pvevY6zYUD https://t.co/rhky5OvmMk

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

来源：https://www.primeintellect.ai/blog/rl-at-1t-scale

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

今天我们发布了 prime-rl 0.6.0 版本。该版本使我们（和您）能够在万亿参数规模的模型上，以最高效率处理高强度的智能体工作负载。我们一直在不懈地优化强化学习基础设施，以最大化大型MoE模型的性能，从而降低对智能体工作流进行开源模型后训练所需的成本、时间和痛苦。我们能够基于 GLM-5 在 SWE 任务上进行训练，最大序列长度达到 131k，每步时间低于 5 分钟，每次 rollout 的批大小为 256，仅需 28 个 H200 节点。

在 GLM-5 SWE 训练中，prime-rl 的步时在 1T 规模下保持平稳本文将涵盖实现这些结果的所有优化，从低精度推理和训练，到 prefill 和 decode 分离的推理部署。我们以 zai-org/GLM-5.1 (https://huggingface.co/zai-org/GLM-5.1) 作为模型示例，但我们的优化同样适用于任何大型混合专家模型，例如 moonshotai/Kimi-K2.7-Code (https://huggingface.co/moonshotai/Kimi-K2.7-Code)、nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 (https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) 等。

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

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

智能体强化学习的基本原则

prime-rl 从零开始构建，旨在高效进行智能体后训练，并采用异步强化学习。智能体任务通常存在长尾异常值；这些 rollout 可能需要几个小时，尤其是长周期编码任务。若等待这些 rollout 完成才更新策略，会导致 GPU 利用率不足并损害性能。异步强化学习通过允许推理策略在训练器部署上的优化器步骤完成后立即更新，解决了这一问题。在异步强化学习中，训练器和推理被解耦，可以独立优化。

异步强化学习：解耦的训练器和推理训练器和推理之间存在一个固有的同步点——策略更新。在每个优化器步骤之后，rollout 策略会用新权重进行更新。在 prime-rl 中，新权重一可用，我们就会更新它们。为了不拖慢推理速度，已调度的 rollout 的活跃前缀缓存不会被重置——这些 rollout 可能由多种策略生成的 token 组成，其 KV 缓存也可能由多个版本产生。但是，新的 rollout，即使与旧 rollout 共享前缀，也会重建自己的 KV 缓存；我们使用 KV 缓存盐值来强制执行这一点。最后，如果请求是由过于陈旧的策略生成的，则会被丢弃；您可以通过 max_off_policy_steps 值来控制这一点。

rollout 过程中的动态权重更新这些交互从系统优化角度带来了一个有趣的问题：如何优化两个系统——训练器和推理，同时保持它们的兼容性。

在接下来的章节中，我们将剖析这两个系统以及我们所做的优化。

推理

推理是强化学习训练生命周期中的关键环节。在这里，模型与其环境交互，产生 rollout，这些 rollout 被评估并赋予奖励。其中一些功能已存在于推理框架中；对于其他功能，我们与 vLLM 和 Dynamo 等框架密切合作，目标只有一个：为社区提供经过验证、易于使用的高性能推理方案。

FP8 推理

推理吞吐量通常是强化学习系统的瓶颈。推理吞吐量在很大程度上受益于 prefill 和 decode 部署中的低精度。我们广泛使用 FP8 推理，并结合来自 DeepEP (https://github.com/deepseek-ai/DeepEP) 和 DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) 的优化内核，以实现更低的延迟和更高的吞吐量。

广泛的专家并行

在其他关于推理性能的文章中，您可能会注意到很多重点放在最小化延迟上，以实现最高的用户交互性。但在强化学习中情况并非如此——我们的主要目标是最大化吞吐量，同时保持延迟在一定范围内（稍后会详细介绍）。

实现此目标的最佳配置之一是 Wide EP——大规模专家并行，通常跨越 ≥32 个 GPU。为了最大化吞吐量，我们将此策略与大的数据并行秩（例如 32）相结合，创建一个大的 GPU 组，每个 GPU 持有不同的专家，每个都充当独立的端点。同步在每层分别通过 dispatch 和 combine 操作进行。

跨 GPU 的广泛专家并行### Prefill 和 Decode 分离

Prefill 吞吐量是智能体 rollout 的一大瓶颈——某些模型↔环境组合会产生高达 4:1 的 prefill:decode token 比率。让相同的推理工作节点同时处理 prefill 和 decode 请求会增加端到端延迟，显著降低 PipelineRL (https://arxiv.org/abs/2509.19128) 的优势。

如前所述，强化学习的优先事项是最大化推理吞吐量，而非最小化延迟。然而，如果由于推理批次被 prefill 请求主导而导致延迟变大，则可以观察到完成的推理 rollout “分组“现象，从而导致训练器和推理步骤的重叠率降低。

使用 prime-rl，您可以无缝使用 P/D 分离。当 prefill 和 decode 工作节点分离时，长的 prefill 请求（例如冗长的工具输出等）不会阻塞 decode 工作节点，使其能够以可预测的延迟继续推进。这可以更快地完成模型回合，工具调用也会更快地到达沙箱执行，然后循环重复，有时会重复数百次。

Prefill 和 decode 分离### KV 缓存管理

最大化吞吐量需要高并发性。这反过来又需要大量的 KV 缓存空间。如果没有足够的空间，可能会导致 KV 缓存抖动和前缀缓存命中率低，从而降低吞吐量。prime-rl 紧跟推理框架的最新功能，支持其端到端集成——其中一个例子就是 KV 缓存卸载。

我们支持将层级式 KV 缓存卸载到 CPU 和磁盘，既可以使用原生的 vLLM 卸载，也可以使用 Mooncake (https://github.com/kvcache-ai/Mooncake/) 进行卸载。有了更多的 KV 缓存空间，我们可以提高并发性，从而更好地分摊训练器的成本。

每种方法之间的两个主要区别是：

• vLLM 原生卸载——一种简单的方法，为每个工作节点（DP rank）创建一个 CPU/磁盘池；只有该工作节点可以从此缓存加载。
• Mooncake Store 则充当一个集中式存储，将所有客户端（节点）的 RAM/磁盘汇聚到一个大型池中，然后任何节点上的任何推理工作节点都可以访问该池——这提供了显著的优势，尤其是在使用更复杂的路由策略时。

请求路由

要将所有组件连接起来，需要高效地路由推理请求，以实现高效的前缀重用、负载感知路由等。

prime-rl 中的默认路由选项是我们 fork 的 vllm-router，一个轻量级的解决方案，在最小的配置开销下提供强大的性能。根据您的需求，您可以选择针对负载均衡、KV 缓存重用或其他目标优化的自定义路由策略。

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

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

跨推理工作节点的请求路由我们还与 vLLM 和 Dynamo 团队积极合作，不断改进路由解决方案，为推理带来更高性能。

路由器重放

训练器↔推理不匹配会悄然破坏您的强化学习训练。为此，您可以使用路由器重放——prime-rl 中的 R3。其工作原理是捕获推理期间做出的路由决策，然后在训练器上直接重放。这有效地将训练器和推理之间的 KL 散度降低了一个数量级，从而带来更稳定的训练。

路由器重放降低了训练器与推理之间的 KL 散度这并非没有代价——在大规模部署中，路由专家数据可能达到每秒数十 Gbps，对处理造成巨大压力。这曾给我们带来很多麻烦，但现在 prime-rl 可以在处理这些数据的同时服务于数千个并发智能体 rollout。

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

路由器重放与包括 P/D 分离在内的其他推理优化完全兼容，使您能够轻松部署生产就绪的堆栈。

训练

我们的训练器基于 torchtitan (https://github.com/pytorch/torchtitan)——一个高性能、PyTorch 原生的大规模训练代码库。我们从 torchtitan 中改编了大量训练器代码，包括 FSDP、EP 和各种其他抽象功能，同时添加了我们自己的特色和改进。

并行策略

prime-rl 主要依赖 3 维并行，具体来说：FSDP、CP 和 EP。每种并行策略都有其各自的用例、优点和缺点。要让大规模运行顺利进行，您需要结合使用不同程度的各种策略。在我们的 GLM-5 案例研究中，我们使用了所有三种策略。

让我们对它们进行简要的回顾。

FSDP。全分片数据并行是我们的基准分布策略。参数、梯度和优化器状态在数据并行秩之间分片，并在前向和后向传递过程中按需收集。对于具有 1T+ 参数的模型，这是分摊完整优化器状态或参数内存占用所必需的。我们使用 PyTorch 的 fully_shard（FSDP2）作为我们的 FSDP 实现；这使其能够轻松与其他策略组合，如下文将进一步讨论的那样。

全分片数据并行 (FSDP)专家并行 (EP)。即使经过 FSDP，大型模型层仍然太大，无法有效地在 FSDP all-gather 后适配到单个 GPU 的 HBM 中。以 78 层、800B 参数和 float32 主权重为例，单层的 all-gather 大约需要 (800B × 4) / 78 ≈ 40GB 缓冲区。如果 FSDP 重叠 1 层，则仅活跃层权重就需要约 80GB 内存。

这就是 EP 的用武之地：我们不是对整个层进行 all-gather，而是设置一个独立的内部 EP 度数，例如 EP=8，在这些专家上不进行收集。相反，token 将使用 all2all 原语进行调度和组合。由于专家是层内存占用量的主要贡献者，这显著减少了活跃内存。

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

使用 all2all dispatch 和 combine 的专家并行上下文并行 (CP)。在 131k+ 序列长度下，中间激活值——而非参数——成为主导的内存成本。上下文并行在多个秩之间分片序列维度，以减少每个 GPU 的激活值。

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

1. Ring Attention——批次在整个模型前向传播过程中按序列分片。当到达核心注意力时，每个秩持有其 Q、K 和 V 的分片，同时以环形模式处理其他秩的 K 和 V。
2. Ulysses——与 ring attention 一样，数据在整个模型前向传播过程中按序列长度分片。当到达注意力时，all2all 操作将布局从序列分片切换为头分片，然后注意力在头维度上计算。注意力计算完成后，布局再次通过另一个 all2all 切换回来。这与大多数非标准注意力（线性注意力、Mamba 等）配合良好，是我们的默认选项。

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

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

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

这种方案每个注意力层只需要一个 all-gather 集合，从而将成本降至最低。

适用于 GLM-5 DSA 的自定义上下文并行### GLM-5 DSA

为了高效计算 DSA，我们使用自定义内核，这些内核基于参考实现并进行了大量调整，以满足我们的需求，同时提供快速的前向和反向传播。

FP8 训练

如前所述，训练器↔推理不匹配会损害您的训练。为了解决这个问题，我们使用 DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) 内核执行块缩放 FP8，正如 DeepSeek V3 所提出的那样。与普遍看法相反，由于量化开销，这实际上并不会提高吞吐量（除了特定配置）；然而，它极大地减少了训练器和推理之间的 KL 散度，因为两者现在使用相同的精度，甚至在某些情况下使用相同的内核。这反过来带来了更稳定的训练。

FP8 训练减少了训练器与推理之间的 KL 散度## 未来工作

我们将继续探索其他方法来提高强化学习引擎的性能，并积极与其他框架合作——特别是 vLLM、Dynamo 和 llm-d 以加速推理端，或与 PyTorch 合作以获得极速的训练器——探索诸如推测解码、NVFP4 训练和推理等技术，以及诸如容错、弹性伸缩和针对大型模型的亚秒级训练器↔推理权重传输等基础设施改进。

我们正在招聘！

在我们看来，大规模智能体强化学习是当今人工智能中最令人兴奋的系统挑战之一。构建高效的强化学习堆栈需要优化广泛的组件，无论是单独优化还是作为统一系统优化：训练、推理、请求路由、权重广播、动态权重更新、环境、代码执行沙箱等等。

在大规模下，每一种开销都很重要。