File size: 3,839 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | #!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"
PYTHON_BIN="${PYTHON_BIN:-/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/self_forcing/bin/python}"
GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}"
WAIT_FOR_GPUS="${WAIT_FOR_GPUS:-1}"
GPU_MEMORY_THRESHOLD_MIB="${GPU_MEMORY_THRESHOLD_MIB:-2048}"
POLL_SECONDS="${POLL_SECONDS:-30}"
RUN_SMOKE_TEST="${RUN_SMOKE_TEST:-1}"
START_TRAINING="${START_TRAINING:-0}"
DATASET_ROOT="${DATASET_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_1000_seed0}"
SMOKE_ROOT="${SMOKE_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_smoke_seed0}"
NVME_TRAIN_ROOT="${NVME_TRAIN_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_training_1000_seed0}"
RUN_DIR="${RUN_DIR:-$ROOT_DIR/checkpoints/predictor_v4_1000_seed0}"
mkdir -p "$RUN_DIR"
exec 9>"$RUN_DIR/end_to_end.lock"
if ! flock -n 9; then
echo "错误：已有 Predictor-v4 end-to-end supervisor 正在运行。" >&2
exit 1
fi
echo "$$" >"$RUN_DIR/end_to_end.pid"
trap 'rm -f "$RUN_DIR/end_to_end.pid"' EXIT
IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
if (( ${#GPU_ARRAY[@]} != 8 )); then
echo "错误：正式构建与训练需要 8 张 GPU，当前 GPU_IDS=$GPU_IDS。" >&2
exit 1
fi
if [[ "$WAIT_FOR_GPUS" != "0" && "$WAIT_FOR_GPUS" != "1" ]]; then
echo "错误：WAIT_FOR_GPUS 必须为 0 或 1。" >&2
exit 1
fi
if [[ "$START_TRAINING" != "0" && "$START_TRAINING" != "1" ]]; then
echo "错误：START_TRAINING 必须为 0 或 1。" >&2
exit 1
fi
gpus_are_free() {
local used
mapfile -t used < <(
nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits
)
if (( ${#used[@]} < 8 )); then
return 1
fi
local gpu_id
for gpu_id in "${GPU_ARRAY[@]}"; do
if (( used[gpu_id] > GPU_MEMORY_THRESHOLD_MIB )); then
return 1
fi
done
}
if [[ "$WAIT_FOR_GPUS" == "1" ]]; then
while ! gpus_are_free; do
echo "[$(date -u +%FT%TZ)] GPU 尚忙，${POLL_SECONDS}s 后重试。" >&2
sleep "$POLL_SECONDS"
done
elif ! gpus_are_free; then
echo "错误：GPU 显存占用超过 ${GPU_MEMORY_THRESHOLD_MIB} MiB。" >&2
exit 1
fi
if [[ "$RUN_SMOKE_TEST" == "1" ]]; then
echo "[$(date -u +%FT%TZ)] 开始单 case GPU 冒烟测试。"
"$PYTHON_BIN" scripts/prepare_predictor_v4_prompts.py \
--output_root "$SMOKE_ROOT" \
--num_prompts 1 \
--sampling_seed 0 \
--inference_seed 0
CUDA_VISIBLE_DEVICES="${GPU_ARRAY[0]}" \
"$PYTHON_BIN" scripts/build_predictor_v4_dataset.py \
--dataset_root "$SMOKE_ROOT" \
--worker_id 0 \
--num_workers 1 \
--case_ids 0 \
--seed 0 \
--blocks 0 1 28 29 \
--validate_kv_rebuild
"$PYTHON_BIN" scripts/merge_predictor_v4_manifests.py \
--dataset_root "$SMOKE_ROOT" \
--num_workers 1
"$PYTHON_BIN" scripts/validate_predictor_v4_dataset.py \
--dataset_root "$SMOKE_ROOT" \
--tensor_records 0 \
--require_kv_metrics 7
echo "[$(date -u +%FT%TZ)] 单 case GPU 冒烟测试通过。"
fi
echo "[$(date -u +%FT%TZ)] 开始 1000 prompts / 8 GPU 离线数据构建。"
GPU_IDS="$GPU_IDS" \
DATASET_ROOT="$DATASET_ROOT" \
PREPARE_PROMPTS=1 \
VALIDATE_KV_REBUILD=0 \
bash scripts/launch_predictor_v4_build.sh
if [[ "$START_TRAINING" == "0" ]]; then
echo "[$(date -u +%FT%TZ)] 数据构建完成；START_TRAINING=0，不启动训练。"
exit 0
fi
echo "[$(date -u +%FT%TZ)] 数据构建完成，开始 8 GPU Predictor-v4 训练。"
NPROC_PER_NODE=8 \
LOG_DIR="$NVME_TRAIN_ROOT" \
bash scripts/run_predictor_v4_training.sh
echo "[$(date -u +%FT%TZ)] 训练完成，开始把 NVMe 产物同步回项目目录。"
DATASET_SOURCE="$DATASET_ROOT" \
TRAIN_SOURCE="$NVME_TRAIN_ROOT" \
bash scripts/sync_predictor_v4_artifacts.sh
echo "[$(date -u +%FT%TZ)] Predictor-v4 数据构建与训练全部完成。"
|