[refactory] add nn.parallel module (#1068)

colossalai/nn/parallel/layers/__init__.py

@@ -0,0 +1,15 @@
+from .colo_module import ColoModule
+from .linear import ColoLinear
+from .embedding import ColoEmbedding
+from .module_utils import register_colo_module, is_colo_module, get_colo_module, init_colo_module, check_colo_module
+
+__all__ = [
+    'ColoModule',
+    'register_colo_module',
+    'is_colo_module',
+    'get_colo_module',
+    'init_colo_module',
+    'check_colo_module',
+    'ColoLinear',
+    'ColoEmbedding',
+]

colossalai/nn/parallel/layers/colo_module.py

@@ -0,0 +1,56 @@
+from colossalai.tensor.distspec import _DistSpec
+from colossalai.tensor import ComputePattern
+from typing import List, Dict
+
+
+class ColoModule(object):
+
+    def __init__(self):
+        self._shard_params: List[str] = []
+        # Example:
+        # {ComputePattern.TP1D:
+        #    'default':
+        #      'weight':
+        #          distspec.shard(xxxxx)
+        #      'bias':
+        #          distspec.shard(xxxxx)
+        #    'row': ...
+        #    'col': ...
+        # }
+        self._allowed_patterns: Dict[ComputePattern, Dict[str, Dict[str, _DistSpec]]] = {}
+
+    def _register_shard_params(self, params: List[str]):
+        self._shard_params = params
+
+    def _register_allowed_patterns(self,
+                                   compute_pattern: ComputePattern,
+                                   dist_specs: Dict[str, _DistSpec],
+                                   mode='default'):
+        assert list(
+            dist_specs.keys()).sort() == self._shard_params.sort(), 'Every registered param should have dist_spec.'
+        if not compute_pattern in self._allowed_patterns:
+            self._allowed_patterns[compute_pattern] = {}
+        self._allowed_patterns[compute_pattern][mode] = dist_specs
+
+    def _set_default(self, compute_pattern: ComputePattern, target_mode):
+        self._allowed_patterns[compute_pattern]['default'] = self._allowed_patterns[compute_pattern][target_mode]
+
+    def has_compute_pattern(self, compute_pattern: ComputePattern):
+        return compute_pattern in self._allowed_patterns
+
+    def get_dist_specs(self, compute_pattern: ComputePattern):
+        assert self.has_compute_pattern(compute_pattern)
+        return self._allowed_patterns[compute_pattern]
+
+    def has_compute_pattern_with_mode(self, compute_pattern: ComputePattern, mode='default'):
+        return compute_pattern in self._allowed_patterns and mode in self._allowed_patterns[compute_pattern]
+
+    def get_dist_specs_with_mode(self, compute_pattern: ComputePattern, mode='default'):
+        assert self.has_compute_pattern_with_mode(compute_pattern, mode)
+        return self._allowed_patterns[compute_pattern][mode]
+
+    def get_param_names(self):
+        return self._shard_params
+
+    def register(self, compute_pattern):
+        raise NotImplementedError

colossalai/nn/parallel/layers/embedding.py

@@ -0,0 +1,42 @@
+from .colo_module import ColoModule
+from colossalai.tensor import ComputePattern, distspec
+from colossalai.core import global_context as gpc
+from colossalai.context.parallel_mode import ParallelMode
+
+
+class ColoEmbedding(ColoModule):
+
+    def __init__(self):
+        super(ColoEmbedding, self).__init__()
+        self._register_shard_params(['weight'])
+
+    def register(self, compute_pattern):
+        if not compute_pattern in self._allowed_patterns:
+            if ComputePattern.TP1D == compute_pattern:
+                self._set_TP1D()
+
+    def _set_TP1D(self):
+        # TP1D Row Linear
+        _compute_pattern = ComputePattern.TP1D
+        self._register_allowed_patterns(
+            compute_pattern=_compute_pattern,
+            dist_specs={
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+            },
+            mode='row',
+        )
+
+        # TP1D Col Linear
+        self._register_allowed_patterns(
+            compute_pattern=_compute_pattern,
+            dist_specs={
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+            },
+            mode='col',
+        )
+
+        self._set_default(compute_pattern=_compute_pattern, target_mode='row')

colossalai/nn/parallel/layers/linear.py

@@ -0,0 +1,47 @@
+from .colo_module import ColoModule
+from colossalai.tensor import ComputePattern, distspec
+from colossalai.core import global_context as gpc
+from colossalai.context.parallel_mode import ParallelMode
+
+
+class ColoLinear(ColoModule):
+
+    def __init__(self):
+        super(ColoLinear, self).__init__()
+        self._register_shard_params(['weight', 'bias'])
+
+    def register(self, compute_pattern):
+        if not compute_pattern in self._allowed_patterns:
+            if ComputePattern.TP1D == compute_pattern:
+                self._set_TP1D()
+
+    def _set_TP1D(self):
+        # TP1D Row Linear
+        _compute_pattern = ComputePattern.TP1D
+        self._register_allowed_patterns(
+            compute_pattern=_compute_pattern,
+            dist_specs={
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+                'bias':
+                    None
+            },
+            mode='row',
+        )
+
+        # TP1D Col Linear
+        self._register_allowed_patterns(
+            compute_pattern=_compute_pattern,
+            dist_specs={
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+                'bias':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)])
+            },
+            mode='col',
+        )
+
+        self._set_default(compute_pattern=_compute_pattern, target_mode='row')

colossalai/nn/parallel/layers/module_utils.py

@@ -0,0 +1,107 @@
+from typing import Dict
+from colossalai.tensor import ColoParameter, ParallelAction, TensorSpec
+from . import ColoModule
+import torch
+
+_COLOSSAL_MODULES: Dict[type, ColoModule] = {}
+
+
+def register_colo_module(module_type: type, colo_module: ColoModule):
+    global _COLOSSAL_MODULES
+    _COLOSSAL_MODULES[module_type] = colo_module
+
+
+def is_colo_module(module: torch.nn.Module):
+    global _COLOSSAL_MODULES
+    for module_type in _COLOSSAL_MODULES.keys():
+        if isinstance(module, module_type):
+            return True
+    return False
+
+
+def get_colo_module(module: torch.nn.Module):
+    global _COLOSSAL_MODULES
+    if is_colo_module(module):
+        for module_type, colo_module in _COLOSSAL_MODULES.items():
+            if isinstance(module, module_type):
+                return colo_module
+    else:
+        return None
+
+
+def check_colo_module(module: torch.nn.Module, recursive=True):
+    if is_colo_module(module):
+        colo_module = get_colo_module(module)
+        param_names = colo_module.get_param_names()
+        compute_pattern = None
+        for param_name in param_names:
+            param = module.get_parameter(param_name)
+            if not isinstance(param, ColoParameter):
+                raise Exception(f'Invalid ColoParameter spec: {param} in {module} is not a ColoParameter.')
+            if param.has_spec():
+                cur_compute_pattern = param.spec.parallel_action.compute_pattern
+                if compute_pattern is None:
+                    compute_pattern = cur_compute_pattern
+                else:
+                    if cur_compute_pattern != compute_pattern:
+                        raise Exception(
+                            f'Invalid ColoParameter spec: Params in {module} have different compute_pattern.')
+            else:
+                continue
+
+        if compute_pattern is not None:
+            colo_module.register(compute_pattern)
+            if not colo_module.has_compute_pattern(compute_pattern):
+                raise Exception(
+                    f'Invalid ColoParameter spec: ComputePattern {compute_pattern} in {module} is not allowed.')
+
+            match_specs = False
+            allowed_specs = colo_module.get_dist_specs(compute_pattern)
+            for _, param_specs in allowed_specs.items():
+                cur_match = True
+                for param_name, dist_spec in param_specs.items():
+                    param = module.get_parameter(param_name)
+                    if param.has_spec():
+                        if dist_spec != param.spec.dist_spec:
+                            cur_match = False
+                            break
+                    else:
+                        if dist_spec is not None:
+                            cur_match = False
+                            break
+                if cur_match == True:
+                    match_specs = True
+                    break
+            if match_specs == False:
+                raise Exception(f'Invalid ColoParameter spec: Params in {module} are incorrectly sharded.')
+    if recursive == True:
+        for submodule in module.children():
+            check_colo_module(submodule, recursive=True)
+
+
+def init_colo_module(module: torch.nn.Module, parallel_action: ParallelAction, recursive=True, mode='default'):
+    compute_pattern = parallel_action.compute_pattern
+    if is_colo_module(module):
+        # for each param
+        # set DistSpec and ParallelAction
+        colo_module = get_colo_module(module)
+        colo_module.register(compute_pattern)
+        if not colo_module.has_compute_pattern_with_mode(compute_pattern, mode=mode):
+            raise NotImplementedError
+        # a set for modules which update at least one param in the init process.
+        # these modules need to be checked whether all params still match one of the valid compute pattern.
+        modules_update_param = {module}
+        for param_name, dist_spec in colo_module.get_dist_specs_with_mode(compute_pattern, mode=mode).items():
+            if dist_spec is None:
+                continue
+            param = module.get_parameter(param_name)
+            if isinstance(param, ColoParameter):
+                spec = TensorSpec(dist_spec, parallel_action)
+                param.set_spec(spec)
+                for mod in param.shared_param_modules:
+                    modules_update_param.add(mod)
+        for mod in modules_update_param:
+            check_colo_module(mod, recursive=False)
+    if recursive == True:
+        for submodule in module.children():
+            init_colo_module(submodule, parallel_action, recursive=True, mode=mode)