File size: 815 Bytes
e0db531 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | #!/bin/bash
# VIPER training launcher -- one stage per invocation.
# scripts/train.sh <stage> [extra args...]
#
# Uses torchrun over the 8 GPUs of one node. Ranks alternate between the
# low-noise and high-noise Wan2.2 experts (see viper/train.py) so both halves of
# the timestep range are covered without holding 56GB of DiT on one card.
set -e
cd "$(dirname "$0")/.."
source env.sh
STAGE=$1; shift
NGPU=${NGPU:-8}
case $STAGE in
1) EXTRA="--stage 1 --steps ${STEPS:-3000}" ;;
2) EXTRA="--stage 2 --steps ${STEPS:-6000} --init_from checkpoints/stage1/final.pt" ;;
3) EXTRA="--stage 3 --steps ${STEPS:-6000} --init_from checkpoints/stage2/final.pt" ;;
*) echo "usage: $0 {1|2|3}"; exit 1 ;;
esac
exec torchrun --nproc_per_node=$NGPU --master_port=${MASTER_PORT:-29571} \
-m viper.train $EXTRA "$@"
|