跳到主要内容
AI News HubLIVE
站内改写3 分钟阅读

全球最快的增量权重同步:770GB 权重同步从36秒缩短到15秒

文章摘要

Baseten 重新设计了增量权重同步路径,直接以 GPU 上的权重作为父版本,不再在 CPU 内存中保留完整检查点。GLM-5.3 的770GB权重同步时间从约36秒降至15秒,请求暂停从6秒缩短到0.8秒,吞吐量提升最高3.2倍,并解锁了此前受限于 CPU 内存的 Kimi K3 在 GB300 上的 RL rollout。

全球最快的增量权重同步:770GB 权重同步从36秒缩短到15秒
报告错误

纠错通道尚未开通,可先复制下方文章信息留存。

查看更正说明
直接读正文

Baseten 工程团队发布了其增量权重同步系统的最新进展:对于约770GB的GLM-5.3检查点,同步时间已从平均约36秒缩短至15秒,请求暂停时间从6秒降至0.8秒,吞吐量提升最高可达3.2倍。这一优化不仅刷新了性能记录,还解锁了此前因CPU内存限制而无法运行的Kimi K3强化学习rollout。

此前,增量权重同步系统需要依靠一个CPU副本作为解码基础。训练器将完整检查点和稀疏增量发布到对象存储,每个增量张量是父版本字节与新版本字节的XOR结果,再经zstd压缩。rollout集群通过Baseten Delivery Network拉取对象,重建新策略并原地更新GPU。然而,在旧设计中,每个副本必须在主机内存中保留完整的父检查点,以便将XOR字节解码为新值,并在每次更新后修补CPU副本。

这个CPU副本带来了三重代价。首先是内存压力:GLM-5.3的完整父版本约770GB,在8块B300的节点上,它与1.1TB的主机KV缓存无法共存,缓存不得不缩减到275GB。其次是时间成本:每次更新后,将CPU副本提升为新策略需要约45秒,阻塞下一个增量的处理。第三是双数据源风险:每次更新必须同时写入GPU和CPU副本,否则下一个增量将基于错误的父版本解码。

新设计的核心思想是:GPU上已经保存了策略N的每一个字节,只需要一种方法读取发生变化的那几个字节。启动时,系统从本地磁盘加载模型,按照引擎的实际GPU布局验证每个服务张量,记录每个张量的Adler-32校验和,然后删除本地副本。此后,GPU内存成为唯一的父版本。

当增量到达时,原生解码器直接遍历zstd命令,只保留非零的XOR字节:对于每个变化的元素,记录其位置和XOR掩码。它从不物化稠密的XOR张量,也从不读取模型权重。在旧策略继续服务的同时,每个GPU工作线程只收集变化位置上的旧值,应用XOR掩码得到绝对新值,用训练器的校验和验证结果,并将新值封存到引擎原生张量布局的写入计划中。

校验和验证也经过了专门优化。训练器发布每个变化张量新字节的Adler-32值。朴素做法是重新哈希整个张量,效率很低。由于Adler-32中每个字节的贡献仅取决于其位置和值,变化字节对校验和的影响可以在常数时间内计算。系统从启动时记录的校验和出发,为每个变化字节添加一次修正,再与训练器的值比较,未变化的字节永远不会被读取。这一优化使整个更新过程最多快2.3倍,例如Kimi K3在30GB增量上的更新时间从约230秒降至约100秒,同时确保任何损坏下载或错误父版本在任何GPU写入前被拒绝。

由于大量准备工作在暂停之前完成,暂停窗口只需执行三个步骤:关闭准入并原地暂停运行中的请求;原地写入新值并刷新派生权重;发布新策略版本并重新开放准入。在GLM-5.3上,暂停时间现为0.8秒,此前为6秒。45秒的CPU提升过程也已消失,下一个增量可以在落地后立即应用。

进一步优化解决了更新工作与在线服务共享GPU和CPU资源的问题。并行感知的增量解码器根据张量并行和专家并行将每个变化分配给负责的GPU,每个GPU将变化直接映射到其打包布局中的地址,包括swizzled FP8和FP4缩放因子。为避免与SGLang调度器的Python事件循环争抢GIL,更新计划在独立的无CUDA辅助进程中运行,通过共享固定内存与GPU工作线程交换批量数据。下载和解码也实现了流水线化,每个分片一落地就立即解码,运行在持续的、低优先级的原生线程上,并预留16个核心用于服务。

Kimi K3的案例尤为关键。其服务态MXFP4布局大小为1.7TB。如果保留CPU检查点,每个GB300节点仅其一半权重就需要约838GB主机内存,几乎耗尽整个服务容器的内存限制。由于GB300硬件预配置,这是一项硬约束。现在,Kimi K3可以运行在两个各四块GPU的GB300节点上,每个节点只解码、验证和暂存自己拥有的张量,绝不持有另一节点的权重。在两个节点都成功写入后,新策略才会发布并恢复请求。对于8.4GB的增量,Kimi K3从同步请求到权重生效仅需75秒,暂停4.1秒。

释放出的主机内存被重新用于KV缓存卸载。被驱逐的GPU KV页面写回主机内存而非丢弃,后续具有相同前缀的请求可以重新加载,而不是重新运行prefill。对于RL rollout的多轮长共享前缀场景,从CPU内存重新加载150k token的驱逐前缀需要0.48秒,而全新prefill需要4.9秒。在AgentX基准测试中,并发为16时KV缓存已完全适配GPU,两种方案无差异;并发为48时,代理前缀超出GPU内存但适合主机内存,卸载将未命中转为命中,缓存命中率从2%跃升至92%,每GPU输出吞吐量提升3.2倍。

所有结果均基于单副本、训练器形状的固定增量测量,从训练器的增量同步请求计到副本开始服务新策略,即以S3标记到第一个新响应计时。团队表示,希望将此方案用于其训练栈的用户可以查阅RL rollout文档或联系团队。

展开要点与分析

文章情报

工程师进阶

要点

  • 以 GPU 权重为父版本,消除约770GB的 CPU 副本,同步时间从36秒降到15秒,暂停从6秒降到0.8秒
  • 利用校验和的 O(1) 增量更新,只读取变化的字节,使整个更新过程最多快2.3倍
  • 释放的主机内存用于 KV 缓存卸载,在并发48时缓存命中率从2%提升到92%,吞吐量提升3.2倍
  • Kimi K3(1.7TB MXFP4)现可在两个 GB300 节点上运行,8.4GB 增量同步仅需75秒,暂停4.1秒

要点与分析由自动化流程生成,可能有误,请结合原始来源核实。