File size: 3,216 Bytes
589ada3 | 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
NPROC_PER_NODE="${1:?Usage: bash scripts/train.sh <nproc_per_node> [hydra_overrides...]}"
shift
EXTRA_ARGS=("$@")
NUM_MACHINES="${NNODES:-1}"
MACHINE_RANK="${NODE_RANK:-0}"
MAIN_PROCESS_IP="${MASTER_ADDR:-127.0.0.1}"
MAIN_PROCESS_PORT="${MASTER_PORT:-29500}"
is_integer() {
[[ "${1}" =~ ^[0-9]+$ ]]
}
if ! is_integer "${NUM_MACHINES}" || ! is_integer "${MACHINE_RANK}"; then
echo "Error: NUM_MACHINES (${NUM_MACHINES}) and MACHINE_RANK (${MACHINE_RANK}) must be integers." >&2
exit 1
fi
extract_task_basename() {
local cfg="$1"
if [[ "${cfg}" == task/* ]]; then
local name="${cfg#task/}"
name="${name%.yaml}"
echo "${name}"
return 0
fi
return 1
}
TASK_BASENAME="train"
for ((i = 0; i < ${#EXTRA_ARGS[@]}; i++)); do
arg="${EXTRA_ARGS[$i]}"
case "${arg}" in
--config-name)
if ((i + 1 < ${#EXTRA_ARGS[@]})); then
next="${EXTRA_ARGS[$((i + 1))]}"
if parsed="$(extract_task_basename "${next}")"; then
TASK_BASENAME="${parsed}"
fi
fi
;;
--config-name=*)
cfg="${arg#--config-name=}"
if parsed="$(extract_task_basename "${cfg}")"; then
TASK_BASENAME="${parsed}"
fi
;;
task=*)
cfg="${arg#task=}"
cfg="${cfg%.yaml}"
TASK_BASENAME="${cfg}"
;;
esac
done
if [[ -z "${RUN_ID:-}" ]]; then
if (( NUM_MACHINES <= 1 )); then
RUN_ID="$(date +%Y-%m-%d_%H-%M-%S)"
else
RUN_ID_SYNC_TIMEOUT="${RUN_ID_SYNC_TIMEOUT:-180}"
RUN_ID_SYNC_PORT="${RUN_ID_SYNC_PORT:-$((MAIN_PROCESS_PORT + 11))}"
export RUN_ID_SYNC_HOST="${MAIN_PROCESS_IP}"
export RUN_ID_SYNC_PORT
export RUN_ID_SYNC_TIMEOUT
export RUN_ID_SYNC_MACHINE_RANK="${MACHINE_RANK}"
export RUN_ID_SYNC_NUM_MACHINES="${NUM_MACHINES}"
export RUN_ID_SYNC_TASK_BASENAME="${TASK_BASENAME}"
RUN_ID="$(
python - <<'PY'
import datetime
import os
from datetime import timedelta
import torch.distributed as dist
host = os.environ["RUN_ID_SYNC_HOST"]
port = int(os.environ["RUN_ID_SYNC_PORT"])
timeout_s = int(os.environ["RUN_ID_SYNC_TIMEOUT"])
machine_rank = int(os.environ["RUN_ID_SYNC_MACHINE_RANK"])
num_machines = int(os.environ["RUN_ID_SYNC_NUM_MACHINES"])
task_basename = os.environ.get("RUN_ID_SYNC_TASK_BASENAME", "train")
store = dist.TCPStore(
host_name=host,
port=port,
world_size=num_machines,
is_master=(machine_rank == 0),
timeout=timedelta(seconds=timeout_s),
)
key = f"run_id::{task_basename}"
if machine_rank == 0:
run_id = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
store.set(key, run_id)
run_id = store.get(key).decode("utf-8")
print(run_id)
PY
)"
echo "[run_id_sync] mode=tcpstore host=${RUN_ID_SYNC_HOST} port=${RUN_ID_SYNC_PORT} timeout_s=${RUN_ID_SYNC_TIMEOUT} run_id=${RUN_ID}"
fi
fi
echo "[launch] nproc_per_node=${NPROC_PER_NODE} num_machines=${NUM_MACHINES} machine_rank=${MACHINE_RANK} run_id=${RUN_ID}"
accelerate launch \
--config_file scripts/accelerate_configs/accelerate_zero2_ds.yaml \
--num_processes "${NPROC_PER_NODE}" \
scripts/train.py \
"output_dir=./runs/${TASK_BASENAME}/${RUN_ID}" \
"wandb.name=${TASK_BASENAME}" \
"${EXTRA_ARGS[@]}"
|