随着世界模型、具身智能模型、生成式模型等新一代 AI 基础模型快速发展，AI Infra 工程优化的重点，正从提升模型训练效率，逐渐转向优化整个训练系统。大模型训练已经不再只是模型前向、反向传播和参数更新。从数据进入训练系统开始，视频编码、数据处理等前置计算共同组成了一条更长、更复杂的训练流水线。当其中某个环节成为瓶颈时，仅仅优化模型训练本身，已经无法继续提升整体训练效率。

在 Cosmos3-Nano-Policy-DROID 的训练实践中，我们首先重新设计了训练架构，将 VAE 编码从训练节点中解耦，构建训练与编码分离的异构 GPU 架构；随后围绕数据流转、任务调度、缓存复用和编码性能等关键环节持续优化，让整套架构真正发挥价值。

最终，基于百度百舸平台，在 1 台主流 GPU 节点和 2 台低成本 GPU 节点组成的异构训练集群上，以约 1.5 倍的硬件投入将训练吞吐提升至 2.1 倍（从 21.6 samples/s 提升至 45 samples/s）。同样的预算可多获得约 35% 的训练吞吐。

这项实践也说明，AI Infra 工程优化的对象，正在从单个模型逐渐扩展到整个训练流水线。算力效率提升并非依靠单纯扩容 GPU，而是可以通过科学的资源分层配置，利用成本更低的编码节点承接前置计算，释放昂贵训练节点的计算资源，从而提升整体投入产出比。

1. 当训练瓶颈不再只是模型本身

长期以来，大模型训练的优化主要围绕模型计算展开。如何提升 Forward、Backward、通信以及 Optimizer 的效率，是训练系统持续演进的重点。

在本系列前两篇文章中，我们的工作正是围绕这一层面展开，两篇实践均基于百度百舸平台 hpas.lgn7ib 实例（搭载国内主流 GPU 型号）完成：第一篇《不用NVLink，如何通过 AI Infra 工程优化拉满 Cosmos 3 训练吞吐》完成了 16B 参数规模 Cosmos3-Nano-Policy-DROID 的全链路训练调优，任务启动速度提升89 倍，单机吞吐提升 99.3%，MFU 达到 0.42，超过官方基准，并在 12 节点规模下验证了 98.3% 的集群扩展效率；

第二篇《Cosmos3-Super 512 卡 Scaling，百度百舸 AI Infra 工程优化实践》将模型升级至 64B 参数规模的 Cosmos3-Super，在实例间无 HPN、实例内无 NVLink 的条件下，把训练规模从 4 节点扩展至 64 节点 512 卡，扩展效率保持在 97.48%，验证了世界模型大规模训练的近线性扩展能力。

至此，从单机吞吐到集群扩展，「模型训练本身」的效率已经得到充分释放。

但随着新一代基础模型的发展，训练流程本身也在发生变化。无论是世界模型、具身智能模型，还是视频生成等生成式模型，训练中都需要处理海量视频数据；视频数据在进入模型之前，需要首先经过 VAE Encoder，将原始视频压缩为 Latent 表示，再送入模型完成后续训练。这类前置计算虽然不参与模型参数更新，却已经成为训练流水线不可或缺的一部分。

因此，影响整体训练效率的，也不再只是模型训练本身，而是整条训练流水线。如果流水线中的某一个环节速度跟不上，即使训练 GPU 仍有充足算力，整体训练吞吐也会受到限制。

2. Cosmos 3 遇到的新瓶颈

在 Cosmos3-Nano-Policy-DROID 的训练过程中，VAE Encoder 正是这样一个典型瓶颈。性能分析发现，在百度百舸平台 hpas.lgn7ib 实例（下文简称训练实例）上，VAE Encoder 单次迭代耗时约 6.4 秒，占整个前向计算时间约 76%。

更关键的是，这一瓶颈在训练 GPU 内部几乎没有优化空间。深入的性能剖析显示，编码期间 GPU 算力已接近满载，且计算完全串行；我们尝试的多种卡内并行方案均无收益：编码计算本身已占满算力，卡内不再有可利用的空闲资源。

进一步分析发现，这部分计算与模型训练存在明显不同。整个训练过程中，VAE Encoder 始终保持冻结，仅执行前向计算，不参与梯度更新，也不存在跨卡参数同步；同时，Encoder 模型规模较小，对显存的需求远低于训练模型本身。也就是说，它位于训练流水线之中，却具有完全不同于训练计算的资源特征。

继续将编码与模型训练部署在同一张训练 GPU 上顺序执行，意味着训练 GPU 需要花费大量时间完成视频编码，而真正的模型训练只能等待编码结束后才能继续进行。当卡内优化已经证明无路可走、VAE 编码又占据大部分前向计算时间时，瓶颈已经不再只是模型，而是整个训练流程的组织方式。

3. 第一步：重新设计训练架构

针对这一问题，我们首先重新设计了 Cosmos3-Nano-Policy-DROID 的训练架构。核心思路并不是继续优化编码本身，而是将训练计算与编码计算彻底解耦。

整套架构由两类实例构成：训练集群沿用前两篇文章中的 hpas.lgn7ib 实例；VAE Encoder 则独立部署到由成本更低的 aihc.gn5rc.c128m476.8gen40 实例（下文简称编码实例）构成的编码集群中，每个编码实例上运行多个无状态 Encoder Actor。训练节点持续发送编码请求，编码集群完成视频编码并返回 Latent；训练 GPU 则专注于模型训练，不再承担编码任务。

这样，原本串行执行的训练流程，被拆分为两条能够独立运行的计算链路：一条负责模型训练，一条负责视频编码，两者通过数据交换协同工作。异构部署不仅使训练 GPU 彻底摆脱了编码计算负担，也让编码算力能够以远低于训练 GPU 的成本独立扩容。

不过，仅仅完成异构部署，还不足以发挥整套架构的价值。真正的挑战，是如何让训练侧和编码集群始终保持高效协同，而不是彼此等待。

4. 第二步：让整个流水线持续运转

为了充分发挥异构架构的优势，我们重新设计了数据流转方式。传统流程通常以 Batch 为单位组织计算：一个 Batch 完成编码之后，训练才能继续。而在新的架构中，我们将 Batch 拆分为独立样本进行流式处理。

训练侧每个节点被拆分为三条并行链路。Producer 线程持续从 DataLoader 获取 Batch，将其中的样本拆分为若干小组、打上标识后流式发送至编码集群，发出即处理下一组；Encoder Actor 收到样本后立即编码，并携带原始标识立即返回，收一组、编一组、回一组，从不等待完整 Batch；Receiver 线程持续接收编码结果，按标识将 Latent 回填至所属 Batch，集齐后组装为完整 Batch 投入队列；训练主循环从队列中取出完整 Batch 执行训练，与前两条线程完全并发。

整个过程中，各模块始终独立运行，没有任何一个阶段需要等待另一个阶段全部结束。通信采用双向异步模式，发送与接收同时进行，编码与网络传输相互重叠。

更重要的是，系统允许多个 Batch 同时处于不同处理阶段，并为同时在途的 Batch 数设置上限：一旦达到上限，数据读取会自动暂缓，等已有 Batch 被消化后再继续，导致内存溢出。当 Batch A 还在等待部分编码结果返回时，Batch B 已经开始发送，Batch C 已经进入编码，而训练主循环可能正在处理更早完成的 Batch。多个 Batch 在流水线中持续流动，使编码、通信和模型训练长期保持重叠执行，编码时间被完整隐藏进训练的 Backward 与 Optimizer 阶段。

异构架构与流水线落地后的实测结果如下：1 台训练实例搭配 3 台编码实例时，稳态迭代时间从单机的约 12 s 降至约 5.8 s，训练吞吐从 21.6 samples/s 提升至 45 samples/s；稳态阶段 dataload 与队列等待时间接近 0，编码不再是瓶颈。

这一提升的成本代价远小于同构扩容：编码实例采用低成本 GPU，单台硬件成本仅约为训练实例的三成。以此折算，「1 台训练实例 + 3 台编码实例」的整体硬件投入约为单台训练实例方案的 1.8 倍，吞吐却达到 2.1 倍。这相当于在相同预算下，可多获得约 15% 的训练吞吐。

而这还只是架构与流水线层面的收益，第五步的算子优化，还将把这一性价比进一步推高。

5. 第三步：让整个编码集群保持高利用率

异步流水线解决了训练侧与编码侧协同的问题，但当部署规模扩大后，新的瓶颈很快出现在编码集群内部。

在真实生产环境中，训练 GPU 与编码 GPU 通常来自不同资源池，两者规模会随业务需求不断调整，往往不存在整除关系。如果采用固定绑定方式，资源利用率很容易受到限制。以两机训练的实际部署为例：训练端共 16 个数据并行 Rank，编码端 5 台编码实例共 40 个 Actor。若为每个 Rank 固定绑定 Actor，每个Rank 只能分到 2 个，最终只有 32 个 Actor 能分到流量，剩余 8 个 Actor（恰好是整整一台编码实例）长期空闲。

针对这一问题，我们设计了全局动态调度机制。其一，每个训练 Rank 与全部 Actor 建立连接，不再固定绑定，任意 Actor 均可服务任意 Rank；其二，任务发送采用轮询方式，并以各 Rank 自身编号作为轮询起始偏移，各 Rank 的请求序列天然错开，避免同时涌向同一节点。配合细粒度的任务拆分，请求由全部 Rank 持续、均匀地撒向所有 Actor，负载在时间维度上自然收敛均匀。接收侧本就按标识重组、不挑节点，无需任何改动。

整个方案无需中心调度器，也无需维护全局状态。实测验证了均衡效果：40 个 Actor 全部有流量，5 台编码实例的请求量完全相等，单 Actor 请求数偏差为 0。在 16 Rank × 40 Actor、global batch size 512 的配置下，稳态迭代时间从约 8 s 降至约 5.95 s，全局吞吐从约 68 samples/s 提升至约 86 samples/s，提升约 26%。

6. 第四步：让相同输入只编码一次

在模型调试和超参数搜索过程中，同一数据集通常需要被反复训练。对于 VAE Encoder 而言，相同输入始终对应相同的 Latent 输出，每一次重新编码，本质上都是重复计算。

针对这一特点，我们在每个编码节点上设计了两级本地 Latent 缓存。缓存键由样本内容唯一确定：只有「相同输入、相同预处理参数」才会命中，从根本上保证命中结果与重新编码完全一致。缓存分为两级：内存热缓存响应最频繁的访问，本地盘冷缓存承载更大的容量。读取时先查内存、再查本地盘，两级都未命中才真正调用 GPU 编码，并将结果写回缓存。

缓存容量按 LRU 策略自动维护：内存缓存写满时，将较久未使用的数据下沉到本地盘；本地盘写满时，再淘汰最旧的数据。同一节点上的多个 Encoder 进程共享同一份缓存，且整个淘汰过程不会阻塞正常编码。

实测中我们发现一个容易被忽略的细节：缓存要真正命中，需要两个条件同时成立。一是每一轮读到的画面像素完全一致（需关闭随机裁剪、色彩抖动等像素级随机增广），二是样本的读取顺序在各轮之间保持一致。只满足其一时，命中率仅约 9%；两个条件齐备后，第二轮训练对首轮已编码样本的命中率可达 100%，几乎不再产生新的编码。单个 Latent 仅约 1.29 MB，所需缓存空间可按数据集规模提前估算。

需要说明的是，缓存命中节省的是编码集群的 GPU 算力，而非训练关键路径。其价值体现在两方面：一是在相同训练吞吐下可部署更少的编码节点；二是对于验证集评估、固定种子复现等输入稳定的场景，可预先完成编码、直接复用。

7. 第五步：持续释放每张编码 GPU 的性能

随着整个异构架构逐渐稳定，瓶颈开始转移到单个编码节点的编码效率。我们围绕底层执行链路，在编码实例上做了一系列实测驱动的优化。

第一层收益来自整图编译。VAE Encoder 的输入尺寸固定，我们将其编码过程整体导出、预编译为一张静态计算图（AOT 编译），消除逐算子的调度开销并完成跨算子融合，单次编码从 512 ms 降至 473 ms。预编译产物可被同节点多个 Actor 共享加载；同一张卡上的多个编码进程共享同一份模型权重，进一步降低显存占用。

第二层收益来自卷积算法的自动选优。由于输入尺寸固定，底层库会在首次运行时实测多种卷积实现、缓存其中最优的一种，此后固定复用，编码时间进一步从 473 ms 降至 402 ms。我们也尝试了更激进的自动调优模式，让编译器把自动生成的候选实现一并纳入比拼，但实测下来最优解仍是原有的卷积实现，收益与前者持平，编译耗时却大幅拉长，因此没有采用。

最后一层收益来自面向硬件特性的定制算子。深入分析编码实例所配 GPU 的微架构特性后，我们发现：通过调整卷积计算的数值累加路径，可以更充分地发挥其 Tensor Core 算力，获得接近翻倍的卷积吞吐。为此，我们基于 CUTLASS 实现了适配该硬件特性的定制卷积算子，端到端编码时间降至 275.5 ms，相对未优化前提升 1.86 倍。

需要说明的是，CUTLASS 定制算子与 AOT 是互斥的两条编译路径：手写的定制算子无法被静态导出，因而不能走 AOT 预编译，只能在运行时即时编译（JIT）为整图。整图融合的收益等价保留，且由于定制算子无需再参与卷积算法搜索，编译时间反而更短。最终的生产路径为：即时整图编译 + 定制卷积算子。

算子优化最终落回集群收益：同为 1 台训练实例搭配 2 台编码实例的配置，稳态迭代时间从约 8.6 s 降至约 5.8 s，吞吐从 30 samples/s 提升至 45 samples/s，达到了优化前 3 台编码实例的水平。相同的目标吞吐，可以少部署一台编码节点。

至此，最优配置收敛为「1 台训练实例 + 2 台编码实例」：总硬件投入仅约为单机的 1.5 倍，吞吐保持单机的 2.1 倍。同样的预算可多获得的训练吞吐，从第二步的约 15% 提升至约 35%。这才是整套优化叠加后的最终性价比。

8. 精度验证：优化前后训练 loss 对齐

我们在相同数据、相同超参数下，对比了基线与异构编码集群两套配置的训练 loss 曲线进行验证。实测结果如下图所示：红色线（异构编码集群版本）与橙色线（Baseline）呈现完全一致的收敛趋势，总体训练 loss 与关键分项 loss 均收敛到相同水平，验证了工程优化的无损性。

9. 从优化模型，到优化训练系统

经过异构架构设计，以及围绕流水线协同、资源调度、缓存复用和底层算子的一系列持续优化，Cosmos3-Nano-Policy-DROID 在保证训练精度一致的前提下，实现了训练与编码计算的高效重叠。

最终，在 1 台主流 GPU 节点之外仅增配 2 台低成本 GPU 节点，以约 1.5 倍的硬件投入获得了 2.1 倍的整体吞吐（21.6 → 45 samples/s）。同样的预算，可多获得约 35% 的训练吞吐。同样的思路也适用于一切「冻结、轻量、无通信依赖」的前置计算，无论是视频 VAE、图像 Tokenizer 还是各类特征提取器。

更重要的是，这项实践也反映出 AI Infra 工程优化思路正在发生变化。过去，我们更多关注如何让模型训练得更快；今天，我们同样需要关注如何重新组织整条训练流水线，让不同类型的计算运行在最合适的资源上，让不同阶段能够持续协同工作。对于 AI Infra 而言，真正需要优化的，不只是模型，而是模型背后的整个训练系统。

更多关于 Cosmos3 集群训练优化、模型效果评测的全方位解读，敬请期待后续文章。
了解百度百舸更多信息