NVIDIA Cosmos · Architecture Field Guide

Cosmos-Framework 架构导览

六张图,一条主线——从代码库布局一路读到分布式并行:数据怎么进来、怎么被一个 Mixture-of-Transformers 处理、训练与推理怎么调用它、又怎么摊到成百上千张 GPU 上。

对象是 Cosmos 3 omnimodal world-model 的训练 + 服务框架(单个 cosmos_framework/ 包)。 每张图从源码还原,关键框都能点开跳到对应 GitHub 源码行(固定在 commit 26a50b8)。

来源 源码遍历 · 固定 commit 26a50b8 图数 6 更新 2026-07-20
一次前向的主线 · 点任意一段跳到对应图

① 数据与网络 · 一次前向怎么跑

🌌 00 · 宏观地图

代码库架构

七层:CLI 入口 → 训练/推理双路径 → 共享 MoT 核心 → 基础设施 → 后端 → 规划中子包。先建立坐标系。

🧠 01 · 完整前向

模型前向(数据 → 输出)

一次完整前向:5 模态各自编码投影 → 打包成序列 → MoT 双专家 ×N → 分流 → 反投影 → 解码。输入 / 输出镜像对称,+ 自回归 / 扩散两条回环。

🔬 02 · 放大到算子

单层展开

MoTDecoderLayer 内部:pre-norm、共享注意力联合 softmax、双通路 MLP/MoE、两处残差。

② 训练、推理与规模 · 同一张网络怎么被调用 & 摊开

⏱️ 03 · 时序对比

训练 vs 推理

同一个 rectified-flow 速度场:训练单点监督(1 次前向)vs 推理多步 ODE 积分(N 次)。

🎲 04 · 推理内部

采样内部

去噪循环放大:每步预测速度 → CFG 有/无条件混合 → 三种求解器(FixedStep/UniPC/EDM)→ 解码。

🕸️ 05 · 分布式

并行 / 分片策略

分区轴(FSDP/HSDP)× 叠加轴(CP/CFG)、装配顺序、三种场景组合——这个框架存在的理由。

建议顺序:00 → 01 → 02 → 03 → 04 → 05,从坐标系到完整前向、放大到单层,再看训练/推理与规模化。