diff --git a/tests/diffusion/models/magi2/test_native_compile.py b/tests/diffusion/models/magi2/test_native_compile.py new file mode 100644 index 000000000..0abac75f8 --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_compile.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest +import torch + +from tests.diffusion.models.magi2.test_native_packing import ( + _longer_text_tensors, + _sampler_tensors, + _tiny_model, + _tiny_sampler, +) +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def test_layer_regions_bracket_the_eager_kernels() -> None: + moe_layer, dense_layer = _tiny_model().block.layers + + assert moe_layer.region_methods == ("_attention_input", "_moe_input", "_moe_output") + assert dense_layer.region_methods == ("_attention_input", "_dense_output") + + +def test_compile_regions_fullgraph_matches_eager_bitwise() -> None: + model = _tiny_model() + sampler = _tiny_sampler(model) + steps = [ + sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig()) + for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0]))) + ] + expected = [sampler.forward(step) for step in steps] + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + first = sampler.forward(steps[0]) + graphs = torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before + with torch._dynamo.config.patch(error_on_recompile=True): + second = sampler.forward(steps[1]) + + for (video, audio), (video_ref, audio_ref) in zip((first, second), expected, strict=True): + assert torch.equal(video, video_ref) + assert torch.equal(audio, audio_ref) + # pre adapter, post adapter, three MoE-layer regions, two dense-layer regions + assert graphs == 7 + + +def test_pipeline_setup_compile_forwards_config_dynamic() -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline + + calls: list[dict[str, object]] = [] + + class _Transformer(torch.nn.Module): + def compile_regions(self, **compile_kwargs: object) -> None: + calls.append(compile_kwargs) + + pipeline = object.__new__(Magi2Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = OmniDiffusionConfig( + model="sand-ai/MAGI-2-preview", + model_class_name="Magi2Pipeline", + diffusion_compile_dynamic=False, + ) + pipeline.transformer = _Transformer() + + pipeline.setup_compile() + + assert calls == [{"fullgraph": True, "dynamic": False, "options": {"emulate_precision_casts": True}}] + + +def test_pipeline_setup_compile_warns_when_granularity_is_not_regional(caplog: pytest.LogCaptureFixture) -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline + + calls: list[dict[str, object]] = [] + + class _Transformer(torch.nn.Module): + def compile_regions(self, **compile_kwargs: object) -> None: + calls.append(compile_kwargs) + + pipeline = object.__new__(Magi2Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = OmniDiffusionConfig( + model="sand-ai/MAGI-2-preview", + model_class_name="Magi2Pipeline", + diffusion_compile_granularity="full", + ) + pipeline.transformer = _Transformer() + + pipeline.setup_compile() + + assert calls == [{"fullgraph": True, "dynamic": True, "options": {"emulate_precision_casts": True}}] + assert "diffusion_compile_granularity='full'" in caplog.text + + +def test_compiled_graphs_are_shared_across_layers_of_one_kind() -> None: + model = _tiny_model(moe_layers=2) + sampler = _tiny_sampler(model) + step = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([900.0]), cfg_config=CFGConfig()) + expected = sampler.forward(step) + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + actual = sampler.forward(step) + + assert torch.equal(actual[0], expected[0]) + # Two MoE layers and two dense layers still compile one graph per region + # kind: pre adapter, post adapter, three MoE regions, two dense regions. + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7 + + +def test_compiled_regions_survive_new_requests_without_recompiling() -> None: + model = _tiny_model() + sampler = _tiny_sampler(model) + t = torch.tensor([900.0]) + first_request = sampler.prepare_model_input(**_sampler_tensors(0), t=t, cfg_config=CFGConfig()) + second_request = sampler.prepare_model_input(**_sampler_tensors(1), t=t, cfg_config=CFGConfig()) + longer_request = sampler.prepare_model_input(**_longer_text_tensors(2), t=t, cfg_config=CFGConfig()) + expected = [sampler.forward(request) for request in (first_request, second_request, longer_request)] + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + actual = [sampler.forward(first_request)] + with torch._dynamo.config.patch(error_on_recompile=True): + actual.append(sampler.forward(second_request)) + actual.append(sampler.forward(longer_request)) + + for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True): + assert torch.equal(video, video_ref) + assert torch.equal(audio, audio_ref) + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7 diff --git a/tests/diffusion/models/magi2/test_native_compile_cuda.py b/tests/diffusion/models/magi2/test_native_compile_cuda.py new file mode 100644 index 000000000..1db394750 --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_compile_cuda.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest +import torch + +from tests.diffusion.models.magi2.test_native_packing import ( + _longer_text_tensors, + _sampler_tensors, + _tiny_model, + _tiny_sampler, +) +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig + +pytestmark = [ + pytest.mark.diffusion, + pytest.mark.core_model, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), + pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"), +] + + +def _cuda_steps(sampler, tensors=_sampler_tensors): + timesteps = (torch.tensor([900.0]), torch.tensor([600.0]), torch.tensor([300.0])) + return [ + sampler.prepare_model_input( + **{name: value.cuda() for name, value in tensors(seed).items()}, + t=t, + cfg_config=CFGConfig(), + ) + for seed, t in enumerate(timesteps) + ] + + +def test_inductor_regions_match_eager_in_deterministic_mode(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MAGI2_DETERMINISTIC", "1") + model = _tiny_model(params_dtype=torch.bfloat16).cuda() + sampler = _tiny_sampler(model) + steps = _cuda_steps(sampler) + expected = [sampler.forward(step) for step in steps] + new_request_step = _cuda_steps(sampler)[0] + longer_text_step = _cuda_steps(sampler, _longer_text_tensors)[0] + expected_new_request = sampler.forward(new_request_step) + expected_longer_text = sampler.forward(longer_text_step) + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + # The pipeline compiles with inductor emulating eager's intermediate + # low-precision casts, so fused bf16 chains keep eager's rounding. + model.compile_regions(fullgraph=True, dynamic=True, options={"emulate_precision_casts": True}) + first = sampler.forward(steps[0]) + with torch._dynamo.config.patch(error_on_recompile=True): + compiled = [first, sampler.forward(steps[1]), sampler.forward(steps[2])] + repeated = [sampler.forward(step) for step in steps] + # A second request repacks the same shapes; a third changes the + # text length and therefore every token count. + compiled_new_request = sampler.forward(new_request_step) + compiled_longer_text = sampler.forward(longer_text_step) + + # Emulated precision casts restore eager's rounding boundaries; what is + # left is one-ulp transcendental noise that some Triton backends carry. + for (video, audio), (video_ref, audio_ref) in zip(compiled, expected, strict=True): + torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0) + torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0) + for (video, audio), (video_again, audio_again) in zip(compiled, repeated, strict=True): + assert torch.equal(video, video_again) + assert torch.equal(audio, audio_again) + for (video, audio), (video_ref, audio_ref) in ( + (compiled_new_request, expected_new_request), + (compiled_longer_text, expected_longer_text), + ): + torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0) + torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0) + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7 diff --git a/tests/diffusion/models/magi2/test_native_compile_distributed.py b/tests/diffusion/models/magi2/test_native_compile_distributed.py new file mode 100644 index 000000000..03f6f601c --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_compile_distributed.py @@ -0,0 +1,179 @@ +# 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 torch.nn as nn +from torch.distributed import init_device_mesh +from torch.distributed.fsdp import MixedPrecisionPolicy + +import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dlo_module +from tests.diffusion.models.magi2.test_native_distributed_parity import ( + _current_group, + _new_groups, + _patched_groups, +) +from tests.diffusion.models.magi2.test_native_packing import _tiny_config, _tiny_model, _tiny_sampler +from tests.diffusion.offloader.helpers import patch_offload_runtime +from vllm_omni.diffusion.distributed.hsdp import shard_model +from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer +from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup +from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig +from vllm_omni.diffusion.offloader.base import OffloadConfig +from vllm_omni.diffusion.offloader.config import OffloadStrategy + +_WORLD_SIZE = 4 +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +@contextmanager +def _cpu_offload_runtime(): + with pytest.MonkeyPatch.context() as monkeypatch: + patch_offload_runtime(monkeypatch, dlo_module.current_omni_platform, synchronize=True) + yield + + +def _enable_dlo( + model: Magi2PreviewTransformer, + *, + dp_group: dist.ProcessGroup | None, + dp_size: int, + dp_rank: int, +) -> None: + pipeline = Magi2Pipeline.__new__(Magi2Pipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = model + backend = dlo_module.DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=dp_size, + dlo_use_allgather=dp_size > 1, + ), + torch.device("cpu"), + ) + backend.dp_group = dp_group + backend.rank = dp_rank + backend.enable(pipeline) + assert backend.enabled + + +def _enable_hsdp(model: Magi2PreviewTransformer) -> None: + mesh = init_device_mesh("cpu", mesh_shape=(1, _WORLD_SIZE), mesh_dim_names=("replicate", "shard")) + ignored = { + parameter for name in model._hsdp_ignored_modules for parameter in model.get_submodule(name).parameters() + } + shard_model( + model, + mesh=mesh, + mp_policy=MixedPrecisionPolicy(cast_forward_inputs=False), + hsdp_shard_conditions=model._hsdp_shard_conditions, + ignored_params=ignored, + ) + model.requires_grad_(False) + + +def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]: + # 14 tokens per CFG branch, so the packed sequence splits evenly across + # two Ulysses ranks; gloo cannot gather uneven shards. + generator = torch.Generator(device="cpu").manual_seed(seed) + return { + "latent": torch.randn(1, 4, 2, 1, 3, generator=generator), + "audio_latent": torch.randn(1, 5, 4, generator=generator), + "txt_feat": torch.randn(1, 3, 4, generator=generator), + "null_txt_feat": torch.randn(1, 3, 4, generator=generator), + } + + +def _steps(sampler): + return [ + sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig()) + for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0]))) + ] + + +def _run(model: Magi2PreviewTransformer, *, compile_regions: bool) -> list[tuple[torch.Tensor, torch.Tensor]]: + sampler = _tiny_sampler(model) + steps = _steps(sampler) + if not compile_regions: + return [sampler.forward(step) for step in steps] + torch._dynamo.reset() + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + outputs = [sampler.forward(steps[0])] + with torch._dynamo.config.patch(error_on_recompile=True): + outputs.append(sampler.forward(steps[1])) + return outputs + + +def _assert_same(actual, expected, *, exact: bool, label: str) -> None: + for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True): + if exact: + assert torch.equal(video, video_ref), label + assert torch.equal(audio, audio_ref), label + else: + torch.testing.assert_close(video, video_ref, atol=2e-5, rtol=2e-4, msg=label) + torch.testing.assert_close(audio, audio_ref, atol=2e-5, rtol=2e-4, msg=label) + + +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 = _tiny_model() + checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()] + expected = _run(oracle, compile_regions=False) + + sp2_groups = _new_groups(((0, 1), (2, 3))) + dp2_groups = _new_groups(((0, 2), (1, 3))) + sp_group = _current_group(rank, sp2_groups) + dp_ranks, dp_process_group = next((ranks, group) for ranks, group in dp2_groups if rank in ranks) + + def dlo_rank_local(model: Magi2PreviewTransformer) -> None: + _enable_dlo(model, dp_group=None, dp_size=1, dp_rank=0) + + def dlo_allgather(model: Magi2PreviewTransformer) -> None: + _enable_dlo(model, dp_group=dp_process_group, dp_size=len(dp_ranks), dp_rank=dp_ranks.index(rank)) + + variants = ( + ("dlo_rank_local_sp2", sp_group, dlo_rank_local), + ("dlo_allgather_dp2sp2", sp_group, dlo_allgather), + ("hsdp4", singleton, _enable_hsdp), + ) + for name, ulysses_group, enable in variants: + with ExitStack() as stack: + stack.enter_context(_patched_groups(singleton, ulysses_group)) + stack.enter_context(_cpu_offload_runtime()) + outputs = {} + for compile_regions in (False, True): + model = Magi2PreviewTransformer(_tiny_config()) + assert model.load_weights(checkpoint) == set(model.state_dict()) + enable(model) + outputs[compile_regions] = _run(model, compile_regions=compile_regions) + _assert_same(outputs[False], expected, exact=False, label=f"{name} eager vs single-rank oracle") + _assert_same(outputs[True], outputs[False], exact=True, label=f"{name} compiled vs eager") + 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_compiled_regions_match_eager_under_dlo_and_hsdp() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}" + with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): + mp.spawn(_distributed_worker, args=(rendezvous,), nprocs=_WORLD_SIZE, join=True) diff --git a/tests/diffusion/models/magi2/test_native_packing.py b/tests/diffusion/models/magi2/test_native_packing.py new file mode 100644 index 000000000..cedabebcb --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_packing.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.diffusion.models.magi2.configuration_magi2 import ( + Magi2MHCConfig, + Magi2MoEConfig, + Magi2PreviewConfig, +) +from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm +from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE +from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer +from vllm_omni.diffusion.models.magi2.preview_data_proxy import ( + Magi2DataProxy, + Magi2PreviewDataProxyConfig, +) +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def _tiny_config(params_dtype: torch.dtype = torch.float32, *, moe_layers: int = 1) -> Magi2PreviewConfig: + # The leading layers are multimodal with MoE, the rest single-modality + # dense, so every layer belongs to one of two kinds. + moe_indices = tuple(range(moe_layers)) + return Magi2PreviewConfig( + num_layers=2 * moe_layers, + hidden_size=16, + head_dim=8, + num_query_groups=2, + video_in_channels=4, + audio_in_channels=4, + text_in_channels=4, + intermediate_factor=2, + multimodal_layers=moe_indices, + params_dtype=params_dtype, + mhc=Magi2MHCConfig(num_streams=2), + moe=Magi2MoEConfig( + num_heads=2, + num_experts=4, + top_k=2, + expert_intermediate_size=8, + shared_expert_intermediate_size=8, + modality_shared_expert_intermediate_size=8, + layers=moe_indices, + ), + ) + + +def _tiny_model( + seed: int = 11, + params_dtype: torch.dtype = torch.float32, + *, + moe_layers: int = 1, +) -> Magi2PreviewTransformer: + model = Magi2PreviewTransformer(_tiny_config(params_dtype, moe_layers=moe_layers)) + generator = torch.Generator(device="cpu").manual_seed(seed) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) + for module in model.modules(): + if isinstance(module, MultiModalityRMSNorm): + module.weight.zero_() + elif isinstance(module, Magi2MultiHeadMoE): + module.router.expert_bias.zero_() + module.router.expert_bias_ema.zero_() + return model + + +def _tiny_sampler(model: torch.nn.Module) -> Magi2PreviewSampler: + return Magi2PreviewSampler(model, Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=8))) + + +def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(seed) + return { + "latent": torch.randn(1, 4, 2, 1, 3, generator=generator), + "audio_latent": torch.randn(1, 5, 4, generator=generator), + "txt_feat": torch.randn(1, 3, 4, generator=generator), + "null_txt_feat": torch.randn(1, 2, 4, generator=generator), + } + + +def _longer_text_tensors(seed: int) -> dict[str, torch.Tensor]: + tensors = _sampler_tensors(seed) + generator = torch.Generator(device="cpu").manual_seed(seed + 100) + tensors["txt_feat"] = torch.randn(1, 4, 4, generator=generator) + return tensors + + +def test_prepare_model_input_keeps_lengths_on_host() -> None: + sampler = Magi2PreviewSampler(torch.nn.Identity()) + model_input = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([500.0]), cfg_config=CFGConfig()) + + assert model_input.audio_feat_len == [5, 5] + assert model_input.txt_feat_len == [3, 2] + assert model_input.ref_audio_feat_len == [0, 0] + assert model_input.ref_video_feat_len == [0, 0] + + positive, negative = Magi2PreviewSampler._split_cfg_model_input(model_input) + assert positive.txt_feat_len == [3] + assert negative.txt_feat_len == [2] + + +def test_pre_adapter_embeds_directly_in_checkpoint_dtype() -> None: + model = _tiny_model(params_dtype=torch.bfloat16) + packed = torch.randn(6, 4) + indices = torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 5]) + + with torch.no_grad(): + hidden = model.pre_adapter(packed, *indices) + + assert hidden.dtype == torch.bfloat16 + assert hidden.shape == (6, model.config.virtual_width) diff --git a/vllm_omni/diffusion/models/magi2/modeling_magi2.py b/vllm_omni/diffusion/models/magi2/modeling_magi2.py index 87d3402d3..44e282aa0 100644 --- a/vllm_omni/diffusion/models/magi2/modeling_magi2.py +++ b/vllm_omni/diffusion/models/magi2/modeling_magi2.py @@ -16,6 +16,7 @@ import os from collections.abc import Iterable from enum import IntEnum from functools import partial +from typing import Any import torch import torch.nn as nn @@ -28,6 +29,7 @@ from .configuration_magi2 import Magi2PreviewConfig from .layers import ( ElementWiseFourierEmbed, MHCHandler, + MHCTensorTuple, ModalityDispatcher, MultiModalityRMSNorm, make_grouped_linear, @@ -130,14 +132,12 @@ class Magi2Attention(nn.Module): start = self.tp_group.rank * self.num_heads_q return checkpoint_tensor[:, start : start + self.num_heads_q] - def forward( + def project( self, hidden_states: torch.Tensor, rope: torch.Tensor, - varlen_handler: VarlenHandler, modality_dispatcher: ModalityDispatcher, - cp_split_sizes: list[int] | torch.Tensor, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: normalized = self.pre_norm(hidden_states, modality_dispatcher) gates = self.linear_g(normalized, modality_dispatcher) qkv = self.linear_qkv(normalized, modality_dispatcher) @@ -159,8 +159,17 @@ class Magi2Attention(nn.Module): q = apply_rotary_emb(q, cos, sin).squeeze(0).to(self.config.params_dtype) k = apply_rotary_emb(k, cos, sin).squeeze(0).to(self.config.params_dtype) v = v.squeeze(0).to(self.config.params_dtype) + return q, k, v, gates - output = self.packed_attention( + def attend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + varlen_handler: VarlenHandler, + cp_split_sizes: list[int] | torch.Tensor, + ) -> torch.Tensor: + return self.packed_attention( q, k, v, @@ -172,7 +181,14 @@ class Magi2Attention(nn.Module): } ), ) - output = modality_dispatcher.permute(output) + + def output( + self, + attention: torch.Tensor, + gates: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + output = modality_dispatcher.permute(attention) output = output * torch.sigmoid(gates) output = output.reshape(-1, self.q_size).to(self.config.params_dtype) return self.linear_proj(output, modality_dispatcher) @@ -303,12 +319,21 @@ class Magi2MultiHeadMoELayer(nn.Module): modality.contiguous(), dispatcher ) - def forward(self, hidden_states: torch.Tensor, dispatcher: ModalityDispatcher) -> torch.Tensor: + def route_input( + self, + hidden_states: torch.Tensor, + dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, torch.Tensor]: normalized = self.pre_norm(hidden_states, dispatcher) - routed = self.split_linear(normalized) - routed = self.moe_mlp(routed) - routed = self.merge_linear(routed) - return routed + self._shared_experts(normalized, dispatcher) + return normalized, self.split_linear(normalized) + + def combine( + self, + normalized: torch.Tensor, + routed: torch.Tensor, + dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + return self.merge_linear(routed) + self._shared_experts(normalized, dispatcher) class Magi2PreAdapter(nn.Module): @@ -339,37 +364,28 @@ class Magi2PreAdapter(nn.Module): def forward( self, packed: torch.Tensor, - coords_mapping: torch.Tensor, video_indices: torch.Tensor, audio_indices: torch.Tensor, text_indices: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - rope = self.rope(coords_mapping) - output = torch.zeros( + ) -> torch.Tensor: + # Every token belongs to exactly one modality, so the three copies + # below cover the whole buffer. Embedding in fp32 and rounding each + # modality once matches the released cast at the block boundary. + output = torch.empty( packed.shape[0], self.adapter_dim, - dtype=torch.float32, + dtype=self.config.params_dtype, device=packed.device, ) - if text_indices.numel(): - output.index_copy_( - 0, - text_indices, - self.text_embedder(packed.index_select(0, text_indices)[:, : self.config.text_in_channels].float()), - ) - if audio_indices.numel(): - output.index_copy_( - 0, - audio_indices, - self.audio_embedder(packed.index_select(0, audio_indices)[:, : self.config.audio_in_channels].float()), - ) - if video_indices.numel(): - output.index_copy_( - 0, - video_indices, - self.video_embedder(packed.index_select(0, video_indices)[:, : self.config.video_in_channels].float()), - ) - return output, rope + for indices, embedder, in_channels in ( + (text_indices, self.text_embedder, self.config.text_in_channels), + (audio_indices, self.audio_embedder, self.config.audio_in_channels), + (video_indices, self.video_embedder, self.config.video_in_channels), + ): + if indices.numel(): + tokens = packed[:, :in_channels].index_select(0, indices).float() + output.index_copy_(0, indices, embedder(tokens).to(output.dtype)) + return output class Magi2PostAdapter(nn.Module): @@ -425,8 +441,10 @@ class Magi2TransformerLayer(nn.Module): self.mlp: nn.Module if layer_index in config.moe.layers: self.mlp = Magi2MultiHeadMoELayer(config) + self.region_methods = ("_attention_input", "_moe_input", "_moe_output") else: self.mlp = Magi2MLP(config, num_modality=num_modality) + self.region_methods = ("_attention_input", "_dense_output") self._init_mhc(num_modality) def _init_mhc(self, num_modality: int) -> None: @@ -525,6 +543,74 @@ class Magi2TransformerLayer(nn.Module): ) return self.mhc_handler.hyper_connect(streams, output, post, residual) + # The layer forward is split into compile regions around the two eager + # kernels: packed attention and the multi-head MoE. + + def _attention_input( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size) + attention_logits = self._branch_logits(streams, "attn", modality_dispatcher) + attention_input = self._branch_input(streams, "attn", attention_logits) + q, k, v, gates = self.attention.project(attention_input, rope, modality_dispatcher) + return streams, attention_logits, q, k, v, gates + + def _mlp_input( + self, + streams: torch.Tensor, + attention: torch.Tensor, + gates: torch.Tensor, + attention_logits: MHCTensorTuple, + modality_dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor]: + attention_output = self.attention.output(attention, gates, modality_dispatcher) + streams = self._connect(streams, attention_output, "attn", attention_logits) + mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher) + return streams, mlp_logits, self._branch_input(streams, "mlp", mlp_logits) + + def _dense_output( + self, + streams: torch.Tensor, + attention: torch.Tensor, + gates: torch.Tensor, + attention_logits: MHCTensorTuple, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + streams, mlp_logits, mlp_input = self._mlp_input( + streams, attention, gates, attention_logits, modality_dispatcher + ) + streams = self._connect(streams, self.mlp(mlp_input, modality_dispatcher), "mlp", mlp_logits) + return streams.reshape(streams.shape[0], -1) + + def _moe_input( + self, + streams: torch.Tensor, + attention: torch.Tensor, + gates: torch.Tensor, + attention_logits: MHCTensorTuple, + modality_dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor, torch.Tensor]: + streams, mlp_logits, mlp_input = self._mlp_input( + streams, attention, gates, attention_logits, modality_dispatcher + ) + normalized, routed = self.mlp.route_input(mlp_input, modality_dispatcher) + return streams, mlp_logits, normalized, routed + + def _moe_output( + self, + streams: torch.Tensor, + mlp_logits: MHCTensorTuple, + normalized: torch.Tensor, + routed: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + mlp_output = self.mlp.combine(normalized, routed, modality_dispatcher) + streams = self._connect(streams, mlp_output, "mlp", mlp_logits) + return streams.reshape(streams.shape[0], -1) + def forward( self, hidden_states: torch.Tensor, @@ -533,22 +619,15 @@ class Magi2TransformerLayer(nn.Module): modality_dispatcher: ModalityDispatcher, cp_split_sizes: list[int] | torch.Tensor, ) -> torch.Tensor: - streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size) - attention_logits = self._branch_logits(streams, "attn", modality_dispatcher) - attention_input = self._branch_input(streams, "attn", attention_logits) - attention_output = self.attention( - attention_input, - rope, - varlen_handler, - modality_dispatcher, - cp_split_sizes, + streams, attention_logits, q, k, v, gates = self._attention_input(hidden_states, rope, modality_dispatcher) + attention = self.attention.attend(q, k, v, varlen_handler, cp_split_sizes) + if not isinstance(self.mlp, Magi2MultiHeadMoELayer): + return self._dense_output(streams, attention, gates, attention_logits, modality_dispatcher) + streams, mlp_logits, normalized, routed = self._moe_input( + streams, attention, gates, attention_logits, modality_dispatcher ) - streams = self._connect(streams, attention_output, "attn", attention_logits) - mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher) - mlp_input = self._branch_input(streams, "mlp", mlp_logits) - mlp_output = self.mlp(mlp_input, modality_dispatcher) - streams = self._connect(streams, mlp_output, "mlp", mlp_logits) - return streams.reshape(streams.shape[0], -1) + routed = self.mlp.moe_mlp(routed) + return self._moe_output(streams, mlp_logits, normalized, routed, modality_dispatcher) class Magi2TransformerBlock(nn.Module): @@ -570,7 +649,6 @@ class Magi2TransformerBlock(nn.Module): modality_dispatcher: ModalityDispatcher, cp_split_sizes: list[int] | torch.Tensor, ) -> torch.Tensor: - hidden_states = hidden_states.to(self.config.params_dtype) hidden_states = modality_dispatcher.permute(hidden_states) for layer in self.layers: hidden_states = layer( @@ -630,6 +708,15 @@ class Magi2PreviewTransformer(nn.Module): return self.block.layers + def compile_regions(self, **compile_kwargs: Any) -> None: + """Compile the dense compute between the eager attention and MoE kernels.""" + + self.pre_adapter.forward = torch.compile(self.pre_adapter.forward, **compile_kwargs) + self.post_adapter.forward = torch.compile(self.post_adapter.forward, **compile_kwargs) + for layer in self.block.layers: + for name in layer.region_methods: + setattr(layer, name, torch.compile(getattr(layer, name), **compile_kwargs)) + def forward( self, x: torch.Tensor, @@ -647,6 +734,7 @@ class Magi2PreviewTransformer(nn.Module): assert dispatcher.split_sizes is not None cp_split_sizes = dispatcher.split_sizes + rope = self.pre_adapter.rope(coords_mapping) time_mask = modality_mapping == int(Modality.TIME) modality_mapping = torch.where(time_mask, int(Modality.TEXT), modality_mapping) modality_dispatcher = ModalityDispatcher(modality_mapping, 3) @@ -654,13 +742,7 @@ class Magi2PreviewTransformer(nn.Module): audio_indices = torch.nonzero(modality_mapping == int(Modality.AUDIO)).flatten() text_indices = torch.nonzero(modality_mapping == int(Modality.TEXT)).flatten() - hidden_states, rope = self.pre_adapter( - x, - coords_mapping, - video_indices, - audio_indices, - text_indices, - ) + hidden_states = self.pre_adapter(x, video_indices, audio_indices, text_indices) if time_token_sequence is not None and time_token_sequence.shape[-1] > 0: hidden_states[:, : time_token_sequence.shape[-1]] = time_token_sequence.to(hidden_states.dtype) hidden_states = self.block( diff --git a/vllm_omni/diffusion/models/magi2/pipeline_magi2.py b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py index c7c480876..910e0d074 100644 --- a/vllm_omni/diffusion/models/magi2/pipeline_magi2.py +++ b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py @@ -695,6 +695,24 @@ class Magi2Pipeline( return self.video_decoder.module + def setup_compile(self) -> None: + """Compile the transformer regions; the attention and MoE kernels stay eager.""" + + granularity = self.od_config.diffusion_compile_granularity + if granularity != "regional": + logger.warning( + "MAGI-2 compiles the transformer regions itself; diffusion_compile_granularity=%r is ignored.", + granularity, + ) + + # The mHC connections chain bf16 ops that eager rounds after every op. + # Emulating those casts keeps the compiled rounding boundaries equal. + self.transformer.compile_regions( + fullgraph=True, + dynamic=self.od_config.diffusion_compile_dynamic, + options={"emulate_precision_casts": True}, + ) + def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], diff --git a/vllm_omni/diffusion/models/magi2/sampler_magi2.py b/vllm_omni/diffusion/models/magi2/sampler_magi2.py index 994fdd3e4..35dd2dab9 100644 --- a/vllm_omni/diffusion/models/magi2/sampler_magi2.py +++ b/vllm_omni/diffusion/models/magi2/sampler_magi2.py @@ -308,29 +308,11 @@ class Magi2PreviewSampler(CFGParallelMixin): per_token_video_t = t_batch.view(-1, 1, 1, 1, 1).expand(cfg_batch, 1, video_t, video_h, video_w) per_token_audio_t = t_batch.view(-1, 1, 1).expand(cfg_batch, audio_t, 1) - audio_lengths = torch.full( - (cfg_batch,), - audio_latent_len, - device=latent.device, - dtype=torch.long, - ) - text_lengths = torch.tensor( - [txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size, - device=latent.device, - dtype=torch.long, - ) - ref_audio_lengths = torch.full( - (cfg_batch,), - ref_audio_feat_len, - device=latent.device, - dtype=torch.long, - ) - ref_video_lengths = torch.full( - (cfg_batch,), - ref_video_feat_len, - device=latent.device, - dtype=torch.long, - ) + # Lengths stay host integers; the data proxy reads them as Python ints. + audio_lengths = [audio_latent_len] * cfg_batch + text_lengths = [txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size + ref_audio_lengths = [ref_audio_feat_len] * cfg_batch + ref_video_lengths = [ref_video_feat_len] * cfg_batch return ModelInput( x_t=torch.cat((latent, latent), dim=0),