78 lines
2.8 KiB
Bash
Executable File
78 lines
2.8 KiB
Bash
Executable File
#!/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"
|