modify scripts + launch qat run

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

@@ -10,7 +10,7 @@ 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
+NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5
 
 
@@ -19,7 +19,7 @@ training_args=(
   --tracker_project_name "wan_t2v_finetune_qat"
   --output_dir "checkpoints/wan_t2v_finetune_qat"
   --max_train_steps 5000
-  --train_batch_size 1
+  --train_batch_size 8
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
   --num_latent_t 12
@@ -33,7 +33,7 @@ training_args=(
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS 
-  --sp_size 1
+  --sp_size $NUM_GPUS
   --tp_size 1
   --hsdp_replicate_dim 1
   --hsdp_shard_dim $NUM_GPUS
@@ -48,14 +48,14 @@ model_args=(
 # Dataset arguments
 dataset_args=(
   --data_path $DATA_DIR
-  --dataloader_num_workers 1
+  --dataloader_num_workers 4
 )
 
 # Validation arguments
 validation_args=(
   --log_validation 
   --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 50
+  --validation_steps 100
   --validation_sampling_steps "50" 
   --validation_guidance_scale "6.0"
 )
@@ -64,8 +64,8 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-6
   --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 200
-  --training_state_checkpointing_steps 200
+  --weight_only_checkpointing_steps 500
+  --training_state_checkpointing_steps 500
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

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

@@ -62,7 +62,7 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-6
+  --learning_rate 1e-7
   --mixed_precision "bf16"
   --weight_only_checkpointing_steps 500
   --training_state_checkpointing_steps 500

fastvideo/attention/backends/sage_attn3.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from curses import KEY_B2
 import os
 import sys
 from pathlib import Path
@@ -502,7 +501,7 @@ class SageAttention3Impl(AttentionImpl):
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         # output = sageattn_blackwell_with_16bit_bwd(query, key, value, is_causal=self.causal)
-        output = sageattn_blackwell_with_triton_bwd(query, key, value, is_causal=self.causal)
+        # output = sageattn_blackwell_with_triton_bwd(query, key, value, is_causal=self.causal)
         # output = attn_forward_4bit_fwd_16bit_bwd(query, key, value)
-        # output = qat_attn(query, key, value, is_causal=self.causal)
+        output = qat_attn(query, key, value, is_causal=self.causal)
         return output