github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

[lazy] support init on cuda (#4269)

Browse Source 
* [lazy] support init on cuda

* [test] update lazy init test

* [test] fix transformer version
This commit is contained in:
Hongxin Liu
2023-07-19 16:43:01 +08:00
committed by GitHub
parent 4b977541a8
commit fc5cef2c79
4 changed files with 31 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/test_lazy/lazy_init_utils.py
View File

@@ -61,14 +61,18 @@ def assert_forward_equal(m1: torch.nn.Module, m2: torch.nn.Module, data_gen_fn:
f'{m1.__class__.__name__} has inconsistent outputs, {out1} vs {out2}'
def check_lazy_init(entry: TestingEntry, seed: int = 42, verbose: bool = False, check_forward: bool = False) -> None:
def check_lazy_init(entry: TestingEntry,
seed: int = 42,
verbose: bool = False,
check_forward: bool = False,
default_device: str = 'cpu') -> None:
model_fn, data_gen_fn, output_transform_fn, _, model_attr = entry
_MyTensor._pre_op_fn = lambda *args: set_seed(seed)
LazyTensor._pre_op_fn = lambda *args: set_seed(seed)
ctx = LazyInitContext(tensor_cls=_MyTensor)
ctx = LazyInitContext(tensor_cls=_MyTensor, default_device=default_device)
with ctx:
model = model_fn()
ctx = LazyInitContext()
ctx = LazyInitContext(default_device=default_device)
with ctx:
deferred_model = model_fn()
copied_deferred_model = deepcopy(deferred_model)

5
tests/test_lazy/test_models.py
View File

@@ -6,13 +6,14 @@ from tests.kit.model_zoo import model_zoo
@pytest.mark.skipif(not SUPPORT_LAZY, reason='requires torch >= 1.12.0')
@pytest.mark.parametrize('subset', ['torchvision', 'diffusers', 'timm', 'transformers', 'torchaudio', 'deepfm', 'dlrm'])
def test_torchvision_models_lazy_init(subset):
@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)
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'):
continue
check_lazy_init(entry, verbose=True)
check_lazy_init(entry, verbose=True, default_device=default_device)
if __name__ == '__main__':