From c0b9954678eefe62e37e55d7daca37397211a164 Mon Sep 17 00:00:00 2001 From: Liang Hao Date: Mon, 6 Jul 2026 14:00:28 +0800 Subject: [PATCH 1/6] cookbook(action/droid): link numerical reproducibility regression test Add a 'Verifying numerical reproducibility' section to the DROID/LIBERO action-policy finetune cookbook, pointing at cosmos-framework's deterministic action_policy_regression_test with the reference H200 LIBERO-10 loss series. Companion to cosmos-framework PR #86. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generator/action/finetune/README.md | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index a21c9341..83ba8861 100644 --- a/cookbooks/cosmos3/generator/action/finetune/README.md +++ b/cookbooks/cosmos3/generator/action/finetune/README.md @@ -121,6 +121,27 @@ Training writes to `outputs/train////`: - `checkpoints/iter_/` — DCP checkpoint (model / optim / scheduler / trainer state); `checkpoints/latest_checkpoint.txt` names the newest. - `config.yaml`, launch metadata, logs, and one directory per registered callback. +## Verifying numerical reproducibility + +Before committing to a long run, you can confirm your environment reproduces the +reference training dynamics bit-for-bit. cosmos-framework ships a deterministic +regression test that re-runs these exact recipes for 10 iterations and asserts the +rank-0 per-step loss series against per-arch goldens: + +```shell +# from your cosmos-framework checkout, on a single 4-GPU node +COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS= \ + pytest -s tests/action_policy_regression_test.py -k action_policy_libero +``` + +The LIBERO-10 spec auto-downloads `libero_10` and needs no extra setup; the DROID +spec runs only when `DROID_ROOT` points at a staged Cosmos3-DROID root. Goldens are +captured on H200 (Hopper) with `torch.compile` on, `--deterministic`, and seed 42; +the reference H200 LIBERO-10 loss series is `[15.81, 15.25, 15.99, 16.53, 14.37, +16.15, 16.61, 14.88, 16.06, 16.64]`. See the test's module docstring for how goldens +are keyed by GPU arch and refreshed. This is the companion to cosmos-framework PR +[#86](https://github.com/NVIDIA/cosmos-framework/pull/86). + ## Export to Hugging Face safetensors ```shell From b76f98f4774a2ccd7f497a1171c879ff170d36d4 Mon Sep 17 00:00:00 2001 From: Liang Hao Date: Mon, 6 Jul 2026 15:51:39 +0800 Subject: [PATCH 2/6] cookbook(action/droid): match GB200 reference config (HSDP 32x8, GBS 8192, lr 2e-4) Pin the DROID repro TOML to the actual reference run shape instead of the single-node placeholder: data_parallel_replicate_degree=32 (HSDP 32x8 = 256 ranks), lr=2e-4, max_samples_per_batch=32, grad_accum_iter=1 -> global batch 8192, matching the LIBERO recipe's convention of encoding the true reference shape. Document the shape in the README/launcher and give the single-node scale-down override. Drops the earlier numerical-reproducibility section. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generator/action/finetune/README.md | 33 +++++-------------- .../launch_sft_action_policy_droid.sh | 9 +++-- .../sft_config/action_policy_droid_repro.toml | 33 +++++++++++-------- 3 files changed, 36 insertions(+), 39 deletions(-) diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index 83ba8861..f4b6c741 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, and the optional `keep_ranges_1_0_1.json` window filter. The reference reproduction runs lr 2e-4 / cycle 10000, global batch 8192 (HSDP 32x8 = 256 ranks; 64x4 on GB200 or 32x8 on H100), for 10000 iters. 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**: @@ -73,13 +73,19 @@ 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 @@ -121,27 +127,6 @@ Training writes to `outputs/train////`: - `checkpoints/iter_/` — DCP checkpoint (model / optim / scheduler / trainer state); `checkpoints/latest_checkpoint.txt` names the newest. - `config.yaml`, launch metadata, logs, and one directory per registered callback. -## Verifying numerical reproducibility - -Before committing to a long run, you can confirm your environment reproduces the -reference training dynamics bit-for-bit. cosmos-framework ships a deterministic -regression test that re-runs these exact recipes for 10 iterations and asserts the -rank-0 per-step loss series against per-arch goldens: - -```shell -# from your cosmos-framework checkout, on a single 4-GPU node -COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS= \ - pytest -s tests/action_policy_regression_test.py -k action_policy_libero -``` - -The LIBERO-10 spec auto-downloads `libero_10` and needs no extra setup; the DROID -spec runs only when `DROID_ROOT` points at a staged Cosmos3-DROID root. Goldens are -captured on H200 (Hopper) with `torch.compile` on, `--deterministic`, and seed 42; -the reference H200 LIBERO-10 loss series is `[15.81, 15.25, 15.99, 16.53, 14.37, -16.15, 16.61, 14.88, 16.06, 16.64]`. See the test's module docstring for how goldens -are keyed by GPU arch and refreshed. This is the companion to cosmos-framework PR -[#86](https://github.com/NVIDIA/cosmos-framework/pull/86). - ## Export to Hugging Face safetensors ```shell 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..090707ac 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, 10000 iters. The remaining recipe +# knobs (scheduler type, 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 (64x4 on GB200, 32x8 on H100) [model.activation_checkpointing] mode = "full" @@ -37,19 +38,25 @@ save_ops_regex = ["fmha"] [model.tokenizer] vae_path = "${oc.env:WAN_VAE_PATH}" +[optimizer] +lr = 2.0e-04 + [scheduler] cycle_lengths = [10000] # match max_iter +[dataloader_train] +max_samples_per_batch = 32 # samples packed into each per-rank batch (res480) + [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 [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 From db26c6154236ea380fd558c6b88d9e4f61f4b4d3 Mon Sep 17 00:00:00 2001 From: Liang Hao Date: Mon, 6 Jul 2026 15:57:50 +0800 Subject: [PATCH 3/6] cookbook(action/droid): note JSON action-prompt format in intro The action_policy_droid_nano experiment trains with format_prompt_as_json=True; call this out in the DROID recipe description and flag that eval must use the matching JSON prompt format. Co-Authored-By: Claude Opus 4.8 (1M context) --- cookbooks/cosmos3/generator/action/finetune/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index f4b6c741..706f6149 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 reference reproduction runs lr 2e-4 / cycle 10000, global batch 8192 (HSDP 32x8 = 256 ranks; 64x4 on GB200 or 32x8 on H100), for 10000 iters. +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 / cycle 10000, global batch 8192 (HSDP 32x8 = 256 ranks; 64x4 on GB200 or 32x8 on H100), 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 From 1bc070c385c9cfe6c6f7c7289442e6f364688aa6 Mon Sep 17 00:00:00 2001 From: Liang Hao Date: Mon, 6 Jul 2026 16:06:14 +0800 Subject: [PATCH 4/6] cookbook(action/droid): drop non-GB200 node-count mapping Keep the HSDP 32x8 reference described as the GB200 shape (64 nodes x 4) without recommending an H100/H200 equivalent layout. Co-Authored-By: Claude Opus 4.8 (1M context) --- cookbooks/cosmos3/generator/action/finetune/README.md | 2 +- .../finetune/toml/sft_config/action_policy_droid_repro.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index 706f6149..4942eea6 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, 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 / cycle 10000, global batch 8192 (HSDP 32x8 = 256 ranks; 64x4 on GB200 or 32x8 on H100), 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 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 / cycle 10000, 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**: 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 090707ac..221d978f 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 @@ -29,7 +29,7 @@ precision = "bfloat16" [model.parallelism] data_parallel_shard_degree = 8 # 8-GPU model shard -data_parallel_replicate_degree = 32 # HSDP 32x8 = 256 ranks (64x4 on GB200, 32x8 on H100) +data_parallel_replicate_degree = 32 # HSDP 32x8 = 256 ranks (GB200 reference: 64 nodes x 4) [model.activation_checkpointing] mode = "full" From 5899d525d2b3e1b9ac79d38d43b4ad545c194e11 Mon Sep 17 00:00:00 2001 From: Liang Hao Date: Mon, 6 Jul 2026 18:51:45 +0800 Subject: [PATCH 5/6] cookbook(action/droid): self-contained TOML (loss_scale, cycle, compile, dataloader) Sync the DROID repro TOML with the framework copy: fold in loss_scale=10, scheduler.cycle_lengths=[100000], compile_tokenizer, and dataloader worker settings so the recipe reproduces the reference run without launcher overrides. Update the README description (cosine cycle 100000, generator loss_scale 10). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generator/action/finetune/README.md | 2 +- .../sft_config/action_policy_droid_repro.toml | 22 ++++++++++++++----- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index 4942eea6..9977668b 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, 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 / cycle 10000, 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 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**: 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 221d978f..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 @@ -4,10 +4,10 @@ # ============================================================================ # DROID action-policy SFT — run config for the `action_policy_droid_nano` # experiment. Reproduces the Cosmos3-Nano-Policy-DROID reference run: -# HSDP 32x8, global batch 8192, lr 2e-4, 10000 iters. The remaining recipe -# knobs (scheduler type, 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. +# 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 @@ -38,20 +38,32 @@ 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 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 From cb68995d63c8c30478b79568db9c7ca9160c9044 Mon Sep 17 00:00:00 2001 From: Liang Hao Date: Mon, 6 Jul 2026 23:14:35 +0800 Subject: [PATCH 6/6] =?UTF-8?q?cookbook(action/droid):=20apply=20review=20?= =?UTF-8?q?suggestion=20=E2=80=94=20multi-line=20smoke=20override?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reformat the DROID single-node smoke-test EXTRA_TAIL_OVERRIDES onto multiple lines for readability, per @ychao-nvidia's suggestion on PR #261. Co-Authored-By: Claude Opus 4.8 (1M context) --- cookbooks/cosmos3/generator/action/finetune/README.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index 9977668b..07098e99 100644 --- a/cookbooks/cosmos3/generator/action/finetune/README.md +++ b/cookbooks/cosmos3/generator/action/finetune/README.md @@ -78,7 +78,13 @@ batch 8192). To run on a single 8-GPU node — e.g. a short smoke test — drop 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 model.config.parallelism.data_parallel_replicate_degree=1 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 ```