Table of Contents

• Step 1 - Prepare data and config
• Step 2 — Prepare checkpoint
• Step 3 — Run training
  • Option A (recommended): the paired launch shell
    • Overriding the defaults
  • Option B: raw torchrun
• Outputs
• Export checkpoint to Hugging Face safetensors
• Config
  • Common Hydra tail overrides

Fine-tune a pre-trained Cosmos3 model on your own dataset using supervised fine-tuning (SFT). Tested on 8× H100 (80 GB).

Prerequisites:

• Setup
• Environment Variables
• FAQ — troubleshooting (OOM during SFT, defaults), common pitfalls.

The runnable artifacts (TOML recipes, paired launch shells, inference helpers) live in examples/.

Some datasets are license gated — visit the repository page and accept any terms, and authenticate with uvx hf@latest auth login (or set HF_TOKEN).

The per-recipe download commands below write to examples/data/<dataset>/ and examples/checkpoints/wan22_vae/Wan2.2_VAE.pth, which match the launcher's default $DATASET_PATH and $WAN_VAE_PATH. See Step 3 → Option A for how to override these defaults if you'd rather keep data on a different filesystem.

Select one of the following recipes:

BASE_CHECKPOINT_NAME=Cosmos3-Nano

# Defaults match the launcher (see Step 3 → Option A to override).
uvx hf@latest download --repo-type dataset nvidia/BridgeData2-Subset-Synthetic-Captions \
    --revision 40d018ac1c1a2a4b9734f17fdb21f3d933c49a01 \
    --local-dir examples/data/BridgeData2-Subset-Synthetic-Captions --quiet
uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth \
    --local-dir examples/checkpoints/wan22_vae --quiet

BASE_CHECKPOINT_NAME=Cosmos3-Super

# Defaults match the launcher (see Step 3 → Option A to override).
uvx hf@latest download --repo-type dataset nvidia/BridgeData2-Subset-Synthetic-Captions \
    --revision 40d018ac1c1a2a4b9734f17fdb21f3d933c49a01 \
    --local-dir examples/data/BridgeData2-Subset-Synthetic-Captions --quiet
uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth \
    --local-dir examples/checkpoints/wan22_vae --quiet

# No required env vars. The first launch will populate the HF Hub cache under
# $HF_HOME (defaults to /tmp/hf_cache inside the wrapper); subsequent launches
# reuse the cached snapshot.
#
# (optional) HF_TOKEN raises HF Hub rate limits for the streamed dataset
# revision lookup — useful if you're running 8-rank fan-out from a single IP:
# export HF_TOKEN=hf_...

# Step 1 (data): materialize the public HF dataset into the canonical local layout
# (videophy2_{train,val}/{meta.json, media/, text/}).
python -m cosmos_framework.scripts.vlm.prepare_videophy2_from_hf \
    --out_root examples/data/videophysics --split both

Convert the base checkpoint to PyTorch Distributed Checkpoint (DCP) format. cosmos_framework.scripts.convert_model_to_dcp ships in the unified cosmos_framework/ package, so this step runs from the repo root (with the environment activated per Setup).

Set BASE_CHECKPOINT_NAME to the value from the recipe block you picked in Step 1 (Cosmos3-Nano or Cosmos3-Super):

BASE_CHECKPOINT_NAME=Cosmos3-Nano   # or Cosmos3-Super — match the recipe in Step 1

# Default output dir matches the launcher (see Step 3 → Option A to override).
python -m cosmos_framework.scripts.convert_model_to_dcp \
  -o examples/checkpoints/$BASE_CHECKPOINT_NAME \
  --checkpoint-path $BASE_CHECKPOINT_NAME

$BASE_CHECKPOINT_NAME (e.g. Cosmos3-Nano, Cosmos3-Super) is a registered name in the checkpoint catalog; the converter downloads the matching repo from the Hugging Face Hub and writes the DCP into examples/checkpoints/$BASE_CHECKPOINT_NAME.

Reasoner Alignment SFT with LLaVA-OneVision (vfm-vlm): Skip this step — the Reasoner alignment SFT loads Qwen/Qwen3-VL-8B-Instruct from the HF Hub at startup (no DCP conversion, no env vars).

Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Nano): Use cosmos_framework.scripts.convert_model_to_vlm_safetensors instead.

# Step 2 (VLM checkpoint): merge Cosmos3-Nano LM onto the Qwen3-VL visual tower.
# Replaces the convert_model_to_dcp step used by the VFM recipes above.
python -m cosmos_framework.scripts.convert_model_to_vlm_safetensors \
    --checkpoint-path Cosmos3-Nano \
    -o examples/checkpoints/Cosmos3-Nano-VLM

Weights & Biases (optional): every recipe TOML defaults to job.wandb_mode = "disabled". To log a run to W&B, flip that field to "online" in the TOML and export WANDB_API_KEY in your environment before launching.

Each recipe ships as a examples/toml/sft_config/<recipe>.toml (validated against the pydantic schema at cosmos_framework/configs/toml_config/sft_config.py) paired with examples/launch_sft_<recipe>.sh; the full catalog is indexed in examples/README.md. Each .sh sources examples/_sft_launcher_common.sh and forwards into cosmos_framework.scripts.train --sft-toml=<recipe-toml>. From the repo root, run the launch shell paired with the recipe you set up in Step 1. The wrapper resolves DATASET_PATH, BASE_CHECKPOINT_PATH, and WAN_VAE_PATH from the default locations under examples/ (populated by Step 1 + Step 2), so no env-var setup is required (see below to override):

# from the repo root, after Step 1 + Step 2:
bash examples/launch_sft_vision_nano.sh

Each launcher's default paths come from the DATASET_PATH + BASE_CHECKPOINT_PATH defaults declared at the top of its .sh (each uses : "${VAR:=…}" so any value you export in the shell before launching wins over the default):

WAN_VAE_PATH defaults to examples/checkpoints/wan22_vae/Wan2.2_VAE.pth for every non-reasoner recipe.

Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Nano):

# Step 3 (launch): export both env vars, then launch.
export VIDEOPHYSICS_ROOT=$PWD/examples/data/videophysics
export VLM_SAFETENSORS_PATH=$PWD/examples/checkpoints/Cosmos3-Nano-VLM
bash examples/launch_sft_videophy2_nano.sh

If you'd rather put data or checkpoints on a different filesystem (e.g. a faster SSD or shared mount), download to your chosen path in Step 1 / convert the DCP to your chosen path in Step 2, then export the matching env var(s) before launching:

# Example: data on /scratch, base DCP on /nfs/ckpts.
export DATASET_PATH=/scratch/BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge
export BASE_CHECKPOINT_PATH=/nfs/ckpts/Cosmos3-Nano
export WAN_VAE_PATH=/nfs/ckpts/wan22_vae/Wan2.2_VAE.pth
bash examples/launch_sft_vision_nano.sh

Each env var falls back to its default if unset, so you only need to export the ones you're moving. The downloads / convert_model_to_dcp commands in Step 1 + Step 2 just need their --local-dir / -o argument pointed at the same path you export here. .gitignore excludes examples/data/ and examples/checkpoints/ so the multi-GB downloads aren't tracked when you keep the defaults.

If you'd rather not use the paired launch shell, invoke torchrun directly with the recipe's TOML. Unlike Option A, raw torchrun does not auto-resolve DATASET_PATH / BASE_CHECKPOINT_PATH / WAN_VAE_PATH from examples/ — they have to come from your shell:

• BASE_CHECKPOINT_PATH and WAN_VAE_PATH are read via ${oc.env:BASE_CHECKPOINT_PATH} / ${oc.env:WAN_VAE_PATH} at the TOML's [checkpoint].load_path / [model.tokenizer].vae_path keys.
• DATASET_PATH is read via ${oc.env:DATASET_PATH} inside the experiment-SKU Python (e.g. cosmos_framework/configs/base/experiment/sft/<recipe>.py), not in the TOML.

You have two options to fill them in (pick either, not both):

1. Export them in the shell before torchrun (whether they point at the default examples/... paths from Step 1+2 or your own overrides) — shown below.
2. Edit the TOML by hand — open examples/toml/sft_config/<recipe>.toml and replace the ${oc.env:BASE_CHECKPOINT_PATH} / ${oc.env:WAN_VAE_PATH} placeholders with literal paths. Useful if you want a self-contained TOML you can hand to a colleague or commit alongside an experiment record. (Hand-editing won't help for DATASET_PATH — that's resolved out of the experiment Python, so you must still export it.)

Run from the repo root (the directory containing pyproject.toml and examples/); the snippet uses $PWD to absolutize the relative paths.

# This example uses the vision_sft_nano recipe end-to-end (same recipe as
# Option A). To switch recipes, swap TOML_FILE + DATASET_PATH per the table in
# Option A, and Cosmos3-Nano → Cosmos3-Super on the LoRA / super recipes.
TOML_FILE="examples/toml/sft_config/vision_sft_nano.toml"

# Match the launcher's defaults — or substitute your own paths.
export DATASET_PATH="$PWD/examples/data/BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge"
export BASE_CHECKPOINT_PATH="$PWD/examples/checkpoints/Cosmos3-Nano"
export WAN_VAE_PATH="$PWD/examples/checkpoints/wan22_vae/Wan2.2_VAE.pth"

IMAGINAIRE_OUTPUT_ROOT=outputs/train PYTHONPATH=. \
torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \
    --sft-toml=$TOML_FILE

To resume from the latest in-progress checkpoint, point BASE_CHECKPOINT_PATH at the run's checkpoints/iter_<N>/ directory under $IMAGINAIRE_OUTPUT_ROOT/<project>/<group>/<name>/ (see Outputs below for the full layout).

Outputs land under $IMAGINAIRE_OUTPUT_ROOT/<project>/<group>/<name>/:

1. config.yaml, config.pkl: Finalized resolved config (YAML for inspection, pickle for re-instantiation).
2. launch_info.yaml, job_env.yaml: Job metadata and captured launch environment.
3. checkpoints/:
   1. latest_checkpoint.txt: Pointer file containing the latest checkpoint directory name (e.g. iter_000000200).
   2. iter_<iter>/: DCP checkpoint saved every [train.ckpt].save_freq iterations (zero-padded 9-digit, e.g. iter_000000200/):
      1. model/: model weights (sharded .distcp).
      2. optim/: optimizer state.
      3. scheduler/: LR scheduler state.
      4. trainer/: training state — includes the iteration counter and per-rank rng_state_<i> (numpy + random + torch + torch_cuda).
      5. dataloader/: optional per-rank pickle shards (rank_<i>.pkl) — only present for dataloaders that implement has_state().
4. <callback_name>/: Callback outputs, one directory per registered callback (e.g. DeviceMonitor/, EveryNDrawSample/, norm_monitor/).
5. wandb/, wandb_id.txt: Wandb run files — only present when [job].wandb_mode is online or offline.

The shorthand $RUN_DIR used in the rest of this page refers to $IMAGINAIRE_OUTPUT_ROOT/<project>/<group>/<name>. For example, with IMAGINAIRE_OUTPUT_ROOT=outputs/train and the vision_sft_nano recipe, $RUN_DIR is outputs/train/cosmos3/sft/vision_sft_nano.

Export the DCP checkpoint produced in Step 3 to a Hugging Face safetensors checkpoint:

RUN_DIR=$IMAGINAIRE_OUTPUT_ROOT/<project>/<group>/<name>

CHECKPOINT_ITER=$(cat $RUN_DIR/checkpoints/latest_checkpoint.txt)
CHECKPOINT_PATH=$RUN_DIR/checkpoints/$CHECKPOINT_ITER

python -m cosmos_framework.scripts.export_model \
  --checkpoint-path $CHECKPOINT_PATH \
  --config-file $RUN_DIR/config.yaml \
  -o $RUN_DIR/model

The exported safetensors land at $RUN_DIR/model and can be used in Inference commands by passing --checkpoint-path $RUN_DIR/model.

The recipe TOML is parsed against the pydantic schema SFTExperimentConfig at load time. Every top-level key listed below maps to a sub-model in that file; unknown keys raise a ValidationError before training starts (extra="forbid" on every sub-model). Values may use OmegaConf env interpolation ${oc.env:NAME} — the recipe TOMLs use this for BASE_CHECKPOINT_PATH ([checkpoint].load_path) and WAN_VAE_PATH ([model.tokenizer].vae_path). DATASET_PATH is consumed the same way but inside the experiment-SKU Python (cosmos_framework/configs/base/experiment/sft/<recipe>.py), not in the TOML.

For the full field-by-field reference (every section, every default, every VFM/VLM applicability note, the "???" MISSING sentinel, env interpolation, the VFM↔VLM path-remap table, and how to extend the schema), see SFT Structured-TOML Config Reference.

The commonly tuned knobs:

1. [job]
   1. task — "vfm" (generator recipes) or "vlm" (Reasoner alignment). Picks the base config: cosmos_framework/configs/base/config.py vs …/vlm/config.py. Also drives PATH_REMAPS in toml_config_helper.py.
   2. experiment — Registered experiment SKU name (e.g. vision_sft_nano). Each SKU is a Python file under cosmos_framework/configs/base/experiment/sft/ that wires up dataloader, model variant, and recipe-specific defaults.
   3. project, group, name — Components of the run output dir $IMAGINAIRE_OUTPUT_ROOT/<project>/<group>/<name>/. Also flow to W&B as the project / group / run name.
   4. wandb_mode — "online" (logs to W&B; WANDB_API_KEY must be set), "offline" (logs locally, sync later with wandb sync), or "disabled".
2. [model]
   1. max_num_tokens_after_packing — VFM token-packing target. -1 disables the cap. VFM only; VLM uses data_setting.max_tokens (tail override).
   2. joint_attn_implementation — VFM attention layout: "two_way" / "three_way" (NATTEN) / "flex".
   3. attn_implementation — VLM attention impl: "cosmos" / "flash_attention_2" / "sdpa" / "eager". VLM only.
   4. lora_enabled, lora_rank, lora_alpha, lora_target_modules — LoRA adapter knobs for the generation pathway. Used by SUPER-tier recipes; NANO-tier leaves lora_enabled=false. VFM only.
3. [model.ema]
   1. enabled, rate, iteration_shift — Exponential moving average of generation-pathway weights. Full fine-tunes typically enable it; LoRA recipes leave it off.
4. [model.parallelism]
   1. data_parallel_shard_degree — FSDP shard degree. data_parallel_shard_degree × data_parallel_replicate_degree × context_parallel_shard_degree must equal WORLD_SIZE. -1 autoselects from torchrun world size.
   2. data_parallel_replicate_degree — HSDP replicate degree (outer replicate loop over the shard topology).
   3. context_parallel_shard_degree — Context-parallel shard degree. >1 splits the sequence dim across ranks (used by super-tier configs: DP=4, CP=2 → 8 GPUs).
   4. cfg_parallel_shard_degree — Classifier-free-guidance shard degree. Almost always 1 for SFT.
   5. fsdp_master_dtype — Master parameter / FSDP reduce dtype: typically "float32".
5. [model.compile]
   1. enabled — Enable torch.compile. Improves speed at the cost of memory.
   2. compile_dynamic — Whether to compile with symbolic-shape (dynamic) kernels. True (default) is appropriate for training; AR inference may prefer False for stable shapes.
6. [model]
   1. precision — Compute dtype for forward/backward: "bfloat16" / "float16" / "float32". Master weights stay fp32 separately.
7. [model.activation_checkpointing]
   1. mode — "none" / "selective" (per-op SAC, MoT-only) / "full" (per-block checkpointing).
   2. save_ops_regex — Regex patterns for ops to keep saved under mode="selective".
   3. preserve_rng_state, determinism_check — Recompute determinism plumbing.
8. [model.tokenizer]
   1. vae_path — Wan2.2 VAE .pth path. Recipe TOMLs use "${oc.env:WAN_VAE_PATH}". VFM only.
9. [optimizer]
   1. lr — Base learning rate.
   2. betas, eps, fused, weight_decay — Standard AdamW knobs. eps is VFM-only.
   3. keys_to_select — Substring allowlist for trainable params. Empty list = train everything; ["lora_"] = adapter-only fine-tune.
10. [optimizer.lr_multipliers]
    1. Inline table of <substring> = <multiplier> pairs that scale the LR of params whose name contains the substring. The shipped vision recipes leave this empty (Hydra default {} stands).
11. [scheduler]
    1. cycle_lengths, warm_up_steps — Cycle length and warmup duration (lists, one entry per cycle), in optimizer steps.
    2. f_max, f_min, f_start — LR multipliers at peak / trough / step-0 (ratios of optimizer.lr).
    3. verbosity_interval — Scheduler-side LR log frequency. VFM only.
12. [trainer]
    1. max_iter — Total optimizer steps.
    2. grad_accum_iter — Micro-batches per optimizer step. Effective global batch = grad_accum_iter × per-rank batch × world_size.
    3. logging_iter — Console / W&B scalar log frequency.
    4. distributed_parallelism — "fsdp" is the only supported value.
13. [trainer.callbacks.compile_tokenizer]
    1. enabled, compile_after_iterations, warmup_resolutions — Lazy torch.compile of the VAE tokenizer. VFM only.
14. [trainer.callbacks.grad_clip]
    1. clip_norm — Max global L2 norm of the gradient (steps with larger norm are rescaled).
    2. force_finite — Replace NaN/Inf grads with zero (default true on VFM, false on VLM).
15. [checkpoint]
    1. load_path — Base DCP checkpoint directory to resume from (Step 2 output, or a prior run's checkpoints/iter_<N>/). Recipe TOMLs use "${oc.env:BASE_CHECKPOINT_PATH}".
    2. save_iter — Save a new DCP checkpoint every N optimizer steps.
    3. keys_to_skip_loading — Substring blocklist applied at load time. Used to mask EMA / LoRA tensors when warm-starting from a checkpoint that doesn't have them yet.
16. [dataloader_train] — Top-level scalars only; the dataloader's class (LazyCall) and pipeline wiring (datasets, packers, …) stay in the experiment Python.
    1. max_samples_per_batch — Per-micro-batch sample cap (remapped to max_batch_size on the VLM packer). null / omitted = no per-count cap.
    2. max_sequence_length — Per-packed-sequence token cap (remapped to max_tokens on the VLM packer).
    3. seed — Dataloader RNG seed (VFM only).

These knobs aren't part of the pydantic schema today; pass them as trailing key.path=value positionals after -- (the cosmos_framework.scripts.train flow forwards them through OmegaConf):

• model.config.policy.backbone.model_name — VLM backbone HF identifier (e.g. Qwen/Qwen3-VL-8B-Instruct). Used by launch_sft_llava_ov.sh.
• data_setting.max_tokens — VLM token-packing cap (the VLM analogue of [model].max_num_tokens_after_packing). Used by launch_sft_llava_ov.sh.

The launchers wire these via TAIL_OVERRIDES=(…); the helper appends -- "${TAIL_OVERRIDES[@]}" after the --sft-toml= argument.