Cosmos3

Cosmos3-Nano Flow-GRPO and PickScore recipe with token-level conditioning and packed single-sample forwards.

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: CondKwargs carries text_ids / text_mask / fps verbatim, 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 1 and CFG batching is disabled by construction.
• Shifted flow-sigma grid. sgl-d builds Cosmos3’s sigma grid as a linear grid warped by flow shift 3.0 for T2I (through its FlowUniPC serving scheduler, which overrides the checkpoint’s scheduler class). Rollout requests do not use UniPC: they take that grid verbatim on a FlowMatchEuler scheduler and step it with Euler (ODE steps) or the SDE kernel (SDE steps). Step sizes are non-uniform, so SDE candidate steps must be derived from the grid.

2. Validated variant

Only Nano has a validated recipe and training curve.

Size	Composition	HF checkpoints	Status
Nano	16 B (8 B UND + 8 B GEN)	Cosmos3-Nano	Validated — canonical recipe

3. Family config

From miles/backends/fsdp_utils/configs/cosmos3.py:

Property	Value	Why
Timestep dtype	fp32, no scaling	Grid timesteps are non-integer and sgl-d conditions on exact fp32 values — bf16 rounds 977.92 → 976
Cond dtype	pass-through	cond carries only token ids, the text mask and fps; the packed forward casts its own inputs
CFG batching	Off (asserted)	Packed forward is single-sample
LoRA targets	GEN attention (add_*_proj, to_add_out)	UND tower and unused sound/action heads stay frozen
Frozen params	Name-fragment allowlist (_GEN_PARAM_FRAGMENTS)	UND sits inside the graph and cannot be detached

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: 🛡️ FG — Fully gated

python3 scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py

The recipe sets SGLANG_DISABLE_COSMOS3_GUARDRAILS=1 in the job environment: RL scores raw samples, so the serving-side guardrail models are skipped.

5. Recipe notes

epoch_global_random_choice draws two steps from --diffusion-sde-candidate-steps 8,9,10,11, shared by each block of --rollout-batch-size consecutive prompt groups (one rollout, unless a dynamic-sampling filter drops groups).

On the 16-step rollout grid (linear, T2I flow shift 3.0, taken verbatim from the serving scheduler), steps 0–7 run σ ≈ 1.00 → 0.79 with |dt| ≈ 0.02–0.04, and the candidates 8–11 sit at σ ≈ 0.75 → 0.58 with |dt| ≈ 0.05–0.08. Step numbers are not transferable across grids: re-derive candidates from σ and |dt| when changing the model, the step count, or the shift.

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.