ColossalAI/colossalai/utils/model/utils.py

116 lines
3.8 KiB
Python
Raw Normal View History

# This code has been adapted from the DeepSpeed library.
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import functools
from typing import Optional
import torch
def substitute_init_recursively(cls, func, visited: set):
for subcls in cls.__subclasses__():
substitute_init_recursively(subcls, func, visited)
if subcls not in visited:
func(subcls)
visited.add(subcls)
def call_to_str(base, *args, **kwargs):
"""Construct a string representation of a call.
Args:
base (str): name of the call
args (tuple, optional): args to ``base``
kwargs (dict, optional): kwargs supplied to ``base``
Returns:
str: A string representation of base(*args, **kwargs)
"""
name = f'{base}('
if args:
name += ', '.join(repr(arg) for arg in args)
if kwargs:
name += ', '
if kwargs:
name += ', '.join(f'{key}={repr(arg)}' for key, arg in kwargs.items())
name += ')'
return name
class InsertPostInitMethodToModuleSubClasses(object):
def __init__(self, default_dtype: Optional[torch.dtype] = None):
self._old_default_dtype = None
self._default_dtype = default_dtype
def __enter__(self):
r"""
Enter the context scope.
"""
if self._default_dtype is not None:
self._old_default_dtype = torch.get_default_dtype()
torch.set_default_dtype(self._default_dtype)
def preprocess_after(f):
@functools.wraps(f)
def wrapper(module: torch.nn.Module, *args, **kwargs):
f(module, *args, **kwargs)
self._post_init_method(module, *args, **kwargs)
return wrapper
def _enable_class(cls):
cls._old_init = cls.__init__
cls.__init__ = preprocess_after(cls.__init__)
# The function is called during init subclass.
def _init_subclass(cls, **kwargs):
cls.__init__ = preprocess_after(cls.__init__)
# Replace .__init__() for all existing subclasses of torch.nn.Module
# Execution self._post_init_method after the default init function.
substitute_init_recursively(torch.nn.modules.module.Module, _enable_class, set())
# holding on to the current __init__subclass__ for exit
torch.nn.modules.module.Module._old_init_subclass = (torch.nn.modules.module.Module.__init_subclass__)
# Replace .__init__() for future subclasses of torch.nn.Module
torch.nn.modules.module.Module.__init_subclass__ = classmethod(_init_subclass)
self._pre_context_exec()
return self
def __exit__(self, exc_type, exc_value, traceback):
if self._default_dtype is not None:
torch.set_default_dtype(self._old_default_dtype)
def _disable_class(cls):
if not hasattr(cls, '_old_init'):
raise AttributeError(
f"_old_init is not found in the {cls.__name__}, please make sure that you have imported {cls.__name__} before entering the context."
)
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, set())
# Replace .__init__() for future subclasses of torch.nn.Module
torch.nn.modules.module.Module.__init_subclass__ = (torch.nn.modules.module.Module._old_init_subclass)
self._post_context_exec()
# Now that we cleaned up the metaclass injection, raise the exception.
if exc_type is not None:
return False
# To be implemented by inheriting classes
def _post_init_method(self, module, *args, **kwargs):
pass
def _pre_context_exec(self):
pass
def _post_context_exec(self):
pass