diff --git a/examples/inference/optimizations/quantization_example.py b/examples/inference/optimizations/quantization_example.py index f6f6db476..171c3f6ce 100644 --- a/examples/inference/optimizations/quantization_example.py +++ b/examples/inference/optimizations/quantization_example.py @@ -5,8 +5,7 @@ from fastvideo.layers.quantization.fp4_config import FP4Config # from fastvideo.layers.quantization.fp8_config import FP8Config from fastvideo.layers.linear import ReplicatedLinear -OUTPUT_PATH = "fp4_video_samples" -# OUTPUT_PATH = "video_samples" +OUTPUT_PATH = "video_samples" def main(): @@ -66,23 +65,23 @@ def main(): print("✓ First video generated successfully with FP4 quantization!") # # Generate a second video to show the model can be reused - # prompt2 = ( - # "A majestic lion strides across the golden savanna, its powerful frame " - # "glistening under the warm afternoon sun. The tall grass ripples gently in " - # "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - # "embodying the raw energy of the wild. Low angle, steady tracking shot, " - # "cinematic." - # ) + prompt2 = ( + "A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic." + ) - # print(f"\nGenerating second video...") - # print(f"Prompt: {prompt2}") + print(f"\nGenerating second video...") + print(f"Prompt: {prompt2}") - # video2 = generator.generate_video( - # prompt2, - # output_path=OUTPUT_PATH, - # save_video=True, - # ) - # print("✓ Second video generated successfully with FP4 quantization!") + video2 = generator.generate_video( + prompt2, + output_path=OUTPUT_PATH, + save_video=True, + ) + print("✓ Second video generated successfully with FP4 quantization!") except Exception as e: print(f"Error during video generation: {e}") diff --git a/examples/inference/sageattn+fp4linear.py b/examples/inference/sageattn+fp4linear.py new file mode 100644 index 000000000..dceb6a3c2 --- /dev/null +++ b/examples/inference/sageattn+fp4linear.py @@ -0,0 +1,49 @@ +from fastvideo import VideoGenerator +import os +# from fastvideo.configs.sample import SamplingParam + +OUTPUT_PATH = "video_samples" +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SAGE_ATTN_THREE" + +def main(): + # FastVideo will automatically use the optimal default arguments for the + # model. + # If a local path is provided, FastVideo will make a best effort + # attempt to identify the optimal arguments. + generator = VideoGenerator.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + # FastVideo will automatically handle distributed setup + num_gpus=1, + use_fsdp_inference=True, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=False, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + # image_encoder_cpu_offload=False, + ) + + # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + # sampling_param.num_frames = 45 + # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" + # Generate videos with the same simple API, regardless of GPU count + prompt = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + ) + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) + # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") + + # Generate another video with a different prompt, without reloading the + # model! + prompt2 = ( + "A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") + video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) + + +if __name__ == "__main__": + main() diff --git a/examples/inference/sageattn_inference_example_custom_weights_5090.py b/examples/inference/sageattn_inference_example_custom_weights_5090.py index 26c312b6f..3f9fef0c4 100644 --- a/examples/inference/sageattn_inference_example_custom_weights_5090.py +++ b/examples/inference/sageattn_inference_example_custom_weights_5090.py @@ -16,7 +16,7 @@ def main(): generator = VideoGenerator.from_pretrained( "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", # FastVideo will automatically handle distributed setup - num_gpus=4, + num_gpus=1, use_fsdp_inference=True, dit_cpu_offload=False, vae_cpu_offload=False, diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_e2e.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_e2e.sh new file mode 100644 index 000000000..61de79062 --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_e2e.sh @@ -0,0 +1,98 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +export TRITON_PRINT_AUTOTUNING=1 # to print the best config +export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR=data/Wan-Syn_77x448x832_600k +VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" +NUM_GPUS=1 +# export CUDA_VISIBLE_DEVICES=4,5 + + +# Training arguments +training_args=( + --tracker_project_name "wan_t2v_finetune_qat" + --output_dir "checkpoints/wan_t2v_finetune_fp4_e2e" + --max_train_steps 5000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 12 + --num_height 480 + --num_width 832 + --num_frames 45 + --enable_gradient_checkpointing_type "full" + --generator_4bit_linear True + --generator_4bit_attn True +) + +# 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 4 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file $VALIDATION_DATASET_FILE + --validation_steps 100 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-6 + --mixed_precision "bf16" + --weight_only_checkpointing_steps 500 + --training_state_checkpointing_steps 500 + --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 \ + fastvideo/training/wan_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_linear.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_linear.sh new file mode 100644 index 000000000..1187c5a47 --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_linear.sh @@ -0,0 +1,97 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +export TRITON_PRINT_AUTOTUNING=1 # to print the best config +export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR=data/Wan-Syn_77x448x832_600k +VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" +NUM_GPUS=1 +# export CUDA_VISIBLE_DEVICES=4,5 + + +# Training arguments +training_args=( + --tracker_project_name "wan_t2v_finetune_qat" + --output_dir "checkpoints/wan_t2v_finetune_fp4_linear" + --max_train_steps 5000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 12 + --num_height 480 + --num_width 832 + --num_frames 45 + --enable_gradient_checkpointing_type "full" + --generator_4bit_linear True +) + +# 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 4 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file $VALIDATION_DATASET_FILE + --validation_steps 100 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-6 + --mixed_precision "bf16" + --weight_only_checkpointing_steps 500 + --training_state_checkpointing_steps 500 + --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 \ + fastvideo/training/wan_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4_5090.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4_5090.sh new file mode 100644 index 000000000..56c7da9fd --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4_5090.sh @@ -0,0 +1,94 @@ +#!/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/Wan-Syn_77x448x832_600k/train/" +VALIDATION_DATASET_FILE="data/Wan-Syn_77x448x832_600k/validation_6.json" +NUM_GPUS=6 +# export CUDA_VISIBLE_DEVICES=4,5 + + +# 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 1 + --num_latent_t 6 + --num_height 480 + --num_width 832 + --num_frames 21 + --enable_gradient_checkpointing_type "full" + --generator_4bit_linear True +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --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 50 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-6 + --mixed_precision "bf16" + --checkpointing_steps 2000 + --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 \ + fastvideo/training/wan_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat.sh new file mode 100644 index 000000000..1138ade3e --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat.sh @@ -0,0 +1,97 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +export TRITON_PRINT_AUTOTUNING=1 # to print the best config +export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR=data/Wan-Syn_77x448x832_600k +VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" +NUM_GPUS=1 +# export CUDA_VISIBLE_DEVICES=4,5 + + +# Training arguments +training_args=( + --tracker_project_name "wan_t2v_finetune_qat" + --output_dir "checkpoints/wan_t2v_finetune_no_high_prec_o" + --max_train_steps 5000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 12 + --num_height 480 + --num_width 832 + --num_frames 45 + --enable_gradient_checkpointing_type "full" + --generator_4bit_attn True +) + +# 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 4 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file $VALIDATION_DATASET_FILE + --validation_steps 100 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-6 + --mixed_precision "bf16" + --weight_only_checkpointing_steps 500 + --training_state_checkpointing_steps 500 + --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 \ + fastvideo/training/wan_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200.sh new file mode 100644 index 000000000..c8c5abc2d --- /dev/null +++ b/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200.sh @@ -0,0 +1,115 @@ +#!/bin/bash +#SBATCH --job-name=wan_t2v_14B_finetune_sage3 +#SBATCH --partition=all +#SBATCH --nodes=2 +#SBATCH --gres=gpu:4 +#SBATCH --ntasks-per-node=4 +#SBATCH --output=logs/wan_t2v_14B_finetune_sage3_%j.out +#SBATCH --error=logs/wan_t2v_14B_finetune_sage3_%j.err + +source .venv/bin/activate + + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +# export TRITON_PRINT_AUTOTUNING=1 # to print the best config +export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f +MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" +DATA_DIR=data/Wan-Syn_77x448x832_600k +VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 + +set -euo pipefail + +# ---- torchrun rendezvous (multi-node) ---- +# Launch ONE torchrun per node (via srun) and let torchrun spawn 4 workers per node. +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" +MASTER_PORT="${MASTER_PORT:-29500}" +export MASTER_ADDR MASTER_PORT + + +# Training arguments +training_args=( + --tracker_project_name "wan_t2v_finetune_qat" + --output_dir "checkpoints/wan_t2v_finetune_14B" + --max_train_steps 4000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 16 + --num_height 448 + --num_width 832 + --num_frames 61 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 8 + --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 4 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file $VALIDATION_DATASET_FILE + --validation_steps 200 + --validation_sampling_steps "50" + --validation_guidance_scale "5.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-6 + --mixed_precision "bf16" + --weight_only_checkpointing_steps 1000 + --training_state_checkpointing_steps 1000 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --dit_precision "fp32" + --ema_start_step 0 + --flow_shift 5 + --seed 1000 +) + +srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ + torchrun \ + --nnodes "$SLURM_NNODES" \ + --nproc_per_node 4 \ + --rdzv_backend c10d \ + --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \ + --node_rank "$SLURM_NODEID" \ + fastvideo/training/wan_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh index f059c3172..22d908f97 100644 --- a/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh +++ b/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh @@ -4,8 +4,8 @@ #SBATCH --nodes=2 #SBATCH --gres=gpu:4 #SBATCH --ntasks-per-node=4 -#SBATCH --output=logs/wan_t2v_14B_finetune_sage3_%j.out -#SBATCH --error=logs/wan_t2v_14B_finetune_sage3_%j.err +#SBATCH --output=logs/wan_t2v_14B_finetune_qat_%j.out +#SBATCH --error=logs/wan_t2v_14B_finetune_qat_%j.err source .venv/bin/activate @@ -18,7 +18,7 @@ export TOKENIZERS_PARALLELISM=false # export TRITON_PRINT_AUTOTUNING=1 # to print the best config export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k +DATA_DIR=FastVideo/Wan2.2-Syn-121x704x1280_32k/train VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" NUM_GPUS=8 # export CUDA_VISIBLE_DEVICES=4,5 @@ -35,15 +35,15 @@ export MASTER_ADDR MASTER_PORT # Training arguments training_args=( --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_sage3_14B" + --output_dir "checkpoints/wan_t2v_finetune_qat_14B" --max_train_steps 4000 --train_batch_size 1 --train_sp_batch_size 1 --gradient_accumulation_steps 1 - --num_latent_t 20 + --num_latent_t 21 --num_height 720 --num_width 1280 - --num_frames 77 + --num_frames 81 --enable_gradient_checkpointing_type "full" --generator_4bit_attn True ) @@ -80,7 +80,7 @@ validation_args=( # Optimizer arguments optimizer_args=( - --learning_rate 1e-5 + --learning_rate 1e-6 --mixed_precision "bf16" --weight_only_checkpointing_steps 1000 --training_state_checkpointing_steps 1000 diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh index c984e5774..8849c7246 100644 --- a/examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh +++ b/examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh @@ -37,10 +37,10 @@ training_args=( --train_batch_size 1 --train_sp_batch_size 1 --gradient_accumulation_steps 1 - --num_latent_t 20 + --num_latent_t 32 --num_height 720 --num_width 1280 - --num_frames 77 + --num_frames 125 --enable_gradient_checkpointing_type "full" --generator_4bit_attn True ) diff --git a/fastvideo/attention/backends/sageattn/api.py b/fastvideo/attention/backends/sageattn/api.py index 28b11f4f1..b679138e8 100644 --- a/fastvideo/attention/backends/sageattn/api.py +++ b/fastvideo/attention/backends/sageattn/api.py @@ -77,7 +77,7 @@ def triton_group_mean(q: torch.Tensor): return q_out, qm -def preprocess_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, per_block_mean: bool = True, disable_delta_s: bool = False): +def preprocess_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, per_block_mean: bool = True, disable_smoothing_q: bool = True): def pad_to_block_size(x): L = x.size(2) @@ -88,14 +88,17 @@ def preprocess_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, per_block_ k -= k.mean(dim=-2, keepdim=True) q, k, v = map(lambda x: pad_to_block_size(x), [q, k, v]) - if per_block_mean: + if per_block_mean and not disable_smoothing_q: q, qm = triton_group_mean(q) - else: + elif not disable_smoothing_q: qm = q.mean(dim=-2, keepdim=True) q = q - qm - delta_s = torch.matmul(qm, k.transpose(-2, -1)).to(torch.float32).contiguous() - if disable_delta_s: # used to disable q smoothing - delta_s = torch.zeros_like(delta_s) + if disable_smoothing_q: # used to disable q smoothing + B, H, L, D = q.shape + delta_s = torch.zeros((B, H, L // BLOCK_M, k.shape[2]), device=q.device, dtype=torch.float32) + else: + _, qm = triton_group_mean(q) + delta_s = torch.matmul(qm, k.transpose(-2, -1)).to(torch.float32).contiguous() return q, k, v, delta_s diff --git a/fastvideo/layers/quantization/fp4_config.py b/fastvideo/layers/quantization/fp4_config.py index f9361bf1d..08be0a392 100644 --- a/fastvideo/layers/quantization/fp4_config.py +++ b/fastvideo/layers/quantization/fp4_config.py @@ -44,7 +44,7 @@ class FP4QuantizeMethod(QuantizeMethodBase): x_fp4, x_scale = nvfp4_quantize(x, x_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) weight_fp4 = layer._fp4_weight weight_scale = layer._fp4_weight_scale - weight_global_sf = layer._weight_global_sf.item() + weight_global_sf = layer._weight_global_sf out = mm_fp4(x_fp4, weight_fp4.T, x_scale, weight_scale.T, 1.0/(x_global_sf * weight_global_sf), torch.bfloat16, None, backend='cutlass') diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 87168e26e..ec1987efb 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -59,7 +59,7 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: USE_FP8 = False -USE_FP4 = False +USE_FP4 = False # TODO: move to FastVideoArgs # TODO(PY): add compile option def maybe_load_fsdp_model( model_cls: type[nn.Module],