NVIDIA的Cosmos 3火了，但企業如何真正用起來？

來源 | 百度智能云技術站

經過百度百舸的 AI Infra 工程優化，Cosmos3-Nano-Policy-DROID 的訓練啟動速度提升 89 倍，國內主流 GPU 型號單機吞吐提升 99.3%，MFU 達到 0.42，顯著優于官方基準；12 節點擴展效率達 98.3%，集群算力得到充分釋放。

從開源社區到生產環境：AI Infra 工程優化的必答題

無論是開源社區的模型，還是企業自研的模型，其核心目標都是算法實現與驗證，而非生產環境下的性能最優。從「能跑」到「好用」，中間橫著一道工程門檻，AI Infra 工程優化是回答這道必答題的關鍵。NVIDIA 最新開源的 Cosmos 3 系列恰好印證了這一現狀。

Cosmos 3 作為全模態的世界模型，能夠統一建模語言、圖像、視頻、音頻以及動作序列等多種模態信息，在多個具身智能 Benchmark 上取得領先結果，已成為當前具身智能后訓練的重要基座模型之一。Cosmos3-Nano-Policy-DROID 作為其在機器人方向的代表模型，更是受到了社區的廣泛關注。

論文地址：https://arxiv.org/pdf/2606.02800

但即便是這樣前沿的模型，在直接部署到生產環境時，依然存在巨大的優化空間。完成性能優化只是挑戰之一，算力平臺差異帶來的資源約束同樣不容忽視。

在官方論文中，Cosmos 3 的訓練運行于由 1024 張 NVIDIA GB200 GPU 組成的超大規模集群環境，并依賴 NVLink 與 HPN 實現高效擴展。對于國內的企業和研究機構而言，這樣的訓練環境并不現實。如何在「通用」的 AI 算力環境下充分釋放 Cosmos 3 的訓練性能，成為模型落地過程中的關鍵挑戰。

Cosmos3-Nano-Policy-DROID 模型訓練優化實踐

百度百舸始終立足國內現有的算力供給格局，通過系統性的 AI Infra 優化與高擴展性集群架構設計，致力于為各類具身智能場景找到高性價比的算力解決方案。

本文以 Cosmos3-Nano-Policy-DROID 為例，基于百度百舸平臺 hpas.lgn7ib 實例（搭載國內主流 GPU 型號），介紹我們的 AI Infra 訓練優化實踐。圍繞數據加載、I/O 流水線、顯存利用、編譯優化以及多機擴展等關鍵環節，我們系統性解決了社區版本在生產環境中的性能瓶頸，并結合該 GPU 特性進一步釋放了模型訓練性能，為具身智能模型在「非旗艦」算力環境下的高效訓練提供了可復用的工程實踐參考。

全鏈路優化效果如下：訓練啟動速度提升 89 倍，單機吞吐提升 99.3%，MFU 達到 0.42，顯著優于官方基準；基于彈性 RDMA 網絡，12 節點擴展效率達 98.3%，集群算力得到充分釋放。

2.1. 任務啟動優化：37 分鐘 OOM → 25 秒正常啟動

具身智能模型通常依賴海量、高維的多模態數據集（如視頻、動作軌跡）。社區開源代碼在數據加載模塊往往采用較為通用的串行或淺層并行邏輯，未針對超大規模、高并發場景下的內存生命周期進行精細化管理，這在處理百兆級幀數的數據集時，極易引發峰值內存溢出。

在使用社區代碼對 Cosmos3-DROID 數據集進行訓練時，我們觀測到訓練進程啟動后長達半小時內無任何日志輸出。通過系統監控發現，CPU 內存持續異常飆升，峰值達到 1734 GB 并被 Linux OOM Killer 強制殺死，導致 8 卡訓練任務直接失敗。

深入剖析社區代碼數據流，定位到 ActionBaseDataset 存在大量重復、冗余的字段讀取。我們結合 Parquet 列式存儲的列裁剪能力進行針對性優化，剔除冗余字段，并重構數據拷貝路徑，從根源上消除內存膨脹風險。

指標

社區代碼

優化結果

指標變化

啟動結果

OOM Killed（無法訓練）

訓練正常啟動

Init 時間

37.2 min

25 s

89 × 加速

峰值系統內存

1,734 GB (OOM)

46 GB

-97%

精度影響

無（數據流路徑不變）

2.2. 突破 I/O 吞吐瓶頸：消滅 GPU 空閑等待間隙

在深度學習訓練架構中，GPU 算力與 CPU 數據預處理速度之間存在著天然的「速度差」。當任務遷移至更強 GPU 或調大 Batch Size 時，DataLoader 需要消耗成倍的 CPU 資源才能喂飽 GPU。一旦 CPU 成為瓶頸，不僅會造成 GPU 算力閑置，甚至可能因 CPU 資源被耗盡而反向拖累 GPU 進程（如影響 CPU Kernel Launch 效率）。

在 Cosmos 3 的實測中，我們發現 GPU 存在明顯的空閑等待間隙（Idle Gap）。社區默認配置中 max_samples_per_batch > prefetch_capacity，導致數據供給斷層。我們嘗試通過調大 DataLoader Worker 數量來緩解，卻發現引發了 CPU 資源擠兌和內存告警。通過 Profiling 深度剖析 CPU 端耗時，定位到根本瓶頸：ColorJitter（圖像數據增強）耗時占比高達 78.5%。

子操作

耗時

占比

ColorJitter（優化項）

1.68 s

78.5%

Video decode（3 路 MP4 視頻流）

0.31 s

15%

Transform / tokenize

0.10 s

5%

Downsample + concat

0.02 s

1%

RandomCrop + Resize

0.01 s

0.5%

優化前的耗時占比剖析

基于此，我們放棄了單純堆疊 CPU DataLoader Worker 的治標方案，而是將高耗時的 ColorJitter 算子從 CPU 遷移至 GPU 執行，徹底釋放 CPU 壓力，打通 I/O 瓶頸。

經過優化，ColorJitter 耗時從 1.68 s 縮短為 0.08 s，每個樣本數據處理總時間從 2.12 s 大幅下降至 0.52 s，徹底消滅 GPU 等待間隙，訓練吞吐從 1.35 sample/s/GPU 提升至 2.03 sample/s/GPU，提升 50%。

2.3. torch.compile 適配國內主流 GPU：解鎖 28.6% 加速

torch.compile 通過將 PyTorch 動態圖編譯為優化后的 Triton/CUDA Kernel，利用算子融合大幅減少 Kernel Launch 開銷和顯存帶寬瓶頸。社區代碼已初步支持 VAE 和 MoT 模塊的 Compile，但在不同架構的 GPU 上，其生成的底層 Kernel 往往需要針對性的適配才能發揮最大效能。

在 hpas.lgn7ib 實例上啟動任務時，系統直接拋出異常，這表明社區默認的算子融合策略超出了該 GPU 硬件的 Shared Memory 資源限制，導致 Compile 功能失效。

根據 torch.compile 的算子融合原理，結合這款 GPU 硬件的底層資源規格，我們通過配置調整禁用 mix-order reduction 策略，成功解鎖 compile 加速。

經過優化，訓練吞吐從 2.03 sample/s/GPU 提升至 2.61 sample/s/GPU，提升 28.6%。

2.4. 分層 Activation Checkpointing：把空閑顯存「兌換」為更高吞吐

Activation Checkpointing（AC）是緩解大模型顯存壓力的經典技術，其核心邏輯是「以計算換顯存」。社區代碼為了追求極致的顯存安全，默認對 Cosmos3-Nano 全部 36 層 Transformer 開啟了 Full AC。常規的優化思路是：利用 Full AC 節省下的顯存去調大 Batch Size（BS），從而尋找整機吞吐的最優解。

在搜索整機最大吞吐 BS 配置的過程中，我們通過硬件監控發現，即使 BS 已經調至當前硬件的極限，GPU 顯存依然存在大量空閑冗余，硬件的絕對算力并未被完全榨干。

在社區原生支持的 Full / Selective AC 基礎上，我們針對 hpas.lgn7ib 規格中 GPU 的顯存容量，進一步開發了分層 AC（Layer-wise AC）策略。通過精細化控制不同 Transformer 層的 AC 開啟狀態，在「顯存占用」與「重算開銷」之間找到更優的平衡點，將閑置顯存轉化為有效吞吐，從 2.61 sample/s/GPU 提升至 2.69 sample/s/GPU，提升 3.1%。

AC 模式

策略說明

vs full (baseline)

full

全 36 層開啟 AC

baseline

selective

層內部分算子關閉 AC

+0.1%（≈無效）

Layer-wise AC

部分層關閉 AC

+3.1%

2.5. 多機集群擴展：12 節點 98.3% 擴展效率

在通用 AI 算力場景下，非 NVLink、非 HPN 的 GPU 單機性能有限，要想加快模型訓練速度，多機協作擴展是核心能力。

我們依托百度百舸的極致 ERI（彈性 RDMA 互聯）網絡性能，配合 HSDP 并行策略并對計算通信 Overlap 做精細化調優，在 12 節點 hpas.lgn7ib 實例（96 卡）上達到 98.3% Scaling Efficiency，支撐客戶在訓練資源布局上的多樣化選擇。

節點數

GPU 數

global sps

擴展效率

1

8

100.0%

2

16

99.6%

4

32

98.7%

8

64

98.4%

12

96

98.3%

2.6. 精度驗證：全鏈路極致優化不損收斂

所有優化均為精確等價變換或不影響數據流的 Pipeline 優化，理論上不改變訓練精度。

Loss 曲線對比分析如下：實測圖中藍線（優化后版本）與紅線（官方 Baseline）呈現完全一致的收斂趨勢，驗證了工程優化的無損性。

2.7. 算力密度的極致壓榨：MFU 達 0.42，超越 B200 官方基準

經過數據加載、I/O 流水線、Activation Checkpointing、torch.compile 等一系列優化后，訓練系統整體效率得到顯著提升，在百度百舸 hpas.lgn7ib 實例上的 MFU 達到了 0.42，超過了官方論文中 GB200 基準為 0.23-0.3。

下圖為 hpas.lgn7ib 實例測試期間的 GPU 硬件監控指標：

• GPU SM 激活時間占比（DCGM_PROF_SM_ACTIVE）：表示在一個時間間隔內，至少一個線程束在一個 SM 上處于 Active 的時間占比，指標表現極高。

• GPU Tensor Pipe 處于激活的周期分數（DCGM_PROF_PIPE_TENSOR_ACTIVE）：表示 Tensor（HMMA/IMMA）Pipe 處于 Active 狀態的周期分數，算力得到充分釋放。

結語

具身智能的大規模落地，既不是單點算法突破能夠完成的，也不是單純堆砌算力就能解決的。它需要模型結構創新、AI Infra 工程優化與算力網絡能力的協同推進。

目前，除了本文介紹的 Cosmos3 之外，百度百舸面向具身智能場景，已經為客戶交付了 OpenPI、DreamZero、GR00T N1.6、Lingbot-VLA、Motus、AHA-WAM、FastWAM、Wan2.2 等眾多模型的訓推加速成果。

Cosmos3 的優化實踐進一步證明：開源社區負責推動算法創新，而真正決定模型能否進入生產環境、能否以合理成本完成訓練的，則是 AI Infra 的工程優化能力。

事實證明，使用國內主流 GPU，同樣能夠實現接近旗艦算力平臺的訓練效率，滿足實際業務需求。本文介紹的訓練優化，只是釋放 Cosmos3 模型性能的第一步。在集群層面，百度百舸已開展更深入的探索，通過計算資源與模型結構的協同設計，進一步釋放集群潛力，降低用戶訓練成本。

關于 Cosmos 3 集群訓練優化、模型評測及更多具身智能 AI Infra 工程實踐，歡迎關注百度智能云技術站公眾號，獲取后續系列技術文章。