ColossalAI/colossalai/nn/layer/parallel_2p5d/_utils.py

26 lines
1.1 KiB
Python
Raw Normal View History

2021-10-28 16:21:23 +00:00
import os
from colossalai.context.parallel_mode import ParallelMode
from colossalai.core import global_context as gpc
def get_tesseract_dim_dep_from_env():
try:
tesseract_dim = int(os.environ['TESSERACT_DIM'])
tesseract_dep = int(os.environ['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 as e:
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'