Making large AI models cheaper, faster and more accessible
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 

32 lines
978 B

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
ALLOWED_MODES = [None, "1d", "2d", "2.5d", "3d", "sequence"]
TENSOR_PARALLEL_MODE = "tensor_parallel_mode"
# initializer
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",
"moe": "Initializer_Moe",
}
# 3D parallelism groups
INPUT_GROUP_3D = "input_group_3d"
WEIGHT_GROUP_3D = "weight_group_3d"
OUTPUT_GROUP_3D = "output_group_3d"
INPUT_X_WEIGHT_3D = "input_x_weight_group_3d"
OUTPUT_X_WEIGHT_3D = "output_x_weight_group_3d"
# Attributes of tensor parallel parameters
IS_TENSOR_PARALLEL = "is_tensor_parallel"
NUM_PARTITIONS = "num_partitions"
TENSOR_PARALLEL_ATTRIBUTES = [IS_TENSOR_PARALLEL, NUM_PARTITIONS]