File size: 1,679 Bytes
93cd35a | 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 | #!/usr/bin/env bash
# Launcher for the PyTorch ELF port.
#
# Single GPU / CPU:
# bash scripts/launch.sh train src/configs/training_configs/train_owt_ELF-B.yml
# bash scripts/launch.sh eval src/configs/training_configs/train_owt_ELF-B.yml embedded-language-flows/ELF-B-owt
#
# Multi-GPU (single-host):
# NGPU=8 bash scripts/launch.sh train src/configs/training_configs/train_owt_ELF-B.yml
#
# Multi-host (torchrun rendezvous):
# NGPU=8 NNODES=2 NODE_RANK=0 MASTER_ADDR=node-0 MASTER_PORT=29500 \
# bash scripts/launch.sh train src/configs/training_configs/train_owt_ELF-B.yml
set -euo pipefail
if [[ $# -lt 2 ]]; then
echo "usage: bash scripts/launch.sh <train|eval> <config.yml> [extra args...]"
exit 1
fi
MODE=$1
CONFIG=$2
shift 2
EXTRA="$@"
# Pick up the script we want to run.
case "$MODE" in
train) ENTRY=src/train.py ;;
eval) ENTRY=src/eval.py ;;
*) echo "Unknown mode: $MODE (expected 'train' or 'eval')"; exit 1 ;;
esac
NGPU=${NGPU:-1}
NNODES=${NNODES:-1}
NODE_RANK=${NODE_RANK:-0}
MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
MASTER_PORT=${MASTER_PORT:-29500}
export PYTHONPATH="$(pwd)/src:${PYTHONPATH:-}"
if [[ "$NGPU" == "1" && "$NNODES" == "1" ]]; then
echo "[launch] single-process: python $ENTRY --config $CONFIG $EXTRA"
exec python "$ENTRY" --config "$CONFIG" $EXTRA
else
echo "[launch] torchrun nproc_per_node=$NGPU nnodes=$NNODES node_rank=$NODE_RANK $ENTRY"
exec torchrun \
--nproc_per_node="$NGPU" \
--nnodes="$NNODES" \
--node_rank="$NODE_RANK" \
--master_addr="$MASTER_ADDR" \
--master_port="$MASTER_PORT" \
"$ENTRY" --config "$CONFIG" $EXTRA
fi
|