vllm-omni/vllm_omni/diffusion/distributed/parallel_state.py

# SPDX-License-Identifier: Apache-2.0
# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py
# Copyright 2023 The vLLM team.
# Adapted from
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
# Copyright 2024 xDiT team.
# Adapted from
# https://github.com/xdit-project/xDiT/blob/main/xfuser/core/distributed/utils.py
# https://github.com/vllm-project/vllm/blob/main/vllm/distributed/parallel_state.py
# Copyright 2023 The vLLM team.
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
"""vLLM-Omni distributed state.

It takes over the control of the distributed environment from PyTorch.
The typical workflow is:

- call `init_distributed_environment` to initialize the distributed environment.
- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to
 initialize the model parallel groups.

- any code dealing with the distributed stuff

- call `destroy_model_parallel` to destroy the model parallel groups.
- call `destroy_distributed_environment` to destroy the distributed environment.

If you only need to use the distributed environment without model parallelism,
 you can skip the model parallel initialization and destruction steps.
"""

import inspect

import torch
import torch.distributed
import vllm.distributed.parallel_state as vllm_parallel_state
from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
from vllm.logger import init_logger

from vllm_omni.diffusion import envs
from vllm_omni.diffusion.forward_context import get_forward_context
from vllm_omni.platforms import current_omni_platform

from .group_coordinator import (
    GroupCoordinator,
    PipelineGroupCoordinator,
    SequenceParallelGroupCoordinator,
)

env_info = envs.PACKAGES_CHECKER.get_packages_info()

HAS_FLASH_ATTN = env_info["has_flash_attn"]

logger = init_logger(__name__)


_WORLD: GroupCoordinator | None = None
# get _TP&_EP from vllm.distributed.parallel_state
_SP: SequenceParallelGroupCoordinator | None = None
_PP: PipelineGroupCoordinator | None = None
_CFG: GroupCoordinator | None = None
_DP: GroupCoordinator | None = None
_FS: GroupCoordinator | None = None  # Fully Sharded (HSDP shard dimension)
_DIT: GroupCoordinator | None = None

# Rank-layout metadata for expert parallelism. This is not a process group;
# it is reused by platform-specific runtimes that must build companion groups
# with the same rank layout as vLLM EP.
_EXPERT_PARALLEL_GROUP_RANKS: list[list[int]] | None = None


def generate_masked_orthogonal_rank_groups(
    world_size: int, parallel_size: list[int], mask: list[bool]
) -> list[list[int]]:
    r"""Generate orthogonal parallel groups based on the parallel size and mask.

    Arguments:
        world_size (int): world size

        parallel_size (list[int]):
            The parallel size of each orthogonal parallel type. For example, if
            tensor_parallel_size = 2, pipeline_model_parallel_group = 3, data_parallel_size = 4,
            and the parallel mapping order is tp-pp-dp, then the parallel_size = [2, 3, 4].

        mask (list[bool]):
            The mask controls which parallel methods the generated groups represent. If mask[i] is
            True, it means the generated group contains the i-th parallelism method. For example,
            if parallel_size = [tp_size, pp_size, dp_size], and mask = [True, False , True], then
            the generated group is the `tp-dp` group, if the mask = [False, True, False], then the
            generated group is the `pp` group.

    Algorithm:
        For orthogonal parallelism, such as tp/dp/pp/cp, the global_rank and
        local_rank satisfy the following equation:
            global_rank = tp_rank + dp_rank * tp_size + pp_rank * tp_size * dp_size (1)
                tp_rank \in [0, tp_size)
                dp_rank \in [0, dp_size)
                pp_rank \in [0, pp_size)

        If we want to get the `dp_group` (tp_size * pp_size groups of dp_size ranks each.
        For example,  if the gpu size is 8 and order is 'tp-pp-dp', size is '2-2-2', and the
        dp_group here is [[0, 4], [1, 5], [2, 6], [3, 7]].)
        The tp_rank and pp_rank will be combined to form the `dp_group_index`.
            dp_group_index = tp_rank + pp_rank * tp_size (2)

        So, Given that tp_rank and pp_rank satisfy equation (2), and dp_rank in
        range(0, dp_size), the ranks in dp_group[dp_group_index] satisfies the
        equation (1).

        This function solve this math problem.

    For example, if the parallel_size = [tp_size, dp_size, pp_size] = [2, 3, 4],
    and the mask = [False, True, False]. Then,
        dp_group_index(0) = tp_rank(0) + pp_rank(0) * 2
        dp_group_index(1) = tp_rank(1) + pp_rank(0) * 2
        ...
        dp_group_index(7) = tp_rank(1) + pp_rank(3) * 2

        dp_group[0] = 0 + range(0, 3) * 2 + 0 = [0, 2, 4]
        dp_group[1] = 1 + range(0, 3) * 2 + 0 = [1, 3, 5]
        ...
        dp_group[7] = 1 + range(0, 3) * 2 + 3 * 2 * 3 = [19, 21, 23]
    """

    def prefix_product(a: list[int], init=1) -> list[int]:
        r = [init]
        for v in a:
            init = init * v
            r.append(init)
        return r

    def inner_product(a: list[int], b: list[int]) -> int:
        return sum([x * y for x, y in zip(a, b)])

    def decompose(index, shape, stride=None):
        """
        This function solve the math problem below:
            There is an equation:
                index = sum(idx[i] * stride[i])
            And given the value of index, stride.
            Return the idx.
        This function will used to get the pp/dp/pp_rank
        from group_index and rank_in_group.
        """
        if stride is None:
            stride = prefix_product(shape)
        idx = [(index // d) % s for s, d in zip(shape, stride)]
        # stride is a prefix_product result. And the value of stride[-1]
        # is not used.
        assert sum([x * y for x, y in zip(idx, stride[:-1])]) == index, (
            f"idx {index} with shape {shape} mismatch the return idx {idx}"
        )
        return idx

    masked_shape = [s for s, m in zip(parallel_size, mask) if m]
    unmasked_shape = [s for s, m in zip(parallel_size, mask) if not m]

    global_stride = prefix_product(parallel_size)
    masked_stride = [d for d, m in zip(global_stride, mask) if m]
    unmasked_stride = [d for d, m in zip(global_stride, mask) if not m]

    group_size = prefix_product(masked_shape)[-1]
    num_of_group = world_size // group_size

    ranks = []
    for group_index in range(num_of_group):
        # get indices from unmaksed for group_index.
        decomposed_group_idx = decompose(group_index, unmasked_shape)
        rank = []
        for rank_in_group in range(group_size):
            # get indices from masked for rank_in_group.
            decomposed_rank_idx = decompose(rank_in_group, masked_shape)
            rank.append(
                inner_product(decomposed_rank_idx, masked_stride) + inner_product(decomposed_group_idx, unmasked_stride)
            )
        ranks.append(rank)
    return ranks


class RankGenerator:
    def __init__(
        self,
        tp: int,
        sp: int,
        pp: int,
        cfg: int,
        dp: int,
        fs: int = 1,
        order: str = "tp-sp-pp-cfg-dp",
        rank_offset: int = 0,
    ) -> None:
        self.tp = tp
        self.sp = sp
        self.pp = pp
        self.cfg = cfg
        self.dp = dp
        self.fs = fs
        self.rank_offset = rank_offset
        self.world_size = tp * sp * pp * cfg * dp

        self.name_to_size = {
            "tp": self.tp,
            "sp": self.sp,
            "pp": self.pp,
            "cfg": self.cfg,
            "dp": self.dp,
            "fs": self.fs,
        }
        order = order.lower()

        for name in self.name_to_size.keys():
            # Skip 'fs' validation - it's handled separately with independent_ranks=True
            # and doesn't participate in the main orthogonal rank generation
            if name == "fs":
                continue
            if name not in order and self.name_to_size[name] != 1:
                raise RuntimeError(
                    f"The size of ({name}) is ({self.name_to_size[name]}), "
                    f"but you haven't specified the order ({order})."
                )
            elif name not in order:
                order = order + "-" + name

        self.order = order
        self.ordered_size = []

        for token in order.split("-"):
            self.ordered_size.append(self.name_to_size[token])

    def get_mask(self, order: str, token: str):
        ordered_token = order.split("-")
        token = token.split("-")
        mask = [False] * len(ordered_token)
        for t in token:
            mask[ordered_token.index(t)] = True
        return mask

    def get_ranks(self, token, independent_ranks: bool = False):
        """Get rank group by input token.

        Arguments:
            token (str):
                Specify the ranks type that want to get. If we want
                to obtain multiple parallel types, we can use a hyphen
                '-' to separate them. For example, if we want to obtain
                the TP_DP group, the token should be 'tp-dp'.
            independent_ranks (bool):
                If True, generate independent rank groups that divide the world
                into groups of the specified size. Used for FS (fully shard) groups
                which operate independently from the main parallelism hierarchy.
        """
        if independent_ranks and token == "fs":
            # FS groups divide world into groups of size fs
            # e.g., world_size=8, fs=4 -> [[0,1,2,3], [4,5,6,7]]
            ranks = []
            num_groups = self.world_size // self.fs
            for i in range(num_groups):
                group = list(range(i * self.fs + self.rank_offset, (i + 1) * self.fs + self.rank_offset))
                ranks.append(group)
            return ranks

        mask = self.get_mask(self.order, token)
        ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
        if self.rank_offset > 0:
            for rank_group in ranks:
                for i in range(len(rank_group)):
                    rank_group[i] += self.rank_offset
        return ranks


# * QUERY
def get_world_group() -> GroupCoordinator:
    assert _WORLD is not None, "world group is not initialized"
    return _WORLD


# SP
def get_sp_group() -> SequenceParallelGroupCoordinator:
    assert _SP is not None, "pipeline model parallel group is not initialized"
    return _SP


def get_sequence_parallel_world_size():
    """Return world size for the sequence parallel group."""
    return get_sp_group().world_size


def get_sequence_parallel_rank():
    """Return my rank for the sequence parallel group."""
    return get_sp_group().rank_in_group


def get_ulysses_parallel_world_size():
    return get_sp_group().ulysses_world_size


def get_ulysses_parallel_rank():
    return get_sp_group().ulysses_rank


def get_ring_parallel_world_size():
    return get_sp_group().ring_world_size


def get_ring_parallel_rank():
    return get_sp_group().ring_rank


def get_allgather_parallel_world_size():
    return get_sp_group().allgather_world_size


def get_allgather_parallel_rank():
    return get_sp_group().allgather_rank


def get_expert_parallel_group_ranks() -> list[list[int]]:
    assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized"
    return _EXPERT_PARALLEL_GROUP_RANKS


# PP
def get_pp_group() -> PipelineGroupCoordinator:
    assert _PP is not None, "pipeline model parallel group is not initialized"
    return _PP


def get_pipeline_parallel_world_size():
    """Return world size for the pipeline model parallel group."""
    return get_pp_group().world_size


def get_pipeline_parallel_rank():
    """Return my rank for the pipeline model parallel group."""
    return get_pp_group().rank_in_group


def is_pipeline_first_stage():
    """Return True if in the first pipeline model parallel stage, False otherwise."""
    return get_pipeline_parallel_rank() == 0


def is_pipeline_last_stage():
    """Return True if in the last pipeline model parallel stage, False otherwise."""
    return get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1)


# CFG
def get_cfg_group() -> GroupCoordinator:
    assert _CFG is not None, "classifier_free_guidance parallel group is not initialized"
    return _CFG


def get_classifier_free_guidance_world_size():
    """Return world size for the classifier_free_guidance parallel group."""
    return get_cfg_group().world_size


def get_classifier_free_guidance_rank():
    """Return my rank for the classifier_free_guidance parallel group."""
    return get_cfg_group().rank_in_group


# DP
def get_dp_group() -> GroupCoordinator:
    assert _DP is not None, "pipeline model parallel group is not initialized"
    return _DP


def get_data_parallel_world_size():
    """Return world size for the data parallel group."""
    return get_dp_group().world_size


def get_data_parallel_rank():
    """Return my rank for the data parallel group."""
    return get_dp_group().rank_in_group


# FS (Fully Shard / HSDP shard dimension)
def get_fs_group() -> GroupCoordinator:
    assert _FS is not None, "fully shard group is not initialized"
    return _FS


def get_fully_shard_world_size():
    """Return world size for the fully shard group."""
    return get_fs_group().world_size


def get_fully_shard_rank():
    """Return my rank for the fully shard group."""
    return get_fs_group().rank_in_group


def is_dp_last_group():
    """Return True if in the last data parallel group, False otherwise."""
    return (
        get_sequence_parallel_rank() == (get_sequence_parallel_world_size() - 1)
        and get_classifier_free_guidance_rank() == (get_classifier_free_guidance_world_size() - 1)
        and get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1)
    )


def get_dit_world_size():
    """Return world size for the DiT model."""
    return (
        get_data_parallel_world_size()
        * get_classifier_free_guidance_world_size()
        * get_sequence_parallel_world_size()
        * get_pipeline_parallel_world_size()
        * get_tensor_model_parallel_world_size()
    )


# * SET


def init_world_group(ranks: list[int], local_rank: int, backend: str) -> GroupCoordinator:
    return GroupCoordinator(
        group_ranks=[ranks],
        local_rank=local_rank,
        torch_distributed_backend=backend,
    )


def init_distributed_environment(
    world_size: int = -1,
    rank: int = -1,
    distributed_init_method: str = "env://",
    local_rank: int = -1,
    backend: str | None = None,
):
    if backend is None:
        backend = current_omni_platform.dist_backend
    logger.debug(
        "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",
        world_size,
        rank,
        local_rank,
        distributed_init_method,
        backend,
    )
    if not torch.distributed.is_initialized():
        assert distributed_init_method is not None, (
            "distributed_init_method must be provided when initializing distributed environment"
        )
        # this backend is used for WORLD
        torch.distributed.init_process_group(
            backend=backend,
            init_method=distributed_init_method,
            world_size=world_size,
            rank=rank,
        )
        device_id = torch.distributed.get_rank() % current_omni_platform.get_device_count()
        current_omni_platform.set_device(current_omni_platform.get_torch_device(device_id))
    # set the local rank
    # local_rank is not available in torch ProcessGroup,
    # see https://github.com/pytorch/pytorch/issues/122816
    if local_rank == -1:
        # local rank not set, this usually happens in single-node
        # setting, where we can use rank as local rank
        if distributed_init_method == "env://":
            local_rank = envs.LOCAL_RANK
        else:
            local_rank = rank
    global _WORLD
    if _WORLD is None:
        ranks = list(range(torch.distributed.get_world_size()))
        _WORLD = init_world_group(ranks, local_rank, backend)
    else:
        assert _WORLD.world_size == torch.distributed.get_world_size(), (
            "world group already initialized with a different world size"
        )


def model_parallel_is_initialized():
    """Check if tensor and pipeline parallel groups are initialized."""
    return (
        _DP is not None
        and _CFG is not None
        and _SP is not None
        and _PP is not None
        and vllm_parallel_state._TP is not None
    )


def init_model_parallel_group(
    group_ranks: list[list[int]],
    local_rank: int,
    backend: str,
    parallel_mode: str,
    **kwargs,
) -> GroupCoordinator:
    assert parallel_mode in [
        "data",
        "pipeline",
        "tensor",
        "expert",
        "sequence",
        "classifier_free_guidance",
        "fully_shard",
    ], f"parallel_mode {parallel_mode} is not supported"
    if parallel_mode == "pipeline":
        return PipelineGroupCoordinator(
            group_ranks=group_ranks,
            local_rank=local_rank,
            torch_distributed_backend=backend,
        )
    elif parallel_mode == "sequence":
        return SequenceParallelGroupCoordinator(
            group_ranks=group_ranks,
            local_rank=local_rank,
            torch_distributed_backend=backend,
            **kwargs,
        )
    else:
        return GroupCoordinator(
            group_ranks=group_ranks,
            local_rank=local_rank,
            torch_distributed_backend=backend,
        )


def init_vllm_model_parallel_group(
    group_ranks: list[list[int]],
    local_rank: int,
    backend: str,
    group_name: str,
    use_all2all: bool = False,
) -> vllm_parallel_state.GroupCoordinator:
    # vLLM 0.27's MoE oracle (make_unquantized_moe_kernel ->
    # get_ep_all2all_manager) asserts that the EP group's device communicator
    # carries an all2all_manager, and vLLM only builds one when the group is
    # created with use_all2all=True. Older vLLM has neither the kwarg nor the
    # requirement, so forward it only when supported.
    kwargs = {}
    if use_all2all and "use_all2all" in inspect.signature(vllm_parallel_state.init_model_parallel_group).parameters:
        kwargs["use_all2all"] = True
    return vllm_parallel_state.init_model_parallel_group(
        group_ranks=group_ranks,
        local_rank=local_rank,
        backend=backend,
        group_name=group_name,
        use_device_communicator=True,
        **kwargs,
    )


def init_dit_group(
    dit_parallel_size: int,
    backend: str,
):
    global _DIT
    _DIT = torch.distributed.new_group(ranks=list(range(dit_parallel_size)), backend=backend)


def get_dit_group():
    assert _DIT is not None, "DIT group is not initialized"
    return _DIT


# adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
def set_seq_parallel_pg(
    sp_ulysses_degree: int,
    sp_ring_degree: int,
    rank: int,
    world_size: int,
    use_ulysses_low: bool = True,
    sp_group_ranks: list[list[int]] | None = None,
    sp_allgather_degree: int = 1,
) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
    """
    Initialize Ulysses, Ring, and AllGather-KV process groups.

    AllGather-KV is mutually exclusive with Ulysses and Ring.

    Args:
        sp_ulysses_degree: Size of each Ulysses subgroup.
        sp_ring_degree: Size of each Ring subgroup.
        rank: Global rank of the current process.
        world_size: Total number of processes.
        use_ulysses_low: If True, Ulysses groups are contiguous chunks and Ring
            groups are strided within each SP group. If False, the opposite.
        sp_group_ranks: Optional explicit SP groups. Each entry must be a list
            of length sp_ulysses_degree * sp_ring_degree. When provided, groups
            are built from these ranks instead of auto-generated contiguous
            ranges.

    Returns:
        ulysses_pg (torch.distributed.ProcessGroup): The Ulysses process group
            for this rank.
        ring_pg (torch.distributed.ProcessGroup): The Ring process group for
            this rank.

    Raises:
        ValueError: If sp_group_ranks length does not match world_size or any
            entry has the wrong size.
        AssertionError: If world_size is not divisible by sp_size.

    Behavior:
        - If sp_group_ranks is provided, groups are built per entry and each
          entry is further split into Ulysses/Ring groups according to
          use_ulysses_low.
        - If sp_group_ranks is None, groups are auto-generated within each DP
          slice using offsets of size sp_size.
    """
    if sp_allgather_degree > 1:
        if sp_ulysses_degree > 1 or sp_ring_degree > 1:
            raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring")
        sp_size = sp_allgather_degree
        if sp_group_ranks is None:
            sp_group_ranks = [list(range(offset, offset + sp_size)) for offset in range(0, world_size, sp_size)]
        if len(sp_group_ranks) * sp_size != world_size or any(len(ranks) != sp_size for ranks in sp_group_ranks):
            raise ValueError(f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}.")

        ulysses_pg = ring_pg = allgather_pg = None
        for group_ranks in sp_group_ranks:
            group = torch.distributed.new_group(group_ranks)
            if rank in group_ranks:
                allgather_pg = group
        for singleton_rank in range(world_size):
            group = torch.distributed.new_group([singleton_rank])
            if rank == singleton_rank:
                ulysses_pg = group
        for singleton_rank in range(world_size):
            group = torch.distributed.new_group([singleton_rank])
            if rank == singleton_rank:
                ring_pg = group
        assert ulysses_pg is not None and ring_pg is not None and allgather_pg is not None
        return ulysses_pg, ring_pg, allgather_pg

    sp_size = sp_ring_degree * sp_ulysses_degree
    dp_size = world_size // sp_size

    assert world_size % sp_size == 0, f"world_size {world_size} % sp_size {sp_ulysses_degree} == 0"

    num_ulysses_pgs = sp_ring_degree  # world_size // sp_ulysses_degree
    num_ring_pgs = sp_ulysses_degree  # world_size // sp_ring_degree

    if sp_group_ranks is not None:
        if len(sp_group_ranks) * sp_size != world_size:
            raise ValueError(
                f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}, "
                f"but got {len(sp_group_ranks)} groups."
            )
        logger.info(
            "Building SP subgroups from explicit sp_group_ranks "
            f"(sp_size={sp_size}, ulysses={sp_ulysses_degree}, ring={sp_ring_degree}, "
            f"use_ulysses_low={use_ulysses_low})."
        )
        local_sp_group = None
        local_ulysses = None
        local_ring = None
        for group_ranks in sp_group_ranks:
            if len(group_ranks) != sp_size:
                raise ValueError(f"Invalid sp_group_ranks entry: expected size {sp_size}, got {len(group_ranks)}.")
            if rank in group_ranks:
                local_sp_group = list(group_ranks)
            if use_ulysses_low:
                # Ulysses groups are contiguous chunks; Ring groups are strided.
                for i in range(num_ulysses_pgs):
                    ulysses_ranks = group_ranks[i * sp_ulysses_degree : (i + 1) * sp_ulysses_degree]
                    group = torch.distributed.new_group(ulysses_ranks)
                    if rank in ulysses_ranks:
                        ulysses_pg = group
                        local_ulysses = list(ulysses_ranks)
                for i in range(num_ring_pgs):
                    ring_ranks = group_ranks[i::num_ring_pgs]
                    group = torch.distributed.new_group(ring_ranks)
                    if rank in ring_ranks:
                        ring_pg = group
                        local_ring = list(ring_ranks)
            else:
                # Ring groups are contiguous chunks; Ulysses groups are strided.
                for i in range(num_ring_pgs):
                    ring_ranks = group_ranks[i * sp_ring_degree : (i + 1) * sp_ring_degree]
                    group = torch.distributed.new_group(ring_ranks)
                    if rank in ring_ranks:
                        ring_pg = group
                        local_ring = list(ring_ranks)
                for i in range(num_ulysses_pgs):
                    ulysses_ranks = group_ranks[i::num_ulysses_pgs]
                    group = torch.distributed.new_group(ulysses_ranks)
                    if rank in ulysses_ranks:
                        ulysses_pg = group
                        local_ulysses = list(ulysses_ranks)
        if local_sp_group is not None:
            logger.info(
                "SP group details for rank %d: sp_group=%s, ulysses_group=%s, ring_group=%s",
                rank,
                local_sp_group,
                local_ulysses,
                local_ring,
            )
    else:
        if use_ulysses_low:
            for dp_rank in range(dp_size):
                offset = dp_rank * sp_size
                for i in range(num_ulysses_pgs):
                    ulysses_ranks = list(
                        range(
                            i * sp_ulysses_degree + offset,
                            (i + 1) * sp_ulysses_degree + offset,
                        )
                    )
                    group = torch.distributed.new_group(ulysses_ranks)
                    if rank in ulysses_ranks:
                        ulysses_pg = group

                for i in range(num_ring_pgs):
                    ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs))
                    group = torch.distributed.new_group(ring_ranks)
                    if rank in ring_ranks:
                        ring_pg = group

        else:
            for dp_rank in range(dp_size):
                offset = dp_rank * sp_size
                for i in range(num_ring_pgs):
                    ring_ranks = list(range(i * sp_ring_degree + offset, (i + 1) * sp_ring_degree + offset))
                    group = torch.distributed.new_group(ring_ranks)
                    if rank in ring_ranks:
                        ring_pg = group

                for i in range(num_ulysses_pgs):
                    ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs))
                    group = torch.distributed.new_group(ulysses_ranks)
                    if rank in ulysses_ranks:
                        ulysses_pg = group

    allgather_pg = None
    for singleton_rank in range(world_size):
        group = torch.distributed.new_group([singleton_rank])
        if rank == singleton_rank:
            allgather_pg = group
    assert allgather_pg is not None
    return ulysses_pg, ring_pg, allgather_pg


def initialize_model_parallel(
    data_parallel_size: int = 1,
    cfg_parallel_size: int = 1,
    sequence_parallel_size: int | None = None,
    ulysses_degree: int = 1,
    ring_degree: int = 1,
    allgather_degree: int = 1,
    tensor_parallel_size: int = 1,
    pipeline_parallel_size: int = 1,
    fully_shard_degree: int = 1,
    hsdp_replicate_size: int = 1,
    enable_expert_parallel: bool = False,
    backend: str | None = None,
) -> None:
    if backend is None:
        backend = current_omni_platform.dist_backend
    """
    Initialize model parallel groups.

    Arguments:
        data_parallel_size: number of data parallelism groups.
        cfg_parallel_size: number of GPUs used for Classifier Free Guidance (CFG) parallelism.
        sequence_parallel_size: number of GPUs used for sequence parallelism.
            Uses allgather_degree when AllGather-KV is enabled, otherwise
            ulysses_degree * ring_degree.
        ulysses_degree: number of GPUs used for ulysses sequence parallelism.
        ring_degree: number of GPUs used for ring sequence parallelism.
        allgather_degree: number of GPUs used for AllGather-KV sequence parallelism
            (causal=False only). Mutually exclusive with ulysses/ring in v1.
        tensor_parallel_size: number of GPUs used for tensor parallelism.
        pipeline_parallel_size: number of GPUs used for pipeline parallelism.
        fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension).
        backend: distributed backend of pytorch collective comm.

    Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
    use 2 groups to parallelize the batch dim(dp), 2 groups to parallelize
    split batch caused by CFG, and 2 GPUs to parallelize sequence.

    dp_size (2) * cfg_size (2) * sp_size (2) * pp_size (2) = 16.

    The present function will create 8 data-parallel groups,
    8 CFG group, 8 pipeline-parallel group, and
    8 sequence-parallel groups:
        8 data-parallel groups:
            [g0, g8], [g1, g9], [g2, g10], [g3, g11],
            [g4, g12], [g5, g13], [g6, g14], [g7, g15]
        8 CFG-parallel groups:
            [g0, g4], [g1, g5], [g2, g6], [g3, g7],
            [g8, g12], [g9, g13], [g10, g14], [g11, g15]
        8 sequence-parallel groups:
            [g0, g1], [g2, g3], [g4, g5], [g6, g7],
            [g8, g9], [g10, g11], [g12, g13], [g14, g15]
        8 pipeline-parallel groups:
            [g0, g2], [g4, g6], [g8, g10], [g12, g14],
            [g1, g3], [g5, g7], [g9, g11], [g13, g15]
    Note that for efficiency, the caller should make sure adjacent ranks
    are on the same DGX box. For example if we are using 2 DGX-1 boxes
    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
    ranks 8 to 15 belong to the second box.
    """
    # Get world size and rank. Ensure some consistencies.
    assert torch.distributed.is_initialized()
    world_size: int = torch.distributed.get_world_size()
    backend = backend or torch.distributed.get_backend(get_world_group().device_group)

    if allgather_degree > 1:
        assert ulysses_degree == 1 and ring_degree == 1, (
            "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
            f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
            f"allgather_degree={allgather_degree}."
        )

    expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree
    if sequence_parallel_size is None:
        sequence_parallel_size = expected_sequence_parallel_size
        logger.info("sequence_parallel_size is not provided, using %d", sequence_parallel_size)

    if sequence_parallel_size != expected_sequence_parallel_size:
        raise ValueError(
            f"sequence_parallel_size must be {expected_sequence_parallel_size} for the configured SP mode, "
            f"but got {sequence_parallel_size}"
        )

    dit_parallel_size = (
        data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
    )

    # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
    is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1

    # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) as dit_parallel_size
    # This ensures orthogonal rank generation works correctly for all HSDP workers
    if is_standalone_hsdp:
        dit_parallel_size = fully_shard_degree * hsdp_replicate_size

    if world_size < dit_parallel_size:
        raise RuntimeError(
            f"world_size ({world_size}) is less than "
            f"tensor_parallel_size ({tensor_parallel_size}) x "
            f"pipeline_parallel_size ({pipeline_parallel_size}) x"
            f"sequence_parallel_size ({sequence_parallel_size}) x"
            f"cfg_parallel_size "
            f"({cfg_parallel_size}) x"
            f"data_parallel_size ({data_parallel_size})"
        )

    # For standalone HSDP, use (fully_shard_degree * hsdp_replicate_size) as data_parallel_size
    # so that RankGenerator.world_size matches the actual number of workers
    effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size

    rank_generator: RankGenerator = RankGenerator(
        tensor_parallel_size,
        sequence_parallel_size,
        pipeline_parallel_size,
        cfg_parallel_size,
        effective_dp_size,
        fs=fully_shard_degree,
        order="tp-sp-pp-cfg-dp",
    )
    use_moe_parallel_mapping = False
    if enable_expert_parallel:
        od_config = get_forward_context().omni_diffusion_config
        use_moe_parallel_mapping = bool(od_config and od_config.is_moe)
        if not use_moe_parallel_mapping:
            raise RuntimeError("Expert parallelism enabled for a non-MoE model")

    sp_group_ranks = rank_generator.get_ranks("sp")
    global _DP
    assert _DP is None, "data parallel group is already initialized"
    _DP = init_model_parallel_group(
        group_ranks=rank_generator.get_ranks("dp"),
        local_rank=get_world_group().local_rank,
        backend=backend,
        parallel_mode="data",
    )
    vllm_parallel_state._DP = _DP

    global _CFG
    assert _CFG is None, "classifier_free_guidance group is already initialized"
    _CFG = init_model_parallel_group(
        group_ranks=rank_generator.get_ranks("cfg"),
        local_rank=get_world_group().local_rank,
        backend=backend,
        parallel_mode="classifier_free_guidance",
    )
    global _PP
    assert _PP is None, "pipeline model parallel group is already initialized"
    _PP = init_model_parallel_group(
        group_ranks=rank_generator.get_ranks("pp"),
        local_rank=get_world_group().local_rank,
        backend=backend,
        parallel_mode="pipeline",
    )
    vllm_parallel_state._PP = _PP

    global _SP
    assert _SP is None, "sequence parallel group is already initialized"
    ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg(
        sp_ulysses_degree=ulysses_degree,
        sp_ring_degree=ring_degree,
        sp_allgather_degree=allgather_degree,
        rank=get_world_group().rank_in_group,
        world_size=dit_parallel_size,
        sp_group_ranks=sp_group_ranks,
    )
    _SP = init_model_parallel_group(
        group_ranks=sp_group_ranks,
        local_rank=get_world_group().local_rank,
        backend=backend,
        parallel_mode="sequence",
        ulysses_group=ulysses_pg,
        ring_group=ring_pg,
        allgather_group=allgather_pg,
    )
    if use_moe_parallel_mapping:
        # Diffusion normally uses its own SP group. Map it to vLLM PCP only for
        # expert-parallel runtimes that rely on vLLM FusedMoE group semantics.
        # vLLM 0.24 MoE kernels require GroupCoordinator.device_communicator
        # and reduce_scatter(), which the diffusion SP coordinator intentionally
        # does not own. Keep the rank layout but build a vLLM coordinator.
        vllm_parallel_state._PCP = init_vllm_model_parallel_group(
            group_ranks=sp_group_ranks,
            local_rank=get_world_group().local_rank,
            backend=backend,
            group_name="pcp",
        )

    assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
    tp_group_ranks = rank_generator.get_ranks("tp")
    if use_moe_parallel_mapping:
        vllm_parallel_state._TP = init_vllm_model_parallel_group(
            group_ranks=tp_group_ranks,
            local_rank=get_world_group().local_rank,
            backend=backend,
            group_name="tp",
        )
    else:
        vllm_parallel_state._TP = init_model_parallel_group(
            group_ranks=tp_group_ranks,
            local_rank=get_world_group().local_rank,
            backend=backend,
            parallel_mode="tensor",
        )
    if use_moe_parallel_mapping:
        # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
        # only when constructing the vLLM EP layout for expert-parallel paths.
        vllm_parallel_state._DP = init_vllm_model_parallel_group(
            group_ranks=rank_generator.get_ranks("cfg-dp"),
            local_rank=get_world_group().local_rank,
            backend=backend,
            group_name="dp",
        )

    global _FS, _EXPERT_PARALLEL_GROUP_RANKS
    assert _FS is None, "fully shard group is already initialized"
    _FS = init_model_parallel_group(
        group_ranks=rank_generator.get_ranks("fs", independent_ranks=True),
        local_rank=get_world_group().local_rank,
        backend=backend,
        parallel_mode="fully_shard",
    )

    _EXPERT_PARALLEL_GROUP_RANKS = rank_generator.get_ranks("tp-sp-cfg-dp")
    if use_moe_parallel_mapping:
        vllm_parallel_state._EP = init_vllm_model_parallel_group(
            group_ranks=_EXPERT_PARALLEL_GROUP_RANKS,
            local_rank=get_world_group().local_rank,
            backend=backend,
            group_name="ep",
            use_all2all=True,
        )

    init_dit_group(dit_parallel_size, backend)


def destroy_model_parallel():
    """Set the groups to none and destroy them."""
    global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS

    if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP:
        vllm_parallel_state._DP.destroy()
    vllm_parallel_state._DP = None

    if _DP:
        _DP.destroy()
    _DP = None

    if _CFG:
        _CFG.destroy()
    _CFG = None

    if vllm_parallel_state._PCP and vllm_parallel_state._PCP is not _SP:
        vllm_parallel_state._PCP.destroy()
    vllm_parallel_state._PCP = None

    if _SP:
        _SP.destroy()
    _SP = None

    if vllm_parallel_state._TP:
        vllm_parallel_state._TP.destroy()
    vllm_parallel_state._TP = None

    if vllm_parallel_state._EP:
        vllm_parallel_state._EP.destroy()
    vllm_parallel_state._EP = None
    _EXPERT_PARALLEL_GROUP_RANKS = None

    if _PP:
        _PP.destroy()
    _PP = None

    if _FS:
        _FS.destroy()
    _FS = None


def destroy_distributed_environment():
    global _WORLD
    if _WORLD:
        _WORLD.destroy()
    _WORLD = None
    if torch.distributed.is_initialized():
        torch.distributed.destroy_process_group()


def destroy_distributed_env():
    if model_parallel_is_initialized():
        destroy_model_parallel()
    destroy_distributed_environment()