# Post-Training (Supervised Fine-Tuning) Fine-tune a pre-trained Cosmos3 model on your own dataset using supervised fine-tuning (SFT). Tested on 8× H100 (80 GB). Prerequisites: - [Setup](https://github.com/NVIDIA/cosmos-framework/blob/main/README.md#setup) - [Environment Variables](environment_variables.md) - [FAQ](faq.md) — troubleshooting (OOM during SFT, defaults), common pitfalls. The runnable artifacts (TOML recipes, paired launch shells, inference helpers) live in [`examples/`](https://github.com/NVIDIA/cosmos-framework/blob/main/examples/README.md). ## Step 1 - Prepare data and config 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//` 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](#option-a-recommended-the-paired-launch-shell) for how to override these defaults if you'd rather keep data on a different filesystem. Select one of the following recipes:
Vision SFT (Cosmos3-Nano) T2V/I2V/V2V SFT on [nvidia/BridgeData2-Subset-Synthetic-Captions](https://huggingface.co/datasets/nvidia/BridgeData2-Subset-Synthetic-Captions/tree/main). `$DATASET_PATH` should be the directory containing `train/video_dataset_file.jsonl`. Each clip carries a structured-JSON caption (`caption_json`) — the model's native prompt format — which the SFT loader trains on by default (the dense narrative is kept as a backup), so training stays aligned with [Inference](dataset_jsonl.md#inference); see [JSONL Dataset → Format](dataset_jsonl.md#format). Launch shell: `examples/launch_sft_vision_nano.sh` ```shell 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 ```
Vision SFT LoRA (Cosmos3-Super) LoRA SFT on Qwen3-VL-32B MoT (Cosmos3-Super), on the same Bridge dataset as **Vision SFT (Cosmos3-Nano)**. Step 2 must convert the Cosmos3-Super checkpoint, not Cosmos3-Nano. Launch shell: `examples/launch_sft_vision_super.sh` ```shell 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 ```
Vision SFT (Cosmos3-Edge) Full fine-tune of the Nemotron-2B-Dense-VL backbone (Cosmos3-Edge), on the same Bridge dataset as **Vision SFT (Cosmos3-Nano)**. Step 2 must convert the Cosmos3-Edge checkpoint, not Cosmos3-Nano. Launch shell: `examples/launch_sft_vision_edge.sh` ```shell BASE_CHECKPOINT_NAME=Cosmos3-Edge # 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 ```
Reasoner Alignment SFT with LLaVA-OneVision (vfm-vlm) Alignment SFT for the Reasoner variant on the [lmms-lab/LLaVA-OneVision-Data](https://huggingface.co/datasets/lmms-lab/LLaVA-OneVision-Data) dataset (streamed from HF Hub). Skips Step 2: by default the backbone `Qwen/Qwen3-VL-8B-Instruct` is fetched from the HF Hub by the model downloader at startup — no DCP conversion needed and no required env vars. To instead start from a merged Cosmos3 reasoner snapshot (Cosmos3-Nano LM merged onto the Qwen3-VL visual tower), build it with `convert_model_to_vlm_safetensors` (see [Step 2](#step-2--prepare-checkpoint)) and point `VLM_SAFETENSORS_PATH` at it — same mechanism as the VideoPhy-2 recipe below. Launch shell: `examples/launch_sft_llava_ov.sh` ```shell # 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_... # # (optional) VLM_SAFETENSORS_PATH starts training from a local pre-converted # Qwen3-VL safetensors snapshot (e.g. Cosmos3-Nano LM merged with the Qwen3-VL # visual tower) instead of the public HF backbone: # export VLM_SAFETENSORS_PATH=$PWD/examples/checkpoints/Cosmos3-Nano-VLM ```
Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Nano) Reasoner alignment SFT for 1–5 physical-plausibility scoring on [videophysics/videophy2_train](https://huggingface.co/datasets/videophysics/videophy2_train) (HF test split renamed to `videophy2_val/`). `[job].task = "vlm"`. Bootstraps from `Cosmos3-Nano`'s language-model weights merged onto the public Qwen3-VL-8B-Instruct visual tower; the merged HF directory is consumed via `[model.backbone].safetensors_path` (plumbed by `VLM_SAFETENSORS_PATH`). Launch shell: `examples/launch_sft_videophy2_nano.sh` ```shell # 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.reasoner.prepare_videophy2_from_hf \ --out_root examples/data/videophysics --split both ```
Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Super) Super-tier counterpart of the VideoPhy-2 recipe above: same 1–5 physical-plausibility scoring on [videophysics/videophy2_train](https://huggingface.co/datasets/videophysics/videophy2_train), but a full fine-tune of the 32B backbone. `[job].task = "vlm"`. Bootstraps from `Cosmos3-Super`'s language-model weights merged onto the public Qwen3-VL-32B-Instruct visual tower; the merged HF directory is consumed via `[model.backbone].safetensors_path` (plumbed by `VLM_SAFETENSORS_PATH`). Full-shard FSDP across all ranks, so it runs on a 4-GPU (e.g. GB200x4) or 8-GPU allocation. Launch shell: `examples/launch_sft_videophy2_super.sh` ```shell # Step 1 (data): same as the nano recipe — materialize the public HF dataset into # the canonical local layout (videophy2_{train,val}/{meta.json, media/, text/}). python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ --out_root examples/data/videophysics --split both ```
Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Edge) Edge-tier counterpart of the VideoPhy-2 recipe above: same 1–5 physical-plausibility scoring on [videophysics/videophy2_train](https://huggingface.co/datasets/videophysics/videophy2_train), but on the compact **Nemotron-2B-Dense-VL** reasoner (SigLIP2 vision tower, `model_type = "cosmos3_edge"` — native HF metadata, no remote code; the model classes are registered in-framework) instead of Qwen3-VL. `[job].task = "vlm"`. The vision tower is frozen; the projector + LM train at a uniform LR. `[model.backbone].model_name` is the **public, ungated** omni release `nvidia/Cosmos3-Edge`, which supplies the arch/config/tokenizer **and** the reasoner weights: the training loader follows the snapshot's root safetensors index into its weight shards, so weights load directly from the download — no conversion step and no required weights env var. Like the nano/super recipes, `VLM_SAFETENSORS_PATH` is an optional override for loading from a local safetensors directory instead. The reasoner weights shipped inside `nvidia/Cosmos3-Edge` are the canonical Edge reasoner weights — there is no separate reasoner repo to download. Only 2B, so it runs on a 4-GPU (e.g. GB200x4) or 8-GPU allocation. Launch shell: `examples/launch_sft_videophy2_edge.sh` ```shell # Step 1 (data): same as the nano recipe — materialize the public HF dataset into # the canonical local layout (videophy2_{train,val}/{meta.json, media/, text/}). python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ --out_root examples/data/videophysics --split both ```
Action-Policy Post-Training (DROID / LIBERO) Robot action-policy recipes: DROID (`joint_pos` 8-D actions + proprioceptive state) and LIBERO (`frame_wise_relative` rot6d 10-D actions), both from Cosmos3-Nano. They follow the same Step 2/Step 3 flow below, with their own data staging (DROID keep-ranges window filter, LIBERO suite selection), and are documented separately: - [DROID action policy](action_policy_droid_posttrain.md) — `examples/launch_sft_action_policy_droid_nano.sh` - [LIBERO action policy](action_policy_libero_posttrain.md) — `examples/launch_sft_action_policy_libero_10_nano.sh` / `examples/launch_sft_action_policy_libero_all_nano.sh`
## Step 2 — Prepare checkpoint Convert the base checkpoint to [PyTorch Distributed Checkpoint (DCP)](https://pytorch.org/docs/stable/distributed.checkpoint.html) 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](setup.md)). Set `BASE_CHECKPOINT_NAME` to the value from the recipe block you picked in Step 1 (`Cosmos3-Nano`, `Cosmos3-Super`, or `Cosmos3-Edge`): ```shell BASE_CHECKPOINT_NAME=Cosmos3-Nano # or Cosmos3-Super / Cosmos3-Edge — 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`, `Cosmos3-Edge`) 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 required). To start from a merged Cosmos3 reasoner snapshot instead, build one with `convert_model_to_vlm_safetensors` (see the VideoPhy-2 note below) and pass it via `VLM_SAFETENSORS_PATH`. **Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Nano):** Use `cosmos_framework.scripts.convert_model_to_vlm_safetensors` instead. ```shell # 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 ``` **Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Super):** Same converter, but merge the `Cosmos3-Super` LM onto the **32B** Qwen3-VL visual tower via `--vlm-model-name Qwen/Qwen3-VL-32B-Instruct`. ```shell # Step 2 (Super VLM checkpoint): merge Cosmos3-Super LM onto the Qwen3-VL-32B visual tower. python -m cosmos_framework.scripts.convert_model_to_vlm_safetensors \ --checkpoint-path Cosmos3-Super \ --vlm-model-name Qwen/Qwen3-VL-32B-Instruct \ -o examples/checkpoints/Cosmos3-Super-VLM ``` **Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Edge):** Skip this step — the recipe loads the reasoner weights directly from the (ungated) `nvidia/Cosmos3-Edge` snapshot at startup: the training loader follows the repo's root safetensors index into its weight shards, so no conversion or extraction is required. To start from a local safetensors directory instead, point `VLM_SAFETENSORS_PATH` at it (optional, same mechanism as the nano recipe). ## Step 3 — Run training **Weights & Biases (optional):** training recipes default to `job.wandb_mode = "disabled"`. To log a run to W&B, set `wandb_mode = "online"` under `[job]` in the TOML (or pass `job.wandb_mode=online` as a command-line override) and export `WANDB_API_KEY` in your environment before launching. ### Option A (recommended): the paired launch shell Each recipe ships as a `examples/toml/sft_config/.toml` (validated against the pydantic schema at [`cosmos_framework/configs/toml_config/sft_config.py`](https://github.com/NVIDIA/cosmos-framework/blob/main/cosmos_framework/configs/toml_config/sft_config.py)) paired with `examples/launch_sft_.sh`; the full catalog is indexed in [`examples/README.md`](https://github.com/NVIDIA/cosmos-framework/blob/main/examples/README.md). Each `.sh` sources [`examples/_sft_launcher_common.sh`](https://github.com/NVIDIA/cosmos-framework/blob/main/examples/_sft_launcher_common.sh) and forwards into `cosmos_framework.scripts.train --sft-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](#overriding-the-defaults) to override): ```shell # 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): | Launch shell | Post-Training Task | Default $DATASET_PATH (under examples/data/) | Default $BASE_CHECKPOINT_PATH (under examples/checkpoints/) | | ------------------------------ | ------------------ | ---------------------------------------------------------- | ------------------------------------------------------------------ | | `launch_sft_vision_nano.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Nano` | | `launch_sft_vision_super.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Super` | | `launch_sft_vision_edge.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Edge` | | `launch_sft_llava_ov.sh` | Reasoner SFT | (none; dataset streams from HF Hub) | (none; backbone fetched at startup, or set `VLM_SAFETENSORS_PATH`) | | `launch_sft_videophy2_nano.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env) | | `launch_sft_videophy2_super.sh`| Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env — Cosmos3-Super-VLM merge) | | `launch_sft_videophy2_edge.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; weights load directly from `nvidia/Cosmos3-Edge`) | `WAN_VAE_PATH` defaults to `examples/checkpoints/wan22_vae/Wan2.2_VAE.pth` for every non-reasoner recipe. Cosmos3-Edge is only 2B and also fits a 4-GPU node — e.g. `NPROC_PER_NODE=4 bash examples/launch_sft_vision_edge.sh` (tested on 4×GB200). **Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Nano):** ```shell # 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 ``` **Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Super):** ```shell # Same as the nano recipe, but point VLM_SAFETENSORS_PATH at the Cosmos3-Super # merge from Step 2. On a 4-GPU node (e.g. GB200x4) also set NPROC_PER_NODE=4. export VIDEOPHYSICS_ROOT=$PWD/examples/data/videophysics export VLM_SAFETENSORS_PATH=$PWD/examples/checkpoints/Cosmos3-Super-VLM bash examples/launch_sft_videophy2_super.sh ``` **Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Edge):** ```shell # Same as the nano recipe, but no VLM_SAFETENSORS_PATH is needed: the reasoner # weights load directly from the (ungated) nvidia/Cosmos3-Edge snapshot. # Only 2B — on a 4-GPU node set NPROC_PER_NODE=4. export VIDEOPHYSICS_ROOT=$PWD/examples/data/videophysics bash examples/launch_sft_videophy2_edge.sh ``` #### Overriding the defaults 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: ```shell # 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. ### Option B: raw `torchrun` 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/.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/.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. ```shell # 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_/` directory under `$IMAGINAIRE_OUTPUT_ROOT////` (see [Outputs](#outputs) below for the full layout). ## Outputs Outputs land under `$IMAGINAIRE_OUTPUT_ROOT////`: 1. `config.yaml`, `config.pkl`: Finalized resolved config (YAML for inspection, pickle for re-instantiation). 1. `launch_info.yaml`, `job_env.yaml`: Job metadata and captured launch environment. 1. `checkpoints/`: 1. `latest_checkpoint.txt`: Pointer file containing the latest checkpoint directory name (e.g. `iter_000000200`). 1. `iter_/`: DCP checkpoint saved every `[train.ckpt].save_freq` iterations (zero-padded 9-digit, e.g. `iter_000000200/`): 1. `model/`: model weights (sharded `.distcp`). 1. `optim/`: optimizer state. 1. `scheduler/`: LR scheduler state. 1. `trainer/`: training state — includes the `iteration` counter and per-rank `rng_state_` (numpy + random + torch + torch_cuda). 1. `dataloader/`: optional per-rank pickle shards (`rank_.pkl`) — only present for dataloaders that implement `has_state()`. 1. `/`: Callback outputs, one directory per registered callback (e.g. `DeviceMonitor/`, `EveryNDrawSample/`, `norm_monitor/`). 1. `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///`. For example, with `IMAGINAIRE_OUTPUT_ROOT=outputs/train` and the `vision_sft_nano` recipe, `$RUN_DIR` is `outputs/train/cosmos3/sft/vision_sft_nano`. ## Hybrid video reasoning SFT The `video_reasoning_sft.toml` recipe accepts existing DAFT `tao-vl-reason-v1.0` annotations directly through Framework's native dataset loader. It also accepts the equivalent `cosmos-video-reasoning-v1.0` records. No dataset conversion or `nvidia-tao-daft` runtime installation is needed. Each JSON object contains an `items` array with `question`, `answer`, optional `reasoning`, and `video_id` fields: ```json { "items": [ { "video_id": "clip.mp4", "question": "Why does the car stop?", "answer": "The light is red.", "reasoning": "The car approaches a red traffic light." } ] } ``` Set annotation paths independently for training and validation. A single path or a JSON array of paths is accepted; multiple files are concatenated in order. For example, with a staged Cosmos3-Nano VLM safetensors checkpoint: ```shell export VLM_SAFETENSORS_PATH=/checkpoints/cosmos3-nano-vlm export COSMOS_VIDEO_TRAIN_ANNOTATIONS='["/data/train-bcq.json", "/data/train-mcq.json"]' export COSMOS_VIDEO_VAL_ANNOTATIONS=/data/val.json export COSMOS_VIDEO_TRAIN_MEDIA_ROOTS=/data/videos export COSMOS_VIDEO_VAL_MEDIA_ROOTS=/data/videos export COSMOS_VIDEO_TRAIN_RESPONSE_MODE=hybrid export COSMOS_VIDEO_VAL_RESPONSE_MODE=answer torchrun --nproc-per-node=4 -m cosmos_framework.scripts.train \ --sft-toml=examples/toml/sft_config/video_reasoning_sft.toml \ -- trainer.seed=42 ``` Media roots may be a shared path or a JSON array matching the annotation paths. When omitted, Framework resolves an annotation's `media_root` relative to its JSON directory, or uses that directory if the field is absent. The older Cosmos-RL DAFT hook uses the annotation directory when no explicit root is given, ignoring an embedded `media_root`; supply explicit roots to preserve those paths. These video recipes do not support `image_id` records. Each split accepts these settings (`TRAIN` or `VAL` replaces `{SPLIT}`): | Environment variable | Meaning | Default | | ------------------------------------ | ------------------------------------------------------------------ | ---------------------------------------------- | | `COSMOS_VIDEO_{SPLIT}_RESPONSE_MODE` | `answer`, `think`, or `hybrid` | `hybrid` for training; `answer` for validation | | `COSMOS_VIDEO_{SPLIT}_SAMPLE_STRIDE` | Select every Nth raw item; positive integer | `1` | | `COSMOS_VIDEO_{SPLIT}_SAMPLE_OFFSET` | Start at this raw index; non-negative integer | `0` | | `COSMOS_VIDEO_{SPLIT}_LIMIT` | Maximum exposed samples after response expansion; positive integer | Unset (unlimited) | | `COSMOS_VIDEO_SYSTEM_PROMPT` | Optional system message for both splits | Empty | `answer` targets contain the answer alone. `think` targets contain `\n{reasoning}\n\n\n{answer}`, falling back to the answer if reasoning is absent or empty. `hybrid` exposes two targets per selected raw item, interleaved as answer, think, answer, think, matching the Cosmos-RL DAFT hook. Stride and offset apply before expansion; the limit applies afterward and may end on an unpaired answer when odd. An offset past the input produces no samples. For example, stride `2`, offset `1`, and limit `3` expose the answer and think targets for raw item `1`, then the answer target for raw item `3`. The reasoning processor applies the Qwen3-VL Instruct chat template to preserve explicit think tags and avoid an injected empty think block. Loss masking trains assistant tokens, including supplied reasoning, while ignoring prompt tokens. Framework uses its native TorchCodec video processor and lifecycle/checkpoint callbacks; this migration preserves the response/sampling contract, not the Cosmos-RL decoder implementation. Adjust frame/pixel budgets and the smoke recipe's training schedule for your workload; its fixed optimizer-step budget does not automatically double in hybrid mode. The Edge experiment is `cosmos_task_aware_video_reasoning_edge` and uses the same dataset settings; use an Edge checkpoint and its corresponding model and token budgets. ## Export checkpoint to Hugging Face safetensors Export the DCP checkpoint produced in Step 3 to a Hugging Face safetensors checkpoint: ```shell RUN_DIR=$IMAGINAIRE_OUTPUT_ROOT/// 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](https://github.com/NVIDIA/cosmos-framework/blob/main/README.md#inference) commands by passing `--checkpoint-path $RUN_DIR/model`. ### ViT / vision tower (Cosmos3-Edge) The plain export above is all any recipe needs — Edge included. For Cosmos3-Edge, `export_model` bundles the SigLIP2 vision tower into `$RUN_DIR/model/vision_encoder/` and the processor/tokenizer files into the export root, so the exported directory is **self-contained**: inference runs from it offline, with no Hub download (see [Inference → Reasoner](inference.md#reasoner)). `export_manifest.json` records where the tower and processor came from. Nano/Super exports bundle their processor files the same way. Optional flags: ```shell # --no-vit: generation-only export (T2V / I2V / V2V / T2I), ~1 GB smaller. # Writes include_visual=false so the checkpoint loads cleanly everywhere. python -m cosmos_framework.scripts.export_model \ --checkpoint-path $CHECKPOINT_PATH \ --config-file $RUN_DIR/config.yaml \ --no-vit \ -o $RUN_DIR/model # --vit-checkpoint-path: use a custom tower from a local directory # (a snapshot root or a vision_encoder/ directory). python -m cosmos_framework.scripts.export_model \ --checkpoint-path $CHECKPOINT_PATH \ --config-file $RUN_DIR/config.yaml \ --vit-checkpoint-path /path/to/snapshot \ -o $RUN_DIR/model # --verify: smoke-test the export on this node's GPU (tiny reasoner query + # tiny generation). A failure exits non-zero but keeps the export. python -m cosmos_framework.scripts.export_model \ --checkpoint-path $CHECKPOINT_PATH \ --config-file $RUN_DIR/config.yaml \ --verify \ -o $RUN_DIR/model ``` ## Convert to Diffusers Convert the Hugging Face export from the previous step into a Diffusers pipeline (`transformer/`, `vae/`, `scheduler/`, `text_tokenizer/`, `model_index.json`, …): ```shell python -m cosmos_framework.scripts.convert_model_to_diffusers \ --checkpoint-path $RUN_DIR/model \ -o $RUN_DIR/diffusers ``` The input is the Hugging Face safetensors directory produced by `export_model` above — not a raw DCP checkpoint, so run the export first. The result at `$RUN_DIR/diffusers` can be loaded with `diffusers` or passed to [Inference](https://github.com/NVIDIA/cosmos-framework/blob/main/README.md#inference) via `--checkpoint-path $RUN_DIR/diffusers`. A generation-only checkpoint with no reasoner vision tower (`include_visual` unset) is converted without a `vision_encoder/` sidecar automatically; pass `--skip-vision-encoder` to force-skip it even when a tower is present. ## Config The recipe TOML is parsed against the pydantic schema [`SFTExperimentConfig`](https://github.com/NVIDIA/cosmos-framework/blob/main/cosmos_framework/configs/toml_config/sft_config.py) 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/.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](sft_config.md). 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`. 1. `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. 1. `project`, `group`, `name` — Components of the run output dir `$IMAGINAIRE_OUTPUT_ROOT////`. Also flow to W&B as the project / group / run name. 1. `wandb_mode` — `"online"` (logs to W&B; `WANDB_API_KEY` must be set), `"offline"` (logs locally, sync later with `wandb sync`), or `"disabled"`. 1. `[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). 1. `joint_attn_implementation` — VFM attention layout: `"two_way"` / `"three_way"` (NATTEN) / `"flex"`. 1. `attn_implementation` — VLM attention impl: `"cosmos"` / `"flash_attention_2"` / `"sdpa"` / `"eager"`. VLM only. 1. `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. 1. `[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. 1. `[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. 1. `data_parallel_replicate_degree` — HSDP replicate degree (outer replicate loop over the shard topology). 1. `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). 1. `cfg_parallel_shard_degree` — Classifier-free-guidance shard degree. Almost always `1` for SFT. 1. `fsdp_master_dtype` — Master parameter / FSDP reduce dtype: typically `"float32"`. 1. `[model.compile]` 1. `enabled` — Enable `torch.compile`. Improves speed at the cost of memory. 1. `compile_dynamic` — Whether to compile with symbolic-shape (dynamic) kernels. `True` (default) is appropriate for training; AR inference may prefer `False` for stable shapes. 1. `[model]` 1. `precision` — Compute dtype for forward/backward: `"bfloat16"` / `"float16"` / `"float32"`. Master weights stay fp32 separately. 1. `[model.activation_checkpointing]` 1. `mode` — `"none"` / `"selective"` (per-op SAC, MoT-only) / `"full"` (per-block checkpointing). 1. `save_ops_regex` — Regex patterns for ops to keep saved under `mode="selective"`. 1. `preserve_rng_state`, `determinism_check` — Recompute determinism plumbing. 1. `[model.tokenizer]` 1. `vae_path` — Wan2.2 VAE `.pth` path. Recipe TOMLs use `"${oc.env:WAN_VAE_PATH}"`. VFM only. 1. `[optimizer]` 1. `lr` — Base learning rate. 1. `betas`, `eps`, `fused`, `weight_decay` — Standard AdamW knobs. `eps` is VFM-only. 1. `keys_to_select` — Substring allowlist for trainable params. Empty list = train everything; `["lora_"]` = adapter-only fine-tune. 1. `[optimizer.lr_multipliers]` 1. Inline table of ` = ` pairs that scale the LR of params whose name contains the substring. The shipped vision recipes leave this empty (Hydra default `{}` stands). 1. `[scheduler]` 1. `cycle_lengths`, `warm_up_steps` — Cycle length and warmup duration (lists, one entry per cycle), in optimizer steps. 1. `f_max`, `f_min`, `f_start` — LR multipliers at peak / trough / step-0 (ratios of `optimizer.lr`). 1. `verbosity_interval` — Scheduler-side LR log frequency. VFM only. 1. `[trainer]` 1. `max_iter` — Total optimizer steps. 1. `grad_accum_iter` — Micro-batches per optimizer step. Effective global batch = `grad_accum_iter × per-rank batch × world_size`. 1. `logging_iter` — Console / W&B scalar log frequency. 1. `distributed_parallelism` — `"fsdp"` is the only supported value. 1. `[trainer.callbacks.compile_tokenizer]` 1. `enabled`, `compile_after_iterations`, `warmup_resolutions` — Lazy `torch.compile` of the VAE tokenizer. VFM only. 1. `[trainer.callbacks.grad_clip]` 1. `clip_norm` — Max global L2 norm of the gradient (steps with larger norm are rescaled). 1. `force_finite` — Replace NaN/Inf grads with zero (default `true` on VFM, `false` on VLM). 1. `[checkpoint]` 1. `load_path` — Base DCP checkpoint directory to resume from (Step 2 output, or a prior run's `checkpoints/iter_/`). Recipe TOMLs use `"${oc.env:BASE_CHECKPOINT_PATH}"`. 1. `save_iter` — Save a new DCP checkpoint every N optimizer steps. 1. `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. 1. `[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. 1. `max_sequence_length` — Per-packed-sequence token cap (remapped to `max_tokens` on the VLM packer). 1. `seed` — Dataloader RNG seed (VFM only). ### Common Hydra tail overrides 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.