SubstanceSHIFT/SeedVR2-3B

1# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates2# //3# // Licensed under the Apache License, Version 2.0 (the "License");4# // you may not use this file except in compliance with the License.5# // You may obtain a copy of the License at6# //7# //     http://www.apache.org/licenses/LICENSE-2.08# //9# // Unless required by applicable law or agreed to in writing, software10# // distributed under the License is distributed on an "AS IS" BASIS,11# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# // See the License for the specific language governing permissions and13# // limitations under the License.14 15"""16Advanced distributed functions for sequence parallel.17"""18 19from typing import Optional, List20import torch21import torch.distributed as dist22from torch.distributed.device_mesh import DeviceMesh, init_device_mesh23from torch.distributed.fsdp import ShardingStrategy24 25from .basic import get_global_rank, get_world_size26 27 28_DATA_PARALLEL_GROUP = None29_SEQUENCE_PARALLEL_GROUP = None30_SEQUENCE_PARALLEL_CPU_GROUP = None31_MODEL_SHARD_CPU_INTER_GROUP = None32_MODEL_SHARD_CPU_INTRA_GROUP = None33_MODEL_SHARD_INTER_GROUP = None34_MODEL_SHARD_INTRA_GROUP = None35_SEQUENCE_PARALLEL_GLOBAL_RANKS = None36 37 38def get_data_parallel_group() -> Optional[dist.ProcessGroup]:39    """40    Get data parallel process group.41    """42    return _DATA_PARALLEL_GROUP43 44 45def get_sequence_parallel_group() -> Optional[dist.ProcessGroup]:46    """47    Get sequence parallel process group.48    """49    return _SEQUENCE_PARALLEL_GROUP50 51 52def get_sequence_parallel_cpu_group() -> Optional[dist.ProcessGroup]:53    """54    Get sequence parallel CPU process group.55    """56    return _SEQUENCE_PARALLEL_CPU_GROUP57 58 59def get_data_parallel_rank() -> int:60    """61    Get data parallel rank.62    """63    group = get_data_parallel_group()64    return dist.get_rank(group) if group else get_global_rank()65 66 67def get_data_parallel_world_size() -> int:68    """69    Get data parallel world size.70    """71    group = get_data_parallel_group()72    return dist.get_world_size(group) if group else get_world_size()73 74 75def get_sequence_parallel_rank() -> int:76    """77    Get sequence parallel rank.78    """79    group = get_sequence_parallel_group()80    return dist.get_rank(group) if group else 081 82 83def get_sequence_parallel_world_size() -> int:84    """85    Get sequence parallel world size.86    """87    group = get_sequence_parallel_group()88    return dist.get_world_size(group) if group else 189 90 91def get_model_shard_cpu_intra_group() -> Optional[dist.ProcessGroup]:92    """93    Get the CPU intra process group of model sharding.94    """95    return _MODEL_SHARD_CPU_INTRA_GROUP96 97 98def get_model_shard_cpu_inter_group() -> Optional[dist.ProcessGroup]:99    """100    Get the CPU inter process group of model sharding.101    """102    return _MODEL_SHARD_CPU_INTER_GROUP103 104 105def get_model_shard_intra_group() -> Optional[dist.ProcessGroup]:106    """107    Get the GPU intra process group of model sharding.108    """109    return _MODEL_SHARD_INTRA_GROUP110 111 112def get_model_shard_inter_group() -> Optional[dist.ProcessGroup]:113    """114    Get the GPU inter process group of model sharding.115    """116    return _MODEL_SHARD_INTER_GROUP117 118 119def init_sequence_parallel(sequence_parallel_size: int):120    """121    Initialize sequence parallel.122    """123    global _DATA_PARALLEL_GROUP124    global _SEQUENCE_PARALLEL_GROUP125    global _SEQUENCE_PARALLEL_CPU_GROUP126    global _SEQUENCE_PARALLEL_GLOBAL_RANKS127    assert dist.is_initialized()128    world_size = dist.get_world_size()129    rank = dist.get_rank()130    data_parallel_size = world_size // sequence_parallel_size131    for i in range(data_parallel_size):132        start_rank = i * sequence_parallel_size133        end_rank = (i + 1) * sequence_parallel_size134        ranks = range(start_rank, end_rank)135        group = dist.new_group(ranks)136        cpu_group = dist.new_group(ranks, backend="gloo")137        if rank in ranks:138            _SEQUENCE_PARALLEL_GROUP = group139            _SEQUENCE_PARALLEL_CPU_GROUP = cpu_group140            _SEQUENCE_PARALLEL_GLOBAL_RANKS = list(ranks)141 142 143def init_model_shard_group(144    *,145    sharding_strategy: ShardingStrategy,146    device_mesh: Optional[DeviceMesh] = None,147):148    """149    Initialize process group of model sharding.150    """151    global _MODEL_SHARD_INTER_GROUP152    global _MODEL_SHARD_INTRA_GROUP153    global _MODEL_SHARD_CPU_INTER_GROUP154    global _MODEL_SHARD_CPU_INTRA_GROUP155    assert dist.is_initialized()156    world_size = dist.get_world_size()157    if device_mesh is not None:158        num_shards_per_group = device_mesh.shape[1]159    elif sharding_strategy == ShardingStrategy.NO_SHARD:160        num_shards_per_group = 1161    elif sharding_strategy in [162        ShardingStrategy.HYBRID_SHARD,163        ShardingStrategy._HYBRID_SHARD_ZERO2,164    ]:165        num_shards_per_group = torch.cuda.device_count()166    else:167        num_shards_per_group = world_size168    num_groups = world_size // num_shards_per_group169    device_mesh = (num_groups, num_shards_per_group)170 171    gpu_mesh_2d = init_device_mesh("cuda", device_mesh, mesh_dim_names=("inter", "intra"))172    cpu_mesh_2d = init_device_mesh("cpu", device_mesh, mesh_dim_names=("inter", "intra"))173 174    _MODEL_SHARD_INTER_GROUP = gpu_mesh_2d.get_group("inter")175    _MODEL_SHARD_INTRA_GROUP = gpu_mesh_2d.get_group("intra")176    _MODEL_SHARD_CPU_INTER_GROUP = cpu_mesh_2d.get_group("inter")177    _MODEL_SHARD_CPU_INTRA_GROUP = cpu_mesh_2d.get_group("intra")178 179def get_sequence_parallel_global_ranks() -> List[int]:180    """181    Get all global ranks of the sequence parallel process group182    that the caller rank belongs to.183    """184    if _SEQUENCE_PARALLEL_GLOBAL_RANKS is None:185        return [dist.get_rank()]186    return _SEQUENCE_PARALLEL_GLOBAL_RANKS187 188 189def get_next_sequence_parallel_rank() -> int:190    """191    Get the next global rank of the sequence parallel process group192    that the caller rank belongs to.193    """194    sp_global_ranks = get_sequence_parallel_global_ranks()195    sp_rank = get_sequence_parallel_rank()196    sp_size = get_sequence_parallel_world_size()197    return sp_global_ranks[(sp_rank + 1) % sp_size]198 199 200def get_prev_sequence_parallel_rank() -> int:201    """202    Get the previous global rank of the sequence parallel process group203    that the caller rank belongs to.204    """205    sp_global_ranks = get_sequence_parallel_global_ranks()206    sp_rank = get_sequence_parallel_rank()207    sp_size = get_sequence_parallel_world_size()208    return sp_global_ranks[(sp_rank + sp_size - 1) % sp_size]