经过百度百舸的 AI Infra 工程优化Cosmos3-Nano-Policy-DROID 的训练启动速度提升 89 倍国内主流 GPU 型号单机吞吐提升 99.3%MFU 达到 0.42显著优于官方基准12 节点扩展效率达 98.3%集群算力得到充分释放。1. 从开源社区到生产环境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 的训练性能成为模型落地过程中的关键挑战。2. 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 min25 s89 × 加速峰值系统内存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 s78.5%Video decode3 路 MP4 视频流0.31 s15%Transform / tokenize0.10 s5%Downsample concat0.02 s1%RandomCrop Resize0.01 s0.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 CheckpointingAC是缓解大模型显存压力的经典技术其核心逻辑是「以计算换显存」。社区代码为了追求极致的显存安全默认对 Cosmos3-Nano 全部 36 层 Transformer 开启了 Full AC。常规的优化思路是利用 Full AC 节省下的显存去调大 Batch SizeBS从而寻找整机吞吐的最优解。在搜索整机最大吞吐 BS 配置的过程中我们通过硬件监控发现即使 BS 已经调至当前硬件的极限GPU 显存依然存在大量空闲冗余硬件的绝对算力并未被完全榨干。在社区原生支持的 Full / Selective AC 基础上我们针对 hpas.lgn7ib 规格中 GPU 的显存容量进一步开发了分层 ACLayer-wise AC策略。通过精细化控制不同 Transformer 层的 AC 开启状态在「显存占用」与「重算开销」之间找到更优的平衡点将闲置显存转化为有效吞吐从 2.61 sample/s/GPU 提升至 2.69 sample/s/GPU提升 3.1%。AC 模式策略说明vs full (baseline)full全 36 层开启 ACbaselineselective层内部分算子关闭 AC0.1%≈无效Layer-wise AC部分层关闭 AC3.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扩展效率1821.30100.0%21642.4299.6%43284.0698.7%864167.6898.4%1296251.2998.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表示 TensorHMMA/IMMAPipe 处于 Active 状态的周期分数算力得到充分释放。3. 结语具身智能的大规模落地既不是单点算法突破能够完成的也不是单纯堆砌算力就能解决的。它需要模型结构创新、AI Infra 工程优化与算力网络能力的协同推进。目前除了本文介绍的 Cosmos3 之外百度百舸面向具身智能场景已经为客户交付了 OpenPI、DreamZero、GR00T N1.6、Lingbot-VLA、Motus、AHA-WAM、FastWAM、Wan2.2 等众多模型的训推加速成果。Cosmos3 的优化实践进一步证明开源社区负责推动算法创新而真正决定模型能否进入生产环境、能否以合理成本完成训练的则是 AI Infra 的工程优化能力。事实证明使用国内主流 GPU同样能够实现接近旗舰算力平台的训练效率满足实际业务需求。本文介绍的训练优化只是释放 Cosmos3 模型性能的第一步。在集群层面百度百舸已开展更深入的探索通过计算资源与模型结构的协同设计进一步释放集群潜力降低用户训练成本。更多关于 Cosmos3 集群训练优化、模型效果评测的全方位解读敬请期待后续文章。了解百度百舸更多信息