A Physical AI system, such as a robot or autonomous vehicle (AV) that translates real-world data into physical actions, can’t be built in a single training job. Instead, it takes a continuous pipeline: a loop of generating synthetic data, post-training perception and policy models, so the system understands its surroundings and can act, and evaluating both in closed-loop simulation. Running that pipeline continuously is the job of a Physical AI model factory, turning a stream of new real-world data into better models, round after round.
This post shows how to build a Physical AI model factory with NVIDIA Cosmos 3 on Amazon SageMaker HyperPod, covering:
- What is unique about Cosmos 3: a Mixture-of-Transformers (MoT) design with per-layer joint attention and a deliberate train-versus-inference asymmetry.
- Why the design choices map cleanly onto Amazon SageMaker HyperPod with Amazon Elastic Kubernetes Service (Amazon EKS).
- Cluster and shared multi-terabyte storage layer setup.
- Distributed post-training for three representative workloads, with a complete end-to-end walkthrough of the robot-policy stage on a public DROID dataset.
The accompanying repository contains the manifests and configuration files for each stage. You can find the runnable code, including the infrastructure templates and job manifests that turn this design into a working cluster, in the awsome-distributed-ai GitHub repository.
Running the loop is a capacity commitment. Acquiring GPUs stage by stage adds variability at this scale: availability and lead times can vary, and the capacity that you do get might land in an Availability Zone or AWS Region away from your data. Committing capacity to the whole loop avoids that churn, whether through a flexible training plan for a bounded campaign or a capacity reservation for an open-ended one. Because you pay for that capacity whether or not the pipeline is making progress on it, the metric that governs cost is not the peak throughput of any one job. It is GPU goodput: the useful pipeline progress per reserved GPU-hour across the whole loop.
Physical AI pipelines often provision separate GPU capacity for each stage: one set of nodes to generate synthetic data, another to post-train, another to evaluate, each with its own lifecycle to stand up and tear down. NVIDIA Cosmos 3 makes that unnecessary. As an open omnimodal world foundation model, Cosmos 3 treats video, image, action, and sound as a single token stream. It runs the same transformer trunk in three modes: a forward-dynamics world model for synthetic video generation, an inverse-dynamics action labeler, and a deployable action policy. Because one model family covers generation, post-training, and evaluation, those stages become three workloads scheduled onto one persistent, resilient GPU node pool under a single cluster control plane. It’s time-shared capacity rather than a separate pool per stage. NVIDIA released it under the Linux Foundation’s OpenMDW-1.1 license, and describes the architecture in the Cosmos 3 technical report.
1. How Cosmos 3 works
A common pattern for world models is to pair a diffusion-transformer video generator with a separate vision-language model that provides text conditioning. Cosmos 3 takes a different approach: one trunk that handles both, integrated at every layer. That integration is what makes it useful as the engine of an end-to-end Physical AI model factory. Three architectural choices define it:
- One token stream. Every modality feeds into a single shared sequence, so one model can both read and generate across modalities. Images the model reads for understanding, pixels it generates, and a compact per-embodiment vector of pose deltas and grasp state each get their own encoder. A vision transformer (ViT) handles image understanding while a frozen Wan2.2 video variational autoencoder (VAE) handles pixel generation. That one action vector is what lets the same model drive both an AV and a robot arm. The sequence puts an autoregressive (AR) zone (the text and vision it reads) ahead of a diffusion zone (the video, audio, and action it produces).
- Two experts, joined at every layer (MoT). Each layer runs a reasoner that predicts the next token and a generator that denoises video, audio, and actions. Dual-stream attention joins them so generation stays grounded in the reasoner’s output at every layer, not only once at the end. The common alternative bolts a diffusion transformer (DiT) onto a vision-language model (VLM) and cross-attends to its final output once. Cosmos 3 grounds generation in the reasoner all the way down.
- Asymmetric at inference. Training and deployment do not run the same amount of work. Training runs the full denoising schedule and decodes video back to pixels, because that predicted video is part of the loss. On the robot, the same model runs a few denoise steps and skips video decoding entirely. The video latents are still produced internally to ground the action, but only the action tokens are decoded into the joint positions the robot executes.
The following diagram shows the first two choices in one view: the shared token stream and the two experts joined by attention. The autoregressive (AR) subsequence (text and the vision tokens the model reads to understand) and the diffusion-model (DM) subsequence (the video, audio, and action tokens it generates) run through the shared Reasoner and Generator towers. The attention mask on the right shows how the two experts differ: DM queries attend over both AR and DM keys (full attention), while AR queries stay causal and never see the diffusion tokens.
Figure 1: The Cosmos 3 shared token stream and its two experts joined by per-layer attention
(source: Cosmos 3: Omnimodal World Models for Physical AI)
Three action modes, one architecture
The mid-trained base checkpoint runs three jobs by changing which tokens start as noise. Post-training then specializes a checkpoint to a single mode and control frequency.
- Forward dynamics (world model). Actions clean, video noisy. “Given this frame and this action, what comes next?” This is the synthetic-data engine, fanned out to generate long-tail driving scenes or rare manipulation interactions that real collection cannot reach affordably.
- Inverse dynamics (action labeler). Video clean, action noisy. “Given these two frames, what action caused the change?” Converts unlabeled video (raw teleoperation recordings, third-person robot video, YouTube driving footage) into action-labeled training data.
- Policy (the deployed robot). Both noisy, conditioned on 3-view image plus proprioception. It outputs 32 future joint positions, with predicted video frames as a byproduct that grounds the action prediction.
The model family has two tiers: Cosmos3-Nano (16B parameters, on a dense 8B parameter Qwen3-VL backbone) and Cosmos3-Super (64B parameters, on a dense 32B parameter Qwen3-VL backbone). Task variants such as Cosmos3-Nano-Policy-DROID build on these tiers. NVIDIA also released Cosmos3-Edge, a compact 4B tier for on-device deployment (benchmarked on Jetson Thor and Orin). Edge shares the same physical-world pretraining data as Nano and Super but is built on a dense ~2B backbone trained from scratch rather than initialized from Qwen3-VL, so it is a separate weight lineage: you post-train Edge directly for the target hardware rather than shrinking a Nano checkpoint into it.
The following diagram shows the three modes side by side, with solid boxes for clean (known) tokens and dashed boxes for noisy tokens the model denoises. Forward dynamics keep the actions and the current frame clean and denoises future video. Inverse dynamics keep the video clean and denoises the actions. The policy sees only the first frame clean and denoises the actions the robot will execute. The same architecture runs all three, and only the pattern of clean versus noisy tokens changes. In the base checkpoint all three modes are available; a post-trained variant such as Cosmos3-Nano-Policy-DROID is specialized to policy mode at 15 Hz with a 32-step horizon.
Figure 2: The three action modes of one checkpoint, set by which tokens start as noise
(source: Cosmos 3: Omnimodal World Models for Physical AI)
2. From one model to a perpetual model factory
A team producing a robot or AV does not run one fine-tuning workload. It runs a loop: ingest real data, curate it, augment it with synthetic data, post-train, evaluate in a closed-loop simulation, deploy the policy, collect more data, and repeat.
Figure 3: The Physical AI model factory as a four-stage flywheel
The loop has four stages. (1) Ingest and curate real-world Physical AI data (DROID, BridgeData2, AV sensor logs) into a shared corpus on Amazon Simple Storage Service (Amazon S3) and Amazon FSx for Lustre. (2) A Cosmos3-Super teacher generates synthetic data to augment that corpus. (3) The combined synthetic and real corpus post-trains a deployable Cosmos3-Nano policy, with vision fine-tuning applied across both the Nano and Super tiers. (4) The policy is evaluated in closed-loop simulation, and its failures become new generation targets that re-enter the corpus for the next round.
Ideally, that loop doesn’t stop, with each stage running again as new data arrives. That cadence makes the cost driver GPU goodput (useful pipeline progress per reserved GPU-hour) rather than the peak throughput of any one job. Goodput is highest when the stages share one pool, so few GPU-hours are lost re-provisioning or moving data between separate clusters. Cosmos 3 makes that possible: it unifies three model classes (a world-sim generator, a policy, a perception model) into one model running in different modes. To support that flywheel, the cluster underneath must match that shape: one persistent pool on one control plane, instead of a disparate compute environment per job.
Amazon SageMaker HyperPod on Amazon EKS delivers exactly that shape. Each of the architectural choices behind Cosmos 3 creates a concrete cluster demand. The single token stream and the 64B MoT structure make training a long-sequence, multi-node job that needs a low-latency interconnect. The train-versus-inference asymmetry keeps generation, post-training, and evaluation on one model and one storage layer, so they can time-share one committed pool of capacity rather than fragmenting it stage by stage. Running the flywheel continuously requires capacity that is reserved and continuously monitored. Four Amazon SageMaker HyperPod properties answer those demands in turn:
- One cluster for all stages. Because Amazon SageMaker HyperPod orchestrates the cluster with EKS, the three engines of the loop run as ordinary Kubernetes workloads on a single shared GPU pool. Generation runs on the vLLM-Omni server, post-training on cosmos-framework under torchrun (Fully Sharded Data Parallel (FSDP2) plus Ulysses context parallelism), and evaluation on a single-GPU policy server. They also share one storage layer. An Amazon FSx for Lustre file system, accessed over Elastic Fabric Adapter (EFA), backed by an Amazon S3 bucket through a data repository association (DRA), mounts once and serves all three stages from the same path. Generation writes synthetic clips, post-training reads them, and the policy server loads its checkpoint off the same volume. There are no re-provisioning steps or terabyte-scale data migrations between stages, and Region-locked AV data stays in one in-Region cluster.
- Health-checked, auto-recovering capacity. A continuous loop wants capacity that is already provisioned and actively monitored: generation is bursty and dominates GPU-hours, and post-training runs for days across many nodes. Amazon SageMaker HyperPod continuously detects faulty nodes and reboots or replaces them automatically, and you can commit that capacity ahead of time with flexible training plans. Its managed job auto-resume then turns a worker failure into a bounded recovery. The Kubeflow PyTorchJob recreates the pod gang, NCCL re-forms, and cosmos-framework resumes from the latest PyTorch Distributed Checkpoint (DCP). A node failure therefore costs at most one checkpoint interval of redone work plus the node-replacement and reschedule latency, rather than a lost run.
- EFA already wired for multi-node NCCL. The long sequences that Cosmos 3 packs together for training (video latents plus text plus action, tens of thousands of tokens each) push the 64B tier into context parallelism on top of FSDP2. Every layer issues cross-node collectives. Standing that up by hand is the usual multi-node time sink: matching the EFA stack, the NCCL plugin, and the exact torch and NCCL versions the cosmos-framework pins. Amazon SageMaker HyperPod ships it pre-configured, and when paired with an AWS Deep Learning Containers (DLC) image, whose torch and aws-ofi-nccl versions match what the framework pins, NCCL over EFA is configured to work out of the box.
- Optional: task governance for many embodiments. If the factory serves several robot types or AV variants at once, Amazon SageMaker HyperPod task governance (built on Kueue) carves the pool into namespace-scoped queues with quotas, priorities, and preemptions. Dozens of heterogeneous jobs then share one capacity reservation instead of contending for it ad hoc, which raises goodput by keeping otherwise-idle GPUs busy across projects. A single-embodiment program can skip it, but task governance pays off once many jobs compete for the same pool.
The choice between Amazon SageMaker HyperPod and a lighter option comes down to the unit of work. A one-shot fine-tuning workload does not necessarily need the resilience and persistence of an Amazon SageMaker HyperPod cluster. An ephemeral managed training job (for example, an Amazon SageMaker AI training job) suffices, because a short run rarely hits a node failure. Amazon SageMaker HyperPod is well-suited for the sustained Cosmos 3 flywheel, where generation runs continuously, post-training is multi-node and long-running, evaluation is co-located on the same storage layer, and failures are statistically frequent.
3. What we are building
The solution post-trains three representative workloads, each a stage of the flywheel, and exercises the generation and evaluation stages end-to-end. All three run end-to-end on p5en.48xlarge (8x NVIDIA H200 GPUs) nodes with real checkpoints. The three are a robot-manipulation policy and two vision-perception fine-tuning workloads.
| Workload | Stage | Model | What it exercises |
| Robot policy (DROID) | Post-train (policy) | Cosmos3-Nano (~16B) | Action-policy post-training on a public LeRobot v3 dataset (droid_policy.toml). The lightest per-step workload |
| Vision Supervised Fine-Tuning (SFT) | Post-train (perception) | Cosmos3-Nano (~16B) | Video plus caption SFT (vision_sft_nano.toml). Substantially heavier per step than the policy workload |
| Vision Low-Rank Adaptation (LoRA) | Post-train (perception) | Cosmos3-Super (~64B) | LoRA fine-tuning of the 64B model with context parallelism (vision_sft_super.toml). The heaviest per-step workload |
Although AV post-training is not specifically covered here, Cosmos3’s base models were trained on a public synthetic-driving corpus (SDG-DriveSim, the nvidia/PhysicalAI-WorldModel-Synthetic-Autonomous-Driving-Scenarios dataset on Hugging Face). Its per-embodiment action projection is designed to extend to an AV ego-pose action space, so the same recipe and cluster setup covered in this post applies to AV post-training too.
The training stack is NVIDIA’s cosmos-framework, run with no forks or source edits to the framework package, so upstream updates drop in cleanly. It trains with FSDP2 and scales to hybrid sharded data parallelism (HSDP) and context parallelism as sequence length and node count grow. Section 6.2 covers how those parallelism choices are set per tier.
This guide uses p5en.48xlarge as the reference instance type t