过去要等一个月的世界模型训练如今有机会压进一周完成。在通往 AGI 的道路上世界模型World Model被视为让 AI 真正理解并预测物理世界的关键拼图。英伟达近期重磅发布的世界动作模型WAM DreamZero 一经发布就在两项机器人基准测试 RoboArena 、MolmoSpaces 上双双登顶在具身智能领域获得极大关注。与传统 VLA 等模型不同WAM 将视频这一具备完整时空信息的载体当作自己的核心学习材料并以一种“先理解世界如何变化再决定自己如何行动”的模式使模型天然获得互联网视频所蕴含的海量物理经验。它不再需要大量重复演示来学习单一动作而是能从多样化的数据中学习世界的物理规律从而在从未见过的环境和任务中依然保持稳定执行能力。〓 当前最优的 VLA 模型与 DreamZero 世界模型在任务成功率、泛化性、跨本体等方面的直观对比上面的表格直观的展示出 DreamZero 模型相比开源最优的 VLA 模型 π0.5在任务成功率、任务泛化性、后训练对成功率的提升效果、以及跨真机本体的泛化性等方面具有明显的优势实现了超过 2x 的成功率提升。它的范式革新不仅大幅降低了学习成本也让机器人的形态适配与技能拓展不再受限于大量专属数据为多机型协同、快速部署与低成本迭代提供了可行路径。然而以 Diffusion 架构为主体的 WAM 多模态模型也给算力和显存带来了巨大的挑战。参考官方开源的 DreamZero 训练代码采用 8 台 H100 训练 24750 万帧数据完整训练周期长达 25 天高昂的训练成本和耗时成为行业复现的主要门槛。为助力前沿研究更高效地落地无问芯穹与清华大学等联合推出的大规模强化学习框架 RLinf 已正式上线了对 DreamZero 训练的深度支持。在实现功能适配的基础之上更进一步依托 RLinf 强大的底层系统优化能力对 DreamZero 的训练管线进行了深度的重构与加速。相比 DreamZero 官方提供的基线训练脚本RLinf 成功实现了近 4 倍的训练吞吐加速且具有更好的收敛效果。RLinf 是如何极致榨干 GPU 的每一滴算力达成 4 倍训练加速的接下来将为您一文拆解背后的核心优化思路与逻辑。代码链接https://github.com/RLinf/RLinfHugging Face链接https://huggingface.co/RLinf使用文档链接https://rlinf.readthedocs.io/zh-cn/latest/rst_source/examples/embodied/sft_dreamzero.html核心揭秘近 4 倍加速背后的 3 大优化维度为了打破官方脚本的性能瓶颈RLinf 系统优化团队从计算图、FSDP2并行优化与全局参数调优、数据处理管线进行了深度优化。极致的算子/计算图优化Torch Compile CUDA GraphPython 层面的算子与调度开销往往是限制 GPU 峰值性能的“隐形杀手”。在 RLinf 中我们深度融合了torch.compile和 CUDA Graph 技术Torch Compile通过底层编译优化对算子进行深度融合Kernel Fusion包括 WanRMSNorm、adaLN-zero 等 Diffusion 架构中的低效算子。CUDA Graph将计算图固化消除 GPU launch 的 CPU 调度瓶颈在DreamZero的训练中CausalWanSelfAttention 部分的 kernel launch 较为密集CUDA Graph 可以做到有效优化。通过该项优化技术DreamZero 5B 和 14B 模型在不改变原有 mbs1此处 mbs 指 mbs per gpu下同的配置下分别获得 50%从 1.8s/step 降到 1.2s/step和 34%从 9s/step 降到 6.7s/step的训练加速。计算与显存的联合优化解锁全方位性能调优支持任意 Microbatch Size、并行方式的参数调优以及 Recompute激活重计算是业界训练大模型时必不可少的性能调优手段。然而在 DreamZero 官方的 baseline 中存在着明显的工程局限例如默认使用 DeepSpeed 的 zero2 offload 并行方法、image encoder 不拼 batch 逐样本执行等大大降低了性能的调优空间。RLinf 团队从底层夯实了工程底座彻底修复了这些痛点交付了一套健壮且高度可配的调优矩阵稳定适配 FSDP2FSDP2 是 PyTorch 官方团队推出的最新 ZeRO 实现也是 RLinf 面向中等规模大模型的默认并行方案。此前在 DreamZero 官方代码中使用的 DeepSpeed 方案存在一定的局限性由于 ZeRO3 与 VAE 模块中 causal conv 的上下文维护机制存在兼容性冲突开发者往往被迫回退至性能较低的 ZeRO2 offload 模式。此外DeepSpeed 在反向传播阶段的 post backward hook 产生了较高的 CPU 侧开销制约了整体训练吞吐。通过向 FSDP2 训练后端的迁移我们彻底解决了上述架构冲突与性能瓶颈。用户现在可以根据显存配置需求在不同的分片策略间灵活切换确保训练过程的高效与稳定。灵活的 Microbatch 设置在 FSDP2 支持 DreamZero 模型训练的初始版本中Microbatch Size (mbs)、Recompute激活重计算与 FSDP2 的策略组合往往会触发复杂的底层计算图冲突而且 image encoder 不拼 batch 会吞掉一部分开大 mbs 的加速收益。RLinf 通过工程上的努力彻底解决了 mbs 1 时与上述特性共存的不兼容问题并且使得 image encoder 能够高效地拼 batch 执行。这一改进使训练系统具备了更高的灵活性用户可以不受约束地配置任意 mbs从而根据硬件资源的显存水位与计算吞吐需求进行精细化的参数调优在显存占用与执行效率之间达成更优的工程平衡。举例来说对 DreamZero 5B 模型的训练在不开启 Recompute 的情况下mbs 开到2相比于原来的 mbs 只能开到1单步耗时几乎没有变化1.2s/step 变到 1.3s/step吞吐增加 85%。Recompute 机制与加速算子的深度协同针对 PyTorch 原生框架在复杂并行策略下的兼容性局限RLinf 通过深度的底层工程优化实现了 Recompute激活重计算与 CUDA Graph、FSDP2 的稳定解耦与协同。这一改进将 Recompute 转化为一个高可靠、可量化的性能调优维度。在显存受限的硬件环境下系统能够以微小的计算耗时为代价换取显著的显存空间释放从而支持更大规模的并行任务大幅提升整体训练吞吐。在 DreamZero 5B 的训练中在不开启 Recompute 情况下单卡 mbs 只能开到2最佳速度约 1.2s/step即1.7 samples/sec/gpu有 Recompute 情况下单卡 mbs 开到 32 可获得 7.2 s/step即 4.4 samples/sec/gpu同等算力下吞吐提升 158%。可以看到开启 Recompute 使 mbs 得以大幅增加从而大大提升算子效率。通过以上 FSDP2、mbs、Recompute 的全局参数调优在 DreamZero 5B 模型训练上我们在第一项算子优化的基础上即 1.2 samples/sec/gpu将训练性能进一步提升了 266%达到 4.4 samples/sec/gpu。突破 I/O 吞吐瓶颈高效视频数据处理管线随着计算密度即上述两项优化的显著提升数据加载效率逐渐成为制约整体训练吞吐的新瓶颈。在 DreamZero 的训练实践中视频数据的解码与预处理过程极其消耗 CPU 资源。传统的方案如 PyAV在解码性能上难以支撑高频的吞吐需求而单纯通过增加dataset的num_workers来尝试“通过数量换速度”往往治标不治本——过多的数据读取进程会剧烈抢占 CPU 资源进而导致训练主线程的内核下发Kernel Launch出现延迟反而拖慢了 GPU 的执行节奏。为了在“解码速度”与“系统资源开销”之间寻找最优解RLinf 团队对主流的视频处理库进行了深度的性能 Benchmark虽然 Decord 在纯解码速度上略胜一筹但 Torchcodec 在保持同梯队性能的同时表现出了更优的 CPU 占用稳定性。这使得我们能够预留出足够的计算余量给训练主线程并支持开启更多的 num_workers 来并发处理数据。相比原生的 PyAV 方案单个视频的解码时间缩短了近 400ms。在 DreamZero 多视角左视角、右视角、腕部视角三个视频的训练场景下视频解码时间累计节省了 1.2s。这一 I/O 端的性能提升为后续进一步压榨 GPU 计算潜力提供了充足的数据“弹药”。性能实测从“能跑通”到“极致高效”的端到端跃迁为了验证上述多维优化的综合成效我们在 Droid 数据集单样本含左、右、腕部三个视角视频规格上对 DreamZero 不同规模的模型进行了严格的端到端测试。DreamZero-14B大参数量下的吞吐飞跃在 14B 大模型上由于显存压力巨大官方基线通常被迫采用 DeepSpeed ZeRO-offload 方案这导致了严重的计算/通信浪费与 CPU 换入换出开销。在 14B 模型上RLinf 相比原生 DeepSpeed 方案实现了2.7 倍的加速即便相比于未经优化的 FSDP2吞吐量也进一步提升了35%。DreamZero-5B算力密度的极致压榨对于 5B 中等规模模型RLinf 的优势在于能够通过高效率的重计算逻辑稳定开启更大的 Microbatch Size (mbs)并配合其他计算图调优彻底释放 GPU 算力。通过 RLinf 调优训练吞吐从官方代码的 1.1 samples/sec/gpu 飙升至4.44 samples/sec/gpu相比于有诸多限制的 FSDP2 Base 更是实现了惊人的5.84 倍性能飞跃。14B 模型和 5B 模型的单步时间和吞吐全程使用 8xH100 测试。其中 14B 模型使用 MBS1 和 GBS8 进行测试这是因为该模型中间维度较大使用 mbs1 也能取得较优的算子效率并掩盖 FSDP2 的通信开销。对于 5B 模型我们使用 GBS256。FSDP2 Base 版本由于一些 PyTorch 的 bug 不能开大 MBS 导致吞吐受限这主要是因为小 MBS 下算子效率不高、CPU 开销显著以及 FSDP2 通信不能被掩盖我们解决了这些问题并且取得了较大的吞吐增长。训练收敛效果测试追求速度更要保证精度在极致的性能优化之外确保训练的正确性与收敛稳定性是框架落地的基石。我们对 RLinf 版本的 DreamZero 进行了严格的收敛性验证。下图展示了 DreamZero 5B 模型在 LIBERO 数据集上的 Loss 曲线对比配置Global Batch Size 2568 卡 H100训练 38 小时。〓 橙线是 RLinf 训练曲线蓝线是 DreamZero 官方代码训练曲线两者在三类 loss 和 grad_norm 指标上的对比Loss 曲线对比分析图中橙线RLinf与蓝线官方 Baseline呈现一致的收敛趋势。值得注意的是官方代码在训练过程中 Loss 波动较为剧烈这源于其以 Episode 为粒度进行数据读取而 RLinf 通过底层重构实现了 Episode 内部的 Step 粒度随机采样有效平滑了训练过程中的噪声提升了梯度更新的稳定性。为了进一步验证性能优化的有效性我们对 RLinf 训练的 9k 到 21k Step 的 Checkpoint 在 LIBERO 仿真器的 Spatial Benchmark 上进行了端到端测评每个 Checkpoint 执行 512 条轨迹单 Episode 最大长度 480 步具体成功率如下表所示实验结果显示模型在 18k Step 处达到了最优的 96.68% 成功率证明 RLinf 在大幅缩短训练耗时的同时完全保持了模型原有的训练效果与收敛质量。模型权重获取https://huggingface.co/RLinf/RLinf-DreamZero-WAN2.2-5B-LIBERO-SFT-Step18000总结让世界模型迭代跑在“快车道”从算子融合到 I/O 调优从并行策略的纠偏到 mbs 自由度的释放RLinf 对 DreamZero 的深度支持并非简单的参数微调而是系统级的重构。近 4 倍的吞吐提升意味着算法研究人员在同等硬件资源下可以将原本需要 1 个月的实验缩短至 1 周内完成。RLinf 不仅仅是一个工具库更是具身智能领域高效迭代的加速器。想要亲身感受 4 倍提速的强大效能欢迎使用 RLinf 工具开启您的 DreamZero 世界模型训练加速之旅代码链接https://github.com/RLinf/RLinfHugging Face链接https://huggingface.co/RLinf使用文档链接https://rlinf.readthedocs.io/zh-cn/latest/rst_source/examples/embodied/sft_dreamzero.html若大家在使用 RLinf 时遇到任何问题与疑惑可扫描下方二维码加入交流答疑群随时在线探讨、咨询解惑。现在在「知乎」也能找到我们了进入知乎首页搜索「PaperWeekly」点击「关注」订阅我们的专栏吧·
