Add Cosmos3-Super VideoPhy-2 Reasoner SFT recipe

cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py

@@ -160,11 +160,33 @@ def _dl(dataset_key, split, num_workers, persistent_workers=False, pin_memory=Fa
             grad_accum_iter=8,
         ),
         optimizer=dict(
+            # Uniform 1e-6 across both tiers. Nano is stable at 1e-5 too, but the 32B
+            # super run spikes transiently (loss->6.6 @ iter 23) at 1e-5 before
+            # recovering to the same endpoint; 1e-6 reaches the identical final loss
+            # with no instability, so both tiers ship at 1e-6. (Super deepcopies this
+            # block, so setting it here covers both.)
             lr=1e-6,
             fused=True,
             weight_decay=0.05,
             betas=[0.9, 0.999],
-            lr_multipliers={"mm_projector": 20.0, "merger": 20.0},
+            # NOTE: the ViT is FROZEN (freeze_vision_encoder=True below), so the only
+            # trainable "model.visual.*" params are the projector (model.visual.merger
+            # + deepstack_merger_list). This key therefore sets the PROJECTOR LR only,
+            # to 1.0x optimizer.lr -- giving a uniform LR across projector + LLM +
+            # lm_head. It is NOT touching the frozen ViT backbone.
+            #
+            # Why the key is "model.visual" (broad) and not "model.visual.merger":
+            # the reasoner default optimizer (configs/base/reasoner/defaults/optimizer.py)
+            # ships lr_multipliers={"model.visual": 0.1}. Hydra merges dicts base-first
+            # (so that key stays PINNED FIRST) and _build_params_with_metadata
+            # (utils/generator/optimizer.py) is first-substring-match-wins -- so a
+            # narrower "model.visual.merger" key is shadowed by "model.visual" and never
+            # matches (verified: the merger still resolves to 0.1x). Overriding the SAME
+            # "model.visual" key to 1.0 is the only way to cancel the 0.1x default; with
+            # the ViT frozen that reaches only the merger. (To give an UNFROZEN ViT its
+            # own LR you'd likewise override this "model.visual" value -- a narrower key
+            # can't out-prioritize it.)
+            lr_multipliers={"model.visual": 1.0},
         ),
         scheduler=dict(
             warm_up_steps=[5],

cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""VideoPhy-2 SFT recipe (Cosmos3-Super tier): Qwen3-VL-32B full fine-tune.
+
+Super-tier counterpart of ``videophy2_sft_nano``. Reuses that recipe's VideoPhy-2
+``LocalSFTDataset`` + ``CosmosDataLoader`` dataflow verbatim (imported + deepcopied,
+the same variant idiom ``llava_ov_vlm.py`` uses) and changes only what the larger
+backbone needs:
+
+  * ``vlm_policy`` ``qwen3_vl_8b_instruct`` -> ``qwen3_vl_32b_instruct``
+    (``Qwen/Qwen3-VL-32B-Instruct``) — the visual tower + config the Cosmos3-Super
+    Reasoner is built on, mirroring how the nano recipe rides the 8B tower.
+  * FSDP full-shard across every rank (``data_parallel_shard_degree=-1``) so the 32B
+    weights + optimizer state fit, instead of the nano recipe's fixed dp=8. This is
+    the same super-tier sharding switch ``vision_sft_super`` makes, and lets the recipe
+    run unchanged on a 4-GPU (e.g. GB200x4) or 8-GPU allocation.
+
+Still a full fine-tune (no LoRA): the freeze config is inherited from the nano recipe
+(vision encoder frozen, LM + mm_projector trained).
+
+Launch via ``examples/launch_sft_videophy2_super.sh`` after
+``prepare_videophy2_from_hf`` populates ``$VIDEOPHYSICS_ROOT``, supplying the merged
+Cosmos3-Super Reasoner checkpoint through ``VLM_SAFETENSORS_PATH`` (see the launch shell).
+"""
+
+from __future__ import annotations
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+# Importing the nano module registers `videophy2_sft_nano` and pulls in the shared
+# VideoPhy-2 dataflow helpers; we clone its LazyDict rather than re-declaring them.
+from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_nano import videophy2_sft_nano
+
+cs = ConfigStore.instance()
+
+
+videophy2_sft_super = copy.deepcopy(videophy2_sft_nano)
+
+# Backbone: nano 8B -> super 32B. The vlm_policy override lives in the Hydra
+# `defaults` list; rewrite that one entry's value in place, leaving the rest of the
+# recipe (checkpoint backend, callbacks, dataflow) untouched.
+for _default in videophy2_sft_super["defaults"]:
+    if not isinstance(_default, str) and "override /vlm_policy" in _default:
+        _default["override /vlm_policy"] = "qwen3_vl_32b_instruct"
+
+# 32B full fine-tune: shard model + optimizer state across every rank (FSDP full
+# shard, auto-sized from WORLD_SIZE) instead of the nano recipe's fixed dp_shard=8,
+# so the recipe fits on 4- or 8-GPU nodes.
+# examples/toml/sft_config/videophy2_sft_super.toml is authoritative at launch and
+# repeats these; keeping them here lets `experiment=videophy2_sft_super` run standalone.
+videophy2_sft_super.model.config.parallelism.data_parallel_shard_degree = -1
+videophy2_sft_super.model.config.parallelism.data_parallel_replicate_degree = 1
+
+
+for _item in [videophy2_sft_super]:
+    experiment_name = [name.lower() for name, value in globals().items() if value is _item][0]
+    if "job" not in _item:
+        _item["job"] = dict(name=experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}")
+    else:
+        _item["job"]["name"] = experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}"
+
+    cs.store(group="experiment", package="_global_", name=experiment_name, node=_item)

docs/training.md

@@ -115,6 +115,21 @@ python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \
 
 </details>
 
+<details><summary><b>Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Super)</b></summary>
+
+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
+```
+
+</details>
+
 ## 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)).
@@ -144,6 +159,16 @@ python -m cosmos_framework.scripts.convert_model_to_vlm_safetensors \
     -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
+```
+
 ## Step 3 — Run training
 
 **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.
@@ -165,6 +190,7 @@ Each launcher's default paths come from the `DATASET_PATH` + `BASE_CHECKPOINT_PA
 | `launch_sft_vision_super.sh`   | Generator SFT      | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Super`                                                    |
 | `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)   |
 
 `WAN_VAE_PATH` defaults to `examples/checkpoints/wan22_vae/Wan2.2_VAE.pth` for every non-reasoner recipe.
 
@@ -177,6 +203,16 @@ 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
+```
+
 #### 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:

examples/README.md

@@ -16,3 +16,4 @@ This directory contains:
 | Vision SFT LoRA (Cosmos3-Super)              | `launch_sft_vision_super.sh`          |
 | Reasoner Alignment SFT                       | `launch_sft_llava_ov.sh`              |
 | Reasoner Alignment SFT (Cosmos3-Nano)        | `launch_sft_videophy2_nano.sh`        |
+| Reasoner Alignment SFT (Cosmos3-Super)       | `launch_sft_videophy2_super.sh`       |

examples/launch_sft_videophy2_super.sh

@@ -0,0 +1,55 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+# Structured-TOML launch for videophy2_sft_super (VLM dialog SFT on VideoPhy-2
+# via CosmosDataLoader, Cosmos3-Super tier / Qwen3-VL-32B full fine-tune). Drives
+# cosmos_framework.scripts.train against
+# examples/toml/sft_config/videophy2_sft_super.toml.
+#
+# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config.
+#
+# Required env:
+#   VIDEOPHYSICS_ROOT  dir containing videophy2_train/ and videophy2_val/
+#                      (each with meta.json + media/ + text/). Populate via
+#                      `python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf`.
+#
+# Optional env:
+#   HF_TOKEN               for gated Qwen3-VL-32B-Instruct downloads.
+#   VLM_SAFETENSORS_PATH   local directory of pre-converted Qwen3-VL safetensors
+#                          (e.g. Cosmos3-Super LM merged with the Qwen3-VL-32B visual
+#                          tower via `cosmos_framework.scripts.convert_model_to_vlm_safetensors
+#                          --checkpoint-path Cosmos3-Super --vlm-model-name Qwen/Qwen3-VL-32B-Instruct`).
+#                          When set, plumbed to backbone.safetensors_path via a
+#                          tail override. When unset, the framework falls back
+#                          to the public Qwen/Qwen3-VL-32B-Instruct HF snapshot.
+#   NPROC_PER_NODE         torchrun GPUs per node; default 8. Set 4 on a GB200x4 node.
+#
+# Usage (8-GPU allocation, inside the training container, from the repo root):
+#   VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_super.sh
+#   # on a 4-GPU node (e.g. GB200x4):
+#   NPROC_PER_NODE=4 VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_super.sh
+
+TOML_FILE="examples/toml/sft_config/videophy2_sft_super.toml"
+
+# Super-variant allocator tweak: expandable_segments so the 32B backbone fits
+# without OOM during compile/decode. (Unlike launch_sft_vision_super.sh we do NOT
+# clear LD_LIBRARY_PATH — this reasoner recipe decodes VideoPhy-2 clips with
+# torchcodec, which dlopen()s the CUDA NPP + FFmpeg libs off LD_LIBRARY_PATH; the
+# nano videophy2 launcher leaves it untouched for the same reason.)
+export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}"
+
+TAIL_OVERRIDES=(
+    ${EXTRA_TAIL_OVERRIDES:-}
+)
+
+# When VLM_SAFETENSORS_PATH is set, plumb it to backbone.safetensors_path so the
+# framework loads weights from the local snapshot (e.g. a Cosmos3-Super LM merged
+# with the Qwen3-VL-32B visual tower via
+# `cosmos_framework.scripts.convert_model_to_vlm_safetensors`) while keeping the
+# public HF model_name for tokenizer/architecture discovery.
+if [[ -n "${VLM_SAFETENSORS_PATH:-}" ]]; then
+    TAIL_OVERRIDES+=("model.config.policy.backbone.safetensors_path=$VLM_SAFETENSORS_PATH")
+fi
+
+source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh"

examples/toml/sft_config/videophy2_sft_nano.toml

@@ -55,6 +55,9 @@ eps           = 1.0e-8
 fused         = true
 lr            = 1.0e-6
 weight_decay  = 0.1
+# NOTE: per-part LR multipliers live in the experiment (videophy2_sft_nano.py),
+# NOT here -- the toml->hydra override path parses a dict key's "." as nesting,
+# so a dotted key like "model.visual" can't be set from toml. Edit the experiment.
 
 [scheduler]
 cycle_lengths      = [50]

examples/toml/sft_config/videophy2_sft_super.toml

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+# videophy2_sft_super — VLM dialog SFT on VideoPhy-2 via CosmosDataLoader,
+# Cosmos3-Super tier (Qwen3-VL-32B full fine-tune).
+# Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm").
+#
+# Super-tier counterpart of videophy2_sft_nano: same dataset + dataflow, but the
+# 32B backbone (Qwen/Qwen3-VL-32B-Instruct) and FSDP full-shard across every rank
+# (data_parallel_shard_degree=-1, auto from WORLD_SIZE) so it fits on a 4-GPU
+# (e.g. GB200x4) or 8-GPU allocation. Still a full fine-tune (no LoRA).
+#
+# Dataset prep:
+#   python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \
+#       --out_root $VIDEOPHYSICS_ROOT --split train  # and again with --split val
+#
+# Required env at launch: VIDEOPHYSICS_ROOT (read by the experiment Python).
+# Supply the merged Cosmos3-Super Reasoner checkpoint via VLM_SAFETENSORS_PATH.
+#
+# Example launch:
+#   bash examples/launch_sft_videophy2_super.sh
+
+[job]
+task         = "vlm"
+experiment   = "videophy2_sft_super"
+project      = "cosmos3"
+group        = "vlm_videophy2_sft"
+name         = "videophy2_sft_super"
+wandb_mode   = "disabled"
+
+[model]
+attn_implementation = "cosmos"
+precision           = "bfloat16"                         # was [model.parallelism].precision
+
+[model.backbone]
+model_name = "Qwen/Qwen3-VL-32B-Instruct"
+
+[model.ema]
+enabled         = false
+rate            = 0.1
+iteration_shift = 0
+
+[model.parallelism]
+data_parallel_shard_degree      = -1                     # super: FSDP full shard, auto from WORLD_SIZE (4- or 8-GPU)
+data_parallel_replicate_degree  = 1
+context_parallel_shard_degree   = 1                      # raise to 2 (needs even GPU count) if the 32B run OOMs
+cfg_parallel_shard_degree       = 1
+
+[model.compile]
+enabled                         = false                  # was [model.parallelism].use_torch_compile
+compile_dynamic                 = true
+
+[model.activation_checkpointing]
+mode                = "full"
+save_ops_regex      = ["fmha"]
+preserve_rng_state  = true
+determinism_check   = "default"
+
+[optimizer]
+betas         = [0.9, 0.95]
+eps           = 1.0e-8
+fused         = true
+lr            = 1.0e-6
+weight_decay  = 0.1
+# NOTE: per-part LR multipliers live in the experiment (videophy2_sft_nano.py),
+# NOT here -- the toml->hydra override path parses a dict key's "." as nesting,
+# so a dotted key like "model.visual" can't be set from toml. Edit the experiment.
+
+[scheduler]
+cycle_lengths      = [50]
+f_max              = [1.0]
+f_min              = [0.1]
+f_start            = [0.05]
+verbosity_interval = 0
+warm_up_steps      = [5]
+
+[trainer]
+distributed_parallelism = "fsdp"
+grad_accum_iter         = 8
+logging_iter            = 1
+max_iter                = 50
+
+[trainer.callbacks.compile_tokenizer]
+compile_after_iterations = 3
+enabled                  = false
+
+[trainer.callbacks.grad_clip]
+clip_norm    = 1.0
+force_finite = false
+
+[checkpoint]
+keys_to_skip_loading = []
+load_path            = "???"
+save_iter            = 100
+
+[dataloader_train]
+# Routed by PATH_REMAPS["vlm"] onto the CosmosDataLoader's nested PoolPackingBatcher:
+#   max_samples_per_batch -> dataloader_train.batcher.max_batch_size
+#   max_sequence_length   -> dataloader_train.batcher.max_tokens
+max_samples_per_batch = 1
+max_sequence_length   = 16000