164 lines
4.7 KiB
Bash
Executable File
164 lines
4.7 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
|
|
ARTIFACT_ROOT=${ARTIFACT_ROOT:-/apps/yi/model_training/artifacts}
|
|
RUN_STATE_DIR="${ARTIFACT_ROOT}/run_state"
|
|
LOG_DIR="${ARTIFACT_ROOT}/logs"
|
|
|
|
usage() {
|
|
cat <<'EOF'
|
|
Usage:
|
|
bash start_multinode_training.sh <model> [mode] [train_name]
|
|
|
|
Default cluster:
|
|
g0033,g0034,g0035,g0036 with 8 GPUs per host.
|
|
|
|
Environment overrides:
|
|
HOSTS="g0033 g0034 g0035 g0036"
|
|
MASTER_ADDR=g0033
|
|
MASTER_PORT=6000
|
|
NPROC_PER_NODE=8
|
|
ZERO_STAGE=0|1|2
|
|
CONTAINER_NAME=megatron-ngc25-training
|
|
NCCL_SOCKET_IFNAME=eth0
|
|
GLOO_SOCKET_IFNAME=eth0
|
|
NCCL_IB_HCA=mlx5_0,mlx5_1
|
|
NCCL_DEBUG=INFO
|
|
CHECKPOINT_KEEP_RECENT=3
|
|
CHECKPOINT_CLEANUP_INTERVAL_SECONDS=300
|
|
EXTRA_ARGS="--exit-duration-in-mins 120"
|
|
|
|
Examples:
|
|
bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu
|
|
CONTAINER_NAME=megatron-ngc25-training bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu
|
|
ZERO_STAGE=1 bash start_multinode_training.sh qwen3_1p7b phase1 qwen3_phase1_zero1
|
|
EOF
|
|
}
|
|
|
|
model=${1:-}
|
|
mode=${2:-}
|
|
train_name=${3:-}
|
|
|
|
if [ -z "$model" ] || [ "$model" = "-h" ] || [ "$model" = "--help" ]; then
|
|
usage
|
|
exit 0
|
|
fi
|
|
|
|
case "$model" in
|
|
gpt_smoke)
|
|
train_script="${SCRIPT_DIR}/training_smoke_gpt2.sh"
|
|
mode=${mode:-smoke}
|
|
train_name=${train_name:-smoke_gpt_multinode}
|
|
;;
|
|
qwen3_1p7b)
|
|
train_script="${SCRIPT_DIR}/training_smoke_qwen3_1p7b.sh"
|
|
mode=${mode:-qwen3_1p7b_smoke_yi}
|
|
train_name=${train_name:-qwen3_1p7b_multinode}
|
|
;;
|
|
*)
|
|
echo "Unknown model: $model" >&2
|
|
usage >&2
|
|
exit 1
|
|
;;
|
|
esac
|
|
|
|
read -r -a HOST_ARRAY <<< "${HOSTS:-g0033 g0034 g0035 g0036}"
|
|
NNODES=${NNODES:-${#HOST_ARRAY[@]}}
|
|
NPROC_PER_NODE=${NPROC_PER_NODE:-8}
|
|
MASTER_ADDR=${MASTER_ADDR:-${HOST_ARRAY[0]}}
|
|
MASTER_PORT=${MASTER_PORT:-6000}
|
|
ZERO_STAGE=${ZERO_STAGE:-0}
|
|
CONTAINER_NAME=${CONTAINER_NAME:-}
|
|
NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-}
|
|
GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-}
|
|
NCCL_IB_HCA=${NCCL_IB_HCA:-}
|
|
NCCL_DEBUG=${NCCL_DEBUG:-}
|
|
NCCL_IB_DISABLE=${NCCL_IB_DISABLE:-}
|
|
CHECKPOINT_KEEP_RECENT=${CHECKPOINT_KEEP_RECENT:-3}
|
|
CHECKPOINT_CLEANUP_INTERVAL_SECONDS=${CHECKPOINT_CLEANUP_INTERVAL_SECONDS:-300}
|
|
EXTRA_ARGS="--exit-signal-handler ${EXTRA_ARGS:-}"
|
|
|
|
mkdir -p "$RUN_STATE_DIR" "$LOG_DIR"
|
|
|
|
echo "Starting multinode training: model=${model}, mode=${mode}, train_name=${train_name}"
|
|
echo "Hosts: ${HOST_ARRAY[*]}"
|
|
echo "Distributed: nnodes=${NNODES}, nproc_per_node=${NPROC_PER_NODE}, master=${MASTER_ADDR}:${MASTER_PORT}, zero_stage=${ZERO_STAGE}"
|
|
if [ -n "$CONTAINER_NAME" ]; then
|
|
echo "Container: ${CONTAINER_NAME}"
|
|
fi
|
|
|
|
for idx in "${!HOST_ARRAY[@]}"; do
|
|
host=${HOST_ARRAY[$idx]}
|
|
node_rank=$idx
|
|
node_train_name="${train_name}_node${node_rank}"
|
|
pid_file="${RUN_STATE_DIR}/${node_train_name}.pid"
|
|
meta_file="${RUN_STATE_DIR}/${node_train_name}.env"
|
|
log_file="${LOG_DIR}/${node_train_name}.log"
|
|
|
|
remote_cmd=$(cat <<EOF
|
|
set -euo pipefail
|
|
run_cmd=\$(cat <<'RUN_CMD'
|
|
mkdir -p "$RUN_STATE_DIR" "$LOG_DIR"
|
|
cd "$SCRIPT_DIR"
|
|
ARTIFACT_ROOT="$ARTIFACT_ROOT" \\
|
|
CHECKPOINT_KEEP_RECENT="$CHECKPOINT_KEEP_RECENT" \\
|
|
CHECKPOINT_CLEANUP_INTERVAL_SECONDS="$CHECKPOINT_CLEANUP_INTERVAL_SECONDS" \\
|
|
NPROC_PER_NODE="$NPROC_PER_NODE" \\
|
|
NNODES="$NNODES" \\
|
|
NODE_RANK="$node_rank" \\
|
|
MASTER_ADDR="$MASTER_ADDR" \\
|
|
MASTER_PORT="$MASTER_PORT" \\
|
|
ZERO_STAGE="$ZERO_STAGE" \\
|
|
NCCL_SOCKET_IFNAME="$NCCL_SOCKET_IFNAME" \\
|
|
GLOO_SOCKET_IFNAME="$GLOO_SOCKET_IFNAME" \\
|
|
NCCL_IB_HCA="$NCCL_IB_HCA" \\
|
|
NCCL_DEBUG="$NCCL_DEBUG" \\
|
|
NCCL_IB_DISABLE="$NCCL_IB_DISABLE" \\
|
|
EXTRA_ARGS="$EXTRA_ARGS" \\
|
|
setsid bash "$train_script" "$mode" "$train_name" > "$log_file" 2>&1 < /dev/null &
|
|
pid=\$!
|
|
pgid=\$(ps -o pgid= -p "\$pid" | tr -d ' ' || true)
|
|
printf '%s\n' "\$pid" > "$pid_file"
|
|
cat > "$meta_file" <<META
|
|
MODEL=$model
|
|
MODE=$mode
|
|
TRAIN_NAME=$train_name
|
|
NODE_TRAIN_NAME=$node_train_name
|
|
HOST=$host
|
|
NODE_RANK=$node_rank
|
|
NNODES=$NNODES
|
|
NPROC_PER_NODE=$NPROC_PER_NODE
|
|
MASTER_ADDR=$MASTER_ADDR
|
|
MASTER_PORT=$MASTER_PORT
|
|
ZERO_STAGE=$ZERO_STAGE
|
|
NCCL_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME
|
|
GLOO_SOCKET_IFNAME=$GLOO_SOCKET_IFNAME
|
|
NCCL_IB_HCA=$NCCL_IB_HCA
|
|
NCCL_DEBUG=$NCCL_DEBUG
|
|
NCCL_IB_DISABLE=$NCCL_IB_DISABLE
|
|
PID=\$pid
|
|
PGID=\$pgid
|
|
LOG_FILE=$log_file
|
|
TRAIN_SCRIPT=$train_script
|
|
META
|
|
echo "host=$host node_rank=$node_rank pid=\$pid pgid=\${pgid:-unknown} log=$log_file"
|
|
RUN_CMD
|
|
)
|
|
if [ -n "$CONTAINER_NAME" ]; then
|
|
docker exec "$CONTAINER_NAME" bash -lc "\$run_cmd"
|
|
else
|
|
bash -lc "\$run_cmd"
|
|
fi
|
|
EOF
|
|
)
|
|
|
|
if [ "$host" = "$(hostname -s)" ] || [ "$host" = "$(hostname)" ]; then
|
|
bash -lc "$remote_cmd"
|
|
else
|
|
ssh "$host" "bash -lc $(printf '%q' "$remote_cmd")"
|
|
fi
|
|
done
|
|
|
|
echo "Launched ${NNODES} nodes. Stop all nodes with: bash ${SCRIPT_DIR}/stop_multinode_training.sh ${train_name}"
|