From ba75e61e4e979ecd969834ce34bd999edcb5f3ed Mon Sep 17 00:00:00 2001 From: Yonghao Zhuang Date: Fri, 28 Nov 2025 07:23:10 +0000 Subject: [PATCH] Add nsys --- .../Wan-Syn-Data-480P/README.salloc.md | 34 +++ .../distill_dmd_VSA_t2v_1.3B.jd.salloc | 201 ++++++++++++++++++ .../distill_dmd_VSA_t2v_1.3B.jd.slurm | 199 +++++++++++++++++ 3 files changed, 434 insertions(+) create mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/README.salloc.md create mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.salloc create mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.slurm diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/README.salloc.md b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/README.salloc.md new file mode 100644 index 000000000..25cef3eda --- /dev/null +++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/README.salloc.md @@ -0,0 +1,34 @@ + + +### Standard way to use `salloc + srun` +```bash +salloc -N 8 -G 64 --cpus-per-task=128 --mem=512G --exclusive --partition=lowprio --qos=lowprio +# Then you enter a special space where you can srun and get all the slurm env var +# 1. To check slurm env var try +env | grep SLURM +# 2. To run the srun command, try +bash ./distill_dmd_VSA_t2v_1.3B.salloc +``` + +### Non-standard way to use `salloc + srun` + +In a tmux, run this command to reserve a terminal space +```bash +salloc -N 8 -G 64 --cpus-per-task=128 --mem=512G --exclusive --partition=lowprio --qos=lowprio +``` + +Get the job id from the output of the `squeue` command: +```bash +squeue -u $USER +``` + +Now in another terminal, run +```bash +conda activate + +# Customize the job id, number of nodes, and number of gpus per node +export JOBID=1045228 +export NUM_NODES=8 +export NUM_GPUS_PER_NODE=8 +bash ./distill_dmd_VSA_t2v_1.3B.salloc +``` \ No newline at end of file diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.salloc b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.salloc new file mode 100644 index 000000000..5d01aa9e0 --- /dev/null +++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.salloc @@ -0,0 +1,201 @@ +#!/bin/bash +#SBATCH --job-name=jd-d2ag +#SBATCH --partition=lowprio +#SBATCH --qos=lowprio +#SBATCH --nodes=2 +#SBATCH --ntasks=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:4 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=outputs/slurm/%j.out +#SBATCH --error=outputs/slurm/%j.err +#SBATCH --exclusive + +# Check required variables +if [[ -z "$JOBID" || -z "$NUM_NODES" || -z "$NUM_GPUS_PER_NODE" || -z "$MASTER_ADDR" ]]; then + echo "ERROR: One or more required environment variables are not set." + echo "You must set JOBID, NUM_NODES, NUM_GPUS_PER_NODE, and MASTER_ADDR before running this script." + echo "Example usage:" + echo ' export JOBID=1045228' + echo ' export NUM_NODES=8' + echo ' export NUM_GPUS_PER_NODE=8' + echo ' export MASTER_ADDR=your.head.node.ip' + echo ' ' + echo 'Current environment variables:' + echo " JOBID: ${JOBID}" + echo " NUM_NODES: ${NUM_NODES}" + echo " NUM_GPUS_PER_NODE: ${NUM_GPUS_PER_NODE}" + echo " MASTER_ADDR: ${MASTER_ADDR}" + echo ' ' + exit 1 +fi + +echo "JOBID is set to ${JOBID}" +echo "NUM_NODES is set to ${NUM_NODES}" +echo "NUM_GPUS_PER_NODE is set to ${NUM_GPUS_PER_NODE}" +echo "MASTER_ADDR is set to ${MASTER_ADDR}" + +MASTER_PORT=${MASTER_PORT:-29500} +echo "MASTER_PORT is set to ${MASTER_PORT}" + +NUM_GPUS=$((NUM_NODES * NUM_GPUS_PER_NODE)) + + +# Environment Setup +# source /mnt/weka/home/yonghao.zhuang/miniconda3/bin/activate +# conda activate jd-fd + +# Basic Info +export WANDB_MODE="offline" +# export WANDB_MODE=online +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# different cache dir for different processes +export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +# export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + + +# Create timestamped log directory +BASE_OUTPUT_DIR="/mnt/sharefs/users/yonghao.zhuang/d2-logs/fd/outputs/$(TZ=America/Los_Angeles date +%Y%m%d_%H%M%S)" +LOG_DIR="$BASE_OUTPUT_DIR/logs" +NSYS_LOG_DIR="$BASE_OUTPUT_DIR/nsys" +mkdir -p $BASE_OUTPUT_DIR +mkdir -p $LOG_DIR +mkdir -p $NSYS_LOG_DIR +echo "Base output directory: $BASE_OUTPUT_DIR" +echo "Log directory: $LOG_DIR" +echo "NSYS log directory: $NSYS_LOG_DIR" + +# Configs + +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train" +# VALIDATION_DATASET_FILE="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/val" +VALIDATION_DATASET_FILE="/mnt/weka/home/yonghao.zhuang/jd/training/FastVideo/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/validation_64.json" +OUTPUT_DIR="/mnt/sharefs/users/yonghao.zhuang/d2-logs/fd/" +mkdir -p $OUTPUT_DIR +echo "Output directory: $OUTPUT_DIR" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name wan_t2v_distill_dmd_VSA + --output_dir $OUTPUT_DIR + --max_train_steps 5 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + # --num_latent_t 21 + --num_latent_t 8 + --num_height 480 + --num_width 832 + --num_frames 81 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + # --num_gpus 16 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 64 + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH + --real_score_model_path $REAL_SCORE_MODEL_PATH + --fake_score_model_path $FAKE_SCORE_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 10 + --validation_sampling_steps "3" + --validation_guidance_scale "6.0" # not used for dmd inference +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 2e-6 + --mixed_precision "bf16" + --training_state_checkpointing_steps 500 + --weight_only_checkpointing_steps 500 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --ema_start_step 0 + --flow_shift 8 + --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 2 + --real_score_guidance_scale 3.5 + --VSA_sparsity 0.8 +) + + +# Generate a random rendezvous id (timestamp-based) +RDZVID=$(date +%s%N) +echo "Rendezvous ID: $RDZVID" + +set -e -x +# TODO: this triton cache dir needs to be different per-process. I want this variable to be set when the process is running. Likely we will need srun to run with `bash -lc ...` to set this inside of a srun command. + # --output="$LOG_DIR/%x_%j_task%t.out" \ + # --error="$LOG_DIR/%x_%j_task%t.err" \ + + +# srun -N $NUM_NODES -G $NUM_GPUS_PER_NODE --jobid $JOBID -w $MASTER_ADDR --cpus-per-task=128 --mem=1440G bash -lc "hostname" + +# srun -N $NUM_NODES -G $NUM_GPUS --jobid $JOBID --cpus-per-task=128 --mem=1440G hostname + +# export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +srun -N $NUM_NODES -G $NUM_GPUS --jobid $JOBID --cpus-per-task=128 --mem=1440G \ + nsys profile --show-output=true --force-overwrite=true -o ${NSYS_LOG_DIR}/%h.nsys-rep -t cuda,nvtx \ + torchrun \ + --nnodes $NUM_NODES \ + --nproc_per_node $NUM_GPUS_PER_NODE \ + --rdzv_id "$RDZVID" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/wan_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" + + diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.slurm b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.slurm new file mode 100644 index 000000000..f54666836 --- /dev/null +++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.jd.slurm @@ -0,0 +1,199 @@ +#!/bin/bash +#SBATCH --job-name=jd-d2ag +#SBATCH --partition=lowprio +#SBATCH --qos=lowprio +#SBATCH --nodes=2 +#SBATCH --ntasks=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:4 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=outputs/slurm/%j.out +#SBATCH --error=outputs/slurm/%j.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +# source /mnt/weka/home/yonghao.zhuang/miniconda3/bin/activate +# conda activate jd-fd + +# Basic Info +export WANDB_MODE="offline" +# 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 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + +# Create timestamped log directory +BASE_OUTPUT_DIR="/mnt/sharefs/users/yonghao.zhuang/d2-logs/fd/outputs/$(TZ=America/Los_Angeles date +%Y%m%d_%H%M%S)" +LOG_DIR="$BASE_OUTPUT_DIR/logs" +NSYS_LOG_DIR="$BASE_OUTPUT_DIR/nsys" +mkdir -p $BASE_OUTPUT_DIR +mkdir -p $LOG_DIR +mkdir -p $NSYS_LOG_DIR +echo "Base output directory: $BASE_OUTPUT_DIR" +echo "Log directory: $LOG_DIR" +echo "NSYS log directory: $NSYS_LOG_DIR" + +# Configs +NUM_GPUS=8 +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train" +# VALIDATION_DATASET_FILE="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/val" +VALIDATION_DATASET_FILE="/mnt/weka/home/yonghao.zhuang/jd/training/FastVideo/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/validation_64.json" +OUTPUT_DIR="/mnt/sharefs/users/yonghao.zhuang/d2-logs/fd/" +mkdir -p $OUTPUT_DIR +echo "Output directory: $OUTPUT_DIR" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name wan_t2v_distill_dmd_VSA + --output_dir $OUTPUT_DIR + --max_train_steps 5 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + # --num_latent_t 21 + --num_latent_t 8 + --num_height 480 + --num_width 832 + --num_frames 81 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus 64 + # --num_gpus 16 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 64 + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH + --real_score_model_path $REAL_SCORE_MODEL_PATH + --fake_score_model_path $FAKE_SCORE_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 10 + --validation_sampling_steps "3" + --validation_guidance_scale "6.0" # not used for dmd inference +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 2e-6 + --mixed_precision "bf16" + --training_state_checkpointing_steps 500 + --weight_only_checkpointing_steps 500 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --ema_start_step 0 + --flow_shift 8 + --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 2 + --real_score_guidance_scale 3.5 + --VSA_sparsity 0.8 +) + +# export FASTVIDEO_TORCH_PROFILER_DIR="${NSYS_LOG_DIR}/" +# export FASTVIDEO_TORCH_PROFILE_REGIONS="profiler_region_model_loading,profiler_region_model_loading,profiler_region_training_save_checkpoint,profiler_region_training_validation,profiler_region_training_train_one_step,profiler_region_training_train,profiler_region_distillation_teacher_forward,profiler_region_distillation_student_forward,profiler_region_distillation_loss,profiler_region_distillation_update" +# export FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=0 +# export FASTVIDEO_TORCH_PROFILER_WITH_STACK=0 +# export FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=0 +# export FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=0 + +# # nsys profile --show-output=true --force-overwrite=true -o ${NSYS_LOG_DIR}/%h.nsys-rep -t cuda,nvtx \ +# srun \ +# --output="$LOG_DIR/%x_%j_task%t.out" \ +# --error="$LOG_DIR/%x_%j_task%t.err" \ +# 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_distillation_pipeline.py \ +# "${parallel_args[@]}" \ +# "${model_args[@]}" \ +# "${dataset_args[@]}" \ +# "${training_args[@]}" \ +# "${optimizer_args[@]}" \ +# "${validation_args[@]}" \ +# "${miscellaneous_args[@]}" \ +# "${dmd_args[@]}" + +# Generate a random rendezvous id (timestamp-based) +RDZVID=$(date +%s%N) +echo "Rendezvous ID: $RDZVID" + + # --output="$LOG_DIR/%x_%j_task%t.out" \ + # --error="$LOG_DIR/%x_%j_task%t.err" \ + +srun \ + nsys profile --show-output=true --force-overwrite=true -o ${NSYS_LOG_DIR}/%h.nsys-rep -t cuda,nvtx \ + torchrun \ + --nnodes $SLURM_JOB_NUM_NODES \ + --nproc_per_node $NUM_GPUS \ + --rdzv_id "$RDZVID" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/wan_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" + + + +# --node_rank $SLURM_PROCID \ \ No newline at end of file