github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 02:51:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

[workflow] fixed oom tests (#5275)

Browse Source 
* [workflow] fixed oom tests

* polish

* polish

* polish
This commit is contained in:
Frank Lee
2024-01-16 18:55:13 +08:00
committed by GitHub
parent 04244aaaf1
commit d69cd2eb89
19 changed files with 50 additions and 582 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/test_lazy/test_models.py
View File

@@ -1,14 +1,19 @@
import pytest
from lazy_init_utils import SUPPORT_LAZY, check_lazy_init
from tests.kit.model_zoo import model_zoo, IS_FAST_TEST, COMMON_MODELS
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"])
@pytest.mark.parametrize(
"subset",
[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):
sub_model_zoo = model_zoo.get_sub_registry(subset)
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(