[update] 新增一键训练 pipeline 脚本 run_pipeline.sh

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"