NVIDIA 开放 Cosmos 框架:用于训练和运行世界模型
当开发者谈论多模态网络时,他们通常指的是文本和图像的结合。NVIDIA 走得更远,发布了 Cosmos 3 系列——世界模型将文本、视频、音频和物理动作序列整合在单一的 Mixture-of-Transformers 架构中。为了使用这个系列,公司发布了一个代码库 cosmos-framework。
框架内部结构
代码库打包为单个 Python 包 cosmos_framework。该项目面向需要为特定任务微调多模态模型或将其部署到生产环境的工程师。如果你正在从事机器人技术或可控行为的视频生成,这个代码库将为你节省大量编写样板代码的时间。
所有逻辑分为两个主要模块。
第一个模块处理分布式训练。该工具支持 FSDP、张量并行(TP)、上下文并行(CP)和流水线并行(PP)。检查点以 DCP 格式保存,但代码包含转换为熟悉的 HuggingFace 格式的转换工具 safetensors。数据准备通过 JSONL、WebDataset、HuggingFace 和 LeRobot 的现成适配器简化。主启动脚本位于 cosmos_framework.scripts.train。
第二个模块处理推理。在底层,你可以使用 Diffusers、Transformers、vLLM 后端或自定义封装器。对于后台生成,提供离线模式;对于 API 创建,部署了 Ray 和 Gradio 技术栈。推理通过 cosmos_framework.scripts.inference 模块启动。
对于机器人任务,该框架包含一个 Policy Server。它提供 Cosmos3-Policy-DROID 模型服务,并将生成的命令直接传递给执行器。
硬件要求与首次启动
该项目的硬件需求相当大。examples 文件夹中的现成配置已在配备 8 张 NVIDIA H100 80GB 卡片的服务器上测试通过。你可以通过 NPROC_PER_NODE 和 FSDP 设置调整并行度和分片参数,但在单张消费级 GPU 上训练模型是不可能的。
该项目使用 uv 包管理器进行环境设置。作者建议从官方 NVIDIA NGC Docker 镜像(nvcr.io/nvidia/pytorch:25.09-py3)开始。
系统依赖安装如下:
sudo apt-get install -y --no-install-recommends curl ffmpeg git-lfs libx11-dev tree wget
为 CUDA 13.0 构建虚拟环境只需一条命令:
uv sync --all-extras --group=cu130-train
source .venv/bin/activate && export LD_LIBRARY_PATH=
为了在单 GPU 上测试推理,代码库包含一个现成的预设:
python -m cosmos_framework.scripts.inference \
--parallelism-preset=latency \
-i "inputs/omni/t2v.json" \
-o outputs/omni_nano \
--checkpoint-path Cosmos3-Nano \
--seed=0
与编码智能体的自动集成
一个有趣的细节:作者在代码库结构中内置了 AI 助手支持。.agents/skills/ 目录包含 Cursor、Claude Code 和 Codex CLI 的现成指令。
其中包含多个 SKILL.md 格式的文件。每条指令涵盖各自的领域:
- 环境设置、权重下载和 Docker 使用
- 浏览
cosmos_framework包文件 - 推理和并行参数配置
- 从数据准备到导出的完整 SFT 微调周期
- 诊断 CUDA 错误、OOM 和容器崩溃
当你让 Claude Code 或 Cursor 修复导入错误或修改启动脚本中的参数时,智能体会自动从 .agents/skills/ 目录加载相应的规则文件。
值得一试吗
该框架对机器人团队和生成式视频系统开发者很有用。代码库还很年轻,但已经包含了模型工作流程的完整工具集:从数据集加载到物理硬件控制。
如果你可以访问 GPU 集群,并且正在寻找一个现成的框架来实验世界模型,请查看 cosmos-framework。examples 文件夹中的脚本将给你一个良好的起点。
相关项目