From c9a694c5de470a423ee8d86e7e1179db5b2fe74b Mon Sep 17 00:00:00 2001 From: Wei-Cheng Tseng Date: Fri, 17 Jul 2026 12:29:42 -0700 Subject: [PATCH 1/3] update action fd scripts --- .../launch_sft_action_fd_droid_posttrain.sh | 39 +++++++++++++++ .../sft_config/action_fd_droid_posttrain.toml | 47 +++++++++++++++++++ 2 files changed, 86 insertions(+) create mode 100644 cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh create mode 100644 cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_fd_droid_posttrain.toml diff --git a/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh new file mode 100644 index 00000000..8dc7dd50 --- /dev/null +++ b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh @@ -0,0 +1,39 @@ +#!/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 action_fd_droid_posttrain. +# +# This trains forward dynamics on the Cosmos3-DROID success + failure splits. See +# docs/action_fd_droid_posttrain.md. +# +# Env vars (override for your filesystem): +# DATASET_PATH Cosmos3-DROID parent dir (success/ + failure/) +# BASE_CHECKPOINT_PATH Base DCP checkpoint +# WAN_VAE_PATH Wan2.2 VAE .pth +# WANDB_API_KEY for online logging (TOML wandb_mode="online") +# NPROC_PER_NODE torchrun --nproc_per_node (default 8) +# EXTRA_TAIL_OVERRIDES space-separated Hydra overrides +# +# Single-node smoke: +# export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10" +# bash examples/launch_sft_action_fd_droid_posttrain.sh +# +# Multi-node: launch on every worker. For HSDP set +# model.parallelism.data_parallel_replicate_degree = (shard stays 8). +# ============================================================================ + +TOML_FILE="examples/toml/sft_config/action_fd_droid_posttrain.toml" +: "${DATASET_PATH:=examples/data/Cosmos3-DROID}" +: "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" + +WAN_VAE_PATH="${WAN_VAE_PATH:-examples/checkpoints/wan22_vae/Wan2.2_VAE.pth}" + +EXTRA_DATASET_CHECK='[[ "$WAN_VAE_PATH" = /* ]] || WAN_VAE_PATH="$WORKDIR/$WAN_VAE_PATH"; export WAN_VAE_PATH; [[ -f "$DATASET_PATH/success/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/success/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID success split under $DATASET_PATH (expected success/meta/info.json or success/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$DATASET_PATH/failure/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/failure/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID failure split under $DATASET_PATH (expected failure/meta/info.json or failure/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$WAN_VAE_PATH" ]] || { echo "ERROR: WAN_VAE_PATH not found: $WAN_VAE_PATH" >&2; exit 1; }' + +TAIL_OVERRIDES=( + ${EXTRA_TAIL_OVERRIDES:-} +) + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_fd_droid_posttrain.toml b/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_fd_droid_posttrain.toml new file mode 100644 index 00000000..65fbdfa1 --- /dev/null +++ b/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_fd_droid_posttrain.toml @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# ============================================================================ +# DROID-Merged LeRobot forward-dynamics post-training. +# +# Env required: +# DATASET_PATH=/path/to/Cosmos3-DROID +# BASE_CHECKPOINT_PATH= +# WAN_VAE_PATH= +# IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root +# +# The registered experiment owns dataset mode/action-space wiring: +# mode="forward_dynamics", action_space="ee_pose", chunk_length=16, res480. +# ============================================================================ + +[job] +task = "vfm" +experiment = "action_fd_droid_posttrain" +project = "cosmos3_action_fd" +group = "action_sft" +name = "action_fd_droid_posttrain" +wandb_mode = "online" + +[model] +precision = "bfloat16" + +[model.parallelism] +data_parallel_shard_degree = 8 +data_parallel_replicate_degree = 1 + +[model.activation_checkpointing] +mode = "selective" + +[model.tokenizer] +vae_path = "${oc.env:WAN_VAE_PATH}" + +[scheduler] +cycle_lengths = [20000] + +[trainer] +max_iter = 20000 +logging_iter = 50 + +[checkpoint] +load_path = "${oc.env:BASE_CHECKPOINT_PATH}" +save_iter = 250 From 270df48c189fb4465088a14324a0470229eb65a5 Mon Sep 17 00:00:00 2001 From: Wei-Cheng Tseng Date: Sun, 19 Jul 2026 20:00:14 -0700 Subject: [PATCH 2/3] update docs --- .../finetune/action_fd_droid_posttrain.md | 173 ++++++++++++++++++ .../launch_sft_action_fd_droid_posttrain.sh | 130 ++++++++++--- 2 files changed, 281 insertions(+), 22 deletions(-) create mode 100644 cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md diff --git a/cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md b/cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md new file mode 100644 index 00000000..99a2c04f --- /dev/null +++ b/cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md @@ -0,0 +1,173 @@ +# Action FD DROID Post-Training + +This document describes how to run the `action_fd_droid_posttrain` experiment. +It trains Cosmos3-DROID forward dynamics in `cosmos_framework`. + +## Overview + +| Piece | Value | +| ------------------------- | ---------------------------------------------------------------------------------------------------------------------- | +| Experiment | `action_fd_droid_posttrain` | +| TOML | `toml/sft_config/action_fd_droid_posttrain.toml` | +| Launch shell | `launch_sft_action_fd_droid_posttrain.sh` | +| Config module | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py` | +| Dataset wrapper | `cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py` | +| Dataset root | [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) parent root containing `success/` and `failure/` | +| Task mode | `forward_dynamics` | +| Action space | `ee_pose`: 10-D `[pos_delta, rot6d_delta, gripper]` | +| Chunk length / resolution | `16` frames at `480` | + +## Prerequisites + +- Install cosmos-framework as described in the shared + [setup instructions](cookbooks/cosmos3/README.md), then activate its + virtual environment. +- Authenticate with Hugging Face (`uvx hf@latest auth login` or `HF_TOKEN`) and + accept the terms for the required model and dataset repositories. +- The TOML defaults to online Weights & Biases logging. Export `WANDB_API_KEY`, + or disable it with `EXTRA_TAIL_OVERRIDES="job.wandb_mode=disabled"`. +- Run commands from `cookbooks/cosmos3/generator/action/finetune/`. +- In NGC / PyTorch containers, set `LD_LIBRARY_PATH=''` before Python commands. + +## Inputs You Provide + +This package ships the training stack — the registered `action_fd_droid_posttrain` +experiment, the dataset class, and the paired TOML/launch shell. Three inputs +are external and must be provided per environment: + +1. **[Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) dataset (in LeRobotDataset v3.0 format)** — pre-download the + dataset and point `DATASET_PATH` at the + resulting `.../Cosmos3-DROID` parent directory. This experiment trains on both + `success/` and `failure/` subsets. +2. **DCP base checkpoint** — prepare a base DCP checkpoint and point + `BASE_CHECKPOINT_PATH` at it. For local smoke runs this can be + `checkpoints/Cosmos3-Nano`. +3. **Wan2.2 VAE** — download `Wan2.2_VAE.pth` from + [`Wan-AI/Wan2.2-TI2V-5B`](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) + and point `WAN_VAE_PATH` at the file. + +## Data Layout + +Set `DATASET_PATH` to the Cosmos3-DROID parent directory: + +```shell +export DATASET_PATH=/path/to/Cosmos3-DROID +``` + +The loader uses both subsets: + +```text +$DATASET_PATH/ +├── success/ +│ └── meta/info.json +└── failure/ + └── meta/info.json +``` + +or institution-sharded split roots: + +```text +$DATASET_PATH/ +├── success//meta/info.json +└── failure//meta/info.json +``` + +The launch shell passes `DATASET_PATH` through to the training config. + +## Full Reproduction + +The flow mirrors the other SFT recipes: + +```shell +# Step 1: prepare Cosmos3-DROID -> $DATASET_PATH. +# If you do not already have the dataset locally, download it from Hugging Face. +uvx hf@latest download --repo-type dataset nvidia/Cosmos3-DROID \ + --local-dir data/Cosmos3-DROID --quiet +export DATASET_PATH=data/Cosmos3-DROID + +# Step 2: prepare the base DCP checkpoint and Wan2.2 VAE. +python -m cosmos_framework.scripts.convert_model_to_dcp \ + -o checkpoints/Cosmos3-Nano \ + --checkpoint-path Cosmos3-Nano +uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth \ + --local-dir checkpoints/wan22_vae +export BASE_CHECKPOINT_PATH=checkpoints/Cosmos3-Nano +export WAN_VAE_PATH=checkpoints/wan22_vae/Wan2.2_VAE.pth + +# Step 3: optionally choose the output root and launch. +export OUTPUT_ROOT=/path/to/output_root +export LD_LIBRARY_PATH='' + +bash launch_sft_action_fd_droid_posttrain.sh +``` + +`BASE_CHECKPOINT_PATH` should point at the base DCP checkpoint. When using the +paired launch shell, it must be a local directory because the shared launcher +validates it before starting `torchrun`. + +## Validate The Config + +Use `--dryrun` before launching workers: + +```shell +PYTHONPATH=. python -m cosmos_framework.scripts.train \ + --sft-toml toml/sft_config/action_fd_droid_posttrain.toml \ + --dryrun +``` + +## Run Training + +Recommended paired launch shell: + +```shell +bash launch_sft_action_fd_droid_posttrain.sh +``` + +Pass short smoke-run overrides through `EXTRA_TAIL_OVERRIDES`: + +```shell +export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10" +bash launch_sft_action_fd_droid_posttrain.sh +``` + +Single-node, 8 GPU: + +```shell +PYTHONPATH=. torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \ + --sft-toml toml/sft_config/action_fd_droid_posttrain.toml +``` + +Multi-node HSDP: + +```shell +PYTHONPATH=. torchrun --nnodes=$NNODES --node_rank=$NODE_RANK --nproc_per_node=8 \ + -m cosmos_framework.scripts.train \ + --sft-toml toml/sft_config/action_fd_droid_posttrain.toml \ + -- model.parallelism.data_parallel_replicate_degree=$NNODES +``` + +Keep `data_parallel_shard_degree=8` and set +`model.parallelism.data_parallel_replicate_degree` to the number of nodes. + +## Outputs + +Training outputs land under: + +```text +$IMAGINAIRE_OUTPUT_ROOT/cosmos3_action_fd/action_sft// +``` + +DCP checkpoints are saved under: + +```text +$RUN_DIR/checkpoints/iter_/ +``` + +The run is resumable by relaunching with the same output directory and job name. + +## Notes + +- This recipe uses `mode="forward_dynamics"`, so actions are conditioning and the + model trains video prediction from the first frame plus action sequence. +- `DROIDMergedLeRobotDataset` uses `action_space="ee_pose"` for the 10-D + end-effector pose action path. diff --git a/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh index 8dc7dd50..02b393be 100644 --- a/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh +++ b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh @@ -5,35 +5,121 @@ # ============================================================================ # Structured-TOML launch for action_fd_droid_posttrain. # -# This trains forward dynamics on the Cosmos3-DROID success + failure splits. See -# docs/action_fd_droid_posttrain.md. +# This script intentionally does not source _sft_launcher_common.sh. It contains +# the recipe-specific defaults, input checks, path anchoring, tail overrides, and +# torchrun invocation in one file. # -# Env vars (override for your filesystem): -# DATASET_PATH Cosmos3-DROID parent dir (success/ + failure/) -# BASE_CHECKPOINT_PATH Base DCP checkpoint -# WAN_VAE_PATH Wan2.2 VAE .pth -# WANDB_API_KEY for online logging (TOML wandb_mode="online") -# NPROC_PER_NODE torchrun --nproc_per_node (default 8) -# EXTRA_TAIL_OVERRIDES space-separated Hydra overrides -# -# Single-node smoke: -# export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10" -# bash examples/launch_sft_action_fd_droid_posttrain.sh +# Run with the cosmos-framework venv active (see action_fd_droid_posttrain.md). +# Training starts from the installed cosmos-framework checkout root so relative +# model config paths (e.g. Qwen3-VL JSON) resolve correctly. # -# Multi-node: launch on every worker. For HSDP set -# model.parallelism.data_parallel_replicate_degree = (shard stays 8). +# Env vars (override for your filesystem): +# DATASET_PATH Cosmos3-DROID parent dir (success/ + failure/) +# BASE_CHECKPOINT_PATH Base DCP checkpoint +# WAN_VAE_PATH Wan2.2 VAE .pth +# COSMOS_FRAMEWORK_ROOT Optional override for the framework checkout root +# NPROC_PER_NODE torchrun --nproc_per_node (default 8) +# EXTRA_TAIL_OVERRIDES space-separated Hydra overrides +# MASTER_PORT torchrun --master_port (default 50012) +# NNODES torchrun --nnodes for multi-node +# NODE_RANK torchrun --node_rank for multi-node +# MASTER_ADDR torchrun --master_addr for multi-node +# OUTPUT_ROOT default output root (default outputs/train) +# LOG_FILENAME override log filename # ============================================================================ -TOML_FILE="examples/toml/sft_config/action_fd_droid_posttrain.toml" -: "${DATASET_PATH:=examples/data/Cosmos3-DROID}" -: "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" +set -euo pipefail + +# Cookbook recipe directory (TOML, logs, default relative paths). +COOKBOOK_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +TOML_FILE="$COOKBOOK_DIR/toml/sft_config/action_fd_droid_posttrain.toml" + +: "${DATASET_PATH:=data/Cosmos3-DROID}" +: "${BASE_CHECKPOINT_PATH:=checkpoints/Cosmos3-Nano}" +: "${WAN_VAE_PATH:=checkpoints/wan22_vae/Wan2.2_VAE.pth}" + +[[ "$DATASET_PATH" = /* ]] || DATASET_PATH="$COOKBOOK_DIR/$DATASET_PATH" +[[ "$BASE_CHECKPOINT_PATH" = /* ]] || BASE_CHECKPOINT_PATH="$COOKBOOK_DIR/$BASE_CHECKPOINT_PATH" +[[ "$WAN_VAE_PATH" = /* ]] || WAN_VAE_PATH="$COOKBOOK_DIR/$WAN_VAE_PATH" +export DATASET_PATH BASE_CHECKPOINT_PATH WAN_VAE_PATH + +OUTPUT_ROOT="${OUTPUT_ROOT:-$COOKBOOK_DIR/outputs/train}" +LOG_DIR="$OUTPUT_ROOT/logs" +TOML_STEM="$(basename "$TOML_FILE" .toml)" +LOG_FILE="$LOG_DIR/${LOG_FILENAME:-${TOML_STEM}_sft.log}" +IMAGINAIRE_OUTPUT_ROOT="${IMAGINAIRE_OUTPUT_ROOT:-$OUTPUT_ROOT}" +mkdir -p "$LOG_DIR" -WAN_VAE_PATH="${WAN_VAE_PATH:-examples/checkpoints/wan22_vae/Wan2.2_VAE.pth}" +echo ">>> $(date '+%H:%M:%S') Checking inputs..." +[[ -f "$TOML_FILE" ]] || { echo "ERROR: TOML not found: $TOML_FILE" >&2; exit 1; } +[[ -d "$DATASET_PATH" ]] || { echo "ERROR: DATASET_PATH not found: $DATASET_PATH" >&2; exit 1; } +[[ -d "$BASE_CHECKPOINT_PATH" ]] || { echo "ERROR: BASE_CHECKPOINT_PATH not found: $BASE_CHECKPOINT_PATH" >&2; exit 1; } +[[ -f "$WAN_VAE_PATH" ]] || { echo "ERROR: WAN_VAE_PATH not found: $WAN_VAE_PATH" >&2; exit 1; } -EXTRA_DATASET_CHECK='[[ "$WAN_VAE_PATH" = /* ]] || WAN_VAE_PATH="$WORKDIR/$WAN_VAE_PATH"; export WAN_VAE_PATH; [[ -f "$DATASET_PATH/success/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/success/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID success split under $DATASET_PATH (expected success/meta/info.json or success/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$DATASET_PATH/failure/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/failure/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID failure split under $DATASET_PATH (expected failure/meta/info.json or failure/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$WAN_VAE_PATH" ]] || { echo "ERROR: WAN_VAE_PATH not found: $WAN_VAE_PATH" >&2; exit 1; }' +[[ -f "$DATASET_PATH/success/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/success/*/meta/info.json")" ]] || { + echo "ERROR: missing Cosmos3-DROID success split under $DATASET_PATH (expected success/meta/info.json or success/*/meta/info.json)" >&2 + exit 1 +} +[[ -f "$DATASET_PATH/failure/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/failure/*/meta/info.json")" ]] || { + echo "ERROR: missing Cosmos3-DROID failure split under $DATASET_PATH (expected failure/meta/info.json or failure/*/meta/info.json)" >&2 + exit 1 +} + +if [[ -n "${COSMOS_FRAMEWORK_ROOT:-}" ]]; then + FRAMEWORK_ROOT="$(cd "$COSMOS_FRAMEWORK_ROOT" && pwd)" +else + FRAMEWORK_ROOT="$( + python -c 'import pathlib, cosmos_framework; print(pathlib.Path(cosmos_framework.__file__).resolve().parents[1])' + )" +fi + +QWEN3_VL_CONFIG="$FRAMEWORK_ROOT/cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" +[[ -f "$QWEN3_VL_CONFIG" ]] || { + cat >&2 <>> $(date '+%H:%M:%S') cookbook: $COOKBOOK_DIR" +echo ">>> $(date '+%H:%M:%S') framework: $FRAMEWORK_ROOT" +echo ">>> $(date '+%H:%M:%S') TOML: $TOML_FILE" +echo ">>> $(date '+%H:%M:%S') dataset: $DATASET_PATH" +echo ">>> $(date '+%H:%M:%S') checkpoint: $BASE_CHECKPOINT_PATH" +echo ">>> $(date '+%H:%M:%S') log: $LOG_FILE" + +IMAGINAIRE_OUTPUT_ROOT="$IMAGINAIRE_OUTPUT_ROOT" \ +PYTHONPATH="${FRAMEWORK_ROOT}${PYTHONPATH:+:$PYTHONPATH}" \ + torchrun "${TORCHRUN_ARGS[@]}" -m cosmos_framework.scripts.train \ + --sft-toml="$TOML_FILE" \ + "${TRAILING_ARGS[@]}" \ + 2>&1 | tee "$LOG_FILE" -source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" +EXIT_CODE=${PIPESTATUS[0]} +echo ">>> $(date '+%H:%M:%S') Done (exit $EXIT_CODE)" +exit $EXIT_CODE From 14229f194448f8caa6ce28001f68c5c052daa118 Mon Sep 17 00:00:00 2001 From: Wei-Cheng Tseng Date: Wed, 22 Jul 2026 11:25:34 -0700 Subject: [PATCH 3/3] update doc link --- .../generator/action/finetune/action_fd_droid_posttrain.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md b/cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md index 99a2c04f..973349c7 100644 --- a/cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md +++ b/cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md @@ -20,7 +20,7 @@ It trains Cosmos3-DROID forward dynamics in `cosmos_framework`. ## Prerequisites - Install cosmos-framework as described in the shared - [setup instructions](cookbooks/cosmos3/README.md), then activate its + [setup instructions](../../../README.md), then activate its virtual environment. - Authenticate with Hugging Face (`uvx hf@latest auth login` or `HF_TOKEN`) and accept the terms for the required model and dataset repositories.