新智元報道

【新智元導讀】英偉達世界動作模型 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 架構中的低效算子。

• CUDAGraph：將計算圖固化，消除 GPU launch 的 CPU 調度瓶頸，在DreamZero的訓練中，CausalWanSelfAttention 部分的kernel launch較為密集，CUDA Graph 可以做到有效優化。

通過該項優化技術，DreamZero 5B 和 14B 模型在不改變原有mbs=1（此處 mbs 指 mbs pergpu，下同）的配置下分別獲得 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 轉化為一個高可靠、可量化的性能調優維度。在顯存受限的硬件環境下，系統能夠以微小的計算耗時為代價，換取顯著的顯存空間釋放，從而支持更大規模的并行任務，大幅提升整體訓練吞吐。在 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計算潛力提供了充足的數據「彈藥」。

02 性能實測

從「能跑通」到「極致高效」

的端到端躍遷

為了驗證上述多維優化的綜合成效，我們在Droid 數據集（單樣本含左、右、腕部三個視角，視頻規格 33 frames × 480 × 640）上，對 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模型使用MBS=1和GBS=8進行測試，這是因為該模型中間維度較大，使用 mbs=1 也能取得較優的算子效率并掩蓋FSDP2的通信開銷。對于 5B 模型，我們使用 GBS=256。FSDP2 Base 版本由于一些 PyTorch 的 bug 不能開大 MBS ，導致吞吐受限，這主要是因為小 MBS 下算子效率不高、CPU開銷顯著以及FSDP2 通信不能被掩蓋；我們解決了這些問題，并且取得了較大的吞吐增長。

訓練收斂效果測試：追求速度，更要保證精度

在極致的性能優化之外，確保訓練的正確性與收斂穩定性是框架落地的基石。

我們對 RLinf 版本的 DreamZero 進行了嚴格的收斂性驗證。

下圖展示了 DreamZero 5B 模型在 LIBERO 數據集上的 Loss 曲線對比（配置：LR = 1e-5，Global Batch Size = 256，8卡H100，訓練38小時）。

Loss 曲線對比分析：圖中橙線（RLinf）與藍線（官方 Baseline）呈現一致的收斂趨勢。值得注意的是，官方代碼在訓練過程中 Loss 波動較為劇烈，這源于其以 Episode 為粒度進行數據讀取；而 RLinf 通過底層重構，實現了 Episode 內部的 Step 粒度隨機采樣，有效平滑了訓練過程中的噪聲，提升了梯度更新的穩定性。

橙線是RLinf訓練曲線，藍線是DreamZero官方代碼訓練曲線，兩者在三類loss和grad_norm指標上的對比

為了進一步驗證性能優化的有效性，我們對 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

03 總結

讓世界模型迭代跑在「快車道」

從算子融合到 I/O 調優，從并行策略的糾偏到 mbs 自由度的釋放，RLinf 對 DreamZero 的深度支持并非簡單的參數微調，而是系統級的重構。

近 4 倍的吞吐提升，意味著算法研究人員在同等硬件資源下，可以將原本需要 1 個月的實驗縮短至 1 周內完成。RLinf 不僅僅是一個工具庫，更是具身智能領域高效迭代的加速器。

想要親身感受 4 倍提速的強大效能？ 歡迎使用 RLinf 工具，開啟您的 DreamZero 世界模型訓練加速之旅！

代碼鏈接：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

若大家在使用 RLinf 時遇到任何問題與疑惑，可掃描下方二維碼加入交流答疑群，隨時在線探討、咨詢解惑。