[Feature] Add support for Pipeline Parallel and integrate it into Wan 2.2 (#2322)

Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -0,0 +1,802 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import socket
+
+import pytest
+import torch
+from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers
+from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
+
+import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    get_classifier_free_guidance_rank,
+    get_pp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.parallel]
+
+
+def _find_free_port() -> str:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return str(s.getsockname()[1])
+
+
+def update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for k, v in envs_dict.items():
+        os.environ[k] = v
+
+
+# ---------------------------------------------------------------------------
+# Shared stubs used by both unit and distributed tests
+# ---------------------------------------------------------------------------
+
+
+class FakeWork:
+    """Drop-in for torch.distributed.Work that records whether wait() was called."""
+
+    def __init__(self):
+        self.waited = False
+
+    def wait(self):
+        self.waited = True
+
+
+class SimpleScheduler:
+    """Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred."""
+
+    def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False):
+        return (latents - 0.1 * noise_pred,)
+
+
+class FakeVAE:
+    def __init__(self, distributed_enabled: bool = False):
+        self.calls = 0
+        self.distributed_enabled = distributed_enabled
+
+    def decode(self, z: torch.Tensor):
+        """Original decode docstring."""
+        self.calls += 1
+        return (z + 1,)
+
+    def is_distributed_enabled(self) -> bool:
+        return self.distributed_enabled
+
+
+class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
+    """Minimal pipeline used to exercise PipelineParallelMixin.
+
+    Uses vLLM's ``make_layers`` for layer partitioning — the same utility used
+    by real DiT models — so the PP layer-split logic is exercised faithfully.
+
+    Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i``
+    is initialized identically on every rank regardless of which ranks are active,
+    allowing the distributed output to be compared against the single-GPU baseline.
+
+    Args:
+        num_layers: Total number of Linear layers.
+        dim:        Input / hidden dimension.
+        seed:       Base RNG seed; layer ``i`` uses ``seed + i``.
+        device:     Target device for layer weights (default: CPU).
+        dtype:      Target dtype for layer weights (default: float32).
+    """
+
+    def __init__(
+        self,
+        num_layers: int = 4,
+        dim: int = 64,
+        seed: int = 42,
+        device: torch.device | None = None,
+        dtype: torch.dtype = torch.float32,
+    ):
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            num_layers,
+            lambda prefix: torch.nn.Linear(dim, dim, bias=False),
+            prefix="layers",
+        )
+
+        for i, layer in enumerate(self.layers):
+            if not isinstance(layer, PPMissingLayer):
+                torch.manual_seed(seed + i)
+                torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02)
+
+        self.layers.to(device=device, dtype=dtype)
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim)
+        self.scheduler = SimpleScheduler()
+
+    def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors:
+        """Layer-split forward pass.
+
+        * First PP rank: uses ``x`` from caller kwargs.
+        * Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]``
+          (which transparently waits for the async receive).
+        * Non-last PP ranks return ``IntermediateTensors``; the last rank
+          returns the plain noise-prediction tensor.
+        """
+        if intermediate_tensors is not None:
+            x = intermediate_tensors["hidden_states"]
+
+        for i in range(self.start_layer, self.end_layer):
+            x = self.layers[i](x)
+
+        pp_group = get_pp_group()
+        if not pp_group.is_last_rank:
+            return IntermediateTensors({"hidden_states": x})
+        return x
+
+
+# ---------------------------------------------------------------------------
+# 1.  AsyncLatents – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestAsyncLatents:
+    """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
+        return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
+
+    def test_resolve_returns_wrapped_tensor(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al._resolve() is t
+
+    def test_attribute_access_resolves(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al.shape == t.shape
+        assert al.dtype == t.dtype
+
+    def test_torch_function_protocol(self):
+        """torch ops that receive an AsyncLatents should see the underlying tensor."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        mask = torch.ones_like(t)
+        result = mask * al  # triggers __torch_function__
+        torch.testing.assert_close(result, mask * t)
+
+    def test_torch_function_with_list_arg(self):
+        """__torch_function__ must unwrap AsyncLatents inside list/tuple args."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.cat([al, al], dim=0)
+        torch.testing.assert_close(result, torch.cat([t, t], dim=0))
+
+    def test_torch_tensor_conversion(self):
+        """torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy)."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.as_tensor(al)
+        assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage"
+
+    def test_handles_are_waited_before_resolve(self):
+        t = torch.randn(2, 4)
+        h1, h2 = FakeWork(), FakeWork()
+        al = self._make(t, handles=[h1, h2])
+        _ = al.shape  # trigger resolution
+        assert h1.waited and h2.waited, "Not all handles were waited on"
+
+    def test_postproc_callbacks_invoked_on_resolve(self):
+        t = torch.randn(2, 4)
+        log: list[int] = []
+        al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)])
+        _ = al.shape
+        assert log == [1, 2], f"postproc not called in order: {log}"
+
+    def test_idempotent_resolve(self):
+        """handle.wait() must not be called twice if _resolve() is called twice."""
+        t = torch.randn(2, 4)
+        h = FakeWork()
+        al = self._make(t, handles=[h])
+        _ = al.shape  # first resolve
+        h.waited = False  # reset sentinel
+        _ = al.dtype  # second resolve
+        assert not h.waited, "handle.wait() was called a second time"
+
+
+# ---------------------------------------------------------------------------
+# 2.  _sync_pp_send / diffuse wrapper – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestSyncPPSend:
+    """Verifies PipelineParallelMixin's internal PP-send flush."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline() -> PipelineParallelMixin:
+        # Instantiate a bare mixin — no layers, no distributed env needed.
+        # _sync_pp_send only touches _pp_send_work, so this is sufficient.
+        class _BarePP(PipelineParallelMixin, CFGParallelMixin):
+            pass
+
+        return _BarePP()
+
+    def test_noop_when_work_list_empty(self):
+        pipeline = self._make_pipeline()
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+    def test_waits_all_pending_handles(self):
+        pipeline = self._make_pipeline()
+        works = [FakeWork(), FakeWork(), FakeWork()]
+        pipeline._pp_send_work = works
+        pipeline._sync_pp_send()
+        assert all(w.waited for w in works), "Some handles were not waited on"
+
+    def test_clears_work_list_after_sync(self):
+        pipeline = self._make_pipeline()
+        pipeline._pp_send_work = [FakeWork()]
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+
+class TestDiffuseWrapper:
+    """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def test_diffuse_flushes_pending_sends_on_success(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                return "done"
+
+        pipeline = _DiffusePP()
+
+        assert pipeline.diffuse() == "done"
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_flushes_pending_sends_on_exception(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                raise RuntimeError("boom")
+
+        pipeline = _DiffusePP()
+
+        with pytest.raises(RuntimeError, match="boom"):
+            pipeline.diffuse()
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_wrapper_preserves_metadata(self):
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                """Original diffuse docstring."""
+                return "done"
+
+        assert _DiffusePP.diffuse.__name__ == "diffuse"
+        assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring."
+
+
+class TestVaeDecodeGuard:
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
+        class _DecodePP(PipelineParallelMixin, CFGParallelMixin):
+            def __init__(self):
+                self.vae = FakeVAE(distributed_enabled=distributed_enabled)
+
+        return _DecodePP()
+
+    @staticmethod
+    def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None:
+        monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size)
+        monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage)
+
+    def test_calls_original_decode_when_pp_disabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=1, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_calls_original_decode_on_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_skips_decode_on_non_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)
+
+        assert pipeline.vae.calls == 0
+        assert output == (None,)
+
+    def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline(distributed_enabled=True)
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_decode_wrapper_preserves_metadata(self):
+        pipeline = self._make_pipeline()
+
+        assert pipeline.vae.decode.__name__ == "decode"
+        assert pipeline.vae.decode.__doc__ == "Original decode docstring."
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_cfg_mixin():
+    with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
+
+        class _MissingCFG(PipelineParallelMixin):
+            pass
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_mro_before_cfg_mixin():
+    with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
+
+        class _WrongOrder(CFGParallelMixin, PipelineParallelMixin):
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Distributed test helpers
+# ---------------------------------------------------------------------------
+
+
+def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
+    """Initialise the distributed environment for a spawned worker."""
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": master_port,
+        }
+    )
+    init_distributed_environment()
+    return device
+
+
+def make_pipeline_and_inputs(
+    test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
+) -> tuple["MockPipelineParallel", dict, dict | None]:
+    """Create a MockPipelineParallel and seeded inputs from a test_config dict.
+
+    Must be called after ``initialize_model_parallel`` so that ``make_layers``
+    can read the PP group to determine this rank's layer slice.
+
+    Returns ``(pipeline, positive_kwargs, negative_kwargs)``.
+    ``negative_kwargs`` is ``None`` when ``do_true_cfg=False``.
+    """
+    pipeline = MockPipelineParallel(
+        num_layers=test_config["num_layers"],
+        dim=test_config["dim"],
+        seed=test_config["model_seed"],
+        device=device,
+        dtype=dtype,
+    )
+
+    torch.manual_seed(test_config["input_seed"])
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(test_config["input_seed"])
+    pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
+
+    negative_kwargs = None
+    if do_true_cfg:
+        torch.manual_seed(test_config["input_seed"] + 1)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
+        neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
+        negative_kwargs = {"x": neg_x}
+
+    return pipeline, pos_x, negative_kwargs
+
+
+# ---------------------------------------------------------------------------
+# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
+# ---------------------------------------------------------------------------
+
+
+def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=world_size)
+    pp_group = get_pp_group()
+
+    if pp_group.is_first_rank:
+        torch.manual_seed(77)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(77)
+        tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
+        handles = pp_group.isend_tensor_dict({"t": tensor})
+        for h in handles:
+            h.wait()
+        result_queue.put(("sent", tensor.cpu()))
+    else:
+        received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
+        result_queue.put(("received", received["t"].cpu()))
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
+@pytest.mark.parametrize("pp_size", [2])
+def test_isend_irecv_tensor_dict(pp_size: int):
+    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
+
+    results = {label: tensor for label, tensor in [q.get(), q.get()]}
+    torch.testing.assert_close(
+        results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 4.  predict_noise_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+_baseline_cache: dict[tuple, torch.Tensor] = {}
+
+
+def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
+    """Compute expected single-GPU output using the same MockPipelineParallel.
+
+    Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
+    Results are cached so identical configs are only computed once.
+    """
+    key = (
+        test_config["num_layers"],
+        test_config["dim"],
+        test_config["batch_size"],
+        test_config["model_seed"],
+        test_config["input_seed"],
+        test_config["cfg_scale"],
+        dtype,
+        do_true_cfg,
+    )
+    if key in _baseline_cache:
+        return _baseline_cache[key]
+
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    destroy_distributed_env()
+
+    _baseline_cache[key] = noise_pred.cpu()
+    return _baseline_cache[key]
+
+
+def predict_noise_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    dtype: torch.dtype,
+    test_config: dict,
+    result_queue,
+):
+    """Generic predict-noise worker parameterized by PP and CFG topology."""
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+    # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the non-last PP rank's async send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_last_rank:
+        assert noise_pred is not None
+        if cfg_rank == 0:
+            result_queue.put(noise_pred.cpu())
+    else:
+        assert noise_pred is None
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.float32,
+            4,
+            100,
+            1e-5,
+            1e-5,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-float32",
+        ),
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            1,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-seq_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-bfloat16",
+        ),
+        pytest.param(
+            3,
+            1,
+            False,
+            torch.bfloat16,
+            6,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
+            id="pp3-no_cfg-bfloat16",
+        ),
+    ],
+)
+def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
+    """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": num_layers,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    pp_q = manager.Queue()
+
+    world_size = pp_size * cfg_size
+    port = _find_free_port()
+    torch.multiprocessing.spawn(
+        predict_noise_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
+        nprocs=world_size,
+    )
+
+    pp_out = pp_q.get()
+
+    assert baseline_out.shape == pp_out.shape
+    torch.testing.assert_close(
+        pp_out,
+        baseline_out,
+        rtol=rtol,
+        atol=atol,
+        msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})",
+    )
+
+
+# ---------------------------------------------------------------------------
+# 5.  scheduler_step_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+
+def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
+    """Single-GPU reference: predict_noise + scheduler_step."""
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        result = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+
+    destroy_distributed_env()
+    return result.cpu()
+
+
+def scheduler_step_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    test_config: dict,
+    result_queue,
+):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        latents = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+    # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the last PP rank's async latent send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_first_rank and cfg_rank == 0:
+        resolved = latents.contiguous()
+        result_queue.put(resolved.cpu())
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            300,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            600,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-true_cfg",
+        ),
+    ],
+)
+def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": 4,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    world_size = pp_size * cfg_size
+    torch.multiprocessing.spawn(
+        scheduler_step_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
+        nprocs=world_size,
+    )
+
+    result = q.get()
+    torch.testing.assert_close(
+        result,
+        baseline,
+        rtol=0,
+        atol=0,
+        msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
+    )