NVIDIA/cosmos-framework

解决的问题

提供一个统一的端到端框架，用于训练和部署多模态世界模型，特别是 Cosmos 3 模型系列。通过将语言、图像、视频、音频和动作序列等多种模态整合到单一架构中，解决了管理这些异构模态的复杂性，适用于物理人工智能（Physical AI）应用，如世界模拟器和动作模型。

工作原理

该框架以单一 Python 包（cosmos_framework）组织，包含两个主要功能路径：

• 训练：使用支持 FSDP、张量并行（TP）、上下文并行（CP）和流水线并行（PP）的分布式训练器。支持原生 DCP 检查点，并兼容多种数据集适配器，包括 JSONL、WebDataset 和 LeRobot。
• 推理：利用 Diffusers、Transformers 和 vLLM 等后端，通过 Ray 和 Gradio 提供离线批量生成和在线服务。

适用人群

适用于从事物理人工智能、世界模型和多模态生成式 AI 的研究人员和开发者，需要一个可扩展的系统来训练和部署能够同时处理和生成多种数据类型模型的用户。

主要亮点

• 多模态支持：联合处理和生成语言、图像、视频、音频和动作序列。
• 统一架构：采用混合 Transformer 架构，整合视觉-语言模型与世界模拟器。
• 可扩展训练：内置对高级分布式训练策略（FSDP/TP/CP/PP）的支持。
• 灵活推理：支持多种后端，并提供用于动作型模型的策略服务器。