[misc] add schedule configurations to pytorch profiler (#934)

examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh

@@ -0,0 +1,107 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=4
+# export CUDA_VISIBLE_DEVICES=4,5
+
+mkdir -p ../profiler_traces/wan_t2v_finetune/
+
+# Torch Profiler Configuration
+export FASTVIDEO_TORCH_PROFILE_REGIONS="profiler_region_training_train_one_step"
+export FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1
+export FASTVIDEO_TORCH_PROFILER_WITH_STACK=1  
+export FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1
+export FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY=1
+export FASTVIDEO_TORCH_PROFILER_WAIT_STEPS=2
+export FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS=1
+export FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS=1
+export FASTVIDEO_TORCH_PROFILER_DIR="../profiler_traces/wan_t2v_finetune/"
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_t2v_finetune"
+  --output_dir "checkpoints/wan_t2v_finetune"
+  --max_train_steps 5000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 8
+  --num_latent_t 20
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS 
+  --sp_size $NUM_GPUS 
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# 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 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation 
+  --validation_dataset_file $VALIDATION_DATASET_FILE
+  --validation_steps 200
+  --validation_sampling_steps "50" 
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-5
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
+  # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+  --master_port 29502 \
+    fastvideo/training/wan_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

fastvideo/envs.py

@@ -32,6 +32,9 @@ if TYPE_CHECKING:
     FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
     FASTVIDEO_TORCH_PROFILER_WITH_STACK: bool = True
     FASTVIDEO_TORCH_PROFILER_WITH_FLOPS: bool = False
+    FASTVIDEO_TORCH_PROFILER_WAIT_STEPS: int = 2
+    FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS: int = 1
+    FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS: int = 2
     FASTVIDEO_TORCH_PROFILE_REGIONS: str = ""
     FASTVIDEO_SERVER_DEV_MODE: bool = False
     FASTVIDEO_STAGE_LOGGING: bool = False
@@ -247,6 +250,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # not profile flops.
     "FASTVIDEO_TORCH_PROFILER_WITH_FLOPS":
     lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_FLOPS", "0") != "0"),
+    # Wait steps per profiling cycle (torch.profiler.schedule wait parameter)
+    # Defaults to 2 if not set.
+    "FASTVIDEO_TORCH_PROFILER_WAIT_STEPS":
+    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_WAIT_STEPS", "2")),
+    # Warmup steps per profiling cycle (torch.profiler.schedule warmup parameter)
+    # Defaults to 1 if not set.
+    "FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS":
+    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS", "1")),
+    # Active steps per profiling cycle (torch.profiler.schedule active parameter)
+    # Defaults to 2 if not set.
+    "FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS":
+    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS", "2")),
     "FASTVIDEO_TORCH_PROFILE_REGIONS":
     lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""),
 

fastvideo/profiler.py

@@ -208,6 +208,11 @@ def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController:
         profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
         with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
         with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
+        schedule=torch.profiler.schedule(
+            wait=envs.FASTVIDEO_TORCH_PROFILER_WAIT_STEPS,
+            warmup=envs.FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS,
+            active=envs.FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS,
+        ),
         on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir,
                                                                 use_gzip=True),
     )