Como construir uma fábrica de modelos de IA Física com NVIDIA Cosmos 3 no SageMaker HyperPod

O que é uma fábrica de modelos de IA Física?

Sistemas de Inteligência Artificial Física — como robôs ou veículos autônomos (VA) que traduzem dados do mundo real em ações físicas — não nascem de um único job de treinamento. Eles exigem um ciclo contínuo: gerar dados sintéticos, realizar pós-treinamento de modelos de percepção e política, e avaliar tudo isso em simulação de malha fechada. Esse ciclo contínuo é o que a AWS chama de Physical AI model factory (fábrica de modelos de IA Física): um pipeline que transforma um fluxo constante de dados reais em modelos progressivamente melhores.

A AWS publicou um guia técnico detalhado explicando como montar essa estrutura usando o NVIDIA Cosmos 3 no Amazon SageMaker HyperPod. O código completo, incluindo templates de infraestrutura e manifestos de job, está disponível no repositório awsome-distributed-ai no GitHub.

O que torna o NVIDIA Cosmos 3 diferente

O Cosmos 3 é um modelo de fundação omnimodal de mundo aberto, lançado pela NVIDIA sob a licença OpenMDW-1.1 da Linux Foundation. Sua arquitetura é descrita no relatório técnico do Cosmos 3 e se diferencia por três escolhas de design principais:

• Um único fluxo de tokens: todas as modalidades — vídeo, imagem, áudio e ações — alimentam uma sequência compartilhada. Um Transformador de Visão (ViT) trata da compreensão de imagens, enquanto um autoencoder variacional (VAE) de vídeo congelado cuida da geração de pixels. O mesmo vetor de ação compacto permite que o modelo controle tanto um veículo autônomo quanto um braço robótico.
• Dois especialistas por camada — Mistura de Transformadores (MoT): cada camada executa um “raciocinador” (que prevê o próximo token) e um “gerador” (que remove ruído de vídeo, áudio e ações). Atenção de fluxo duplo os conecta em cada camada, mantendo a geração ancorada na saída do raciocinador ao longo de todo o processo — não apenas no final.
• Assimetria entre treinamento e inferência: durante o treinamento, o modelo executa o cronograma completo de remoção de ruído e decodifica o vídeo até os pixels. No robô implantado, o mesmo modelo executa apenas alguns passos de remoção de ruído e ignora completamente a decodificação de vídeo. Os latentes de vídeo ainda são produzidos internamente para ancorar a ação, mas apenas os tokens de ação são decodificados nas posições de junta que o robô executa.

Três modos de operação, uma única arquitetura

O checkpoint base do Cosmos 3 opera em três modos distintos, definidos apenas por quais tokens começam como ruído:

• Dinâmica direta (modelo de mundo): ações limpas, vídeo ruidoso. Responde à pergunta “dado este frame e esta ação, o que vem a seguir?” — é o motor de geração de dados sintéticos para cenas raras ou de cauda longa.
• Dinâmica inversa (rotulador de ações): vídeo limpo, ação ruidosa. Converte vídeos não rotulados (gravações de teleoperação, vídeos de robôs em terceira pessoa, filmagens de direção) em dados de treinamento com rótulos de ação.
• Política (o robô implantado): ambos ruidosos, condicionados em imagem de 3 visões mais propriocepção. Produz 32 posições futuras de junta, com frames de vídeo previstos como subproduto que ancora a predição de ação.

A família de modelos tem dois níveis: Cosmos3-Nano (16B parâmetros, sobre um backbone Qwen3-VL denso de 8B) e Cosmos3-Super (64B parâmetros, sobre um backbone Qwen3-VL denso de 32B). A NVIDIA também lançou o Cosmos3-Edge, um nível compacto de 4B para implantação em dispositivos como Jetson Thor e Orin.

Do modelo único à fábrica perpétua de modelos

Uma equipe desenvolvendo um robô ou VA não executa um único job de fine-tuning. Ela opera um ciclo de quatro estágios:

• (1) Ingestão e curadoria de dados reais (DROID, BridgeData2, logs de sensores de VA) em um corpus compartilhado no Amazon S3 e Amazon FSx for Lustre.
• (2) Geração sintética: um professor Cosmos3-Super gera dados sintéticos para aumentar esse corpus.
• (3) Pós-treinamento: o corpus combinado pós-treina uma política Cosmos3-Nano implantável, com fine-tuning de visão aplicado em ambos os níveis Nano e Super.
• (4) Avaliação: a política é avaliada em simulação de malha fechada, e suas falhas se tornam novos alvos de geração que reingressam no corpus para a próxima rodada.

O que torna esse ciclo economicamente viável é o conceito de GPU goodput: o progresso útil do pipeline por hora de GPU reservada em todo o ciclo. O goodput é maximizado quando os estágios compartilham um único pool de GPUs, evitando horas perdidas com re-provisionamento ou migração de dados entre clusters separados.

Por que o SageMaker HyperPod é a infraestrutura ideal

O Amazon SageMaker HyperPod com Amazon Elastic Kubernetes Service (EKS) entrega exatamente a forma que esse ciclo exige. Quatro propriedades são centrais:

• Um cluster para todos os estágios: geração roda no servidor vLLM-Omni, pós-treinamento no cosmos-framework via torchrun com Paralelismo de Dados Totalmente Fragmentado (FSDP2) mais paralelismo de contexto Ulysses, e avaliação em um servidor de política de GPU única — tudo como workloads Kubernetes comuns em um único pool de GPUs compartilhado.
• Camada de armazenamento unificada: um sistema de arquivos Amazon FSx for Lustre, acessado via Elastic Fabric Adapter (EFA) e vinculado a um bucket S3 por uma associação de repositório de dados (DRA), é montado uma vez e serve todos os três estágios no mesmo caminho. Sem migrações de dados em escala de terabytes entre estágios.
• Capacidade com verificação de saúde e auto-recuperação: o HyperPod detecta nós com falha continuamente e os reinicializa ou substitui automaticamente. O auto-resume gerenciado transforma uma falha de worker em uma recuperação limitada: o PyTorchJob recria o gang de pods, o NCCL se reforma, e o cosmos-framework retoma do último checkpoint.
• EFA pré-configurado para NCCL multi-nó: as longas sequências do Cosmos 3 (latentes de vídeo mais texto mais ação, dezenas de milhares de tokens cada) exigem paralelismo de contexto sobre o FSDP2. O HyperPod vem com a pilha EFA pré-configurada, e quando combinado com uma imagem AWS Deep Learning Containers (DLC), o NCCL sobre EFA funciona sem configuração adicional.

Configuração do cluster e da camada de armazenamento

A maioria dos pré-requisitos pode ser satisfeita com os módulos Terraform do HyperPod-EKS, que provisionam o cluster HyperPod orquestrado por EKS, a Nuvem Privada Virtual (VPC) com grupos de segurança habilitados para EFA, o sistema de arquivos FSx for Lustre e o driver CSI, o operador de treinamento Kubeflow e o add-on de observabilidade. Alternativamente, é possível usar o console do Amazon SageMaker AI para criar esses recursos via AWS CloudFormation.

Antes de executar o primeiro job, é necessário ter em mãos: um cluster HyperPod com grupo de instâncias p5en.48xlarge em uma única Zona de Disponibilidade; cota de serviço suficiente para o tipo de instância escolhido, solicitada via AWS Service Quotas; o driver CSI do FSx for Lustre instalado; e um token de leitura do Hugging Face com a licença nvidia/Cosmos-Guardrail1 aceita.

Escolha da imagem base de treinamento

A escolha da imagem base é uma decisão deliberada, não uma suposição. O ambiente virtual do cosmos-framework fixa torch==2.10.0+cu130 (CUDA 13), e suas dependências CUDA são publicadas apenas para CPython 3.13. Uma imagem incompatível pode bloquear o NCCL multi-nó: o plugin aws-ofi-nccl da imagem precisa ter sido construído contra a mesma versão de NCCL que o framework usa.

A imagem AWS Deep Learning Containers para PyTorch entrega torch 2.10.0+cu130 — correspondência exata com o pin do cosmos-framework — e inclui uma pilha EFA ajustada pela AWS e com versões compatíveis (EFA 1.47.0, libfabric 2.4, aws-ofi-nccl 1.18.0, GDRCopy 2.5.1). Dois problemas adicionais surgiram no caminho de decodificação de vídeo do DROID na DLC: uma versão do FFmpeg antiga demais para o torchcodec e uma libpython compartilhada ausente. Ambos têm correções incorporadas ao Dockerfile do repositório.

Camada de armazenamento: dois regimes de I/O

Os workloads não estressam o armazenamento da mesma forma. Os dados de política robótica (dataset LeRobot/DROID) representam um regime de metadados e arquivos pequenos, onde taxa de requisições e latência importam mais do que largura de banda bruta. Já os dados de Ajuste Fino Supervisionado (SFT) de visão representam um regime de largura de banda e arquivos grandes, onde o throughput sustentado domina.

Para o FSx for Lustre com EFA, um sistema de arquivos PERSISTENT_2 sem EFA limita a 100 Gbps por instância cliente, enquanto um sistema de arquivos FSx for Lustre habilitado para EFA alcança 700 Gbps por cliente via EFA. Com GPUDirect Storage em instâncias GPU NVIDIA habilitadas para EFA como a p5en, chega a até 1200 Gbps. A AWS recomenda habilitar EFA para qualquer sistema de arquivos acima de 10 GBps. O ajuste por diretório do Lustre — incluindo contagem e tamanho de stripe, layouts de arquivo progressivos e read-ahead do lado do cliente — é uma otimização adicional que pode ser aplicada por padrão de acesso.

Pós-treinamento distribuído: três workloads representativos

A solução pós-treina três workloads representativos, cada um correspondendo a um estágio do ciclo, todos rodando em nós p5en.48xlarge (8x NVIDIA H200 GPUs):

• Política robótica (DROID): pós-treinamento de política de ação no dataset público LeRobot v3 (droid_policy.toml). O workload mais leve por passo.
• Ajuste Fino Supervisionado (SFT) de visão: SFT de vídeo mais legenda no Cosmos3-Nano (vision_sft_nano.toml). Substancialmente mais pesado por passo.
• Adaptação de Baixo Rank (LoRA) de visão: fine-tuning LoRA do modelo de 64B com paralelismo de contexto no Cosmos3-Super (vision_sft_super.toml). O workload mais pesado por passo.

O Nano de 16B é totalmente fine-tunado porque é o modelo que vai para o robô e é pequeno o suficiente para treinamento de parâmetros completos. O Super de 64B é adaptado com LoRA em vez de retreinado completamente: congelar o backbone e treinar adaptadores de rank-16 reduz a memória do otimizador, encolhe os checkpoints de um snapshot de 64B para megabytes de tensores de adaptadores, e permite que uma única base Super congelada sirva a muitos domínios apenas trocando adaptadores.

Checkpointing resiliente

O cosmos-framework escreve checkpoints como Checkpoint Distribuído PyTorch (DCP). Três configurações no bloco de checkpoint do experimento action_policy_public_lerobot são importantes para um warm-start resiliente: dcp_async_mode_enabled, strict_resume e keys_to_skip_loading. O modo DCP assíncrono mantém o overhead de checkpoint em estado estacionário negligível enquanto limita o trabalho perdido em caso de falha. Combinado com o auto-resume do HyperPod, uma reinicialização de pod ou substituição de nó continua automaticamente a partir do último checkpoint.

Medindo o que importa: metodologia de goodput

A AWS descreve uma configuração de métricas reproduzível construída sobre o add-on de observabilidade do HyperPod EKS, unificando telemetria de GPU e métricas de treinamento do cosmos-framework em um único painel do Amazon Managed Grafana. O dashboard cosmos3-goodput-dashboard.json pode ser importado conforme descrito no README de observabilidade.

Métricas de infraestrutura e GPU fluem para o Amazon Managed Service for Prometheus via exportadores DCGM e de nó, sem instrumentação customizada. O add-on de observabilidade do SageMaker HyperPod fornece as visões de cluster, nó e job. As métricas do cosmos-framework são conectadas via uma ponte OpenTelemetry Protocol (OTLP) que espelha os escalares do framework para o coletor OTLP in-cluster do add-on de observabilidade (o serviço hyperpod-otel-collector, acessível em http://hyperpod-otel-collector.hyperpod-observability.svc:4317 via gRPC). O repositório inclui o dashboard Grafana e a ponte OTLP que renderizam ambas as fontes de métricas juntas.

Intervalo de checkpoint ótimo para goodput

A fórmula de Young/Daly fornece o intervalo ótimo que equilibra os dois custos — salvar raramente demais desperdiça trabalho em falhas; salvar com frequência demais consome computação útil:

intervalo ≈ √(2 × C × MTBF)

Onde C é o custo de salvamento do checkpoint e MTBF é o tempo médio entre falhas. Por exemplo, com um custo de salvamento de 30 segundos e um MTBF de 24 horas (86.400 s), o intervalo ótimo é aproximadamente 2.300 s, ou cerca de 40 minutos. O overhead de checkpoint fica em torno de 1% do tempo de parede. Como referência de ordem de grandeza, o pré-treinamento do Llama 3 405B da Meta registrou uma interrupção a cada três horas em 16.384 GPUs H100 — mas trate isso apenas como ilustração, não como número a reutilizar. O MTBF é uma quantidade de nível de frota: para N nós, é aproximadamente o MTBF por nó dividido por N.

Os outros estágios do ciclo: geração e avaliação

O mesmo cluster, imagem e camada de armazenamento executam os estágios de geração e avaliação, fechando o ciclo.

Geração (limitada por escrita): a geração roda em uma imagem separada do treinamento — o motor oficial vllm/vllm-omni:cosmos3 — como um servidor compatível com OpenAI que aceita requisições de vídeo-para-vídeo (V2V) do Cosmos3-Super. Após o servidor estar pronto, basta fazer um port-forward e enviar um clipe de condicionamento para receber uma continuação gerada. O clipe retornado é um exemplo sintético que alimenta a próxima rodada de pós-treinamento. Consulte o manifesto de geração para os detalhes de configuração.

Avaliação (limitada por latência): um Deployment leve de GPU única serve uma política de ação do Cosmos 3 via HTTP para avaliação em malha fechada, usando o mesmo volume FSx para que qualquer checkpoint produzido no cluster esteja diretamente disponível. O loop de controle do simulador confirma que o servidor está ativo com GET /info, depois envia uma observação para POST /predict a cada passo e recebe o próximo chunk de ação. Consulte o manifesto de serving de política para os detalhes.

Considerações de custo e limpeza

O SageMaker HyperPod é um cluster persistente: instâncias são cobradas enquanto fazem parte do cluster, o FSx for Lustre cobra por hora por capacidade provisionada, e qualquer pod de serving ou visualização mantém um nó de GPU ativo enquanto estiver rodando. Consulte as páginas de preços do Amazon SageMaker AI e do Amazon FSx for Lustre para valores atuais. Para pausar os custos de instância entre sessões mantendo o cluster configurado, basta escalar o grupo de instâncias de GPU para zero. Veja como gerenciar um cluster SageMaker HyperPod para detalhes.

Ao deletar o sistema de arquivos FSx, a cópia local de checkpoints e logs é removida. Dados escritos sob o DRA do S3 são exportados de volta ao bucket vinculado e sobrevivem, mas qualquer dado fora desse caminho (ou ainda não exportado) é perdido.

Resultados e boas práticas

Nos testes realizados em nós p5en (H200), o workload LoRA do Super (64B) manteve eficiência de escalonamento forte quase linear ao escalar de 1 a 4 nós (8 a 32 GPUs), ficando em aproximadamente 0,97–0,99 do linear. A Utilização de FLOPs do Modelo (MFU) ficou próxima de 0,50 para o workload Super (mais intensivo computacionalmente) e próxima de 0,24 para o workload de visão Nano (mais leve).

As principais lições práticas são: dimensione a contagem de nós pelo seu objetivo de tempo de execução por workload; rastreie o tempo por passo (não iterações por hora) como sinal de throughput; escolha o backend de armazenamento pelo padrão de acesso, não por uma preferência única; e expresse o retorno da resiliência gerenciada como tempo recuperado e melhoria na fração de goodput.

Por onde começar

Para explorar a solução na prática, acesse o repositório awsome-distributed-ai no GitHub e use a receita LeRobotV3ActionDataset como ponto de partida para adaptar ao seu próprio dataset. Para aprofundar o conhecimento, consulte o site do NVIDIA Cosmos, o repositório cosmos-framework e a documentação do Amazon SageMaker HyperPod.

Fonte

Build a Physical AI model factory with NVIDIA Cosmos 3 on SageMaker HyperPod (https://aws.amazon.com/blogs/machine-learning/build-a-physical-ai-model-factory-with-nvidia-cosmos-3-on-sagemaker-hyperpod/)