NVIDIA Cosmos · System Map

Cosmos-Framework 架构图

面向 Cosmos 3 omnimodal world-model 家族的端到端训练 + 服务框架。 全部代码收敛在单个 cosmos_framework/ 包内,围绕一次大规模分布式训练/推理的工作流分层: CLI 入口分派出训练推理两条路径,二者共享同一套 model/ 核心(Mixture-of-Transformers),并坐落在跨领域的基础设施之上。

已落地(当前 release 存在) [planned] 规划中(目录尚未创建) 训练路径 推理路径
Layer 1 · CLI 入口

Entry-point scripts

python -m cosmos_framework.scripts.<name>
trainSFT 训练主入口,--sft-toml 驱动
inference离线批量 / 在线生成入口
convert_model_to_dcpHF → DCP 基座转换
export_modelDCP → HF safetensors 导出
action_policy_serverDROID / robolab 策略服务
Layer 2 · 配置系统

Configs

cosmos_framework/configs/
toml_config/用户面 pydantic schema:SFTExperimentConfigextra="forbid")+ VFM↔VLM 路径重映射
base/内部 LazyConfig 实验 SKU:experiment/sft/*.pyvlm/defaults/
训练路径 ◄─────── scripts.train ┊ scripts.inference ───────► 推理路径
Layer 3A · 训练基础设施

Training

FSDP / TP / CP / PP
data/JSONL · WebDataset · LeRobot 适配器;imaginaire · generator · reasoner
trainer/训练循环、梯度累积、optimizer/scheduler step、混合精度、callback 分派
callbacks/日志、grad_clip、MoE 稳定性监控、hf_export、checkpoint 节奏
checkpoint/DCP ↔ HF safetensors、s3_filesystem、schema 迁移与断点续训
Layer 3B · 推理子包

Inference

cosmos_framework/inference/
model · inference模型 + 推理引擎入口
args · defaults/SamplingArgs 与逐模态默认 sample_args.json
ray/Ray Serve · Submit · Gradio 在线服务
modalitiesvision · action · sound · transfer · interactive
训练产出 checkpoint推理载入 checkpoint
Layer 4 · 共享模型核心

Model Core — Mixture-of-Transformers

cosmos_framework/model/ generator/mot/
generator/mot/MoT 主干:统一处理 language·image·video·audio·action
generator/diffusion/扩散生成 + tokenizers · upsampler · reasoner · algorithm
attention/flash2 · flash3 · cudnn · natten 后端
tokenizer/Wan2.2 VAE 等视觉 tokenizer(models · evaluation)
Layer 5 · 横贯基础设施

Cross-cutting foundation

utils · tools · auxiliary · simulation
utils/日志、config 加载、分布式 setup、EMA、easy_io、vfm/vlm 助手
tools/包内 CLI:flops/ · visualize/
auxiliary/guardrail内容安全护栏
simulation/liberoLIBERO 仿真环境(策略评测)
Layer 6 · 生态后端 shim(独立可安装)

Backend packages

packages/
transformers-cosmos3把 Cosmos3 checkpoint 适配进 🤗 Transformers 生态
vllm-cosmos3vLLM 推理后端 shim
diffusers(原生)diffusers ≥ 0.39 内置 Cosmos3OmniPipeline,无需 shim
Layer 7 · 规划中子包(RL / 多 worker 编排)

Planned subpackages

目录当前故意缺席,功能落地时创建
algorithm/loss · reward · rlplanned
controller/head-node 编排planned
workers/rollout · reward · referenceplanned
communicator/进程间通信planned
evaluation/离线评测 harnessplanned
launcher/Slurm · torchrun · k8splanned