Continuing on topic of Cosmos 3.0 I have also been comparing its training infra with some other world models but as a baseline this is the training setup. As I posted earlier the model is a 16B (Nano) or 64B (Super) Mixture-of-Transformers with a Reasoner tower and a Generator tower in every layer. Nano's Generator was pretrained on 31T tokens on 1,024 GB200 GPUs while Super on 18T tokens on 2,048.
The stack is plain PyTorch, released as cosmos-framework.
Two things have to be split across GPUs 1. Model weights 2. Activations for a 74,000-token sequence.
Weights are split w/ FSDP. Each GPU stores a slice of every layer's weights. When a layer runs, the GPUs gather that layer's weights, compute, and drop the gathered copy. At larger scale they switch to HSDP, which keeps the gathering inside one rack's NVLink domain and only averages gradients across racks.
Activations are split with context parallelism. Each GPU takes a slice of the 74,000 tokens. Layer norms and MLPs work on one token at a time, so they run on the slice with no communication.
Attention is the exception, because every token needs to see every other token. Cosmos 3 uses the Ulysses scheme for this, just before attention, one collective reshuffles data bringing whole sequence for few heads on one GPU.
They chose Ulysses over ring attention for one reason. Reasoner tokens attend causally to Reasoner tokens only. Generator tokens attend fully to both towers, and the two towers produce keys and values with different weights. With Ulysses, each GPU holds the whole sequence when attention runs, so it concatenates the two key/value streams and builds the mask locally, and the standard attention kernel runs unchanged.
Sequence packing is done at a fixed 74,000-token budget. Clips at 256p, 480p and 720p with different frame counts fill each sequence with no padding. Constant token count per step, constant shapes for torch.compile, and the same variable-length attention kernel serves both packing and the two-tower mask.
The Wan2.2 video tokenizer is frozen and since it runs on every batch it is compiled. It sees 45 shapes: 3 resolutions × 5 aspect ratios × 3 call modes. Every rank compiling all 45 at startup cost about 15 minutes per run. They partition the 45 graphs across ranks, each rank compiles one and writes the artifact to shared storage, every rank loads all 45.
Checkpointing is asynchronous through a dedicated process group, so checkpoint writes never share a communicator w/ training collectives. The save plan is computed on first save and reused which cut total training time by 4% for Nano and 9% for Super over a pure synchronous save plan
On serving front one checkpoint holds an autoregressive workload and an iterative denoising workload. They run on different engines, vLLM or TensorRT-LLM for the Reasoner and vLLM-Omni for the Generator, and the Reasoner output is computed once and cached across denoising steps because the Reasoner never reads Generator tokens.