115 lines
4.5 KiB
Bash
Executable File
115 lines
4.5 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Submit a multi-node Slurm training job.
|
|
#
|
|
# Usage:
|
|
# bash examples/train/run_slurm.sh <config.yaml> <num_nodes> [--dotted.key value ...]
|
|
#
|
|
# Examples:
|
|
# bash examples/train/run_slurm.sh examples/train/configs/example.yaml 2
|
|
# bash examples/train/run_slurm.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml 4 \
|
|
# --training.optimizer.learning_rate 1e-5
|
|
# bash examples/train/run_slurm.sh examples/train/configs/example.yaml 8 \
|
|
# --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
|
|
#
|
|
# Environment variables (override defaults):
|
|
# PARTITION Slurm partition (default: main)
|
|
# NUM_GPUS GPUs per node (default: 8)
|
|
# CPUS_PER_TASK CPUs per task (default: 128)
|
|
# MEM Memory per node (default: 1440G)
|
|
# JOB_NAME Slurm job name (default: derived from config)
|
|
# OUTPUT_DIR Directory for slurm logs (default: slurm_logs)
|
|
# MASTER_PORT Rendezvous port (default: 29500)
|
|
# EXCLUDE Nodes to exclude (default: "")
|
|
# WANDB_API_KEY W&B API key (default: "")
|
|
# WANDB_MODE W&B mode (default: online)
|
|
|
|
set -euo pipefail
|
|
|
|
CONFIG="${1:?Usage: $0 <config.yaml> <num_nodes> [extra flags...]}"
|
|
NUM_NODES="${2:?Usage: $0 <config.yaml> <num_nodes> [extra flags...]}"
|
|
shift 2
|
|
|
|
# ── Defaults ──────────────────────────────────────────────────────
|
|
PARTITION="${PARTITION:-main}"
|
|
NUM_GPUS="${NUM_GPUS:-8}"
|
|
CPUS_PER_TASK="${CPUS_PER_TASK:-128}"
|
|
MEM="${MEM:-1440G}"
|
|
MASTER_PORT="${MASTER_PORT:-29500}"
|
|
EXCLUDE="${EXCLUDE:-}"
|
|
WANDB_API_KEY="${WANDB_API_KEY:-}"
|
|
WANDB_MODE="${WANDB_MODE:-online}"
|
|
|
|
TOTAL_GPUS=$(( NUM_NODES * NUM_GPUS ))
|
|
|
|
CONFIG_NAME="$(basename "${CONFIG}" .yaml)"
|
|
JOB_NAME="${JOB_NAME:-${CONFIG_NAME}}"
|
|
OUTPUT_DIR="${OUTPUT_DIR:-logs/slurm}"
|
|
mkdir -p "${OUTPUT_DIR}"
|
|
|
|
# ── Build sbatch args ─────────────────────────────────────────────
|
|
SBATCH_ARGS=(
|
|
--job-name="${JOB_NAME}"
|
|
--partition="${PARTITION}"
|
|
--nodes="${NUM_NODES}"
|
|
--ntasks="${NUM_NODES}"
|
|
--ntasks-per-node=1
|
|
--gres="gpu:${NUM_GPUS}"
|
|
--cpus-per-task="${CPUS_PER_TASK}"
|
|
--mem="${MEM}"
|
|
--output="${OUTPUT_DIR}/${JOB_NAME}_%j.out"
|
|
--error="${OUTPUT_DIR}/${JOB_NAME}_%j.err"
|
|
--exclusive
|
|
)
|
|
|
|
if [[ -n "${EXCLUDE}" ]]; then
|
|
SBATCH_ARGS+=(--exclude="${EXCLUDE}")
|
|
fi
|
|
|
|
# ── Collect extra overrides for the training script ───────────────
|
|
EXTRA_ARGS=("$@")
|
|
|
|
echo "=== Slurm Training Submission ==="
|
|
echo "Config: ${CONFIG}"
|
|
echo "Nodes: ${NUM_NODES}"
|
|
echo "GPUs/node: ${NUM_GPUS}"
|
|
echo "Total GPUs: ${TOTAL_GPUS}"
|
|
echo "Partition: ${PARTITION}"
|
|
echo "Job name: ${JOB_NAME}"
|
|
echo "Extra args: ${EXTRA_ARGS[*]:-}"
|
|
echo "================================="
|
|
|
|
# ── Submit ────────────────────────────────────────────────────────
|
|
sbatch "${SBATCH_ARGS[@]}" <<EOF
|
|
#!/bin/bash
|
|
set -e -x
|
|
|
|
# ── Environment ───────────────────────────────────────────────────
|
|
export NCCL_P2P_DISABLE=1
|
|
export TORCH_NCCL_ENABLE_MONITORING=0
|
|
export TRITON_CACHE_DIR=/tmp/triton_cache_\${SLURM_PROCID}
|
|
export TOKENIZERS_PARALLELISM=false
|
|
export WANDB_API_KEY="${WANDB_API_KEY}"
|
|
export WANDB_MODE="${WANDB_MODE}"
|
|
|
|
# ── Rendezvous ────────────────────────────────────────────────────
|
|
export MASTER_PORT=${MASTER_PORT}
|
|
nodes=( \$(scontrol show hostnames \$SLURM_JOB_NODELIST) )
|
|
export MASTER_ADDR=\${nodes[0]}
|
|
export NODE_RANK=\$SLURM_PROCID
|
|
|
|
echo "MASTER_ADDR: \$MASTER_ADDR"
|
|
echo "NODE_RANK: \$NODE_RANK"
|
|
|
|
# ── Launch ────────────────────────────────────────────────────────
|
|
srun torchrun \\
|
|
--nnodes \$SLURM_JOB_NUM_NODES \\
|
|
--nproc_per_node ${NUM_GPUS} \\
|
|
--node_rank \$SLURM_PROCID \\
|
|
--rdzv_backend=c10d \\
|
|
--rdzv_endpoint="\$MASTER_ADDR:\$MASTER_PORT" \\
|
|
fastvideo/train/entrypoint/train.py \\
|
|
--config ${CONFIG} \\
|
|
--training.distributed.num_gpus ${TOTAL_GPUS} \\
|
|
${EXTRA_ARGS[*]:-}
|
|
EOF
|