近 4 倍训练吞吐加速，RLinf 让 DreamZero 训练时长从 1 个月缩短至 1 周！

📌 一句话摘要

RLinf 框架通过算子优化、FSDP2 并行策略、Recompute 和 I/O 管线重构，在 DreamZero 世界模型训练上实现近 4 倍吞吐加速，将 8 卡 H100 训练周期从 1 个月缩短至 1 周。

📝 详细摘要

本文详细介绍了无问芯穹与清华大学联合推出的大规模强化学习框架 RLinf 对英伟达 DreamZero 世界动作模型训练的系统级优化。DreamZero 作为基于 Diffusion 架构的 WAM 模型，在机器人任务上表现优异，但官方训练代码在 8 台 H100 上完整训练周期长达 25 天。RLinf 团队从三个维度进行了深度重构：算子/计算图优化（Torch Compile + CUDA Graph）使 5B 模型单步耗时从 1.8s 降至 1.2s；FSDP2 并行策略配合灵活的 Microbatch Size 和 Recompute 机制，在 5B 模型上进一步将吞吐提升 266%；视频数据管线改用 Torchcodec 解码，单样本解码时间节省 1.2s。实测结果显示，5B 模型训练吞吐从官方基线的 1.1 samples/sec/gpu 提升至 4.44 samples/sec/gpu，14B 模型实现 2.7 倍加速。收敛性验证表明，优化后的训练 Loss 曲线更平滑，18k Step 的 Checkpoint 在 LIBERO 仿真器上达到 96.68% 成功率。

💡 主要观点

- RLinf 通过算子融合和计算图固化实现 34%-50% 的加速。 Torch Compile 对 WanRMSNorm、adaLN-zero 等算子进行深度融合，CUDA Graph 消除 CausalWanSelfAttention 的 kernel launch 瓶颈，5B 模型单步耗时从 1.8s 降至 1.2s。

FSDP2 并行策略配合 Recompute 和灵活 Microbatch 设置释放显存和算力。 修复了官方 DeepSpeed ZeRO3 与 VAE 模块的兼容性冲突，支持任意 mbs 配置。5B 模型开启 Recompute 后 mbs 可从 2 提升至 32，算子效率大幅提升，吞吐再增 266%。

视频数据管线改用 Torchcodec 解码，缓解 I/O 瓶颈。 相比官方 PyAV 方案，Torchcodec 在保持解码速度的同时 CPU 占用更平稳，单视频解码时间缩短近 400ms，三视角场景累计节省 1.2s。

优化后的训练收敛质量不降反升，Loss 曲线更平滑。 RLinf 实现了 Episode 内部的 Step 粒度随机采样，有效平滑训练噪声。18k Step 的 Checkpoint 在 LIBERO 仿真器上达到 96.68% 成功率，验证了加速方案的可靠性。

💬 文章金句

- RLinf 对 DreamZero 的深度支持并非简单的参数微调，而是系统级的重构。

• 近 4 倍的吞吐提升，意味着算法研究人员在同等硬件资源下，可以将原本需要 1 个月的实验缩短至 1 周内完成。
• 通过向 FSDP2 训练后端的迁移，我们彻底解决了上述架构冲突与性能瓶颈。
• 在 DreamZero 多视角的训练场景下，视频解码时间累计节省了 1.2s。

📊 文章信息

AI 初评：83

来源：青稞AI

作者：青稞AI

分类：人工智能

语言：中文

阅读时间：23 分钟

字数：5517

标签： RLinf, DreamZero, 世界模型, 训练加速, FSDP2

阅读完整文章