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

[context]use meta tensor to init model lazily. (#1187)

Browse Source 
* [CLI] add CLI launcher

* Revert "[CLI] add CLI launcher"

This reverts commit df7e6506d4.

* [context]use meta tensor to init model lazily.

* polish

* make module with device kwargs bypass the normal init.

* change unit test to adapt updated context.
This commit is contained in:
YuliangLiu0306
2022-06-29 21:02:30 +08:00
committed by GitHub
parent 2c8c05675d
commit 2053e138a2
3 changed files with 77 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/utils/model/utils.py
View File

@@ -3,9 +3,9 @@ import functools
from typing import Optional
def _substitute_init_recursively(cls, func):
def substitute_init_recursively(cls, func):
for subcls in cls.__subclasses__():
_substitute_init_recursively(subcls, func)
substitute_init_recursively(subcls, func)
func(subcls)
@@ -64,7 +64,7 @@ class InsertPostInitMethodToModuleSubClasses(object):
# Replace .__init__() for all existing subclasses of torch.nn.Module
# Excution self._post_init_method after the default init function.
_substitute_init_recursively(torch.nn.modules.module.Module, _enable_class)
substitute_init_recursively(torch.nn.modules.module.Module, _enable_class)
# holding on to the current __init__subclass__ for exit
torch.nn.modules.module.Module._old_init_subclass = (torch.nn.modules.module.Module.__init_subclass__)
@@ -87,7 +87,7 @@ class InsertPostInitMethodToModuleSubClasses(object):
cls.__init__ = cls._old_init
# Replace .__init__() for all existing subclasses of torch.nn.Module
_substitute_init_recursively(torch.nn.modules.module.Module, _disable_class)
substitute_init_recursively(torch.nn.modules.module.Module, _disable_class)
# Replace .__init__() for future subclasses of torch.nn.Module
torch.nn.modules.module.Module.__init_subclass__ = (torch.nn.modules.module.Module._old_init_subclass)