github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 18:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[lazy] support torch 2.0 (#4763)

Browse Source 
* [lazy] support _like methods and clamp

* [lazy] pass transformers models

* [lazy] fix device move and requires grad

* [lazy] fix requires grad and refactor api

* [lazy] fix requires grad
This commit is contained in:
Hongxin Liu
2023-09-21 16:30:23 +08:00
committed by GitHub
parent 901ab1eedd
commit 3e05c07bb8
5 changed files with 273 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/test_lazy/test_models.py
View File

@@ -11,14 +11,12 @@ def test_torchvision_models_lazy_init(subset, default_device):
sub_model_zoo = model_zoo.get_sub_registry(subset)
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_llama")
or name.startswith(("transformers_vit", "transformers_blip2"))
if name in ("torchaudio_wav2vec2_base", "torchaudio_hubert_base") or name.startswith(
("transformers_vit", "transformers_blip2")
):
continue
check_lazy_init(entry, verbose=True, default_device=default_device)
if __name__ == "__main__":
test_torchvision_models_lazy_init("torchvision")
test_torchvision_models_lazy_init("transformers", "cpu")