Peiyuan Zhang
5680039dbe
Update dependencies; Setup code for debugging Delete unused files and code Remove all npu code Update torchvision imports Update EMA model and t2v_debug.sh script Delete npu related stuff and remove inpaint module Remove compress kv Fix warning with dataset handling and model loading Update PyTorch index URLs and video length tolerance range Remove UDIT and inpaint fix typo for dataset download include pretrained open-sora Add pretrained model for OpenSoraT2V-ROPE-L Update max height and width for video processing
45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
import torch
|
|
import torch.distributed as dist
|
|
import os
|
|
|
|
class COMM_INFO:
|
|
def __init__(self):
|
|
self.group = None
|
|
self.world_size = 0
|
|
self.rank = -1
|
|
|
|
nccl_info = COMM_INFO()
|
|
_SEQUENCE_PARALLEL_STATE = False
|
|
def initialize_sequence_parallel_state(sequence_parallel_size):
|
|
global _SEQUENCE_PARALLEL_STATE
|
|
if sequence_parallel_size > 1:
|
|
_SEQUENCE_PARALLEL_STATE = True
|
|
initialize_sequence_parallel_group(sequence_parallel_size)
|
|
|
|
def set_sequence_parallel_state(state):
|
|
global _SEQUENCE_PARALLEL_STATE
|
|
_SEQUENCE_PARALLEL_STATE = state
|
|
|
|
def get_sequence_parallel_state():
|
|
return _SEQUENCE_PARALLEL_STATE
|
|
|
|
def initialize_sequence_parallel_group(sequence_parallel_size):
|
|
"""Initialize the sequence parallel group."""
|
|
rank = int(os.getenv('RANK', '0'))
|
|
world_size = int(os.getenv("WORLD_SIZE", '1'))
|
|
assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format(world_size, sequence_parallel_size)
|
|
# hccl
|
|
nccl_info.world_size = sequence_parallel_size
|
|
nccl_info.rank = rank
|
|
num_sequence_parallel_groups: int = world_size // sequence_parallel_size
|
|
for i in range(num_sequence_parallel_groups):
|
|
ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size)
|
|
group = dist.new_group(ranks)
|
|
if rank in ranks:
|
|
nccl_info.group = group
|
|
|
|
|
|
def destroy_sequence_parallel_group():
|
|
"""Destroy the sequence parallel group."""
|
|
dist.destroy_process_group()
|