diff --git a/fastvideo/attention/distributed_attn.py b/fastvideo/attention/distributed_attn.py index 3b92668dc..169578e1a 100644 --- a/fastvideo/attention/distributed_attn.py +++ b/fastvideo/attention/distributed_attn.py @@ -94,6 +94,7 @@ class DistributedAttention(nn.Module): if replicated_q is not None: replicated_output = output[:, seq_len*world_size:] output = output[:, :seq_len*world_size] + # TODO: make this asynchronous replicated_output = sequence_model_parallel_all_gather(replicated_output, dim=2) output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2) return output, replicated_output diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index 0fb82d9d4..db2e7d4b6 100644 --- a/fastvideo/distributed/communication_op.py +++ b/fastvideo/distributed/communication_op.py @@ -35,6 +35,7 @@ def broadcast_tensor_dict(tensor_dict: Optional[Dict[Any, Union[torch.Tensor, return get_tp_group().broadcast_tensor_dict(tensor_dict, src) +# TODO: remove model, make it sequence_parallel def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1) -> torch.Tensor: diff --git a/fastvideo/distributed/device_communicators/base_device_communicator.py b/fastvideo/distributed/device_communicators/base_device_communicator.py index aa7e20ca7..22b2e6000 100644 --- a/fastvideo/distributed/device_communicators/base_device_communicator.py +++ b/fastvideo/distributed/device_communicators/base_device_communicator.py @@ -6,7 +6,7 @@ from typing import Optional import torch import torch.distributed as dist from torch.distributed import ProcessGroup - +from einops import rearrange class DeviceCommunicatorBase: """ @@ -121,20 +121,16 @@ class DeviceCommunicatorBase: shard_hc = hc // self.world_size # Reshape and transpose for scattering - input_t = (input_.reshape(bs, shard_seqlen, self.world_size, shard_hc, hs) - .transpose(0, 2) - .contiguous()) + input_t = rearrange(input_, 'b s (w h) d -> w (b s h) d', w=self.world_size, h=shard_hc) output = torch.empty_like(input_t) - if self.world_size > 1: - torch.distributed.all_to_all_single(output, input_t, group=self.device_group) - else: - output = input_t + + torch.distributed.all_to_all_single(output, input_t, group=self.device_group) + # Reshape and transpose back - output = output.reshape(seqlen, bs, shard_hc, hs) - output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + output = rearrange(output, '(s) b h d -> b s h d', s=seqlen, b=bs, h=shard_hc) return output @@ -145,22 +141,16 @@ class DeviceCommunicatorBase: shard_seqlen = seqlen // self.world_size # Reshape and transpose for scattering - input_t = (input_.reshape(bs, self.world_size, shard_seqlen, shard_hc, hs) - .transpose(0, 3) - .transpose(0, 1) - .contiguous() - .reshape(self.world_size, shard_hc, shard_seqlen, bs, hs)) + input_t = rearrange(input_, 'b (w s) h d -> w h s b d', w=self.world_size, s=shard_seqlen) output = torch.empty_like(input_t) - if self.world_size > 1: - torch.distributed.all_to_all_single(output, input_t, group=self.device_group) - else: - output = input_t + + torch.distributed.all_to_all_single(output, input_t, group=self.device_group) + # Reshape and transpose back - output = output.reshape(hc, shard_seqlen, bs, hs) - output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + output = rearrange(output, 'w h s b d -> b (w s) h d', w=self.world_size, s=shard_seqlen) return output else: diff --git a/fastvideo/layers/mlp.py b/fastvideo/layers/mlp.py index ed17b3b27..db2d2ce01 100644 --- a/fastvideo/layers/mlp.py +++ b/fastvideo/layers/mlp.py @@ -6,7 +6,10 @@ from typing import Optional class MLP(nn.Module): - """MLP for DiT blocks, NO gated linear units""" + """ + MLP for DiT blocks, NO gated linear units + TODO: add Tensor Parallel + """ def __init__( self, diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/models/dits/hunyuanvideo.py index 48c667166..0a6f1334f 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/models/dits/hunyuanvideo.py @@ -24,16 +24,16 @@ class MMDoubleStreamBlock(nn.Module): def __init__( self, hidden_size: int, - heads_num: int, - mlp_width_ratio: float, + num_attention_heads: int, + mlp_ratio: float, dtype: Optional[torch.dtype] = None, ): super().__init__() self.deterministic = False - self.heads_num = heads_num - head_dim = hidden_size // heads_num - mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.num_attention_heads = num_attention_heads + head_dim = hidden_size // num_attention_heads + mlp_hidden_dim = int(hidden_size * mlp_ratio) # Image modulation components self.img_mod = ModulateProjection( @@ -155,7 +155,7 @@ class MMDoubleStreamBlock(nn.Module): batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] # Split QKV - img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.heads_num, -1) + img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1) img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] # Apply QK-Norm if needed @@ -173,7 +173,7 @@ class MMDoubleStreamBlock(nn.Module): batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] # Split QKV - txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.heads_num, -1) + txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1) txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] # Apply QK-Norm if needed @@ -217,17 +217,17 @@ class MMSingleStreamBlock(nn.Module): def __init__( self, hidden_size: int, - heads_num: int, - mlp_width_ratio: float = 4.0, + num_attention_heads: int, + mlp_ratio: float = 4.0, dtype: Optional[torch.dtype] = None, ): super().__init__() self.deterministic = False self.hidden_size = hidden_size - self.heads_num = heads_num - head_dim = hidden_size // heads_num - mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.num_attention_heads = num_attention_heads + head_dim = hidden_size // num_attention_heads + mlp_hidden_dim = int(hidden_size * mlp_ratio) self.mlp_hidden_dim = mlp_hidden_dim # Combined QKV and MLP input projection @@ -296,7 +296,7 @@ class MMSingleStreamBlock(nn.Module): # Process QKV batch_size, seq_len = qkv.shape[0], qkv.shape[1] - qkv = qkv.view(batch_size, seq_len, 3, self.heads_num, -1) + qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] # Apply QK-Norm @@ -342,45 +342,46 @@ class HunyuanVideoDiT(nn.Module): - Flux.1: https://github.com/black-forest-labs/flux - MMDiT: http://arxiv.org/abs/2403.03206 """ - + # PY: we make the input args the same as HF config def __init__( self, - patch_size: list = [1, 2, 2], - in_channels: int = 4, - out_channels: int = None, + patch_size: int = 2, + patch_size_t: int = 1, + in_channels: int = 16, + out_channels: int = 16, hidden_size: int = 3072, - heads_num: int = 24, - mlp_width_ratio: float = 4.0, - mm_double_blocks_depth: int = 20, - mm_single_blocks_depth: int = 40, - rope_dim_list: List[int] = [16, 56, 56], - guidance_embed: bool = False, + num_attention_heads: int = 24, + mlp_ratio: float = 4.0, + num_layers: int = 20, + num_single_layers: int = 40, + rope_axes_dim: List[int] = [16, 56, 56], + guidance_embeds: bool = False, dtype: Optional[torch.dtype] = None, - text_states_dim: int = 4096, - text_states_dim_2: int = 768, + text_embed_dim: int = 4096, + pooled_projection_dim: int = 768, rope_theta: int = 256, ): super().__init__() - self.patch_size = patch_size + self.patch_size = [patch_size_t, patch_size, patch_size] self.in_channels = in_channels self.out_channels = in_channels if out_channels is None else out_channels self.unpatchify_channels = self.out_channels - self.guidance_embed = guidance_embed - self.rope_dim_list = rope_dim_list + self.guidance_embeds = guidance_embeds + self.rope_dim_list = rope_axes_dim self.rope_theta = rope_theta - self.text_states_dim = text_states_dim - self.text_states_dim_2 = text_states_dim_2 + self.text_states_dim = text_embed_dim + self.text_states_dim_2 = pooled_projection_dim - if hidden_size % heads_num != 0: - raise ValueError(f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}") + if hidden_size % num_attention_heads != 0: + raise ValueError(f"Hidden size {hidden_size} must be divisible by num_attention_heads {num_attention_heads}") - pe_dim = hidden_size // heads_num - if sum(rope_dim_list) != pe_dim: - raise ValueError(f"Got {rope_dim_list} but expected positional dim {pe_dim}") + pe_dim = hidden_size // num_attention_heads + if sum(rope_axes_dim) != pe_dim: + raise ValueError(f"Got {rope_axes_dim} but expected positional dim {pe_dim}") self.hidden_size = hidden_size - self.heads_num = heads_num + self.num_attention_heads = num_attention_heads # Image projection self.img_in = PatchEmbed( @@ -394,7 +395,7 @@ class HunyuanVideoDiT(nn.Module): self.txt_in = SingleTokenRefiner( self.text_states_dim, hidden_size, - heads_num, + num_attention_heads, depth=2, dtype=dtype ) @@ -419,27 +420,27 @@ class HunyuanVideoDiT(nn.Module): # Guidance modulation self.guidance_in = ( TimestepEmbedder(self.hidden_size, act_layer="silu", dtype=dtype) - if guidance_embed else None + if guidance_embeds else None ) # Double blocks self.double_blocks = nn.ModuleList([ MMDoubleStreamBlock( hidden_size, - heads_num, - mlp_width_ratio=mlp_width_ratio, + num_attention_heads, + mlp_ratio=mlp_ratio, dtype=dtype, - ) for _ in range(mm_double_blocks_depth) + ) for _ in range(num_layers) ]) # Single blocks self.single_blocks = nn.ModuleList([ MMSingleStreamBlock( hidden_size, - heads_num, - mlp_width_ratio=mlp_width_ratio, + num_attention_heads, + mlp_ratio=mlp_ratio, dtype=dtype, - ) for _ in range(mm_single_blocks_depth) + ) for _ in range(num_single_layers) ]) self.final_layer = FinalLayer( @@ -489,7 +490,7 @@ class HunyuanVideoDiT(nn.Module): ) # Get rotary embeddings - freqs_cos, freqs_sin = get_rotary_pos_embed((tt * get_sequence_model_parallel_world_size(), th, tw), self.hidden_size, self.heads_num, self.rope_dim_list, self.rope_theta) + freqs_cos, freqs_sin = get_rotary_pos_embed((tt * get_sequence_model_parallel_world_size(), th, tw), self.hidden_size, self.num_attention_heads, self.rope_dim_list, self.rope_theta) freqs_cos = freqs_cos.to(x.device) freqs_sin = freqs_sin.to(x.device) # Prepare modulation vectors @@ -499,7 +500,7 @@ class HunyuanVideoDiT(nn.Module): vec = vec + self.vector_in(text_states_2) # Add guidance modulation if needed - if self.guidance_embed and guidance is not None: + if self.guidance_embeds and guidance is not None: vec = vec + self.guidance_in(guidance) # Embed image and text img = self.img_in(img) @@ -557,7 +558,7 @@ class SingleTokenRefiner(nn.Module): self, in_channels, hidden_size, - heads_num, + num_attention_heads, depth=2, qkv_bias=True, dtype=None, @@ -592,7 +593,7 @@ class SingleTokenRefiner(nn.Module): self.refiner_blocks = nn.ModuleList([ IndividualTokenRefinerBlock( hidden_size, - heads_num, + num_attention_heads, qkv_bias=qkv_bias, dtype=dtype, ) for _ in range(depth) @@ -624,14 +625,14 @@ class IndividualTokenRefinerBlock(nn.Module): def __init__( self, hidden_size, - heads_num, - mlp_width_ratio=4.0, + num_attention_heads, + mlp_ratio=4.0, qkv_bias=True, dtype=None, ): super().__init__() - self.heads_num = heads_num - mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.num_attention_heads = num_attention_heads + mlp_hidden_dim = int(hidden_size * mlp_ratio) # Normalization and attention self.norm1 = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype) @@ -679,7 +680,7 @@ class IndividualTokenRefinerBlock(nn.Module): qkv, _ = self.self_attn_qkv(norm_x) batch_size, seq_len = qkv.shape[0], qkv.shape[1] - qkv = qkv.view(batch_size, seq_len, 3, self.heads_num, -1) + qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] diff --git a/fastvideo/v1/test_hunyuanvideo.py b/fastvideo/v1/test_hunyuanvideo.py index 6996365cf..65948747a 100644 --- a/fastvideo/v1/test_hunyuanvideo.py +++ b/fastvideo/v1/test_hunyuanvideo.py @@ -111,13 +111,15 @@ def test_hunyuanvideo_distributed(): # Initialize the two model implementations model1 = HunyuanVideoDiT( - patch_size=patch_size, + patch_size=2, + patch_size_t=1, in_channels=4, + out_channels=4, hidden_size=hidden_size, - heads_num=heads_num, - mm_double_blocks_depth=mm_double_blocks_depth, - mm_single_blocks_depth=mm_single_blocks_depth, - rope_dim_list=[8, 16, 8], # sum = hidden_size // heads_num = 32 + num_attention_heads=heads_num, + num_layers=mm_double_blocks_depth, + num_single_layers=mm_single_blocks_depth, + rope_axes_dim=[8, 16, 8], # sum = hidden_size // heads_num = 32 dtype=torch.bfloat16 ).to(torch.bfloat16)