1df513922f
Co-authored-by: alexzms <3036648523@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
72 lines
2.7 KiB
Bash
Executable File
72 lines
2.7 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Launch training from a YAML config.
|
|
#
|
|
# Usage:
|
|
# bash examples/train/run.sh <config.yaml> [--dotted.key value ...]
|
|
#
|
|
# Examples:
|
|
# bash examples/train/run.sh examples/train/configs/fine_tuning/wan/t2v.yaml
|
|
# bash examples/train/run.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml --dry-run
|
|
# bash examples/train/run.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \
|
|
# --training.distributed.num_gpus 4 \
|
|
# --training.optimizer.learning_rate 1e-5
|
|
# bash examples/train/run.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \
|
|
# --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
|
|
#
|
|
# Logs are written to logs/<config_name>_<timestamp>.log (and also printed to stdout).
|
|
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
|
|
|
|
CONFIG="${1:?Usage: $0 <config.yaml> [extra flags...]}"
|
|
shift
|
|
|
|
if [[ "${CONFIG}" != /* ]]; then
|
|
CONFIG="$(pwd)/${CONFIG}"
|
|
fi
|
|
|
|
# ── GPU / node settings ──────────────────────────────────────────
|
|
NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | wc -l)}"
|
|
NUM_GPUS="${NUM_GPUS:-1}"
|
|
NNODES="${NNODES:-1}"
|
|
NODE_RANK="${NODE_RANK:-0}"
|
|
MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
|
|
MASTER_PORT="${MASTER_PORT:-29501}"
|
|
export TOKENIZERS_PARALLELISM=false
|
|
export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
|
|
|
|
cd "${REPO_ROOT}"
|
|
# ── W&B ──────────────────────────────────────────────────────────
|
|
export WANDB_API_KEY="${WANDB_API_KEY:-}"
|
|
export WANDB_MODE="${WANDB_MODE:-online}"
|
|
|
|
# ── Log file ─────────────────────────────────────────────────────
|
|
CONFIG_NAME="$(basename "${CONFIG}" .yaml)"
|
|
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
|
LOG_DIR="${LOG_DIR:-examples/train}"
|
|
mkdir -p "${LOG_DIR}"
|
|
LOG_FILE="${LOG_DIR}/${CONFIG_NAME}_${TIMESTAMP}.log"
|
|
|
|
echo "=== Train Training ==="
|
|
echo "Config: ${CONFIG}"
|
|
echo "Num GPUs: ${NUM_GPUS}"
|
|
echo "Num Nodes: ${NNODES}"
|
|
echo "Node Rank: ${NODE_RANK}"
|
|
echo "Master: ${MASTER_ADDR}:${MASTER_PORT}"
|
|
echo "Extra args: $*"
|
|
echo "Log file: ${LOG_FILE}"
|
|
echo "=============================="
|
|
|
|
python -m torch.distributed.run \
|
|
--nnodes "${NNODES}" \
|
|
--node_rank "${NODE_RANK}" \
|
|
--nproc_per_node "${NUM_GPUS}" \
|
|
--master_addr "${MASTER_ADDR}" \
|
|
--master_port "${MASTER_PORT}" \
|
|
-m fastvideo.train.entrypoint.train \
|
|
--config "${CONFIG}" \
|
|
"$@" \
|
|
2>&1 | tee "${LOG_FILE}"
|