随着大规模语言模型(LLM)和强化学习(RL)技术的不断发展,如何高效地完成训练参数的更新,成为业界关注的焦点。Kimi K2作为今年发布的前沿模型,凭借其独创的Checkpoint Engine解决方案,成功将1万亿参数规模模型的RL参数更新时间从传统的10分钟大幅缩短至约20秒,显著突破了强化学习训练效率的瓶颈。本文将全面揭示Kimi K2在参数同步优化方面的技术难题与创新手段,探讨其背后的架构设计、性能优化和工程权衡,为行业提供宝贵的实践经验和参考。 强化学习训练中,参数更新环节一直是制约整体性能的关键瓶颈。传统RL训练通常采用按轮次完成权重同步,训练与推理阶段需共享或快速切换参数。在不同架构设计中,常见的有共置(Colocate)和分离(Disaggregation)两种模式。
共置模式中训练和推理共享GPU资源并交替使用,分离模式则将训练与推理部署在独立的GPU设备上。无论哪种架构,高效完成训练模型的权重同步,避免GPU因等待参数更新而空闲,都是提升端到端训练性能的关键。Kimi K2团队基于公司内部实验经验,主要聚焦于共置场景优化,打造了独具特色的容器隔离部署方式,实现训练和推理各自运行于不同容器,但共用一台机器。这一设计既保证了两者环境以及镜像的完全解耦,又允许独立迭代开发,增强了系统的灵活性和稳定性。同时,也带来了进程组难以直接交互、参数同步复杂度加大的挑战。 针对训练与推理进程间权限和通信限制,Kimi K2创新设计了Checkpoint Engine作为中间层桥梁,专职负责参数的同步与更新,最大限度减少对现有训练及推理流程的侵入。
这种设计权衡了理论最优与工程可行性。理论上最优方案是让RL框架统一管理训练与推理的并行度,采用精细的重分片(reshard)策略来减少冗余传输,提高通信效率。但实际工程中为了避免改动训练推理的核心逻辑,Kimi K2选择了更为解耦、但数据传输更完全的方式 - - 每个设备均完整接收权重。考虑到H800/H20等GPU设备间通过RDMA网络或NVLink可提供超过100GiB/s的带宽,20秒内同步TB级权重成为完全可以接受且高效的解决方案,这体现了团队在性能与复杂度间的平衡。 Kimi K2在初期版本(k1.5)中,基于CUDA IPC技术完成Tensor数据传输,推理端通过接口接受GPU间共享的IPC句柄,将训练端传出的权重广播给推理端所有rank。由于GPU显存无法同时存放全部权重,模型权重通过按层、按专家(EP,Expert Parallel)分批载入CPU,逐步加载到GPU。
这种方案虽然满足了基础需求,但随着模型参数规模和GPU数量剧增,传输开销和内存管理中暴露的问题逐渐显现,参数更新时间最高达十分钟,严重影响训练效率。 深入剖析性能瓶颈后,团队发现分层、分专家粒度的频繁小批量异步广播造成了GPU内存使用的不稳定,甚至时常罹患CUDA显存溢出(OOM)风险。同时,每个张量对应一个IPC句柄的高频序列化和反序列化过程,极大增加了CPU和GPU间通信负担。vLLM推理引擎在处理IPC句柄时,还需汇集所有rank信息,进一步加重了数据传输负载。 为提升整体效率,Kimi K2提出了共享缓冲区与两阶段Pipeline机制。通过将分散零碎的小张量归并打包至固定大小的缓冲桶中,实现批量数据的广播,极大降低通信次数和内存碎片化风险。
该缓冲区自初始即作为vLLM的共享通道,减少每次更新时因IPC句柄传输带来的额外开销。改进后的方案还优化了IPC句柄收集逻辑,只在首次请求时汇集必要句柄,并取消了所有rank的全部收集,改为针对性收集,大幅简化操作。为了进一步减少通信和权重更新流程的串行阻塞,团队引入双缓冲区设计,构建双阶段流水线,广播和权重更新流程得以实现重叠执行。 这一系列改良将参数更新时间从超过10分钟优化至约2分钟,满足了团队当时的大规模H800 GPU集群强化学习需求。然而,理论上1TB权重文件仅依靠内存带宽和网络极限性能,完成同步同步应低于10秒。基于此,团队继续挑战性能极限,从主机到设备(H2D)传输速度和推理引擎权重更新细节中挖掘优化空间。
发现当前设计中Checkpoint Engine必须等待单个rank完成H2D操作后,才能发起广播,导致实际上带宽被单点主机到设备传输限制(40-50GiB/s)束缚,无法充分利用集群间RDMA或NVLink高达100GiB/s的网络带宽。vLLM推理引擎内部也存在性能隐患:每次更新均需动态计算指定参数字典,且专家权重频繁调用.item触发GPU与CPU的同步阻塞,导致更新速度不稳定。 为此,Kimi K2团队实施了三大关键优化措施。首先,通过重叠H2D和广播,实现主机到设备的传输和设备间广播的流水线并行,极大提升了整体传输吞吐。其次,缓存vLLM中dict(self.named_parameters()),避免重复CPU计算开销。最后,将专家映射表(expert_map)缓存至CPU端,规避了频繁的CUDA同步延迟。
尽管理想方案设想H2D和广播完全并行,但集群中的PCIe带宽争用使得二者相互影响,降低各自速率。Kimi K2创新亮点在于,让所有节点先并行完成H2D操作,聚合完成权重的设备内存数据后,再通过快速的设备到设备(D2D)操作完成广播准备。因D2D传输速度远高于PCIe,故仍能显著提升整体更新效率。 经过上述优化后,Kimi K2成功在16个或更大规模的H800集群中实现了1万亿参数模型约20秒内的参数同步,速度稳定且几乎不出现长尾延迟现象。 强化学习过程中,推理引擎偶发故障常导致训练中断,严重影响训练连续性和资源利用率。单纯重启推理服务通常需重新从检查点磁盘加载权重,耗时且易触发IO瓶颈。
Kimi K2大胆引入故障自愈机制,利用Mooncake Transfer Engine实现对失败节点的权重远程拉取和同步。该方案基于CPU RDMA直连,Rank0节点能够通过点对点(P2P)方式从远程机器的CPU内存快速读取全部权重,随后在本机内广播至故障机器内所有rank,完成权重快速更新,完成单点故障的自愈恢复,时间缩短至约40秒。 此外,针对大规模推理服务的启动延迟问题,Kimi K2带来了革新思路。传统做法通过将权重预热到/dev/shm内存文件系统,虽然加速了文件读取,但占用大量内存且依然存在等待时间。凭借Checkpoint Engine高速参数同步能力,团队设计了并行启动机制:推理引擎先以虚拟模式(dummy)快速启动,同时Checkpoint Engine后台完成权重注册,待推理引擎准备就绪时立即触发权重更新。在此过程中,参数同步工作与推理引擎的Torch Compile及CUDA Capture Graph等预热活动并行执行,实现了启动时间接近虚拟启动时长,大幅提升用户体验,现已在多项内部服务中大规模应用。
为了惠及更广泛社区,Kimi K2团队将解耦后的Checkpoint Engine核心组件ParameterServer进行了开源,提供了灵活且便捷的接口。用户可以轻松集成自定义的权重更新逻辑,与不同推理引擎对接,且无需关心底层NCCL组通讯细节。Checkpoint Engine支持广播和P2P两种参数同步模式,分别适用于高性能大规模同步和灵活容错的场景切换。同时,支持多检查点注册与切换,方便实现多版本模型并行验证与测试。 当前P2P同步方案尚处初级阶段,仍存在串行读取节点权重的提升空间,团队计划深耕重叠传输流水线设计,欢迎社区伙伴联合优化。开源进展同时得到了vLLM社区的积极响应,借助科学的接口设计,Checkpoint Engine与vLLM实现无缝集成,推动强化学习训练效率的进一步跃升。
回顾Kimi K2在强化学习参数更新领域的技术升级历程,不难发现其由内而外的工程思考和创新实践。精准权衡理论性能与工程复杂度,结合先进的硬件带宽特性与软件架构设计,协同优化传输管线和推理引擎负载,共同突破了1万亿参数同步的时间壁垒,助力实现更高效、更稳定的模型训练。未来,伴随硬件性能的持续提升及社群生态的共同促进,基于Checkpoint Engine的RL参数更新技术必将为行业带来更多创新可能,推动语言模型向更高维度蜕变。 。
