wehub-resource-sync
52 lines
1.1 KiB
Python
52 lines
1.1 KiB
Python
# Adopted from https://github.com/Wan-Video/Wan2.2
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Wan distributed utility compatibility layer.
|
|
|
|
LongLive-specific SP/DP group routing and autograd all-to-all live in
|
|
``wan_5b.distributed.sp_training``. This module keeps Wan2.2's public import
|
|
paths intact for the rest of the codebase.
|
|
"""
|
|
|
|
import torch.distributed as dist
|
|
|
|
from .sp_training import (
|
|
all_gather,
|
|
all_to_all,
|
|
all_to_all_with_grad,
|
|
gather_forward,
|
|
get_data_parallel_group,
|
|
get_sp_rank,
|
|
get_sp_world_size,
|
|
set_data_parallel_group,
|
|
set_sequence_parallel_group,
|
|
)
|
|
|
|
|
|
def init_distributed_group():
|
|
"""Initialize the default distributed group when it is not yet ready."""
|
|
if not dist.is_initialized():
|
|
dist.init_process_group(backend="nccl")
|
|
|
|
|
|
def get_rank():
|
|
return get_sp_rank()
|
|
|
|
|
|
def get_world_size():
|
|
return get_sp_world_size()
|
|
|
|
|
|
__all__ = [
|
|
"all_gather",
|
|
"all_to_all",
|
|
"all_to_all_with_grad",
|
|
"gather_forward",
|
|
"get_data_parallel_group",
|
|
"get_rank",
|
|
"get_world_size",
|
|
"init_distributed_group",
|
|
"set_data_parallel_group",
|
|
"set_sequence_parallel_group",
|
|
]
|