From b8ee60be2531b7760cb52445c3f7ab84ba9819fa Mon Sep 17 00:00:00 2001 From: Simon Zhang Date: Thu, 9 Jul 2026 20:58:14 -0700 Subject: [PATCH] Add Cosmos3-Super VideoPhy-2 Reasoner SFT recipe Add the Super-tier counterpart to the Nano VideoPhy-2 recipe. It fully fine-tunes Qwen3-VL-32B using FSDP full sharding and initializes from the Cosmos3-Super language model merged with the 32B visual tower. The new experiment mirrors the Nano configuration but uses the 32B backbone and sets data_parallel_shard_degree=-1. It also adds a launcher and TOML recipe. Use a learning rate of 1e-6 for both tiers. In a 150-iteration Super run, 1e-5 caused a temporary loss spike to about 6.6 at iteration 23. Both learning rates reached the same last-20 average loss of 0.437, so 1e-6 provides the same result without the instability. Set lr_multipliers={"model.visual": 1.0} to override the reasoner default of 0.1 and train the visual projector at the base learning rate. The previous mm_projector and merger keys matched no Qwen3- VL parameters because the projector is named model.visual.merger, so the intended 20x multiplier was never applied. Document the Super recipe in docs/training.md and examples/README.md. --- .../reasoner/experiment/videophy2_sft_nano.py | 24 ++++- .../experiment/videophy2_sft_super.py | 65 +++++++++++ docs/training.md | 36 +++++++ examples/README.md | 1 + examples/launch_sft_videophy2_super.sh | 55 ++++++++++ .../toml/sft_config/videophy2_sft_nano.toml | 3 + .../toml/sft_config/videophy2_sft_super.toml | 101 ++++++++++++++++++ 7 files changed, 284 insertions(+), 1 deletion(-) create mode 100644 cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py create mode 100755 examples/launch_sft_videophy2_super.sh create mode 100644 examples/toml/sft_config/videophy2_sft_super.toml diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py index 1405ac7f..1123e5fb 100644 --- a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py +++ b/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], diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py new file mode 100644 index 00000000..83d5d7e2 --- /dev/null +++ b/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) diff --git a/docs/training.md b/docs/training.md index 5b11e5ff..056ceb0a 100644 --- a/docs/training.md +++ b/docs/training.md @@ -115,6 +115,21 @@ python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +
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 +``` + +
+ ## 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: diff --git a/examples/README.md b/examples/README.md index 0e51140f..08c219a8 100644 --- a/examples/README.md +++ b/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` | diff --git a/examples/launch_sft_videophy2_super.sh b/examples/launch_sft_videophy2_super.sh new file mode 100755 index 00000000..785b896b --- /dev/null +++ b/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" diff --git a/examples/toml/sft_config/videophy2_sft_nano.toml b/examples/toml/sft_config/videophy2_sft_nano.toml index 35b8e785..81f43964 100644 --- a/examples/toml/sft_config/videophy2_sft_nano.toml +++ b/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] diff --git a/examples/toml/sft_config/videophy2_sft_super.toml b/examples/toml/sft_config/videophy2_sft_super.toml new file mode 100644 index 00000000..9029e246 --- /dev/null +++ b/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