fywc
9f81a0a4b0
Signed-off-by: hanzheli <hanzheli@kuaishou.com> Signed-off-by: fywc <hanzheli@kuaishou.com>
28 lines
815 B
Python
28 lines
815 B
Python
import pytest
|
|
import torch.nn as nn
|
|
|
|
pytest.importorskip("dreamid_omni.modules.model")
|
|
|
|
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
|
|
model = object.__new__(FusionModel)
|
|
nn.Module.__init__(model)
|
|
model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
|
|
model.attn = nn.Identity()
|
|
|
|
conditions = getattr(model, "_hsdp_shard_conditions", None)
|
|
|
|
assert conditions is not None
|
|
assert len(conditions) == 1
|
|
|
|
matched = []
|
|
for name, module in model.named_modules():
|
|
if any(cond(name, module) for cond in conditions):
|
|
matched.append(name)
|
|
|
|
assert matched == ["fused_blocks.0", "fused_blocks.1"]
|