133 lines
4.0 KiB
Python
133 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for GLM-Image Sequence Parallelism support."""
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
def setup_sp_groups(mocker):
|
|
"""Set up SP and TP groups for each test function."""
|
|
mock_get_sp_group = mocker.patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group")
|
|
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
|
|
mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
|
|
|
|
mock_sp_group = mocker.MagicMock()
|
|
mock_sp_group.world_size = 4
|
|
mock_get_sp_group.return_value = mock_sp_group
|
|
|
|
mock_tp_group = mocker.MagicMock()
|
|
mock_tp_group.world_size = 1
|
|
mock_get_tp_group.return_value = mock_tp_group
|
|
yield
|
|
|
|
|
|
def test_glm_image_sp_plan_defined():
|
|
"""Test that _sp_plan is properly defined on GlmImageTransformer2DModel."""
|
|
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
|
GlmImageTransformer2DModel,
|
|
)
|
|
|
|
assert hasattr(GlmImageTransformer2DModel, "_sp_plan")
|
|
plan = GlmImageTransformer2DModel._sp_plan
|
|
assert plan is not None
|
|
|
|
# Verify plan structure
|
|
assert "prepare" in plan
|
|
assert "proj_out" in plan
|
|
|
|
|
|
def test_glm_image_sp_plan_valid():
|
|
"""Validate _sp_plan structure."""
|
|
from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan
|
|
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
|
GlmImageTransformer2DModel,
|
|
)
|
|
|
|
plan = GlmImageTransformer2DModel._sp_plan
|
|
validate_sp_plan(plan)
|
|
|
|
|
|
def test_glm_image_prepare_module_exists():
|
|
"""Test that GlmImagePrepare module exists."""
|
|
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
|
GlmImagePrepare,
|
|
)
|
|
|
|
assert GlmImagePrepare is not None
|
|
|
|
|
|
def test_glm_image_attention_accepts_parallel_config():
|
|
"""Test that GlmImageAttention accepts parallel_config parameter."""
|
|
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
|
GlmImageAttention,
|
|
)
|
|
|
|
parallel_config = DiffusionParallelConfig(
|
|
ulysses_degree=2,
|
|
ring_degree=2,
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=4,
|
|
)
|
|
|
|
attn = GlmImageAttention(
|
|
dim=2560,
|
|
num_heads=64,
|
|
head_dim=40,
|
|
parallel_config=parallel_config,
|
|
)
|
|
|
|
assert attn.parallel_config is not None
|
|
assert attn.parallel_config.sequence_parallel_size == 4
|
|
|
|
|
|
def test_glm_image_transformer_block_accepts_parallel_config():
|
|
"""Test that GlmImageTransformerBlock accepts parallel_config parameter."""
|
|
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
|
GlmImageTransformerBlock,
|
|
)
|
|
|
|
parallel_config = DiffusionParallelConfig(
|
|
ulysses_degree=2,
|
|
ring_degree=2,
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=4,
|
|
)
|
|
|
|
block = GlmImageTransformerBlock(
|
|
dim=2560,
|
|
num_attention_heads=64,
|
|
attention_head_dim=40,
|
|
time_embed_dim=512,
|
|
parallel_config=parallel_config,
|
|
)
|
|
|
|
assert block.attn1.parallel_config is not None
|
|
assert block.attn1.parallel_config.sequence_parallel_size == 4
|
|
|
|
|
|
def test_glm_image_has_sp_support():
|
|
"""Test that GLM-Image has SP support implemented."""
|
|
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
|
GlmImageTransformer2DModel,
|
|
)
|
|
|
|
# Check that the model has parallel_config support
|
|
assert hasattr(GlmImageTransformer2DModel, "__init__")
|
|
|
|
# Verify the model can be instantiated with SP config
|
|
|
|
# This test just verifies the structure exists
|
|
# Actual SP testing requires multi-GPU setup
|
|
|
|
|
|
@pytest.mark.cuda
|
|
@pytest.mark.sp
|
|
def test_glm_image_sp_inference():
|
|
"""Test SP inference (requires multi-GPU setup)."""
|
|
pytest.skip("Requires multi-GPU SP setup")
|