WeChat-YATT：微信强化学习大模型训练库

推荐语

微信团队开源WeChat-YATT训练库，突破多模态与强化学习训练瓶颈，大幅提升大模型训练效率。

核心内容：
1. 解决多模态场景下的可扩展性瓶颈
2. 优化动态采样与生成式奖励下的训练效率
3. 支持两种资源放置模式提升集群利用率

杨芳贤

53AI创始人/腾讯云(TVP)最具价值专家

项目背景

我们基于 Megatron-Core 和 SGLang/vLLM 研发了大模型训练库 WeChat-YATT（YATT,Yet Another Transformer Trainer），内部项目名为 gCore，专注于强化学习和多模态模型的训练，旨在提供易扩展、简洁、高效、可靠的大模型训练能力。通过定制化的并行计算策略，我们的训练库能够处理大尺寸模型、长序列输入和大数据集场景，解决了微信中多个实际场景的痛点问题，显著提升了业务训练大模型的效率。此工具为研究人员和开发者提供了灵活且可扩展的解决方案，以推动多模态和强化学习领域的创新发展。

我们提出的 WeChat-YATT 训练库，解决了大模型分布式训练过程中面临的两大核心痛点：

● 多模态场景下的可扩展性瓶颈：随着多模态数据（如图像、视频）规模的不断增长，传统架构中由 SingleController 进行数据管理，容易成为通讯和内存的瓶颈，导致系统吞吐量受限，甚至引发训练流程异常中断。WeChat-YATT 通过引入 Parallel Controller 的并行管理机制，有效分散压力，大幅提升系统的可扩展性和稳定性，更好地应对多模态、大数据量的复杂场景。

● 动态采样与生成式奖励下的效率短板：在需频繁动态采样或生成式奖励计算的训练流程中，模型频繁切换和“长尾”任务容易引发大量额外开销，导致无法充分利用 GPU 算力，影响整体训练效率。WeChat-YATT 通过部分共存策略和异步交互，大幅度减轻模型切换损耗和长尾任务影响，实现了训练过程中的高吞吐量和高资源利用，更好地支撑大规模 RLHF 任务的高效迭代。

架构介绍

WeChat-YATT 针对不同业务场景，支持了两种资源放置模式：全员共存与部分共存，以最大化提升集群的资源利用率。通过灵活的调度策略，WeChat-YATT 能够有效适应不同的训练需求和计算环境。

与此同时，WeChat-YATT 采用了 Parallel Controller 模式，由多个 Controller 协同管理数据任务，显著降低了单节点的内存压力，尤其为多模态训练场景提供了更优的系统支持，相较于传统的 Single Controller 架构具备更强的可靠性。

● 全员共存模式采用串行调度机制，Actor Rollouts、GenRM(Generative Reward Model) 与 Train 依次串行执行。每个角色完成任务后主动释放计算资源，系统加载下一个任务所需模型。该策略适配绝大多数常规训练场景。值得一提的是，在每个阶段，相关组件均可独占全部 GPU 资源，这极大缩短了资源空闲“气泡”时间，显著提升总体训练吞吐量和效率。

● 部分共存模式下，Actor Rollouts 与 GenRM 独立部署，并通过异步方式进行高效交互。Actor 训练阶段会占用全部 GPU 资源，在 Rollouts 生成阶段，Actor 将 GPU 资源释放并唤醒 Actor Rollouts 及 GenRM 两大组件协同工作。并通过动态的负载评估，进行资源分配与均衡。当 Rollouts 生成完毕，这两者会释放资源，Actor 随之加载到 GPU 上，进入下一轮训练流程。部分共存模式非常适合 Rollouts 与 GenRM 需要高频交互、动态采样的任务场景。

多元的资源放置模式和灵活的调度机制，使 WeChat-YATT 在复杂多变的实际环境下都能实现资源的高效利用，助力大模型在微信内部多个场景的应用落地。

项目特点

● 高效内存利用：项目采用 Parallel Controller，有效降低了单节点的内存消耗，更适合多模态场景下的大模型训练，提升了系统的扩展性和稳定性。

● GenRM 高效支持：对于 GenRM 场景实现了不同资源放置策略，供使用者根据场景进行高效训练。

● 智能 Checkpoint 策略： WeChat-YATT 支持异步 Checkpoint 保存，并针对微信业务场景，根据调度流程，实现断点自动保存，进一步保障训练安全与高可用性。

● 负载均衡优化： 在训练过程中，WeChat-YATT 实现了各个数据并行组间的负载均衡，有效减少资源空闲时间，显著提升整体训练吞吐量。

实验效果

图1.随着训练进行训练时间的变化

图2.训练过程中每个阶段的用时

在 GenRM 任务场景下，我们对 WeChat-YATT 与开源训练框架 VeRL 进行了实验对比。WeChat-YATT 采用全员共存架构，在给定 GPU 数量的前提下，其他实验参数严格按照相关论文设置。实验结果如图1所示，WeChat-YATT 相较于 VeRL，整体平均训练时间缩短约 60%。

训练效率的大幅提升，主要得益于 WeChat-YATT 独特的全员共存架构：所有关键组件能够高效共享计算资源，使每一角色在其所处阶段均能充分利用全部硬件算力，极大提升了资源使用率。进一步观察图2的分阶段耗时分解，可以发现 WeChat-YATT 在各阶段的训练速度均较 VeRL 快 50% 以上。相比之下，VeRL 采用外置 GenRM 策略，在除 GenRM reward 阶段之外，GenRM 所占据的资源处于空闲状态，导致计算资源利用效率低下。

综上，WeChat-YATT 在 GenRM 场景下展现出卓越的训练效率，为复杂任务提供了更快速、资源利用更充分的分布式训练方案。