mirror of https://github.com/hpcaitech/ColossalAI
31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
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
|
|
|
|
|
|
def get_tesseract_dim_dep_from_env():
|
|
try:
|
|
tesseract_dim = env.tesseract_dim
|
|
tesseract_dep = env.tesseract_dep
|
|
assert tesseract_dim > 0, "TESSERACT_DIM must be larger than zero"
|
|
assert tesseract_dep > 0, "TESSERACT_DEP must be larger than zero"
|
|
return tesseract_dim, tesseract_dep
|
|
|
|
except KeyError:
|
|
raise EnvironmentError(
|
|
"TESSERACT_DIM or TESSERACT_DEP is not found in the current environment, "
|
|
"please make sure that you have used the correct process group initializer"
|
|
)
|
|
|
|
|
|
def assert_tesseract_initialization():
|
|
assert (
|
|
gpc.is_initialized(ParallelMode.PARALLEL_2P5D_COL)
|
|
and gpc.is_initialized(ParallelMode.PARALLEL_2P5D_ROW)
|
|
and gpc.is_initialized(ParallelMode.PARALLEL_2P5D_DEP)
|
|
and gpc.is_initialized(ParallelMode.PARALLEL_2P5D_XZ)
|
|
), (
|
|
"Both PARALLEL_2P5D_COL, PARALLEL_2P5D_ROW, PARALLEL_2P5D_DEP and PARALLEL_2P5D_XZ "
|
|
"must be initialized by the process group initializer"
|
|
)
|