[Bugfix] Fix tp+sp accuracy, incorrect process group mapping (#1178)

Signed-off-by: David Chen <530634352@qq.com>

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SP subgroup construction in parallel_state.py."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+
+def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
+    def _fake_new_group(ranks, *args, **kwargs):
+        group = SimpleNamespace(ranks=list(ranks))
+        created_groups.append(group)
+        return group
+
+    return _fake_new_group
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "rank, expected_ulysses, expected_ring",
+    [
+        (0, [0, 2], [0]),
+        (1, [1, 3], [1]),
+        (2, [0, 2], [2]),
+        (3, [1, 3], [3]),
+    ],
+)
+def test_set_seq_parallel_pg_uses_explicit_sp_groups(rank, expected_ulysses, expected_ring, monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+
+    # tp=2, sp=2 -> SP groups are non-contiguous: [0,2] and [1,3]
+    sp_group_ranks = RankGenerator(2, 2, 1, 1, 1, "tp-sp-pp-cfg-dp").get_ranks("sp")
+
+    ulysses_pg, ring_pg = set_seq_parallel_pg(
+        sp_ulysses_degree=2,
+        sp_ring_degree=1,
+        rank=rank,
+        world_size=4,
+        sp_group_ranks=sp_group_ranks,
+    )
+
+    assert ulysses_pg.ranks == expected_ulysses
+    assert ring_pg.ranks == expected_ring
+
+
+@pytest.mark.cpu
+def test_set_seq_parallel_pg_validates_sp_group_ranks(monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+
+    # world_size=4, sp_size=2 -> expect 2 groups, provide 1 to trigger validation
+    with pytest.raises(ValueError, match="Invalid sp_group_ranks"):
+        set_seq_parallel_pg(
+            sp_ulysses_degree=2,
+            sp_ring_degree=1,
+            rank=0,
+            world_size=4,
+            sp_group_ranks=[[0, 2]],
+        )

vllm_omni/diffusion/distributed/parallel_state.py

@@ -503,10 +503,50 @@ def init_vae_group(
 
 
 # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
-def set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size, use_ulysses_low=True):
+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,
+) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
     """
-    sp_ulysses_degree x sp_ring_degree = seq_parallel_size
-    (ulysses_degree, dp_size)
+    Initialize sequence-parallel Ulysses and Ring process groups.
+
+    This builds sequence-parallel (SP) subgroups inside each data-parallel (DP)
+    slice. The SP group size is sp_ulysses_degree * sp_ring_degree, and
+    world_size must be divisible by that size.
+
+    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:
+        ulyssess_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.
     """
     sp_size = sp_ring_degree * sp_ulysses_degree
     dp_size = world_size // sp_size
@@ -516,40 +556,96 @@ def set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size, use
     num_ulysses_pgs = sp_ring_degree  # world_size // sp_ulysses_degree
     num_ring_pgs = sp_ulysses_degree  # world_size // sp_ring_degree
 
-    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:
-                    ulyssess_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
-
+    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:
+                        ulyssess_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:
+                        ulyssess_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:
-        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
+        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:
+                        ulyssess_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:
-                    ulyssess_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:
+                        ulyssess_pg = group
 
     return ulyssess_pg, ring_pg
 
@@ -652,6 +748,7 @@ def initialize_model_parallel(
         data_parallel_size,
         "tp-sp-pp-cfg-dp",
     )
+    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(
@@ -685,9 +782,10 @@ def initialize_model_parallel(
         sp_ring_degree=ring_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=rank_generator.get_ranks("sp"),
+        group_ranks=sp_group_ranks,
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="sequence",