95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def test_grouped_qkv_checkpoint_reorder():
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
_reorder_grouped_qkv_to_qkv,
|
|
)
|
|
|
|
# Two groups with rows [q, k, v] become [q0, q1, k0, k1, v0, v1].
|
|
grouped = torch.arange(6, dtype=torch.float32).reshape(6, 1)
|
|
reordered = _reorder_grouped_qkv_to_qkv(
|
|
grouped,
|
|
num_query_groups=2,
|
|
heads_per_group=1,
|
|
head_dim=1,
|
|
)
|
|
|
|
assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
|
|
|
|
|
|
def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
|
|
from cache_dit import ForwardPattern
|
|
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
MiniMaxH3DiTModel,
|
|
)
|
|
|
|
assert MiniMaxH3DiTModel._repeated_blocks == ["MiniMaxH3DiTBlock"]
|
|
assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
|
|
assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
|
|
assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
|
|
assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
|
|
|
|
model = object.__new__(MiniMaxH3DiTModel)
|
|
nn.Module.__init__(model)
|
|
model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
|
|
model.token_refiner = nn.Module()
|
|
model.token_refiner.blocks = nn.ModuleList([nn.Linear(4, 4)])
|
|
model.final_layer = nn.Linear(4, 4)
|
|
|
|
matched = [
|
|
name
|
|
for name, module in model.named_modules()
|
|
if any(condition(name, module) for condition in MiniMaxH3DiTModel._hsdp_shard_conditions)
|
|
]
|
|
assert matched == ["blocks.0", "blocks.1"]
|
|
|
|
|
|
def test_packed_attention_is_a_regional_compile_boundary():
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
MiniMaxH3Attention,
|
|
)
|
|
|
|
assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("tp_size", "message"),
|
|
[
|
|
(3, "num_attention_heads"),
|
|
(5, "num_attention_heads"),
|
|
],
|
|
)
|
|
def test_tp_rejects_non_divisible_head_counts(tp_size, message):
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
MiniMaxH3DiTArchConfig,
|
|
MiniMaxH3DiTModel,
|
|
)
|
|
|
|
model = object.__new__(MiniMaxH3DiTModel)
|
|
with pytest.raises(ValueError, match=message):
|
|
model._validate_tp_config(
|
|
arch=MiniMaxH3DiTArchConfig(),
|
|
tp_size=tp_size,
|
|
)
|
|
|
|
|
|
def test_tp_accepts_checkpoint_supported_sizes():
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
MiniMaxH3DiTArchConfig,
|
|
MiniMaxH3DiTModel,
|
|
)
|
|
|
|
model = object.__new__(MiniMaxH3DiTModel)
|
|
arch = MiniMaxH3DiTArchConfig()
|
|
for tp_size in (1, 2, 4, 7):
|
|
model._validate_tp_config(arch=arch, tp_size=tp_size)
|