diff --git a/run_pipeline.sh b/run_pipeline.sh new file mode 100755 index 0000000..a35bca6 --- /dev/null +++ b/run_pipeline.sh @@ -0,0 +1,77 @@ +#!/bin/bash +# MiniMind 一键训练流水线:pretrain -> full_sft -> eval(自动测试) +# 用于在新环境/新硬件上快速验证完整训练流程,各阶段参数默认值与对应训练脚本一致。 +# +# 用法: +# bash run_pipeline.sh +# PRETRAIN_EPOCHS=1 SFT_BATCH=32 bash run_pipeline.sh # 用环境变量覆盖参数 +# RUN_EVAL=0 bash run_pipeline.sh # 跳过最后的自动测试 +set -uo pipefail + +cd "$(dirname "$0")" +ROOT=$(pwd) +LOG_DIR="$ROOT/logs" +TS=$(date +%Y%m%d_%H%M%S) +PLOG="$LOG_DIR/pipeline_$TS.log" +mkdir -p "$LOG_DIR" + +PY=${PYTHON:-python3} + +# ---- 可调参数(默认值与各训练脚本一致)---- +PRETRAIN_EPOCHS=${PRETRAIN_EPOCHS:-2} +PRETRAIN_BATCH=${PRETRAIN_BATCH:-32} +PRETRAIN_ACCUM=${PRETRAIN_ACCUM:-8} +SFT_EPOCHS=${SFT_EPOCHS:-2} +SFT_BATCH=${SFT_BATCH:-16} +SFT_ACCUM=${SFT_ACCUM:-1} +NUM_WORKERS=${NUM_WORKERS:-8} +LOG_INTERVAL=${LOG_INTERVAL:-100} +RESUME=${RESUME:-0} # 1=检测到上次中断的检查点时自动续训 +RUN_EVAL=${RUN_EVAL:-1} +# ------------------------------------------ + +# stdout 经管道/tee 时是块缓冲,必须关闭缓冲才能实时看到训练日志 +export PYTHONUNBUFFERED=1 + +# 自动检测训练设备:cuda > mps > cpu(也可用 DEVICE=xxx 覆盖) +DEVICE=${DEVICE:-$("$PY" -c "import torch; print('cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu'))" 2>/dev/null || echo cpu)} +export PYTORCH_ENABLE_MPS_FALLBACK=1 + +echo "[pipeline] device=$DEVICE start $(date '+%F %T')" | tee "$PLOG" + +run_stage() { + local name="$1"; shift + echo "[pipeline] === stage: $name START $(date '+%F %T') ===" | tee -a "$PLOG" + local t0 + t0=$(date +%s) + "$@" 2>&1 | tee "$LOG_DIR/${name}_${TS}.log" + local rc=${PIPESTATUS[0]} + local t1 + t1=$(date +%s) + echo "[pipeline] === stage: $name END rc=$rc elapsed=$(( (t1 - t0) / 60 ))min $(date '+%F %T') ===" | tee -a "$PLOG" + return $rc +} + +cd "$ROOT/trainer" + +run_stage pretrain "$PY" train_pretrain.py \ + --device "$DEVICE" --epochs "$PRETRAIN_EPOCHS" \ + --batch_size "$PRETRAIN_BATCH" --accumulation_steps "$PRETRAIN_ACCUM" \ + --num_workers "$NUM_WORKERS" --log_interval "$LOG_INTERVAL" \ + --from_resume "$RESUME" || exit 1 + +run_stage full_sft "$PY" train_full_sft.py \ + --device "$DEVICE" --epochs "$SFT_EPOCHS" \ + --batch_size "$SFT_BATCH" --accumulation_steps "$SFT_ACCUM" \ + --num_workers "$NUM_WORKERS" --log_interval "$LOG_INTERVAL" \ + --from_resume "$RESUME" || exit 1 + +if [ "$RUN_EVAL" = "1" ]; then + cd "$ROOT" + echo "[pipeline] === stage: eval START $(date '+%F %T') ===" | tee -a "$PLOG" + echo 0 | "$PY" eval_llm.py --weight full_sft --device "$DEVICE" \ + 2>&1 | tee "$LOG_DIR/eval_${TS}.log" + echo "[pipeline] === stage: eval END $(date '+%F %T') ===" | tee -a "$PLOG" +fi + +echo "[pipeline] === ALL DONE $(date '+%F %T') ===" | tee -a "$PLOG"