82
Using Cosmos 3 as an example, this article introduces how Baidu Baige achieved 89x boot acceleration, 99.3% throughput improvement, and MFU surpassing the official baseline on domestic GPUs through AI Infra optimizations such as data loading, operator offloading, and compilation adaptation. 
Yesterday 3481 words (about 14 min) View Source →
Sign in to highlight text and take notes as you read. Sign in now
CSDN 2026-07-02 15:24 江苏
从「能跑」到「好用」，中间横着一道工程门槛，AI Infra 工程优化是回答这道必答题的关键。
来源 | 百度智能云技术站
经过百度百舸的 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 21.30 100.0% 2 16 42.42 99.6% 4 32 84.06 98.7% 8 64 167.68 98.4% 12 96 251.29 98.3% !Image 7
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 工程实践，欢迎关注百度智能云技术站公众号，获取后续系列技术文章。 阅读原文