当前位置：首页 > news >正文

不用 NVLink，如何通过 AI Infra 工程优化拉满 Cosmos 3 训练吞吐

news 2026/6/30 6:28:41

经过百度百舸的 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 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%