From 032016f7ddbfb2c4b0c20b901b377d28d9efff55 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Fri, 3 Oct 2025 16:15:00 +0000 Subject: [PATCH] update --- .../distill_dmd_t2v_1.3B_sage3.sh | 2 +- .../crush_smol/finetune_t2v_nvfp4.sh | 8 +- .../crush_smol/finetune_t2v_sage3.sh | 9 +- fastvideo/attention/backends/sage_attn3.py | 141 +++++++++++++----- .../nightly/test_sage3_with_16bit_bwd.py | 5 +- fastvideo/tests/nightly/test_torch_fp4_qk.py | 85 +++++++++++ fastvideo/training/training_pipeline.py | 33 ++-- fastvideo/training/training_utils.py | 11 +- 8 files changed, 219 insertions(+), 75 deletions(-) create mode 100644 fastvideo/tests/nightly/test_torch_fp4_qk.py diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh index 20a358931..b4da98b1c 100644 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh +++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh @@ -40,6 +40,7 @@ training_args=( --num_width 832 --num_frames 81 --enable_gradient_checkpointing_type "full" + --generator_4bit_attn True ) # Parallel arguments (adjusted to 6 GPUs) @@ -100,7 +101,6 @@ dmd_args=( --max_timestep_ratio 0.98 --generator_update_interval 5 --real_score_guidance_scale 3.5 - --generator_4bit_attn True ) ############################################ diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh index ecec54409..56c7da9fd 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh @@ -19,11 +19,11 @@ training_args=( --max_train_steps 5000 --train_batch_size 1 --train_sp_batch_size 1 - --gradient_accumulation_steps 4 - --num_latent_t 20 + --gradient_accumulation_steps 1 + --num_latent_t 6 --num_height 480 --num_width 832 - --num_frames 77 + --num_frames 21 --enable_gradient_checkpointing_type "full" --generator_4bit_linear True ) @@ -53,7 +53,7 @@ dataset_args=( validation_args=( --log_validation --validation_dataset_file $VALIDATION_DATASET_FILE - --validation_steps 10 + --validation_steps 50 --validation_sampling_steps "50" --validation_guidance_scale "6.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 bb8012c33..fd11bb370 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 @@ -19,14 +19,13 @@ training_args=( --max_train_steps 5000 --train_batch_size 1 --train_sp_batch_size 1 - --gradient_accumulation_steps 2 - --num_latent_t 20 + --gradient_accumulation_steps 1 + --num_latent_t 6 --num_height 480 --num_width 832 - --num_frames 77 + --num_frames 21 --enable_gradient_checkpointing_type "full" --generator_4bit_attn True - ) # Parallel arguments @@ -54,7 +53,7 @@ dataset_args=( validation_args=( --log_validation --validation_dataset_file $VALIDATION_DATASET_FILE - --validation_steps 100 + --validation_steps 50 --validation_sampling_steps "50" --validation_guidance_scale "6.0" ) diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index 910ffac87..186c1e76b 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/fastvideo/attention/backends/sage_attn3.py @@ -15,11 +15,108 @@ from math import sqrt from flash_attn.flash_attn_interface import _wrapped_flash_attn_forward, _wrapped_flash_attn_backward logger = init_logger(__name__) +from flashinfer import SfLayout, mm_fp4, nvfp4_quantize + +def fp_16_qk(q, k): + qk = torch.matmul(q, k) / torch.sqrt(torch.tensor(q.shape[-1], device=q.device, dtype=q.dtype)) + p = torch.softmax(qk, dim=-1) + return p + +def sage_attn_qk_torch(q, k, per_block_mean=False): + """ + Inputs: q/k/v in [B, H, L, D] + Returns: out in [B, H, L, D] + """ + B, H, qL, D = q.shape + B, H, kL, D = k.shape + assert B == 1 + q, k = q.squeeze(0), k.squeeze(0) + k = k - k.mean(dim=-2, keepdim=True) # [H, L, D] + qm = q.mean(dim=-2, keepdim=True) # [H, 1, D] + q = q - qm # [B, H, L, D] + delta_s = torch.matmul(qm, k.transpose(-2, -1)).to(torch.float32).contiguous() # [H, 1, L] + + alpha = torch.ones(1, device=q.device, dtype=torch.float32) + qk = torch.empty((H, qL, kL), device=q.device, dtype=q.dtype) + for h in range(H): + q_fp4, q_inv_s = nvfp4_quantize(q[h], torch.ones(1, device=q.device, dtype=torch.float32), sfLayout=SfLayout.layout_128x4, do_shuffle=False) + k_fp4, k_inv_s = nvfp4_quantize(k[h], torch.ones(1, device=q.device, dtype=torch.float32), sfLayout=SfLayout.layout_128x4, do_shuffle=False) + out = torch.empty((qL, kL), device=q.device, dtype=q.dtype) + mm_fp4( + q_fp4, k_fp4.T, q_inv_s, k_inv_s.T, alpha, + q.dtype, out, + block_size=16, + use_8x4_sf_layout=False, + backend="cutlass", + ) + qk[h] = out + qk = qk + delta_s + qk = qk / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) + p = torch.softmax(qk, dim=-1) + return p.unsqueeze(0) + +def sage3_torch_16bit_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal=False): + assert is_causal == False + q, k, v, o, do = q_BLHD.transpose(1 ,2).contiguous(), k_BLHD.transpose(1 ,2).contiguous(), v_BLHD.transpose(1 ,2).contiguous(), out_BLHD.transpose(1 ,2).contiguous(), grad_out_BLHD.transpose(1 ,2).contiguous() + B, H, _, D = q.shape + p = sage_attn_qk_torch(q, k, per_block_mean=False).to(do.dtype) + dV = p.transpose(-2, -1) @ do # [B,H,L,D] + + dP = do @ v.transpose(-2, -1) # [B,H,L,L] + + do_dot_o = (do.float() * o.float()).sum(dim=-1, keepdim=True).to(dP.dtype) # [B,H,L,1] + dS = p * (dP - do_dot_o) # [B,H,L,L] + dQ = (dS @ k) / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) # [B,H,L,D] + dK = (dS.transpose(-2, -1) @ q) / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) # [B,H,L,D] + + # grads w.r.t inputs, None for non-tensor inputs + return dQ.transpose(1, 2).contiguous(), dK.transpose(1, 2).contiguous(), dV.transpose(1, 2).contiguous(), None, None, None + + +def flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal=False): + D = q_BLHD.shape[-1] + softmax_scale = 1.0 / sqrt(D) + + # FA forward only to get softmax_lse (rng_state is also returned) + _, softmax_lse, _S_dmask, rng_state = _wrapped_flash_attn_forward( + q_BLHD, k_BLHD, v_BLHD, + 0.0, # dropout_p + softmax_scale, # softmax_scale + is_causal, # causal + -1, -1, # window_size_left/right + 0.0, # softcap + None, # alibi_slopes + False # return_softmax + ) + + # Allocate grads and call FA backward using the SAGE output + dq_BLHD = torch.empty_like(q_BLHD) + dk_BLHD = torch.empty_like(k_BLHD) + dv_BLHD = torch.empty_like(v_BLHD) + + _wrapped_flash_attn_backward( + grad_out_BLHD, # dout (BLHD, 16-bit) + q_BLHD, k_BLHD, v_BLHD, + out_BLHD, # use SAGE forward output here + softmax_lse, + dq_BLHD, dk_BLHD, dv_BLHD, + 0.0, # dropout_p + softmax_scale, # softmax_scale + is_causal, # causal + -1, -1, # window_size_left/right + 0.0, # softcap + None, # alibi_slopes + False, # deterministic + rng_state=rng_state, + ) + return dq_BLHD, dk_BLHD, dv_BLHD, None, None + +USE_TORCH_16BIT_BWD = True class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function): @staticmethod - def forward(ctx, q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True): + def forward(ctx, q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=False): """ Inputs: q/k/v in [B, L, H, D] Returns: out in [B, L, H, D] @@ -51,45 +148,13 @@ class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function): q_BLHD, k_BLHD, v_BLHD, out_BLHD = ctx.saved_tensors is_causal = ctx.is_causal - D = q_BLHD.shape[-1] - softmax_scale = 1.0 / sqrt(D) - - # FA forward only to get softmax_lse (rng_state is also returned) - _, softmax_lse, _S_dmask, rng_state = _wrapped_flash_attn_forward( - q_BLHD, k_BLHD, v_BLHD, - 0.0, # dropout_p - softmax_scale, # softmax_scale - is_causal, # causal - -1, -1, # window_size_left/right - 0.0, # softcap - None, # alibi_slopes - False # return_softmax - ) - - # Allocate grads and call FA backward using the SAGE output - dq_BLHD = torch.empty_like(q_BLHD) - dk_BLHD = torch.empty_like(k_BLHD) - dv_BLHD = torch.empty_like(v_BLHD) - - _wrapped_flash_attn_backward( - grad_out_BLHD, # dout (BLHD, 16-bit) - q_BLHD, k_BLHD, v_BLHD, - out_BLHD, # use SAGE forward output here - softmax_lse, - dq_BLHD, dk_BLHD, dv_BLHD, - 0.0, # dropout_p - softmax_scale, # softmax_scale - is_causal, # causal - -1, -1, # window_size_left/right - 0.0, # softcap - None, # alibi_slopes - False, # deterministic - rng_state=rng_state, - ) - return dq_BLHD, dk_BLHD, dv_BLHD, None, None + if USE_TORCH_16BIT_BWD: + return sage3_torch_16bit_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal) + else: + return flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal) -def sageattn_blackwell_with_16bit_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True): +def sageattn_blackwell_with_16bit_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=False): """ Forward: uses sageattn_blackwell under the hood. Backward: recomputes FlashAttention fwd+bwd in 16bit directly. diff --git a/fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py b/fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py index 0975c1c12..129846a1a 100644 --- a/fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py +++ b/fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py @@ -117,13 +117,12 @@ def run_case(case, fwd_cos_min=0.95, fwd_rel_l2_max=0.25, bwd_cos_min=0.95, bwd_ metric_report(dq_ref, dq_sage, "grad_q", cos_min=bwd_cos_min, rel_l2_max=bwd_rel_l2_max) metric_report(dk_ref, dk_sage, "grad_k", cos_min=bwd_cos_min, rel_l2_max=bwd_rel_l2_max) metric_report(dv_ref, dv_sage, "grad_v", cos_min=bwd_cos_min, rel_l2_max=bwd_rel_l2_max) - import pdb; pdb.set_trace() def main(): # Keep sizes moderate so it runs quickly; adjust as you like. cases = [ - {"name": "bf16_noncausal_L16384", "B": 1, "L": 16384, "H": 8, "D": 128, "dtype": torch.bfloat16, "causal": False}, - {"name": "bf16_noncausal_L32768", "B": 1, "L": 32768, "H": 8, "D": 128, "dtype": torch.bfloat16, "causal": False}, + {"name": "bf16_noncausal_L16384", "B": 1, "L": 1024, "H": 8, "D": 128, "dtype": torch.bfloat16, "causal": False}, + {"name": "bf16_noncausal_L2048", "B": 1, "L": 2048, "H": 8, "D": 128, "dtype": torch.bfloat16, "causal": False}, ] for case in cases: diff --git a/fastvideo/tests/nightly/test_torch_fp4_qk.py b/fastvideo/tests/nightly/test_torch_fp4_qk.py new file mode 100644 index 000000000..3493c3dc0 --- /dev/null +++ b/fastvideo/tests/nightly/test_torch_fp4_qk.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +import math +import sys +import torch +import torch.nn.functional as F +from fastvideo.attention.backends.sage_attn3 import fp_16_qk, sage_attn_qk_torch + + +def fail(msg: str): + print(f"[FAIL] {msg}") + sys.exit(1) + + +def metric_report(ref, test, name, cos_min=None, rel_l2_max=None): + diff = (test - ref).float() + ref32 = ref.float() + test32 = test.float() + + mean_abs_diff = diff.abs().mean().item() + max_abs_diff = diff.abs().max().item() + rel_l2 = (diff.norm() / (ref32.norm() + 1e-12)).item() + cos = F.cosine_similarity(ref32.reshape(-1), test32.reshape(-1), dim=0).item() + + print(f" [{name}] cos: {cos:.6f}, rel_l2: {rel_l2:.6e}, " + f"mean_abs: {mean_abs_diff:.6e}, max_abs: {max_abs_diff:.6e}") + + if cos_min is not None and cos < cos_min: + fail(f"{name}: cosine similarity too low: {cos:.6f} < {cos_min}") + if rel_l2_max is not None and rel_l2 > rel_l2_max: + fail(f"{name}: relative L2 too high: {rel_l2:.6e} > {rel_l2_max}") + + +def run_case(case, fwd_cos_min=0.98, fwd_rel_l2_max=2.0e-1): + device = torch.device("cuda") + torch.manual_seed(0) + + B, H, L, D = case["B"], case["H"], case["L"], case["D"] + NAME = case["name"] + assert B == 1, "sage_attn_qk_torch requires B == 1" + + print(f"\n== Case: {NAME} | B={B}, H={H}, L={L}, D={D} ==") + + dtype = torch.bfloat16 + q = torch.randn(B, H, L, D, device=device, dtype=dtype) + k = torch.randn(B, H, L, D, device=device, dtype=dtype) + + with torch.inference_mode(): + # Reference computes q @ k^T / sqrt(D) + p_ref = fp_16_qk(q, k.transpose(-1, -2)) # [B,H,L,L] + p_sage = sage_attn_qk_torch(q, k, False) # [B,H,L,L] + print("Forward metrics (sage_attn FP4 vs fp16/bf16 reference):") + metric_report(p_ref, p_sage, "forward", fwd_cos_min, fwd_rel_l2_max) + + +def main(): + try: + _ = fp_16_qk + _ = sage_attn_qk_torch + except NameError: + print("Please import/define fp_16_qk and sage_attn_qk_torch before running.") + sys.exit(1) + + if not torch.cuda.is_available(): + print("CUDA is required.") + sys.exit(1) + + configs = [ + {"B": 1, "H": 4, "L": 256, "D": 128, "name": "H4_L256_D128"}, + {"B": 1, "H": 8, "L": 512, "D": 128, "name": "H8_L512_D128"}, + {"B": 1, "H": 8, "L": 256, "D": 256, "name": "H8_L256_D256"}, + {"B": 1, "H": 16, "L": 384, "D": 128, "name": "H16_L384_D128"}, + ] + + fwd_cos_min = 0.98 + fwd_rel_l2_max = 2.0e-1 # relax for FP4 probabilities + + for case in configs: + run_case(case, fwd_cos_min=fwd_cos_min, fwd_rel_l2_max=fwd_rel_l2_max) + torch.cuda.synchronize() + + print("\n[OK] All forward diff checks passed for all cases.") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 43f9d557a..e0d7f3899 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -49,10 +49,9 @@ from fastvideo.training.training_utils import ( from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict) from fastvideo.training.training_utils import traverse_swap_module -try: - from fastvideo.attention.backends.sage_attn3 import SageAttention3Impl -except ImportError: - SageAttention3Impl = None + +from fastvideo.attention.backends.sage_attn3 import SageAttention3Impl + import wandb # isort: skip @@ -86,20 +85,18 @@ def swap_sage_attn3(obj: Any, obj_path: str) -> int: causal = getattr(old_impl, "causal", False) softmax_scale = getattr(old_impl, "softmax_scale", 1.0) - try: - new_impl = SageAttention3Impl( - num_heads=None, - head_size=None, - causal=causal, - softmax_scale=softmax_scale, - num_kv_heads=None, - prefix="", - ) - setattr(obj, "attn_impl", new_impl) - return 1 - except Exception: - # If assignment fails (frozen dataclass/property), skip - return 0 + + new_impl = SageAttention3Impl( + num_heads=None, + head_size=None, + causal=causal, + softmax_scale=softmax_scale, + num_kv_heads=None, + prefix="", + ) + setattr(obj, "attn_impl", new_impl) + return 1 + def swap_fp4_linear(obj: Any, obj_path: str) -> int: """ diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index c493a3c07..f1cd544a3 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -1840,12 +1840,11 @@ def traverse_swap_module( # Then, traverse children # 1) PyTorch modules: prefer named_children() for proper module graph traversal if nn is not None and isinstance(obj, nn.Module): - try: - for name, child in obj.named_children(): - if is_traversable(child): - swaps += walk(child, f"{obj_path}.{name}", depth + 1) - except Exception: - pass # keep going + + for name, child in obj.named_children(): + if is_traversable(child): + swaps += walk(child, f"{obj_path}.{name}", depth + 1) + # Also traverse non-module attributes living on the module d = getattr(obj, "__dict__", None)