/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
tests/test_lazy/test_models.py
39 строк
1 KB
Hanks
[fix] fix_lazy_init for deepseek model in transformers
02 июн 2025, 06:31
Не верифицирован
02 июн 2025, 06:31
6f19618
Код
Авторство
О чём код?
import pytest from lazy_init_utils import SUPPORT_LAZY, check_lazy_init 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("default_device", ["cpu", "cuda"]) 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", "timm_beit", "timm_vision_transformer", "timm_deit", "timm_beitv2", "timm_deit3", "timm_convit", "timm_tnt_b_patch16_224", ) or name.startswith( ("transformers_vit", "transformers_blip2", "transformers_whisper", "transformers_deepseek") ): continue check_lazy_init(entry, verbose=True, default_device=default_device) if __name__ == "__main__": test_models_lazy_init("transformers", "cpu")