mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-29 02:14:29 +00:00
96780e6ee4
* add pipeline shared module wrapper and update load batch * added model parallel process group for amp and clip grad (#86) * added model parallel process group for amp and clip grad * update amp and clip with model parallel process group * remove pipeline_prev/next group (#88) * micro batch offload * optimize pipeline gpu memory usage * pipeline can receive tensor shape (#93) * optimize pipeline gpu memory usage * fix grad accumulation step counter * rename classes and functions Co-authored-by: Frank Lee <somerlee.9@gmail.com>
41 lines
1.0 KiB
Python
41 lines
1.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
ALLOWED_MODES = [None, '1d', '2d', '2.5d', '3d', 'sequence']
|
|
TENSOR_PARALLEL_MODE = 'tensor_parallel_mode'
|
|
|
|
# intializer
|
|
INITIALIZER_MAPPING = {
|
|
'data': 'Initializer_Data',
|
|
'tensor': 'Initializer_Tensor',
|
|
'pipeline': 'Initializer_Pipeline',
|
|
'embedding': 'Initializer_Embedding',
|
|
'1d': 'Initializer_1D',
|
|
'2d': 'Initializer_2D',
|
|
'2.5d': 'Initializer_2p5D',
|
|
'3d': 'Initializer_3D',
|
|
'sequence': 'Initializer_Sequence',
|
|
'model': 'Initializer_Model'
|
|
}
|
|
|
|
# 1D parallel
|
|
PARALLEL_INPUT_1D = 'parallel_input_1d'
|
|
|
|
# 2D paralllel
|
|
SUMMA_DIM = 'SUMMA_DIM'
|
|
|
|
# 2.5D paralllel
|
|
TESSERACT_DIM = 'TESSERACT_DIM'
|
|
TESSERACT_DEP = 'TESSERACT_DEP'
|
|
|
|
# 3D parallel
|
|
DEPTH_3D = 'DEPTH_3D'
|
|
INPUT_GROUP_3D = 'PARALLEL_3D_INPUT'
|
|
WEIGHT_GROUP_3D = 'PARALLEL_3D_WEIGHT'
|
|
OUTPUT_GROUP_3D = 'PARALLEL_3D_OUTPUT'
|
|
|
|
# Tensor parallel attributes
|
|
IS_TENSOR_PARALLEL = 'is_tensor_parallel'
|
|
NUM_PARTITIONS = 'num_partitions'
|
|
TENSOR_PARALLEL_ATTRIBUTES = [IS_TENSOR_PARALLEL, NUM_PARTITIONS]
|