LoongForge/examples/embodied/pi05/run_pi05_fsdp_finetune.sh at e7a9a590ec64252984b857f550f7d58d5bb4edaa · XueSongTap/LoongForge · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#!/usr/bin/env bash
# ============================================================================
# run_pi05_fsdp_finetune.sh - pi05 VLA SFT Launch Script (FSDP, Embodied)
#
# Mirrors run_pi05_ddp_zero1_finetune.sh, but uses the embodied FSDP strategy.
#
# Usage:
#   bash run_pi05_fsdp_finetune.sh
#   bash run_pi05_fsdp_finetune.sh --train-iters 50000                   # override a training param (flag form)
#   bash run_pi05_fsdp_finetune.sh model.state_dim=8 data.image_size=448 # override YAML model:/data: fields (dotlist form)
# ============================================================================
set -euo pipefail

export LOONGFORGE_PATH=${LOONGFORGE_PATH:-"/workspace/LoongForge"}

# -- Paths --------------------------------------------------------------------
TOKENIZER_PATH=${TOKENIZER_PATH:-"/workspace/paligemma-3b-pt-224"}
CHECKPOINT_PATH=${CHECKPOINT_PATH:-"/workspace/pi05_base"}
DATA_PATH=${DATA_PATH:-"/workspace/libero"}
OUTPUT_DIR=${OUTPUT_DIR:-"/workspace/outputs/"}
TENSORBOARD_PATH=${TENSORBOARD_PATH:-"/workspace/tensorboard-log"}

export CUDA_DEVICE_MAX_CONNECTIONS=8
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# -- Distributed --------------------------------------------------------------
GPUS_PER_NODE=${GPUS_PER_NODE:-8}
MASTER_ADDR=${MASTER_ADDR:-"localhost"}
MASTER_PORT=${MASTER_PORT:-"29500"}
NNODES=${WORLD_SIZE:-"1"}
NODE_RANK=${RANK:-"0"}

DISTRIBUTED_ARGS=(
    --nproc_per_node $GPUS_PER_NODE
    --nnodes $NNODES
    --node_rank $NODE_RANK
    --master_addr $MASTER_ADDR
    --master_port $MASTER_PORT
)

# -- Model config -------------------------------------------------------------
MODEL_NAME=${MODEL_NAME:-"pi05"}
MODEL_CONFIG_ARGS=(
    --model-name $MODEL_NAME
)

# -- Data params --------------------------------------------------------------
DATA_ARGS=(
    --dataset-format lerobot_datasets
    --dataset-path $DATA_PATH
    --tokenizer-path $TOKENIZER_PATH
    --robot-type libero_franka
    --num-workers 16
    --batch-drop-last
)

# -- Training params (aligned with examples/pi05/finetuning/sft_pi05.sh) ------
# Old Megatron config: micro-batch-size=16, global-batch-size=128, 8 GPUs
# => gradient-accumulation-steps = 128 / (16 * 8) = 1
TRAINING_ARGS=(
    --trainer-type FinetuneTrainer
    --train-iters 30000
    --per-device-batch-size 16
    --gradient-accumulation-steps 1
    --seed 1234
    --output-dir $OUTPUT_DIR
    # Learning rate
    --lr-base 2.5e-8
    --min-lr 0
    --lr-decay-style cosine
    --lr-warmup-iters 0
    # Optimizer
    --optimizer AdamW
    --clip-grad 1.0
    --weight-decay 0.01
    --adam-beta1 0.9
    --adam-beta2 0.95
    --adam-eps 1e-8
    # Checkpoint
    --save-interval 30000
    --pretrained-checkpoint $CHECKPOINT_PATH
)

FSDP_NO_WRAP_MODULES=(
    PaliGemmaWithExpertModel
    PaliGemmaForConditionalGenerationWithPiGemma
    PaliGemmaModelWithPiGemma
    PiGemmaForCausalLM
    PiGemmaModel
    _PiGemmaDecoderLayerBase
    GemmaAttention
)

DISTRIBUTED_TRAINING_ARGS=(
    --distributed-strategy fsdp
    --dtype bfloat16
    --fsdp-reshard-root false
    --fsdp-no-wrap-modules "$(
        IFS=,
        echo "${FSDP_NO_WRAP_MODULES[*]}"
    )"
)

# -- Logging params -----------------------------------------------------------
LOGGING_ARGS=(
    --log-interval 1
    --tensorboard-dir ${TENSORBOARD_PATH}
)

# -- Launch -------------------------------------------------------------------
echo "========================================================================"
echo "  LoongForge pi05 Training (FSDP)"
echo "  Model:      $MODEL_NAME"
echo "  GPUs:       $GPUS_PER_NODE x $NNODES nodes"
echo "  Data:       $DATA_PATH"
echo "  Checkpoint: $CHECKPOINT_PATH"
echo "  Output:     $OUTPUT_DIR"
echo "========================================================================"

PYTHONPATH=$LOONGFORGE_PATH:${PYTHONPATH:-} \
    torchrun "${DISTRIBUTED_ARGS[@]}" \
    "$LOONGFORGE_PATH/loongforge/embodied/train.py" \
    "${MODEL_CONFIG_ARGS[@]}" \
    "${DATA_ARGS[@]}" \
    "${TRAINING_ARGS[@]}" \
    "${DISTRIBUTED_TRAINING_ARGS[@]}" \
    "${LOGGING_ARGS[@]}" \
    "$@"