156 lines
6.3 KiB
Python
156 lines
6.3 KiB
Python
import torch
|
|
from fastvideo.model.pipeline_mochi import MochiPipeline
|
|
import torch.distributed as dist
|
|
|
|
from diffusers.utils import export_to_video
|
|
from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
|
|
import argparse
|
|
import os
|
|
from diffusers.models.transformers.transformer_mochi import MochiTransformerBlock
|
|
from fastvideo.model.modeling_mochi import MochiTransformer3DModel
|
|
|
|
|
|
|
|
import sys
|
|
import pdb
|
|
|
|
class ForkedPdb(pdb.Pdb):
|
|
"""
|
|
PDB Subclass for debugging multi-processed code
|
|
Suggested in: https://stackoverflow.com/questions/4716533/how-to-attach-debugger-to-a-python-subproccess
|
|
"""
|
|
def interaction(self, *args, **kwargs):
|
|
_stdin = sys.stdin
|
|
try:
|
|
sys.stdin = open('/dev/stdin')
|
|
pdb.Pdb.interaction(self, *args, **kwargs)
|
|
finally:
|
|
sys.stdin = _stdin
|
|
def assert_all_close_list(input_list):
|
|
for i in range(len(input_list) - 1):
|
|
assert torch.allclose(input_list[i], input_list[i + 1]), f"input_list[{i}]: {input_list[i]}, input_list[{i+1}]: {input_list[i+1]}"
|
|
|
|
weight_dtype = torch.float32
|
|
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)
|
|
return world_size
|
|
|
|
def main_print(content):
|
|
if int(os.getenv('RANK', 0)) <= 0:
|
|
print(content)
|
|
|
|
@torch.inference_mode
|
|
def test_single_block(batch_size, device, seed):
|
|
# set manual seed
|
|
torch.manual_seed(seed)
|
|
device = torch.cuda.current_device()
|
|
block = MochiTransformerBlock(
|
|
dim=768,
|
|
num_attention_heads=12,
|
|
attention_head_dim=64,
|
|
pooled_projection_dim=256,
|
|
qk_norm="rms_norm",
|
|
activation_fn="swiglu",
|
|
context_pre_only=False,
|
|
).to(device)
|
|
hidden_states = torch.randn(1, 16, 768).to(device).repeat(batch_size, 1, 1)
|
|
encoder_hidden_states = torch.randn(1, 4, 256).to(device).repeat(batch_size, 1, 1)
|
|
temb = torch.randn(1, 768).to(device).repeat(batch_size, 1)
|
|
# shard hiddent_states according to world_size
|
|
local_seq_length = hidden_states.shape[1] // nccl_info.sp_size
|
|
hidden_states = hidden_states.narrow(1, nccl_info.global_rank * local_seq_length, local_seq_length)
|
|
main_print(hidden_states.shape)
|
|
hidden_states, encoder_hidden_states = block(
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
temb=temb,
|
|
)
|
|
mean = hidden_states[0].mean()
|
|
torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM)
|
|
mean = mean / nccl_info.sp_size
|
|
return mean
|
|
|
|
@torch.inference_mode
|
|
def test_DiT(batch_size, transformer, seed):
|
|
generator = torch.Generator(torch.cuda.current_device()).manual_seed(seed)
|
|
device = torch.cuda.current_device()
|
|
latent = torch.randn((1, 12, 8, 12, 8), device=device, dtype=weight_dtype, generator=generator).repeat(batch_size, 1, 1, 1, 1)
|
|
prompt_embeds = torch.randn((1, 20, 4096), device=device, dtype=weight_dtype, generator=generator).repeat(batch_size, 1, 1)
|
|
prompt_attention_mask = torch.ones((1, 20), device=device, dtype=weight_dtype).repeat(batch_size, 1)
|
|
timestep = 0
|
|
timestep = torch.tensor(timestep, device=device, dtype=weight_dtype).unsqueeze(0).repeat(batch_size)
|
|
local_seq_length = latent.shape[2] // nccl_info.sp_size
|
|
latent = latent.narrow(2, nccl_info.global_rank * local_seq_length, local_seq_length)
|
|
# main_print(latent.shape)
|
|
hidden_states = transformer(
|
|
hidden_states=latent,
|
|
encoder_hidden_states=prompt_embeds,
|
|
encoder_attention_mask=prompt_attention_mask,
|
|
timestep=timestep,
|
|
return_dict=False,
|
|
)[0]
|
|
def calculate_mean(states):
|
|
mean = states.mean()
|
|
torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM)
|
|
mean = mean / int(os.getenv('WORLD_SIZE', 1))
|
|
return mean
|
|
mean1 = calculate_mean(hidden_states[0])
|
|
main_print(hidden_states.shape)
|
|
if hidden_states.shape[0] > 1:
|
|
mean2 = calculate_mean(hidden_states[1])
|
|
return mean1, mean2
|
|
return mean1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
world_size = initialize_distributed()
|
|
device = torch.cuda.current_device()
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--test_single_block", action="store_true")
|
|
args = parser.parse_args()
|
|
seed = args.seed
|
|
|
|
if args.test_single_block:
|
|
pass
|
|
single_no_patch_bs_1 = test_single_block(1)
|
|
single_no_patch_bs_2 = test_single_block(2)
|
|
# check all close
|
|
assert torch.allclose(single_no_patch_bs_1, single_no_patch_bs_2)
|
|
single_patch_bs_1 = test_single_block(1)
|
|
single_patch_bs_2 = test_single_block(2)
|
|
assert torch.allclose(single_patch_bs_1, single_patch_bs_2)
|
|
assert torch.allclose(single_no_patch_bs_1, single_patch_bs_2)
|
|
initialize_sequence_parallel_state(world_size)
|
|
|
|
sp_patch_bs_1 = test_single_block(1)
|
|
sp_patch_bs_2 = test_single_block(2)
|
|
|
|
assert torch.allclose(sp_patch_bs_1, sp_patch_bs_2)
|
|
assert torch.allclose(single_no_patch_bs_1, sp_patch_bs_2)
|
|
|
|
else:
|
|
transformer = MochiTransformer3DModel.from_pretrained("data/mochi/transformer", torch_dtype=weight_dtype).to(device)
|
|
single_no_patch_bs_1 = test_DiT(1, transformer, seed)
|
|
single_no_patch_bs_2_a, single_no_patch_bs_2_b = test_DiT(2, transformer, seed)
|
|
|
|
|
|
single_patch_bs_1 = test_DiT(1, transformer, seed)
|
|
single_patch_bs_2_a, single_patch_bs_2_b = test_DiT(2, transformer, seed)
|
|
|
|
initialize_sequence_parallel_state(world_size)
|
|
|
|
sp_patch_bs_1 = test_DiT(1, transformer, seed)
|
|
sp_patch_bs_2_a, sp_patch_bs_2_b = test_DiT(2, transformer, seed)
|
|
|
|
assert_all_close_list([single_no_patch_bs_1, single_no_patch_bs_2_a, single_no_patch_bs_2_b, single_patch_bs_1, single_patch_bs_2_a, single_patch_bs_2_b, sp_patch_bs_1, sp_patch_bs_2_a, sp_patch_bs_2_b])
|
|
|
|
main_print(sp_patch_bs_1)
|
|
|