File size: 1,906 Bytes
2847d0b | 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 | #!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
PYTHON_BIN="${PYTHON_BIN:-/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/self_forcing/bin/python}"
CONFIG_PATH="${CONFIG_PATH:?CONFIG_PATH is required}"
OUTPUT_DIR="${OUTPUT_DIR:?OUTPUT_DIR is required}"
GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}"
MASTER_PORT="${MASTER_PORT:-29741}"
USE_SWANLAB="${USE_SWANLAB:-true}"
MAX_STUDENT_STEPS="${MAX_STUDENT_STEPS:-}"
FORCED_EXIT_STEP="${FORCED_EXIT_STEP:-}"
cd "$ROOT_DIR"
IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
NPROC_PER_NODE="${#GPU_ARRAY[@]}"
if (( NPROC_PER_NODE != 8 )); then
echo "错误：正式 Predictor-v4 DMD 训练需要 8 张 GPU，当前为 $NPROC_PER_NODE。" >&2
exit 1
fi
for required in "$PYTHON_BIN" "$CONFIG_PATH"; do
if [[ ! -e "$required" ]]; then
echo "错误：找不到 $required" >&2
exit 1
fi
done
mkdir -p "$OUTPUT_DIR"
TRAIN_TMPDIR="${TRAIN_TMPDIR:-/mnt/local_nvme/zoubin/cz/tmp/predictor_v4_dmd}"
mkdir -p "$TRAIN_TMPDIR"
overrides=(
--config_override "predictor_v4_dmd.output_dir=${OUTPUT_DIR}"
--config_override "predictor_v4_dmd.use_swanlab=${USE_SWANLAB}"
)
if [[ -n "$MAX_STUDENT_STEPS" ]]; then
overrides+=(
--config_override "predictor_v4_dmd.max_student_steps=${MAX_STUDENT_STEPS}"
)
fi
if [[ -n "$FORCED_EXIT_STEP" ]]; then
overrides+=(
--config_override "predictor_v4_dmd.forced_exit_step=${FORCED_EXIT_STEP}"
)
fi
echo "[$(date -u +%FT%TZ)] 启动 Predictor-v4 Wan14B DMD：$CONFIG_PATH"
echo "[$(date -u +%FT%TZ)] 输出目录：$OUTPUT_DIR"
TMPDIR="$TRAIN_TMPDIR" \
CUDA_VISIBLE_DEVICES="$GPU_IDS" \
"$PYTHON_BIN" -m torch.distributed.run \
--standalone \
--nproc_per_node="$NPROC_PER_NODE" \
--master_port="$MASTER_PORT" \
train.py \
--config_path "$CONFIG_PATH" \
--no_visualize \
--disable-wandb \
"${overrides[@]}" \
2>&1 | tee -a "$OUTPUT_DIR/launcher.log"
|