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

update

Browse Source
This commit is contained in:
Peiyuan Zhang
2025-12-23 01:28:08 +00:00
committed by RandNMR73
parent 2a28c081be
commit 032016f7dd
8 changed files with 219 additions and 75 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_sage3.sh
+1 -1
View File
@@ -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
)
############################################
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh
+4 -4
View File
@@ -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"
)
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh
+4 -5
View File
@@ -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"
)
fastvideo/attention/backends/sage_attn3.py
+103 -38
View File
@@ -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.
fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py
+2 -3
View File
@@ -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:
fastvideo/tests/nightly/test_torch_fp4_qk.py
+85
View File
@@ -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()
fastvideo/training/training_pipeline.py
+15 -18
View File
@@ -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:
"""
fastvideo/training/training_utils.py
+5 -6
View File
@@ -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)