腾讯近期正式发布了基于Megatron-Core和SGLang/vLLM自主研发的大模型训练库——WeChat-YATT(Yet Another Transformer Trainer),内部代号为gCore。这一创新训练库专注于强化学习与多模态模型训练,旨在为开发者打造一套易扩展、简洁高效且高度可靠的大模型训练解决方案。WeChat-YATT通过定制化的并行计算策略,能够精准应对大尺寸模型、长序列输入及海量数据集等复杂训练场景,成功攻克了微信内部多个关键业务场景的技术难题,显著提升了大模型训练的整体效率。该工具为研究人员和开发者提供了灵活且可扩展的技术平台,有望推动多模态与强化学习领域的创新发展进程。
WeChat-YATT重点解决了大模型分布式训练过程中的两大核心技术瓶颈。首先是多模态场景下的可扩展性瓶颈问题。随着图像、视频等多模态数据规模的持续增长,传统架构中单一控制器管理数据的模式容易引发通信和内存瓶颈,导致系统吞吐量受限甚至训练流程异常中断。WeChat-YATT通过引入并行控制器(Parallel Controller)的并行管理机制,有效分散了系统压力,大幅提升了系统的可扩展性和稳定性,从而更好地应对多模态、大数据量的复杂应用场景需求。其次是动态采样与生成式奖励计算下的效率短板。在需要频繁动态采样或生成式奖励计算的训练流程中,模型频繁切换和”长尾”任务容易产生大量额外开销,导致GPU算力无法充分利用,严重影响整体训练效率。WeChat-YATT通过部分共存策略和异步交互机制,大幅减轻了模型切换损耗和长尾任务影响,实现了训练过程中的高吞吐量和高资源利用率,为大规模RLHF任务的高效迭代提供了有力支撑。
针对不同业务场景的差异化需求,WeChat-YATT支持全员共存和部分共存两种资源放置模式,以最大化提升集群资源利用率。全员共存模式采用串行调度机制,Actor Rollouts、GenRM(生成式奖励模型)与Train依次串行执行。每个角色完成任务后主动释放计算资源,系统随即加载下一个任务所需模型。该策略适用于绝大多数常规训练场景。值得注意的是,在每个阶段,相关组件均可独占全部GPU资源,这极大缩短了资源空闲”气泡”时间,显著提升总体训练吞吐量和效率。部分共存模式下,Actor Rollouts与GenRM独立部署,并通过异步方式进行高效交互。Actor训练阶段占用全部GPU资源,在Rollouts生成阶段释放GPU资源并唤醒Actor Rollouts及GenRM两大组件协同工作。系统通过动态负载评估进行资源分配与均衡,当Rollouts生成完毕后,这两者释放资源,Actor随之重新加载到GPU上进入下一轮训练流程。部分共存模式特别适合Rollouts与GenRM需要高频交互、动态采样的复杂任务场景。
WeChat-YATT还具备多项技术特色。在内存利用方面,项目采用并行控制器架构,有效降低了单节点内存消耗,更适合多模态场景下的大模型训练,提升了系统扩展性和稳定性。在GenRM支持方面,针对生成式奖励模型场景实现了不同的资源放置策略,用户可根据具体场景选择最优的训练方案。智能检查点策略是另一大亮点。WeChat-YATT支持异步Checkpoint保存,并针对微信业务场景特点,根据调度流程实现断点自动保存,进一步保障了训练安全性和高可用性。此外,系统在训练过程中实现了各个数据并行组间的负载均衡,有效减少资源空闲时间,显著提升整体训练吞吐量。
这一训练库的发布标志着腾讯在大模型技术基础设施建设方面取得重要进展,也为业界提供了处理复杂多模态训练场景的有效解决方案。WeChat-YATT凭借其创新的技术架构和实用的功能特性,有望成为多模态与强化学习领域的重要技术参考,推动大模型训练技术的持续进步与发展。