Make nsys selective and also add annotations

Add nsys
small fix on the wan2.1 T2V 1.3B readme
2025-11-28 08:27:32 +00:00 · 2025-11-28 07:23:10 +00:00 · 2025-11-26 23:51:02 +00:00
6 changed files with 613 additions and 98 deletions
@@ -70,3 +70,8 @@ docs/distillation/examples/

 dmd_t2v_output/
 preprocess_output_text/
+
+
+outputs 
+outputs/
+*.log
@@ -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,236 @@
+#!/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 --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" "$LOG_DIR" "$NSYS_LOG_DIR"
+echo "Base output directory: $BASE_OUTPUT_DIR"
+echo "Log directory: $LOG_DIR"
+echo "NSYS log directory: $NSYS_LOG_DIR"
+
+ENABLE_NSYS_PROFILING=${ENABLE_NSYS_PROFILING:-1}
+NSYS_RANK0_ONLY=${NSYS_RANK0_ONLY:-1}
+echo "ENABLE_NSYS_PROFILING=${ENABLE_NSYS_PROFILING}"
+echo "NSYS_RANK0_ONLY=${NSYS_RANK0_ONLY}"
+
+# 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"
+
+
+# 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.
+  
+
+
+# Pre-compose the torchrun command as an array
+torchrun_cmd=(
+  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[@]}"
+)
+
+nsys_cmd=(
+  nsys
+    profile
+    --show-output=true
+    --force-overwrite=true
+    -o "${NSYS_LOG_DIR}/%h.nsys-rep"
+    -t cuda,nvtx
+)
+
+nsys_cmd_decl=$(declare -p nsys_cmd)
+
+  # --output="$LOG_DIR/%x_%j_task%t.out" \
+  # --error="$LOG_DIR/%x_%j_task%t.err" \
+
+set -e -x
+srun \
+    -N "$NUM_NODES" -G "$NUM_GPUS" --jobid "$JOBID" \
+    --cpus-per-task=128 --mem=1440G \
+  bash -c '
+    set -e
+    '"$nsys_cmd_decl"'
+    export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+
+    should_profile=${ENABLE_NSYS_PROFILING}
+    if [[ "${ENABLE_NSYS_PROFILING}" == "1" && "${NSYS_RANK0_ONLY}" == "1" && "${SLURM_PROCID}" != "0" ]]; then
+      should_profile=0
+    fi
+
+    log_file="'"${LOG_DIR}"'"/${SLURM_PROCID}.out
+    if [[ "${should_profile}" == "1" ]]; then
+      echo "Profiling with NSYS on rank ${SLURM_PROCID}"
+      "${nsys_cmd[@]}" "$@" 2>&1 | tee "${log_file}"
+    else
+      echo "Running without NSYS on rank ${SLURM_PROCID}"
+      "$@" 2>&1 | tee "${log_file}"
+    fi
+  ' _ "${torchrun_cmd[@]}"
+
+
+# log_file="'"${LOG_DIR}"'"/${SLURM_PROCID}.out
@@ -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 \
@@ -5,12 +5,12 @@ These are e2e example scripts for finetuning Wan2.1 T2V 1.3B on the crush-smol d

 ### Download crush-smol dataset:

-`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh`
+`bash examples/training/finetune/wan_t2v_1.3B/crush_smol/download_dataset.sh`

 ### Preprocess the videos and captions into latents:

-`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh`
+`bash examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh`

 ### Edit the following file and run finetuning:

-`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh`
+`bash examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh`
@@ -7,6 +7,7 @@ import time
 from abc import abstractmethod
 from collections import deque
 from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import Any

 import imageio
@@ -43,11 +44,28 @@ from fastvideo.training.training_utils import (
 from fastvideo.utils import (is_vsa_available, maybe_download_model,
                             set_random_seed, verify_model_config_and_directory)

+try:
+    from torch.cuda import nvtx
+except ImportError:  # pragma: no cover - CPU-only environments
+    nvtx = None
+
 vsa_available = is_vsa_available()

 logger = init_logger(__name__)


+@contextmanager
+def nvtx_range(message: str):
+    """
+    Lightweight NVTX range helper that is a no-op when torch.cuda.nvtx is unavailable.
+    """
+    if nvtx is None:
+        yield
+    else:
+        with nvtx.range(message):
+            yield
+
+
 class DistillationPipeline(TrainingPipeline):
    """
    A distillation pipeline for training a 3 step model.
@@ -947,120 +965,143 @@ class DistillationPipeline(TrainingPipeline):
                                              'gradient_accumulation_steps', 1)
        batches = []
        # Collect N batches for gradient accumulation
-        for _ in range(gradient_accumulation_steps):
-            batch = self._prepare_distillation(training_batch)
-            batch = self._get_next_batch(batch)
-            batch = self._normalize_dit_input(batch)
-            batch = self._prepare_dit_inputs(batch)
-            batch = self._build_attention_metadata(batch)
-            batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata)
-            if batch.attn_metadata is not None:
-                batch.attn_metadata.VSA_sparsity = 0.0  # type: ignore
-            batches.append(batch)
+        with nvtx_range("prepare_microbatches"):
+            for _ in range(gradient_accumulation_steps):
+                batch = self._prepare_distillation(training_batch)
+                batch = self._get_next_batch(batch)
+                batch = self._normalize_dit_input(batch)
+                batch = self._prepare_dit_inputs(batch)
+                batch = self._build_attention_metadata(batch)
+                batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata)
+                if batch.attn_metadata is not None:
+                    batch.attn_metadata.VSA_sparsity = 0.0  # type: ignore
+                batches.append(batch)

-        self.optimizer.zero_grad()
+        with nvtx_range("generator_zero_grad"):
+            self.optimizer.zero_grad()
        total_dmd_loss = 0.0
        dmd_latent_vis_dict = {}
        fake_score_latent_vis_dict = {}
        if (self.current_trainstep % self.generator_update_interval == 0):
-            for batch in batches:
-                batch_gen = copy.deepcopy(batch)
+            with nvtx_range("generator_step"):
+                for batch_idx, batch in enumerate(batches):
+                    batch_gen = copy.deepcopy(batch)
+                    with nvtx_range(f"generator_microbatch[{batch_idx}]"):
+                        with set_forward_context(
+                                current_timestep=batch_gen.timesteps,
+                                attn_metadata=batch_gen.attn_metadata_vsa):
+                            if self.training_args.simulate_generator_forward:
+                                with nvtx_range("generator_forward_sim"):
+                                    generator_pred_video = self._generator_multi_step_simulation_forward(
+                                        batch_gen)
+                            else:
+                                with nvtx_range("generator_forward"):
+                                    generator_pred_video = self._generator_forward(
+                                        batch_gen)

-                with set_forward_context(
-                        current_timestep=batch_gen.timesteps,
-                        attn_metadata=batch_gen.attn_metadata_vsa):
-                    if self.training_args.simulate_generator_forward:
-                        generator_pred_video = self._generator_multi_step_simulation_forward(
-                            batch_gen)
-                    else:
-                        generator_pred_video = self._generator_forward(
-                            batch_gen)
+                        with nvtx_range("dmd_forward"), set_forward_context(
+                                current_timestep=batch_gen.timesteps,
+                                attn_metadata=batch_gen.attn_metadata):
+                            dmd_loss = self._dmd_forward(
+                                generator_pred_video=generator_pred_video,
+                                training_batch=batch_gen)

-                with set_forward_context(current_timestep=batch_gen.timesteps,
-                                         attn_metadata=batch_gen.attn_metadata):
-                    dmd_loss = self._dmd_forward(
-                        generator_pred_video=generator_pred_video,
-                        training_batch=batch_gen)
+                        with nvtx_range("generator_backward"), set_forward_context(
+                                current_timestep=batch_gen.timesteps,
+                                attn_metadata=batch_gen.attn_metadata_vsa):
+                            (dmd_loss / gradient_accumulation_steps).backward()
+                        total_dmd_loss += dmd_loss.detach().item()

-                with set_forward_context(
-                        current_timestep=batch_gen.timesteps,
-                        attn_metadata=batch_gen.attn_metadata_vsa):
-                    (dmd_loss / gradient_accumulation_steps).backward()
-                total_dmd_loss += dmd_loss.detach().item()
+                with nvtx_range("generator_post_step"):
+                    # Only clip gradients for the model that is currently training
+                    self._clip_model_grad_norm_(batch_gen, self.transformer)
+                    for param in self.transformer.parameters():
+                        # check if the gradient is not None and not zero
+                        assert param.grad is not None and param.grad.abs(
+                        ).sum() > 0
+                    self.optimizer.step()
+                    self.optimizer.zero_grad(set_to_none=True)

-            # Only clip gradients for the model that is currently training
-            self._clip_model_grad_norm_(batch_gen, self.transformer)
-            for param in self.transformer.parameters():
-                # check if the gradient is not None and not zero
-                assert param.grad is not None and param.grad.abs().sum() > 0
-            self.optimizer.step()
-            self.optimizer.zero_grad(set_to_none=True)
+                    if self.generator_ema is not None:
+                        self.generator_ema.update(self.transformer)
+                    if self.generator_ema_2 is not None:
+                        self.generator_ema_2.update(self.transformer_2)

-            if self.generator_ema is not None:
-                self.generator_ema.update(self.transformer)
-            if self.generator_ema_2 is not None:
-                self.generator_ema_2.update(self.transformer_2)
-
-            avg_dmd_loss = torch.tensor(total_dmd_loss /
-                                        gradient_accumulation_steps,
-                                        device=self.device)
-            world_group = get_world_group()
-            world_group.all_reduce(avg_dmd_loss,
-                                   op=torch.distributed.ReduceOp.AVG)
-            training_batch.generator_loss = avg_dmd_loss.item()
-            dmd_latent_vis_dict = batch_gen.dmd_latent_vis_dict
+                with nvtx_range("generator_reduce_loss"):
+                    avg_dmd_loss = torch.tensor(total_dmd_loss /
+                                                gradient_accumulation_steps,
+                                                device=self.device)
+                    world_group = get_world_group()
+                    world_group.all_reduce(avg_dmd_loss,
+                                           op=torch.distributed.ReduceOp.AVG)
+                training_batch.generator_loss = avg_dmd_loss.item()
+                dmd_latent_vis_dict = batch_gen.dmd_latent_vis_dict
        else:
            training_batch.generator_loss = 0.0

-        self.fake_score_optimizer.zero_grad()
-        if self.fake_score_transformer_2 is not None:
-            self.fake_score_optimizer_2.zero_grad()
+        with nvtx_range("fake_score_zero_grad"):
+            self.fake_score_optimizer.zero_grad()
+            if self.fake_score_transformer_2 is not None:
+                self.fake_score_optimizer_2.zero_grad()
        total_fake_score_loss = 0.0
-        for batch in batches:
-            batch_fake = copy.deepcopy(batch)
-            batch_fake, fake_score_loss = self.faker_score_forward(batch_fake)
-            with set_forward_context(current_timestep=batch_fake.timesteps,
-                                     attn_metadata=batch_fake.attn_metadata):
-                (fake_score_loss / gradient_accumulation_steps).backward()
-            total_fake_score_loss += fake_score_loss.detach().item()
-            fake_score_latent_vis_dict.update(
-                batch_fake.fake_score_latent_vis_dict)
-        if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
-            self._clip_model_grad_norm_(batch_fake,
-                                        self.fake_score_transformer_2)
-        else:
-            self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer)
+        with nvtx_range("fake_score_step"):
+            for batch_idx, batch in enumerate(batches):
+                batch_fake = copy.deepcopy(batch)
+                with nvtx_range(f"fake_score_microbatch_{batch_idx}"):
+                    with nvtx_range("fake_score_forward"):
+                        batch_fake, fake_score_loss = self.faker_score_forward(
+                            batch_fake)
+                    with nvtx_range("fake_score_backward"), set_forward_context(
+                            current_timestep=batch_fake.timesteps,
+                            attn_metadata=batch_fake.attn_metadata):
+                        (fake_score_loss /
+                         gradient_accumulation_steps).backward()
+                    total_fake_score_loss += fake_score_loss.detach().item()
+                    fake_score_latent_vis_dict.update(
+                        batch_fake.fake_score_latent_vis_dict)
+            with nvtx_range("fake_score_post_step"):
+                if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
+                    self._clip_model_grad_norm_(batch_fake,
+                                                self.fake_score_transformer_2)
+                else:
+                    self._clip_model_grad_norm_(batch_fake,
+                                                self.fake_score_transformer)

-        # Check gradients for fake score transformer
-        for param in self.fake_score_transformer.parameters():
-            if param.requires_grad:
-                assert param.grad is not None and param.grad.abs().sum() > 0
+                # Check gradients for fake score transformer
+                for param in self.fake_score_transformer.parameters():
+                    if param.requires_grad:
+                        assert param.grad is not None and param.grad.abs(
+                        ).sum() > 0

-        # Check gradients for fake score transformer_2 if available
-        if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
-            for param in self.fake_score_transformer_2.parameters():
-                if param.requires_grad:
-                    assert param.grad is not None and param.grad.abs().sum() > 0
+                # Check gradients for fake score transformer_2 if available
+                if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
+                    for param in self.fake_score_transformer_2.parameters():
+                        if param.requires_grad:
+                            assert param.grad is not None and param.grad.abs(
+                            ).sum() > 0

-        if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
-            self.fake_score_optimizer_2.step()
-            self.fake_score_lr_scheduler_2.step()
-        else:
-            self.fake_score_optimizer.step()
-            self.fake_score_lr_scheduler.step()
+                if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
+                    self.fake_score_optimizer_2.step()
+                    self.fake_score_lr_scheduler_2.step()
+                else:
+                    self.fake_score_optimizer.step()
+                    self.fake_score_lr_scheduler.step()

-        # Step the appropriate scheduler
-        self.lr_scheduler.step()
+        with nvtx_range("scheduler_step"):
+            # Step the appropriate scheduler
+            self.lr_scheduler.step()

-        self.fake_score_optimizer.zero_grad(set_to_none=True)
-        if self.fake_score_transformer_2 is not None:
-            self.fake_score_optimizer_2.zero_grad(set_to_none=True)
-        avg_fake_score_loss = torch.tensor(total_fake_score_loss /
-                                           gradient_accumulation_steps,
-                                           device=self.device)
-        world_group = get_world_group()
-        world_group.all_reduce(avg_fake_score_loss,
-                               op=torch.distributed.ReduceOp.AVG)
+        with nvtx_range("fake_score_zero_grad_post"):
+            self.fake_score_optimizer.zero_grad(set_to_none=True)
+            if self.fake_score_transformer_2 is not None:
+                self.fake_score_optimizer_2.zero_grad(set_to_none=True)
+        with nvtx_range("fake_score_reduce_loss"):
+            avg_fake_score_loss = torch.tensor(total_fake_score_loss /
+                                               gradient_accumulation_steps,
+                                               device=self.device)
+            world_group = get_world_group()
+            world_group.all_reduce(avg_fake_score_loss,
+                                   op=torch.distributed.ReduceOp.AVG)
        training_batch.fake_score_loss = avg_fake_score_loss.item()
        training_batch.dmd_latent_vis_dict = dmd_latent_vis_dict
        training_batch.fake_score_latent_vis_dict = batch_fake.fake_score_latent_vis_dict
Author	SHA1	Message	Date
Yonghao Zhuang	7979a2c420	Make nsys selective and also add annotations	2025-11-28 08:27:32 +00:00
Yonghao Zhuang	ba75e61e4e	Add nsys	2025-11-28 07:23:10 +00:00
Yonghao Zhuang	b52a5ae237	small fix on the wan2.1 T2V 1.3B readme	2025-11-26 23:51:02 +00:00