741 lines
26 KiB
Python
741 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Unit tests for GLM-Image quantization support (W4A16/AutoRound).
|
|
|
|
These tests verify that the GLM-Image DiT transformer correctly accepts and uses
|
|
quantization configs for W4A16/AutoRound quantization support.
|
|
"""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
import torch
|
|
from pytest_mock import MockerFixture
|
|
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
|
ColumnParallelGELU,
|
|
ColumnParallelSiLU,
|
|
GlmImageAdaLayerNormContinuous,
|
|
GlmImageAdaLayerNormZero,
|
|
GlmImageAttention,
|
|
GlmImageFeedForward,
|
|
GlmImageImageProjector,
|
|
GlmImagePrepare,
|
|
GlmImageRotaryPosEmbed,
|
|
GlmImageTransformer2DModel,
|
|
GlmImageTransformerBlock,
|
|
_positive_divisors,
|
|
validate_glm_image_tp_constraints,
|
|
)
|
|
from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
def _init_distributed():
|
|
"""Initialize the minimal distributed environment required by
|
|
ReplicatedLinear (tensor-parallel group must exist)."""
|
|
from vllm.distributed.parallel_state import (
|
|
cleanup_dist_env_and_memory,
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
|
|
os.environ.setdefault("MASTER_ADDR", "localhost")
|
|
os.environ.setdefault("MASTER_PORT", "29502")
|
|
init_distributed_environment(
|
|
world_size=1,
|
|
rank=0,
|
|
local_rank=0,
|
|
distributed_init_method="env://",
|
|
)
|
|
initialize_model_parallel()
|
|
yield
|
|
cleanup_dist_env_and_memory()
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
def _force_default_gemm(monkeypatch):
|
|
"""Force CPU-compatible GEMM dispatch for tests using CPU tensors.
|
|
|
|
vLLM's dispatch_unquantized_gemm() selects the backend by platform
|
|
(e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
|
|
CPU test tensors crash with NotImplementedError on ROCm. Monkeypatch
|
|
the dispatcher to always return the default (torch.nn.functional.linear)
|
|
implementation which works on any device."""
|
|
from vllm.model_executor.layers.utils import default_unquantized_gemm
|
|
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
|
|
lambda *_args, **_kwargs: default_unquantized_gemm,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
def setup_mocks(mocker: MockerFixture):
|
|
"""Set up common mocks for all tests."""
|
|
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_tp_group = mocker.MagicMock()
|
|
mock_tp_group.world_size = 1
|
|
mock_get_tp_group.return_value = mock_tp_group
|
|
yield
|
|
|
|
|
|
class TestPositiveDivisors:
|
|
"""Test _positive_divisors helper function."""
|
|
|
|
def test_divisors_of_1(self):
|
|
assert _positive_divisors(1) == {1}
|
|
|
|
def test_divisors_of_12(self):
|
|
assert _positive_divisors(12) == {1, 2, 3, 4, 6, 12}
|
|
|
|
def test_divisors_of_prime(self):
|
|
assert _positive_divisors(7) == {1, 7}
|
|
|
|
def test_divisors_of_0_returns_empty(self):
|
|
assert _positive_divisors(0) == set()
|
|
|
|
def test_divisors_of_negative_returns_empty(self):
|
|
assert _positive_divisors(-5) == set()
|
|
|
|
|
|
class TestValidateGlmImageTpConstraints:
|
|
"""Test TP constraint validation for GLM-Image."""
|
|
|
|
def test_valid_tp_size_1(self):
|
|
"""TP=1 should always be valid."""
|
|
result = validate_glm_image_tp_constraints(
|
|
dim=2560,
|
|
num_heads=64,
|
|
ffn_hidden_dim=10240,
|
|
tensor_parallel_size=1,
|
|
)
|
|
assert 1 in result
|
|
|
|
def test_valid_tp_size_2_for_divisible_dim(self):
|
|
"""TP=2 is valid when all dims are divisible by 2."""
|
|
result = validate_glm_image_tp_constraints(
|
|
dim=2560,
|
|
num_heads=64,
|
|
ffn_hidden_dim=10240,
|
|
tensor_parallel_size=2,
|
|
)
|
|
assert 1 in result
|
|
assert 2 in result
|
|
|
|
def test_valid_tp_size_4_for_divisible_dim(self):
|
|
"""TP=4 is valid when all dims are divisible by 4."""
|
|
result = validate_glm_image_tp_constraints(
|
|
dim=2560,
|
|
num_heads=64,
|
|
ffn_hidden_dim=10240,
|
|
tensor_parallel_size=4,
|
|
)
|
|
assert 1 in result
|
|
assert 2 in result
|
|
assert 4 in result
|
|
|
|
def test_invalid_tp_size_3_for_divisible_dim(self):
|
|
"""TP=3 is invalid when dim is not divisible by 3."""
|
|
with pytest.raises(ValueError, match="dim % tensor_parallel_size == 0"):
|
|
validate_glm_image_tp_constraints(
|
|
dim=2560, # 2560 % 3 != 0
|
|
num_heads=64,
|
|
ffn_hidden_dim=10240,
|
|
tensor_parallel_size=3,
|
|
)
|
|
|
|
def test_invalid_tp_size_zero(self):
|
|
"""TP=0 should raise error."""
|
|
with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
|
|
validate_glm_image_tp_constraints(
|
|
dim=2560,
|
|
num_heads=64,
|
|
ffn_hidden_dim=10240,
|
|
tensor_parallel_size=0,
|
|
)
|
|
|
|
def test_invalid_tp_size_negative(self):
|
|
"""Negative TP size should raise error."""
|
|
with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
|
|
validate_glm_image_tp_constraints(
|
|
dim=2560,
|
|
num_heads=64,
|
|
ffn_hidden_dim=10240,
|
|
tensor_parallel_size=-1,
|
|
)
|
|
|
|
|
|
class TestGlmImageAdaLayerNormZeroQuantization:
|
|
"""Test GlmImageAdaLayerNormZero with quantization config."""
|
|
|
|
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
|
"""Verify the class accepts quant_config parameter."""
|
|
mock_quant_config = mocker.MagicMock()
|
|
layer = GlmImageAdaLayerNormZero(
|
|
embedding_dim=512,
|
|
dim=2560,
|
|
quant_config=mock_quant_config,
|
|
prefix="test.norm1",
|
|
)
|
|
assert layer.linear.quant_config is mock_quant_config
|
|
|
|
def test_accepts_none_quant_config(self):
|
|
"""Verify quant_config=None is accepted."""
|
|
layer = GlmImageAdaLayerNormZero(
|
|
embedding_dim=512,
|
|
dim=2560,
|
|
quant_config=None,
|
|
)
|
|
assert layer.linear.quant_config is None
|
|
|
|
def test_forward_handles_tuple_return_from_linear(self):
|
|
"""Verify forward handles tuple returns from ReplicatedLinear."""
|
|
layer = GlmImageAdaLayerNormZero(
|
|
embedding_dim=512,
|
|
dim=2560,
|
|
quant_config=None,
|
|
)
|
|
|
|
batch_size = 2
|
|
seq_len = 10
|
|
hidden_states = torch.randn(batch_size, seq_len, 2560)
|
|
encoder_hidden_states = torch.randn(batch_size, seq_len, 2560)
|
|
temb = torch.randn(batch_size, 512)
|
|
|
|
# This should work regardless of whether linear returns tuple or tensor
|
|
result = layer(hidden_states, encoder_hidden_states, temb)
|
|
assert len(result) == 10 # Should return 10 chunks
|
|
|
|
|
|
class TestGlmImageAdaLayerNormContinuousQuantization:
|
|
"""Test GlmImageAdaLayerNormContinuous with quantization config."""
|
|
|
|
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
|
"""Verify the class accepts quant_config parameter."""
|
|
mock_quant_config = mocker.MagicMock()
|
|
layer = GlmImageAdaLayerNormContinuous(
|
|
embedding_dim=2560,
|
|
conditioning_embedding_dim=512,
|
|
quant_config=mock_quant_config,
|
|
prefix="test.norm_out",
|
|
)
|
|
assert layer.linear.quant_config is mock_quant_config
|
|
|
|
def test_accepts_none_quant_config(self):
|
|
"""Verify quant_config=None is accepted."""
|
|
layer = GlmImageAdaLayerNormContinuous(
|
|
embedding_dim=2560,
|
|
conditioning_embedding_dim=512,
|
|
quant_config=None,
|
|
)
|
|
assert layer.linear.quant_config is None
|
|
|
|
def test_forward_handles_tuple_return_from_linear(self):
|
|
"""Verify forward handles tuple returns from ReplicatedLinear."""
|
|
layer = GlmImageAdaLayerNormContinuous(
|
|
embedding_dim=2560,
|
|
conditioning_embedding_dim=512,
|
|
quant_config=None,
|
|
)
|
|
|
|
batch_size = 2
|
|
seq_len = 10
|
|
x = torch.randn(batch_size, seq_len, 2560)
|
|
conditioning_embedding = torch.randn(batch_size, 512)
|
|
|
|
result = layer(x, conditioning_embedding)
|
|
assert result.shape == (batch_size, seq_len, 2560)
|
|
|
|
|
|
class TestGlmImageAttentionQuantization:
|
|
"""Test GlmImageAttention with quantization config."""
|
|
|
|
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
|
"""Verify GlmImageAttention accepts quant_config parameter."""
|
|
mock_quant_config = mocker.MagicMock()
|
|
attn = GlmImageAttention(
|
|
dim=2560,
|
|
num_heads=64,
|
|
head_dim=40,
|
|
quant_config=mock_quant_config,
|
|
prefix="test.attn1",
|
|
)
|
|
assert attn.to_qkv.quant_config is mock_quant_config
|
|
|
|
def test_accepts_none_quant_config(self):
|
|
"""Verify quant_config=None is accepted."""
|
|
attn = GlmImageAttention(
|
|
dim=2560,
|
|
num_heads=64,
|
|
head_dim=40,
|
|
quant_config=None,
|
|
)
|
|
assert attn.to_qkv.quant_config is None
|
|
|
|
|
|
class TestColumnParallelModulesQuantization:
|
|
"""Test ColumnParallelGELU and ColumnParallelSiLU with quantization config."""
|
|
|
|
def test_column_parallel_gelu_accepts_quant_config(self, mocker: MockerFixture):
|
|
"""Verify ColumnParallelGELU accepts quant_config."""
|
|
mock_quant_config = mocker.MagicMock()
|
|
layer = ColumnParallelGELU(
|
|
dim_in=2560,
|
|
dim_out=10240,
|
|
quant_config=mock_quant_config,
|
|
prefix="test.gelu",
|
|
)
|
|
assert layer.proj.quant_config is mock_quant_config
|
|
|
|
def test_column_parallel_silu_accepts_quant_config(self, mocker: MockerFixture):
|
|
"""Verify ColumnParallelSiLU accepts quant_config."""
|
|
mock_quant_config = mocker.MagicMock()
|
|
layer = ColumnParallelSiLU(
|
|
dim_in=2560,
|
|
dim_out=10240,
|
|
quant_config=mock_quant_config,
|
|
prefix="test.silu",
|
|
)
|
|
assert layer.proj.quant_config is mock_quant_config
|
|
|
|
|
|
class TestGlmImageFeedForwardQuantization:
|
|
"""Test GlmImageFeedForward with quantization config."""
|
|
|
|
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
|
"""Verify GlmImageFeedForward accepts quant_config parameter."""
|
|
mock_quant_config = mocker.MagicMock()
|
|
ff = GlmImageFeedForward(
|
|
dim=2560,
|
|
dim_out=2560,
|
|
inner_dim=10240,
|
|
activation_fn="gelu-approximate",
|
|
quant_config=mock_quant_config,
|
|
prefix="test.ff",
|
|
)
|
|
# Check that the first layer (ColumnParallelGELU) has quant_config
|
|
gelu_layer = ff.net[0]
|
|
assert gelu_layer.proj.quant_config is mock_quant_config
|
|
|
|
def test_accepts_none_quant_config(self):
|
|
"""Verify quant_config=None is accepted."""
|
|
ff = GlmImageFeedForward(
|
|
dim=2560,
|
|
dim_out=2560,
|
|
inner_dim=10240,
|
|
activation_fn="gelu-approximate",
|
|
quant_config=None,
|
|
)
|
|
gelu_layer = ff.net[0]
|
|
assert gelu_layer.proj.quant_config is None
|
|
|
|
def test_linear_silu_activation(self):
|
|
"""Test linear-silu activation function initialization works."""
|
|
# This test verifies the module can be instantiated with linear-silu activation.
|
|
# Full forward testing requires proper TP group setup, so we just verify construction.
|
|
ff = GlmImageFeedForward(
|
|
dim=2560,
|
|
dim_out=2560,
|
|
inner_dim=10240,
|
|
activation_fn="linear-silu",
|
|
quant_config=None,
|
|
)
|
|
# Verify the FFN has the correct structure
|
|
assert len(ff.net) == 3 # ColumnParallelSiLU, Identity, RowParallelLinear
|
|
assert isinstance(ff.net[0], ColumnParallelSiLU)
|
|
|
|
|
|
class TestGlmImageTransformerBlockQuantization:
|
|
"""Test GlmImageTransformerBlock with quantization config."""
|
|
|
|
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
|
"""Verify GlmImageTransformerBlock accepts quant_config parameter."""
|
|
mock_quant_config = mocker.MagicMock()
|
|
parallel_config = DiffusionParallelConfig(
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
)
|
|
block = GlmImageTransformerBlock(
|
|
dim=2560,
|
|
num_attention_heads=64,
|
|
attention_head_dim=40,
|
|
time_embed_dim=512,
|
|
parallel_config=parallel_config,
|
|
quant_config=mock_quant_config,
|
|
prefix="test.block",
|
|
)
|
|
# Check that inner modules have quant_config
|
|
assert block.norm1.linear.quant_config is mock_quant_config
|
|
assert block.attn1.to_qkv.quant_config is mock_quant_config
|
|
assert block.ff.net[0].proj.quant_config is mock_quant_config
|
|
|
|
def test_accepts_none_quant_config(self):
|
|
"""Verify quant_config=None is accepted."""
|
|
block = GlmImageTransformerBlock(
|
|
dim=2560,
|
|
num_attention_heads=64,
|
|
attention_head_dim=40,
|
|
time_embed_dim=512,
|
|
quant_config=None,
|
|
)
|
|
assert block.norm1.linear.quant_config is None
|
|
|
|
|
|
class TestGlmImagePrepareModule:
|
|
"""Test GlmImagePrepare module."""
|
|
|
|
def test_prepare_module_exists(self):
|
|
"""Verify GlmImagePrepare module exists and works."""
|
|
projector = GlmImageImageProjector(in_channels=16, hidden_size=2560, patch_size=2)
|
|
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2)
|
|
prepare = GlmImagePrepare(
|
|
image_projector=projector,
|
|
rope=rope,
|
|
)
|
|
|
|
# Test forward pass
|
|
hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W]
|
|
result = prepare(hidden_states)
|
|
|
|
assert len(result) == 3
|
|
hidden_out, rope_cos, rope_sin = result
|
|
assert hidden_out.shape[0] == 1 # batch size
|
|
assert hidden_out.shape[1] == 1024 # seq_len (32 * 32)
|
|
assert hidden_out.shape[2] == 2560 # hidden_dim
|
|
|
|
|
|
class TestGlmImagePipelineConfig:
|
|
"""Test GLM_IMAGE_PIPELINE configuration."""
|
|
|
|
def test_glm_image_pipeline_config_exists(self):
|
|
"""Verify GLM_IMAGE_PIPELINE config exists."""
|
|
assert GLM_IMAGE_PIPELINE is not None
|
|
|
|
def test_pipeline_has_correct_model_type(self):
|
|
"""Verify pipeline has correct model_type."""
|
|
assert GLM_IMAGE_PIPELINE.model_type == "glm_image"
|
|
|
|
def test_pipeline_has_correct_model_arch(self):
|
|
"""Verify pipeline has correct model_arch."""
|
|
assert GLM_IMAGE_PIPELINE.model_arch == "GlmImageForConditionalGeneration"
|
|
|
|
def test_pipeline_has_two_stages(self):
|
|
"""Verify pipeline has two stages (AR + DiT)."""
|
|
assert len(GLM_IMAGE_PIPELINE.stages) == 2
|
|
|
|
def test_stage_0_is_ar_llm(self):
|
|
"""Verify stage 0 is the AR LLM stage."""
|
|
stage_0 = GLM_IMAGE_PIPELINE.stages[0]
|
|
assert stage_0.stage_id == 0
|
|
assert stage_0.model_stage == "ar"
|
|
assert stage_0.execution_type.value == "llm_ar"
|
|
assert stage_0.owns_tokenizer is True
|
|
|
|
def test_stage_1_is_diffusion(self):
|
|
"""Verify stage 1 is the DiT diffusion stage."""
|
|
stage_1 = GLM_IMAGE_PIPELINE.stages[1]
|
|
assert stage_1.stage_id == 1
|
|
assert stage_1.model_stage == "dit"
|
|
assert stage_1.execution_type.value == "diffusion"
|
|
assert stage_1.final_output is True
|
|
assert stage_1.final_output_type == "image"
|
|
assert stage_1.input_sources == (0,) # Takes input from stage 0
|
|
|
|
def test_pipeline_has_diffusers_class_name(self):
|
|
"""Verify pipeline has diffusers_class_name."""
|
|
assert GLM_IMAGE_PIPELINE.diffusers_class_name == "GlmImagePipeline"
|
|
|
|
|
|
class TestGlmImageImageProjector:
|
|
"""Test GlmImageImageProjector module."""
|
|
|
|
def test_projector_output_shape(self):
|
|
"""Verify projector produces correct output shape."""
|
|
projector = GlmImageImageProjector(
|
|
in_channels=16,
|
|
hidden_size=2560,
|
|
patch_size=2,
|
|
)
|
|
|
|
# Input: [B, C, H, W] = [1, 16, 64, 64]
|
|
hidden_states = torch.randn(1, 16, 64, 64)
|
|
output = projector(hidden_states)
|
|
|
|
# After patchify: [B, H/2*W/2, D] = [1, 32*32, 2560] = [1, 1024, 2560]
|
|
assert output.shape == (1, 1024, 2560)
|
|
|
|
def test_projector_with_different_patch_sizes(self):
|
|
"""Verify projector works with different patch sizes."""
|
|
for patch_size in [2, 4]:
|
|
projector = GlmImageImageProjector(
|
|
in_channels=16,
|
|
hidden_size=2560,
|
|
patch_size=patch_size,
|
|
)
|
|
h, w = 64, 64
|
|
hidden_states = torch.randn(1, 16, h, w)
|
|
output = projector(hidden_states)
|
|
expected_seq_len = (h // patch_size) * (w // patch_size)
|
|
assert output.shape[1] == expected_seq_len
|
|
|
|
|
|
class TestGlmImageRotaryPosEmbed:
|
|
"""Test GlmImageRotaryPosEmbed module."""
|
|
|
|
def test_rope_output_shape(self):
|
|
"""Verify RoPE produces correct output shape."""
|
|
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
|
|
hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W]
|
|
|
|
cos, sin = rope(hidden_states)
|
|
|
|
# After patchify: height=32, width=32
|
|
# Output: [32*32, dim] = [1024, 40]
|
|
assert cos.shape[0] == 1024
|
|
assert cos.shape[1] == 40
|
|
assert sin.shape == cos.shape
|
|
|
|
def test_rope_value_range(self):
|
|
"""Verify RoPE produces values in valid range."""
|
|
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
|
|
hidden_states = torch.randn(2, 16, 32, 32)
|
|
|
|
cos, sin = rope(hidden_states)
|
|
|
|
# cos and sin should be in [-1, 1]
|
|
assert cos.min() >= -1.0 and cos.max() <= 1.0
|
|
assert sin.min() >= -1.0 and sin.max() <= 1.0
|
|
|
|
def test_rope_consistency(self):
|
|
"""Verify RoPE is consistent for same input."""
|
|
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
|
|
hidden_states = torch.randn(1, 16, 32, 32)
|
|
|
|
cos1, sin1 = rope(hidden_states)
|
|
cos2, sin2 = rope(hidden_states)
|
|
|
|
# Same input should produce same output
|
|
assert torch.allclose(cos1, cos2)
|
|
assert torch.allclose(sin1, sin2)
|
|
|
|
|
|
class TestGlmImageTransformer2DModelQuantization:
|
|
"""Test GlmImageTransformer2DModel with quantization config."""
|
|
|
|
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
|
"""Verify the model accepts quant_config parameter."""
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
mock_quant_config = mocker.MagicMock()
|
|
parallel_config = DiffusionParallelConfig(
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
)
|
|
|
|
# Create a minimal mock od_config
|
|
mock_tf_config = mocker.MagicMock()
|
|
mock_tf_config.patch_size = 2
|
|
mock_tf_config.in_channels = 16
|
|
mock_tf_config.out_channels = 16
|
|
mock_tf_config.num_attention_heads = 64
|
|
mock_tf_config.attention_head_dim = 40
|
|
mock_tf_config.time_embed_dim = 512
|
|
mock_tf_config.condition_dim = 256
|
|
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
|
mock_tf_config.text_embed_dim = 1024
|
|
mock_tf_config.num_layers = 2 # Small number for testing
|
|
|
|
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
|
mock_od_config.tf_model_config = mock_tf_config
|
|
mock_od_config.parallel_config = parallel_config
|
|
|
|
model = GlmImageTransformer2DModel(
|
|
od_config=mock_od_config,
|
|
quant_config=mock_quant_config,
|
|
)
|
|
|
|
# Check that quantization config was passed to transformer blocks
|
|
for block in model.transformer_blocks:
|
|
assert block.norm1.linear.quant_config is mock_quant_config
|
|
assert block.attn1.to_qkv.quant_config is mock_quant_config
|
|
assert block.ff.net[0].proj.quant_config is mock_quant_config
|
|
|
|
def test_accepts_none_quant_config(self, mocker: MockerFixture):
|
|
"""Verify quant_config=None is accepted."""
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
parallel_config = DiffusionParallelConfig(
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
)
|
|
|
|
mock_tf_config = mocker.MagicMock()
|
|
mock_tf_config.patch_size = 2
|
|
mock_tf_config.in_channels = 16
|
|
mock_tf_config.out_channels = 16
|
|
mock_tf_config.num_attention_heads = 64
|
|
mock_tf_config.attention_head_dim = 40
|
|
mock_tf_config.time_embed_dim = 512
|
|
mock_tf_config.condition_dim = 256
|
|
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
|
mock_tf_config.text_embed_dim = 1024
|
|
mock_tf_config.num_layers = 2
|
|
|
|
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
|
mock_od_config.tf_model_config = mock_tf_config
|
|
mock_od_config.parallel_config = parallel_config
|
|
|
|
model = GlmImageTransformer2DModel(
|
|
od_config=mock_od_config,
|
|
quant_config=None,
|
|
)
|
|
|
|
# Check that quantization config is None
|
|
for block in model.transformer_blocks:
|
|
assert block.norm1.linear.quant_config is None
|
|
assert block.attn1.to_qkv.quant_config is None
|
|
|
|
def test_norm_out_has_no_quantization(self, mocker: MockerFixture):
|
|
"""Verify norm_out (output layer) does NOT use quantization to preserve precision."""
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
mock_quant_config = mocker.MagicMock()
|
|
parallel_config = DiffusionParallelConfig(
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
)
|
|
|
|
mock_tf_config = mocker.MagicMock()
|
|
mock_tf_config.patch_size = 2
|
|
mock_tf_config.in_channels = 16
|
|
mock_tf_config.out_channels = 16
|
|
mock_tf_config.num_attention_heads = 64
|
|
mock_tf_config.attention_head_dim = 40
|
|
mock_tf_config.time_embed_dim = 512
|
|
mock_tf_config.condition_dim = 256
|
|
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
|
mock_tf_config.text_embed_dim = 1024
|
|
mock_tf_config.num_layers = 2
|
|
|
|
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
|
mock_od_config.tf_model_config = mock_tf_config
|
|
mock_od_config.parallel_config = parallel_config
|
|
|
|
model = GlmImageTransformer2DModel(
|
|
od_config=mock_od_config,
|
|
quant_config=mock_quant_config,
|
|
)
|
|
|
|
# norm_out.linear should NOT have quant_config to preserve output precision
|
|
assert model.norm_out.linear.quant_config is None
|
|
|
|
def test_model_has_sp_plan(self, mocker: MockerFixture):
|
|
"""Verify model has _sp_plan defined for sequence parallelism."""
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
parallel_config = DiffusionParallelConfig(
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
)
|
|
|
|
mock_tf_config = mocker.MagicMock()
|
|
mock_tf_config.patch_size = 2
|
|
mock_tf_config.in_channels = 16
|
|
mock_tf_config.out_channels = 16
|
|
mock_tf_config.num_attention_heads = 64
|
|
mock_tf_config.attention_head_dim = 40
|
|
mock_tf_config.time_embed_dim = 512
|
|
mock_tf_config.condition_dim = 256
|
|
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
|
mock_tf_config.text_embed_dim = 1024
|
|
mock_tf_config.num_layers = 2
|
|
|
|
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
|
mock_od_config.tf_model_config = mock_tf_config
|
|
mock_od_config.parallel_config = parallel_config
|
|
|
|
model = GlmImageTransformer2DModel(
|
|
od_config=mock_od_config,
|
|
quant_config=None,
|
|
)
|
|
|
|
# Verify _sp_plan exists
|
|
assert hasattr(model, "_sp_plan")
|
|
assert "prepare" in model._sp_plan
|
|
assert "proj_out" in model._sp_plan
|
|
|
|
def test_model_has_hsdp_shard_conditions(self, mocker: MockerFixture):
|
|
"""Verify model has _hsdp_shard_conditions for HSDP parallelism."""
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
parallel_config = DiffusionParallelConfig(
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
)
|
|
|
|
mock_tf_config = mocker.MagicMock()
|
|
mock_tf_config.patch_size = 2
|
|
mock_tf_config.in_channels = 16
|
|
mock_tf_config.out_channels = 16
|
|
mock_tf_config.num_attention_heads = 64
|
|
mock_tf_config.attention_head_dim = 40
|
|
mock_tf_config.time_embed_dim = 512
|
|
mock_tf_config.condition_dim = 256
|
|
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
|
mock_tf_config.text_embed_dim = 1024
|
|
mock_tf_config.num_layers = 2
|
|
|
|
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
|
mock_od_config.tf_model_config = mock_tf_config
|
|
mock_od_config.parallel_config = parallel_config
|
|
|
|
model = GlmImageTransformer2DModel(
|
|
od_config=mock_od_config,
|
|
quant_config=None,
|
|
)
|
|
|
|
assert hasattr(model, "_hsdp_shard_conditions")
|
|
assert len(model._hsdp_shard_conditions) > 0
|
|
|
|
def test_model_creates_kv_cache(self, mocker: MockerFixture):
|
|
"""Verify model can create KV cache for image editing."""
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
parallel_config = DiffusionParallelConfig(
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
)
|
|
|
|
mock_tf_config = mocker.MagicMock()
|
|
mock_tf_config.patch_size = 2
|
|
mock_tf_config.in_channels = 16
|
|
mock_tf_config.out_channels = 16
|
|
mock_tf_config.num_attention_heads = 64
|
|
mock_tf_config.attention_head_dim = 40
|
|
mock_tf_config.time_embed_dim = 512
|
|
mock_tf_config.condition_dim = 256
|
|
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
|
mock_tf_config.text_embed_dim = 1024
|
|
mock_tf_config.num_layers = 2
|
|
|
|
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
|
mock_od_config.tf_model_config = mock_tf_config
|
|
mock_od_config.parallel_config = parallel_config
|
|
|
|
model = GlmImageTransformer2DModel(
|
|
od_config=mock_od_config,
|
|
quant_config=None,
|
|
)
|
|
|
|
kv_cache = model.create_kv_cache()
|
|
assert kv_cache is not None
|
|
assert len(kv_cache) == 2 # num_layers
|