File size: 4,122 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | #!/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}"
NUM_WORKERS="${NUM_WORKERS:-8}"
DATASET_ROOT="${DATASET_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_1000_seed0}"
CONFIG_PATH="${CONFIG_PATH:-$ROOT_DIR/configs/self_forcing_dmd.yaml}"
CHECKPOINT_PATH="${CHECKPOINT_PATH:-$ROOT_DIR/checkpoints/self_forcing_dmd.pt}"
PREPARE_PROMPTS="${PREPARE_PROMPTS:-1}"
VALIDATE_KV_REBUILD="${VALIDATE_KV_REBUILD:-0}"
DRY_RUN="${DRY_RUN:-0}"
IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
if (( ${#GPU_ARRAY[@]} != NUM_WORKERS )); then
echo "错误：GPU_IDS 数量 ${#GPU_ARRAY[@]} 与 NUM_WORKERS=$NUM_WORKERS 不一致。" >&2
exit 1
fi
if (( NUM_WORKERS != 8 )); then
echo "警告：正式数据构建约定为 8 workers；当前 NUM_WORKERS=$NUM_WORKERS。" >&2
fi
if [[ "$PREPARE_PROMPTS" != "0" && "$PREPARE_PROMPTS" != "1" ]]; then
echo "错误：PREPARE_PROMPTS 必须为 0 或 1。" >&2
exit 1
fi
if [[ "$VALIDATE_KV_REBUILD" != "0" && "$VALIDATE_KV_REBUILD" != "1" ]]; then
echo "错误：VALIDATE_KV_REBUILD 必须为 0 或 1。" >&2
exit 1
fi
for required_path in "$PYTHON_BIN" "$CONFIG_PATH" "$CHECKPOINT_PATH"; do
if [[ ! -e "$required_path" ]]; then
echo "错误：找不到 $required_path" >&2
exit 1
fi
done
PREPARE_COMMAND=(
"$PYTHON_BIN" scripts/prepare_predictor_v4_prompts.py
--output_root "$DATASET_ROOT"
--num_prompts 1000
--sampling_seed 0
--inference_seed 0
)
echo "Predictor v4 离线数据构建："
echo " Dataset: $DATASET_ROOT"
echo " GPUs: $GPU_IDS"
echo " Workers: $NUM_WORKERS"
echo " Teacher: $CHECKPOINT_PATH [generator_ema]"
echo " Schedule: F-F-F-F"
echo " Seed: 每个 case 重置为 0"
echo " Blocks: 0 1 28 29"
echo " KV check: $VALIDATE_KV_REBUILD"
if [[ "$DRY_RUN" == "1" ]]; then
printf ' Prepare: '
printf '%q ' "${PREPARE_COMMAND[@]}"
printf '\n'
for worker_id in $(seq 0 $((NUM_WORKERS - 1))); do
echo " Worker $worker_id: CUDA_VISIBLE_DEVICES=${GPU_ARRAY[$worker_id]} ... --worker_id $worker_id"
done
echo "DRY_RUN=1：未创建 prompt split，也未启动 GPU 构建。"
exit 0
fi
if [[ "$PREPARE_PROMPTS" == "1" && ! -f "$DATASET_ROOT/cases.jsonl" ]]; then
"${PREPARE_COMMAND[@]}"
fi
if [[ ! -f "$DATASET_ROOT/cases.jsonl" ]]; then
echo "错误：缺少 $DATASET_ROOT/cases.jsonl，请先运行 prompt 准备脚本。" >&2
exit 1
fi
mkdir -p "$DATASET_ROOT/logs"
pids=()
for worker_id in $(seq 0 $((NUM_WORKERS - 1))); do
WORKER_COMMAND=(
"$PYTHON_BIN" scripts/build_predictor_v4_dataset.py
--dataset_root "$DATASET_ROOT"
--config "$CONFIG_PATH"
--checkpoint "$CHECKPOINT_PATH"
--worker_id "$worker_id"
--num_workers "$NUM_WORKERS"
--seed 0
--blocks 0 1 28 29
)
if [[ "$VALIDATE_KV_REBUILD" == "1" ]]; then
WORKER_COMMAND+=(--validate_kv_rebuild)
fi
CUDA_VISIBLE_DEVICES="${GPU_ARRAY[$worker_id]}" \
"${WORKER_COMMAND[@]}" \
>"$DATASET_ROOT/logs/worker_${worker_id}.log" 2>&1 &
worker_pid="$!"
pids+=("$worker_pid")
echo "已启动 worker $worker_id，GPU=${GPU_ARRAY[$worker_id]}，PID=$worker_pid"
done
terminate_workers() {
for pid in "${pids[@]}"; do
kill "$pid" 2>/dev/null || true
done
}
trap terminate_workers INT TERM
failed=0
for index in "${!pids[@]}"; do
if ! wait "${pids[$index]}"; then
echo "错误：worker $index 失败，见 $DATASET_ROOT/logs/worker_${index}.log" >&2
failed=1
fi
done
trap - INT TERM
if (( failed != 0 )); then
exit 1
fi
"$PYTHON_BIN" scripts/merge_predictor_v4_manifests.py \
--dataset_root "$DATASET_ROOT" \
--num_workers "$NUM_WORKERS"
VALIDATE_ARGS=(
scripts/validate_predictor_v4_dataset.py
--dataset_root "$DATASET_ROOT"
--tensor_records 16
)
if [[ "$VALIDATE_KV_REBUILD" == "1" ]]; then
VALIDATE_ARGS+=(--require_kv_metrics 1)
fi
"$PYTHON_BIN" "${VALIDATE_ARGS[@]}"
echo "Predictor v4 离线数据构建和验证完成：$DATASET_ROOT"
|