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

fix sp issue

Browse Source
This commit is contained in:
Peiyuan Zhang
2025-02-15 08:30:20 +00:00
parent a468e82fd6
commit 90c1a6710c
2 changed files with 20 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py
-8
View File
@@ -46,14 +46,6 @@ from ...modules import HYVideoDiffusionTransformer
from ...text_encoder import TextEncoder
from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
import torch.nn.functional as F
import json
from fastvideo.utils.sliding_block_attention import get_sliding_block_attention_mask
import os
from torch.nn.attention.flex_attention import flex_attention
from functools import partial
from tqdm import tqdm
from collections import defaultdict
import matplotlib.pyplot as plt
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
fastvideo/models/hunyuan/modules/models.py
+20 -44
View File
@@ -173,28 +173,16 @@ class MMDoubleStreamBlock(nn.Module):
# Apply RoPE if needed.
if freqs_cis is not None:
if mask_param[0] is not None:
def shrink(freqs):
freqs = freqs.view(30, nccl_info.sp_size, -1, 80, 128)
local_freqs = freqs[:,nccl_info.rank_within_group].reshape(-1, 128)
return local_freqs
def shrink_head(encoder_state, dim):
local_heads = encoder_state.shape[dim] // nccl_info.sp_size
return encoder_state.narrow(
dim, nccl_info.rank_within_group * local_heads,
local_heads)
freqs_cis = (
# 122880, 128]
shrink(freqs_cis[0]),
shrink(freqs_cis[1]),
)
else:
def shrink_head(encoder_state, dim):
local_heads = encoder_state.shape[dim] // nccl_info.sp_size
return encoder_state.narrow(
dim, nccl_info.rank_within_group * local_heads,
local_heads)
freqs_cis = (
shrink_head(freqs_cis[0], dim=0),
shrink_head(freqs_cis[1], dim=0),
)
freqs_cis = (
shrink_head(freqs_cis[0], dim=0),
shrink_head(freqs_cis[1], dim=0),
)
img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False)
assert (
@@ -345,28 +333,16 @@ class MMSingleStreamBlock(nn.Module):
q = self.q_norm(q).to(v)
k = self.k_norm(k).to(v)
if mask_param[0] is not None:
def shrink(freqs):
freqs = freqs.view(30, nccl_info.sp_size, -1, 80, 128)
local_freqs = freqs[:,nccl_info.rank_within_group].reshape(-1, 128)
return local_freqs
def shrink_head(encoder_state, dim):
local_heads = encoder_state.shape[dim] // nccl_info.sp_size
return encoder_state.narrow(
dim, nccl_info.rank_within_group * local_heads,
local_heads)
freqs_cis = (
# 122880, 128]
shrink(freqs_cis[0]),
shrink(freqs_cis[1]),
)
else:
def shrink_head(encoder_state, dim):
local_heads = encoder_state.shape[dim] // nccl_info.sp_size
return encoder_state.narrow(
dim, nccl_info.rank_within_group * local_heads,
local_heads)
freqs_cis = (
shrink_head(freqs_cis[0], dim=0),
shrink_head(freqs_cis[1], dim=0),
)
freqs_cis = (
shrink_head(freqs_cis[0], dim=0),
shrink_head(freqs_cis[1], dim=0),
)
img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :]
img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :]
@@ -642,8 +618,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
oh // self.patch_size[1], # codespell:ignore
ow // self.patch_size[2], # codespell:ignore
)
original_th = nccl_info.sp_size * th
freqs_cos, freqs_sin = self.get_rotary_pos_embed((tt, original_th, tw))
original_tt = nccl_info.sp_size * tt
freqs_cos, freqs_sin = self.get_rotary_pos_embed((original_tt, th, tw))
# Prepare modulation vectors.
vec = self.time_in(t)