diff --git a/examples/distill/Wan-Syn-480P/distill_dmd_VSA_t2v_14B_480P.slurm b/examples/distill/Wan-Syn-480P/distill_dmd_VSA_t2v_14B_480P.slurm new file mode 100644 index 000000000..eb128348b --- /dev/null +++ b/examples/distill/Wan-Syn-480P/distill_dmd_VSA_t2v_14B_480P.slurm @@ -0,0 +1,137 @@ +#!/bin/bash +#SBATCH --job-name=t2v +#SBATCH --partition=main +#SBATCH --nodes=8 +#SBATCH --ntasks=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=dmd_t2v_output/t2v_%j.out +#SBATCH --error=dmd_t2v_output/t2v_%j.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate your_env + +# Basic Info +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + +# Configs +NUM_GPUS=8 +MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" +DATA_DIR=your_data_dir +VALIDATION_DATASET_FILE=your_validation_dataset_file +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name wan_t2v_distill_dmd_VSA + --output_dir="checkpoints/wan_t2v_finetune" + --max_train_steps=4000 + --train_batch_size=1 + --train_sp_batch_size 1 + --gradient_accumulation_steps=1 + --num_latent_t 16 + --num_height 448 + --num_width 832 + --num_frames 61 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus 64 + --sp_size 4 + --tp_size 1 + --hsdp_replicate_dim 8 + --hsdp_shard_dim 8 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 200 + --validation_sampling_steps "3" + --validation_guidance_scale "1.0" # not used for dmd inference +) + +# Optimizer arguments +optimizer_args=( + --learning_rate=1e-5 + --mixed_precision="bf16" + --checkpointing_steps=500 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --allow_tf32 + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "bf16" + --ema_start_step 0 + --flow_shift 3 + --seed 1000 +) + +# DMD arguments +dmd_args=( + --dmd_denoising_steps '1000,757,522' + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --generator_update_interval 5 + --real_score_guidance_scale 3.5 + --VSA_sparsity 0.9 +) + +srun 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/training/wan_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ No newline at end of file diff --git a/examples/distillation/dmd/Wan2.1-Fun-1.3B-InP/crush_smol/README.md b/examples/distillation/dmd/Wan2.1-Fun-1.3B-InP/crush_smol/README.md deleted file mode 100644 index c2bfd39da..000000000 --- a/examples/distillation/dmd/Wan2.1-Fun-1.3B-InP/crush_smol/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# DMD Distillation Wan2.1-I2V-14B-480P Crush-Smol Example - -Coming soon! \ No newline at end of file