84 lines
2.2 KiB
Python
84 lines
2.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.z_image.z_image_transformer import (
|
|
UnifiedPrepare,
|
|
validate_zimage_tp_constraints,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_unified_prepare_applies_padding_masks_before_sharding():
|
|
x = torch.zeros((2, 4, 2))
|
|
cap = torch.zeros((2, 3, 2))
|
|
x_rope = torch.zeros((2, 4, 1))
|
|
cap_rope = torch.zeros((2, 3, 1))
|
|
x_mask = torch.tensor([[1, 0, 1, 0], [1, 1, 0, 0]], dtype=torch.bool)
|
|
cap_mask = torch.tensor([[1, 1, 0], [1, 0, 1]], dtype=torch.bool)
|
|
|
|
*_, unified_mask = UnifiedPrepare()(
|
|
x,
|
|
x_rope,
|
|
x_rope,
|
|
cap,
|
|
cap_rope,
|
|
cap_rope,
|
|
[4, 2],
|
|
[3, 3],
|
|
x_mask,
|
|
cap_mask,
|
|
)
|
|
|
|
assert torch.equal(
|
|
unified_mask,
|
|
torch.tensor(
|
|
[[1, 0, 1, 0, 1, 1, 0], [1, 1, 1, 0, 1, 0, 0]],
|
|
dtype=torch.bool,
|
|
),
|
|
)
|
|
|
|
|
|
def test_validate_zimage_tp_constraints_tp2_ok():
|
|
ffn_hidden_dim, final_out_dims, supported_tp = validate_zimage_tp_constraints(
|
|
dim=3840,
|
|
n_heads=30,
|
|
n_kv_heads=30,
|
|
in_channels=16,
|
|
all_patch_size=(2,),
|
|
all_f_patch_size=(1,),
|
|
tensor_parallel_size=2,
|
|
)
|
|
assert ffn_hidden_dim == 10240
|
|
assert final_out_dims == [64]
|
|
assert supported_tp == [1, 2]
|
|
|
|
|
|
def test_validate_zimage_tp_constraints_tp4_fails_on_heads():
|
|
with pytest.raises(ValueError, match=r"n_heads % tensor_parallel_size"):
|
|
validate_zimage_tp_constraints(
|
|
dim=3840,
|
|
n_heads=30,
|
|
n_kv_heads=30,
|
|
in_channels=16,
|
|
all_patch_size=(2,),
|
|
all_f_patch_size=(1,),
|
|
tensor_parallel_size=4,
|
|
)
|
|
|
|
|
|
def test_validate_zimage_tp_constraints_tp3_fails_on_ffn_hidden_dim():
|
|
with pytest.raises(ValueError, match=r"ffn_hidden_dim % tensor_parallel_size"):
|
|
validate_zimage_tp_constraints(
|
|
dim=3840,
|
|
n_heads=30,
|
|
n_kv_heads=30,
|
|
in_channels=16,
|
|
all_patch_size=(2,),
|
|
all_f_patch_size=(1,),
|
|
tensor_parallel_size=3,
|
|
)
|