NVIDIA/cosmos-framework

Our inference and training framework to run on the Cosmos Models

What it solves

It provides a unified end-to-end framework for training and serving omnimodal world models, specifically the Cosmos 3 model family. It addresses the complexity of managing diverse modalities (language, images, video, audio, and action sequences) by consolidating them into a single architecture for Physical AI applications, such as world simulators and action models.

How it works

The framework is organized into a single Python package (cosmos_framework) with two primary functional paths:

  • Training: Uses a distributed trainer supporting FSDP, Tensor Parallelism (TP), Context Parallelism (CP), and Pipeline Parallelism (PP). It handles native DCP checkpoints and supports various dataset adapters including JSONL, WebDataset, and LeRobot.
  • Inference: Utilizes backends like Diffusers, Transformers, and vLLM to provide both offline batch generation and online serving via Ray and Gradio.

Who it’s for

It is designed for researchers and developers working on Physical AI, world models, and multimodal generative AI who need a scalable system to train and deploy models that process and generate multiple data types simultaneously.

Highlights

  • Omnimodal Support: Jointly processes and generates language, images, video, audio, and action sequences.
  • Unified Architecture: Uses a Mixture-of-Transformers architecture to subsume vision-language models and world simulators.
  • Scalable Training: Built-in support for advanced distributed training strategies (FSDP/TP/CP/PP).
  • Flexible Inference: Supports multiple backends and provides a policy server for action-based models.

Related

  • Project
    NVIDIA/cosmos

    An open platform of omnimodal world models and tools for building Physical AI, enabling the joint processing and generation of text, vision, audio, and action sequences for robotics and autonomous systems.

  • Project
    nvidia-cosmos/cosmos-rl

    A scalable Reinforcement Learning framework specialized for Physical AI applications, providing a decoupled architecture for large-scale training workloads.

  • Project
    huggingface/transformers

    A model-definition framework providing a unified API to access, train, and deploy state-of-the-art pretrained models across text, vision, audio, and multimodal domains.

  • Project
    EvolvingLMMs-Lab/lmms-engine

    A unified, high-performance training engine for multimodal models that integrates distributed training, kernel fusion, and advanced optimizers to enable scalable pretraining and fine-tuning.

  • Project
    nvidia-cosmos/cosmos-predict1

    A suite of world foundation models for future state prediction that generates visual simulations from text or video prompts to support Physical AI development.