mirror of https://github.com/hpcaitech/ColossalAI
[lazyinit] skip whisper test (#5653)
parent
8b7d535977
commit
2082852f3f
|
@ -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.skipif(not SUPPORT_LAZY, reason="requires torch >= 1.12.0")
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"subset",
|
"subset",
|
||||||
[COMMON_MODELS]
|
(
|
||||||
if IS_FAST_TEST
|
[COMMON_MODELS]
|
||||||
else ["torchvision", "diffusers", "timm", "transformers", "torchaudio", "deepfm", "dlrm"],
|
if IS_FAST_TEST
|
||||||
|
else ["torchvision", "diffusers", "timm", "transformers", "torchaudio", "deepfm", "dlrm"]
|
||||||
|
),
|
||||||
)
|
)
|
||||||
@pytest.mark.parametrize("default_device", ["cpu", "cuda"])
|
@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)
|
sub_model_zoo = model_zoo.get_sub_registry(subset, allow_empty=True)
|
||||||
for name, entry in sub_model_zoo.items():
|
for name, entry in sub_model_zoo.items():
|
||||||
# TODO(ver217): lazy init does not support weight norm, skip these models
|
# TODO(ver217): lazy init does not support weight norm, skip these models
|
||||||
if name in ("torchaudio_wav2vec2_base", "torchaudio_hubert_base") or name.startswith(
|
if name in ("torchaudio_wav2vec2_base", "torchaudio_hubert_base") or name.startswith(
|
||||||
("transformers_vit", "transformers_blip2")
|
("transformers_vit", "transformers_blip2", "transformers_whisper")
|
||||||
):
|
):
|
||||||
continue
|
continue
|
||||||
check_lazy_init(entry, verbose=True, default_device=default_device)
|
check_lazy_init(entry, verbose=True, default_device=default_device)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
test_torchvision_models_lazy_init("transformers", "cpu")
|
test_models_lazy_init("transformers", "cpu")
|
||||||
|
|
Loading…
Reference in New Issue