[Model] Add native MAGI-2 Preview diffusion support (#5918)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu authored and GitHub committed
1 parent
13c617a140
commit
d4096e01fe
78 files changed
+9896
-459
No files matched your search
@@ -0,0 +1,229 @@
|
||||
# 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,
|
||||
)
|
||||
Reference in new issue
Block a user