From eb160a2e232d7bb1a492d9b8680bbaf86501a4a7 Mon Sep 17 00:00:00 2001
From: RandNMR73 <notomatthew31@gmail.com>
Date: Thu, 25 Dec 2025 09:30:27 +0000
Subject: [PATCH] modify scripts + launch qat run

---
 .../wan_t2v_1.3B/crush_smol/finetune_t2v_qat.sh    | 14 +++++++-------
 .../wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh  |  2 +-
 fastvideo/attention/backends/sage_attn3.py         |  5 ++---
 3 files changed, 10 insertions(+), 11 deletions(-)

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
index d98e039eb..3ff8498bb 100644
--- 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
@@ -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
 )
diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh
index c50feae70..74421c010 100644
--- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh
+++ b/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
diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py
index a9bd941c5..580c800af 100644
--- a/fastvideo/attention/backends/sage_attn3.py
+++ b/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