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

[Refactor] SP Mask --> original seq len; HunyuanVideo 1.5 does not need mask (#1142)

Browse Source
This commit is contained in:
Zhang Peiyuan
2026-03-04 11:44:47 +08:00
committed by GitHub
parent e63b7d8ac4
commit cd1b7cf139
17 changed files with 154 additions and 360 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/inference/basic/basic_wan2_2.py
+1 -1
View File
@@ -11,7 +11,7 @@ def main():
generator = VideoGenerator.from_pretrained(
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
# FastVideo will automatically handle distributed setup
num_gpus=1,
num_gpus=2,
use_fsdp_inference=False, # set to True if GPU is out of memory
dit_cpu_offload=True, # DiT need to be offloaded for MoE
vae_cpu_offload=False,
fastvideo/attention/layer.py
+15 -26
View File
@@ -66,11 +66,11 @@ class DistributedAttention(nn.Module):
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
original_seq_len: int | None = None,
replicated_q: torch.Tensor | None = None,
replicated_k: torch.Tensor | None = None,
replicated_v: torch.Tensor | None = None,
freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Forward pass for distributed attention.
@@ -78,10 +78,10 @@ class DistributedAttention(nn.Module):
q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
original_seq_len (int): Original (unpadded) full sequence length
replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
replicated_k (Optional[torch.Tensor]): Replicated key tensor
replicated_v (Optional[torch.Tensor]): Replicated value tensor
attention_mask (Optional[torch.Tensor]): Attention mask [batch_size, seq_len]
Returns:
Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
@@ -91,7 +91,7 @@ class DistributedAttention(nn.Module):
# Check input shapes
assert q.dim() == 4 and k.dim() == 4 and v.dim(
) == 4, "Expected 4D tensors"
batch_size, seq_len, num_heads, head_dim = q.shape
batch_size, _, num_heads, _ = q.shape
local_rank = get_sp_parallel_rank()
world_size = get_sp_world_size()
@@ -107,15 +107,11 @@ class DistributedAttention(nn.Module):
scatter_dim=2,
gather_dim=1)
# After all-to-all, each rank has the full sequence but only a subset of heads
# The attention mask should now apply to the full sequence length
# Since mask is [batch, full_seq_len], it's already in the correct format
# LOAY TODO, instead of slicing repeatedly maintain an original qkv and rewrite into that
valid_seq_len = None
if attention_mask is not None:
valid_seq_len = (attention_mask[0] == 1).sum().item()
qkv = qkv[:, :valid_seq_len, :, :]
# After all-to-all, each rank has the full sequence but only a subset of heads.
# Trim away SP padding for attention compute, then pad back before returning.
original_seq_len = original_seq_len or qkv.shape[1]
pad_seq_len = qkv.shape[1] - original_seq_len
qkv = qkv[:, :original_seq_len, :, :]
if freqs_cis is not None:
cos, sin = freqs_cis
@@ -145,7 +141,7 @@ class DistributedAttention(nn.Module):
# Redistribute back if using sequence parallelism
replicated_output = None
if replicated_q is not None:
split_idx = seq_len * world_size if valid_seq_len is None else valid_seq_len
split_idx = original_seq_len
replicated_output = output[:, split_idx:]
output = output[:, :split_idx]
# TODO: make this asynchronous
@@ -154,9 +150,7 @@ class DistributedAttention(nn.Module):
# Apply backend-specific postprocess_output
output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
if attention_mask is not None:
pad_len = (attention_mask[0] == 0).sum().item()
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
output = sequence_model_parallel_all_to_all_4D(output,
scatter_dim=1,
@@ -175,12 +169,12 @@ class DistributedAttention_VSA(DistributedAttention):
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
original_seq_len: int,
replicated_q: torch.Tensor | None = None,
replicated_k: torch.Tensor | None = None,
replicated_v: torch.Tensor | None = None,
gate_compress: torch.Tensor | None = None,
freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Forward pass for distributed attention.
@@ -188,11 +182,11 @@ class DistributedAttention_VSA(DistributedAttention):
q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
original_seq_len (int): Original (unpadded) full sequence length
gate_compress (torch.Tensor): Gate compress tensor [batch_size, seq_len, num_heads, head_dim]
replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
replicated_k (Optional[torch.Tensor]): Replicated key tensor
replicated_v (Optional[torch.Tensor]): Replicated value tensor
attention_mask (Optional[torch.Tensor]): Attention mask [batch_size, seq_len]
Returns:
Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
@@ -221,11 +215,8 @@ class DistributedAttention_VSA(DistributedAttention):
gather_dim=1)
# After all-to-all, each rank has the full sequence but only a subset of heads
# The attention mask should now apply to the full sequence length
if attention_mask is not None:
valid_seq_len = (attention_mask[0] == 1).sum().item()
qkvg = qkvg[:, :valid_seq_len, :, :]
pad_seq_len = qkvg.shape[1] - original_seq_len
qkvg = qkvg[:, :original_seq_len, :, :]
if freqs_cis is not None:
cos, sin = freqs_cis
@@ -246,9 +237,7 @@ class DistributedAttention_VSA(DistributedAttention):
# Apply backend-specific postprocess_output
output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
if attention_mask is not None:
pad_len = (attention_mask[0] == 0).sum().item()
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
output = sequence_model_parallel_all_to_all_4D(output,
scatter_dim=1,
fastvideo/configs/models/encoders/qwen2_5.py
-1
View File
@@ -78,7 +78,6 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig):
"add_generation_prompt": True,
"tokenize": True,
"return_dict": True,
"padding": "max_length",
"max_length": 1000 + 108,
"truncation": True,
"return_tensors": "pt",
fastvideo/configs/models/encoders/t5.py
-1
View File
@@ -63,7 +63,6 @@ class T5ArchConfig(TextEncoderArchConfig):
self.dense_act_fn = "gelu_new"
self.tokenizer_kwargs = {
"padding": "max_length",
"truncation": True,
"max_length": self.text_len,
"add_special_tokens": True,
fastvideo/configs/sample/base.py
-1
View File
@@ -58,7 +58,6 @@ class SamplingParam:
# Original dimensions (before VAE scaling)
num_frames: int = 125
num_frames_round_down: bool = False # Whether to round down num_frames if it's not divisible by num_gpus
height: int = 720
width: int = 1280
height_sr: int = 1072
fastvideo/configs/sample/hunyuan15.py
-2
View File
@@ -15,8 +15,6 @@ class Hunyuan15_480P_SamplingParam(SamplingParam):
fps: int = 24
guidance_scale: float = 6.0
prompt_attention_mask: list = field(default_factory=list)
negative_attention_mask: list = field(default_factory=list)
sigmas: list[float] | None = field(
default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))
fastvideo/distributed/utils.py
-37
View File
@@ -81,43 +81,6 @@ def compute_padding_for_sp(seq_len: int, sp_world_size: int) -> tuple[int, int]:
return padded_seq_len, padding_amount
def create_attention_mask_for_padding(
seq_len: int,
padded_seq_len: int,
batch_size: int,
device: torch.device,
dtype: torch.dtype = torch.bool,
) -> torch.Tensor | None:
"""
Create attention mask to ignore padded tokens.
Args:
seq_len: Original sequence length (before padding)
padded_seq_len: Padded sequence length
batch_size: Batch size
device: Device to create mask on
dtype: Data type for the mask (default: bool)
Returns:
Tensor: Boolean mask [B, padded_seq_len] where True = valid token,
or None if no padding is needed
"""
if seq_len == padded_seq_len:
return None
# Create mask: True for valid tokens, False for padding
attention_mask = torch.ones(
(batch_size, padded_seq_len),
dtype=dtype,
device=device,
)
# Mask out padding tokens
attention_mask[:, seq_len:] = 0
return attention_mask
def pad_sequence_tensor(
tensor: torch.Tensor,
target_seq_len: int,
fastvideo/models/dits/hunyuanvideo.py
+6 -3
View File
@@ -196,6 +196,7 @@ class MMDoubleStreamBlock(nn.Module):
txt: torch.Tensor,
vec: torch.Tensor,
freqs_cis: tuple,
original_seq_len: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
# Process modulation vectors
img_mod_outputs = self.img_mod(vec)
@@ -253,7 +254,7 @@ class MMDoubleStreamBlock(nn.Module):
txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
# Run distributed attention
img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)
img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)
img_attn_out, _ = self.img_attn_proj(
img_attn.view(batch_size, image_seq_len, -1))
# Use fused operation for residual connection, normalization, and modulation
@@ -355,6 +356,7 @@ class MMSingleStreamBlock(nn.Module):
vec: torch.Tensor,
txt_len: int,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
original_seq_len: int | None = None,
) -> torch.Tensor:
# Process modulation
mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
@@ -386,7 +388,7 @@ class MMSingleStreamBlock(nn.Module):
# Run distributed attention
img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, txt_q,
img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, original_seq_len, txt_q,
txt_k, txt_v, freqs_cis = freqs_cis)
attn_output = torch.cat((img_attn_output, txt_attn_output),
dim=1).view(batch_size, seq_len, -1)
@@ -594,7 +596,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
# Process through double stream blocks
for index, block in enumerate(self.double_blocks):
double_block_args = [img, txt, vec, freqs_cis]
double_block_args = [img, txt, vec, freqs_cis, original_seq_len]
img, txt = block(*double_block_args)
# Merge txt and img to pass through single stream blocks
x = torch.cat((img, txt), 1)
@@ -607,6 +609,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
vec,
txt_seq_len,
freqs_cis,
original_seq_len,
]
x = block(*single_block_args)
fastvideo/models/dits/hunyuanvideo15.py
+21 -113
View File
@@ -16,9 +16,9 @@ from typing import Any, Dict, Optional, List
import torch
import torch.nn as nn
import torch.nn.functional as F
from fastvideo.attention import DistributedAttention, LocalAttention
from fastvideo.forward_context import set_forward_context
from fastvideo.distributed.communication_op import (
sequence_model_parallel_all_gather_with_unpad,
sequence_model_parallel_shard)
@@ -34,11 +34,8 @@ from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed,
from fastvideo.models.dits.base import BaseDiT
from fastvideo.platforms import AttentionBackendEnum
from fastvideo.logger import init_logger
from fastvideo.forward_context import set_forward_context
from fastvideo.attention.backends.abstract import AttentionMetadata
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.distributed.utils import create_attention_mask_for_padding
logger = init_logger(__name__)
@@ -290,10 +287,9 @@ class MMDoubleStreamBlock(nn.Module):
self,
img: torch.Tensor,
txt: torch.Tensor,
encoder_attention_mask: torch.Tensor,
vec: torch.Tensor,
freqs_cis: tuple,
seq_attention_mask: torch.Tensor,
original_seq_len: int,
) -> tuple[torch.Tensor, torch.Tensor]:
# Process modulation vectors
img_mod_outputs = self.img_mod(vec)
@@ -348,21 +344,8 @@ class MMDoubleStreamBlock(nn.Module):
txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
# seq_len = txt_q.shape[1] + img_q.shape[1]
# attention_mask = F.pad(encoder_attention_mask, (seq_len - encoder_attention_mask.shape[1], 0), value=True)
# attention_mask = attention_mask.bool()
# self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
# self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
# attention_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
attn_metadata = FlashAttnMetadataBuilder().build(
current_timestep=0,
attn_mask=encoder_attention_mask,
)
# Run distributed attention
with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis, attention_mask=seq_attention_mask)
img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)
img_attn_out, _ = self.img_attn_proj(
img_attn.view(batch_size, image_seq_len, -1))
@@ -469,15 +452,12 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
encoder_hidden_states: List[torch.Tensor],
timestep: torch.LongTensor,
encoder_hidden_states_image: List[torch.Tensor],
encoder_attention_mask: List[torch.Tensor],
guidance: Optional[torch.Tensor] = None,
timestep_r: Optional[torch.LongTensor] = None,
attention_kwargs: Optional[Dict[str, Any]] = None,
):
encoder_hidden_states_image = encoder_hidden_states_image[0]
encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
post_patch_num_frames = num_frames // p_t
@@ -503,18 +483,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
sp_world_size = get_sp_world_size()
padded_seq_len = current_seq_len * sp_world_size
if padded_seq_len > original_seq_len:
seq_attention_mask = create_attention_mask_for_padding(
seq_len=original_seq_len,
padded_seq_len=padded_seq_len,
batch_size=batch_size,
device=hidden_states.device,
)
else:
seq_attention_mask = None
# qwen text embedding
encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep, encoder_attention_mask)
encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep)
encoder_hidden_states_cond_emb = self.cond_type_embed(
torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)
@@ -532,19 +503,6 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
# image embed
encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
is_t2v = torch.all(encoder_hidden_states_image == 0)
if is_t2v:
encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0
encoder_attention_mask_3 = torch.zeros(
(batch_size, encoder_hidden_states_3.shape[1]),
dtype=encoder_attention_mask.dtype,
device=encoder_attention_mask.device,
)
else:
encoder_attention_mask_3 = torch.ones(
(batch_size, encoder_hidden_states_3.shape[1]),
dtype=encoder_attention_mask.dtype,
device=encoder_attention_mask.device,
)
encoder_hidden_states_3_cond_emb = self.cond_type_embed(
2
* torch.ones_like(
@@ -554,52 +512,7 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
)
encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb
# reorder and combine text tokens: combine valid tokens first, then padding
encoder_attention_mask = encoder_attention_mask.bool()
encoder_attention_mask_2 = encoder_attention_mask_2.bool()
encoder_attention_mask_3 = encoder_attention_mask_3.bool()
new_encoder_hidden_states = []
new_encoder_attention_mask = []
for text, text_mask, text_2, text_mask_2, image, image_mask in zip(
encoder_hidden_states,
encoder_attention_mask,
encoder_hidden_states_2,
encoder_attention_mask_2,
encoder_hidden_states_3,
encoder_attention_mask_3,
):
# Concatenate: [valid_image, valid_byt5, valid_mllm, invalid_image, invalid_byt5, invalid_mllm]
new_encoder_hidden_states.append(
torch.cat(
[
image[image_mask], # valid image
text_2[text_mask_2], # valid byt5
text[text_mask], # valid mllm
image[~image_mask], # invalid image (zeroed)
torch.zeros_like(text_2[~text_mask_2]), # invalid byt5 (zeroed)
torch.zeros_like(text[~text_mask]), # invalid mllm (zeroed)
],
dim=0,
)
)
# Apply same reordering to attention masks
new_encoder_attention_mask.append(
torch.cat(
[
image_mask[image_mask],
text_mask_2[text_mask_2],
text_mask[text_mask],
image_mask[~image_mask],
text_mask_2[~text_mask_2],
text_mask[~text_mask],
],
dim=0,
)
)
encoder_hidden_states = torch.stack(new_encoder_hidden_states)
encoder_attention_mask = torch.stack(new_encoder_attention_mask)
encoder_hidden_states = torch.cat([encoder_hidden_states_2, encoder_hidden_states], dim=1) if is_t2v else torch.cat([encoder_hidden_states_3, encoder_hidden_states_2, encoder_hidden_states], dim=1)
# 4. Transformer blocks
if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -608,10 +521,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
block,
hidden_states,
encoder_hidden_states,
encoder_attention_mask,
temb,
freqs_cis,
seq_attention_mask
original_seq_len
)
else:
@@ -619,10 +531,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
hidden_states, encoder_hidden_states = block(
hidden_states,
encoder_hidden_states,
encoder_attention_mask,
temb,
freqs_cis,
seq_attention_mask
original_seq_len
)
# Final layer processing
@@ -691,7 +602,6 @@ class SingleTokenRefiner(nn.Module):
# Get timestep embeddings
timestep_aware_representations = self.t_embedder(t)
# Get context-aware representations
original_dtype = x.dtype
if mask is None:
context_aware_representations = x.mean(dim=1)
@@ -706,7 +616,7 @@ class SingleTokenRefiner(nn.Module):
x = self.input_embedder(x)
# Process through refiner blocks
for block in self.refiner_blocks:
x = block(x, c, mask)
x = block(x, c, mask=mask)
return x
class IndividualTokenRefinerBlock(nn.Module):
@@ -776,9 +686,6 @@ class IndividualTokenRefinerBlock(nn.Module):
)
def forward(self, x, c, mask=None):
if mask is not None:
mask = mask.clone().bool()
mask[:, 0] = True # Prevent attention weights from becoming NaN
# Get modulation parameters
gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1)
@@ -789,18 +696,19 @@ class IndividualTokenRefinerBlock(nn.Module):
batch_size, seq_len = qkv.shape[0], qkv.shape[1]
qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
# Run scaled dot product attention
from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
attn_metadata = FlashAttnMetadataBuilder().build(
current_timestep=0,
attn_mask=mask,
)
# Run distributed attention
with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
attn_output = self.attn(q, k, v) # [B, L, H, D]
if mask is None:
attn_output = self.attn(q, k, v)
else:
from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
attn_metadata = FlashAttnMetadataBuilder().build(
current_timestep=0,
attn_mask=mask,
)
# Run distributed attention
with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
attn_output = self.attn(q, k, v) # [B, L, H, D]
attn_output = attn_output.reshape(batch_size, seq_len,
-1) # [B, L, H*D]
-1)
# Project and apply residual connection with gating
attn_out, _ = self.self_attn_proj(attn_output)
fastvideo/models/dits/hyworld/hyworld.py
+22 -21
View File
@@ -35,7 +35,6 @@ from fastvideo.platforms import AttentionBackendEnum
from fastvideo.logger import init_logger
from fastvideo.forward_context import set_forward_context
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.distributed.utils import create_attention_mask_for_padding
from .camera_rope import prope_qkv
@@ -88,7 +87,7 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
vec: torch.Tensor,
vec_txt: torch.Tensor,
freqs_cis: tuple,
seq_attention_mask: torch.Tensor,
original_seq_len: int,
viewmats: torch.Tensor,
Ks: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -101,7 +100,7 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
encoder_attention_mask: Text attention mask
vec: Modulation vector
freqs_cis: Rotary embedding frequencies
seq_attention_mask: Sequence attention mask
original_seq_len: Original (unpadded) image sequence length
viewmats: Camera view matrices for ProPE [B, T, 4, 4]
Ks: Camera intrinsics for ProPE [B, T, 3, 3]
@@ -181,7 +180,16 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
)
# Run distributed attention
with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis, attention_mask=seq_attention_mask)
img_attn, txt_attn = self.attn(
img_q,
img_k,
img_v,
original_seq_len,
txt_q,
txt_k,
txt_v,
freqs_cis=freqs_cis,
)
# begin hyworld
# attention with prope
@@ -193,8 +201,14 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
# NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope
with set_forward_context(current_timestep=0, attn_metadata=attn_metadata_prope):
img_attn_prope, _ = self.attn(
img_q_prope, img_k_prope, img_v_prope, txt_q, txt_k, txt_v,
freqs_cis=None, attention_mask=seq_attention_mask # No RoPE for prope attention
img_q_prope,
img_k_prope,
img_v_prope,
original_seq_len,
txt_q,
txt_k,
txt_v,
freqs_cis=None, # No RoPE for prope attention
)
img_attn_prope = img_attn_prope.reshape(batch_size, image_seq_len, -1)
img_attn_prope = rearrange(
@@ -399,20 +413,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
hidden_states = self.img_in(hidden_states)
hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
current_seq_len = hidden_states.shape[1]
sp_world_size = get_sp_world_size()
padded_seq_len = current_seq_len * sp_world_size
if padded_seq_len > original_seq_len:
seq_attention_mask = create_attention_mask_for_padding(
seq_len=original_seq_len,
padded_seq_len=padded_seq_len,
batch_size=batch_size,
device=hidden_states.device,
)
else:
seq_attention_mask = None
viewmats_seq = repeat(
viewmats, "B T M N->B (T H W) M N",
@@ -535,7 +536,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
temb,
temb_txt,
freqs_cis,
seq_attention_mask,
original_seq_len,
viewmats_seq,
Ks_seq,
)
@@ -548,7 +549,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
temb,
temb_txt,
freqs_cis,
seq_attention_mask,
original_seq_len,
viewmats_seq,
Ks_seq,
)
fastvideo/models/dits/lingbotworld/model.py
+10 -26
View File
@@ -30,7 +30,6 @@ from fastvideo.models.dits.wanvideo import (
from fastvideo.platforms import AttentionBackendEnum, current_platform
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.distributed.utils import create_attention_mask_for_padding
logger = init_logger(__name__)
@@ -145,7 +144,7 @@ class LingBotWorldTransformerBlock(nn.Module):
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
attention_mask: torch.Tensor | None = None,
original_seq_len: int,
c2ws_plucker_emb: torch.Tensor | None = None,
) -> torch.Tensor:
if hidden_states.dim() == 4:
@@ -189,7 +188,13 @@ class LingBotWorldTransformerBlock(nn.Module):
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
attn_output, _ = self.attn1(query, key, value, freqs_cis=freqs_cis, attention_mask=attention_mask)
attn_output, _ = self.attn1(
query,
key,
value,
original_seq_len,
freqs_cis=freqs_cis,
)
attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
attn_output = attn_output.squeeze(1)
@@ -292,7 +297,6 @@ class LingBotWorldTransformer3DModel(BaseDiT):
torch.randn(1, 2, inner_dim) / inner_dim**0.5)
self.gradient_checkpointing = False
self._logged_attention_mask = False
self.__post_init__()
def forward(self,
@@ -350,26 +354,6 @@ class LingBotWorldTransformer3DModel(BaseDiT):
if c2ws_plucker_emb is not None:
c2ws_plucker_emb, _ = sequence_model_parallel_shard(c2ws_plucker_emb, dim=1)
# Create attention mask for padded tokens if padding was applied
current_seq_len = hidden_states.shape[1]
sp_world_size = get_sp_world_size()
padded_seq_len = current_seq_len * sp_world_size
if padded_seq_len > original_seq_len:
if not self._logged_attention_mask:
logger.info(f"Padding applied, original seq len: {original_seq_len}, padded seq len: {padded_seq_len}")
self._logged_attention_mask = True
attention_mask = create_attention_mask_for_padding(
seq_len=original_seq_len,
padded_seq_len=padded_seq_len,
batch_size=batch_size,
device=hidden_states.device,
)
else:
if not self._logged_attention_mask:
logger.info(f"Padding not applied")
self._logged_attention_mask = True
attention_mask = None
# timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
if timestep.dim() == 2:
@@ -406,11 +390,11 @@ class LingBotWorldTransformer3DModel(BaseDiT):
for block in self.blocks:
hidden_states = self._gradient_checkpointing_func(
block, hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, attention_mask, c2ws_plucker_emb)
timestep_proj, freqs_cis, original_seq_len, c2ws_plucker_emb)
else:
for block in self.blocks:
hidden_states = block(hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, attention_mask,
timestep_proj, freqs_cis, original_seq_len,
c2ws_plucker_emb)
# 5. Output norm, projection & unpatchify
if temb.dim() == 3:
fastvideo/models/dits/ltx2.py
+53 -84
View File
@@ -27,7 +27,6 @@ from fastvideo.distributed.communication_op import (
sequence_model_parallel_shard,
)
from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
from fastvideo.distributed.utils import create_attention_mask_for_padding
from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context
from fastvideo.logger import init_logger
from fastvideo.models.dits.base import BaseDiT
@@ -1020,11 +1019,11 @@ class LTXDistributedAttention(DistributedAttention):
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
original_seq_len: int,
replicated_q: torch.Tensor | None = None,
replicated_k: torch.Tensor | None = None,
replicated_v: torch.Tensor | None = None,
gate_compress: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Forward pass with LTX-2 style RoPE application.
@@ -1036,14 +1035,14 @@ class LTXDistributedAttention(DistributedAttention):
replicated_q: Replicated query tensor for text tokens
replicated_k: Replicated key tensor
replicated_v: Replicated value tensor
attention_mask: Attention mask [batch_size, seq_len]
original_seq_len: Original (unpadded) full sequence length
ltx_freqs_cis: LTX-2 style RoPE (cos, sin) with shape [B, H, T, D]
Returns:
Tuple of output tensor and optional replicated output
"""
assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
batch_size, seq_len, num_heads, head_dim = q.shape
batch_size, _, num_heads, _ = q.shape
local_rank = get_sp_parallel_rank()
world_size = get_sp_world_size()
@@ -1067,19 +1066,16 @@ class LTXDistributedAttention(DistributedAttention):
qkvg = sequence_model_parallel_all_to_all_4D(qkvg,
scatter_dim=2,
gather_dim=1)
valid_seq_len = None
if attention_mask is not None:
valid_seq_len = (attention_mask[0] == 1).sum().item()
qkvg = qkvg[:, :valid_seq_len, :, :]
pad_seq_len = qkvg.shape[1] - original_seq_len
qkvg = qkvg[:, :original_seq_len, :, :]
if ltx_freqs_cis is not None:
cos, sin = ltx_freqs_cis
heads_per_rank = num_heads // world_size
head_start = local_rank * heads_per_rank
head_end = head_start + heads_per_rank
cos_local = cos[:, head_start:head_end, :valid_seq_len, :]
sin_local = sin[:, head_start:head_end, :valid_seq_len, :]
cos_local = cos[:, head_start:head_end, :original_seq_len, :]
sin_local = sin[:, head_start:head_end, :original_seq_len, :]
qk_part = qkvg[:batch_size * 2].transpose(1, 2)
qk_part = apply_ltx_rotary_emb_4d(
qk_part, (cos_local, sin_local), self.rope_type)
@@ -1090,10 +1086,7 @@ class LTXDistributedAttention(DistributedAttention):
output = self.attn_impl.forward( # type: ignore[call-arg]
q, k, v, gate_compress, ctx_attn_metadata)
output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
if attention_mask is not None:
pad_len = (attention_mask[0] == 0).sum().item()
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
output = sequence_model_parallel_all_to_all_4D(output,
scatter_dim=1,
@@ -1107,10 +1100,8 @@ class LTXDistributedAttention(DistributedAttention):
qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
# After all-to-all, each rank has the full sequence but only a subset of heads
valid_seq_len = None
if attention_mask is not None:
valid_seq_len = (attention_mask[0] == 1).sum().item()
qkv = qkv[:, :valid_seq_len, :, :]
pad_seq_len = qkv.shape[1] - original_seq_len
qkv = qkv[:, :original_seq_len, :, :]
# Apply LTX-2 style RoPE after all-to-all (when we have full sequence)
if ltx_freqs_cis is not None:
@@ -1119,8 +1110,8 @@ class LTXDistributedAttention(DistributedAttention):
head_start = local_rank * heads_per_rank
head_end = head_start + heads_per_rank
# Slice to this rank's heads: [B, H/SP, T, D]
cos_local = cos[:, head_start:head_end, :valid_seq_len, :]
sin_local = sin[:, head_start:head_end, :valid_seq_len, :]
cos_local = cos[:, head_start:head_end, :original_seq_len, :]
sin_local = sin[:, head_start:head_end, :original_seq_len, :]
# Apply RoPE to Q and K together (first 2*batch_size in dim 0)
qk_part = qkv[:batch_size * 2]
@@ -1152,7 +1143,7 @@ class LTXDistributedAttention(DistributedAttention):
# Redistribute back if using sequence parallelism
replicated_output = None
if replicated_q is not None:
split_idx = seq_len * world_size if valid_seq_len is None else valid_seq_len
split_idx = original_seq_len
replicated_output = output[:, split_idx:]
output = output[:, :split_idx]
replicated_output = sequence_model_parallel_all_gather(
@@ -1161,9 +1152,7 @@ class LTXDistributedAttention(DistributedAttention):
# Apply backend-specific postprocess_output
output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
if attention_mask is not None:
pad_len = (attention_mask[0] == 0).sum().item()
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)
@@ -1407,7 +1396,7 @@ class LTXDistributedSelfAttention(nn.Module):
context: torch.Tensor | None = None,
pe: tuple[torch.Tensor, torch.Tensor] | None = None,
k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
attention_mask: torch.Tensor | None = None,
original_seq_len: int | None = None,
) -> torch.Tensor:
"""Forward pass for distributed self-attention.
@@ -1417,7 +1406,7 @@ class LTXDistributedSelfAttention(nn.Module):
pe: Rotary position embeddings for Q (cos, sin) with shape [B, H, T_full, D]
NOTE: For SP, this must be the FULL sequence RoPE, not sharded!
k_pe: Rotary position embeddings for K (cos, sin), or None to use pe
attention_mask: Attention mask for padding [B, padded_seq_len]
original_seq_len: Original (unpadded) full sequence length
"""
q = self.to_q(x)
context = x if context is None else context
@@ -1442,10 +1431,14 @@ class LTXDistributedSelfAttention(nn.Module):
# Pass full RoPE to distributed attention - it will apply after all-to-all
# and slice to this rank's heads
if original_seq_len is None:
raise ValueError("original_seq_len must be provided for distributed attention")
out, _ = self.attn(
q, k, v,
q,
k,
v,
original_seq_len,
gate_compress=gate_compress,
attention_mask=attention_mask,
ltx_freqs_cis=pe,
)
@@ -1633,8 +1626,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
self,
video: TransformerArgs | None,
audio: TransformerArgs | None,
video_attention_mask: torch.Tensor | None = None,
audio_attention_mask: torch.Tensor | None = None,
video_original_seq_len: int | None = None,
audio_original_seq_len: int | None = None,
skip_cross_modal_attn: bool = False,
skip_video_self_attn: bool = False,
skip_audio_self_attn: bool = False,
@@ -1644,8 +1637,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
Args:
video: Video transformer args
audio: Audio transformer args
video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
video_original_seq_len: Original unpadded video sequence length
audio_original_seq_len: Original unpadded audio sequence length
skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
attention (used for the modality-isolated CFG pass).
skip_video_self_attn: If True, skip video self-attention
@@ -1667,9 +1660,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
)
if not skip_video_self_attn:
norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
# Self-attention: pass SP attention mask for distributed attention
if self.use_distributed_attention:
vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings, attention_mask=video_attention_mask) * vgate_msa
vx = vx + self.attn1(
norm_vx,
pe=video.positional_embeddings,
original_seq_len=video_original_seq_len,
) * vgate_msa
else:
vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa
# Text cross-attention: no SP mask needed (text is replicated, uses local attention)
@@ -1685,9 +1681,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
)
if not skip_audio_self_attn:
norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
# Self-attention: pass SP attention mask for distributed attention
if self.use_distributed_attention:
ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings, attention_mask=audio_attention_mask) * agate_msa
ax = ax + self.audio_attn1(
norm_ax,
pe=audio.positional_embeddings,
original_seq_len=audio_original_seq_len,
) * agate_msa
else:
ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa
# Text cross-attention: no SP mask needed (text is replicated, uses local attention)
@@ -2116,8 +2115,8 @@ class LTXModel(torch.nn.Module):
self,
video: TransformerArgs | None,
audio: TransformerArgs | None,
video_attention_mask: torch.Tensor | None = None,
audio_attention_mask: torch.Tensor | None = None,
video_original_seq_len: int | None = None,
audio_original_seq_len: int | None = None,
skip_cross_modal_attn: bool = False,
skip_video_self_attn_blocks: list[int] | None = None,
skip_audio_self_attn_blocks: list[int] | None = None,
@@ -2132,8 +2131,8 @@ class LTXModel(torch.nn.Module):
video, audio = block(
video=video,
audio=audio,
video_attention_mask=video_attention_mask,
audio_attention_mask=audio_attention_mask,
video_original_seq_len=video_original_seq_len,
audio_original_seq_len=audio_original_seq_len,
skip_cross_modal_attn=skip_cross_modal_attn,
skip_video_self_attn=skip_v_sa,
skip_audio_self_attn=skip_a_sa,
@@ -2161,8 +2160,8 @@ class LTXModel(torch.nn.Module):
self,
video: Modality | None,
audio: Modality | None,
video_attention_mask: torch.Tensor | None = None,
audio_attention_mask: torch.Tensor | None = None,
video_original_seq_len: int | None = None,
audio_original_seq_len: int | None = None,
skip_cross_modal_attn: bool = False,
skip_video_self_attn_blocks: list[int] | None = None,
skip_audio_self_attn_blocks: list[int] | None = None,
@@ -2172,8 +2171,8 @@ class LTXModel(torch.nn.Module):
Args:
video: Video modality input
audio: Audio modality input
video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
video_original_seq_len: Original unpadded video sequence length
audio_original_seq_len: Original unpadded audio sequence length
skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
attention in all transformer blocks (modality-isolated
CFG pass).
@@ -2202,8 +2201,8 @@ class LTXModel(torch.nn.Module):
video_out, audio_out = self._process_transformer_blocks(
video_args,
audio_args,
video_attention_mask=video_attention_mask,
audio_attention_mask=audio_attention_mask,
video_original_seq_len=video_original_seq_len,
audio_original_seq_len=audio_original_seq_len,
skip_cross_modal_attn=skip_cross_modal_attn,
skip_video_self_attn_blocks=skip_video_self_attn_blocks,
skip_audio_self_attn_blocks=skip_audio_self_attn_blocks,
@@ -2310,7 +2309,6 @@ class LTX2Transformer3DModel(BaseDiT):
self.hidden_size = arch.num_attention_heads * arch.attention_head_dim
self.num_attention_heads = arch.num_attention_heads
self.num_channels_latents = arch.num_channels_latents
self._logged_attention_mask = False
if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1":
detail_path = os.getenv("LTX2_PIPELINE_DEBUG_DETAIL_PATH", "")
@@ -2388,11 +2386,8 @@ class LTX2Transformer3DModel(BaseDiT):
) -> torch.Tensor:
if isinstance(encoder_hidden_states, list):
encoder_hidden_states = encoder_hidden_states[0]
# Get SP parameters
sp_world_size = get_sp_world_size()
sp_rank = get_sp_parallel_rank()
batch_size = hidden_states.shape[0]
# Get fps for position computation
fps = None
@@ -2413,29 +2408,14 @@ class LTX2Transformer3DModel(BaseDiT):
# Shard video latents and timestep across SP ranks
video_original_seq_len = latents.shape[1]
video_attention_mask = None
video_padded_seq_len = video_original_seq_len
video_timestep = timestep
if sp_world_size > 1:
latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1)
# Shard timestep along sequence dimension (timestep has shape [batch, seq_len])
video_timestep, _ = sequence_model_parallel_shard(timestep, dim=1)
# Create attention mask for padded tokens
current_seq_len = latents.shape[1]
padded_seq_len = current_seq_len * sp_world_size
if padded_seq_len > video_original_seq_len:
if not self._logged_attention_mask:
logger.info(
f"Video padding applied, original seq len: {video_original_seq_len}, "
f"padded seq len: {padded_seq_len}"
)
self._logged_attention_mask = True
video_attention_mask = create_attention_mask_for_padding(
seq_len=video_original_seq_len,
padded_seq_len=padded_seq_len,
batch_size=batch_size,
device=hidden_states.device,
)
video_padded_seq_len = current_seq_len * sp_world_size
# Compute RoPE positions for the FULL sequence (before sharding)
# This is necessary because sequence sharding may not align to frame boundaries.
# The full RoPE will be passed to DistributedAttention which applies it after
@@ -2451,8 +2431,8 @@ class LTX2Transformer3DModel(BaseDiT):
).to(hidden_states.dtype)
# Pad positions to match padded sequence length for SP cross-attention
if sp_world_size > 1 and padded_seq_len > video_original_seq_len:
padding_needed = padded_seq_len - video_original_seq_len
if sp_world_size > 1 and video_padded_seq_len > video_original_seq_len:
padding_needed = video_padded_seq_len - video_original_seq_len
last_pos = positions[:, :, -1:, :].expand(-1, -1, padding_needed, -1)
positions = torch.cat([positions, last_pos], dim=2)
@@ -2482,7 +2462,6 @@ class LTX2Transformer3DModel(BaseDiT):
audio_modality = None
audio_shape = None
audio_original_seq_len = 0
audio_attention_mask = None
if audio_hidden_states is not None and audio_encoder_hidden_states is not None and audio_timestep is not None:
audio_shape = AudioLatentShape.from_torch_shape(audio_hidden_states.shape)
@@ -2495,16 +2474,6 @@ class LTX2Transformer3DModel(BaseDiT):
audio_latents, audio_original_seq_len = sequence_model_parallel_shard(audio_latents, dim=1)
# Shard audio timestep along sequence dimension
sharded_audio_timestep, _ = sequence_model_parallel_shard(audio_timestep, dim=1)
# Create attention mask for padded tokens
audio_current_seq_len = audio_latents.shape[1]
audio_padded_seq_len = audio_current_seq_len * sp_world_size
if audio_padded_seq_len > audio_original_seq_len:
audio_attention_mask = create_attention_mask_for_padding(
seq_len=audio_original_seq_len,
padded_seq_len=audio_padded_seq_len,
batch_size=batch_size,
device=audio_hidden_states.device,
)
# Compute audio RoPE positions for the FULL sequence (before sharding)
# Same as video: full RoPE is applied after all-to-all in DistributedAttention
@@ -2534,12 +2503,12 @@ class LTX2Transformer3DModel(BaseDiT):
f"latent_checksum={audio_checksum:.6f}"
)
# Run transformer with attention masks
# Run transformer with original sequence lengths
video_out, audio_out = self.model(
video=video_modality,
audio=audio_modality,
video_attention_mask=video_attention_mask,
audio_attention_mask=audio_attention_mask,
video_original_seq_len=video_original_seq_len,
audio_original_seq_len=audio_original_seq_len,
skip_cross_modal_attn=skip_cross_modal_attn,
skip_video_self_attn_blocks=skip_video_self_attn_blocks,
skip_audio_self_attn_blocks=skip_audio_self_attn_blocks,
fastvideo/models/dits/wanvideo.py
+19 -29
View File
@@ -28,7 +28,6 @@ from fastvideo.models.dits.base import BaseDiT
from fastvideo.platforms import AttentionBackendEnum, current_platform
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.distributed.utils import create_attention_mask_for_padding
logger = init_logger(__name__)
@@ -324,7 +323,7 @@ class WanTransformerBlock(nn.Module):
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
attention_mask: torch.Tensor | None = None,
original_seq_len: int,
) -> torch.Tensor:
if hidden_states.dim() == 4:
hidden_states = hidden_states.squeeze(1)
@@ -367,7 +366,13 @@ class WanTransformerBlock(nn.Module):
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
attn_output, _ = self.attn1(query, key, value, freqs_cis=freqs_cis, attention_mask=attention_mask)
attn_output, _ = self.attn1(
query,
key,
value,
original_seq_len,
freqs_cis=freqs_cis,
)
attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
attn_output = attn_output.squeeze(1)
@@ -479,7 +484,7 @@ class WanTransformerBlock_VSA(nn.Module):
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
attention_mask: torch.Tensor | None = None,
original_seq_len: int,
) -> torch.Tensor:
if hidden_states.dim() == 4:
hidden_states = hidden_states.squeeze(1)
@@ -510,12 +515,14 @@ class WanTransformerBlock_VSA(nn.Module):
gate_compress = gate_compress.squeeze(1).unflatten(
2, (self.num_attention_heads, -1))
attn_output, _ = self.attn1(query,
key,
value,
freqs_cis = freqs_cis,
gate_compress=gate_compress,
attention_mask=attention_mask)
attn_output, _ = self.attn1(
query,
key,
value,
original_seq_len,
freqs_cis=freqs_cis,
gate_compress=gate_compress,
)
attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
attn_output = attn_output.squeeze(1)
@@ -610,7 +617,6 @@ class WanTransformer3DModel(BaseDiT):
torch.randn(1, 2, inner_dim) / inner_dim**0.5)
self.gradient_checkpointing = False
self._logged_attention_mask = False
self.__post_init__()
def forward(self,
@@ -656,26 +662,10 @@ class WanTransformer3DModel(BaseDiT):
# Shard with padding support - returns (sharded_tensor, original_seq_len)
hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
# Create attention mask for padded tokens if padding was applied
current_seq_len = hidden_states.shape[1]
sp_world_size = get_sp_world_size()
padded_seq_len = current_seq_len * sp_world_size
if padded_seq_len > original_seq_len:
if not self._logged_attention_mask:
logger.info(f"Padding applied, original seq len: {original_seq_len}, padded seq len: {padded_seq_len}")
self._logged_attention_mask = True
attention_mask = create_attention_mask_for_padding(
seq_len=original_seq_len,
padded_seq_len=padded_seq_len,
batch_size=batch_size,
device=hidden_states.device,
)
else:
if not self._logged_attention_mask:
logger.info(f"Padding not applied")
self._logged_attention_mask = True
attention_mask = None
# timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
if timestep.dim() == 2:
@@ -712,11 +702,11 @@ class WanTransformer3DModel(BaseDiT):
for block in self.blocks:
hidden_states = self._gradient_checkpointing_func(
block, hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, attention_mask)
timestep_proj, freqs_cis, original_seq_len)
else:
for block in self.blocks:
hidden_states = block(hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, attention_mask)
timestep_proj, freqs_cis, original_seq_len)
# 5. Output norm, projection & unpatchify
if temb.dim() == 3:
# batch_size, seq_len, inner_dim (wan 2.2 ti2v)
fastvideo/pipelines/pipeline_batch_info.py
-1
View File
@@ -150,7 +150,6 @@ class ForwardBatch:
height_latents: list[int] | int | None = None
width_latents: list[int] | int | None = None
num_frames: list[int] | int = 1 # Default for image models
num_frames_round_down: bool = False # Whether to round down num_frames if it's not divisible by num_gpus
# Original dimensions (before VAE scaling)
height: list[int] | int | None = None
fastvideo/pipelines/stages/base.py
-7
View File
@@ -203,10 +203,3 @@ class PipelineStage(ABC):
The updated batch information after this stage's processing.
"""
raise NotImplementedError
def backward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
raise NotImplementedError
fastvideo/pipelines/stages/text_encoding.py
+6 -7
View File
@@ -230,11 +230,12 @@ class TextEncodingStage(PipelineStage):
if processed_text is not None:
processed_texts.append(processed_text)
else:
# Assuming batch_size = 1
prompt_embeds = torch.zeros((1, tok_kwargs["max_length"],
encoder_config.hidden_size),
device=target_device)
attention_mask = torch.zeros((1, tok_kwargs["max_length"]),
# Assuming batch_size = 1, special case for hunyuanvideo1.5 where there is no glyph text
assert len(texts) == 1
prompt_embeds = torch.zeros(
(1, 0, encoder_config.hidden_size),
device=target_device)
attention_mask = torch.zeros((1, 0),
device=target_device,
dtype=torch.int64)
embeds_list.append(prompt_embeds)
@@ -274,13 +275,11 @@ class TextEncodingStage(PipelineStage):
if dtype is not None:
audio_embed = audio_embed.to(dtype=dtype)
audio_embeds_list.append(audio_embed)
if dtype is not None:
prompt_embeds = prompt_embeds.to(dtype=dtype)
embeds_list.append(prompt_embeds)
if return_attention_mask:
attn_masks_list.append(attention_mask)
self._last_audio_embeds = audio_embeds_list if is_ltx2 else None
return self.return_embeds(embeds_list, attn_masks_list, return_type,
return_attention_mask, indices)
fastvideo/training/training_pipeline.py
+1
View File
@@ -304,6 +304,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
# Make sure that the timesteps are the same across all sp processes.
sp_group = get_sp_group()
sp_group.broadcast(timesteps, src=0)
sp_group.broadcast(noise, src=0)
sigmas = get_sigmas(
self.noise_scheduler,
latents.device,