[Tensor] add module handler for linear (#1021)

* add module spec for linear

* polish

* polish

* polish

colossalai/tensor/__init__.py

@@ -8,8 +8,12 @@ from ._ops import *
 from .optim.colo_optimizer import ColoOptimizer
 from . import distspec
 from .dist_spec_mgr import DistSpecManager
+from .module_utils import register_colo_module, is_colo_module, get_colo_module, init_colo_module, check_colo_module
+from .modules import ColoLinear
 
 __all__ = [
     'ColoTensor', 'convert_parameter', 'colo_op_impl', 'ComputePattern', 'TensorSpec', 'ParallelAction',
-    'named_params_with_colotensor', 'ColoOptimizer', 'ColoParameter', 'distspec', 'DistSpecManager'
+    'named_params_with_colotensor', 'ColoOptimizer', 'ColoParameter', 'distspec', 'DistSpecManager',
+    'register_colo_module', 'is_colo_module', 'get_colo_module', 'init_colo_module', 'check_colo_module',
+    'ColoLinear'
 ]

colossalai/tensor/module_utils.py

@@ -0,0 +1,92 @@
+from typing import Dict
+from colossalai.tensor import ColoParameter, ParallelAction, TensorSpec
+from .modules 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
+    return type(module) in _COLOSSAL_MODULES
+
+def get_colo_module(module: torch.nn.Module):
+    global _COLOSSAL_MODULES
+    if is_colo_module(module):
+        colo_module = _COLOSSAL_MODULES[type(module)]
+        colo_module.register()
+        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:
+            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, label='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)
+        if not colo_module.has_compute_pattern_with_label(compute_pattern, label=label):
+            raise NotImplementedError
+        for param_name, dist_spec in colo_module.get_dist_specs_with_label(compute_pattern, label=label).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)
+        check_colo_module(module, recursive=False)
+    if recursive == True:
+        for submodule in module.children():
+            init_colo_module(submodule, parallel_action, recursive=True, label=label)
+    

colossalai/tensor/modules/__init__.py

@@ -0,0 +1,2 @@
+from .colo_module import ColoModule
+from .linear import ColoLinear

colossalai/tensor/modules/colo_module.py

@@ -0,0 +1,51 @@
+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], label='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][label] = dist_specs
+
+    def _set_default(self, compute_pattern: ComputePattern, target_label):
+        self._allowed_patterns[compute_pattern]['default'] = self._allowed_patterns[compute_pattern][target_label]
+    
+    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_label(self, compute_pattern: ComputePattern, label='default'):
+        return compute_pattern in self._allowed_patterns and label in self._allowed_patterns[compute_pattern]
+        
+    def get_dist_specs_with_label(self, compute_pattern: ComputePattern, label='default'):
+        assert self.has_compute_pattern_with_label(compute_pattern, label)
+        return self._allowed_patterns[compute_pattern][label]
+
+    def get_param_names(self):
+        return self._shard_params
+
+    def register(self):
+        raise NotImplementedError

colossalai/tensor/modules/linear.py

@@ -0,0 +1,39 @@
+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'])
+        self._register = False
+        
+    def register(self):
+        if self._register == False:
+            self._set_TP1D()
+            self._register = True
+        
+    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
+            },
+            label='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)])
+            },
+            label='col',
+        )
+
+        self._set_default(compute_pattern=_compute_pattern, target_label='row')