From d277361aae666bc506aeb1e4e4dc7cb339148d58 Mon Sep 17 00:00:00 2001 From: Ohm-Rishabh <119029179+Ohm-Rishabh@users.noreply.github.com> Date: Wed, 17 Dec 2025 23:45:23 -0800 Subject: [PATCH] [misc] add schedule configurations to pytorch profiler (#934) --- .../crush_smol/profile_finetune_t2v.sh | 107 ++++++++++++++++++ fastvideo/envs.py | 15 +++ fastvideo/profiler.py | 5 + 3 files changed, 127 insertions(+) create mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh new file mode 100644 index 000000000..0802d54db --- /dev/null +++ b/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[@]}" diff --git a/fastvideo/envs.py b/fastvideo/envs.py index fffd6955c..e73d926f3 100644 --- a/fastvideo/envs.py +++ b/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", ""), diff --git a/fastvideo/profiler.py b/fastvideo/profiler.py index 914f63a06..71aaecb22 100644 --- a/fastvideo/profiler.py +++ b/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), )