github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[lazyinit] skip whisper test (#5653)

pull/5655/head
Hongxin Liu 2024-04-26 14:03:12 +08:00 committed by GitHub
parent 8b7d535977
commit 2082852f3f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tests/test_lazy/test_models.py
View File

@ -7,21 +7,23 @@ from tests.kit.model_zoo import COMMON_MODELS, IS_FAST_TEST, model_zoo
@pytest.mark.skipif(not SUPPORT_LAZY, reason="requires torch >= 1.12.0")
@pytest.mark.parametrize(
"subset",
[COMMON_MODELS]
if IS_FAST_TEST
else ["torchvision", "diffusers", "timm", "transformers", "torchaudio", "deepfm", "dlrm"],
(
[COMMON_MODELS]
if IS_FAST_TEST
else ["torchvision", "diffusers", "timm", "transformers", "torchaudio", "deepfm", "dlrm"]
),
)
@pytest.mark.parametrize("default_device", ["cpu", "cuda"])
def test_torchvision_models_lazy_init(subset, default_device):
def test_models_lazy_init(subset, default_device):
sub_model_zoo = model_zoo.get_sub_registry(subset, allow_empty=True)
for name, entry in sub_model_zoo.items():
# TODO(ver217): lazy init does not support weight norm, skip these models
if name in ("torchaudio_wav2vec2_base", "torchaudio_hubert_base") or name.startswith(
("transformers_vit", "transformers_blip2")
("transformers_vit", "transformers_blip2", "transformers_whisper")
):
continue
check_lazy_init(entry, verbose=True, default_device=default_device)
if __name__ == "__main__":
test_torchvision_models_lazy_init("transformers", "cpu")
test_models_lazy_init("transformers", "cpu")