veomni/distributed/parallel_state.py

# Copyright 2025 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# Adapted from https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/parallel_dims.py

import math
import os
from dataclasses import dataclass, field
from functools import cached_property, wraps
from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional, Tuple

import torch
from torch import distributed as dist
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh

from ..utils import logging
from ..utils.device import get_device_type


if TYPE_CHECKING:
    from torch.distributed import ProcessGroup
    from torch.distributed.device_mesh import DeviceMesh


logger = logging.get_logger(__name__)

_PARALLEL_STATE: "ParallelState" = None


def requires_mesh(fn: Callable) -> Callable:
    @wraps(fn)
    def _inner(self: "ParallelState", *args, **kwargs):
        if self.device_mesh is None:
            raise ValueError("Device mesh is not initialized.")

        return fn(self, *args, **kwargs)

    return _inner


def init_para_mesh_matrix(para_size: int, para_fsdp_size: int, para_outside: bool = False) -> "DeviceMesh":
    """
    Initialize the device mesh matrix for the ExtraParallel.
    Args:
        para_size (int): The size of the ExtraParallel.
        para_fsdp_size (int): The size of the ExtraParallel-FSDP.
        para_outside (bool): Whether the ExtraParallel is outside in para-fsdp group.
    """
    if para_outside:
        with torch.device("cpu"):
            mesh = torch.arange(math.prod((para_size, para_fsdp_size)), dtype=torch.int).view(
                para_size, para_fsdp_size
            )
    else:
        with torch.device("cpu"):
            mesh = (
                torch.arange(math.prod((para_size, para_fsdp_size)), dtype=torch.int)
                .view(para_fsdp_size, para_size)
                .transpose(0, 1)
            )
    return mesh


@dataclass(frozen=True)
class ParallelState:
    dp_size: int = 1
    dp_replicate_size: int = 1
    dp_shard_size: int = 1
    tp_size: int = 1
    pp_size: int = 1
    cp_size: int = 1
    ulysses_size: int = 1
    dp_mode: Literal["ddp", "fsdp2"] = "fsdp2"
    device_type: str = get_device_type()
    include_sp_in_fsdp: bool = True
    device_mesh: Optional["DeviceMesh"] = None
    extra_parallel_names: Tuple[str] = ("ep",)
    extra_parallel_sizes: Dict[str, int] = field(default_factory=lambda: {"ep": 1})
    extra_parallel_fsdp_device_mesh: Dict[str, Optional["DeviceMesh"]] = field(default_factory=lambda: {"ep": None})
    async_enabled: Optional[bool] = False

    def __post_init__(self):
        if not self.include_sp_in_fsdp:
            raise NotImplementedError("Decoupled sequence parallel has not been implemented.")

        if self.cp_size > 1:
            raise NotImplementedError("Ring attention is not supported yet.")

        if self.pp_size * self.dp_size * self.cp_size * self.ulysses_size * self.tp_size != self.world_size:
            raise ValueError("The product of parallel sizes should be equal to the world size.")

        if self.dp_replicate_size * self.dp_shard_size != self.dp_size:
            raise ValueError(
                f"The product of dp_replicate_size: {self.dp_replicate_size} and dp_shard_size: {self.dp_shard_size} should be equal to dp_size: {self.dp_size}."
            )

        if self.sp_enabled:
            from ..distributed.sequence_parallel import (
                init_sequence_parallel,
                set_context_parallel_group,
                set_data_parallel_group,
                set_ulysses_sequence_parallel_group,
                set_unified_sequence_parallel_group,
            )

            if self.device_mesh is not None:
                set_data_parallel_group(self.device_mesh.get_group("dp"))
                if self.ulysses_size > 1:
                    set_ulysses_sequence_parallel_group(self.device_mesh.get_group("ulysses"))
                if self.cp_size > 1:
                    set_context_parallel_group(self.device_mesh.get_group("cp"))
                # set unified sequence parallel group
                set_unified_sequence_parallel_group(self.device_mesh.get_group("sp"))
            else:
                init_sequence_parallel(
                    ulysses_size=self.ulysses_size,
                    sep_dp=True,
                    ulysses_group_key="default",
                    cp_size=self.cp_size,
                )

    @property
    def is_initialized(self) -> bool:
        return dist.is_initialized()

    @property
    def local_rank(self) -> int:
        return int(os.getenv("LOCAL_RANK", "-1"))

    @property
    def global_rank(self) -> int:
        if self.is_initialized:
            return dist.get_rank()
        return -1

    @property
    def world_size(self) -> int:
        if self.is_initialized:
            return dist.get_world_size()
        return 1

    # ------------------------------ DP ------------------------------ #
    @property
    def dp_group(self) -> Optional["ProcessGroup"]:
        if self.device_mesh is not None:
            return self.device_mesh.get_group("dp")

        if self.sp_enabled:
            from ..distributed.sequence_parallel import get_data_parallel_group

            return get_data_parallel_group()

        return self.fsdp_group

    @property
    def dp_rank(self) -> int:
        if self.device_mesh is not None:
            return self.device_mesh.get_local_rank("dp")

        if self.sp_enabled:
            from ..distributed.sequence_parallel import get_data_parallel_rank

            return get_data_parallel_rank()

        return self.fsdp_rank

    @property
    @requires_mesh
    def dp_mesh(self) -> "DeviceMesh":
        if self.device_mesh is not None:
            return self.device_mesh["dp"]

        raise self.fsdp_mesh

    @property
    def dp_enabled(self) -> bool:
        return self.dp_size > 1

    # ------------------------------ DP replicate ------------------------------ #
    @property
    def dp_replicate_group(self) -> Optional["ProcessGroup"]:
        if self.device_mesh is not None:
            return self.device_mesh.get_group("dp_replicate")

    @property
    def dp_replicate_rank(self) -> int:
        if self.device_mesh is not None:
            return self.device_mesh.get_local_rank("dp_replicate")

    @property
    @requires_mesh
    def dp_replicate_mesh(self) -> "DeviceMesh":
        if self.device_mesh is not None:
            return self.device_mesh["dp_replicate"]

    @property
    def dp_replicate_enabled(self) -> bool:
        return self.dp_replicate_size > 1

    # ------------------------------ DP shard ------------------------------ #
    @property
    def dp_shard_group(self) -> Optional["ProcessGroup"]:
        if self.device_mesh is not None:
            return self.device_mesh.get_group("dp_shard")

    @property
    def dp_shard_rank(self) -> int:
        if self.device_mesh is not None:
            return self.device_mesh.get_local_rank("dp_shard")

    @property
    @requires_mesh
    def dp_shard_mesh(self) -> "DeviceMesh":
        if self.device_mesh is not None:
            return self.device_mesh["dp_shard"]

    @property
    def dp_shard_enabled(self) -> bool:
        return self.dp_shard_size >= 1

    # ----------------------------- FSDP ----------------------------- #
    @property
    def fsdp_group(self) -> Optional["ProcessGroup"]:
        if self.device_mesh is not None:
            return self.device_mesh.get_group("dp_sp")

    @property
    def fsdp_rank(self) -> int:
        if self.device_mesh is not None:
            return self.device_mesh.get_local_rank("dp_sp")

        return self.global_rank

    @property
    def dp_shard_sp_enabled(self) -> bool:
        return self.dp_shard_enabled and self.sp_enabled

    @property
    @requires_mesh
    def fsdp_mesh(self) -> "DeviceMesh":
        if self.dp_replicate_enabled:
            # HSDP
            if self.dp_shard_sp_enabled:
                return self.device_mesh["dp_replicate", "dp_shard_sp"]
            elif self.dp_shard_enabled:
                return self.device_mesh["dp_replicate", "dp_shard"]
            else:
                # DDP
                return self.device_mesh["dp_replicate"]
        # FSDP
        elif self.dp_shard_sp_enabled:
            return self.device_mesh["dp_shard_sp"]
        elif self.dp_shard_enabled:
            return self.device_mesh["dp_shard"]
        else:
            return self.device_mesh["dp"]

    @property
    def fsdp_enabled(self) -> bool:
        return self.fsdp_size > 1

    @property
    def fsdp_size(self) -> int:
        return self.world_size // (self.pp_size * self.tp_size)

    # ------------------------------ TP ------------------------------ #
    @property
    @requires_mesh
    def tp_rank(self) -> int:
        return self.device_mesh.get_local_rank("tp")

    @property
    @requires_mesh
    def tp_mesh(self) -> "DeviceMesh":
        return self.device_mesh["tp"]

    @property
    def tp_enabled(self) -> bool:
        return self.tp_size > 1

    # ------------------------------ PP ------------------------------ #
    @property
    @requires_mesh
    def pp_rank(self) -> int:
        return self.device_mesh.get_local_rank("pp")

    @property
    @requires_mesh
    def pp_mesh(self) -> "DeviceMesh":
        return self.device_mesh["pp"]

    @property
    def pp_enabled(self) -> bool:
        return self.pp_size > 1

    @property
    @requires_mesh
    def is_first_pp_stage(self) -> bool:
        return self.pp_rank == 0

    @property
    @requires_mesh
    def is_last_pp_stage(self) -> bool:
        return self.pp_rank == (self.pp_size - 1)

    # ------------------------------ EP ------------------------------ #
    @property
    @requires_mesh
    def ep_mesh(self) -> "DeviceMesh":
        return self.extra_parallel_mesh("ep")

    @property
    @requires_mesh
    def ep_fsdp_mesh(self) -> "DeviceMesh":
        return self.extra_parallel_fsdp_mesh("ep")

    @cached_property
    def ep_group(self) -> "ProcessGroup":
        return self.extra_parallel_group("ep")

    @property
    def ep_enabled(self) -> bool:
        return self.extra_parallel_enabled("ep")

    @property
    def ep_size(self) -> int:
        return self.extra_parallel_sizes["ep"]

    @property
    def ep_rank(self) -> int:
        return self.extra_parallel_rank("ep")

    @property
    def ep_fsdp_size(self) -> int:
        return self.extra_parallel_fsdp_size("ep")

    @property
    def ep_gradient_divide_factor(self) -> int:
        return self.extra_parallel_gradient_divide_factor("ep")

    # ------------------------------ Parallel list ------------------------------ #
    @requires_mesh
    def extra_parallel_mesh(self, para_name) -> "DeviceMesh":
        return self.extra_parallel_fsdp_device_mesh[para_name][para_name]

    @requires_mesh
    def extra_parallel_fsdp_mesh(self, para_name) -> "DeviceMesh":
        return self.extra_parallel_fsdp_device_mesh[para_name][para_name, f"{para_name}_fsdp"]

    @requires_mesh
    def extra_parallel_group(self, para_name) -> "ProcessGroup":
        if self.extra_parallel_enabled(para_name):
            return self.extra_parallel_mesh(para_name).get_group()
        else:
            return None

    def extra_parallel_enabled(self, para_name) -> bool:
        return self.extra_parallel_sizes[para_name] > 1

    def extra_parallel_rank(self, para_name) -> int:
        return self.extra_parallel_fsdp_device_mesh[para_name].get_local_rank(para_name)

    def extra_parallel_fsdp_size(self, para_name) -> int:
        assert self.extra_parallel_enabled(para_name), (
            f"{para_name}_fsdp_size is only available when {para_name} is enabled ({para_name}_size > 1)"
        )
        return self.fsdp_size // self.extra_parallel_sizes[para_name]

    def extra_parallel_gradient_divide_factor(self, para_name) -> int:
        # We assume the world size is the total dp size by now
        # TP and PP would make this assumption not true
        assert self.tp_size == 1
        assert self.pp_size == 1
        # For ep+fsdp2, the grad divide factor should alwasy be world size (no matter HSDP or not)
        # SP does not affect this since SP groups still replicate params
        # and their grads are all-reduced which would match grads for the same data without SP.
        return self.world_size

    @property
    def any_extra_parallel_enabled(self) -> bool:
        return any(self.extra_parallel_enabled(para_name) for para_name in self.extra_parallel_names)

    # ------------------------------ SP ------------------------------ #
    @property
    def sp_group(self) -> Optional["ProcessGroup"]:
        if self.device_mesh is not None:
            return self.device_mesh.get_group("sp")

        if self.sp_enabled:
            from .sequence_parallel import get_unified_sequence_parallel_group

            return get_unified_sequence_parallel_group()

        return None

    @property
    def sp_rank(self) -> int:
        if self.device_mesh is not None:
            return self.device_mesh.get_local_rank("sp")

        if self.sp_enabled:
            from .sequence_parallel import get_unified_sequence_parallel_rank

            return get_unified_sequence_parallel_rank()

        return -1

    @property
    def sp_enabled(self) -> bool:
        return self.cp_size > 1 or self.ulysses_size > 1

    @property
    def sp_size(self) -> int:
        return self.ulysses_size * self.cp_size

    @property
    def ulysses_group(self) -> Optional["ProcessGroup"]:
        if self.device_mesh is not None:
            return self.device_mesh.get_group("ulysses")

        if self.sp_enabled:
            from .sequence_parallel import get_ulysses_sequence_parallel_group

            return get_ulysses_sequence_parallel_group()

        return None

    @property
    def ulysses_rank(self) -> int:
        if self.device_mesh is not None:
            return self.device_mesh.get_local_rank("ulysses")

        if self.sp_enabled:
            from .sequence_parallel import get_ulysses_sequence_parallel_rank

            return get_ulysses_sequence_parallel_rank()

        return -1

    @property
    def ulysses_enabled(self) -> bool:
        return self.ulysses_size > 1

    @property
    def cp_group(self) -> Optional["ProcessGroup"]:
        if self.device_mesh is not None:
            return self.device_mesh.get_group("cp")

        if self.sp_enabled:
            from .sequence_parallel import get_context_parallel_group

            return get_context_parallel_group()

        return None

    @property
    def cp_rank(self) -> int:
        if self.device_mesh is not None:
            return self.device_mesh.get_local_rank("cp")

        if self.sp_enabled:
            from .sequence_parallel import get_context_parallel_rank

            return get_context_parallel_rank()

        return -1

    @property
    def cp_enabled(self) -> bool:
        return self.cp_size > 1


def init_parallel_state(
    dp_size: int = 1,
    dp_replicate_size: int = 1,
    dp_shard_size: int = 1,
    tp_size: int = 1,
    pp_size: int = 1,
    cp_size: int = 1,
    ulysses_size: int = 1,
    dp_mode: Literal["ddp", "fsdp2"] = "fsdp2",
    device_type: str = None,
    include_sp_in_fsdp: bool = True,
    extra_parallel_sizes: Tuple[int] = (1,),
    extra_parallel_placement_innermost: Tuple[bool] = (False,),
    extra_parallel_names: Tuple[str] = ("ep",),
    async_enabled: Optional[bool] = False,
) -> None:
    """
    Initializes global parallel state.
    """
    global _PARALLEL_STATE
    if _PARALLEL_STATE is not None:
        logger.warning("Parallel state has already been initialized.")
        return

    if device_type is None:
        device_type = get_device_type()

    # Set dp_shard_size to dp_size if dp_shard_size and dp_replicate_size are not set when dp enabled
    if dp_size > 1 and dp_shard_size == 1 and dp_replicate_size == 1:
        dp_shard_size = dp_size

    # Note that Expert Parallel is included into Extra Parallel
    assert len(extra_parallel_sizes) == len(extra_parallel_placement_innermost) == len(extra_parallel_names), (
        "each extra parallel should correspond to a size, a placement and a name"
    )

    logger.info_rank0(
        f"Initializing parallel state: dp_size {dp_size}, dp_replicate_size {dp_replicate_size}, "
        + f"dp_shard_size {dp_shard_size},tp_size {tp_size}, pp_size {pp_size}, cp_size {cp_size}, ulysses_size {ulysses_size}, "
        + ", ".join(
            [
                f"{para_name}_size {para_size}"
                for para_name, para_size in zip(extra_parallel_names, extra_parallel_sizes)
            ]
        )
    )

    device_mesh = None

    extra_parallel_fsdp_device_mesh = {f"{para_name}": None for para_name in extra_parallel_names}

    mesh_shape = []
    mesh_dim_names = []
    for d, name in zip(
        [pp_size, dp_replicate_size, dp_shard_size, ulysses_size, cp_size, tp_size],
        ["pp", "dp_replicate", "dp_shard", "ulysses", "cp", "tp"],
    ):
        if d > 1 or name in ["dp_shard"]:
            mesh_shape.append(d)
            mesh_dim_names.append(name)

    device_mesh = init_device_mesh(
        device_type=device_type,
        mesh_shape=tuple(mesh_shape),
        mesh_dim_names=tuple(mesh_dim_names),
    )

    # Mesh for data loading (no communication on this mesh)
    dp_mesh_dim_names = []
    # Mesh for param sharding
    dp_shard_sp_mesh_dim_names = []
    # Mesh for loss all-reduce
    dp_sp_mesh_dim_names = []
    # Mesh for sequence parallel
    sp_mesh_dim_names = []

    if dp_replicate_size > 1:
        dp_mesh_dim_names.append("dp_replicate")
        dp_sp_mesh_dim_names.append("dp_replicate")
    if dp_shard_size >= 1:
        dp_mesh_dim_names.append("dp_shard")
        dp_shard_sp_mesh_dim_names.append("dp_shard")
        dp_sp_mesh_dim_names.append("dp_shard")
    if ulysses_size > 1:
        dp_shard_sp_mesh_dim_names.append("ulysses")
        sp_mesh_dim_names.append("ulysses")
        dp_sp_mesh_dim_names.append("ulysses")
    if cp_size > 1:
        dp_shard_sp_mesh_dim_names.append("cp")
        sp_mesh_dim_names.append("cp")
        dp_sp_mesh_dim_names.append("cp")

    if dp_mesh_dim_names != []:
        device_mesh[tuple(dp_mesh_dim_names)]._flatten(mesh_dim_name="dp")

    if dp_shard_sp_mesh_dim_names != []:
        device_mesh[tuple(dp_shard_sp_mesh_dim_names)]._flatten(mesh_dim_name="dp_shard_sp")

    if dp_sp_mesh_dim_names != []:
        device_mesh[tuple(dp_sp_mesh_dim_names)]._flatten(mesh_dim_name="dp_sp")

    if sp_mesh_dim_names != []:
        device_mesh[tuple(sp_mesh_dim_names)]._flatten(mesh_dim_name="sp")

    for para_size, para_outside, para_name in zip(
        extra_parallel_sizes, extra_parallel_placement_innermost, extra_parallel_names
    ):
        if para_size > 1:
            world_size = dist.get_world_size()
            assert world_size % para_size == 0, f"{para_name}_size must be a factor of world_size"
            para_fsdp_size = world_size // para_size
            mesh = init_para_mesh_matrix(para_size=para_size, para_fsdp_size=para_fsdp_size, para_outside=para_outside)
            extra_parallel_fsdp_device_mesh[f"{para_name}"] = DeviceMesh(
                device_type=device_type,
                mesh=mesh,
                mesh_dim_names=(para_name, f"{para_name}_fsdp"),
            )

    logger.info_rank0(f"Device mesh: {device_mesh}")
    for para_name in extra_parallel_names:
        logger.info_rank0(f"{para_name} FSDP device mesh: {extra_parallel_fsdp_device_mesh[para_name]}")

    _PARALLEL_STATE = ParallelState(
        dp_size=dp_size,
        dp_replicate_size=dp_replicate_size,
        dp_shard_size=dp_shard_size,
        tp_size=tp_size,
        pp_size=pp_size,
        cp_size=cp_size,
        ulysses_size=ulysses_size,
        dp_mode=dp_mode,
        device_type=device_type,
        include_sp_in_fsdp=include_sp_in_fsdp,
        device_mesh=device_mesh,
        extra_parallel_names=extra_parallel_names,
        extra_parallel_sizes=dict(zip(extra_parallel_names, extra_parallel_sizes)),
        extra_parallel_fsdp_device_mesh=extra_parallel_fsdp_device_mesh,
        async_enabled=async_enabled,
    )


def get_parallel_state() -> "ParallelState":
    """
    Returns global parallel state.
    """
    if _PARALLEL_STATE is None:
        logger.warning_once("Parallel state has not been initialized. returning default Single-process state.")
        return ParallelState()

    return _PARALLEL_STATE