论文

NeMo-DCR: 面向万亿参数规模可扩展 Agentic RL 的位精确增量压缩重拟合

NeMo-DCR: 面向万亿参数规模可扩展 Agentic RL 的位精确增量压缩重拟合

Agentic RL 将训练与 rollout 解耦,因此每次策略更新都必须在下个 batch 前送达 rollout 集群。在万亿参数规模下,同步一整个 checkpoint(refit)在两个 AWS 区域之间需 87.5 分钟。对 BF16 训练的测量显示,每步约 1% 的权重会改变其存储值。 已有系统虽利用了这一稀疏性,但在布局、精确性与效率上均有不足:它们重新实现布局规则、拼装完整张量、以算术方式重建数值,或使用跨集群集合通信,且都无法从 refit 中途失败中完整恢复。 NeMo-DCR(Delta-Compressed Refit)只发送变化量,同时保持位精确:接收端获得与稠密 refit 相同的参数与 buffer 位。 1. 布局:固定仿射映射将变化从训练 shard 投影到 checkpoint 的规范坐标,残差转换覆盖其余变化,serving 运行时的原生 loader 将所有变化放入接收端存储。 2. 精确性:可压缩 XOR mask 承载投影与 loader 均保留存储位的仿射变化,其余由覆写承载;接收端原地应用,重试覆写部分写入,联合提交将策略与基线绑定以便下一步 delta。 3. 效率:对象存储或中继树在 delta 构建期间流式传输负载,无需跨集群集合通信。 即便变化率达 3% 与 5%,30B-1T 模型的 NeMo-DCR refit 也比仅传输的完整 checkpoint 参考快 12-40 倍。1T 中继树在 3% 变化率下仅需 150 秒(而非 87.5 分钟),使跨集群 agentic RL 在万亿参数规模下的 refit 变得切实可行。

论文精读

TL;DR NeMo-DCR 利用稀疏权重变化与位精确 delta 压缩,将万亿参数 agentic RL 的跨集群 refit 从 87.5 分钟压缩到 150 秒,实现 12–40 倍加速。

问题

问题背景

Agentic RL 将训练与 rollout 解耦,每次策略更新后的权重同步(refit) 必须在下一次 batch 前推送到 rollout 集群,否则训练流程被阻塞。

现有方法局限

实测 BF16 训练中每步仅约 1% 的权重存储值发生变化。已有 delta 压缩系统利用了这一稀疏性,但存在四点具体不足:

  • 重新实现 placement 规则,与原生加载器不一致,导致接收端无法直接使用标准恢复路径;
  • 组装完整 tensor 后才能应用变化,内存开销大,难以扩展到万亿参数;
  • 通过算术重建值而非直接复制存储位,可能破坏位精确性;
  • 依赖跨集群 collective 同步,引入额外通信拓扑约束。 且现有系统均无法在 refit 中途故障后完整恢复,通常需要重传整个 checkpoint。

为什么这个问题难/重要

万亿参数规模下,跨区域传输全量 checkpoint 需 87.5 分钟,远超过单步训练时间,refit 成为系统瓶颈。实现 bit-exact delta 同步要求发送端准确识别权重存储位变化,接收端通过原生加载器原地应用,且任意一步失败都能回退到一致状态——这涉及 placement 映射、压缩编码、流水线传输与恢复协议的协同设计。Agentic RL 的规模化部署高度依赖高效的策略分发机制。

行业类比

类似 LLM 在线推理服务中,每次模型热更新只需推送 LoRA 增量而非全量权重,从而降低传输与加载延迟,提升迭代频率。

核心洞察

  • Bit-exact 增量压缩实现模型同步的“无损”与“高效”兼得。传统增量更新方法通常牺牲精度换取压缩率,但 NeMo-DCR 通过 XOR 掩码和 overwrite 编码,确保接收端权重与全量 checkpoint 逐位一致,同时只传输约 1%-5% 的变更,避免了近似误差在 RL 训练中累积导致策略漂移,对于万亿参数规模的 agentic RL 稳定性至关重要。
  • 将变化投影到 checkpoint canonical 坐标并复用 serving runtime 的 native loader,实现无缝部署与高效恢复。与重新实现 placement 规则或组装全张量的系统不同,NeMo-DCR 利用固定仿射映射将训练分片的变化直接映射到 checkpoint 布局,接收端通过原生加载器就地应用增量,无需额外内存拷贝,并支持中途失败重试和联合提交,保证策略与基线一致,使跨区域大规模 RL pipeline 具备生产级可靠性。

方法

输入与变更检测

NeMo-DCR 在每次 policy 更新后,从训练集群的 BF16 检查点中检测相对上一个版本发生存储值变化的权重。实测表明,在 BF16 训练中每步仅约 1% 的权重存储值发生变化。检测在训练分片的所有权粒度上进行,每个分片记录本地变化元素及对应版本号。

关键模块与处理流程

  1. 直接投影(Direct Projection)
    对于大部分变化元素,通过固定仿射映射将训练分片中的索引变换到检查点的规范坐标。该映射仅依赖静态的分片布局,无需跨分片通信,生成仿射变化记录。

  2. 残差转换(Residual Conversion)
    不满足直接投影条件的少量变化进入残差路径。训练侧按规范坐标重新组装对应的张量片段,计算出与上一版本的差异值,形成残差变化记录。两种记录合并为统一的 delta 流。

  3. 混合编码与压缩
    对仿射变化采用 XOR 掩码:比较新值与旧值,仅发送不同 bit 的位置。掩码本身高度稀疏,使用压缩算法(如 LZ4/ZSTD)进一步减小体积。对残差变化采用直接覆写,因为其模式不适合 XOR。这种混合策略让 delta 体积比纯逐值压缩小 1.7–2.2 倍。

  4. 可恢复的就地应用
    接收端(rollout 集群)通过原生模型加载器将 delta 应用到本地检查点存储中。XOR 掩码在原有参数缓冲区上按位异或,覆写记录直接替换对应位置。应用过程支持部分写入重试:若中途失败,重试仅覆盖未完成的分区。全部完成后执行联合提交,将新版本标记为 active,并绑定基础版本信息,作为下一次 delta 计算的基线。

  5. 流水线传输
    训练侧在 delta 构建过程中即开始通过对象存储或中继树向接收端推送已就绪的数据块,无需等待完整 delta 生成,也不引入跨集群全规约(collective)。多个 rollouts 集群可以通过中继树扩展,实现并行接收。

输出与差异点

最终输出是接收端与发送端位完全一致的全量检查点,但实际传输的数据量仅为变化部分。与同类 delta 系统相比,NeMo-DCR 同时满足位精确(bit-exact)、原生加载器放置、无跨集群集合以及中途失败可恢复,而此前系统至少缺失其中一项。

实验

实验设计

论文以 跨区域 Agentic RL 重构 为评估场景,用量化后的 BF16 检查点 构造 30B 至 1T 参数模型,模拟权重同步。通过统计每步训练中存储值发生变化的权重比例,验证稀疏性假设;再分别在 3% 和 5% 变化率下,对比 NeMo-DCR 与仅传输完整检查点的基线。重构路径覆盖 对象存储 与 中继树,并注入接收端故障以测试恢复能力。

关键发现

  • BF16 训练中仅约 1% 权重每步改变存储值,为 delta 压缩提供依据。
  • XOR 掩码 比残余 overwrite 编码小 1.7–2.2 倍,直接投影节省计算时间。
  • 在 30B–1T 范围,NeMo-DCR 重构比完整 checkpoint 传输快 12–40 倍;其中 1T 模型、3% 变化率、中继树传输仅需 150 秒,而基线为 87.5 分钟。
  • 在接收端被 kill 的情况下,训练曲线与密集 NCCL 参考一致,证明 bit-exact 与可恢复性。

与基线的深度解读

基线是只做传输的完整 checkpoint 方案,不包含任何稀疏性利用。NeMo-DCR 的加速不仅来自少传数据,更来自投影与原生加载器的结合:直接投影避免了在训练分片和 checkpoint 规范坐标之间做全量搬运,残余转换兜底非仿射变化,两者都不用跨集群 collective。这解释了为何即使在 5% 变化率下仍能保持显著加速,而同类 delta 系统往往因重新实现放置规则或重建数值而丢失 bit-exact 属性。对工程落地而言,该方案把 refit 从分钟级压缩到秒级,使万亿参数 Agentic RL 的跨集群策略同步变得可行。

行业影响

落地场景

NeMo-DCR 针对大规模 agentic RL 场景,训练与 rollout 集群分离,每次策略更新需快速同步权重到推理节点。该技术适用于需要高频策略迭代的在线学习产品:

  • 对话式 AI agent:如客服系统、虚拟助手,根据用户反馈快速调整策略,缩短生效时间。
  • 推荐系统:实时个性化推荐,基于在线探索数据频繁更新模型。
  • 自动驾驶仿真:大规模并行 rollout 需要最新策略网络,快速 refit 加速迭代。

商业价值

主要带来 降本 与 体验提升:

  • 降本:万亿参数模型全量 checkpoint 传输耗时 87.5 分钟,NeMo-DCR 在 3% 变化率下最快 150 秒,显著减少跨区域网络带宽成本,降低 rollout 集群闲置等待时间。
  • 体验提升:策略更新延迟从分钟级降到秒级,产品更快响应环境变化,提升用户满意度(如聊天机器人更快学习新知识)。
  • 增收:按 API 调用或推荐效果计费的系统,更快的模型迭代可提升点击率与转化率。

与现有工作流集成

NeMo-DCR 避免重新实现 placement 规则、不组装完整张量、不重建值,而是通过固定仿射映射将训练分片变化投影到 checkpoint 规范坐标,残差转换处理其余变化,并复用 serving runtime 原生 loader 应用更新。因此可与现有训练框架(如 NeMo、PyTorch)和对象存储 / relay tree 直接集成,无需修改模型代码。

  • 增量构建过程中通过对象存储或 relay tree 流式传输 payload,无需跨集群 collective,降低集成门槛。
  • 故障恢复通过 overwrite 重试与 joint commit 保证 bit-exact,可在现有 MLops 流水线中直接部署。

具体落地案例

  1. 电商智能客服:大型电商平台使用基于 LLM 的客服 agent,通过 RLHF 持续优化回答质量。每次在线用户反馈后需要更新策略,NeMo-DCR 使策略更新从几十分钟缩短到分钟级,客服 bot 快速适应新政策或产品信息,减少人工转接率。
  2. 内容平台推荐:短视频/新闻平台使用千亿参数推荐模型进行在线探索,rollout 集群分布在多个区域。利用 NeMo-DCR 的增量 refit,策略每周更新次数提升数倍,同时降低跨区域传输成本,直接提升推荐点击率与用户留存。

局限

  • **依赖训练并行布局的静态性**:NeMo-DCR 的直接投影要求训练分片与 checkpoint 规范坐标之间存在固定仿射映射,这假设了并行拓扑在 refit 周期内保持不变。在弹性训练、动态重分片或频繁调整张量并行/流水线并行度数的场景中,该映射可能失效,导致大量变化只能走残差转换路径,残差转换虽能保证正确性但压缩效率会显著下降。论文未验证在动态并行策略下的性能退化程度,也未讨论如何自动维护映射以适应常见的大规模训练弹性需求。
  • **变化率敏感且未覆盖极端情况**:性能优势高度依赖权重存储值变化的稀疏性,摘要指出 BF16 训练每步约 1% 权重变化。当学习率较高、训练初期或任务分布剧烈变化时,变化率可能远超 5%,此时 XOR 掩码的压缩收益快速衰减,传输量趋向全量 checkpoint,加速比可能大幅缩水。论文仅在 3% 和 5% 变化率下测试,未探索 10% 以上的场景,也未讨论对于其他精度(如 FP8、FP16)或需要同步优化器状态的扩展 refit 场景的适用性。
  • **耦合原生加载器与特定基础设施**:NeMo-DCR 依赖 serving runtime 的原生加载器实现就地应用和 bit-exact 恢复,这限制了它在非 NeMo 或自定义推理运行时上的直接部署。对象存储或中继树的传输方式可能引入额外的存储延迟和成本,特别是在跨区域高延迟网络中,对象存储的读写吞吐可能成为新瓶颈。联合提交机制要求多集群间协调一致,增加了系统的故障域和运维复杂度,论文附录也提及在某些异常情况下需要操作员介入,说明全自动恢复仍有边界。
论文Songlin Jiang2026-10-06原文

相关内容