135 lines
5.7 KiB
Bash
Executable File
135 lines
5.7 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 Exported W&B API key required for online logging
|
|
# WANDB_MODE W&B mode (default: online)
|
|
# SBATCH_CONSTRAINT Slurm node constraint (default: "")
|
|
# SBATCH_TIMELIMIT Slurm time limit (default: "")
|
|
|
|
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_MODE="${WANDB_MODE:-online}"
|
|
SBATCH_CONSTRAINT="${SBATCH_CONSTRAINT:-}"
|
|
SBATCH_TIMELIMIT="${SBATCH_TIMELIMIT:-}"
|
|
|
|
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 ─────────────────────────────────────────────
|
|
# Forward the exported W&B API key through Slurm's environment so the generated
|
|
# job script and its xtrace output do not contain the credential value.
|
|
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"
|
|
--export="ALL,WANDB_MODE=${WANDB_MODE}"
|
|
--exclusive
|
|
)
|
|
|
|
if [[ -n "${EXCLUDE}" ]]; then
|
|
SBATCH_ARGS+=(--exclude="${EXCLUDE}")
|
|
fi
|
|
# Apply hardware and wall-clock requirements selected by an experiment launcher.
|
|
if [[ -n "${SBATCH_CONSTRAINT}" ]]; then
|
|
SBATCH_ARGS+=(--constraint="${SBATCH_CONSTRAINT}")
|
|
fi
|
|
if [[ -n "${SBATCH_TIMELIMIT}" ]]; then
|
|
SBATCH_ARGS+=(--time="${SBATCH_TIMELIMIT}")
|
|
fi
|
|
|
|
# ── Collect extra overrides for the training script ───────────────
|
|
EXTRA_ARGS=("$@")
|
|
# Shell-escape values before embedding them in the generated Bash script so
|
|
# config paths and override arguments retain their original boundaries.
|
|
printf -v CONFIG_SHELL_ARG '%q' "${CONFIG}"
|
|
EXTRA_ARGS_SHELL=""
|
|
if (( ${#EXTRA_ARGS[@]} > 0 )); then
|
|
printf -v EXTRA_ARGS_SHELL ' %q' "${EXTRA_ARGS[@]}"
|
|
fi
|
|
|
|
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 TOKENIZERS_PARALLELISM=false
|
|
|
|
# ── Rendezvous ────────────────────────────────────────────────────
|
|
export MASTER_PORT=${MASTER_PORT}
|
|
nodes=( \$(scontrol show hostnames \$SLURM_JOB_NODELIST) )
|
|
export MASTER_ADDR=\${nodes[0]}
|
|
# ── Launch ────────────────────────────────────────────────────────
|
|
# Resolve SLURM_PROCID inside each srun task so every node receives a distinct
|
|
# torchrun node rank and Triton cache directory.
|
|
launch_training_node() {
|
|
export TRITON_CACHE_DIR="/tmp/triton_cache_\${SLURM_PROCID}"
|
|
echo "MASTER_ADDR: \${MASTER_ADDR}"
|
|
echo "NODE_RANK: \${SLURM_PROCID}"
|
|
exec 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_SHELL_ARG} \\
|
|
--training.distributed.num_gpus ${TOTAL_GPUS}${EXTRA_ARGS_SHELL}
|
|
}
|
|
export -f launch_training_node
|
|
srun bash -c launch_training_node
|
|
EOF
|