[Model][Perf] Optimize LTX Ulysses sequence parallelism (#7079)

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/models/ltx2/test_ltx25_pipeline.py

@@ -5,6 +5,7 @@
 
 import json
 from dataclasses import replace
+from threading import Lock
 from types import SimpleNamespace
 from typing import Any
 
@@ -500,6 +501,32 @@ def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypa
     assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
 
 
+def test_ltx_runtime_profiles_phase_methods_separately(monkeypatch):
+    pipe = object.__new__(LTX2DistilledTwoStagePipeline)
+    pipe.enable_diffusion_pipeline_profiler = True
+    pipe._profiler_lock = Lock()
+    pipe._stage_durations = {}
+    monkeypatch.setattr(ltx2_runtime.current_omni_platform, "is_available", lambda: False)
+
+    with pipe._profile_phase_method("run_phase", "generate_lowres"):
+        pass
+    with pipe._profile_phase_method("run_phase", "refine"):
+        pass
+    with pipe._profile_phase_method("decode_phase"):
+        pass
+    output = DiffusionOutput()
+    assert pipe._attach_profiled_stage_durations(output) is output
+
+    expected_keys = {
+        "LTX2DistilledTwoStagePipeline.generate_lowres.run_phase",
+        "LTX2DistilledTwoStagePipeline.refine.run_phase",
+        "LTX2DistilledTwoStagePipeline.decode_phase",
+    }
+    assert set(pipe.stage_durations) == expected_keys
+    assert set(output.stage_durations) == expected_keys
+    assert all(duration >= 0.0 for duration in pipe.stage_durations.values())
+
+
 def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -103,7 +103,7 @@ def test_prepare_audio_latents_samples_directly_in_packed_token_space():
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
 
 
-def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
+def test_prepare_audio_latents_pads_short_audio_for_sp():
     pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
 
     latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for LTX guidance and forward-parallel behavior."""
 
@@ -85,14 +85,14 @@ class TestCFGParallelHelpers:
 
         LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size)
 
-    def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch):
+    def test_guidance_parallel_preserves_full_guidance_plan(self, monkeypatch):
         from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
         from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
         pipe = object.__new__(LTX2Pipeline)
-        req = SimpleNamespace(is_dummy_run=lambda: True)
+        req = SimpleNamespace()
         request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
 
         guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -175,7 +175,19 @@ def test_ltx_device_uint8_output_rejects_invalid_contract(video, error):
         _prepare_ltx2_video_output(video, do_normalize=True)
 
 
-def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
+@pytest.mark.parametrize(
+    "parallel_config",
+    [
+        None,
+        SimpleNamespace(ulysses_degree=4),
+        SimpleNamespace(tensor_parallel_size=2),
+        SimpleNamespace(use_hsdp=True, hsdp_shard_size=2),
+        SimpleNamespace(ulysses_degree=2, tensor_parallel_size=2),
+        SimpleNamespace(ulysses_degree=2, use_hsdp=True, hsdp_shard_size=2),
+    ],
+    ids=["single", "ulysses", "tp", "hsdp", "tp-ulysses", "hsdp-ulysses"],
+)
+def test_ltx_supported_parallel_configs_keep_output_profiling(tmp_path, monkeypatch, parallel_config):
     from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
 
     (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
@@ -193,7 +205,13 @@ def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
     monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
     monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup)
 
-    LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+    LTX2Pipeline(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            parallel_config=parallel_config,
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
 
     assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"]
     assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"]
@@ -383,14 +401,42 @@ def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expecte
         )
 
 
-def test_ltx_rejects_advanced_uaa_before_component_initialization():
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+        LTX2T2VDMD2Pipeline,
+        LTX2I2VDMD2Pipeline,
+    ],
+)
+@pytest.mark.parametrize(
+    ("parallel_config", "error"),
+    [
+        (SimpleNamespace(ulysses_mode="advanced_uaa"), "does not support ulysses_mode='advanced_uaa'"),
+        (SimpleNamespace(ring_degree=2), "pure Ulysses sequence parallelism only"),
+        (SimpleNamespace(ulysses_degree=2, ring_degree=2), "pure Ulysses sequence parallelism only"),
+        (SimpleNamespace(allgather_degree=2), "pure Ulysses sequence parallelism only"),
+    ],
+    ids=["advanced-uaa", "ring", "hybrid", "allgather-kv"],
+)
+def test_ltx_rejects_unsupported_sp_before_component_initialization(pipeline_cls, parallel_config, error, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    def unexpected_model_access(*_args, **_kwargs):
+        pytest.fail("Unsupported SP must be rejected before accessing model files or initializing components")
+
+    monkeypatch.setattr(ltx2_runtime, "detect_ltx_model_version", unexpected_model_access)
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", unexpected_model_access)
     od_config = SimpleNamespace(
         model="unused",
-        parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
+        parallel_config=parallel_config,
     )
 
-    with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
-        LTX2Pipeline(od_config=od_config)
+    with pytest.raises(ValueError, match=error):
+        pipeline_cls(od_config=od_config)
 
 
 def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
@@ -1409,39 +1455,19 @@ def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sam
         pipeline.forward(req, **direct_kwargs)
 
 
-def test_ltx_distilled_dummy_run_uses_fixed_recipe_values():
-    from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
-    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-    req = DiffusionRequestBatch(
-        [
-            OmniDiffusionRequest(
-                prompt="dummy run",
-                sampling_params=OmniDiffusionSamplingParams(
-                    height=512,
-                    width=512,
-                    num_frames=LTX2DistilledPipeline.dummy_run_num_frames,
-                    num_inference_steps=1,
-                    guidance_scale=0.0,
-                ),
-                request_id=DUMMY_DIFFUSION_REQUEST_ID,
-            )
-        ]
-    )
-    pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline)
-    seen = {}
-
-    def fake_run_recipe(_req, request_inputs, **_kwargs):
-        seen["request_inputs"] = request_inputs
-        return "dummy-output"
-
-    object.__setattr__(pipeline, "_run_recipe", fake_run_recipe)
-
-    assert pipeline.forward(req) == "dummy-output"
-    assert seen["request_inputs"].num_frames == 1
-    assert seen["request_inputs"].num_inference_steps == 8
-    assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only()
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+        LTX2T2VDMD2Pipeline,
+        LTX2I2VDMD2Pipeline,
+    ],
+)
+def test_ltx_pipelines_disable_startup_dummy_run(pipeline_cls):
+    assert pipeline_cls.dummy_run_num_frames == 0
 
 
 def test_ltx_custom_sigmas_bypass_scheduler_shifting():

tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.models.ltx2.ltx2_sequence_parallel import (
+    LTX2VideoToAudioParallelAttention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _run_video_to_audio_parity(rank: int, world_size: int, master_port: int) -> None:
+    dist.init_process_group(
+        "gloo",
+        init_method=f"tcp://127.0.0.1:{master_port}",
+        rank=rank,
+        world_size=world_size,
+        timeout=timedelta(seconds=60),
+    )
+    try:
+        # Non-contiguous SP subgroups match TP-fastest rank ordering. Different
+        # seeds across groups catch collectives that accidentally use WORLD.
+        group_count = world_size // 2
+        ulysses_group = dist.group.WORLD
+        if group_count > 1:
+            for offset in range(group_count):
+                ranks = list(range(offset, world_size, group_count))
+                group = dist.new_group(ranks)
+                if rank in ranks:
+                    ulysses_group = group
+        ulysses_rank = dist.get_rank(ulysses_group)
+        sp_group = SimpleNamespace(
+            ring_world_size=1,
+            ulysses_group=ulysses_group,
+            ulysses_world_size=2,
+            ulysses_rank=ulysses_rank,
+        )
+        strategy = LTX2VideoToAudioParallelAttention(sp_group)
+
+        for batch_size, masked in [(1, False), (2, True)]:
+            torch.manual_seed(17 + rank % group_count)
+            query = torch.randn(batch_size, 3, 4, 8)
+            global_key = torch.randn(batch_size, 6, 4, 8)
+            global_value = torch.randn(batch_size, 6, 4, 8)
+            key = global_key.chunk(2, dim=1)[ulysses_rank].contiguous()
+            value = global_value.chunk(2, dim=1)[ulysses_rank].contiguous()
+            mask = None
+            if masked:
+                mask = torch.ones(batch_size, 1, 1, 6, dtype=torch.bool)
+                mask[0, ..., -1] = False
+                mask[1, ..., -2:] = False
+            metadata = AttentionMetadata(attn_mask=mask) if masked else None
+
+            local_query, local_key, local_value, local_metadata, ctx = strategy.pre_attention(
+                query, key, value, metadata
+            )
+            assert local_metadata is metadata
+            local_output = F.scaled_dot_product_attention(
+                local_query.transpose(1, 2),
+                local_key.transpose(1, 2),
+                local_value.transpose(1, 2),
+                attn_mask=local_metadata.attn_mask if local_metadata is not None else None,
+            ).transpose(1, 2)
+            actual = strategy.post_attention(local_output, ctx)
+            expected = F.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                global_key.transpose(1, 2),
+                global_value.transpose(1, 2),
+                attn_mask=mask,
+            ).transpose(1, 2)
+
+            torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("world_size", [2, 4], ids=["world", "subgroups"])
+def test_ltx_video_to_audio_sp_matches_replicated_attention(unused_tcp_port, world_size):
+    torch.multiprocessing.spawn(
+        _run_video_to_audio_parity,
+        args=(world_size, unused_tcp_port),
+        nprocs=world_size,
+    )
+
+
+def test_ltx_video_to_audio_sp_rejects_ring_parallelism():
+    sp_group = SimpleNamespace(ring_world_size=2)
+
+    with pytest.raises(NotImplementedError, match="ring_degree must be 1"):
+        LTX2VideoToAudioParallelAttention(sp_group)

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -141,16 +141,54 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
 
 
 @pytest.mark.parametrize("rope_type", ["interleaved", "split"])
-def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
-    root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""]
+def test_ltx_sp_plan_only_shards_video_stream(rope_type):
+    plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)
+    root_plan = plan[""]
 
-    video_timestep = root_plan["timestep"]
-    audio_timestep = root_plan["audio_timestep"]
+    assert set(root_plan) == {"hidden_states", "keyframes_mask", "timestep"}
+    assert root_plan["hidden_states"].split_dim == 1
+    assert root_plan["keyframes_mask"].split_dim == 1
+    assert root_plan["timestep"].split_dim == 1
+    assert root_plan["timestep"].expected_dims == 2
+    assert not root_plan["timestep"].split_output
+    assert set(plan) == {"", "rope", "cross_attn_rope", "proj_out"}
 
-    assert video_timestep.split_dim == audio_timestep.split_dim == 1
-    assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
-    assert not video_timestep.split_output
-    assert not audio_timestep.split_output
+
+def test_ltx_attention_assigns_video_only_sp_modes(monkeypatch):
+    modes = {}
+
+    class FakeAttention(nn.Module):
+        def __init__(self, *, prefix, sequence_parallel_mode="local", **_kwargs):
+            super().__init__()
+            modes[prefix.rsplit(".", 1)[-1]] = sequence_parallel_mode
+
+    class FakeFeedForward(nn.Module):
+        def __init__(self, *_args, **_kwargs):
+            super().__init__()
+
+    monkeypatch.setattr(ltx2_transformer, "LTX2Attention", FakeAttention)
+    monkeypatch.setattr(ltx2_transformer, "LTX2FeedForward", FakeFeedForward)
+
+    LTX2VideoTransformerBlock(
+        dim=8,
+        num_attention_heads=2,
+        attention_head_dim=4,
+        cross_attention_dim=8,
+        audio_dim=8,
+        audio_num_attention_heads=2,
+        audio_attention_head_dim=4,
+        audio_cross_attention_dim=8,
+        prefix="transformer_blocks.0",
+    )
+
+    assert modes == {
+        "attn1": "video_self",
+        "audio_attn1": "local",
+        "attn2": "local",
+        "audio_attn2": "local",
+        "audio_to_video_attn": "local",
+        "video_to_audio_attn": "video_to_audio",
+    }
 
 
 def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):

tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Full-transformer numeric parity for LTX Ulysses sequence parallelism.
+
+This FP32/SDPA check verifies that Ulysses preserves the complete transformer
+mathematical path. It does not guarantee pixel-level E2E parity: production
+BF16 and compiled attention kernels can change rounding order, and diffusion
+can amplify that drift across denoising steps.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.parallel]
+
+_FP32_REL_L2 = 1e-5
+_TINY_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 4,
+    "cross_attention_dim": 8,
+    "vae_scale_factors": (1, 1, 1),
+    "pos_embed_max_pos": 4,
+    "base_height": 2,
+    "base_width": 2,
+    "audio_in_channels": 4,
+    "audio_out_channels": 4,
+    "audio_num_attention_heads": 2,
+    "audio_attention_head_dim": 4,
+    "audio_cross_attention_dim": 8,
+    "audio_scale_factor": 1,
+    "audio_pos_embed_max_pos": 4,
+    "audio_sampling_rate": 16,
+    "audio_hop_length": 4,
+    "num_layers": 1,
+    "caption_channels": 8,
+    "rope_type": "split",
+    "use_prompt_embeddings": False,
+    "use_keyframes_abs_pos_embedding": True,
+}
+
+
+def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float:
+    return ((actual - expected).double().norm() / expected.double().norm()).item()
+
+
+def _omni_config(sp_size: int) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig.from_kwargs(
+        model="test",
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            tensor_parallel_size=1,
+            sequence_parallel_size=sp_size,
+            ulysses_degree=sp_size,
+            ring_degree=1,
+            cfg_parallel_size=1,
+        ),
+        diffusion_attention_backend="TORCH_SDPA",
+    )
+
+
+def _inputs(device: torch.device) -> dict[str, torch.Tensor | int | bool]:
+    generator = torch.Generator().manual_seed(123)
+    return {
+        "hidden_states": torch.randn(1, 8, 4, generator=generator, device="cpu").to(device),
+        "audio_hidden_states": torch.randn(1, 4, 4, generator=generator, device="cpu").to(device),
+        "encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
+        "audio_encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
+        "timestep": torch.linspace(0.1, 0.8, 8, device=device).view(1, 8),
+        "audio_timestep": torch.linspace(0.2, 0.5, 4, device=device).view(1, 4),
+        "keyframes_mask": torch.tensor([[[1.0], [0.0], [1.0], [0.0], [0.0], [1.0], [0.0], [0.0]]], device=device),
+        "sigma": torch.tensor([0.1], device=device),
+        "audio_sigma": torch.tensor([0.2], device=device),
+        "num_frames": 2,
+        "height": 2,
+        "width": 2,
+        "audio_num_frames": 4,
+        "return_dict": False,
+    }
+
+
+def _run_transformer_parity(rank: int, world_size: int, master_port: int) -> None:
+    device = torch.device(f"{current_omni_platform.device_type}:{rank}")
+    current_omni_platform.set_device(device)
+    os.environ.update(
+        {
+            "RANK": str(rank),
+            "LOCAL_RANK": str(rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+
+    try:
+        init_distributed_environment()
+        initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
+
+        sp_config = _omni_config(world_size)
+        torch.manual_seed(42)
+        with set_forward_context(omni_diffusion_config=sp_config), set_current_diffusion_config(sp_config):
+            model = LTX2VideoTransformer3DModel(**_TINY_CONFIG).to(device).eval()
+        with torch.no_grad():
+            for _, parameter in sorted(model.named_parameters()):
+                torch.nn.init.normal_(parameter, mean=0.0, std=0.02)
+
+        inputs = _inputs(device)
+        reference_config = _omni_config(1)
+        with torch.no_grad(), set_forward_context(omni_diffusion_config=reference_config):
+            expected_video, expected_audio = model(**inputs)
+
+        apply_sequence_parallel(
+            model,
+            SequenceParallelConfig(ulysses_degree=world_size, ring_degree=1),
+            model._sp_plan,
+        )
+        with torch.no_grad(), set_forward_context(omni_diffusion_config=sp_config):
+            get_forward_context().sp_plan_hooks_applied = True
+            actual_video, actual_audio = model(**inputs)
+
+        v2a_attention = model.transformer_blocks[0].video_to_audio_attn.attn
+        assert v2a_attention._video_to_audio_strategy is not None
+        assert torch.isfinite(expected_video).all()
+        assert torch.isfinite(expected_audio).all()
+        for modality, actual, expected in (
+            ("video", actual_video, expected_video),
+            ("audio", actual_audio, expected_audio),
+        ):
+            drift = _relative_l2(actual, expected)
+            assert drift <= _FP32_REL_L2, (
+                f"LTX SP{world_size} {modality} output drifted rel_l2={drift:.3e} from SP1 (bound {_FP32_REL_L2:.1e})"
+            )
+    finally:
+        destroy_distributed_env()
+
+
+@pytest.mark.full_model
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ltx_ulysses_transformer_matches_sp1(unused_tcp_port) -> None:
+    """Exercise real plan hooks, RoPE/timestep sharding, and attention routing."""
+    world_size = 2
+    torch.multiprocessing.spawn(
+        _run_transformer_parity,
+        args=(world_size, unused_tcp_port),
+        nprocs=world_size,
+    )