vllm-omni/tests/diffusion/models/magi2/test_native_distributed_parity.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

import os
import tempfile
from contextlib import ExitStack, contextmanager
from unittest.mock import patch

import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

import vllm_omni.diffusion.models.magi2.attention as attention_module
import vllm_omni.diffusion.models.magi2.layers as layers_module
import vllm_omni.diffusion.models.magi2.mh_moe as mh_moe_module
import vllm_omni.diffusion.models.magi2.parallel as parallel_module
from vllm_omni.diffusion.models.magi2.attention import VarlenHandler
from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
    Magi2MHCConfig,
    Magi2MoEConfig,
    Magi2PreviewConfig,
)
from vllm_omni.diffusion.models.magi2.modeling_magi2 import (
    Magi2PreviewTransformer,
    Modality,
)
from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup

_WORLD_SIZE = 4
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]


def _tiny_config() -> Magi2PreviewConfig:
    # Four attention, KV, and MoE heads make every head axis divisible by each
    # requested TP/SP layout. The MLP and shared-expert widths are likewise
    # divisible by TP=4, so this exercises all native row/column collectives.
    return Magi2PreviewConfig(
        num_layers=1,
        hidden_size=32,
        head_dim=8,
        num_query_groups=4,
        video_in_channels=4,
        audio_in_channels=4,
        text_in_channels=4,
        intermediate_factor=1.5,
        multimodal_layers=(0,),
        params_dtype=torch.float32,
        mhc=Magi2MHCConfig(num_streams=2),
        moe=Magi2MoEConfig(
            num_heads=4,
            num_experts=2,
            top_k=2,
            expert_intermediate_size=8,
            shared_expert_intermediate_size=8,
            modality_shared_expert_intermediate_size=8,
            layers=(0,),
        ),
    )


def _initialize_model(model: Magi2PreviewTransformer) -> None:
    generator = torch.Generator(device="cpu").manual_seed(71)
    with torch.no_grad():
        for name, parameter in model.named_parameters():
            if name == "pre_adapter.rope.bands":
                continue
            parameter.copy_(
                torch.randn(
                    parameter.shape,
                    dtype=parameter.dtype,
                    generator=generator,
                )
                * 0.025
            )
        moe = model.block.layers[0].mlp.moe_mlp
        moe.router.expert_bias.zero_()
        moe.router.expert_bias_ema.zero_()


def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, VarlenHandler]:
    generator = torch.Generator(device="cpu").manual_seed(113)
    packed = torch.randn(8, 4, generator=generator) * 0.2
    coordinates = torch.tensor(
        [
            [0, 0, 0, 2, 2, 2, 2, 2, 2],
            [0, 1, 1, 2, 2, 2, 2, 2, 2],
            [0, 0, 0, 1, 2, 2, 1, 2, 2],
            [0, 1, 1, 1, 2, 2, 1, 2, 2],
            [1, 0, 0, 2, 2, 2, 2, 2, 2],
            [1, 1, 1, 2, 2, 2, 2, 2, 2],
            [0, 0, 0, 1, 1, 1, 1, 1, 1],
            [0, 0, 0, 1, 1, 1, 1, 1, 1],
        ],
        dtype=torch.float32,
    )
    modalities = torch.tensor(
        [
            Modality.VIDEO,
            Modality.AUDIO,
            Modality.TEXT,
            Modality.TIME,
            Modality.VIDEO,
            Modality.AUDIO,
            Modality.TEXT,
            Modality.TIME,
        ]
    )
    cumulative = torch.tensor([0, packed.shape[0]], dtype=torch.int32)
    varlen = VarlenHandler(cumulative, cumulative, packed.shape[0], packed.shape[0])
    return packed, coordinates, modalities, varlen


@contextmanager
def _patched_groups(
    tp_group: Magi2ParallelGroup,
    sp_group: Magi2ParallelGroup,
):
    # TP owns native weight and MoE-head shards when enabled. SP owns Ulysses
    # tokens and is also the MoE-head group in the released SP-only layout.
    ep_group = tp_group if tp_group.world_size > 1 else sp_group
    with ExitStack() as stack:
        stack.enter_context(patch.object(layers_module, "get_magi2_tp_group", return_value=tp_group))
        stack.enter_context(patch.object(mh_moe_module, "get_magi2_ep_group", return_value=ep_group))
        stack.enter_context(patch.object(parallel_module, "get_magi2_ulysses_group", return_value=sp_group))
        stack.enter_context(patch.object(attention_module, "get_magi2_ulysses_group", return_value=sp_group))
        yield


def _current_group(
    rank: int,
    rank_groups: tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...],
) -> Magi2ParallelGroup:
    for ranks, process_group in rank_groups:
        if rank in ranks:
            return Magi2ParallelGroup(
                process_group,
                world_size=len(ranks),
                rank=ranks.index(rank),
            )
    return Magi2ParallelGroup(None, world_size=1, rank=0)


def _new_groups(
    rank_sets: tuple[tuple[int, ...], ...],
) -> tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...]:
    groups = []
    for ranks in rank_sets:
        process_group = dist.new_group(ranks=list(ranks), backend="gloo")
        groups.append((ranks, process_group))
    return tuple(groups)


def _distributed_worker(rank: int, rendezvous: str) -> None:
    torch.set_num_threads(1)
    dist.init_process_group(
        "gloo",
        init_method=rendezvous,
        rank=rank,
        world_size=_WORLD_SIZE,
    )
    try:
        singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
        with _patched_groups(singleton, singleton):
            oracle = Magi2PreviewTransformer(_tiny_config())
            _initialize_model(oracle)
            checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()]
            inputs = _inputs()
            with torch.no_grad():
                expected = oracle(*inputs)

        tp2_groups = _new_groups(((0, 1), (2, 3)))
        sp2_groups = _new_groups(((0, 2), (1, 3)))
        world_group = Magi2ParallelGroup(
            dist.group.WORLD,
            world_size=_WORLD_SIZE,
            rank=rank,
        )
        layouts = (
            ("tp4", world_group, singleton),
            ("tp2sp2", _current_group(rank, tp2_groups), _current_group(rank, sp2_groups)),
            ("sp4", singleton, world_group),
        )

        for layout_name, tp_group, sp_group in layouts:
            tp_group = Magi2ParallelGroup(
                tp_group.group,
                tp_group.world_size,
                tp_group.rank,
                replicated_sequence=True,
            )
            with _patched_groups(tp_group, sp_group):
                model = Magi2PreviewTransformer(_tiny_config())
                assert model.load_weights(checkpoint) == set(model.state_dict())
                with torch.no_grad():
                    actual = model(*inputs)

            max_abs_error = (actual - expected).abs().max()
            dist.all_reduce(max_abs_error, op=dist.ReduceOp.MAX)
            close = torch.tensor(
                int(torch.allclose(actual, expected, atol=2e-5, rtol=2e-4)),
                dtype=torch.int32,
            )
            dist.all_reduce(close, op=dist.ReduceOp.MIN)
            if not close.item():
                raise AssertionError(
                    f"{layout_name} differs from the single-rank oracle; "
                    f"global max absolute error={max_abs_error.item():.8g}"
                )
            dist.barrier()
    finally:
        dist.destroy_process_group()


@pytest.mark.skipif(
    not dist.is_available() or not dist.is_gloo_available(),
    reason="requires torch.distributed gloo",
)
def test_four_rank_tp_sp_layouts_match_single_rank_oracle() -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
        mp.spawn(
            _distributed_worker,
            args=(rendezvous,),
            nprocs=_WORLD_SIZE,
            join=True,
        )