mirror of https://github.com/hpcaitech/ColossalAI
113 lines
4.6 KiB
Python
113 lines
4.6 KiB
Python
from typing import Dict
|
|
from colossalai.tensor import ColoParameter, ComputeSpec, ProcessGroup
|
|
from colossalai.tensor import distspec
|
|
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, pg: ProcessGroup, 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_compute_spec():
|
|
cur_compute_pattern = param.compute_spec.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, pg)
|
|
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_compute_spec():
|
|
if dist_spec != param.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, pg=pg, recursive=True)
|
|
|
|
|
|
def init_colo_module(module: torch.nn.Module,
|
|
compute_spec: ComputeSpec,
|
|
pg: ProcessGroup,
|
|
recursive=True,
|
|
mode='default'):
|
|
compute_pattern = compute_spec.compute_pattern
|
|
if is_colo_module(module):
|
|
# for each param
|
|
# set DistSpec and ComputeSpec
|
|
colo_module = get_colo_module(module)
|
|
colo_module.register(compute_pattern, pg)
|
|
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):
|
|
param.set_dist_spec(dist_spec)
|
|
param.compute_spec = compute_spec
|
|
for mod in param.shared_param_modules:
|
|
modules_update_param.add(mod)
|
|
for mod in modules_update_param:
|
|
check_colo_module(mod, pg, recursive=False)
|
|
if recursive == True:
|
|
for submodule in module.children():
|
|
init_colo_module(submodule, compute_spec, pg=pg, recursive=True, mode=mode)
|