anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Browse Source 
Signed-off-by: wangyu <410167048@qq.com>
This commit is contained in:
wangyu authored and GitHub committed 2026-07-29 23:53:03 +08:00
1 parent 925347332f
commit 7c17111ff5
158 files changed
+3479 -4130
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/distributed/test_cfg_parallel.py
+774 -560
View File
File diff suppressed because it is too large. Load diff
tests/diffusion/distributed/test_comm.py
+333 -276
View File
@@ -1,12 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for SeqAllToAll4D and SeqAllToAll5D communication primitives."""
"""Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.
CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
checks on real multi-GPU NCCL collectives.
"""
from __future__ import annotations
import os
from typing import Literal
import pytest
import torch
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
@@ -16,56 +24,51 @@ from vllm_omni.diffusion.distributed.parallel_state import (
)
from vllm_omni.platforms import current_omni_platform
_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
def update_environment_variables(envs_dict: dict[str, str]):
"""Update multiple environment variables with logging."""
for k, v in envs_dict.items():
os.environ[k] = v
DeviceKind = Literal["cpu", "cuda"]
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("batch_size", [2])
@pytest.mark.parametrize("seq_len_per_rank", [8])
@pytest.mark.parametrize("num_heads", [8])
@pytest.mark.parametrize("head_size", [32])
@pytest.mark.parametrize("use_sync", [False, True])
def test_4d_identity(
def _update_environment_variables(envs_dict: dict[str, str]) -> None:
for key, value in envs_dict.items():
os.environ[key] = value
def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
if device_kind == "cpu":
return torch.device("cpu")
return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
def _init_worker(
local_rank: int,
world_size: int,
dtype: torch.dtype,
batch_size: int,
seq_len_per_rank: int,
num_heads: int,
head_size: int,
use_sync: bool,
):
"""Test that two consecutive all-to-all operations return the original input."""
# Skip if not enough GPUs available
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
# Ensure num_heads is divisible by world_size
if num_heads % world_size != 0:
pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
# Run test with multiprocessing spawn
torch.multiprocessing.spawn(
_test_4d_identity_worker,
args=(
world_size,
dtype,
batch_size,
seq_len_per_rank,
num_heads,
head_size,
use_sync,
),
nprocs=world_size,
master_port: int,
device_kind: DeviceKind,
) -> None:
_update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": str(master_port),
}
)
backend = "gloo" if device_kind == "cpu" else None
init_distributed_environment(backend=backend)
def _test_4d_identity_worker(
def _close_tolerance(dtype: torch.dtype) -> tuple[float, float]:
if dtype == torch.bfloat16:
return 1e-3, 1e-3
if dtype == torch.float16:
return 1e-3, 1e-3
return 1e-5, 1e-5
def _run_4d_identity(
local_rank: int,
world_size: int,
dtype: torch.dtype,
@@ -74,29 +77,17 @@ def _test_4d_identity_worker(
num_heads: int,
head_size: int,
use_sync: bool,
):
"""Worker function for test_4d_identity."""
# Set device
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
device_kind: DeviceKind,
master_port: int,
) -> None:
device = _worker_device(local_rank, device_kind)
if device_kind == "cuda":
current_omni_platform.set_device(device)
# Set environment variables for distributed training
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "29500",
}
)
_init_worker(local_rank, world_size, master_port, device_kind)
initialize_model_parallel(ulysses_degree=world_size)
sp_group = get_sp_group().ulysses_group
# Initialize distributed environment
init_distributed_environment()
initialize_model_parallel(ulysses_degree=world_size) # test ulysses sp by default
sp_group = get_sp_group().ulysses_group # get ulysses sp group not ring sp group
# Create input tensor: (bs, seqlen/P, hc, hs)
torch.manual_seed(42 + local_rank)
input_tensor = torch.randn(
batch_size,
@@ -106,100 +97,36 @@ def _test_4d_identity_worker(
dtype=dtype,
device=device,
)
# Save original input for comparison
original_input = input_tensor.clone()
# First all-to-all: (bs, seqlen/P, hc, hs) -> (bs, seqlen, hc/P, hs)
intermediate = SeqAllToAll4D.apply(
sp_group,
input_tensor,
2, # scatter head dimension
1, # gather sequence dimension
use_sync,
)
# Verify intermediate shape
expected_shape = (
intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync)
expected_intermediate_shape = (
batch_size,
seq_len_per_rank * world_size,
num_heads // world_size,
head_size,
)
assert intermediate.shape == expected_shape, (
f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
assert intermediate.shape == expected_intermediate_shape, (
f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
)
# Second all-to-all: (bs, seqlen, hc/P, hs) -> (bs, seqlen/P, hc, hs)
output = SeqAllToAll4D.apply(
sp_group,
intermediate,
1, # scatter sequence dimension
2, # gather head dimension
use_sync,
)
# Verify output shape matches input
output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync)
assert output.shape == original_input.shape, (
f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
)
# Verify output matches original input
rtol, atol = _close_tolerance(dtype)
torch.testing.assert_close(
output,
original_input,
rtol=1e-5,
atol=1e-5,
rtol=rtol,
atol=atol,
msg="Output does not match original input after two all-to-all operations",
)
# Cleanup distributed environment
destroy_distributed_env()
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("batch_size", [2])
@pytest.mark.parametrize("seq_len_per_rank", [8])
@pytest.mark.parametrize("num_heads", [8])
@pytest.mark.parametrize("head_size", [32])
@pytest.mark.parametrize("use_sync", [False, True])
def test_5d_identity(
world_size: int,
dtype: torch.dtype,
batch_size: int,
seq_len_per_rank: int,
num_heads: int,
head_size: int,
use_sync: bool,
):
"""Test that two consecutive all-to-all operations return the original input."""
# Skip if not enough GPUs available
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
# Ensure num_heads is divisible by world_size
if num_heads % world_size != 0:
pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
# Run test with multiprocessing spawn
torch.multiprocessing.spawn(
_test_5d_identity_worker,
args=(
world_size,
dtype,
batch_size,
seq_len_per_rank,
num_heads,
head_size,
use_sync,
),
nprocs=world_size,
)
def _test_5d_identity_worker(
def _run_5d_identity(
local_rank: int,
world_size: int,
dtype: torch.dtype,
@@ -208,199 +135,329 @@ def _test_5d_identity_worker(
num_heads: int,
head_size: int,
use_sync: bool,
):
"""Worker function for test_5d_identity."""
# Set device
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
device_kind: DeviceKind,
master_port: int,
) -> None:
device = _worker_device(local_rank, device_kind)
if device_kind == "cuda":
current_omni_platform.set_device(device)
# Set environment variables for distributed training
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "29500",
}
)
_init_worker(local_rank, world_size, master_port, device_kind)
initialize_model_parallel(ulysses_degree=world_size)
sp_group = get_sp_group().ulysses_group
# Initialize distributed environment
init_distributed_environment()
initialize_model_parallel(ulysses_degree=world_size) # test ulysses sp by default
sp_group = get_sp_group().ulysses_group # get ulysses sp group not ring sp group
# Create input tensor: (bs, seqlen/P, 3, hc, hs)
# The '3' dimension is for Q, K, V
torch.manual_seed(42 + local_rank)
input_tensor = torch.randn(
batch_size,
seq_len_per_rank,
3, # Q, K, V
3,
num_heads,
head_size,
dtype=dtype,
device=device,
)
# Save original input for comparison
original_input = input_tensor.clone()
# First all-to-all: (bs, seqlen/P, 3, hc, hs) -> (bs, seqlen, 3, hc/P, hs)
intermediate = SeqAllToAll5D.apply(
sp_group,
input_tensor,
3, # scatter head dimension
1, # gather sequence dimension
use_sync,
)
# Verify intermediate shape
expected_shape = (
intermediate = SeqAllToAll5D.apply(sp_group, input_tensor, 3, 1, use_sync)
expected_intermediate_shape = (
batch_size,
seq_len_per_rank * world_size,
3,
num_heads // world_size,
head_size,
)
assert intermediate.shape == expected_shape, (
f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
assert intermediate.shape == expected_intermediate_shape, (
f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
)
# Second all-to-all: (bs, seqlen, 3, hc/P, hs) -> (bs, seqlen/P, 3, hc, hs)
output = SeqAllToAll5D.apply(
sp_group,
intermediate,
1, # scatter sequence dimension
3, # gather head dimension
use_sync,
)
# Verify output shape matches input
output = SeqAllToAll5D.apply(sp_group, intermediate, 1, 3, use_sync)
assert output.shape == original_input.shape, (
f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
)
# Verify output matches original input
rtol, atol = _close_tolerance(dtype)
torch.testing.assert_close(
output,
original_input,
rtol=1e-5,
atol=1e-5,
rtol=rtol,
atol=atol,
msg="Output does not match original input after two all-to-all operations",
)
# Cleanup distributed environment
destroy_distributed_env()
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("batch_size", [2])
@pytest.mark.parametrize("num_heads", [8])
@pytest.mark.parametrize("head_size", [128])
def test_ring_p2p(
world_size: int,
dtype: torch.dtype,
batch_size: int,
num_heads: int,
head_size: int,
):
"""Test Ring P2P communication (send_recv)."""
# Skip if not enough GPUs available
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
torch.multiprocessing.spawn(
_test_ring_p2p_worker,
args=(world_size, dtype, batch_size, num_heads, head_size),
nprocs=world_size,
)
def _test_ring_p2p_worker(
def _run_ring_p2p(
local_rank: int,
world_size: int,
dtype: torch.dtype,
batch_size: int,
num_heads: int,
head_size: int,
):
"""Worker for Ring P2P test."""
import sys
device_kind: DeviceKind,
master_port: int,
) -> None:
device = _worker_device(local_rank, device_kind)
if device_kind == "cuda":
current_omni_platform.set_device(device)
# Set device
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
# Set env vars
# Use a different port to avoid conflict with other tests if run in parallel
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "29501",
}
)
# Init distributed
try:
init_distributed_environment()
# Ring degree = world_size to test ring group
_init_worker(local_rank, world_size, master_port, device_kind)
initialize_model_parallel(ring_degree=world_size)
sp_group = get_sp_group()
print(f"[Rank {local_rank}] Initialized. Ring group size: {sp_group.ring_group.size()}")
sys.stdout.flush()
# Create RingComm
comm = RingComm(sp_group.ring_group)
# Create tensor: rank-specific data
# (batch, num_heads, head_size)
# Fill with rank value + 1 to avoid 0 and make verification easy
input_tensor = torch.full(
(batch_size, num_heads, head_size), fill_value=float(local_rank + 1), dtype=dtype, device=device
(batch_size, num_heads, head_size),
fill_value=float(local_rank + 1),
dtype=dtype,
device=device,
)
print(f"[Rank {local_rank}] Input sum: {input_tensor.sum().item()}")
sys.stdout.flush()
# Send input, receive from prev
# RingComm.send_recv sends to next, receives from prev
t0 = __import__("time").time()
recv_tensor = comm.send_recv(input_tensor)
comm.commit()
comm.wait()
t1 = __import__("time").time()
print(f"[Rank {local_rank}] Communication done in {t1 - t0:.4f}s")
# Verify
# Expected value: from (rank - 1) % world_size
prev_rank = (local_rank - 1 + world_size) % world_size
expected_value = float(prev_rank + 1)
recv_sum = recv_tensor.sum().item()
print(f"[Rank {local_rank}] Received sum: {recv_sum}, Expected value: {expected_value}")
sys.stdout.flush()
expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value)
# Use a slightly loose tolerance for bfloat16
rtol, atol = _close_tolerance(dtype)
torch.testing.assert_close(
recv_tensor, expected_tensor, rtol=1e-3, atol=1e-3, msg=f"[Rank {local_rank}] Data mismatch!"
recv_tensor,
expected_tensor,
rtol=rtol,
atol=atol,
msg=f"[Rank {local_rank}] Ring P2P data mismatch",
)
print(f"[Rank {local_rank}] Verification PASSED")
except Exception as e:
print(f"[Rank {local_rank}] FAILED with error: {e}")
import traceback
traceback.print_exc()
raise e
finally:
destroy_distributed_env()
def _spawn_4d_identity(
*,
world_size: int,
dtype: torch.dtype,
batch_size: int,
seq_len_per_rank: int,
num_heads: int,
head_size: int,
use_sync: bool,
device_kind: DeviceKind,
master_port: int,
) -> None:
torch.multiprocessing.spawn(
_run_4d_identity,
args=(
world_size,
dtype,
batch_size,
seq_len_per_rank,
num_heads,
head_size,
use_sync,
device_kind,
master_port,
),
nprocs=world_size,
)
def _spawn_5d_identity(
*,
world_size: int,
dtype: torch.dtype,
batch_size: int,
seq_len_per_rank: int,
num_heads: int,
head_size: int,
use_sync: bool,
device_kind: DeviceKind,
master_port: int,
) -> None:
torch.multiprocessing.spawn(
_run_5d_identity,
args=(
world_size,
dtype,
batch_size,
seq_len_per_rank,
num_heads,
head_size,
use_sync,
device_kind,
master_port,
),
nprocs=world_size,
)
def _spawn_ring_p2p(
*,
world_size: int,
dtype: torch.dtype,
batch_size: int,
num_heads: int,
head_size: int,
device_kind: DeviceKind,
master_port: int,
) -> None:
torch.multiprocessing.spawn(
_run_ring_p2p,
args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
nprocs=world_size,
)
def _require_heads_divisible(num_heads: int, world_size: int) -> None:
if num_heads % world_size != 0:
pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
def _require_gpus(world_size: int) -> None:
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
# ---------------------------------------------------------------------------
# CPU: gloo collectives on CPU tensors (no GPU)
# ---------------------------------------------------------------------------
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("use_sync", [False, True])
def test_4d_identity(world_size: int, use_sync: bool):
_require_heads_divisible(8, world_size)
_spawn_4d_identity(
world_size=world_size,
dtype=torch.float32,
batch_size=2,
seq_len_per_rank=8,
num_heads=8,
head_size=32,
use_sync=use_sync,
device_kind="cpu",
master_port=29610,
)
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
@pytest.mark.parametrize("use_sync", [False, True])
def test_5d_identity(world_size: int, use_sync: bool):
_require_heads_divisible(8, world_size)
_spawn_5d_identity(
world_size=world_size,
dtype=torch.float32,
batch_size=2,
seq_len_per_rank=8,
num_heads=8,
head_size=32,
use_sync=use_sync,
device_kind="cpu",
master_port=29611,
)
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
def test_ring_p2p(world_size: int):
_spawn_ring_p2p(
world_size=world_size,
dtype=torch.float32,
batch_size=2,
num_heads=8,
head_size=128,
device_kind="cpu",
master_port=29612,
)
# ---------------------------------------------------------------------------
# Nightly: same checks on real multi-GPU NCCL collectives
# ---------------------------------------------------------------------------
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
"world_size",
[
pytest.param(2, marks=_L4_TWO_GPU),
pytest.param(4, marks=_L4_FOUR_GPU),
],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("use_sync", [False, True])
def test_4d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
_require_gpus(world_size)
_require_heads_divisible(8, world_size)
_spawn_4d_identity(
world_size=world_size,
dtype=dtype,
batch_size=2,
seq_len_per_rank=8,
num_heads=8,
head_size=32,
use_sync=use_sync,
device_kind="cuda",
master_port=29500,
)
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
"world_size",
[
pytest.param(2, marks=_L4_TWO_GPU),
pytest.param(4, marks=_L4_FOUR_GPU),
],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("use_sync", [False, True])
def test_5d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
_require_gpus(world_size)
_require_heads_divisible(8, world_size)
_spawn_5d_identity(
world_size=world_size,
dtype=dtype,
batch_size=2,
seq_len_per_rank=8,
num_heads=8,
head_size=32,
use_sync=use_sync,
device_kind="cuda",
master_port=29502,
)
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
"world_size",
[
pytest.param(2, marks=_L4_TWO_GPU),
pytest.param(4, marks=_L4_FOUR_GPU),
],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
_require_gpus(world_size)
_spawn_ring_p2p(
world_size=world_size,
dtype=dtype,
batch_size=2,
num_heads=8,
head_size=128,
device_kind="cuda",
master_port=29501,
)
tests/diffusion/distributed/test_pipeline_parallel.py
+366 -185
View File
@@ -1,8 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for PipelineParallelMixin and related PP communication helpers.
CPU unit tests exercise mixin logic without a process group. Distributed CPU
tests use gloo on CPU tensors; nightly parity tests run the same checks on
real multi-GPU NCCL collectives.
"""
from __future__ import annotations
import os
import socket
from typing import Literal
import pytest
import torch
@@ -11,6 +20,7 @@ from vllm.sequence import IntermediateTensors
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
@@ -22,7 +32,12 @@ from vllm_omni.diffusion.distributed.parallel_state import (
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.parallel]
_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
DeviceKind = Literal["cpu", "cuda"]
_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _find_free_port() -> str:
@@ -36,6 +51,22 @@ def update_environment_variables(envs_dict: dict[str, str]) -> None:
os.environ[k] = v
def _cleanup_distributed() -> None:
"""Destroy omni distributed state and clear test-process leftovers.
``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP``
aliased to the destroyed group. Parent-process baselines must clear that
dangling reference (and RANK/MASTER_* env) so later unit tests that call
``vllm.distributed.initialize_model_parallel`` do not see stale state.
"""
destroy_distributed_env()
import vllm.distributed.parallel_state as vllm_parallel_state
vllm_parallel_state._PP = None
for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
os.environ.pop(key, None)
# ---------------------------------------------------------------------------
# Shared stubs used by both unit and distributed tests
# ---------------------------------------------------------------------------
@@ -143,7 +174,7 @@ class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
class TestAsyncLatents:
"""Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
pytestmark = [pytest.mark.cpu]
pytestmark = _UNIT_MARKS
def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
@@ -214,7 +245,7 @@ class TestAsyncLatents:
class TestSyncPPSend:
"""Verifies PipelineParallelMixin's internal PP-send flush."""
pytestmark = [pytest.mark.cpu]
pytestmark = _UNIT_MARKS
@staticmethod
def _make_pipeline() -> PipelineParallelMixin:
@@ -247,7 +278,7 @@ class TestSyncPPSend:
class TestDiffuseWrapper:
"""Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
pytestmark = [pytest.mark.cpu]
pytestmark = _UNIT_MARKS
def test_diffuse_flushes_pending_sends_on_success(self):
work = FakeWork()
@@ -289,7 +320,7 @@ class TestDiffuseWrapper:
class TestVaeDecodeGuard:
pytestmark = [pytest.mark.cpu]
pytestmark = _UNIT_MARKS
@staticmethod
def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
@@ -351,6 +382,8 @@ class TestVaeDecodeGuard:
assert pipeline.vae.decode.__doc__ == "Original decode docstring."
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_pipeline_parallel_requires_cfg_mixin():
with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
@@ -359,6 +392,8 @@ def test_pipeline_parallel_requires_cfg_mixin():
pass
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_pipeline_parallel_requires_mro_before_cfg_mixin():
with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
@@ -372,10 +407,22 @@ def test_pipeline_parallel_requires_mro_before_cfg_mixin():
# ---------------------------------------------------------------------------
def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
if device_kind == "cpu":
return torch.device("cpu")
return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
def init_dist(
local_rank: int,
world_size: int,
master_port: str,
device_kind: DeviceKind,
) -> 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)
device = _worker_device(local_rank, device_kind)
if device_kind == "cuda":
current_omni_platform.set_device(device)
update_environment_variables(
{
"RANK": str(local_rank),
@@ -385,13 +432,31 @@ def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.devic
"MASTER_PORT": master_port,
}
)
init_distributed_environment()
backend = "gloo" if device_kind == "cpu" else None
init_distributed_environment(backend=backend)
return device
def _mp_backend(device_kind: DeviceKind) -> str | None:
"""Backend for model-parallel process groups.
``GroupCoordinator.all_gather`` always uses ``device_group``. CPU tensors
therefore require gloo device groups; otherwise a CUDA host's default NCCL
backend breaks PP+CFG CPU tests at CFG all_gather (pp2-cfg2).
isend/irecv of CPU tensors can already fall back to ``cpu_group``.
"""
return "gloo" if device_kind == "cpu" else None
def _seed_device(input_seed: int, device: torch.device) -> None:
torch.manual_seed(input_seed)
if device.type == "cuda":
torch.cuda.manual_seed_all(input_seed)
def make_pipeline_and_inputs(
test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
) -> tuple["MockPipelineParallel", dict, dict | None]:
) -> 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``
@@ -408,16 +473,12 @@ def make_pipeline_and_inputs(
dtype=dtype,
)
torch.manual_seed(test_config["input_seed"])
if torch.cuda.is_available():
torch.cuda.manual_seed_all(test_config["input_seed"])
_seed_device(test_config["input_seed"], device)
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)
_seed_device(test_config["input_seed"] + 1, device)
neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
negative_kwargs = {"x": neg_x}
@@ -425,19 +486,23 @@ def make_pipeline_and_inputs(
# ---------------------------------------------------------------------------
# 3. isend_tensor_dict / irecv_tensor_dict (2 GPUs)
# 3. isend_tensor_dict / irecv_tensor_dict
# ---------------------------------------------------------------------------
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)
def isend_irecv_worker(
local_rank: int,
world_size: int,
master_port: str,
device_kind: DeviceKind,
result_queue,
):
device = init_dist(local_rank, world_size, master_port, device_kind)
initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind))
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)
_seed_device(77, device)
tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
handles = pp_group.isend_tensor_dict({"t": tensor})
for h in handles:
@@ -449,27 +514,42 @@ def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, resul
if torch.distributed.is_initialized():
torch.distributed.barrier()
destroy_distributed_env()
_cleanup_distributed()
@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."""
def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None:
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)
torch.multiprocessing.spawn(
isend_irecv_worker,
args=(pp_size, master_port, device_kind, 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"
)
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@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."""
_run_isend_irecv(pp_size, device_kind="cpu", master_port=_find_free_port())
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)])
def test_isend_irecv_tensor_dict_parity(pp_size: int):
"""Nightly: isend/irecv on real multi-GPU NCCL."""
_run_isend_irecv(pp_size, device_kind="cuda", master_port=_find_free_port())
# ---------------------------------------------------------------------------
# 4. predict_noise_maybe_with_cfg
# ---------------------------------------------------------------------------
@@ -477,12 +557,13 @@ def test_isend_irecv_tensor_dict(pp_size: int):
_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.
"""
def compute_single_gpu_baseline(
test_config: dict,
dtype: torch.dtype,
do_true_cfg: bool,
device_kind: DeviceKind,
) -> torch.Tensor:
"""Compute expected single-process output using the same MockPipelineParallel."""
key = (
test_config["num_layers"],
test_config["dim"],
@@ -492,27 +573,29 @@ def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_c
test_config["cfg_scale"],
dtype,
do_true_cfg,
device_kind,
)
if key in _baseline_cache:
return _baseline_cache[key]
device = init_dist(0, 1, _find_free_port())
initialize_model_parallel(pipeline_parallel_size=1)
device = init_dist(0, 1, _find_free_port(), device_kind)
try:
initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
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,
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
test_config, dtype, device, do_true_cfg=do_true_cfg
)
destroy_distributed_env()
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,
)
finally:
_cleanup_distributed()
_baseline_cache[key] = noise_pred.cpu()
return _baseline_cache[key]
@@ -527,11 +610,16 @@ def predict_noise_worker(
do_true_cfg: bool,
dtype: torch.dtype,
test_config: dict,
device_kind: DeviceKind,
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)
device = init_dist(local_rank, world_size, master_port, device_kind)
initialize_model_parallel(
pipeline_parallel_size=pp_size,
cfg_parallel_size=cfg_size,
backend=_mp_backend(device_kind),
)
pp_group = get_pp_group()
cfg_rank = get_classifier_free_guidance_rank()
@@ -548,8 +636,6 @@ def predict_noise_worker(
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:
@@ -561,77 +647,22 @@ def predict_noise_worker(
if torch.distributed.is_initialized():
torch.distributed.barrier()
destroy_distributed_env()
_cleanup_distributed()
@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."""
def _run_predict_noise(
*,
pp_size: int,
cfg_size: int,
do_true_cfg: bool,
dtype: torch.dtype,
num_layers: int,
input_seed: int,
rtol: float,
atol: float,
device_kind: DeviceKind,
master_port: str,
) -> None:
test_config = {
"num_layers": num_layers,
"dim": 64,
@@ -641,22 +672,20 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
"input_seed": input_seed,
}
baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
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),
args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
nprocs=world_size,
)
pp_out = pp_q.get()
assert baseline_out.shape == pp_out.shape
torch.testing.assert_close(
pp_out,
@@ -667,35 +696,157 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
)
_CPU_PREDICT_NOISE_CASES = [
pytest.param(2, 1, False, 4, 100, 1e-5, 1e-5, id="pp2-no_cfg"),
pytest.param(2, 1, True, 4, 100, 1e-5, 1e-5, id="pp2-seq_cfg"),
pytest.param(2, 2, True, 4, 100, 1e-5, 1e-5, id="pp2-cfg2"),
pytest.param(3, 1, False, 6, 100, 1e-5, 1e-5, id="pp3-no_cfg"),
]
_GPU_PREDICT_NOISE_CASES = [
pytest.param(
2,
1,
False,
torch.float32,
4,
100,
1e-5,
1e-5,
marks=_L4_TWO_GPU,
id="pp2-no_cfg-float32",
),
pytest.param(
2,
1,
False,
torch.bfloat16,
4,
100,
1e-2,
1e-2,
marks=_L4_TWO_GPU,
id="pp2-no_cfg-bfloat16",
),
pytest.param(
2,
1,
True,
torch.bfloat16,
4,
100,
1e-2,
1e-2,
marks=_L4_TWO_GPU,
id="pp2-seq_cfg-bfloat16",
),
pytest.param(
2,
2,
True,
torch.bfloat16,
4,
100,
1e-2,
1e-2,
marks=_L4_FOUR_GPU,
id="pp2-cfg2-bfloat16",
),
pytest.param(
3,
1,
False,
torch.bfloat16,
6,
100,
1e-2,
1e-2,
marks=_L4_THREE_GPU,
id="pp3-no_cfg-bfloat16",
),
]
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize(
"pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol",
_CPU_PREDICT_NOISE_CASES,
)
def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol):
"""predict_noise_maybe_with_cfg output matches baseline across PP / CFG topologies."""
_run_predict_noise(
pp_size=pp_size,
cfg_size=cfg_size,
do_true_cfg=do_true_cfg,
dtype=torch.float32,
num_layers=num_layers,
input_seed=input_seed,
rtol=rtol,
atol=atol,
device_kind="cpu",
master_port=_find_free_port(),
)
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
"pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
_GPU_PREDICT_NOISE_CASES,
)
def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
"""Nightly: predict_noise parity on real multi-GPU NCCL."""
_run_predict_noise(
pp_size=pp_size,
cfg_size=cfg_size,
do_true_cfg=do_true_cfg,
dtype=dtype,
num_layers=num_layers,
input_seed=input_seed,
rtol=rtol,
atol=atol,
device_kind="cuda",
master_port=_find_free_port(),
)
# ---------------------------------------------------------------------------
# 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)
def compute_scheduler_step_baseline(
test_config: dict,
do_true_cfg: bool,
device_kind: DeviceKind,
) -> torch.Tensor:
"""Single-process reference: predict_noise + scheduler_step."""
device = init_dist(0, 1, _find_free_port(), device_kind)
try:
initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
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
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)
destroy_distributed_env()
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
)
finally:
_cleanup_distributed()
return result.cpu()
@@ -707,10 +858,15 @@ def scheduler_step_worker(
cfg_size: int,
do_true_cfg: bool,
test_config: dict,
device_kind: DeviceKind,
result_queue,
):
device = init_dist(local_rank, world_size, master_port)
initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
device = init_dist(local_rank, world_size, master_port, device_kind)
initialize_model_parallel(
pipeline_parallel_size=pp_size,
cfg_parallel_size=cfg_size,
backend=_mp_backend(device_kind),
)
pp_group = get_pp_group()
cfg_rank = get_classifier_free_guidance_rank()
@@ -732,43 +888,25 @@ def scheduler_step_worker(
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())
result_queue.put(latents.contiguous().cpu())
if torch.distributed.is_initialized():
torch.distributed.barrier()
destroy_distributed_env()
_cleanup_distributed()
@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."""
def _run_scheduler_step(
*,
pp_size: int,
cfg_size: int,
do_true_cfg: bool,
input_seed: int,
device_kind: DeviceKind,
master_port: str,
) -> None:
test_config = {
"num_layers": 4,
"dim": 64,
@@ -778,17 +916,16 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
"input_seed": input_seed,
}
baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
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),
args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
nprocs=world_size,
)
@@ -800,3 +937,47 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
atol=0,
msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
)
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@pytest.mark.parametrize(
"pp_size, cfg_size, do_true_cfg, input_seed",
[
pytest.param(2, 1, False, 300, id="pp2-no_cfg"),
pytest.param(2, 2, True, 600, 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 baseline across PP / CFG topologies."""
_run_scheduler_step(
pp_size=pp_size,
cfg_size=cfg_size,
do_true_cfg=do_true_cfg,
input_seed=input_seed,
device_kind="cpu",
master_port=_find_free_port(),
)
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.parametrize(
"pp_size, cfg_size, do_true_cfg, input_seed",
[
pytest.param(2, 1, False, 300, marks=_L4_TWO_GPU, id="pp2-no_cfg"),
pytest.param(2, 2, True, 600, marks=_L4_FOUR_GPU, id="pp2-cfg2-true_cfg"),
],
)
def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed):
"""Nightly: scheduler_step parity on real multi-GPU NCCL."""
_run_scheduler_step(
pp_size=pp_size,
cfg_size=cfg_size,
do_true_cfg=do_true_cfg,
input_seed=input_seed,
device_kind="cuda",
master_port=_find_free_port(),
)
tests/diffusion/distributed/test_zimage_parallelism.py → tests/diffusion/distributed/test_tensor_parallel.py
+20 -110
View File
@@ -1,16 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Z-Image end-to-end tests for diffusion parallelism.
"""DiT tensor-parallel parity and perf checks for diffusion models.
This file currently covers:
- DiT tensor parallelism (TP=2) vs TP=1.
- VAE patch parallelism (vae_patch_parallel_size=2) vs baseline on TP=2.
Note: CUDA-only (>=2 GPUs). We use `enforce_eager=False` (default) to enable
`torch.compile`.
Compares TP=2 against TP=1 on output correctness, latency, and peak memory.
Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
@@ -31,22 +29,15 @@ from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
PROMPT = "a photo of a cat sitting on a laptop keyboard"
def _get_zimage_model() -> str:
# Allow overriding the model for local/offline environments.
# Can be either a HuggingFace repo id or a local path.
return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo"
def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
"""Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
return torch.from_numpy(arr)
def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
"""Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1]."""
ta = _pil_to_float_rgb_tensor(a)
tb = _pil_to_float_rgb_tensor(b)
assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
@@ -71,7 +62,7 @@ def _extract_single_image(outputs) -> Image.Image:
return images[0]
def _run_zimage_generate(
def _run_generate(
*,
tp_size: int,
height: int,
@@ -79,8 +70,6 @@ def _run_zimage_generate(
num_inference_steps: int,
seed: int,
enforce_eager: bool,
vae_use_tiling: bool = False,
vae_patch_parallel_size: int = 1,
num_requests: int = 4,
) -> tuple[Image.Image, float, float]:
if num_requests < 2:
@@ -91,25 +80,11 @@ def _run_zimage_generate(
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
try:
# Each run needs a distinct DiffusionParallelConfig; use OmniRunner per call (not the
# parametrized omni_runner fixture, which is fixed per module).
with OmniRunner(
_get_zimage_model(),
parallel_config=DiffusionParallelConfig(
tensor_parallel_size=tp_size,
vae_patch_parallel_size=vae_patch_parallel_size,
),
DEFAULT_MODEL,
parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
enforce_eager=enforce_eager,
vae_use_tiling=vae_use_tiling,
) as runner:
# NOTE: Omni closes itself when a generate() call is exhausted.
# To avoid measuring teardown time (process shutdown, memory cleanup),
# we measure the latency to produce *subsequent* outputs within a single
# generator run.
#
# This also serves as a warmup: the first output may include extra
# compilation/caching overhead, while later outputs are closer to
# steady-state inference.
gen = runner.omni.generate(
[PROMPT] * num_requests,
OmniDiffusionSamplingParams(
@@ -134,14 +109,11 @@ def _run_zimage_generate(
per_request_times_s.append(t_now - t_prev)
t_prev = t_now
# Ensure the generator is fully consumed so it can clean up.
for _ in gen:
pass
median_time_s = float(np.median(per_request_times_s))
peak_memory_mb = monitor.peak_used_mb
return _extract_single_image([last_output]), median_time_s, peak_memory_mb
finally:
monitor.stop()
@@ -151,20 +123,20 @@ def _run_zimage_generate(
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
def test_zimage_tensor_parallel_tp2(tmp_path: Path):
def test_tensor_parallel_tp2(tmp_path: Path):
if current_omni_platform.is_npu():
pytest.skip("Z-Image TP e2e test is only supported on CUDA and ROCm for now.")
pytest.skip("Tensor parallel parity test is only supported on CUDA and ROCm for now.")
if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
pytest.skip("Z-Image TP=2 requires >= 2 devices.")
pytest.skip("Tensor parallel TP=2 requires >= 2 devices.")
model_name = DEFAULT_MODEL
enforce_eager = False
height = 512
width = 512
num_inference_steps = 2
seed = 42
tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
tp1_img, tp1_time_s, tp1_peak_mem = _run_generate(
tp_size=1,
height=height,
width=width,
@@ -172,7 +144,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
seed=seed,
enforce_eager=enforce_eager,
)
tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
tp2_img, tp2_time_s, tp2_peak_mem = _run_generate(
tp_size=2,
height=height,
width=width,
@@ -181,8 +153,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
enforce_eager=enforce_eager,
)
tp1_path = tmp_path / "zimage_tp1.png"
tp2_path = tmp_path / "zimage_tp2.png"
tp1_path = tmp_path / "tp1.png"
tp2_path = tmp_path / "tp2.png"
tp1_img.save(tp1_path)
tp2_img.save(tp2_path)
@@ -193,7 +165,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
mean_threshold = 3e-2
p99_threshold = 2.5e-1
print(
"Z-Image TP image diff stats (TP=1 vs TP=2): "
f"{model_name} TP image diff (TP=1 vs TP=2): "
f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
f"tp1_img={tp1_path}, tp2_img={tp2_path}"
@@ -203,73 +175,11 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
)
print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
# ROCm is not optimized TP2 can be slower than TP1
print(f"TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
if not current_omni_platform.is_rocm():
assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
print(f"TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
assert tp2_peak_mem < tp1_peak_mem, (
f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
)
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
if current_omni_platform.is_npu():
pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA and ROCm for now.")
if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.")
enforce_eager = False
# Use a larger image to ensure there are multiple VAE tiles.
height = 1152
width = 1152
num_inference_steps = 2
seed = 42
baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate(
tp_size=2,
height=height,
width=width,
num_inference_steps=num_inference_steps,
seed=seed,
enforce_eager=enforce_eager,
vae_use_tiling=True,
vae_patch_parallel_size=1,
num_requests=2,
)
pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate(
tp_size=2,
height=height,
width=width,
num_inference_steps=num_inference_steps,
seed=seed,
enforce_eager=enforce_eager,
vae_use_tiling=True,
vae_patch_parallel_size=2,
num_requests=2,
)
baseline_path = tmp_path / "zimage_tp2_vae_pp1.png"
pp2_path = tmp_path / "zimage_tp2_vae_pp2.png"
baseline_img.save(baseline_path)
pp2_img.save(pp2_path)
mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
mean_threshold = 5e-3
p99_threshold = 1e-1
print(
"Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): "
f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
f"pp1_img={baseline_path}, pp2_img={pp2_path}"
)
assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
)
tests/diffusion/distributed/test_vae_decode_parallelism.py
+153 -123
View File
@@ -1,151 +1,181 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""VAE patch-parallel decode parity tests for diffusion models.
Compares ``vae_patch_parallel_size=2`` against ``vae_patch_parallel_size=1`` on
TP=2 with VAE tiling enabled. Covers one video model (Wan T2V) and one image
model (Qwen-Image). CUDA/ROCm-only (>=2 devices).
"""
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import Any
import numpy as np
import pytest
import torch
from PIL import Image
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
import time
from tests.helpers.runtime import OmniRunner
from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
all_model_configs = [
{
"model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
"out_width": 1280,
"out_height": 720,
"out_frames": 16,
},
{
"model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
"out_width": 1280,
"out_height": 720,
"out_frames": 16,
},
{
"model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
"out_width": 1280,
"out_height": 720,
"out_frames": 16,
},
{
"model_name": "Qwen/Qwen-Image",
"out_width": 3072,
"out_height": 3072,
"out_frames": 1,
},
# {
# "model_name": "stepfun-ai/NextStep-1.1",
# "out_width": 3072,
# "out_height": 3072,
# "out_frames": 1,
# },
PROMPT = "A cat sitting on a table"
MODEL_CASES = [
pytest.param(
{
"model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
"height": 512,
"width": 512,
"num_frames": 5,
"needs_image": False,
"is_moe": False,
"mean_threshold": 3e-2,
"p99_threshold": 3e-2,
},
id="wan22_t2v_a14b",
),
pytest.param(
{
"model_name": "Qwen/Qwen-Image",
"height": 1152,
"width": 1152,
"num_frames": 1,
"needs_image": False,
"is_moe": False,
"mean_threshold": 5e-3,
"p99_threshold": 1e-1,
},
id="qwen_image",
),
]
def is_nextstep_model(model_name: str) -> bool:
"""Check if the model is a NextStep model by reading its config."""
from vllm.transformers_utils.config import get_hf_file_to_dict
try:
cfg = get_hf_file_to_dict("config.json", model_name)
if cfg and cfg.get("model_type") == "nextstep":
return True
except Exception:
pass
return False
def _to_float_array(data: Any) -> np.ndarray:
if isinstance(data, Image.Image):
return np.asarray(data.convert("RGB"), dtype=np.float32) / 255.0
if isinstance(data, torch.Tensor):
arr = data.detach().cpu().numpy()
else:
arr = np.asarray(data)
if arr.dtype == np.uint8:
return arr.astype(np.float32) / 255.0
return arr.astype(np.float32)
def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
def _extract_output_array(output: Any) -> np.ndarray:
req_out = output.request_output
payload = req_out.images[0]
return _to_float_array(payload)
def _diff_metrics(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
assert a.shape == b.shape, f"Output shapes differ: {a.shape} vs {b.shape}"
abs_diff = np.abs(a - b)
return float(abs_diff.mean()), float(np.quantile(abs_diff, 0.99))
def _save_preview(arr: np.ndarray, path: Path) -> None:
"""Save a single RGB frame preview; video tensors use the first frame."""
preview = np.asarray(arr)
# (B, T, H, W, C) video
if preview.ndim == 5:
preview = preview[0, 0]
# (T, H, W, C) video or (B, H, W, C) image batch
elif preview.ndim == 4:
preview = preview[0]
if preview.ndim != 3 or preview.shape[-1] not in (3, 4):
raise ValueError(f"Expected HWC RGB(A) preview, got shape {arr.shape} -> {preview.shape}")
preview = np.clip(preview[..., :3], 0.0, 1.0)
Image.fromarray((preview * 255.0).astype(np.uint8)).save(path)
def _build_sampling_params(model_case: dict[str, Any], seed: int) -> OmniDiffusionSamplingParams:
kwargs: dict[str, Any] = {
"height": model_case["height"],
"width": model_case["width"],
"num_frames": model_case["num_frames"],
"num_inference_steps": 2,
"seed": seed,
}
if model_case["model_name"].startswith("Wan-AI/"):
kwargs["guidance_scale"] = 4.0
if model_case["is_moe"]:
kwargs["guidance_scale_2"] = 1.0
kwargs["boundary_ratio"] = 0.5
return OmniDiffusionSamplingParams(**kwargs)
def _run_generate(
*,
model_case: dict[str, Any],
vae_patch_parallel_size: int,
seed: int,
) -> np.ndarray:
current_omni_platform.empty_cache()
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=tp,
tensor_parallel_size=2,
vae_patch_parallel_size=vae_patch_parallel_size,
)
omni_kwargs = {
"vae_use_tiling": using_tile,
"parallel_config": parallel_config,
}
use_nextstep = is_nextstep_model(model_configs["model_name"])
if use_nextstep:
# NextStep-1.1 requires explicit pipeline class
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
with OmniRunner(model_configs["model_name"], **omni_kwargs) as runner:
m = runner.omni
image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
start = time.perf_counter()
outputs = m.generate(
{
"prompt": "A cat sitting on a table",
"multi_modal_data": {"image": image},
},
sampling_params_list=OmniDiffusionSamplingParams(
height=model_configs["out_height"],
width=model_configs["out_width"],
num_frames=model_configs["out_frames"],
num_inference_steps=2,
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
),
with OmniRunner(
model_case["model_name"],
parallel_config=parallel_config,
vae_use_tiling=True,
enforce_eager=False,
) as runner:
request: dict[str, Any] = {"prompt": model_case.get("prompt", PROMPT)}
if model_case["needs_image"]:
request["multi_modal_data"] = {
"image": Image.new("RGB", (model_case["width"], model_case["height"]), (0, 0, 0))
}
outputs = runner.omni.generate(
request,
sampling_params_list=_build_sampling_params(model_case, seed),
)
end = time.perf_counter()
first_output = outputs[0]
req_out = first_output.request_output
frames = req_out.images[0]
if isinstance(frames, torch.Tensor):
frames = frames.detach().cpu().numpy()
elif isinstance(frames, Image.Image):
frames = np.array(frames)
# frames shape: (batch, num_frames, height, width, channels)
cost = (end - start) * 1000
return frames, cost
cleanup_dist_env_and_memory()
return _extract_output_array(outputs[0])
@pytest.mark.parametrize("model_configs", all_model_configs)
def test_vae_parallel_model(model_configs: str):
parallel_result, parallel_time = model_run(
model_configs=model_configs,
tp=2,
out_width=1280,
out_height=704,
out_frames=5,
using_tile=True,
vae_patch_parallel_size=2,
)
non_parallel_result, non_parallel_time = model_run(
model_configs=model_configs,
tp=2,
out_width=1280,
out_height=704,
out_frames=5,
using_tile=True,
vae_patch_parallel_size=1,
)
result_diff = np.abs(non_parallel_result - parallel_result)
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("model_case", MODEL_CASES)
def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
if current_omni_platform.is_npu():
pytest.skip("VAE patch parallel test is only supported on CUDA and ROCm for now.")
if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
pytest.skip("VAE patch parallel TP=2 requires >= 2 devices.")
mean_threshold = 3e-2
max_threshold = 3e-2 # they should be totally same
seed = 42
model_id = model_case["model_name"].split("/")[-1]
baseline = _run_generate(model_case=model_case, vae_patch_parallel_size=1, seed=seed)
parallel = _run_generate(model_case=model_case, vae_patch_parallel_size=2, seed=seed)
baseline_path = tmp_path / f"{model_id}_tp2_vae_pp1.png"
parallel_path = tmp_path / f"{model_id}_tp2_vae_pp2.png"
_save_preview(baseline, baseline_path)
_save_preview(parallel, parallel_path)
mean_abs_diff, p99_abs_diff = _diff_metrics(baseline, parallel)
mean_threshold = model_case["mean_threshold"]
p99_threshold = model_case["p99_threshold"]
print(
f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
f"{model_case['model_name']} VAE patch parallel diff (TP=2, pp=1 vs pp=2): "
f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
f"pp1_preview={baseline_path}, pp2_preview={parallel_path}"
)
assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
f"Output diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, "
f"p99_abs_diff={p99_abs_diff:.6e} "
f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
)
# assert non_parallel_time > parallel_time
assert result_diff.mean() < mean_threshold
assert result_diff.max() < max_threshold