一个月的活一周干完！英伟达世界模型训练速度飙升400%

新智元报道 【新智元导读】 英伟达世界动作模型 DreamZero 训练一次要烧 8 张 H100 整整 25 天，RLinf 从算子融合到 I/O 全链路系统级重构，把训练吞吐拉高近 4 倍——1 个月的活，1 周就能干完。 在通往 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/RLinf Hugging 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 Graph Python 层面的算子与调度开销往往是限制 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 模型在不改变原有mbs=1（此处 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 团队从底层夯实了工程底座，彻底修复了这些痛点，交付了一套健壮且高度可配的调优矩阵： 稳定适配 FSDP2 ：FSDP2 是 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 转化为一个高可靠、可量化的