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..973349c7 --- /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](../../../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 new file mode 100644 index 00000000..02b393be --- /dev/null +++ b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh @@ -0,0 +1,125 @@ +#!/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 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. +# +# 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. +# +# 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 +# ============================================================================ + +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" + +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; } + +[[ -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" + +EXIT_CODE=${PIPESTATUS[0]} +echo ">>> $(date '+%H:%M:%S') Done (exit $EXIT_CODE)" +exit $EXIT_CODE 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