Skip to content

Commit 6284eaa

Browse files
BrianChen1129BrianChen1129
authored
[Feature][Distill]Add DMD+VSA joint training example (hao-ai-lab#654)
1 parent 636524e commit 6284eaa

5 files changed

Lines changed: 811 additions & 0 deletions

File tree

‎examples/distill/Wan-Syn-480P/README.md‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# Wan2.1-T2V-1.3B Distill Example
2+
These are end-to-end example scripts for distilling Wan2.1 T2V 1.3B model using DMD-only and DMD+VSA methods.
3+
4+
### 1. Download dataset:
5+
```bash
6+
bash examples/distill/Wan-Syn-480P/download_dataset.sh
7+
```
8+
9+
### 2. Configure and run distillation:
10+
11+
#### For DMD-only distillation:
12+
```bash
13+
sbatch examples/distill/Wan-Syn-480P/distill_dmd_t2v.slurm
14+
```
15+
16+
#### For DMD+VSA distillation:
17+
```bash
18+
sbatch examples/distill/Wan-Syn-480P/distill_dmd_VSA_t2v.slurm
19+
```

‎examples/distill/Wan-Syn-480P/distill_dmd_VSA_t2v.slurm‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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-1.3B-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 1
66+
--tp_size 1
67+
--hsdp_replicate_dim 64
68+
--hsdp_shard_dim 1
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 8
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.8
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[@]}"

‎examples/distill/Wan-Syn-480P/distill_dmd_t2v.slurm‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
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=FLASH_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-1.3B-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
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 1
66+
--tp_size 1
67+
--hsdp_replicate_dim 64
68+
--hsdp_shard_dim 1
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 8
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+
)
121+
122+
srun torchrun \
123+
--nnodes $SLURM_JOB_NUM_NODES \
124+
--nproc_per_node $NUM_GPUS \
125+
--node_rank $SLURM_PROCID \
126+
--rdzv_backend=c10d \
127+
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
128+
fastvideo/training/wan_training_pipeline.py \
129+
"${parallel_args[@]}" \
130+
"${model_args[@]}" \
131+
"${dataset_args[@]}" \
132+
"${training_args[@]}" \
133+
"${optimizer_args[@]}" \
134+
"${validation_args[@]}" \
135+
"${miscellaneous_args[@]}" \
136+
"${dmd_args[@]}"

‎examples/distill/Wan-Syn-480P/download_dataset.sh‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
#!/bin/bash
2+
3+
python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"

0 commit comments

Comments
 (0)