diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index a21c9341..07098e99 100644 --- a/cookbooks/cosmos3/generator/action/finetune/README.md +++ b/cookbooks/cosmos3/generator/action/finetune/README.md @@ -11,7 +11,7 @@ This example demonstrates supervised fine-tuning (SFT) of [Cosmos3-Nano](https:/ | Policy-LIBERO-10 SFT (A) | `launch_sft_action_policy_libero.sh` | Cosmos3-Nano | [LIBERO_LeRobot_v3](https://huggingface.co/datasets/nvidia/LIBERO_LeRobot_v3) `libero_10` | | Policy-LIBERO-all SFT (B) | `launch_sft_action_policy_libero_all.sh` | Cosmos3-Nano | [LIBERO_LeRobot_v3](https://huggingface.co/datasets/nvidia/LIBERO_LeRobot_v3) all 4 suites | -The DROID recipe uses the registered `action_policy_droid_nano` experiment: `joint_pos` 8-D actions, proprioceptive state, `concat_view` 480p video, chunk length 32, episode-shuffle streaming, and the optional `keep_ranges_1_0_1.json` window filter. +The DROID recipe uses the registered `action_policy_droid_nano` experiment: `joint_pos` 8-D actions, proprioceptive state, `concat_view` 480p video, chunk length 32, episode-shuffle streaming, JSON-formatted action prompts (`format_prompt_as_json=True`), and the optional `keep_ranges_1_0_1.json` window filter. The reference reproduction runs lr 2e-4 (cosine, cycle 100000), generator loss_scale 10, global batch 8192 (HSDP 32x8 = 256 ranks; GB200 reference, 64 nodes x 4), for 10000 iters. The action prompt is serialized as JSON at both train and eval time, so evaluation must use the matching JSON prompt format. The LIBERO recipe uses `frame_wise_relative` rot6d 10-D actions, `quantile_rot` normalization, `concat_view` (third-person + wrist) at 20 fps, lr 5e-5 / warmup 500 / cycle 16000, global batch 2048 (HSDP 2x8). To match the LIBERO-10 results reported in Cosmos3, we provide **two presets**: @@ -23,7 +23,7 @@ For a runnable egocentric hand-pose data conversion example, see converts a sample video and 3D hand-pose annotation pair into the raw 57D hand Action format used by the dataset path. -The recipe uses `[job].task = "vfm"` with the registered `action_policy_droid_nano` experiment. It trains a DROID policy model with `joint_pos` 8-D actions, proprioceptive state, `concat_view` 480p video, chunk length 32, episode-shuffle streaming, and the optional `keep_ranges_1_0_1.json` window filter. +The recipe uses `[job].task = "vfm"` with the registered `action_policy_droid_nano` experiment. It trains a DROID policy model with `joint_pos` 8-D actions, proprioceptive state, `concat_view` 480p video, chunk length 32, episode-shuffle streaming, JSON-formatted action prompts, and the optional `keep_ranges_1_0_1.json` window filter. ## Prerequisites @@ -73,13 +73,25 @@ export FILTER_PATH=/scratch/droid/keep_ranges_1_0_1.json bash launch_sft_action_policy_droid.sh ``` -To run a short smoke test, keep the same inputs and override the iteration/batch knobs: +The committed TOML pins the GB200 reference shape (HSDP 32x8 = 256 ranks, global +batch 8192). To run on a single 8-GPU node — e.g. a short smoke test — drop the +replicate degree to 1 alongside the iteration/batch knobs: ```shell -export EXTRA_TAIL_OVERRIDES="job.wandb_mode=disabled trainer.max_iter=10 checkpoint.save_iter=10 dataloader_train.max_samples_per_batch=32" +export EXTRA_TAIL_OVERRIDES=" \ + job.wandb_mode=disabled \ + trainer.max_iter=10 \ + checkpoint.save_iter=10 \ + model.config.parallelism.data_parallel_replicate_degree=1 \ + dataloader_train.max_samples_per_batch=32 \ +" bash launch_sft_action_policy_droid.sh ``` +To reproduce the reference at full global batch 8192 on fewer GPUs, keep +`data_parallel_replicate_degree=1` and raise `trainer.grad_accum_iter` (32 on one +8-GPU node) instead of shrinking the batch. + ## LIBERO quick start Each launcher stages its dataset (auto-downloaded if missing), downloads the Wan diff --git a/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_policy_droid.sh b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_policy_droid.sh index a31be61f..e13fe0e2 100755 --- a/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_policy_droid.sh +++ b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_policy_droid.sh @@ -2,7 +2,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -# Complete recipe: DROID action-policy SFT on Cosmos3-Nano (8x H100). +# Complete recipe: DROID action-policy SFT on Cosmos3-Nano. +# The TOML pins the GB200 reference shape (HSDP 32x8 = 256 ranks, global batch +# 8192); on fewer GPUs override data_parallel_replicate_degree / grad_accum_iter +# (see README). Set NNODES / NODE_RANK / MASTER_ADDR for multi-node. # Run from this folder with the cosmos-framework venv active (see README): # bash launch_sft_action_policy_droid.sh # It prepares the small dependencies, checks for the staged DROID dataset, and trains. @@ -48,7 +51,9 @@ if [[ ! -f "$FILTER_PATH" ]]; then uvx hf@latest download KarlP/droid keep_ranges_1_0_1.json --local-dir "$(dirname "$FILTER_PATH")" fi -# 5. Train (8-GPU FSDP by default). The TOML reads these paths from the environment. +# 5. Train. torchrun uses NPROC_PER_NODE GPUs (8 by default); the TOML's HSDP 32x8 +# shape needs 256 ranks, so scale nodes or override parallelism (see README). +# The TOML reads these paths from the environment. export DROID_ROOT="${DROID_ROOT:-$DATASET_PATH}" export BASE_CHECKPOINT_PATH export WAN_VAE_PATH diff --git a/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_policy_droid_repro.toml b/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_policy_droid_repro.toml index 661f7bdb..2326d407 100644 --- a/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_policy_droid_repro.toml +++ b/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_policy_droid_repro.toml @@ -3,10 +3,11 @@ # ============================================================================ # DROID action-policy SFT — run config for the `action_policy_droid_nano` -# experiment. The recipe knobs (optimizer/lr, scheduler type, grad_clip, -# count-based batch, action-head skip-on-load, dataset knobs) live in the -# registered experiment; this file only sets run-level scalars (iters, ckpt -# cadence, parallelism shape, wandb, VAE path). +# experiment. Reproduces the Cosmos3-Nano-Policy-DROID reference run: +# HSDP 32x8, global batch 8192, lr 2e-4, loss_scale 10, 10000 iters. The +# remaining recipe knobs (grad_clip, count-based batch, action-head +# skip-on-load, dataset knobs) live in the registered experiment; this file +# pins the run-level scalars that define the reference reproduction. # # Env required: # DROID_ROOT=/path/to/droid_lerobot_640x360/success @@ -27,8 +28,8 @@ wandb_mode = "online" precision = "bfloat16" [model.parallelism] -data_parallel_shard_degree = 8 # 8-GPU model shard; set replicate for multi-node HSDP -data_parallel_replicate_degree = 1 +data_parallel_shard_degree = 8 # 8-GPU model shard +data_parallel_replicate_degree = 32 # HSDP 32x8 = 256 ranks (GB200 reference: 64 nodes x 4) [model.activation_checkpointing] mode = "full" @@ -37,19 +38,37 @@ save_ops_regex = ["fmha"] [model.tokenizer] vae_path = "${oc.env:WAN_VAE_PATH}" +[model.rectified_flow_training_config] +loss_scale = 10.0 # generator (diffusion) loss weight, reference value + +[optimizer] +lr = 2.0e-04 + [scheduler] -cycle_lengths = [10000] # match max_iter +cycle_lengths = [100000] # long cosine cycle: lr decays slowly across the 10000-iter run + +[dataloader_train] +max_samples_per_batch = 32 # samples packed into each per-rank batch (res480) + +[dataloader_train.dataloader] +num_workers = 16 # decode-worker throughput tuning; adjust to host CPU count +batch_size = 16 +prefetch_factor = 2 [trainer] -max_iter = 10000 -logging_iter = 50 +max_iter = 10000 +logging_iter = 50 +grad_accum_iter = 1 # global batch = max_samples 32 x (shard 8 x replicate 32) x 1 = 8192 + +[trainer.callbacks.compile_tokenizer] +enabled = true # torch.compile the Wan VAE tokenizer (speed); warms up at res480 +warmup_resolutions = ["480"] [checkpoint] load_path = "${oc.env:BASE_CHECKPOINT_PATH}" save_iter = 1000 -# max_samples_per_batch is 128 in the experiment — samples packed into each per-rank batch -# (the num_workers x prefetch_factor workers just decode in parallel to keep it fed); res480, -# reference recipe, validated multi-node on GB200. -# On lower-memory GPUs, reduce it at launch, e.g.: -# --opts dataloader_train.max_samples_per_batch=32 +# The 256-rank HSDP 32x8 shape (global batch 8192) is the GB200 reference. To fit +# fewer GPUs while keeping global batch 8192, lower data_parallel_replicate_degree +# and raise grad_accum_iter at launch — e.g. one 8-GPU node: +# --opts model.parallelism.data_parallel_replicate_degree=1 trainer.grad_accum_iter=32