2023-09-18 08:31:06 +00:00
|
|
|
from colossalai.legacy.context.parallel_mode import ParallelMode
|
|
|
|
|
from colossalai.legacy.core import global_context as gpc
|
|
|
|
|
from colossalai.legacy.global_variables import tensor_parallel_env as env
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_summa_dim_from_env() -> int:
|
|
|
|
|
try:
|
2022-02-14 03:15:02 +00:00
|
|
|
summa_dim = env.summa_dim
|
2023-09-19 06:20:26 +00:00
|
|
|
assert summa_dim > 0, "SUMMA_DIM must be larger than zero"
|
2021-10-28 16:21:23 +00:00
|
|
|
return summa_dim
|
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
except KeyError:
|
|
|
|
|
raise EnvironmentError(
|
|
|
|
|
"SUMMA_DIM is not found in the current environment, "
|
|
|
|
|
"please make sure that you have used the correct process group initializer"
|
|
|
|
|
)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
|
|
|
|
|
|
def assert_summa_initialization():
|
2023-09-19 06:20:26 +00:00
|
|
|
assert gpc.is_initialized(ParallelMode.PARALLEL_2D_COL) and gpc.is_initialized(
|
|
|
|
|
ParallelMode.PARALLEL_2D_ROW
|
|
|
|
|
), "Both TWO_DIMENSION_COL and TWO_DIMENSION_ROW must be initialized by the process group initializer"
|
