nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

5090 testing

Browse Source
This commit is contained in:
RandNMR73
2025-12-24 06:04:36 +00:00
parent 3b70702af9
commit 7afe836f28
12 changed files with 681 additions and 589 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.sh
+18 -13
View File
@@ -4,10 +4,11 @@ set -euo pipefail
############################################
# Single node, 6 GPUs
############################################
NUM_GPUS=6
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5
NUM_GPUS=1
export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f
# export CUDA_VISIBLE_DEVICES=0,1,2,3
export TRITON_PRINT_AUTOTUNING=1 # to print the best config
export WANDB_MODE="online"
export WANDB_BASE_URL="https://api.wandb.ai"
export TOKENIZERS_PARALLELISM=false
@@ -23,44 +24,48 @@ echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
# Configs
############################################
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
OUTPUT_DIR="checkpoints/full_disitlle"
REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
FAKE_SCORE_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"
OUTPUT_DIR="checkpoints/base_distill"
# Training arguments
training_args=(
--tracker_project_name wan_t2v_distill_dmd
--tracker_project_name wan_t2v_distill_dmd_qat
--output_dir "$OUTPUT_DIR"
--max_train_steps 4000
--train_batch_size 1
--train_batch_size 64
--train_sp_batch_size 1
--gradient_accumulation_steps 1
--num_latent_t 21
--num_latent_t 12
--num_height 480
--num_width 832
--num_frames 81
--num_frames 45
--enable_gradient_checkpointing_type "full"
)
# Parallel arguments (adjusted to 6 GPUs)
parallel_args=(
--num_gpus 6
--num_gpus $NUM_GPUS
--sp_size 1
--tp_size 1
--hsdp_replicate_dim 1
--hsdp_shard_dim 6
--hsdp_shard_dim $NUM_GPUS
)
# Model arguments
model_args=(
--model_path "$MODEL_PATH"
--pretrained_model_name_or_path "$MODEL_PATH"
--real_score_model_path "$REAL_SCORE_MODEL_PATH"
--fake_score_model_path "$FAKE_SCORE_MODEL_PATH"
)
# Dataset arguments
dataset_args=(
--data_path "$DATA_DIR"
--dataloader_num_workers 4
--dataloader_num_workers 1
)
# Validation arguments
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_qat.sh
+125
View File
@@ -0,0 +1,125 @@
#!/usr/bin/env bash
set -euo pipefail
############################################
# Single node, 6 GPUs
############################################
NUM_GPUS=1
export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f
# export CUDA_VISIBLE_DEVICES=0,1,2,3
export TRITON_PRINT_AUTOTUNING=1 # to print the best config
export WANDB_MODE="online"
export WANDB_BASE_URL="https://api.wandb.ai"
export TOKENIZERS_PARALLELISM=false
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
export TRITON_CACHE_DIR=/tmp/triton_cache
echo "NUM_GPUS: $NUM_GPUS"
echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
############################################
# Configs
############################################
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"
OUTPUT_DIR="checkpoints/qat_distill"
REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
# Training arguments
training_args=(
--tracker_project_name wan_t2v_distill_dmd_qat
--output_dir "$OUTPUT_DIR"
--max_train_steps 4000
--train_batch_size 64
--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 (adjusted to 6 GPUs)
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"
--real_score_model_path "$REAL_SCORE_MODEL_PATH"
--fake_score_model_path "$FAKE_SCORE_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 100
--validation_sampling_steps "3"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 2e-6
--mixed_precision "bf16"
--training_state_checkpointing_steps 2000
--weight_only_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.0
--dit_precision "fp32"
--ema_start_step 0
--flow_shift 8
--seed 1000
)
# DMD arguments
dmd_args=(
--dmd_denoising_steps '1000,757,522'
--min_timestep_ratio 0.02
--max_timestep_ratio 0.98
--generator_update_interval 5
--real_score_guidance_scale 3.5
)
############################################
# Launch (no rendezvous flags needed)
############################################
set -x
torchrun \
--nproc_per_node="$NUM_GPUS" \
fastvideo/training/wan_distillation_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
"${dataset_args[@]}" \
"${training_args[@]}" \
"${optimizer_args[@]}" \
"${validation_args[@]}" \
"${miscellaneous_args[@]}" \
"${dmd_args[@]}"
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh
+8 -4
View File
@@ -4,7 +4,7 @@ set -euo pipefail
############################################
# Single node, 6 GPUs
############################################
NUM_GPUS=1
NUM_GPUS=4
export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f
# export CUDA_VISIBLE_DEVICES=0,1,2,3
@@ -24,16 +24,18 @@ echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
# Configs
############################################
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
DATA_DIR=data/Wan-Syn_77x448x832_600k
VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json"
OUTPUT_DIR="checkpoints/sage3_distill"
REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
# Training arguments
training_args=(
--tracker_project_name wan_t2v_distill_dmd_qat
--output_dir "$OUTPUT_DIR"
--max_train_steps 4000
--train_batch_size 1
--train_batch_size 16
--train_sp_batch_size 1
--gradient_accumulation_steps 1
--num_latent_t 12
@@ -47,7 +49,7 @@ training_args=(
# Parallel arguments (adjusted to 6 GPUs)
parallel_args=(
--num_gpus $NUM_GPUS
--sp_size 1
--sp_size 4
--tp_size 1
--hsdp_replicate_dim 1
--hsdp_shard_dim $NUM_GPUS
@@ -57,6 +59,8 @@ parallel_args=(
model_args=(
--model_path "$MODEL_PATH"
--pretrained_model_name_or_path "$MODEL_PATH"
--real_score_model_path "$REAL_SCORE_MODEL_PATH"
--fake_score_model_path "$FAKE_SCORE_MODEL_PATH"
)
# Dataset arguments
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
+12 -9
View File
@@ -5,25 +5,27 @@ 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/train/"
VALIDATION_DATASET_FILE="data/Wan-Syn_77x448x832_600k/validation_6.json"
NUM_GPUS=6
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"
--tracker_project_name "wan_t2v_finetune_qat"
--output_dir "checkpoints/wan_t2v_finetune_full_precision"
--max_train_steps 5000
--train_batch_size 1
--train_sp_batch_size 1
--gradient_accumulation_steps 1
--num_latent_t 6
--num_latent_t 12
--num_height 480
--num_width 832
--num_frames 21
--num_frames 45
--enable_gradient_checkpointing_type "full"
)
@@ -52,16 +54,17 @@ dataset_args=(
validation_args=(
--log_validation
--validation_dataset_file $VALIDATION_DATASET_FILE
--validation_steps 200
--validation_steps 50
--validation_sampling_steps "50"
--validation_guidance_scale "3.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 1e-6
--mixed_precision "bf16"
--checkpointing_steps 200
--weight_only_checkpointing_steps 200
--training_state_checkpointing_steps 200
--weight_decay 1e-4
--max_grad_norm 1.0
)
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat.sh
+2 -1
View File
@@ -64,7 +64,8 @@ validation_args=(
optimizer_args=(
--learning_rate 1e-6
--mixed_precision "bf16"
--checkpointing_steps 200
--weight_only_checkpointing_steps 200
--training_state_checkpointing_steps 200
--weight_decay 1e-4
--max_grad_norm 1.0
)
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh
+7 -6
View File
@@ -8,18 +8,18 @@ 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-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
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
# Training arguments
training_args=(
--tracker_project_name "wan_t2v_finetune_qat"
--output_dir "checkpoints/wan_t2v_finetune_qat"
--output_dir "checkpoints/wan_t2v_finetune_handwritten_sage3"
--max_train_steps 5000
--train_batch_size 1
--train_batch_size 16
--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 4
--tp_size 1
--hsdp_replicate_dim 1
--hsdp_shard_dim $NUM_GPUS
@@ -64,7 +64,8 @@ validation_args=(
optimizer_args=(
--learning_rate 1e-6
--mixed_precision "bf16"
--checkpointing_steps 200
--weight_only_checkpointing_steps 200
--training_state_checkpointing_steps 200
--weight_decay 1e-4
--max_grad_norm 1.0
)
examples/training/finetune/wan_t2v_14B/finetune_t2v_qat.sh
+5 -4
View File
@@ -17,9 +17,9 @@ NUM_GPUS=4
# Training arguments
training_args=(
--tracker_project_name "wan_t2v_finetune_qat"
--output_dir "checkpoints/wan_t2v_finetune_qat"
--output_dir "checkpoints/wan_t2v_finetune_qat_14B"
--max_train_steps 4000
--train_batch_size 1
--train_batch_size 16
--train_sp_batch_size 1
--gradient_accumulation_steps 1
--num_latent_t 21
@@ -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
@@ -64,7 +64,8 @@ validation_args=(
optimizer_args=(
--learning_rate 1e-5
--mixed_precision "bf16"
--checkpointing_steps 1000
--weight_only_checkpointing_steps 1000
--training_state_checkpointing_steps 1000
--weight_decay 0.01
--max_grad_norm 1.0
)
examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh
+3 -2
View File
@@ -32,7 +32,7 @@ export MASTER_ADDR MASTER_PORT
# Training arguments
training_args=(
--tracker_project_name "wan_t2v_finetune_qat"
--output_dir "checkpoints/wan_t2v_finetune_qat"
--output_dir "checkpoints/wan_t2v_finetune_sage3_14B"
--max_train_steps 4000
--train_batch_size 1
--train_sp_batch_size 1
@@ -79,7 +79,8 @@ validation_args=(
optimizer_args=(
--learning_rate 1e-5
--mixed_precision "bf16"
--checkpointing_steps 1000
--weight_only_checkpointing_steps 1000
--training_state_checkpointing_steps 1000
--weight_decay 0.01
--max_grad_norm 1.0
)
fastvideo/attention/backends/sage_attn3.py
+2 -2
View File
@@ -375,7 +375,7 @@ class SageAttention3Impl(AttentionImpl):
value: torch.Tensor,
attn_metadata: AttentionMetadata,
) -> torch.Tensor:
output = sageattn_blackwell_with_16bit_bwd(query, key, value, is_causal=self.causal)
# output = sageattn_blackwell_with_16bit_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
nvfp4_utils.py
+22 -8
View File
@@ -9,8 +9,7 @@ from triton.language.target_info import cuda_capability_geq
MXFP_BLOCK_SIZE = tl.constexpr(16)
@triton.jit
def _compute_quant_and_scale(src_tensor, valid_src_mask, mx_tensor_dtype: tl.constexpr):
is_fp8: tl.constexpr = mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5
def _compute_quant_and_scale(src_tensor, valid_src_mask, use_global_sf=True, two_level_quant_P=False):
BLOCK_SIZE_OUT_DIM: tl.constexpr = src_tensor.shape[0]
BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1]
BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // MXFP_BLOCK_SIZE
@@ -19,17 +18,32 @@ def _compute_quant_and_scale(src_tensor, valid_src_mask, mx_tensor_dtype: tl.con
f32_tensor = src_tensor.to(tl.float32)
abs_tensor = tl.abs(f32_tensor)
abs_tensor = tl.where(valid_src_mask, abs_tensor, -1.0) # Don't consider padding tensors in scale computation
abs_tensor = tl.reshape(abs_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
global_max_val = tl.max(abs_tensor)
# Avoid division by zero: if all values are padding (max is 0), use a default scale
global_max_val = tl.maximum(global_max_val, 1e-8)
s_enc = (6 * 448) / global_max_val
s_dec = (1 / s_enc)
if two_level_quant_P:
global_max_val = tl.max(f32_tensor, axis=1, keep_dims=True) # (BLOCK_SIZE_OUT_DIM, 1)
global_max_val = tl.maximum(global_max_val, 1e-8)
s_enc = ((6 * 448) / global_max_val).reshape([BLOCK_SIZE_OUT_DIM, 1, 1])
s_dec = (1 / s_enc)
abs_tensor = tl.reshape(abs_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
if use_global_sf and not two_level_quant_P:
global_max_val = tl.max(abs_tensor)
# Avoid division by zero: if all values are padding (max is 0), use a default scale
global_max_val = tl.maximum(global_max_val, 1e-8)
s_enc = (6 * 448) / global_max_val
s_dec = (1 / s_enc)
elif not two_level_quant_P and not use_global_sf:
s_dec = 1.0
s_enc = 1.0
max_val = tl.max(abs_tensor, axis=2, keep_dims=True) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) # per block maxima
s_dec_b = max_val / 6 # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
s_dec_b_e4m3 = (s_dec_b * s_enc).to(tl.float8e4nv) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
s_enc_b = 1 / (s_dec_b_e4m3.to(tl.float32) * s_dec) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
# # Handle 0 case to avoid Inf/NaN.
# s_dec_b_f32 = s_dec_b_e4m3.to(tl.float32)
# s_enc_b = tl.where(s_dec_b_f32 != 0, s_enc / s_dec_b_f32, 0.0)
f32_tensor = tl.reshape(f32_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
quant_tensor = f32_tensor * s_enc_b
qat_attn.py
+471 -537
View File
File diff suppressed because it is too large Load Diff
quant_utils.py
+6 -3
View File
@@ -5,10 +5,13 @@ from nvfp4_utils import _compute_quant_and_scale, _compute_dequant
@triton.jit
def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr,
BLOCK_SIZE_QUANT_DIM: tl.constexpr,
dst_dtype: tl.constexpr):
dst_dtype: tl.constexpr,
two_level_quant_P: tl.constexpr = False):
high_prec_src_tensor = src_tensor
src_tensor, src_scale, src_s_dec = _compute_quant_and_scale(src_tensor=src_tensor,
valid_src_mask=valid_src_mask)
valid_src_mask=valid_src_mask,
use_global_sf=True,
two_level_quant_P=two_level_quant_P)
src_tensor = _compute_dequant(mx_tensor=src_tensor,
scale=src_scale,
s_dec=src_s_dec,
@@ -68,4 +71,4 @@ def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv,
k, _ = fake_quantize(src_tensor=k_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=k_block.dtype)
v, _ = fake_quantize(src_tensor=v_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=v_block.dtype)
tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, k, mask=kv_valid[:, None])
tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, v, mask=kv_valid[:, None])
tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, v, mask=kv_valid[:, None])