cookbook(action/droid): align DROID recipe with GB200 reference reproduction

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

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

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