2022-05-26 03:50:44 +00:00
|
|
|
from typing import Dict
|
2023-09-11 08:24:28 +00:00
|
|
|
|
2022-05-26 03:50:44 +00:00
|
|
|
import torch
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
from colossalai.legacy.tensor import ComputeSpec, ProcessGroup
|
2023-09-18 08:31:06 +00:00
|
|
|
from colossalai.tensor import ColoParameter
|
2023-09-11 08:24:28 +00:00
|
|
|
|
|
|
|
from . import ColoModule
|
|
|
|
|
2022-05-26 03:50:44 +00:00
|
|
|
_COLOSSAL_MODULES: Dict[type, ColoModule] = {}
|
|
|
|
|
|
|
|
|
|
|
|
def register_colo_module(module_type: type, colo_module: ColoModule):
|
|
|
|
global _COLOSSAL_MODULES
|
|
|
|
_COLOSSAL_MODULES[module_type] = colo_module
|
|
|
|
|
2022-06-03 10:04:22 +00:00
|
|
|
|
2022-05-26 03:50:44 +00:00
|
|
|
def is_colo_module(module: torch.nn.Module):
|
|
|
|
global _COLOSSAL_MODULES
|
2022-06-03 04:09:49 +00:00
|
|
|
for module_type in _COLOSSAL_MODULES.keys():
|
2022-06-06 04:11:48 +00:00
|
|
|
if isinstance(module, module_type):
|
2022-06-03 04:09:49 +00:00
|
|
|
return True
|
|
|
|
return False
|
2022-05-26 03:50:44 +00:00
|
|
|
|
2022-06-03 10:04:22 +00:00
|
|
|
|
2022-05-26 03:50:44 +00:00
|
|
|
def get_colo_module(module: torch.nn.Module):
|
|
|
|
global _COLOSSAL_MODULES
|
|
|
|
if is_colo_module(module):
|
2022-06-03 04:09:49 +00:00
|
|
|
for module_type, colo_module in _COLOSSAL_MODULES.items():
|
2022-06-06 04:11:48 +00:00
|
|
|
if isinstance(module, module_type):
|
2022-06-03 04:09:49 +00:00
|
|
|
return colo_module
|
2022-05-26 03:50:44 +00:00
|
|
|
else:
|
|
|
|
return None
|
|
|
|
|
2022-06-03 10:04:22 +00:00
|
|
|
|
2022-07-04 10:54:37 +00:00
|
|
|
def check_colo_module(module: torch.nn.Module, pg: ProcessGroup, recursive=True):
|
2022-05-26 03:50:44 +00:00
|
|
|
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):
|
2023-09-19 06:20:26 +00:00
|
|
|
raise Exception(f"Invalid ColoParameter spec: {param} in {module} is not a ColoParameter.")
|
2022-06-24 05:08:54 +00:00
|
|
|
if param.has_compute_spec():
|
2022-07-06 08:15:16 +00:00
|
|
|
cur_compute_pattern = param.compute_spec.compute_pattern
|
2022-05-26 03:50:44 +00:00
|
|
|
if compute_pattern is None:
|
|
|
|
compute_pattern = cur_compute_pattern
|
|
|
|
else:
|
|
|
|
if cur_compute_pattern != compute_pattern:
|
2022-06-03 10:04:22 +00:00
|
|
|
raise Exception(
|
2023-09-19 06:20:26 +00:00
|
|
|
f"Invalid ColoParameter spec: Params in {module} have different compute_pattern."
|
|
|
|
)
|
2022-05-26 03:50:44 +00:00
|
|
|
else:
|
|
|
|
continue
|
2022-06-03 10:04:22 +00:00
|
|
|
|
2022-05-26 03:50:44 +00:00
|
|
|
if compute_pattern is not None:
|
2022-07-04 10:54:37 +00:00
|
|
|
colo_module.register(compute_pattern, pg)
|
2022-05-26 03:50:44 +00:00
|
|
|
if not colo_module.has_compute_pattern(compute_pattern):
|
2022-06-03 10:04:22 +00:00
|
|
|
raise Exception(
|
2023-09-19 06:20:26 +00:00
|
|
|
f"Invalid ColoParameter spec: ComputePattern {compute_pattern} in {module} is not allowed."
|
|
|
|
)
|
2022-05-26 03:50:44 +00:00
|
|
|
|
|
|
|
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)
|
2022-06-24 05:08:54 +00:00
|
|
|
if param.has_compute_spec():
|
2022-07-06 08:15:16 +00:00
|
|
|
if dist_spec != param.dist_spec:
|
2022-05-26 03:50:44 +00:00
|
|
|
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:
|
2023-09-19 06:20:26 +00:00
|
|
|
raise Exception(f"Invalid ColoParameter spec: Params in {module} are incorrectly sharded.")
|
2022-05-26 03:50:44 +00:00
|
|
|
if recursive == True:
|
|
|
|
for submodule in module.children():
|
2022-07-04 10:54:37 +00:00
|
|
|
check_colo_module(submodule, pg=pg, recursive=True)
|
2022-05-26 03:50:44 +00:00
|
|
|
|
2022-06-03 10:04:22 +00:00
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def init_colo_module(
|
|
|
|
module: torch.nn.Module, compute_spec: ComputeSpec, pg: ProcessGroup, recursive=True, mode="default"
|
|
|
|
):
|
2022-06-27 02:04:45 +00:00
|
|
|
compute_pattern = compute_spec.compute_pattern
|
2022-05-26 03:50:44 +00:00
|
|
|
if is_colo_module(module):
|
|
|
|
# for each param
|
2022-07-15 06:02:32 +00:00
|
|
|
# set its process_group, dist_spec and compute_spec
|
2022-05-26 03:50:44 +00:00
|
|
|
colo_module = get_colo_module(module)
|
2022-07-04 10:54:37 +00:00
|
|
|
colo_module.register(compute_pattern, pg)
|
2022-05-26 10:15:42 +00:00
|
|
|
if not colo_module.has_compute_pattern_with_mode(compute_pattern, mode=mode):
|
2022-05-26 03:50:44 +00:00
|
|
|
raise NotImplementedError
|
2022-05-26 10:15:42 +00:00
|
|
|
# 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():
|
2022-05-26 03:50:44 +00:00
|
|
|
if dist_spec is None:
|
|
|
|
continue
|
|
|
|
param = module.get_parameter(param_name)
|
|
|
|
if isinstance(param, ColoParameter):
|
2022-07-15 06:02:32 +00:00
|
|
|
param.set_process_group(pg)
|
2022-07-06 08:15:16 +00:00
|
|
|
param.set_dist_spec(dist_spec)
|
|
|
|
param.compute_spec = compute_spec
|
2022-05-26 10:15:42 +00:00
|
|
|
for mod in param.shared_param_modules:
|
|
|
|
modules_update_param.add(mod)
|
|
|
|
for mod in modules_update_param:
|
2022-07-04 10:54:37 +00:00
|
|
|
check_colo_module(mod, pg, recursive=False)
|
2022-05-26 03:50:44 +00:00
|
|
|
if recursive == True:
|
|
|
|
for submodule in module.children():
|
2022-07-04 10:54:37 +00:00
|
|
|
init_colo_module(submodule, compute_spec, pg=pg, recursive=True, mode=mode)
|