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

update sp

Browse Source
This commit is contained in:
foreverpiano
2024-12-10 13:59:33 +00:00
parent 20b03a45d7
commit f5f5903aa4
7 changed files with 282 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/hunyuan/inference.py
+17 -56
View File
@@ -18,23 +18,12 @@ from fastvideo.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_emb
from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler
from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline
try:
import xfuser
from xfuser.core.distributed import (
get_sequence_parallel_world_size,
get_sequence_parallel_rank,
get_sp_group,
initialize_model_parallel,
init_distributed_environment
)
except:
xfuser = None
get_sequence_parallel_world_size = None
get_sequence_parallel_rank = None
get_sp_group = None
initialize_model_parallel = None
init_distributed_environment = None
from fastvideo.utils.parallel_states import (
initialize_sequence_parallel_state,
nccl_info,
)
from fastvideo.utils.communications import all_gather, all_to_all_4D
def parallelize_transformer(pipe):
transformer = pipe.transformer
@@ -53,34 +42,30 @@ def parallelize_transformer(pipe):
guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000.
return_dict: bool = True,
):
if x.shape[-2] // 2 % get_sequence_parallel_world_size() == 0:
if x.shape[-2] // 2 % nccl_info.sp_size == 0:
# try to split x by height
split_dim = -2
elif x.shape[-1] // 2 % get_sequence_parallel_world_size() == 0:
elif x.shape[-1] // 2 % nccl_info.sp_size == 0:
# try to split x by width
split_dim = -1
else:
raise ValueError(f"Cannot split video sequence into ulysses_degree x ring_degree ({get_sequence_parallel_world_size()}) parts evenly")
raise ValueError(f"Cannot split video sequence into ulysses_degree x ring_degree ({nccl_info.sp_size}) parts evenly")
# patch sizes for the temporal, height, and width dimensions are 1, 2, and 2.
temporal_size, h, w = x.shape[2], x.shape[3] // 2, x.shape[4] // 2
x = torch.chunk(x, get_sequence_parallel_world_size(),dim=split_dim)[get_sequence_parallel_rank()]
x = torch.chunk(x, nccl_info.sp_size,dim=split_dim)[nccl_info.rank_within_group]
print("### x.shape", x.shape)
dim_thw = freqs_cos.shape[-1]
freqs_cos = freqs_cos.reshape(temporal_size, h, w, dim_thw)
freqs_cos = torch.chunk(freqs_cos, get_sequence_parallel_world_size(),dim=split_dim - 1)[get_sequence_parallel_rank()]
freqs_cos = torch.chunk(freqs_cos, nccl_info.sp_size,dim=split_dim - 1)[nccl_info.rank_within_group]
freqs_cos = freqs_cos.reshape(-1, dim_thw)
dim_thw = freqs_sin.shape[-1]
freqs_sin = freqs_sin.reshape(temporal_size, h, w, dim_thw)
freqs_sin = torch.chunk(freqs_sin, get_sequence_parallel_world_size(),dim=split_dim - 1)[get_sequence_parallel_rank()]
freqs_sin = torch.chunk(freqs_sin, nccl_info.sp_size,dim=split_dim - 1)[nccl_info.rank_within_group]
freqs_sin = freqs_sin.reshape(-1, dim_thw)
from xfuser.core.long_ctx_attention import xFuserLongContextAttention
for block in transformer.double_blocks + transformer.single_blocks:
block.hybrid_seq_parallel_attn = xFuserLongContextAttention()
output = original_forward(
x,
t,
@@ -95,7 +80,7 @@ def parallelize_transformer(pipe):
return_dict = not isinstance(output, tuple)
sample = output["x"]
sample = get_sp_group().all_gather(sample, dim=split_dim)
sample = all_gather(sample, dim=split_dim)
output["x"] = sample
return output
@@ -153,28 +138,8 @@ class Inference(object):
logger.info(f"Got text-to-video model root path: {pretrained_model_path}")
# ==================== Initialize Distributed Environment ================
if False:
if args.ulysses_degree > 1 or args.ring_degree > 1:
assert xfuser is not None, \
"Ulysses Attention and Ring Attention requires xfuser package."
assert args.use_cpu_offload is False, \
"Cannot enable use_cpu_offload in the distributed environment."
dist.init_process_group("nccl")
assert dist.get_world_size() == args.ring_degree * args.ulysses_degree, \
"number of GPUs should be equal to ring_degree * ulysses_degree."
init_distributed_environment(rank=dist.get_rank(), world_size=dist.get_world_size())
initialize_model_parallel(
sequence_parallel_degree=dist.get_world_size(),
ring_degree=args.ring_degree,
ulysses_degree=args.ulysses_degree,
)
device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}")
#else:
if nccl_info.sp_size > 1:
device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}")
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -522,12 +487,8 @@ class HunyuanVideoSampler(Inference):
num_images_per_prompt (int): The number of images per prompt. Default is 1.
infer_steps (int): The number of inference steps. Default is 100.
"""
if False:
if self.parallel_args['ulysses_degree'] > 1 or self.parallel_args['ring_degree'] > 1:
assert seed is not None, \
"You have to set a seed in the distributed environment, please rerun with --seed <your-seed>."
parallelize_transformer(self.pipeline)
if nccl_info.sp_size > 1:
parallelize_transformer(self.pipeline)
out_dict = dict()
fastvideo/models/hunyuan/modules/attenion.py
+59 -13
View File
@@ -13,6 +13,10 @@ except ImportError:
flash_attn = None
flash_attn_varlen_func = None
_flash_attn_forward = None
from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
from fastvideo.utils.communications import all_gather, all_to_all_4D
from fastvideo.utils.logging import ForkedPdb
MEMORY_LAYOUT = {
@@ -157,7 +161,6 @@ def attention(
def parallel_attention(
hybrid_seq_parallel_attn,
q,
k,
v,
@@ -166,18 +169,61 @@ def parallel_attention(
cu_seqlens_q,
cu_seqlens_kv
):
attn1 = hybrid_seq_parallel_attn(
None,
q[:, :img_q_len, :, :],
k[:, :img_kv_len, :, :],
v[:, :img_kv_len, :, :],
dropout_p=0.0,
causal=False,
joint_tensor_query=q[:,img_q_len:cu_seqlens_q[1]],
joint_tensor_key=k[:,img_kv_len:cu_seqlens_kv[1]],
joint_tensor_value=v[:,img_kv_len:cu_seqlens_kv[1]],
joint_strategy="rear",
)
# 1GPU torch.Size([1, 11264, 24, 128]) tensor([ 0, 11275, 11520], device='cuda:0', dtype=torch.int32)
# 2GPU torch.Size([1, 5632, 24, 128]) tensor([ 0, 5643, 5888], device='cuda:0', dtype=torch.int32)
query, encoder_query = q[:, :img_q_len, :, :], q[:,img_q_len:cu_seqlens_q[1]]
key, encoder_key = k[:, :img_kv_len, :, :], k[:,img_kv_len:cu_seqlens_kv[1]]
value, encoder_value = v[:, :img_kv_len, :, :], v[:,img_kv_len:cu_seqlens_kv[1]]
print(query.shape, encoder_query.shape)
if True:
# batch_size, seq_len, attn_heads, head_dim
query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
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)
encoder_query = shrink_head(encoder_query, dim=2)
encoder_key = shrink_head(encoder_key, dim=2)
encoder_value = shrink_head(encoder_value, dim=2)
query, key, value = (
query.transpose(1, 2),
key.transpose(1, 2),
value.transpose(1, 2)
)
encoder_query, encoder_key, encoder_value = (
encoder_query.transpose(1, 2),
encoder_key.transpose(1, 2),
encoder_value.transpose(1, 2),
)
# [b, h, s, d]
sequence_length = query.size(2)
encoder_sequence_length = encoder_query.size(2)
# Hint: please check encoder_query.shape
query = torch.cat([query, encoder_query], dim=2)
key = torch.cat([key, encoder_key], dim=2)
value = torch.cat([value, encoder_value], dim=2)
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
(sequence_length, encoder_sequence_length), dim=2
)
# B, H, S, D
hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1)
encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous()
hidden_states = hidden_states.transpose(1,2)#.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
encoder_hidden_states = encoder_hidden_states.transpose(1, 2)#.flatten(2, 3)
encoder_hidden_states = encoder_hidden_states.to(query.dtype)
attn1 = torch.cat([hidden_states, encoder_hidden_states], dim=1)
#ForkedPdb().set_trace()
if flash_attn.__version__ >= '2.7.0':
attn2, *_ = _flash_attn_forward(
q[:,cu_seqlens_q[1]:],
fastvideo/models/hunyuan/modules/models.py
+26 -24
View File
@@ -17,6 +17,9 @@ from .mlp_layers import MLP, MLPEmbedder, FinalLayer
from .modulate_layers import ModulateDiT, modulate, apply_gate
from .token_refiner import SingleTokenRefiner
from fastvideo.utils.parallel_states import (
nccl_info,
)
class MMDoubleStreamBlock(nn.Module):
"""
@@ -200,7 +203,17 @@ class MMDoubleStreamBlock(nn.Module):
), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, img.shape[0]:{img.shape[0]}"
# attention computation start
if not self.hybrid_seq_parallel_attn:
if nccl_info.sp_size > 1:
attn = parallel_attention(
q,
k,
v,
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv
)
else:
attn = attention(
q,
k,
@@ -211,17 +224,6 @@ class MMDoubleStreamBlock(nn.Module):
max_seqlen_kv=max_seqlen_kv,
batch_size=img_k.shape[0],
)
else:
attn = parallel_attention(
self.hybrid_seq_parallel_attn,
q,
k,
v,
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv
)
# attention computation end
@@ -364,7 +366,17 @@ class MMSingleStreamBlock(nn.Module):
), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, x.shape[0]:{x.shape[0]}"
# attention computation start
if not self.hybrid_seq_parallel_attn:
if nccl_info.sp_size > 1:
attn = parallel_attention(
q,
k,
v,
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv
)
else:
attn = attention(
q,
k,
@@ -375,17 +387,7 @@ class MMSingleStreamBlock(nn.Module):
max_seqlen_kv=max_seqlen_kv,
batch_size=x.shape[0],
)
else:
attn = parallel_attention(
self.hybrid_seq_parallel_attn,
q,
k,
v,
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv
)
# attention computation end
# Compute activation in mlp stream, cat again and run second linear layer.
fastvideo/sample/sample_t2v_hunyuan.py
+39 -15
View File
@@ -15,7 +15,29 @@ from diffusers.utils import export_to_video
from fastvideo.models.hunyuan.utils.file_utils import save_videos_grid
from fastvideo.models.hunyuan.inference import HunyuanVideoSampler
import torch.distributed as dist
from fastvideo.utils.parallel_states import (
initialize_sequence_parallel_state,
nccl_info,
)
def initialize_distributed():
local_rank = int(os.getenv("RANK", 0))
world_size = int(os.getenv("WORLD_SIZE", 1))
print("world_size", world_size)
torch.cuda.set_device(local_rank)
dist.init_process_group(
backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
)
initialize_sequence_parallel_state(world_size)
def main(args):
initialize_distributed()
print(nccl_info.sp_size)
device = torch.cuda.current_device()
print(args)
models_root_path = Path(args.model_path)
if not models_root_path.exists():
@@ -32,21 +54,23 @@ def main(args):
args = hunyuan_video_sampler.args
# Start sampling
outputs = hunyuan_video_sampler.predict(
prompt=args.prompts[0],
height=args.height,
width=args.width,
video_length=args.num_frames,
seed=args.seed,
negative_prompt=args.neg_prompt,
infer_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
num_videos_per_prompt=args.num_videos,
flow_shift=args.flow_shift,
batch_size=args.batch_size,
embedded_guidance_scale=args.embedded_cfg_scale
)
samples = outputs['samples']
samples = []
for prompt in args.prompts:
outputs = hunyuan_video_sampler.predict(
prompt=prompt,
height=args.height,
width=args.width,
video_length=args.num_frames,
seed=args.seed,
negative_prompt=args.neg_prompt,
infer_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
num_videos_per_prompt=args.num_videos,
flow_shift=args.flow_shift,
batch_size=args.batch_size,
embedded_guidance_scale=args.embedded_cfg_scale
)
samples.append(outputs['samples'][0])
for prompt, video in zip(args.prompts, samples):
videos = rearrange(video.unsqueeze(0), "b c t h w -> t b c h w")
fastvideo/sample/sample_t2v_hunyuan_no_sp.py
+126
View File
@@ -0,0 +1,126 @@
import os
import imageio
import time
from einops import rearrange
import torch
import torchvision
import numpy as np
from pathlib import Path
from loguru import logger
from datetime import datetime
import argparse
from diffusers.utils import export_to_video
from fastvideo.models.hunyuan.utils.file_utils import save_videos_grid
from fastvideo.models.hunyuan.inference import HunyuanVideoSampler
def main(args):
print(args)
models_root_path = Path(args.model_path)
if not models_root_path.exists():
raise ValueError(f"`models_root` not exists: {models_root_path}")
# Create save folder to save the samples
save_path = args.output_path
os.makedirs(os.path.dirname(save_path), exist_ok=True)
# Load models
hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args)
# Get the updated args
args = hunyuan_video_sampler.args
# Start sampling
samples = []
for prompt in args.prompts:
outputs = hunyuan_video_sampler.predict(
prompt=prompt,
height=args.height,
width=args.width,
video_length=args.num_frames,
seed=args.seed,
negative_prompt=args.neg_prompt,
infer_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
num_videos_per_prompt=args.num_videos,
flow_shift=args.flow_shift,
batch_size=args.batch_size,
embedded_guidance_scale=args.embedded_cfg_scale
)
samples.append(outputs['samples'][0])
for prompt, video in zip(args.prompts, samples):
videos = rearrange(video.unsqueeze(0), "b c t h w -> t b c h w")
outputs = []
for x in videos:
x = torchvision.utils.make_grid(x, nrow=6)
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
outputs.append((x * 255).numpy().astype(np.uint8))
os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
imageio.mimsave(args.output_path + f"{prompt[:100]}.mp4", outputs, fps=args.fps)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Basic parameters
parser.add_argument("--prompts", nargs="+", default=[])
parser.add_argument("--num_frames", type=int, default=16)
parser.add_argument("--height", type=int, default=256)
parser.add_argument("--width", type=int, default=256)
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--model_path", type=str, default="data/hunyuan")
parser.add_argument("--output_path", type=str, default="./outputs/video")
parser.add_argument("--fps", type=int, default=24)
# Additional parameters
parser.add_argument("--denoise-type", type=str, default="flow", help="Denoise type for noised inputs.")
parser.add_argument("--seed", type=int, default=None, help="Seed for evaluation.")
parser.add_argument("--neg_prompt", type=str, default=None, help="Negative prompt for sampling.")
parser.add_argument("--guidance_scale", type=float, default=1.0, help="Classifier free guidance scale.")
parser.add_argument("--embedded_cfg_scale", type=float, default=6.0, help="Embedded classifier free guidance scale.")
parser.add_argument("--flow_shift", type=int, default=7, help="Flow shift parameter.")
parser.add_argument("--batch_size", type=int, default=1, help="Batch size for inference.")
parser.add_argument("--num_videos", type=int, default=1, help="Number of videos to generate per prompt.")
parser.add_argument("--load-key", type=str, default="module", help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.")
parser.add_argument("--use-cpu-offload", action="store_true", help="Use CPU offload for the model load.")
parser.add_argument("--dit-weight", type=str, default="data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt")
parser.add_argument("--reproduce", action="store_true", help="Enable reproducibility by setting random seeds and deterministic algorithms.")
parser.add_argument("--disable-autocast", action="store_true", help="Disable autocast for denoising loop and vae decoding in pipeline sampling.")
# Flow Matching
parser.add_argument("--flow-reverse", action="store_true", help="If reverse, learning/sampling from t=1 -> t=0.")
parser.add_argument("--flow-solver", type=str, default="euler", help="Solver for flow matching.")
parser.add_argument("--use-linear-quadratic-schedule", action="store_true",
help="Use linear quadratic schedule for flow matching. Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)")
parser.add_argument("--linear-schedule-end", type=int, default=25,
help="End step for linear quadratic schedule for flow matching.")
# Model parameters
parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill")
parser.add_argument("--latent-channels", type=int, default=16)
parser.add_argument("--precision", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.")
parser.add_argument("--vae", type=str, default="884-16c-hy")
parser.add_argument("--vae-precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--vae-tiling", action="store_true", default=True)
parser.add_argument("--text-encoder", type=str, default="llm")
parser.add_argument("--text-encoder-precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--text-states-dim", type=int, default=4096)
parser.add_argument("--text-len", type=int, default=256)
parser.add_argument("--tokenizer", type=str, default="llm")
parser.add_argument("--prompt-template", type=str, default="dit-llm-encode")
parser.add_argument("--prompt-template-video", type=str, default="dit-llm-encode-video")
parser.add_argument("--hidden-state-skip-layer", type=int, default=2)
parser.add_argument("--apply-final-norm", action="store_true")
parser.add_argument("--text-encoder-2", type=str, default="clipL")
parser.add_argument("--text-encoder-precision-2", type=str, default="fp16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--text-states-dim-2", type=int, default=768)
parser.add_argument("--tokenizer-2", type=str, default="clipL")
parser.add_argument("--text-len-2", type=int, default=77)
args = parser.parse_args()
main(args)
inference_hunyuan_no_sp.sh
+2 -1
View File
@@ -1,4 +1,4 @@
python3 fastvideo/sample/sample_t2v_hunyuan.py \
python3 fastvideo/sample/sample_t2v_hunyuan_no_sp.py \
--height 500 \
--width 700 \
--num_frames 29 \
@@ -7,4 +7,5 @@ python3 fastvideo/sample/sample_t2v_hunyuan.py \
--embedded_cfg_scale 6 \
--flow-reverse \
--prompts "A cat walks on the grass, realistic style." \
--prompts "A dog runs in the park, realistic style." \
--output_path outputs_video/hunyuan/
inference_hunyuan_sp.sh
+13
View File
@@ -0,0 +1,13 @@
num_gpus=2
torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
fastvideo/sample/sample_t2v_hunyuan.py \
--height 500 \
--width 700 \
--num_frames 29 \
--num_inference_steps 50 \
--guidance_scale 1 \
--embedded_cfg_scale 6 \
--flow-reverse \
--prompts "A cat walks on the grass, realistic style." \
--output_path outputs_video/hunyuan/