finetune.sh

#!/usr/bin/env bash

set -x -euo pipefail

NUM_GPUS="${NUM_GPUS:-1}"
MASTER_PORT="${MASTER_PORT:-29500}"
SAVE_STEPS="${SAVE_STEPS:-1000}"
MAX_STEPS="${MAX_STEPS:-10000}"
USE_WANDB="${USE_WANDB:-1}"
DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-32}"
# Micro-batches accumulated per optimizer step. The effective batch size is
# GLOBAL_BATCH_SIZE x GRADIENT_ACCUMULATION_STEPS, so lowering GLOBAL_BATCH_SIZE
# while raising this keeps optimization identical at a lower peak VRAM
# (activation memory scales with GLOBAL_BATCH_SIZE / NUM_GPUS).
GRADIENT_ACCUMULATION_STEPS="${GRADIENT_ACCUMULATION_STEPS:-1}"
SHARD_SIZE="${SHARD_SIZE:-1024}"
NUM_SHARDS_PER_EPOCH="${NUM_SHARDS_PER_EPOCH:-100000}"
EPISODE_SAMPLING_RATE="${EPISODE_SAMPLING_RATE:-0.1}"
DS_WEIGHTS_ALPHA="${DS_WEIGHTS_ALPHA:-}"

BASE_MODEL_PATH=""
DATASET_PATH=""
MODALITY_CONFIG_PATH=""
EMBODIMENT_TAG=""
OUTPUT_DIR=""
EXPERIMENT_NAME=""
WANDB_PROJECT=""
STATE_DROPOUT_PROB=""
COLOR_JITTER_PARAMS="${COLOR_JITTER_PARAMS:-brightness 0.3 contrast 0.4 saturation 0.5 hue 0.08}"
USE_PERCENTILES=""
SHORTEST_IMAGE_EDGE=""
CROP_FRACTION=""
EXTRA_ARGS=()

usage() {
    cat <<'EOF'
Usage: bash examples/finetune.sh \
  --base-model-path <path> \
  --dataset-path <path> \
  --embodiment-tag <tag> \
  --output-dir <path> \
  [--modality-config-path <path>] \
  [--state-dropout-prob <value>] \
  [--color-jitter-params "brightness 0.3 contrast 0.4 saturation 0.5 hue 0.08"] \
  [--use-percentiles <true|false>] \
  [--shortest-image-edge <pixels>] \
  [--crop-fraction <fraction>] \
  [--ds-weights-alpha <value>] \
  [--save-only-model] \
  [--resume-from-checkpoint] \
  [-- <extra launch_finetune.py args>...]

Environment variables:
  NUM_GPUS, MASTER_PORT, SAVE_STEPS, MAX_STEPS, USE_WANDB,
  DATALOADER_NUM_WORKERS, GLOBAL_BATCH_SIZE, GRADIENT_ACCUMULATION_STEPS,
  SHARD_SIZE, NUM_SHARDS_PER_EPOCH, EPISODE_SAMPLING_RATE, DS_WEIGHTS_ALPHA

Low-VRAM example (same effective batch as the defaults, lower peak memory):
  GLOBAL_BATCH_SIZE=8 GRADIENT_ACCUMULATION_STEPS=4 bash examples/finetune.sh ...
EOF
}

while [ "$#" -gt 0 ]; do
    case "$1" in
        --base-model-path)
            BASE_MODEL_PATH="$2"
            shift 2
            ;;
        --dataset-path)
            DATASET_PATH="$2"
            shift 2
            ;;
        --modality-config-path)
            MODALITY_CONFIG_PATH="$2"
            shift 2
            ;;
        --embodiment-tag)
            EMBODIMENT_TAG="$2"
            shift 2
            ;;
        --output-dir)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        --experiment-name)
            EXPERIMENT_NAME="$2"
            shift 2
            ;;
        --wandb-project)
            WANDB_PROJECT="$2"
            shift 2
            ;;
        --state-dropout-prob)
            STATE_DROPOUT_PROB="$2"
            shift 2
            ;;
        --color-jitter-params)
            COLOR_JITTER_PARAMS="$2"
            shift 2
            ;;
        --use-percentiles)
            USE_PERCENTILES="$2"
            shift 2
            ;;
        --shortest-image-edge)
            SHORTEST_IMAGE_EDGE="$2"
            shift 2
            ;;
        --crop-fraction)
            CROP_FRACTION="$2"
            shift 2
            ;;
        --ds-weights-alpha)
            DS_WEIGHTS_ALPHA="$2"
            shift 2
            ;;
        --save-only-model)
            SAVE_ONLY_MODEL=1
            shift
            ;;
        --resume-from-checkpoint)
            RESUME_FROM_CHECKPOINT=1
            shift
            ;;
        --help|-h)
            usage
            exit 0
            ;;
        --)
            shift
            EXTRA_ARGS=("$@")
            break
            ;;
        *)
            echo "Unknown argument: $1" >&2
            usage >&2
            exit 1
            ;;
    esac
done

for required_var in BASE_MODEL_PATH DATASET_PATH EMBODIMENT_TAG OUTPUT_DIR; do
    if [ -z "${!required_var}" ]; then
        echo "Missing required argument: ${required_var}" >&2
        usage >&2
        exit 1
    fi
done

WANDB_FLAG=()
if [ "$USE_WANDB" = "1" ]; then
    WANDB_FLAG+=(--use_wandb)
fi

LAUNCH_CMD=(
    gr00t/experiment/launch_finetune.py
    --base_model_path "$BASE_MODEL_PATH"
    --dataset_path "$DATASET_PATH"
    --embodiment_tag "$EMBODIMENT_TAG"
    --num_gpus "$NUM_GPUS"
    --output_dir "$OUTPUT_DIR"
    --save_steps "$SAVE_STEPS"
    --save_total_limit 5
    --max_steps "$MAX_STEPS"
    --warmup_ratio 0.05
    --weight_decay 1e-5
    --learning_rate 1e-4
    "${WANDB_FLAG[@]}"
    --global_batch_size "$GLOBAL_BATCH_SIZE"
    --gradient_accumulation_steps "$GRADIENT_ACCUMULATION_STEPS"
    --dataloader_num_workers "$DATALOADER_NUM_WORKERS"
    --shard_size "$SHARD_SIZE"
    --num_shards_per_epoch "$NUM_SHARDS_PER_EPOCH"
    --episode_sampling_rate "$EPISODE_SAMPLING_RATE"
)

if [ -n "$MODALITY_CONFIG_PATH" ]; then
    LAUNCH_CMD+=(--modality_config_path "$MODALITY_CONFIG_PATH")
fi
if [ -n "$EXPERIMENT_NAME" ]; then
    LAUNCH_CMD+=(--experiment_name "$EXPERIMENT_NAME")
fi
if [ -n "$WANDB_PROJECT" ]; then
    LAUNCH_CMD+=(--wandb_project "$WANDB_PROJECT")
fi

if [ -n "$STATE_DROPOUT_PROB" ]; then
    LAUNCH_CMD+=(--state_dropout_prob "$STATE_DROPOUT_PROB")
fi
if [ -n "$COLOR_JITTER_PARAMS" ]; then
    read -r -a COLOR_JITTER_ARGS <<< "$COLOR_JITTER_PARAMS"
    LAUNCH_CMD+=(--color_jitter_params "${COLOR_JITTER_ARGS[@]}")
fi
if [ -n "$USE_PERCENTILES" ]; then
    USE_PERCENTILES_NORMALIZED="$(printf '%s' "$USE_PERCENTILES" | tr '[:upper:]' '[:lower:]')"
    case "$USE_PERCENTILES_NORMALIZED" in
        1|true|yes|on)
            LAUNCH_CMD+=(--use-percentiles)
            ;;
        0|false|no|off)
            LAUNCH_CMD+=(--no-use-percentiles)
            ;;
        *)
            echo "Invalid --use-percentiles value: $USE_PERCENTILES" >&2
            exit 1
            ;;
    esac
fi
if [ -n "$SHORTEST_IMAGE_EDGE" ]; then
    LAUNCH_CMD+=(--shortest-image-edge "$SHORTEST_IMAGE_EDGE")
fi
if [ -n "$CROP_FRACTION" ]; then
    LAUNCH_CMD+=(--crop-fraction "$CROP_FRACTION")
fi
if [ -n "$DS_WEIGHTS_ALPHA" ]; then
    LAUNCH_CMD+=(--ds_weights_alpha "$DS_WEIGHTS_ALPHA")
fi
if [ -n "${SAVE_ONLY_MODEL:-}" ]; then
    LAUNCH_CMD+=(--save_only_model)
fi
if [ -n "${RESUME_FROM_CHECKPOINT:-}" ]; then
    LAUNCH_CMD+=(--resume_from_checkpoint)
fi

if [ "${#EXTRA_ARGS[@]}" -gt 0 ]; then
    LAUNCH_CMD+=("${EXTRA_ARGS[@]}")
fi

if [ "$NUM_GPUS" = "1" ]; then
    # Restrict to a single GPU so HF Trainer doesn't wrap the model in DataParallel,
    # which crashes with a StopIteration error in the model's device property.
    export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
    exec python "${LAUNCH_CMD[@]}"
fi

exec torchrun --nproc_per_node="$NUM_GPUS" --master_port="$MASTER_PORT" "${LAUNCH_CMD[@]}"