1+ #! /bin/bash
2+ # SBATCH --job-name=t2v
3+ # SBATCH --partition=main
4+ # SBATCH --nodes=8
5+ # SBATCH --ntasks=8
6+ # SBATCH --ntasks-per-node=1
7+ # SBATCH --gres=gpu:8
8+ # SBATCH --cpus-per-task=128
9+ # SBATCH --mem=1440G
10+ # SBATCH --output=dmd_t2v_output/t2v_%j.out
11+ # SBATCH --error=dmd_t2v_output/t2v_%j.err
12+ # SBATCH --exclusive
13+ set -e -x
14+
15+ # Environment Setup
16+ source ~ /conda/miniconda/bin/activate
17+ conda activate your_env
18+
19+ # Basic Info
20+ export WANDB_MODE=" online"
21+ export NCCL_P2P_DISABLE=1
22+ export TORCH_NCCL_ENABLE_MONITORING=0
23+ # different cache dir for different processes
24+ export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
25+ export MASTER_PORT=29500
26+ export NODE_RANK=$SLURM_PROCID
27+ nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST )
28+ export MASTER_ADDR=${nodes[0]}
29+ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
30+ export TOKENIZERS_PARALLELISM=false
31+ export WANDB_BASE_URL=" https://api.wandb.ai"
32+ export WANDB_MODE=online
33+ export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
34+ # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
35+
36+ echo " MASTER_ADDR: $MASTER_ADDR "
37+ echo " NODE_RANK: $NODE_RANK "
38+
39+ # Configs
40+ NUM_GPUS=8
41+ MODEL_PATH=" Wan-AI/Wan2.1-T2V-14B-Diffusers"
42+ DATA_DIR=your_data_dir
43+ VALIDATION_DATASET_FILE=your_validation_dataset_file
44+ # export CUDA_VISIBLE_DEVICES=4,5
45+ # IP=[MASTER NODE IP]
46+
47+ # Training arguments
48+ training_args=(
49+ --tracker_project_name wan_t2v_distill_dmd_VSA
50+ --output_dir=" checkpoints/wan_t2v_finetune"
51+ --max_train_steps=4000
52+ --train_batch_size=1
53+ --train_sp_batch_size 1
54+ --gradient_accumulation_steps=1
55+ --num_latent_t 16
56+ --num_height 448
57+ --num_width 832
58+ --num_frames 61
59+ --enable_gradient_checkpointing_type " full"
60+ )
61+
62+ # Parallel arguments
63+ parallel_args=(
64+ --num_gpus 64
65+ --sp_size 4
66+ --tp_size 1
67+ --hsdp_replicate_dim 8
68+ --hsdp_shard_dim 8
69+ )
70+
71+ # Model arguments
72+ model_args=(
73+ --model_path $MODEL_PATH
74+ --pretrained_model_name_or_path $MODEL_PATH
75+ )
76+
77+ # Dataset arguments
78+ dataset_args=(
79+ --data_path " $DATA_DIR "
80+ --dataloader_num_workers 4
81+ )
82+
83+ # Validation arguments
84+ validation_args=(
85+ --log_validation
86+ --validation_dataset_file " $VALIDATION_DATASET_FILE "
87+ --validation_steps 200
88+ --validation_sampling_steps " 3"
89+ --validation_guidance_scale " 1.0" # not used for dmd inference
90+ )
91+
92+ # Optimizer arguments
93+ optimizer_args=(
94+ --learning_rate=1e-5
95+ --mixed_precision=" bf16"
96+ --checkpointing_steps=500
97+ --weight_decay 0.01
98+ --max_grad_norm 1.0
99+ )
100+
101+ # Miscellaneous arguments
102+ miscellaneous_args=(
103+ --inference_mode False
104+ --allow_tf32
105+ --checkpoints_total_limit 3
106+ --training_cfg_rate 0.0
107+ --dit_precision " bf16"
108+ --ema_start_step 0
109+ --flow_shift 3
110+ --seed 1000
111+ )
112+
113+ # DMD arguments
114+ dmd_args=(
115+ --dmd_denoising_steps ' 1000,757,522'
116+ --min_timestep_ratio 0.02
117+ --max_timestep_ratio 0.98
118+ --generator_update_interval 5
119+ --real_score_guidance_scale 3.5
120+ --VSA_sparsity 0.9
121+ )
122+
123+ srun torchrun \
124+ --nnodes $SLURM_JOB_NUM_NODES \
125+ --nproc_per_node $NUM_GPUS \
126+ --node_rank $SLURM_PROCID \
127+ --rdzv_backend=c10d \
128+ --rdzv_endpoint=" $MASTER_ADDR :$MASTER_PORT "
129+ fastvideo/training/wan_training_pipeline.py \
130+ " ${parallel_args[@]} "
131+ " ${model_args[@]} "
132+ " ${dataset_args[@]} "
133+ " ${training_args[@]} "
134+ " ${optimizer_args[@]} "
135+ " ${validation_args[@]} "
136+ " ${miscellaneous_args[@]} "
137+ " ${dmd_args[@]} "
0 commit comments