mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-26 12:14:02 +00:00
[misc] update pre-commit and run all files (#4752)
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
This commit is contained in:
@@ -106,14 +106,12 @@ class Policy(ABC):
|
||||
This method is made abstractmethod with no default implementation because we want to the policy writer
|
||||
to take note of the feature supported by his/her model and policy.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def preprocess(self) -> nn.Module:
|
||||
r"""
|
||||
Perform some preprocessing of the model, like reshaping the embedding layer.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
|
||||
@@ -122,7 +120,6 @@ class Policy(ABC):
|
||||
and the value is the ModulePolicyDescription object. The ModulePolicyDescription object describes how the module
|
||||
will be transformed.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def postprocess(self) -> nn.Module:
|
||||
@@ -130,13 +127,13 @@ class Policy(ABC):
|
||||
Perform some postprocessing of the model, like binding the weight of embedding layer with
|
||||
the classifier layer
|
||||
"""
|
||||
pass
|
||||
|
||||
def append_or_create_submodule_replacement(
|
||||
self, description: Union[SubModuleReplacementDescription,
|
||||
List[SubModuleReplacementDescription]], policy: Dict[Union[str, nn.Module],
|
||||
ModulePolicyDescription],
|
||||
target_key: Union[str, nn.Module]) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
|
||||
self,
|
||||
description: Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]],
|
||||
policy: Dict[Union[str, nn.Module], ModulePolicyDescription],
|
||||
target_key: Union[str, nn.Module],
|
||||
) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
|
||||
r"""
|
||||
Append or create a new submodule replacement description to the policy for the given key.
|
||||
|
||||
@@ -161,8 +158,11 @@ class Policy(ABC):
|
||||
return policy
|
||||
|
||||
def append_or_create_method_replacement(
|
||||
self, description: Dict[str, Callable], policy: Dict[Union[str, nn.Module], ModulePolicyDescription],
|
||||
target_key: Union[str, nn.Module]) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
|
||||
self,
|
||||
description: Dict[str, Callable],
|
||||
policy: Dict[Union[str, nn.Module], ModulePolicyDescription],
|
||||
target_key: Union[str, nn.Module],
|
||||
) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
|
||||
r"""
|
||||
Append or create a new method replacement description to the policy for the given key.
|
||||
|
||||
@@ -199,9 +199,7 @@ class Policy(ABC):
|
||||
|
||||
@staticmethod
|
||||
def distribute_layers(num_layers: int, num_stages: int) -> List[int]:
|
||||
"""Divide layers into stages
|
||||
|
||||
"""
|
||||
"""Divide layers into stages"""
|
||||
quotient = num_layers // num_stages
|
||||
remainder = num_layers % num_stages
|
||||
|
||||
|
Reference in New Issue
Block a user