1. Model introduction
Cosmos3 is NVIDIA’s Mixture-of-Transformers (MoT) omni family: an UND (understanding) tower and a GEN (generation) tower over a joint text+vision packed sequence, with the Wan2.2 VAE (4× temporal compression). This guide covers Cosmos3-Nano, the variant validated by the canonical recipe and training curve. Key highlights for RL training:- No separate text encoder. Conditioning is token-level:
CondKwargscarriestext_ids/text_mask/fpsverbatim, which eliminates the text-replay-consistency failure class other families guard against. - UND tower frozen inside the training graph. The UND tower participates in the packed forward, so it is frozen by
parameter-name fragments rather than dropped; LoRA targets are GEN attention only (
add_q_proj,add_k_proj,add_v_proj,to_add_out). - Packed single-sample forward. The transformer consumes one packed text+vision sequence per forward — one request
cannot batch multiple outputs, so recipes run
--rollout-microgroup-size 1and CFG batching is disabled by construction. - Karras flow-sigma grid. Checkpoints ship a non-uniform sigma grid; SDE candidate steps must be derived from it.
2. Supported variants
3. Family config
Frommiles/backends/fsdp_utils/configs/cosmos3.py:
4. Launch
Canonical recipe:scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py — train, rollout,
and PickScore colocated on 4 GPUs; T2I (832×480, 1 frame).
Status: 📈 V — Verified
5. Recipe notes
epoch_global_random_choice draws two steps per epoch from
--diffusion-sde-candidate-steps 8,9,10,11.
The Cosmos3 checkpoint’s Karras flow-sigma grid puts head steps 1–7 at
sigma > 0.96 with |dt| < 0.02; steps 8–11 are the useful high-noise segment.
Step numbers are not transferable across sigma-grid families: re-derive candidates
from |dt| when changing model or grid.
6. Reference results
The 4-GPU colocated Cosmos3-Nano recipe raises PickScore (rollout/reward/raw_mean) from ~0.77 to ~0.85 over 250 rollouts:
7. Pairs well with
- LoRA Training and Weight Sync — GEN-tower LoRA sync.
- Rewards — PickScore worker pool configuration.