2023-05-22 07:02:17 +00:00
|
|
|
# part of code modified from https://github.com/tunib-ai/parallelformers
|
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
from abc import ABC, abstractmethod
|
2023-06-07 08:09:40 +00:00
|
|
|
from dataclasses import dataclass
|
2023-07-05 06:16:55 +00:00
|
|
|
from typing import Any, Callable, Dict, List, Optional, Union
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-07-07 07:41:00 +00:00
|
|
|
import numpy as np
|
2023-05-22 07:02:17 +00:00
|
|
|
import torch.nn as nn
|
2023-07-05 06:16:55 +00:00
|
|
|
from torch import Tensor
|
|
|
|
from torch.nn import Module
|
|
|
|
|
|
|
|
from colossalai.pipeline.stage_manager import PipelineStageManager
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
from ..shard.shard_config import ShardConfig
|
2023-05-22 07:02:17 +00:00
|
|
|
|
2023-06-30 02:56:29 +00:00
|
|
|
__all__ = ["ParallelModule", "SubModuleReplacementDescription", "ModulePolicyDescription", "Policy"]
|
|
|
|
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
class ParallelModule():
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
def __init__(self):
|
|
|
|
pass
|
2023-05-22 07:02:17 +00:00
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
2023-06-15 09:55:42 +00:00
|
|
|
class SubModuleReplacementDescription:
|
2023-05-24 02:26:46 +00:00
|
|
|
r"""
|
2023-06-15 09:55:42 +00:00
|
|
|
Describe how a submodule will be replaced
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-07-03 07:29:11 +00:00
|
|
|
Args:
|
|
|
|
suffix (str): used to get the submodule object
|
|
|
|
target_module (ParallelModule): specifies the module class used to replace to submodule
|
|
|
|
kwargs (Dict[str, Any]): the dictionary used to pass extra arguments to the `ParallelModule.from_native_module` method.
|
|
|
|
ignore_if_not_exist (bool): if the submodule does not exist, ignore it or raise an exception
|
2023-05-22 07:02:17 +00:00
|
|
|
"""
|
2023-06-15 09:55:42 +00:00
|
|
|
suffix: str
|
|
|
|
target_module: ParallelModule
|
|
|
|
kwargs: Dict[str, Any] = None
|
2023-06-19 09:57:37 +00:00
|
|
|
ignore_if_not_exist: bool = False
|
2023-05-22 07:02:17 +00:00
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
2023-06-15 09:55:42 +00:00
|
|
|
class ModulePolicyDescription:
|
2023-05-24 02:26:46 +00:00
|
|
|
r"""
|
2023-07-03 07:29:11 +00:00
|
|
|
Describe how the attributes and parameters will be transformed in a policy.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
attribute_replacement (Dict[str, Any]): key is the attribute name, value is the attribute value after sharding
|
|
|
|
param_replacement (List[Callable]): a list of functions to perform in-place param replacement. The function
|
|
|
|
must receive only one arguments: module. One example is
|
|
|
|
|
|
|
|
```python
|
|
|
|
def example_replace_weight(module: torch.nn.Module):
|
|
|
|
weight = module.weight
|
|
|
|
new_weight = shard_rowwise(weight, process_group)
|
|
|
|
module.weight = torch.nn.Parameter(new_weight)
|
|
|
|
```
|
|
|
|
sub_module_replacement (List[SubModuleReplacementDescription]): each element in the list is a ParamReplacementDescription
|
|
|
|
object which specifies the module to be replaced and the target module used to replacement.
|
|
|
|
method_replace (Dict[str, Callable]): key is the method name, value is the method for replacement
|
2023-06-15 08:50:08 +00:00
|
|
|
"""
|
2023-07-03 07:29:11 +00:00
|
|
|
attribute_replacement: Dict[str, Any] = None
|
|
|
|
param_replacement: List[Callable] = None
|
|
|
|
sub_module_replacement: List[SubModuleReplacementDescription] = None
|
|
|
|
method_replacement: Dict[str, Callable] = None
|
2023-06-15 08:50:08 +00:00
|
|
|
|
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
class Policy(ABC):
|
2023-05-24 02:26:46 +00:00
|
|
|
r"""
|
2023-07-03 07:29:11 +00:00
|
|
|
The base class for all the policies. For each different model, it should have a different policy class,
|
|
|
|
like BertPolicy for Bert Model or OPTPolicy for OPT model.
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-07-03 07:29:11 +00:00
|
|
|
Shardformer has provided many built-in sharding policies for the mainstream models. You can use the
|
2023-07-04 09:53:39 +00:00
|
|
|
built-in policies by setting `policy = None`, which is already the default argument for `Shardformer.optimize`.
|
2023-07-03 07:29:11 +00:00
|
|
|
If you want to define your own policy, you can inherit from this class and overwrite the methods you want to modify.
|
2023-05-22 07:02:17 +00:00
|
|
|
"""
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
def __init__(self) -> None:
|
2023-07-05 06:16:55 +00:00
|
|
|
self.shard_config: Optional[ShardConfig] = None
|
|
|
|
self.model: Optional[Module] = None
|
2023-06-15 09:55:42 +00:00
|
|
|
|
|
|
|
def set_model(self, model: nn.Module) -> None:
|
2023-05-24 02:26:46 +00:00
|
|
|
r"""
|
2023-06-15 09:55:42 +00:00
|
|
|
Set model as an attribute of the Policy object so that we can access the model's attributes.
|
2023-05-22 07:02:17 +00:00
|
|
|
|
|
|
|
Args:
|
2023-06-15 09:55:42 +00:00
|
|
|
model (:class:`nn.Module`): The model to be perform
|
|
|
|
"""
|
|
|
|
self.model = model
|
|
|
|
|
2023-06-19 02:47:16 +00:00
|
|
|
def set_shard_config(self, shard_config: ShardConfig) -> None:
|
|
|
|
r"""
|
|
|
|
Set shard config as an attribute of the Policy object.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
shard_config (:class:`ShardConfig`): The shard config to be perform
|
|
|
|
"""
|
|
|
|
self.shard_config = shard_config
|
2023-06-30 01:32:37 +00:00
|
|
|
self.config_sanity_check()
|
|
|
|
|
2023-07-05 06:16:55 +00:00
|
|
|
@property
|
|
|
|
def pipeline_stage_manager(self) -> Optional[PipelineStageManager]:
|
|
|
|
if self.shard_config is not None:
|
|
|
|
return self.shard_config.pipeline_stage_manager
|
|
|
|
return None
|
|
|
|
|
2023-06-30 01:32:37 +00:00
|
|
|
@abstractmethod
|
|
|
|
def config_sanity_check(self):
|
|
|
|
"""
|
|
|
|
Check if the shard config is valid for the model. Raise an exception if the config is invalid.
|
2023-07-03 07:29:11 +00:00
|
|
|
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.
|
2023-06-30 01:32:37 +00:00
|
|
|
"""
|
|
|
|
pass
|
2023-06-19 02:47:16 +00:00
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
@abstractmethod
|
2023-06-19 02:47:16 +00:00
|
|
|
def preprocess(self) -> nn.Module:
|
2023-06-15 09:55:42 +00:00
|
|
|
r"""
|
2023-07-03 07:29:11 +00:00
|
|
|
Perform some preprocessing of the model, like reshaping the embedding layer.
|
2023-06-15 09:55:42 +00:00
|
|
|
"""
|
2023-06-19 05:53:17 +00:00
|
|
|
pass
|
2023-06-15 09:55:42 +00:00
|
|
|
|
|
|
|
@abstractmethod
|
2023-06-19 02:47:16 +00:00
|
|
|
def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
|
2023-06-15 09:55:42 +00:00
|
|
|
r"""
|
2023-07-03 07:29:11 +00:00
|
|
|
This method returns the module policy, which is a dictionary. The key is the module name or the module object,
|
|
|
|
and the value is the ModulePolicyDescription object. The ModulePolicyDescription object describes how the module
|
|
|
|
will be transformed.
|
2023-05-22 07:02:17 +00:00
|
|
|
"""
|
2023-06-19 05:53:17 +00:00
|
|
|
pass
|
2023-05-24 02:26:46 +00:00
|
|
|
|
2023-06-15 09:55:42 +00:00
|
|
|
@abstractmethod
|
|
|
|
def postprocess(self) -> nn.Module:
|
2023-05-24 02:26:46 +00:00
|
|
|
r"""
|
2023-06-15 09:55:42 +00:00
|
|
|
Perform some postprocessing of the model, like binding the weight of embedding layer with
|
|
|
|
the classifier layer
|
2023-05-22 07:02:17 +00:00
|
|
|
"""
|
2023-06-19 05:53:17 +00:00
|
|
|
pass
|
2023-07-04 01:57:03 +00:00
|
|
|
|
|
|
|
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]:
|
|
|
|
r"""
|
|
|
|
Append or create a new submodule replacement description to the policy for the given key.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
submodule_replace_desc (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended
|
|
|
|
policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated
|
|
|
|
target_key (Union[str, nn.Module]): the key of the policy to be updated
|
|
|
|
"""
|
|
|
|
# convert to list
|
|
|
|
if isinstance(description, SubModuleReplacementDescription):
|
|
|
|
description = [description]
|
|
|
|
|
|
|
|
# append or create a new description
|
|
|
|
if target_key in policy:
|
2023-07-14 01:51:53 +00:00
|
|
|
if policy[target_key].sub_module_replacement is None:
|
|
|
|
policy[target_key].sub_module_replacement = description
|
|
|
|
else:
|
|
|
|
policy[target_key].sub_module_replacement.extend(description)
|
2023-07-04 01:57:03 +00:00
|
|
|
else:
|
|
|
|
policy[target_key] = ModulePolicyDescription(sub_module_replacement=description)
|
|
|
|
|
|
|
|
return policy
|
2023-07-05 06:16:55 +00:00
|
|
|
|
2023-07-11 07:23:33 +00:00
|
|
|
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]:
|
|
|
|
r"""
|
|
|
|
Append or create a new method replacement description to the policy for the given key.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
description (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended
|
|
|
|
policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated
|
|
|
|
target_key (Union[str, nn.Module]): the key of the policy to be updated
|
|
|
|
"""
|
|
|
|
if target_key in policy:
|
2023-07-14 01:51:53 +00:00
|
|
|
if policy[target_key].method_replacement is None:
|
|
|
|
policy[target_key].method_replacement = description
|
|
|
|
else:
|
|
|
|
policy[target_key].method_replacement.update(description)
|
2023-07-11 07:23:33 +00:00
|
|
|
else:
|
|
|
|
policy[target_key] = ModulePolicyDescription(method_replacement=description)
|
|
|
|
|
|
|
|
return policy
|
|
|
|
|
2023-07-05 06:16:55 +00:00
|
|
|
def get_held_layers(self) -> List[Module]:
|
|
|
|
"""Get layers that should be held in current stage. This method should be implemented by subclass.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
List[Module]: List of layers that should be hold in current stage
|
|
|
|
"""
|
|
|
|
raise NotImplementedError
|
|
|
|
|
|
|
|
def get_shared_params(self) -> List[Dict[int, Tensor]]:
|
|
|
|
"""Get parameters that should be shared across stages. This method should be implemented by subclass.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
List[Dict[int, Tensor]]: List of parameters that should be shared across stages. E.g. [{0: module.model.embed_tokens.weight, 3: module.lm_head.weight}]
|
|
|
|
"""
|
|
|
|
return []
|
2023-07-07 07:41:00 +00:00
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
def distribute_layers(num_layers: int, num_stages: int) -> List[int]:
|
|
|
|
"""Divide layers into stages
|
|
|
|
|
|
|
|
"""
|
|
|
|
quotient = num_layers // num_stages
|
|
|
|
remainder = num_layers % num_stages
|
|
|
|
|
|
|
|
# calculate the num_layers per stage
|
|
|
|
layers_per_stage = [quotient] * num_stages
|
|
|
|
|
|
|
|
# deal with the rest layers
|
|
|
|
if remainder > 0:
|
2023-07-10 05:58:58 +00:00
|
|
|
start_position = num_stages // 2 - remainder // 2
|
2023-07-07 07:41:00 +00:00
|
|
|
for i in range(start_position, start_position + remainder):
|
|
|
|
layers_per_stage[i] += 1
|
|
|
|
return layers_per_stage
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
def get_stage_index(layers_per_stage: List[int], stage: int) -> List[int]:
|
|
|
|
"""
|
|
|
|
get the start index and end index of layers for each stage.
|
|
|
|
"""
|
|
|
|
num_layers_per_stage_accumulated = np.insert(np.cumsum(layers_per_stage), 0, 0)
|
|
|
|
|
|
|
|
start_idx = num_layers_per_stage_accumulated[stage]
|
|
|
|
end_idx = num_layers_per_stage_accumulated[stage + 1]
|
|
|
|
|
|
|
|
return [start_idx, end_idx]
|