| # 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 | |