1assets/pi05_libero/physical-intelligence/libero/norm_stats.json

1/root/workspace/openpi/.venv/bin/python3 scripts/compute_norm_stats.py --config-name pi05_libero

1# 模型权重
2export CHECKPOINT_PATH=
3
4# 数据集
5export DATA_PATH=

1JAX_PLATFORMS=cpu \
2BATCH_SIZE=256 \
3GRAD_ACC=2 \
4NUM_TRAIN_STEPS=50 \ CHECKPOINT_PATH=/mnt/pfs/models/pi05_base \
5DATA_PATH=/mnt/pfs/datasets bash \
6bash /root/workspace/run_train.sh

1#!/usr/bin/env bash
2set -o pipefail
3set -x
4
5WORKSPACE_DIR=${WORKSPACE_DIR:-"/root/workspace/openpi"}
6cd "$WORKSPACE_DIR"
7
8PY=${PY:-"$WORKSPACE_DIR/.venv/bin/python"}
9
10CHECKPOINT_PATH=${CHECKPOINT_PATH:-"/root/workspace/checkpoints/pi05_base"}
11DATA_PATH=${DATA_PATH:-"/root/workspace/datasets/libero"}
12ASSETS_BASE_DIR=${ASSETS_BASE_DIR:-"$WORKSPACE_DIR/assets"}
13OPENPI_DATA_HOME=${OPENPI_DATA_HOME:-"$WORKSPACE_DIR/.cache_openpi"}
14OUTPUT_DIR=${OUTPUT_DIR:-"/root/workspace/outputs"}
15TENSORBOARD_PATH=${TENSORBOARD_PATH:-"/root/workspace/tensorboard-log"}
16TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-"$WORKSPACE_DIR/torchinductor_cache"}
17
18mkdir -p "$OUTPUT_DIR" "$TENSORBOARD_PATH"
19
20BATCH_SIZE=${BATCH_SIZE:-384}
21GRAD_ACC=${GRAD_ACC:-2}
22GRAD_ACCUM_MODE=${GRAD_ACCUM_MODE:-always_sync}
23NUM_TRAIN_STEPS=${NUM_TRAIN_STEPS:-12}
24EXP_NAME=${EXP_NAME:-openpi_pi05_libero_gbs${BATCH_SIZE}_acc${GRAD_ACC}}
25
26COMPILE_SCOPE=${COMPILE_SCOPE:-multi_group}
27GEMMA_MODE=${GEMMA_MODE:-max-autotune-no-cudagraphs}
28GEMMA_GC_EVERY=${GEMMA_GC_EVERY:-0}
29BUCKET_MB=${BUCKET_MB:-200}
30SAVE_INTERVAL=${SAVE_INTERVAL:-999999}
31
32export JAX_PLATFORMS=cpu
33export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
34export WANDB_MODE=offline WANDB_DISABLED=true
35export HF_LEROBOT_HOME="$DATA_PATH"
36export OPENPI_DATA_HOME
37export TORCHINDUCTOR_CACHE_DIR
38export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
39export TORCHINDUCTOR_MIX_ORDER_REDUCTION=0
40export OMP_NUM_THREADS=${OMP_NUM_THREADS:-8}
41export TORCH_DYNAMO_RECOMPILE_LIMIT=64
42export TORCHDYNAMO_CACHE_SIZE_LIMIT=256
43
44NNODES=${NNODES:-${WORLD_SIZE:-1}}
45NODE_RANK=${NODE_RANK:-${RANK:-0}}
46MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
47MASTER_PORT=${MASTER_PORT:-23456}
48NPROC=${NPROC:-${NPROC_PER_NODE:-8}}
49unset RANK WORLD_SIZE LOCAL_RANK
50$PY -m torch.distributed.run --nnodes=$NNODES --node_rank=$NODE_RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT --nproc_per_node=$NPROC \
51  scripts/train_pytorch.py pi05_libero \
52  --exp_name "$EXP_NAME" \
53  --ddp-grad-accum-mode "$GRAD_ACCUM_MODE" \
54  --batch_size "$BATCH_SIZE" --gradient_accumulation_steps "$GRAD_ACC" \
55  --num_train_steps "$NUM_TRAIN_STEPS" --log_interval 1 --save_interval "$SAVE_INTERVAL" --no-wandb_enabled \
56  --no-ddp-find-unused-parameters --ddp-bucket-cap-mb "$BUCKET_MB" \
57  --model.pytorch-compile-scope "$COMPILE_SCOPE" \
58  --model.pytorch-compile-gemma-mode "$GEMMA_MODE" \
59  --gemma-gc-every-n-layers "$GEMMA_GC_EVERY" \
60  --pytorch_weight_path "$CHECKPOINT_PATH" \
61  --assets_base_dir "$ASSETS_BASE_DIR" \
62  --checkpoint_base_dir "$OUTPUT_DIR" \
63  "$@"