Add nsys
This commit is contained in:
Yonghao Zhuang
@@ -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 <your_conda_env>
|
||||
|
||||
# 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
|
||||
```
|
||||
@@ -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[@]}"
|
||||
|
||||
|
||||
@@ -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 \
|
||||
Reference in New Issue
Block a user