一句话定位
Cosmos 世界基础模型平台第二代:统一的扩散 Transformer 架构,覆盖 Text2Image 与 Video2World 两大能力,提供 0.6B/2B/14B 三档规模,原生支持 720p/480p 分辨率与 10/16fps 帧率的自由组合,并配套独立训练的 NATTEN 稀疏注意力 checkpoint 把 720p 推理成本压到基线的 38%-59%;内置 Cosmos-Reason1-7B 同时充当 prompt refiner 与 rejection-sampling 的 critic。GitHub / Hugging Face / PyPI 三线开源(NVIDIA Open Model License),但未发布独立技术报告(README 长期标注”Paper coming soon”),继任者 Cosmos-Predict2.5 于 2025 年 10 月发布后,GitHub 仓库 README 顶部换成了迁移横幅(README 最后一次修改记录为 2025-10-07,与 Predict2.5 发布时间点接近)。
背景与定位
延续 [cosmos-predict]确立的”预训练通才世界模型 → 下游小数据 post-training 专才”两段式平台叙事。Predict1 同时押注扩散(DiT + Cosmos-Tokenize1-CV8x8x8 连续 tokenizer)与自回归两条技术路线;Predict2 收敛为只保留并升级扩散 DiT 这一支——tokenizer 从 8x8x8 换成时间压缩率更低的 8x8x4(细节见下文架构),latent 时间粒度更细,并新增可选的 NATTEN 稀疏注意力 checkpoint 把推理成本打下来。
发布约四个月后(2025-10),NVIDIA 推出继任者 Cosmos-Predict2.5(cosmos-predict2-5-world-simulation),把 Text2World/Image2World/Video2World 三任务收进单一 flow-matching 模型、换成 WAN2.1 VAE tokenizer、用 Cosmos-Reason1 替换 T5 文本编码器。相比之下 Predict2 是 NVIDIA 把 Text2Image、Video2World 及后续 multiview / action-conditioned 变体分开训练、各自独立发布 checkpoint 的最后一代:仓库自 2025 年 6 月发布骨架架构后只增加了 0.6B 规模与 NATTEN 变体,未再迭代核心设计。
同属 Cosmos 生态、可交叉参考的是 cosmos-drive-dreams(Transfer1 驱动的驾驶合成数据)。Predict2 用来自证”世界模型”身份的评测基线是 Wan2.1-I2V-14B-720P、HunyuanVideo-I2V、CogVideoX-5B-I2V、LTX-Video 这些同期开源视频生成模型,评测维度沿用 NVIDIA 自建的 PBench(Domain/Quality 双轴),而非通用视频生成常用的 FID/美学评分。
模型架构
统一 DiT 骨干(GitHub 代码里的类名 MinimalV1LVGDiT),Text2Image 与 Video2World 共享同一套 backbone / patchify / 位置编码设计,仅输入输出的时间维不同:
| 规模 | model_channels | num_blocks | num_heads | head dim | 用途 |
|---|---|---|---|---|---|
| 0.6B | 1,280 | 20 | 20 | 64 | 仅 Text2Image |
| 2B | 2,048 | 28 | 16 | 128 | Text2Image + Video2World |
| 14B | 5,120 | 36 | 40 | 128 | Text2Image + Video2World |
(数字取自 GitHub cosmos_predict2/configs/base/config_video2world.py 与 config_text2image.py,两文件中 2B/14B 的这组数字完全一致,说明两个任务共享同一套 backbone 规格)
- Block 结构:自注意力 + 交叉注意力(文本条件)+ 前馈层交替堆叠;每层前置 adaLN(用 LoRA 压缩,
adaln_lora_dim=256)注入扩散时间步嵌入。 - 位置编码:3D RoPE(
rope3d),可学习,extrapolation ratio 因规模而异——2B 为 (h,w,t)=(3.0, 3.0, 1.0),14B 为 (2.0, 2.0, 0.8333);patch 化为 spatial 2 × temporal 1;latent 通道数(in/out_channels)均为 16。 - 视觉 tokenizer:连续、因果视频 VAE,spatial 压缩 8×8、temporal 压缩 4×(代码里的公式
get_latent_num_frames(n) = 1 + (n-1)//4),latent 通道数 16。训练用 clip 标准长度 93 帧 @ 704×1280(720p,16:9,约 5.8 秒 @16fps)→ 24 个 latent 帧(state_t=24),与上式吻合。推理 API 暴露的--num_conditional_frames只能取 1 或 5:1 帧图像对应 1 个 latent 帧,5 帧视频经同一公式对应 2 个 latent 帧,正好等于配置里的max_num_conditional_frames=2。 - 文本编码器:T5-11B(仓库代码与文档里也称”T5-XXL”)预计算 embedding,训练时文本条件 dropout 率 0.2(为 classifier-free guidance 训练做准备);推理时 guidance scale 默认 7.0。
- 条件策略:FRAME_REPLACE——图像/视频条件帧的 latent 直接替换生成序列开头对应位置的 latent,并对条件 latent 加”augment noise”弥合训练/推理 gap(HF 模型卡原话)。
- 精度:仅验证过 bfloat16(fp16/fp32 未获官方支持);扩散损失用 rectified-flow 参数化(配置字段
rectified_flow_t_scaling_factor=1.0、rectified_flow_loss_weight_uniform=True、sigma_data=1.0、sigma_conditional=0.0001)。 - 安全与增强模块(均可选、默认开启、可 offload 到 CPU):Llama Guard 3(8B)作 guardrail 过滤有害内容;Cosmos-Reason1-7B 身兼二职——prompt refiner(自动扩写过短的 prompt)与 rejection-sampling(best-of-N)的 critic 模型。
- 动作条件变体(Cosmos-Predict2-2B-Sample-Action-Conditioned,专用类
ActionConditionedMinimalV1LVGDiT):动作维度action_dim = 7 × 12 = 84,即每帧 7 维(末端执行器 Δpose [x,y,z,roll,pitch,yaw] + 二值夹爪开合)乘 12 个未来帧;给定 1 帧图像 + 12 步动作,预测未来 12 帧,480p/4fps,在 Bridge(经 IRASim 整理)数据集上 post-train 而成。 - NATTEN 稀疏注意力变体:是独立训练并单独发布的 checkpoint(HF 模型卡把它列为 6 个可下载变体之一),而非纯推理期插件,仅限 720p、仅限 Hopper(cc9.0)/ Blackwell(cc10.0)硬件。
数据
Predict2 本身预训练用的数据规模、来源、配比、过滤流程均未披露——GitHub README 发布时注明”Paper (coming soon!)“,这份技术报告始终未以 Predict2 的名义单独发表;继任者 Predict2.5 的技术报告披露了自己的预训练规模(>2 亿条视频、约 3500 万小时素材),但无法确认与 Predict2 是否共享同一批语料。以下是 GitHub 仓库中可验证的、post-training/微调层面的具体数据:
- 训练 clip 标准格式:93 帧、704×1280(720p,16:9)。
- Cosmos-NeMo-Assets:4 条视频的最小 demo 数据集,只用于验证 post-training 流程能跑通,不代表真实规模。
- AgiBotWorld-Alpha 鱼眼子集(task_id 327):原始约 100GB,按 5 秒窗口切分并清理后剩约 2GB,视频与文本 caption 配对,T5-XXL embedding 预计算存为 pickle。
- GR00T-Dreams(GR1、DROID 两个子集):用于两个”Sample-GR00T-Dreams”发布 checkpoint 的 post-training 语料,来自 NVIDIA Isaac GR00T 项目的合成/遥操作数据,具体规模未披露。
- Bridge(经 IRASim 整理版):action-conditioned post-training 用,每条视频配 JSON 标注(末端执行器 pose、夹爪状态、7 维 delta 动作),数据集规模未披露。
- 官方推荐的通用 post-training 数据配方(开发者博客):采集 100+ 小时遥操作视频,用 Cosmos-Curate 做分段整理,用任意 VLM(含 Cosmos-Reason1)生成 caption。
训练方法
GitHub 公开的是 post-training/微调 recipe,预训练配方本身未披露:
- 分布式:FSDP + context parallelism;LoRA 可选,一个 flag(
model.config.train_architecture=lora)在全量微调与 LoRA 之间切换。 - 优化器:FusedAdamW;示例配置(
cosmos_nemo_assetsrecipe)lr=2^(-14.5)≈4.3e-5,线性 warmup +lambdalinear调度,max_iter=1000,每 500 iter 存 checkpoint;训练时开启 EMA,关闭 prompt refiner 与 guardrail(两者只在推理期启用)。 - 2B 与 14B 的并行度不同:2B 用 context-parallel size=2(单机 8 卡),14B 用 context-parallel size=8(4 节点 × 8 卡 = 32 卡),14B 单步全局 batch 因而比 2B 低 4 倍(GitHub 文档原话)。
- Action-conditioned 变体:新增专用 conditioner 把 7×12 维动作张量接入模型,训练配置单独维护在
cosmos_predict2/configs/action_conditioned/目录下。 - NATTEN 变体:作为独立 checkpoint 训练/发布,而非纯推理期插件。
Infra(训练 / 推理工程)
-
预训练 GPU 数与 GPU-hours:未披露(无独立技术报告)。
-
Post-training(微调)算力规模有具体公开例子:2B 用 8 卡(单机);14B 用 32 卡(4 节点 × 8 卡),对应 context-parallel size 从 2 升到 8。
-
Post-training 迭代速度(AgiBot-Fisheye 示例,秒/迭代):
GPU 2B-Video2World 14B-Video2World NVIDIA B200 6.05s 6.27s NVIDIA H100 NVL 10.07s 8.72s NVIDIA A100 22.5s 22.14s -
推理显存(仅 bf16 已验证):2B-Text2Image 26.02GB、14B-Text2Image 48.93GB、2B-Video2World 32.54GB、14B-Video2World 56.38GB。
-
推理单次生成耗时(无 NATTEN 稀疏化),节选自 GitHub
performance.md的完整 11 款 GPU 表:GPU 2B-T2I 14B-T2I 2B-V2W(720p16fps) 14B-V2W(720p16fps) GB200 3.39s 8.5s 25.61s 85.26s B200 3.24s 8.68s 30.7s 92.59s H200 SXM 9.02s 15.96s 50.2s 176.19s H100 PCIe 11.12s 23.83s 79.87s 286.46s DGX Spark 24.87s 138.94s 344.64s 1902.26s -
NATTEN 稀疏注意力加速(仅 Hopper/Blackwell、仅 720p),节选 720p16fps:H100 SXM 上 14B-Video2World 从 856.9s 降到 426.0s(2.0×);2B-Video2World 从 228.8s 降到 94.2s(2.4×);B200 上 14B 从 439.4s 降到 223.1s(2.0×)。跨全部机型/帧率组合的加速幅度在 1.7×-2.6× 之间。
-
多卡推理:context parallelism(
--num_gpus+ torchrun),可选 CUDA Graphs 加速 DiT 前向。 -
长视频生成:自回归分 chunk 方式,每个 chunk 用前一 chunk 最后
num_conditional_frames帧续接,推荐配合多卡 context parallelism 使用(GitHub 演示过 6-chunk 的 14B 样例)。 -
Control-Hz / 端侧硬件延迟:未披露。
评测 benchmark
PBench(video2world,NVIDIA 自建的物理 AI 视频评测,research.nvidia.com/labs/dir/pbench),数字来自 HF 模型卡与项目页,两处数值完全一致但列名标注不同(HF 模型卡标”Overall/Domain/Quality Score”,项目页标”Domain/Quality/PBench Score”——同一组数字被贴了两套不同表头,以下按 HF 模型卡的标注呈现):
| 模型 | Overall | Domain | Quality |
|---|---|---|---|
| LTX-Video | 74.0 | 77.2 | 70.8 |
| HunyuanVideo-I2V | 74.0 | 77.4 | 70.6 |
| CogVideoX-5B-I2V | 74.2 | 79.5 | 69.0 |
| Wan2.1-I2V-14B-720P | 75.8 | 81.9 | 69.7 |
| Cosmos-Predict1-7B-Video2World | 73.2 | 77.4 | 69.0 |
| Cosmos-Predict1-14B-Video2World | 73.3 | 77.6 | 69.0 |
| Cosmos-Predict2-2B-Video2World | 77.2 | 84.8 | 69.6 |
| Cosmos-Predict2-14B-Video2World | 77.4 | 84.9 | 69.9 |
(以上为 720p/16fps、无 NATTEN 的默认变体;HF 模型卡另给出 2B/14B × {480p,720p} × {10fps,16fps} × {NATTEN 开/关} 共 12 组变体的完整 sweep,Overall 分布在 76.3-77.4 之间,4 组开启 NATTEN 的变体与对应基线相比 Overall 最多降 0.6 分)
GenEval(text2image),与 Stable Diffusion XL、DALL-E 3、Flux 1-Dev 对比(HF 模型卡与项目页数字一致):
| 模型 | Overall | Single Obj | Two Obj | Counting | Colors | Position | Color Attr |
|---|---|---|---|---|---|---|---|
| Stable Diffusion XL | 0.55 | 0.98 | 0.74 | 0.39 | 0.85 | 0.15 | 0.23 |
| DALL-E 3 | 0.67 | 0.96 | 0.87 | 0.47 | 0.83 | 0.43 | 0.45 |
| Flux 1-Dev | 0.66 | 0.98 | 0.79 | 0.73 | 0.77 | 0.22 | 0.45 |
| Cosmos-Predict2-2B | 0.83 | 1.00 | 0.99 | 0.73 | 0.89 | 0.65 | 0.73 |
| Cosmos-Predict2-14B | 0.84 | 1.00 | 0.98 | 0.79 | 0.90 | 0.64 | 0.72 |
Position 与 Color attribution 两个子项领先基线幅度最大(0.64-0.65 vs 基线最高 0.43;0.72-0.73 vs 基线最高 0.45),对应它作为物理场景合成数据源需要的精确空间/属性绑定能力。
Action-conditioned 与两个 GR00T-Dreams post-train 样例 checkpoint 未公开对应的量化 benchmark 数字。
创新点与影响
- 相对 Predict1 的两处收敛:砍掉自回归路线只留扩散 DiT 一支;tokenizer 从 8×8×8 换成 8×8×4,latent 时间粒度更细。
- 把分辨率(480p/720p)与帧率(10/16fps)做成可自由组合的矩阵,而不是像 Predict1 那样绑定固定规格,是相对上一代更明显的产品化改进。
- NATTEN 稀疏注意力 checkpoint 把 720p 推理成本压到基线的 38%-59%(即 1.7×-2.6× 加速),PBench 质量几乎不掉(Overall 最多降 0.6 分),是本作在”speed”这条自我定位承诺上的主要交付。
- 把 prompt refiner(Cosmos-Reason1-7B 扩写短 prompt)与 rejection sampling(同一模型当 critic 做 best-of-N)打包进标准推理脚本,是不需要额外训练即可提升产出质量的工程化手段。
- 自陈局限(HF 模型卡原话大意):长时高分辨率视频仍有伪影;时间不一致、相机/物体运动不稳、交互不精确;可能不符合 3D/4D 空间或物理定律(物体消失/变形、不合理运动);因此不适合要求严格物理律仿真或复杂多智能体动态的场景。
- 平台层面最大的局限是自陈的:承诺的独立技术报告始终未兑现(“Paper coming soon”),GitHub 仓库 README 目前(本页写作时)已换成迁移到 Cosmos-Predict2.5 的横幅,数据规模/配比/预训练 GPU 预算等一手工程细节因此永久缺失——除非 Predict2.5 报告事后确认二者共享哪些训练资产,但目前没有找到这样的确认。
原始链接
- GitHub: https://github.com/nvidia-cosmos/cosmos-predict2
- Hugging Face 合集: https://huggingface.co/collections/nvidia/cosmos-predict2-68028efc052239369a0f2959 (现已归档并跳转到 Predict2.5 合集;14B 等 gated 模型页需登录同意协议后仍可查看完整 model card)
- 项目页(Cosmos Lab): https://research.nvidia.com/labs/dir/cosmos-predict2/ (跳转到 https://research.nvidia.com/labs/cosmos-lab/cosmos-predict2/)
- 官方开发者博客: https://developer.nvidia.com/blog/develop-custom-physical-ai-foundation-models-with-nvidia-cosmos-predict-2/
- HF 社区博客(NVIDIA 官方账号发布): https://huggingface.co/blog/nvidia/cosmos-predict-2
- PyPI: https://pypi.org/project/cosmos-predict2/
一手源存档(sources/)
- cosmos-predict2—github-readme — GitHub README 快照(模型列表、News 时间线、许可证)
- cosmos-predict2—github-docs — setup / performance / post-training(video2world、action-conditioned、AgiBot-fisheye)/ inference 系列 doc 合并快照
- cosmos-predict2—github-configs — config_video2world.py、config_text2image.py、tokenizer.py、action_conditioned/config.py 等关键代码片段
- cosmos-predict2—hf-card-14b-video2world — 14B-Video2World 模型卡(架构、输入输出规格、推理表、PBench)
- cosmos-predict2—hf-card-14b-text2image — 14B-Text2Image 模型卡(GenEval 表)
- cosmos-predict2—hf-blog — Hugging Face 社区博客(NVIDIA 官方账号)
- cosmos-predict2—project-page — Cosmos Lab 项目页快照(摘要 + PBench/GenEval 表 + 引用)
- cosmos-predict2—devblog — NVIDIA 开发者博客原文