vla(fastwam): add ZeRO-1 SFT example

Author: XueSongTap

examples/embodied/fastwam/run_fastwam_sft_ddp.sh

@@ -21,8 +21,10 @@ TOKENIZER_PATH=${TOKENIZER_PATH:-"/path/to/tokenizer"}
 
 MASTER_PORT=${MASTER_PORT:-29519}
 
+MODEL_NAME=${MODEL_NAME:-"fastwam"}
+
 TRAINING_ARGS=(
-  --model-name fastwam
+  --model-name "$MODEL_NAME"
   --trainer-type FinetuneTrainer
   --distributed-strategy ddp
   --dtype bfloat16

examples/embodied/fastwam/run_fastwam_sft_zero1.sh

@@ -0,0 +1,140 @@
+#!/usr/bin/env bash
+# Copyright 2026 The LoongForge Authors.
+# SPDX-License-Identifier: Apache-2.0
+#
+# FastWAM SFT with ZeRO Stage-1 (multi-GPU DDP).
+#
+# Delta versus run_fastwam_sft_ddp.sh:
+#   --zero-optimizer                    wrap the optimizer in
+#                                       ZeroRedundancyOptimizer, sharding
+#                                       optimizer states across ranks. Only
+#                                       effective with --distributed-strategy ddp.
+#   --no-ddp-find-unused-parameters     skip the unused-parameter scan; FastWAM's
+#                                       forward graph has no conditional branches.
+#   --ddp-static-graph                  graph is identical every iteration, lets
+#                                       DDP reuse its bucket/reduction plan.
+#   --ddp-gradient-as-bucket-view       expose grads as views into the comm
+#                                       buckets instead of separate allocations.
+#   --no-ddp-broadcast-buffers          no BN-style buffers to sync each forward.
+#   --ddp-bucket-cap-mb                 larger buckets: fewer, bigger all-reduces.
+#
+# The memory saved by ZeRO-1 is what makes the larger --per-device-batch-size
+# below affordable relative to the plain DDP script.
+#
+# Two optional ZeRO knobs are left off by default:
+#   --zero-parameters-as-bucket-view    further cuts peak memory, but can clash
+#                                       with torch.compile + the DDP reducer.
+#   --zero-master-param-dtype fp32      rank-local fp32 master params, broadcast
+#                                       after each step. Better numerics under
+#                                       bf16 training at some bandwidth cost.
+#
+# Usage:
+#   DATASET_PATH=/path/to/libero TOKENIZER_PATH=/path/to/tokenizer \
+#     bash run_fastwam_sft_zero1.sh
+#   GPUS_PER_NODE=4 bash run_fastwam_sft_zero1.sh   # override via env
+#   ... bash run_fastwam_sft_zero1.sh --train-iters 50   # override a flag
+
+set -euo pipefail
+
+export LOONGFORGE_PATH="${LOONGFORGE_PATH:-/workspace/LoongForge}"
+
+DATASET_PATH=${DATASET_PATH:-/path/to/libero}
+TOKENIZER_PATH=${TOKENIZER_PATH:-/path/to/tokenizer}
+OUTPUT_DIR=${OUTPUT_DIR:-"outputs/fastwam_sft_zero1_$(date +%Y%m%d_%H%M%S)"}
+
+export CUBLAS_WORKSPACE_CONFIG=${CUBLAS_WORKSPACE_CONFIG:-:4096:8}
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+
+PRETRAINED_CHECKPOINT=${PRETRAINED_CHECKPOINT:-}
+ACTION_DIT_PRETRAINED_PATH=${ACTION_DIT_PRETRAINED_PATH:-}
+TEXT_EMBEDDING_CACHE_DIR=${TEXT_EMBEDDING_CACHE_DIR:-}
+
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+MASTER_PORT=${MASTER_PORT:-29519}
+
+MODEL_NAME=${MODEL_NAME:-"fastwam"}
+
+TRAINING_ARGS=(
+  --model-name "$MODEL_NAME"
+  --trainer-type FinetuneTrainer
+  --dtype bfloat16
+  --train-iters 20000
+  --save-interval 2000
+  --clip-grad 1.0
+  --gradient-accumulation-steps 1
+  --log-interval 1
+  --seed 3047
+  --output-dir "$OUTPUT_DIR"
+  --tokenizer-path "$TOKENIZER_PATH"
+)
+
+# ── ZeRO Stage-1 + DDP tuning (the point of this script) ──────
+ZERO1_ARGS=(
+  --distributed-strategy ddp
+  --zero-optimizer
+  --no-ddp-find-unused-parameters
+  --ddp-static-graph
+  --ddp-gradient-as-bucket-view
+  --no-ddp-broadcast-buffers
+  --ddp-bucket-cap-mb 200
+)
+
+if [[ -n "$PRETRAINED_CHECKPOINT" ]]; then
+  TRAINING_ARGS+=(--pretrained-checkpoint "$PRETRAINED_CHECKPOINT")
+fi
+
+LR_ARGS=(
+  --lr-base 1.0e-8
+  --lr-decay-style cosine_warmup_with_min_lr
+  --lr-warmup-iters 0
+  --min-lr 1.0e-9
+  --weight-decay 0.01
+  --adam-beta1 0.9
+  --adam-beta2 0.95
+)
+
+DATA_ARGS=(
+  --dataset-format lerobot_datasets
+  --dataset-strategy fastwam
+  --dataset-path "$DATASET_PATH"
+  --robot-type libero_franka
+  --per-device-batch-size 16
+  --num-workers 16
+  --lerobotdataset-version v2.1
+  --video-backend pyav
+)
+
+LOGGING_ARGS=(
+  --wandb-project loongforge-vla
+  --wandb-mode disabled
+)
+
+# ── Model/data dotlist overrides ──────────────────────────────
+MODEL_DATA_OVERRIDES=()
+if [[ -n "$ACTION_DIT_PRETRAINED_PATH" ]]; then
+  MODEL_DATA_OVERRIDES+=("model.action_dit_pretrained_path=$ACTION_DIT_PRETRAINED_PATH")
+fi
+if [[ -n "$TEXT_EMBEDDING_CACHE_DIR" ]]; then
+  MODEL_DATA_OVERRIDES+=("data.text_embedding_cache_dir=$TEXT_EMBEDDING_CACHE_DIR")
+fi
+
+mkdir -p "$OUTPUT_DIR"
+
+echo "════════════════════════════════════════════════════════════"
+echo "  LoongForge FastWAM SFT (DDP + ZeRO-1)"
+echo "  Model:    $MODEL_NAME"
+echo "  GPUs:     $GPUS_PER_NODE"
+echo "  Data:     $DATASET_PATH"
+echo "  Output:   $OUTPUT_DIR"
+echo "════════════════════════════════════════════════════════════"
+
+PYTHONPATH="$LOONGFORGE_PATH:${PYTHONPATH:-}" \
+  torchrun --nproc_per_node "$GPUS_PER_NODE" --master_port "$MASTER_PORT" \
+  "$LOONGFORGE_PATH/loongforge/embodied/train.py" \
+  "${TRAINING_ARGS[@]}" \
+  "${ZERO1_ARGS[@]}" \
+  "${LR_ARGS[@]}" \
+  "${DATA_ARGS[@]}" \
+  "${LOGGING_ARGS[@]}" \
+  "${MODEL_DATA_OVERRIDES[@]+"${MODEL_DATA_OVERRIDES[@]}"}" \
+  "$@" 2>&1 | tee "$OUTPUT_DIR/$(basename "${BASH_SOURCE[0]}" .sh)_$(date +%Y%m%d_%H%M%S).log"