A NVIDIA Abre o Framework Cosmos para Treinar e Executar World Models

Quando desenvolvedores falam sobre redes multimodais, geralmente querem dizer uma combinação de texto e imagens. A NVIDIA foi além e lançou a família Cosmos 3 — world models onde texto, vídeo, áudio e sequências de ações físicas são combinados em uma única arquitetura Mixture-of-Transformers. Para trabalhar com essa família, a empresa lançou um repositório cosmos-framework.

O Que Está Dentro do Framework

O código é empacotado em um único pacote Python cosmos_framework. O projeto é projetado para engenheiros que precisam fazer fine-tuning de modelos multimodais para tarefas específicas ou implantá-los em produção. Se você está trabalhando com robótica ou geração de vídeo com comportamento controlável, este repositório vai economizar muito tempo escrevendo código boilerplate.

Toda a lógica é dividida em dois blocos principais.

O primeiro bloco lida com treinamento distribuído. A ferramenta suporta paralelismo FSDP, tensor (TP), contexto (CP) e pipeline (PP). Checkpoints são salvos no formato DCP, mas o código inclui conversão para o familiar HuggingFace safetensors. O preparo de dados é simplificado com adaptadores prontos para JSONL, WebDataset, HuggingFace e LeRobot. O script principal de execução está localizado em cosmos_framework.scripts.train.

O segundo bloco lida com inferência. Nos bastidores, você pode usar os backends Diffusers, Transformers, vLLM ou wrappers personalizados. Para geração em segundo plano, um modo offline está disponível, e para criação de API, uma stack Ray e Gradio é implantada. A inferência é iniciada através do módulo cosmos_framework.scripts.inference.

Para tarefas de robótica, o framework inclui um Policy Server. Ele serve o modelo Cosmos3-Policy-DROID e passa os comandos gerados diretamente para os atuadores.

Requisitos de Hardware e Primeira Execução

O apetite de hardware do projeto é substancial. Configs prontos da pasta examples foram testados em servidores com 8 placas NVIDIA H100 80GB. Você pode ajustar o grau de paralelismo e parâmetros de sharding através de NPROC_PER_NODE e configurações FSDP, mas treinar um modelo em uma única GPU de consumidor não será possível.

O projeto usa o gerenciador de pacotes uv para configuração do ambiente. Os autores recomendam começar com a imagem Docker oficial NVIDIA NGC (nvcr.io/nvidia/pytorch:25.09-py3).

A instalação de dependências do sistema fica assim:

sudo apt-get install -y --no-install-recommends curl ffmpeg git-lfs libx11-dev tree wget

Construir um ambiente virtual para CUDA 13.0 é feito com um único comando:

uv sync --all-extras --group=cu130-train
source .venv/bin/activate && export LD_LIBRARY_PATH=

Para testar inferência em uma única GPU, o repositório inclui um preset pronto:

python -m cosmos_framework.scripts.inference \
    --parallelism-preset=latency \
    -i "inputs/omni/t2v.json" \
    -o outputs/omni_nano \
    --checkpoint-path Cosmos3-Nano \
    --seed=0

Integração Automática com Agentes de Codificação

Um detalhe interessante: os autores construíram suporte a assistentes de IA na estrutura do repositório. O diretório .agents/skills/ contém instruções prontas para Cursor, Claude Code e Codex CLI.

Lá dentro você encontrará vários arquivos no formato SKILL.md. Cada instrução cobre sua própria área:

• Configuração do ambiente, download de pesos e trabalho com Docker
• Navegando pelos arquivos do pacote cosmos_framework
• Configuração de inferência e parâmetros de paralelismo
• Ciclo completo de fine-tuning SFT desde o preparo de dados até a exportação
• Diagnosticando erros CUDA, OOM e crashes de container

Quando você pedir ao Claude Code ou Cursor para corrigir um erro de import ou modificar um parâmetro no script de execução, o agente carrega automaticamente o arquivo de regras correspondente do diretório .agents/skills/.

Vale a Pena Experimentar

O framework será útil para equipes de robótica e desenvolvedores de sistemas de geração de vídeo. O repositório ainda é jovem, mas já inclui um conjunto completo de ferramentas para o fluxo de trabalho do modelo: desde o carregamento de datasets até o controle de hardware físico.

Se você tem acesso a clusters de GPU e está procurando um framework pronto para experimentar com world models, dê uma olhada em cosmos-framework. Scripts da pasta examples vão te dar um bom começo.