File size: 3,427 Bytes
59630ba | 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 | #!/bin/bash
# WorldMem diffusion multi-node training, launched MANUALLY on each node.
#
# Run this on every node, passing that node's rank as the only argument.
# Node 0 must be the machine whose IP matches MASTER_ADDR.
#
# Usage (run on each node):
# bash scripts/train_multinode.sh <node_rank> [exp_name]
#
# Example with 2 nodes:
# # on master node:
# bash scripts/train_multinode.sh 0 decmem
# # on worker node:
# bash scripts/train_multinode.sh 1 decmem
#
# Env overrides:
# EXP_NAME experiment name / config stem (default: decmem)
# NPROC_PER_NODE GPUs per node (default 8)
# MASTER_PORT torchrun rendezvous port (default 29500)
# RESUME_CKPT path to a checkpoint to fully resume from (optional)
# ─── Topology (EDIT THESE) ───────────────────────────────────────────────────
NNODES=2 # TODO: total number of nodes
MASTER_ADDR="" # TODO: IP of node_rank=0
MASTER_PORT=${MASTER_PORT:-29500}
NPROC_PER_NODE=${NPROC_PER_NODE:-8}
# ─── Node rank / experiment from CLI ─────────────────────────────────────────
NODE_RANK=$1
EXP_NAME=${2:-${EXP_NAME:-decmem}}
if [ -z "$NODE_RANK" ]; then
echo "Error: NODE_RANK is required. Usage: bash $0 <node_rank> [exp_name]"
exit 1
fi
# ─── Paths / config ──────────────────────────────────────────────────────────
CONFIG=configs/${EXP_NAME}.yaml
LOGDIR=exps/${EXP_NAME}
WANDB_SAVE_DIR=${LOGDIR}/wandb
TORCHRUN=torchrun
RESUME_CKPT=${RESUME_CKPT:-}
mkdir -p "${LOGDIR}"
# Optional resume: `RESUME_CKPT=/path/to/model.pt bash <this_script> <node_rank>`
# Must be set on EVERY node (this script doesn't propagate env across nodes).
RESUME_ARGS=""
if [ -n "${RESUME_CKPT}" ]; then
RESUME_ARGS="--resume-ckpt ${RESUME_CKPT}"
echo "[resume] RESUME_CKPT = ${RESUME_CKPT}"
fi
# ─── NCCL (match train_worldmem_multinode_torchrun.sh) ───────────────────────
export NCCL_IB_DISABLE=0
export NCCL_IB_GID_INDEX=3
export NCCL_MIN_NCHANNELS=16
export NCCL_IB_HCA=mlx5
export NCCL_IB_QPS_PER_CONNECTION=4
export NCCL_IB_TIMEOUT=22
export NCCL_DEBUG=WARN
echo "EXP_NAME = ${EXP_NAME}"
echo "CONFIG = ${CONFIG}"
echo "LOGDIR = ${LOGDIR}"
echo "NNODES = ${NNODES}"
echo "NPROC_PER_NODE = ${NPROC_PER_NODE}"
echo "WORLD_SIZE = $((NNODES * NPROC_PER_NODE))"
echo "MASTER_ADDR = ${MASTER_ADDR}"
echo "MASTER_PORT = ${MASTER_PORT}"
echo "RESUME_CKPT = ${RESUME_CKPT:-<none>}"
echo "Starting node with NODE_RANK=${NODE_RANK}"
# ─── Launch torchrun ─────────────────────────────────────────────────────────
${TORCHRUN} \
--nnodes=${NNODES} \
--nproc_per_node=${NPROC_PER_NODE} \
--node_rank=${NODE_RANK} \
--master_addr=${MASTER_ADDR} \
--master_port=${MASTER_PORT} \
train.py \
--config_path ${CONFIG} \
--logdir ${LOGDIR} \
--wandb-save-dir ${WANDB_SAVE_DIR} \
${RESUME_ARGS} \
> "${LOGDIR}/node_${NODE_RANK}_$(date +%Y%m%d_%H%M%S).log" 2>&1
|