Construct a Bodily AI mannequin manufacturing unit with NVIDIA Cosmos 3 on SageMaker HyperPod

A Bodily AI system, equivalent to a robotic or autonomous automobile (AV) that interprets real-world information into bodily actions, can’t be inbuilt a single coaching job. As a substitute, it takes a steady pipeline: a loop of producing artificial information, post-training notion and coverage fashions, so the system understands its environment and may act, and evaluating each in closed-loop simulation. Working that pipeline repeatedly is the job of a Bodily AI mannequin manufacturing unit, turning a stream of recent real-world information into higher fashions, spherical after spherical.

This publish reveals the way to construct a Bodily AI mannequin manufacturing unit with NVIDIA Cosmos 3 on Amazon SageMaker HyperPod, overlaying:

• What is exclusive about Cosmos 3: a Combination-of-Transformers (MoT) design with per-layer joint consideration and a deliberate train-versus-inference asymmetry.
• Why the design selections 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 3 consultant workloads, with an entire end-to-end walkthrough of the robot-policy stage on a public DROID dataset.

The accompanying repository incorporates the manifests and configuration information for every stage. You could find the runnable code, together with the infrastructure templates and job manifests that flip this design right into a working cluster, within the awsome-distributed-ai GitHub repository.

Working the loop is a capability dedication. Buying GPUs stage by stage provides variability at this scale: availability and lead instances can range, and the capability that you simply do get would possibly land in an Availability Zone or AWS Area away out of your information. Committing capability to the entire loop avoids that churn, whether or not via a versatile coaching plan for a bounded marketing campaign or a capability reservation for an open-ended one. Since you pay for that capability whether or not or not the pipeline is making progress on it, the metric that governs value is just not the height throughput of anyone job. It’s GPU goodput: the helpful pipeline progress per reserved GPU-hour throughout the entire loop.

Bodily AI pipelines typically provision separate GPU capability for every stage: one set of nodes to generate artificial information, one other to post-train, one other to guage, every with its personal lifecycle to face up and tear down. NVIDIA Cosmos 3 makes that pointless. As an open omnimodal world basis mannequin, Cosmos 3 treats video, picture, motion, and sound as a single token stream. It runs the identical transformer trunk in three modes: a forward-dynamics world mannequin for artificial video era, an inverse-dynamics motion labeler, and a deployable motion coverage. As a result of one mannequin household covers era, post-training, and analysis, these levels turn out to be three workloads scheduled onto one persistent, resilient GPU node pool below a single cluster management aircraft. It’s time-shared capability relatively than a separate pool per stage. NVIDIA launched it below the Linux Basis’s OpenMDW-1.1 license, and describes the structure within the Cosmos 3 technical report.

1. How Cosmos 3 works

A standard sample for world fashions is to pair a diffusion-transformer video generator with a separate vision-language mannequin that gives textual content conditioning. Cosmos 3 takes a unique strategy: one trunk that handles each, built-in at each layer. That integration is what makes it helpful because the engine of an end-to-end Bodily AI mannequin manufacturing unit. Three architectural selections outline it:

• One token stream. Each modality feeds right into a single shared sequence, so one mannequin can each learn and generate throughout modalities. Photos the mannequin reads for understanding, pixels it generates, and a compact per-embodiment vector of pose deltas and grasp state every get their very own encoder. A imaginative and prescient transformer (ViT) handles picture understanding whereas a frozen Wan2.2 video variational autoencoder (VAE) handles pixel era. That one motion vector is what lets the identical mannequin drive each an AV and a robotic arm. The sequence places an autoregressive (AR) zone (the textual content and imaginative and prescient it reads) forward of a diffusion zone (the video, audio, and motion it produces).
• Two consultants, joined at each layer (MoT). Every layer runs a reasoner that predicts the following token and a generator that denoises video, audio, and actions. Twin-stream consideration joins them so era stays grounded within the reasoner’s output at each layer, not solely as soon as on the finish. The frequent different bolts a diffusion transformer (DiT) onto a vision-language mannequin (VLM) and cross-attends to its ultimate output as soon as. Cosmos 3 grounds era within the reasoner all the best way down.
• Uneven at inference. Coaching and deployment don’t run the identical quantity of labor. Coaching runs the complete denoising schedule and decodes video again to pixels, as a result of that predicted video is a part of the loss. On the robotic, the identical mannequin runs a couple of denoise steps and skips video decoding fully. The video latents are nonetheless produced internally to floor the motion, however solely the motion tokens are decoded into the joint positions the robotic executes.

The next diagram reveals the primary two selections in a single view: the shared token stream and the 2 consultants joined by consideration. The autoregressive (AR) subsequence (textual content and the imaginative and prescient tokens the mannequin reads to know) and the diffusion-model (DM) subsequence (the video, audio, and motion tokens it generates) run via the shared Reasoner and Generator towers. The eye masks on the appropriate reveals how the 2 consultants differ: DM queries attend over each AR and DM keys (full consideration), whereas AR queries keep causal and by no means see the diffusion tokens.

(supply: Cosmos 3: Omnimodal World Models for Physical AI)

Three motion modes, one structure

The mid-trained base checkpoint runs three jobs by altering which tokens begin as noise. Publish-training then specializes a checkpoint to a single mode and management frequency.

• Ahead dynamics (world mannequin). Actions clear, video noisy. “Given this body and this motion, what comes subsequent?” That is the synthetic-data engine, fanned out to generate long-tail driving scenes or uncommon manipulation interactions that actual assortment can not attain affordably.
• Inverse dynamics (motion labeler). Video clear, motion noisy. “Given these two frames, what motion induced the change?” Converts unlabeled video (uncooked teleoperation recordings, third-person robotic video, YouTube driving footage) into action-labeled coaching information.
• Coverage (the deployed robotic). Each noisy, conditioned on 3-view picture plus proprioception. It outputs 32 future joint positions, with predicted video frames as a byproduct that grounds the motion prediction.

The mannequin household has two tiers: Cosmos3-Nano (16B parameters, on a dense 8B parameter Qwen3-VL spine) and Cosmos3-Tremendous (64B parameters, on a dense 32B parameter Qwen3-VL spine). Activity variants equivalent to Cosmos3-Nano-Coverage-DROID construct on these tiers. NVIDIA additionally launched Cosmos3-Edge, a compact 4B tier for on-device deployment (benchmarked on Jetson Thor and Orin). Edge shares the identical physical-world pretraining information as Nano and Tremendous however is constructed on a dense ~2B spine educated from scratch relatively than initialized from Qwen3-VL, so it’s a separate weight lineage: you post-train Edge immediately for the goal {hardware} relatively than shrinking a Nano checkpoint into it.

The next diagram reveals the three modes aspect by aspect, with stable bins for clear (recognized) tokens and dashed bins for noisy tokens the mannequin denoises. Ahead dynamics hold the actions and the present body clear and denoises future video. Inverse dynamics hold the video clear and denoises the actions. The coverage sees solely the primary body clear and denoises the actions the robotic will execute. The identical structure runs all three, and solely the sample of fresh versus noisy tokens adjustments. Within the base checkpoint all three modes can be found; a post-trained variant equivalent to Cosmos3-Nano-Coverage-DROID is specialised to coverage mode at 15 Hz with a 32-step horizon.

Determine 2: The three motion modes of 1 checkpoint, set by which tokens begin as noise

(supply: Cosmos 3: Omnimodal World Models for Physical AI)

2. From one mannequin to a perpetual mannequin manufacturing unit

A staff producing a robotic or AV doesn’t run one fine-tuning workload. It runs a loop: ingest actual information, curate it, increase it with artificial information, post-train, consider in a closed-loop simulation, deploy the coverage, gather extra information, and repeat.

Determine 3: The Bodily AI mannequin manufacturing unit as a four-stage flywheel

The loop has 4 levels. (1) Ingest and curate real-world Bodily AI information (DROID, BridgeData2, AV sensor logs) right into a shared corpus on Amazon Easy Storage Service (Amazon S3) and Amazon FSx for Lustre. (2) A Cosmos3-Tremendous instructor generates artificial information to enhance that corpus. (3) The mixed artificial and actual corpus post-trains a deployable Cosmos3-Nano coverage, with imaginative and prescient fine-tuning utilized throughout each the Nano and Tremendous tiers. (4) The coverage is evaluated in closed-loop simulation, and its failures turn out to be new era targets that re-enter the corpus for the following spherical.

Ideally, that loop doesn’t cease, with every stage working once more as new information arrives. That cadence makes the fee driver GPU goodput (helpful pipeline progress per reserved GPU-hour) relatively than the height throughput of anyone job. Goodput is highest when the levels share one pool, so few GPU-hours are misplaced re-provisioning or shifting information between separate clusters. Cosmos 3 makes that doable: it unifies three mannequin courses (a world-sim generator, a coverage, a notion mannequin) into one mannequin working in several modes. To assist that flywheel, the cluster beneath should match that form: one persistent pool on one management aircraft, as a substitute of a disparate compute surroundings per job.

Amazon SageMaker HyperPod on Amazon EKS delivers precisely that form. Every of the architectural selections behind Cosmos 3 creates a concrete cluster demand. The only token stream and the 64B MoT construction make coaching a long-sequence, multi-node job that wants a low-latency interconnect. The train-versus-inference asymmetry retains era, post-training, and analysis on one mannequin and one storage layer, to allow them to time-share one dedicated pool of capability relatively than fragmenting it stage by stage. Working the flywheel repeatedly requires capability that’s reserved and repeatedly monitored. 4 Amazon SageMaker HyperPod properties reply these calls for in flip:

• One cluster for all levels. As a result of Amazon SageMaker HyperPod orchestrates the cluster with EKS, the three engines of the loop run as strange Kubernetes workloads on a single shared GPU pool. Technology runs on the vLLM-Omni server, post-training on cosmos-framework below torchrun (Absolutely Sharded Knowledge Parallel (FSDP2) plus Ulysses context parallelism), and analysis on a single-GPU coverage server. In addition they share one storage layer. An Amazon FSx for Lustre file system, accessed over Elastic Material Adapter (EFA), backed by an Amazon S3 bucket via a knowledge repository affiliation (DRA), mounts as soon as and serves all three levels from the identical path. Technology writes artificial clips, post-training reads them, and the coverage server hundreds its checkpoint off the identical quantity. There are not any re-provisioning steps or terabyte-scale information migrations between levels, and Area-locked AV information stays in a single in-Area cluster.
• Well being-checked, auto-recovering capability. A steady loop desires capability that’s already provisioned and actively monitored: era is bursty and dominates GPU-hours, and post-training runs for days throughout many nodes. Amazon SageMaker HyperPod repeatedly detects defective nodes and reboots or replaces them mechanically, and you’ll commit that capability forward of time with versatile coaching plans. Its managed job auto-resume then turns a employee failure right into a bounded restoration. The Kubeflow PyTorchJob recreates the pod gang, NCCL re-forms, and cosmos-framework resumes from the most recent PyTorch Distributed Checkpoint (DCP). A node failure subsequently prices at most one checkpoint interval of redone work plus the node-replacement and reschedule latency, relatively than a misplaced run.
• EFA already wired for multi-node NCCL. The lengthy sequences that Cosmos 3 packs collectively for coaching (video latents plus textual content plus motion, tens of 1000’s of tokens every) push the 64B tier into context parallelism on high of FSDP2. Each layer points cross-node collectives. Standing that up by hand is the standard multi-node time sink: matching the EFA stack, the NCCL plugin, and the precise torch and NCCL variations the cosmos-framework pins. Amazon SageMaker HyperPod ships it pre-configured, and when paired with an AWS Deep Studying Containers (DLC) picture, whose torch and aws-ofi-nccl variations match what the framework pins, NCCL over EFA is configured to work out of the field.
• Non-compulsory: activity governance for a lot of embodiments. If the manufacturing unit serves a number of robotic sorts or AV variants directly, Amazon SageMaker HyperPod activity governance (constructed on Kueue) carves the pool into namespace-scoped queues with quotas, priorities, and preemptions. Dozens of heterogeneous jobs then share one capability reservation as a substitute of contending for it advert hoc, which raises goodput by conserving otherwise-idle GPUs busy throughout initiatives. A single-embodiment program can skip it, however activity governance pays off as soon as many roles compete for a similar pool.

The selection between Amazon SageMaker HyperPod and a lighter choice comes right down to the unit of labor. A one-shot fine-tuning workload doesn’t essentially want the resilience and persistence of an Amazon SageMaker HyperPod cluster. An ephemeral managed coaching job (for instance, an Amazon SageMaker AI coaching job) suffices, as a result of a brief run hardly ever hits a node failure. Amazon SageMaker HyperPod is well-suited for the sustained Cosmos 3 flywheel, the place era runs repeatedly, post-training is multi-node and long-running, analysis is co-located on the identical storage layer, and failures are statistically frequent.

3. What we’re constructing

The answer post-trains three consultant workloads, every a stage of the flywheel, and workouts the era and analysis levels end-to-end. All three run end-to-end on p5en.48xlarge (8x NVIDIA H200 GPUs) nodes with actual checkpoints. The three are a robot-manipulation coverage and two vision-perception fine-tuning workloads.

Though AV post-training is just not particularly lined right here, Cosmos3’s base fashions have been educated on a public synthetic-driving corpus (SDG-DriveSim, the nvidia/PhysicalAI-WorldModel-Synthetic-Autonomous-Driving-Scenarios dataset on Hugging Face). Its per-embodiment motion projection is designed to increase to an AV ego-pose motion area, so the identical recipe and cluster setup lined on this publish applies to AV post-training too.

The coaching stack is NVIDIA’s cosmos-framework, run with no forks or supply edits to the framework package deal, so upstream updates drop in cleanly. It trains with FSDP2 and scales to hybrid sharded information parallelism (HSDP) and context parallelism as sequence size and node depend develop. Part 6.2 covers how these parallelism selections are set per tier.

This information makes use of p5en.48xlarge because the reference occasion kind all through. Relatively than publishing a cross-instance rating, we provide you with a goodput methodology you may run by yourself {hardware} (Part 7). It’s constructed on per-step time, GPU saturation, and a configurable Mannequin FLOPs Utilization (MFU). With it you may measurement your chosen platform, together with the NVIDIA Blackwell platform (B200, B300, and rack-scale GB200/GB300 NVL72), in opposition to your individual measurements.

4. Cluster setup on SageMaker HyperPod EKS

Organising the cluster breaks down into satisfying the stipulations, enabling NCCL over EFA, turning on deep well being checks and auto-recovery, selecting a base coaching picture, and staging and validating the end result.

4.1 Stipulations

You’ll be able to fulfill most of those stipulations with the Amazon SageMaker HyperPod-EKS Terraform modules. These modules provision the EKS-orchestrated Amazon SageMaker HyperPod cluster, the digital personal cloud (VPC) and EFA-enabled safety teams, the Amazon FSx for Lustre file system and CSI driver, the Kubeflow coaching operator, and the observability add-on. In the event you favor, it’s also possible to use the Amazon SageMaker AI console to create these assets utilizing AWS CloudFormation. Alternatively, you may carry your individual equivalents.

You want the next in place earlier than the primary job runs:

• Cluster. An Amazon SageMaker HyperPod cluster orchestrated by Amazon EKS, with a GPU occasion group of p5en.48xlarge nodes in a single Availability Zone (see Creating an Amazon SageMaker HyperPod cluster with Amazon EKS orchestration). At this scale, GPU capability is the binding constraint: p5en isn’t obtainable on demand, so plan on a versatile coaching plan or a capability reservation to safe the nodes.
• Service quota. As a result of GPU capability constrains availability, a enough service quota for the chosen occasion kind within the goal Area, requested via AWS Service Quotas earlier than you scale up the cluster.
• Job submission. kubectl configured in opposition to the cluster and the Kubeflow Coaching Operator put in, in order that PyTorchJob customized assets are acknowledged.
• Storage. The FSx for Lustre CSI driver put in and an Amazon FSx for Lustre file system hooked up in the identical VPC and subnet because the GPU nodes. The Terraform modules provision each whenever you activate the FSx module, or you may connect an current file system.
• Credentials. A Hugging Face learn token with the nvidia/Cosmos-Guardrail1 license accepted on the token’s account, saved as a Kubernetes secret named hf-token, as a result of the era and policy-serving paths pull this gated guardrail repository at startup.

4.2 Enabling NCCL over EFA

EFA provides NCCL a kernel-bypass, distant direct reminiscence entry (RDMA) succesful transport for multi-node collectives, and every p5en.48xlarge node advertises 16 EFA Community Interface Playing cards (NICs). On Amazon SageMaker HyperPod EKS the EFA drivers (from the Deep Studying AMI) and the EFA system plugin (pre-installed by the HyperPod service) are already in place. The pod spec solely requests vpc.amazonaws.com/efa assets alongside the GPUs (see the sample manifests). The coaching picture should carry an aws-ofi-nccl plugin constructed in opposition to the identical NCCL model the framework makes use of, which is strictly why this pattern builds on the AWS DLC (Part 4.4).

EFA being current on the {hardware} is just not the identical as NCCL truly utilizing it, so confirm the transport relatively than assume it. The coaching manifest runs a brief diagnostic preamble earlier than coaching begins, and with NCCL debug logging on, a wholesome multi-node run experiences EFA with GPUDirect RDMA as the chosen transport. A fallback to TCP seems as NET/Socket within the logs, which means the collectives are working over the unsuitable transport. To validate the interconnect end-to-end earlier than an actual run, run an ordinary multi-node NCCL check (for instance, all_reduce_perf) and make sure the achieved bus bandwidth (see the NCCL tests guide). For the precise diagnostic instructions and the complete log signature, see the repository README.

4.3 Deep well being checks and auto-recovery

A per-node health-monitoring agent repeatedly runs fundamental, passive checks (DCGM coverage violations, nvidia-smi errors, GPU-count validation), whereas deep well being checks (DCGM level-4 diagnostics and NCCL/EFA benchmarks) run when nodes be part of or the cluster is up to date. If you activate automated node restoration, a fault from any of those sources triggers Amazon SageMaker HyperPod to reboot or change the defective occasion, and auto-resume restarts the job from the final checkpoint as soon as the substitute is prepared.

4.4 Selecting the bottom coaching picture

Getting distributed coaching to run throughout nodes can eat a shocking quantity of setup time, so it’s price treating the base-image selection as a deliberate resolution relatively than an assumption. The selection hinges on one query: do NCCL collectives trip EFA throughout nodes, and do the cosmos-framework pinned CUDA wheels load? The framework’s digital surroundings (venv) pins torch==2.10.0+cu130 (CUDA 13), and its CUDA wheels (flash-attn, transformer-engine, natten) are revealed for CPython 3.13 solely. These pins drive the base-image selection in two methods: the picture’s NCCL should match the torch wheel’s bundled NCCL for EFA to work, and the picture should present a CPython 3.13 surroundings for the wheels to put in in any respect.

A mismatched base picture might block multi-node. A general-purpose GPU PyTorch base picture can validate single-node but fail cross-node NCCL over EFA at initialization (for instance, fi_getinfo() No information obtainable) even when EFA itself is totally purposeful. The foundation trigger is a version-matrix mismatch. The cosmos-framework venv’s torch bundles a particular NCCL (right here, 2.28.9), but when the bottom picture’s bundled aws-ofi-nccl plugin was constructed in opposition to a unique NCCL, the plugin and the runtime don’t line up. Setting NCCL_NET_PLUGIN=none sidesteps the error, however solely by dropping cross-node site visitors onto TCP as a substitute of EFA, a non-starter for multi-node efficiency.

A version-matched AWS Deep Studying Containers picture removes this work. The AWS Deep Learning Containers (DLC) for PyTorch ships torch 2.10.0+cu130, a precise match to the cosmos-framework pin. It additionally bundles an AWS tuned, version-matched EFA stack (EFA 1.47.0, libfabric 2.4, aws-ofi-nccl 1.18.0, GDRCopy 2.5.1). As a result of the DLC ships the identical torch wheel the venv installs, NCCL and the aws-ofi-nccl construct within the DLC line up. Multi-node EFA is then configured to work and not using a plugin rebuild or a model mismatch to work round.

Two additional build-time points surfaced within the DROID video-decode path on the DLC: an FFmpeg model too outdated for torchcodec, and a lacking shared libpython. Each are packaging issues with clear fixes baked into the Dockerfile within the accompanying repository.

In brief, decide the bottom picture by cosmos-framework model compatibility and verified NCCL over EFA, not by model or familiarity. For this framework model, the version-matched AWS DLC generally is a lower-effort path.

4.5 Staging and validation

With the bottom picture chosen and the prerequisite cluster in place, you stage the picture and storage and validate the end result earlier than working a workload. Every step is backed by code within the accompanying repository, so that you run templates relatively than hand-assemble assets.

Construct and push the coaching picture to Amazon Elastic Container Registry (Amazon ECR), and apply the storage class and the non-compulsory Amazon S3 information repository affiliation so datasets and base checkpoints hydrate into /fsx on first entry:

./build-push.sh
envsubst < storage/storage-fsx-efa-sc.yaml | kubectl apply -f -
envsubst < storage/storage-fsx-dra.yaml | kubectl apply -f -

Earlier than submitting a job, you validate that the cluster is prepared: verify that each GPU node reveals Prepared and the Kubeflow training-operator pod reveals Working.

kubectl get nodes
kubectl get pods -n kubeflow

With provisioning and validation performed, Part 6.1 walks via getting ready information, launching the robot-policy job, monitoring it, and validating its output, and Part 10 covers the way to tear the workloads and cluster again down whenever you’re completed.

5. Wiring the storage layer

The flywheel strikes multi-terabyte datasets between levels, so the storage layer is a first-class design resolution relatively than an afterthought.

Staging: Hugging Face to S3 to FSx for Lustre. Datasets and base checkpoints stage from Hugging Face into an in-Area Amazon S3 bucket, which is then hooked up to an FSx for Lustre filesystem via a DRA. FSx for Lustre presents one POSIX namespace at /fsx to every pod, and the DRA lazily hundreds objects from S3 on first entry or preloads them on demand.

Two I/O regimes. The workloads don’t stress storage the identical manner, so it helps to deal with them as two distinct regimes. The robot-policy information (the LeRobot/DROID dataset) is a metadata and small-file regime: many small Parquet shards and brief video clips, the place request charge and latency matter greater than uncooked bandwidth. The video SFT information is a bandwidth and large-file regime, the place sustained throughput dominates. As a result of the 2 regimes pull in several instructions, the appropriate backend is dependent upon the entry sample, not on a single best-choice verdict. Each regimes are served from one FSx for Lustre filesystem right here. Per-directory Lustre tuning (stripe depend and measurement, progressive file layouts, and client-side read-ahead) is an additional optimization you may layer on per entry sample relatively than one thing this pattern presets.

FSx for Lustre plus EFA, with chilly versus heat tradeoffs. Provision an EFA-enabled PERSISTENT_2 FSx for Lustre file system sized for a throughput goal (1000 MBps/TiB), as a result of capability governs the mixture throughput ceiling: the pattern’s 9.6 TiB filesystem tops out round 9.4 GB/s combination. EFA is what lifts the per-client ceiling. A non-EFA file system caps at 100 Gbps per consumer occasion, whereas an EFA-enabled file system reaches 700 Gbps per consumer over EFA. With GPUDirect Storage on EFA-enabled NVIDIA GPU situations equivalent to p5en, it reaches as much as 1200 Gbps. AWS recommends enabling EFA for any file system above 10 GBps for that reason. Deal with these as documented ceilings relatively than measured throughput: benchmark with fio (put in within the coaching picture) by yourself cluster, and notice that site visitors to a single object storage server (OST) caps at 5 Gbps, so a excessive per-client charge requires striping throughout many OSTs. The backends additionally differ most on uncooked throughput at first contact: native NVMe is quickest, FSx for Lustre sits beneath it, and a chilly learn from Amazon S3 is slower nonetheless. That hole is a cold-start and first-touch value. After a working set is page-cache heat, coaching reads are served from RAM, so the backend ceases to be the bottleneck. Backend selection subsequently issues for cold-start and for working units bigger than RAM, not for heat reuse of a cache-resident dataset. The cosmos-framework dataloader reinforces this: background staff prefetch and decode upcoming batches whereas the GPU computes the present step, so as soon as the working set is heat, steps keep compute-bound relatively than I/O-bound.

6. Launching distributed post-training

This part runs the robot-policy stage finish to finish, then explains how the job maps onto Kubernetes, how parallelism is configured on H200, and the way resilient checkpointing works.

6.1 Working the robot-policy stage finish to finish

Publish-training the robot-policy workload comes right down to a brief sequence: put together information and the bottom checkpoint, launch the distributed job, monitor it, and validate the output. The walkthrough makes use of the SageMaker HyperPod manifests, which render their surroundings variables with envsubst earlier than kubectl applies them. The subsections that observe clarify the equipment behind every step.

You first put together information and the bottom checkpoint. Level the Amazon FSx for Lustre DRA on the Amazon S3 bucket holding the general public DROID dataset, so the dataset hydrates into /fsx on first learn. The Hugging Face-released checkpoint ships as Diffusers/safetensors, so you change it as soon as to the DCP format the cosmos-framework hundreds (a CPU job is okay, although the 64B Tremendous wants a number of hundred GB of RAM):

python -m cosmos_framework.scripts.convert_model_to_dcp 
    --checkpoint-path Cosmos3-Nano 
    -o $BASE_CHECKPOINT_PATH

You then launch the distributed job by rendering and making use of its manifest, which wires the picture, the /fsx quantity, and the torchrun launch right into a Kubeflow PyTorchJob (Part 6.2):

envsubst < hyperpod-eks/train-multi-node-dlc.yaml | kubectl apply -f -

You monitor progress in two methods. kubectl reveals scheduling and pod well being, and the goodput dashboard from Part 7 reveals loss, per-step time, and GPU saturation because the run proceeds:

kubectl get pytorchjob cosmos3-droid-policy-hp
kubectl logs -f cosmos3-droid-policy-hp-worker-0

You validate output by confirming the run writes DCP checkpoints to $IMAGINAIRE_OUTPUT_ROOT on /fsx on the configured interval, and that coaching loss on the dashboard traits down as step time holds regular. As a result of the checkpoint lands on the shared quantity, the analysis server in Part 9 can load it immediately, which closes the loop.

6.2 How the job maps onto Kubernetes

The manifest launched in Part 6.1 wraps torchrun in a Kubeflow PyTorchJob, one main duplicate and N−1 employee replicas for an N-node run. The PyTorchJob controller injects the usual PyTorch rendezvous variables (coordinator deal with and port, per-pod rank, and world measurement) into every pod, and torchrun then spawns one course of per GPU to type the worldwide course of group over EFA. On a 2-node p5en run that is 16 ranks throughout 2 nodes.

6.3 Parallelism configuration on H200

The three post-training workloads break up by tier. The 16B Nano is totally fine-tuned into the deployable coverage, as a result of it’s the mannequin that ships to the robotic and is sufficiently small that full-parameter coaching is reasonably priced. The 64B Tremendous is the synthetic-data generator, and it’s tailored with LoRA relatively than totally retrained: you hardly ever must relearn a 64B instructor, solely shift it towards your area (your cameras, lighting, object courses, or state of affairs combine). Freezing the spine and coaching rank-16 adapters collapses optimizer and exponential-moving-average (EMA) reminiscence, shrinks checkpoints from a 64B snapshot to megabytes of adapter tensors, and lets one frozen Tremendous base serve many domains by swapping adapters. All of this raises goodput on the massive tier.

These tiers additionally drive the parallelism technique, with the H200’s 141 GB of high-bandwidth reminiscence (HBM) because the lever for minimizing communication. Extra HBM per GPU means much less aggressive sharding and fewer collectives per step. The Nano workloads run pure FSDP2 with a shard diploma set to the world measurement. Tremendous provides context parallelism diploma 2 on high of FSDP2, as a result of the long-packed sequences make consideration activation reminiscence the limiter. Ulysses context parallelism splits the sequence throughout GPUs with solely a few all-to-alls per consideration layer. At bigger node counts, set the replicate diploma above 1 to modify to HSDP when cross-cluster all-gather site visitors turns into the bottleneck.

The pattern additionally applies two small runtime patches at import time relatively than modifying the framework: a guard in opposition to an empty-shard edge case in gradient-norm monitoring at excessive rank counts, and the OpenTelemetry metrics bridge lined in Part 7. Each are utilized from the pattern, so upstream framework updates nonetheless drop in cleanly.

6.4 Resilient checkpointing

The cosmos-framework writes checkpoints as PyTorch Distributed Checkpoint (DCP), and the action-policy checkpoint conduct is preset within the checkpoint block of the action_policy_public_lerobot experiment. Three settings there matter for a resilient warm-start:

• dcp_async_mode_enabled=False, so saves are synchronous by default. Set it to True for async DCP, which retains steady-state checkpoint stall negligible whereas bounding misplaced work on failure.
• strict_resume=False, so a freshly initialized motion head (or LoRA adapters) can initialize whereas the remainder of the mannequin warm-starts from the transformed base checkpoint (Part 6.1).
• keys_to_skip_loading lists the tensors not anticipated throughout load (the motion heads and the bottom mannequin’s EMA weights), so that they initialize contemporary relatively than erroring out.

A fourth setting, ckpt_type, defaults to dcp in the identical experiment however is supposed to be overridden per run (for instance, dummy for a smoke check). Altering any of the previous three settings means modifying the experiment config, not passing a runtime flag.

Async DCP saves pin roughly model-size in host shared reminiscence, so set the pod’s /dev/shm quantity sizeLimit generously. For instance, 256Gi works nicely since P5en has roughly 2 TiB of host reminiscence. At 64Gi the save can throw an out of reminiscence (OOM) error.

Mixed with HyperPod auto-resume, a pod restart or node substitute continues from the final checkpoint mechanically. The coaching manifest additionally checks for an current checkpoint at startup and resumes from it relatively than ranging from scratch.

7. Measuring what issues: A goodput methodology you may run

This part provides you a metrics setup you may reproduce by yourself cluster. It builds on the Amazon SageMaker HyperPod EKS observability add-on and unifies GPU telemetry and cosmos-framework coaching metrics in a single Amazon Managed Grafana pane, utilizing a dashboard and bridge that ship with the accompanying repository. The dashboard is an importable Grafana mannequin, cosmos3-goodput-dashboard.json. To load it, open your Grafana workspace, select Dashboards, New, Import, and add the file, as described within the observability README.

Infrastructure and GPU metrics, by default. Streaming multiprocessor (SM) and HBM utilization, NCCL and EFA site visitors, and node well being all move into Amazon Managed Service for Prometheus via the DCGM exporter and node exporters, with no customized instrumentation. The Amazon SageMaker HyperPod observability add-on offers the cluster, node, and job views with out further configuration, rendered in Grafana alongside the cosmos-framework metrics.

Cosmos framework metrics, bridged into the native stack. The cosmos-framework emits loss, per-step timers, MFU, gradient norm, and sequence-packing statistics via its callbacks, which default to Weights & Biases (W&B). To land these in the identical Prometheus workspace because the GPU metrics with no W&B dependency, the repository ships a small OpenTelemetry Protocol (OTLP) bridge. The bridge mirrors these framework scalars to the Amazon SageMaker HyperPod observability add-on’s in-cluster OTLP collector (the hyperpod-otel-collector service, reached over gRPC). It additionally provides a Cosmos-specific MFU callback that exposes MFU in opposition to a peak-FLOPS fixed you configure. To show it on, set one surroundings variable on the coaching pod: set OTEL_EXPORTER_OTLP_ENDPOINT to the observability add-on’s in-cluster OTLP endpoint (http://hyperpod-otel-collector.hyperpod-observability.svc:4317). Leaving it unset retains the default path unchanged. The repository additionally ships the Grafana dashboard and OTLP bridge that render each metric sources collectively.

The next two panes are an illustrative seize from a single run, not a efficiency end result to learn numbers from. The primary pane is the cosmos-framework coach view. It reveals coaching loss, step time, achieved TFLOPS per GPU, Mannequin FLOPs Utilization (right here computed in opposition to the framework’s default per-GPU peak), iteration throughput, gradient norm, and sequence-packing token lengths, all bridged from the framework callbacks.

Determine 4: The cosmos-framework coach pane, bridged from the framework callbacks

The second pane is the GPU and infrastructure view from the Amazon SageMaker HyperPod observability add-on (utilizing the DCGM exporter): graphics-engine-active, GPU utilization, framebuffer (HBM) used, and GPU energy, reported per GPU so an idle or under-driven system is straight away seen.

Determine 5: The GPU and infrastructure pane from the observability add-on, utilizing the DCGM exporter

This setup surfaces two layers of sign, as methodology relatively than a broadcast rating:

• Micro layer: Per-GPU saturation and step effectivity. DCGM graphics-engine-active and HBM utilization come straight from the observability add-on as immediately measured {hardware} indicators. The cosmos-framework callbacks add achieved TFLOPS per GPU, iteration throughput, gradient norm, and sequence-packing token lengths. MFU can also be surfaced, computed in opposition to a configurable per-GPU peak-FLOPS fixed you set to your accelerator and precision.
• Macro layer: Goodput framing. Mix the micro-layer saturation with the goodput fraction to motive about efficient utilization throughout the entire flywheel, relatively than the height throughput of anyone step. The goodput fraction is the share of GPU-hours spent on helpful ahead progress, after initialization and scheduling, checkpoint stall, and restart and restoration.

7.1 How fault restoration works

On this stack, a single worker-pod or node failure throughout a long term resolves with out guide intervention, and the trail is price tracing as a result of it determines the price of a failure. After a node fault, the Amazon SageMaker HyperPod well being monitoring agent detects the unhealthy node, and its node restoration system reboots or replaces it. The PyTorchJob’s managed auto-resume (with the sagemaker.amazonaws.com/enable-job-auto-resume annotation) then recreates the pod gang onto wholesome capability. torchrun re-rendezvous, and cosmos-framework auto-resumes from the most recent checkpoint. Coaching continues from the final checkpoint relatively than from scratch.

Restoration decomposes into two buckets: (b1) node substitute and pod reschedule, and (b2) checkpoint reload plus catch-up to the failure level. Though b2 is framework-level and behaves the identical throughout platforms, the node auto-replacement and job auto-resume of Amazon SageMaker HyperPod deal with the b1 portion, the half that in any other case requires a human to detect, change, and restart. With b1 automated and b2 bounded by the checkpoint interval, a failure prices at most one checkpoint interval of redone work plus the b1 reschedule latency, relatively than a misplaced run, and it’s the foundation for optimizing checkpoint interval sizing.

7.2 Choosing the goodput-optimal checkpoint interval

The checkpoint interval is how typically coaching saves its state. Save too hardly ever and a failure throws away quite a lot of work. Save too typically and the saves themselves eat into helpful compute. The cosmos-framework checkpoints in an async mode: the GPUs copy the mannequin state to host reminiscence shortly, and coaching continues whereas the write to FSx finishes within the background. The steady-state value of a save is subsequently small, and a single-worker failure prices roughly one checkpoint interval of misplaced work plus the time to reschedule and reload.

There’s a known-optimal interval that balances these two prices, given by the Younger/Daly system: interval ≈ √(2 × C × MTBF), the place C is your checkpoint save value and MTBF is the imply time between failures. Each inputs are stuff you measure by yourself cluster relatively than guess. For a broadcast anchor, Meta’s Llama 3 405B pre-training noticed one interruption roughly each three hours on 16,384 H100 GPUs. Deal with it as an order-of-magnitude illustration, not a quantity to reuse, since it’s H100 at far bigger scale than this pattern’s p5en cluster. C is the save period you may learn from the cosmos-framework step metrics. MTBF is the common run-time between node failures noticed throughout your fleet, an estimate that sharpens as you accumulate run historical past.

For instance, with a 30-second save value and a 24-hour MTBF (86,400 s), the optimum interval is √(2 × 30 × 86,400) ≈ 2,300 s, or about each 40 minutes. At that interval, the checkpoint overhead works out to roughly 30 s of saving each 2,300 s, on the order of 1% of wall-clock, which is the goodput the interval is tuned to guard. Plug in your individual C and MTBF to get the interval that maximizes goodput in your cluster, and re-derive it if both adjustments. Observe MTBF is a fleet-level amount: for N nodes it’s roughly the per-node MTBF divided by N, so a 1024-GPU (128-node) cluster fails way more typically than a 16-GPU (2-node) cluster and desires a correspondingly shorter interval. Scale the per-node determine by your node depend earlier than fixing.

8. Findings and greatest practices

These findings are intentionally relative and regime-based (which workload is heaviest, the place a backend selection issues), so that they maintain by yourself cluster whatever the absolute numbers you measure there.

• Measured scaling and saturation. In our runs on p5en (H200) nodes, recorded with the goodput dashboard from Part 7, the Tremendous (64B) LoRA workload held near-flat strong-scaling effectivity because it scaled from 1–4 nodes (8–32 GPUs). It stayed at roughly 0.97–0.99 of linear, with per-step time inside about 3% throughout the ladder. Mannequin FLOPs Utilization, computed in opposition to the H200 BF16 peak, landed close to 0.50 for the compute-dense Tremendous workload and close to 0.24 for the lighter Nano imaginative and prescient workload, within the anticipated ordering (the bigger, extra compute-bound mannequin saturates the GPUs extra totally). These are relative figures from this setup that ought to reproduce in form on comparable {hardware}, not a leaderboard quantity.
• Node depend and parallelism per workload. Dimension the node depend to your wall-clock goal per workload, and let the workload’s per-step value drive the selection: the 64B Tremendous LoRA tier dominates per-step value, whereas the robot-policy workload is the lightest. As a result of the cosmos-framework packing dataloader holds a hard and fast per-rank token price range, per-step time, not iterations per hour, is the throughput sign to trace as you scale out. The parallelism technique additionally shifts from pure FSDP2 towards HSDP as cross-node all-gather site visitors grows (Part 6.2).
• Storage by entry sample. The metadata regime favors chilly and bulk staging from S3 and heat reuse from native NVMe. The bandwidth regime favors FSx plus EFA because the shared multi-stage aircraft. The backend issues at chilly first-touch, not for heat reuse of a cache-resident dataset.
• What managed resilience buys. Categorical the return as recovered time and goodput-fraction enchancment relatively than absolute {dollars}. Auto-replacement plus auto-resume removes the guide detect-replace-restart loop from each node failure, so the price of a failure drops to at most one checkpoint interval of redone work plus the node-replacement and reschedule latency.
• When does Amazon SageMaker HyperPod repay? For a one-shot small fine-tune, an ephemeral managed coaching job, equivalent to an Amazon SageMaker AI coaching job that provisions per run and tears down on completion, is enough as a result of auto-recovery hardly ever triggers. Amazon SageMaker HyperPod is well-suited for the sustained, at-scale flywheel: many concurrent jobs, steady artificial era, reserved capability, and the tens-of-nodes tail the place failures are statistically frequent.

9. The flywheel’s different levels: Technology and analysis

The identical cluster, picture, and storage layer run the era and analysis levels, which closes the loop.

Technology (write-bound). Technology runs on a separate picture from coaching: the official vllm/vllm-omni:cosmos3 engine (a cp312 / vLLM 0.23 stack), not the cosmos-framework DLC coaching picture. It comes up as an OpenAI-compatible server (vllm serve nvidia/Cosmos3-Tremendous --omni) that takes Cosmos3-Tremendous video-to-video (V2V) requests at POST /v1/movies/sync (see the generation manifest). After the server is Prepared, port-forward and POST a conditioning clip to get a generated continuation again:

curl -F input_reference=@clip.mp4 http://localhost:8000/v1/movies/sync --output generated.mp4

The returned clip is one artificial instance that feeds the following post-training spherical. That is the stage the place AV groups fan out long-tail driving eventualities (the identical class of safety-critical nook circumstances that corpora like SDG-DriveSim goal) to enhance actual fleet information. Technology makes use of its personal parallelism axes, Classifier-Free Steering (CFG) parallel x Ulysses x HSDP (--cfg-parallel-size, --ulysses-degree, --use-hsdp --hsdp-shard-size), which differ from the training-side FSDP2 + context-parallel knobs and should multiply to the per-node GPU depend. The server is single-node, so era wants no cross-node EFA and scales out as impartial servers. Guardrails (nvidia/Cosmos-Guardrail1) are toggled per request (extra_params.guardrails), not via a server flag, and the gated guardrail license should be accepted on the Hugging Face token’s account or startup fails. Delete the era Job when the batch is finished.

Analysis (latency-bound). A light-weight single-GPU Deployment serves a Cosmos 3 motion coverage over HTTP for closed-loop analysis, utilizing the identical FSx quantity so any checkpoint produced on the cluster is immediately obtainable (see the policy-serving manifest). The checkpoint should be an area listing (a naked org/repo Hugging Face id is just not resolved by the server), so pre-download or export your individual post-trained checkpoint to FSx first. A simulator’s management loop confirms the server is up with GET /information, then posts an statement to POST /predict every step and receives the following motion chunk. The request carries the present digicam body and the duty immediate:

# readiness / mannequin metadata
curl http://localhost:8000/information

# one management step: statement in, motion chunk out
curl -X POST http://localhost:8000/predict 
    -H "Content material-Kind: utility/json" 
    -d '{"picture": "<base64_png>", "immediate": "decide up the cup", "domain_name": "droid", "image_size": 256}'

The server returns the anticipated motion chunk (the 32 future joint positions described in Part 1), which the simulator applies earlier than sending the following statement. That request/response cycle is the closed loop: the checkpoint that post-training produced is similar one the analysis server hundreds off FSx, so a newly educated coverage will be evaluated with out shifting information between clusters.

10. Price concerns and cleanup

Amazon SageMaker HyperPod is a persistent cluster: situations are billed whereas they’re a part of the cluster, FSx for Lustre payments hourly per provisioned capability, and any serving or visualization pod holds a GPU node for so long as it runs. See the Amazon SageMaker AI and Amazon FSx for Lustre pricing pages for present charges. To pause occasion prices between periods whereas conserving the cluster configured, scale the GPU occasion group to zero, or delete the cluster fully. See Handle a SageMaker HyperPod cluster. Delete coaching workloads and the policy-serving Deployment when completed:

kubectl delete pytorchjob cosmos3-droid-policy-hp
kubectl delete -f hyperpod-eks/serve-policy.yaml

Deleting the FSx file system removes its native copy of checkpoints and logs. Knowledge written below the S3 DRA is exported again to the linked bucket and survives, however something exterior that path (or not but exported) is misplaced. Affirm your DRA has completed exporting (or obtain what you wish to hold) earlier than deleting and notice that the underlying S3 bucket persists independently and payments individually till you delete it too.

aws fsx delete-file-system --file-system-id <FSX-FILE-SYSTEM-ID>

11. Conclusion

As Bodily AI strikes into manufacturing, the problem shifts from a single coaching run to working the entire lifecycle repeatedly and economically. This publish described how you should utilize NVIDIA Cosmos 3 on Amazon SageMaker HyperPod (EKS) because the substrate for that flywheel. That substrate is a persistent, at-scale cluster with one shared storage layer for artificial era, post-training of coverage and notion fashions, and closed-loop analysis, with goodput relatively than single-job throughput because the metric that issues. We ran the robot-policy post-training stage on a public DROID dataset, exercised Nano and Tremendous imaginative and prescient fine-tuning throughout 1–4 nodes, validated multi-node EFA on a version-matched AWS Deep Studying Containers picture, and shipped a reproducible goodput dashboard on native observability that unifies cosmos-framework and GPU metrics in a single pane. In these runs, the Tremendous workload held near-flat strong-scaling effectivity (roughly 0.97–0.99 of linear) throughout that vary. As a result of each stage shares the identical cluster and storage, a checkpoint that post-training writes to Amazon FSx for Lustre is similar one the analysis server hundreds. The artificial clips that era produces land on the identical quantity the following post-training spherical reads.

This publish offers a reference structure and a reproducible methodology you may level at your individual embodiment and information, not a leaderboard you tackle religion. To get began, discover the awsome-distributed-ai GitHub repository, use the LeRobotV3ActionDataset recipe as a place to begin to adapt to your individual dataset, and prolong the flywheel to your individual robots or autos. To study extra, see the NVIDIA Cosmos website, the cosmos-framework repository, and the Amazon SageMaker HyperPod documentation.

In regards to the authors