vllm-omni/tests/diffusion/distributed/test_comm.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.

CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
checks on real multi-GPU NCCL collectives.
"""

from __future__ import annotations

import os
from typing import Literal

import pytest
import torch

from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
from vllm_omni.diffusion.distributed.parallel_state import (
    destroy_distributed_env,
    get_sp_group,
    init_distributed_environment,
    initialize_model_parallel,
)
from vllm_omni.platforms import current_omni_platform

_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)

DeviceKind = Literal["cpu", "cuda"]


def _update_environment_variables(envs_dict: dict[str, str]) -> None:
    for key, value in envs_dict.items():
        os.environ[key] = value


def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
    if device_kind == "cpu":
        return torch.device("cpu")
    return torch.device(f"{current_omni_platform.device_type}:{local_rank}")


def _init_worker(
    local_rank: int,
    world_size: int,
    master_port: int,
    device_kind: DeviceKind,
) -> None:
    _update_environment_variables(
        {
            "RANK": str(local_rank),
            "LOCAL_RANK": str(local_rank),
            "WORLD_SIZE": str(world_size),
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": str(master_port),
        }
    )
    backend = "gloo" if device_kind == "cpu" else None
    init_distributed_environment(backend=backend)


def _close_tolerance(dtype: torch.dtype) -> tuple[float, float]:
    if dtype == torch.bfloat16:
        return 1e-3, 1e-3
    if dtype == torch.float16:
        return 1e-3, 1e-3
    return 1e-5, 1e-5


def _run_4d_identity(
    local_rank: int,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)

    _init_worker(local_rank, world_size, master_port, device_kind)
    initialize_model_parallel(ulysses_degree=world_size)
    sp_group = get_sp_group().ulysses_group

    torch.manual_seed(42 + local_rank)
    input_tensor = torch.randn(
        batch_size,
        seq_len_per_rank,
        num_heads,
        head_size,
        dtype=dtype,
        device=device,
    )
    original_input = input_tensor.clone()

    intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync)
    expected_intermediate_shape = (
        batch_size,
        seq_len_per_rank * world_size,
        num_heads // world_size,
        head_size,
    )
    assert intermediate.shape == expected_intermediate_shape, (
        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
    )

    output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync)
    assert output.shape == original_input.shape, (
        f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
    )

    rtol, atol = _close_tolerance(dtype)
    torch.testing.assert_close(
        output,
        original_input,
        rtol=rtol,
        atol=atol,
        msg="Output does not match original input after two all-to-all operations",
    )
    destroy_distributed_env()


def _run_5d_identity(
    local_rank: int,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)

    _init_worker(local_rank, world_size, master_port, device_kind)
    initialize_model_parallel(ulysses_degree=world_size)
    sp_group = get_sp_group().ulysses_group

    torch.manual_seed(42 + local_rank)
    input_tensor = torch.randn(
        batch_size,
        seq_len_per_rank,
        3,
        num_heads,
        head_size,
        dtype=dtype,
        device=device,
    )
    original_input = input_tensor.clone()

    intermediate = SeqAllToAll5D.apply(sp_group, input_tensor, 3, 1, use_sync)
    expected_intermediate_shape = (
        batch_size,
        seq_len_per_rank * world_size,
        3,
        num_heads // world_size,
        head_size,
    )
    assert intermediate.shape == expected_intermediate_shape, (
        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
    )

    output = SeqAllToAll5D.apply(sp_group, intermediate, 1, 3, use_sync)
    assert output.shape == original_input.shape, (
        f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
    )

    rtol, atol = _close_tolerance(dtype)
    torch.testing.assert_close(
        output,
        original_input,
        rtol=rtol,
        atol=atol,
        msg="Output does not match original input after two all-to-all operations",
    )
    destroy_distributed_env()


def _run_ring_p2p(
    local_rank: int,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    num_heads: int,
    head_size: int,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    device = _worker_device(local_rank, device_kind)
    if device_kind == "cuda":
        current_omni_platform.set_device(device)

    try:
        _init_worker(local_rank, world_size, master_port, device_kind)
        initialize_model_parallel(ring_degree=world_size)
        sp_group = get_sp_group()
        comm = RingComm(sp_group.ring_group)

        input_tensor = torch.full(
            (batch_size, num_heads, head_size),
            fill_value=float(local_rank + 1),
            dtype=dtype,
            device=device,
        )
        recv_tensor = comm.send_recv(input_tensor)
        comm.commit()
        comm.wait()

        prev_rank = (local_rank - 1 + world_size) % world_size
        expected_value = float(prev_rank + 1)
        expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value)
        rtol, atol = _close_tolerance(dtype)
        torch.testing.assert_close(
            recv_tensor,
            expected_tensor,
            rtol=rtol,
            atol=atol,
            msg=f"[Rank {local_rank}] Ring P2P data mismatch",
        )
    finally:
        destroy_distributed_env()


def _spawn_4d_identity(
    *,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    torch.multiprocessing.spawn(
        _run_4d_identity,
        args=(
            world_size,
            dtype,
            batch_size,
            seq_len_per_rank,
            num_heads,
            head_size,
            use_sync,
            device_kind,
            master_port,
        ),
        nprocs=world_size,
    )


def _spawn_5d_identity(
    *,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    seq_len_per_rank: int,
    num_heads: int,
    head_size: int,
    use_sync: bool,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    torch.multiprocessing.spawn(
        _run_5d_identity,
        args=(
            world_size,
            dtype,
            batch_size,
            seq_len_per_rank,
            num_heads,
            head_size,
            use_sync,
            device_kind,
            master_port,
        ),
        nprocs=world_size,
    )


def _spawn_ring_p2p(
    *,
    world_size: int,
    dtype: torch.dtype,
    batch_size: int,
    num_heads: int,
    head_size: int,
    device_kind: DeviceKind,
    master_port: int,
) -> None:
    torch.multiprocessing.spawn(
        _run_ring_p2p,
        args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
        nprocs=world_size,
    )


def _require_heads_divisible(num_heads: int, world_size: int) -> None:
    if num_heads % world_size != 0:
        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")


def _require_gpus(world_size: int) -> None:
    available_gpus = current_omni_platform.get_device_count()
    if available_gpus < world_size:
        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")


# ---------------------------------------------------------------------------
# CPU: gloo collectives on CPU tensors (no GPU)
# ---------------------------------------------------------------------------


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("use_sync", [False, True])
def test_4d_identity(world_size: int, use_sync: bool):
    _require_heads_divisible(8, world_size)
    _spawn_4d_identity(
        world_size=world_size,
        dtype=torch.float32,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cpu",
        master_port=29610,
    )


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("use_sync", [False, True])
def test_5d_identity(world_size: int, use_sync: bool):
    _require_heads_divisible(8, world_size)
    _spawn_5d_identity(
        world_size=world_size,
        dtype=torch.float32,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cpu",
        master_port=29611,
    )


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
def test_ring_p2p(world_size: int):
    _spawn_ring_p2p(
        world_size=world_size,
        dtype=torch.float32,
        batch_size=2,
        num_heads=8,
        head_size=128,
        device_kind="cpu",
        master_port=29612,
    )


# ---------------------------------------------------------------------------
# Nightly: same checks on real multi-GPU NCCL collectives
# ---------------------------------------------------------------------------


@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
    "world_size",
    [
        pytest.param(2, marks=_L4_TWO_GPU),
        pytest.param(4, marks=_L4_FOUR_GPU),
    ],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("use_sync", [False, True])
def test_4d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
    _require_gpus(world_size)
    _require_heads_divisible(8, world_size)
    _spawn_4d_identity(
        world_size=world_size,
        dtype=dtype,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cuda",
        master_port=29500,
    )


@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
    "world_size",
    [
        pytest.param(2, marks=_L4_TWO_GPU),
        pytest.param(4, marks=_L4_FOUR_GPU),
    ],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("use_sync", [False, True])
def test_5d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
    _require_gpus(world_size)
    _require_heads_divisible(8, world_size)
    _spawn_5d_identity(
        world_size=world_size,
        dtype=dtype,
        batch_size=2,
        seq_len_per_rank=8,
        num_heads=8,
        head_size=32,
        use_sync=use_sync,
        device_kind="cuda",
        master_port=29502,
    )


@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
    "world_size",
    [
        pytest.param(2, marks=_L4_TWO_GPU),
        pytest.param(4, marks=_L4_FOUR_GPU),
    ],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
    _require_gpus(world_size)
    _spawn_ring_p2p(
        world_size=world_size,
        dtype=dtype,
        batch_size=2,
        num_heads=8,
        head_size=128,
        device_kind="cuda",
        master_port=29501,
    )