该代码结构如下图所示：
在initialize_megatron初始化megatron的过程中，有关于数据并行、流水线并行、张量并行的实现，简介及其实现如下：
模型分布式环境初始化：
以两台分别有8个GPU服务器为例，训练具有12层的transformer layers，
图二
本示例将模型纵向切割为4部分，每部分3个layers，实现pipeline-parallel（流水线并行），将模型横向切割实现tensor-parallel（向量并行），把图二中的“1，2，3层”切割成两部分。
图三
上图说明了以model1为例，如何切割一个模型为八个部分，分别放入八个gpu的过程。
一个完整的模型model1的含义：
纵向三刀，把transformer layers的一共12层，切割成了四个部分，每个部分3个layers，其目的是实现pipeline-parallel；【需要pipeline_model_parallel_size=4】
横向的一刀，代表了tensor-parallel，是把（1，2，3）直到（10，11，12）这样的每三层layers，都切割成上下两个部分。【需要tensor_model_parallel_size=2】
tensor model-parallel groups：代表有多少个包含向量并行的groups，由图可知：
model1:[0, 1; 8, 9; 4, 5; 12, 13]
Model2:[2, 3; 10, 11; 6, 7; 14, 15]
对应代码示例中的：
8 tensor model-parallel groups:
[g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]
pipeline model-parallel groups：代表有多少个包含流水线并行的模型，由图可知：
模型model1先纵向切割为4份为流水线并行关系，然后横向切分，故有两个groups，第一个，[0,4,8,12]，第二个：[1,5,9,13]
同理model2。
data_parallel groups：数据并行groups，数据并行，是”含有相同参数的模型的子块“之间进行数据并行，有图可以看到两台服务器中的模型结构，（0、2相同），（1、3相同）,（4、6相同）,对应代码示例中的：
8 data_parallel groups:
[g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]
代码实现：
initialize_model_parallel(
tensor_model_parallel_size: int = 1,
pipeline_model_parallel_size: int = 1,
virtual_pipeline_model_parallel_size: Optional[int] = None,
pipeline_model_parallel_split_rank: Optional[int] = None,
use_fp8: bool = False,
)
tensor_model_parallel_size = 4
pipeline_model_parallel_size = 2
world_size = 16
data_parallel_size: int = world_size // (tensor_model_parallel_size * pipeline_model_parallel_size)
num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size = 4
num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size = 8
# Build the data-parallel groups.
#构建数据并行groups
all_data_parallel_group_ranks = []
for i in range(pipeline_model_parallel_size):
start_rank = i * num_pipeline_model_parallel_groups
end_rank = (i + 1) * num_pipeline_model_parallel_groups
for i in range(pipeline_model_parallel_size):
start_rank = i * num_pipeline_model_parallel_groups
end_rank = (i + 1) * num_pipeline_model_parallel_groups
for j in range(tensor_model_parallel_size):
ranks = range(start_rank + j, end_rank, tensor_model_parallel_size)
all_data_parallel_group_ranks.append(list(ranks))
group = torch.distributed.new_group(ranks)
group_gloo = torch.distributed.new_group(ranks, backend="gloo")
if rank in ranks:
_DATA_PARALLEL_GROUP = group
_DATA_PARALLEL_GROUP_GLOO = group_gloo
_DATA_PARALLEL_GLOBAL_RANKS = ranks
print(all_data_parallel_group_ranks)
all_data_parallel_group_ranks
[[0, 2], [1, 3], [4, 6], [5, 7], [8, 10], [9, 11], [12, 14], [13, 15]]
# Build the model-parallel groups.
#构建模型并行占用groups，即模型占用了哪些GPU
global _MODEL_PARALLEL_GROUP
assert _MODEL_PARALLEL_GROUP is None, 'model parallel group is already initialized'
for i in range(data_parallel_size):
ranks = [data_parallel_group_ranks[i] for data_parallel_group_ranks in all_data_parallel_group_ranks]
group = torch.distributed.new_group(ranks)
print(ranks)
if rank in ranks:
_MODEL_PARALLEL_GROUP = group
ranks
[0, 1, 4, 5, 8, 9, 12, 13]
[2, 3, 6, 7, 10, 11, 14, 15]
# Build the tensor model-parallel groups.
#构建张量并行groups
global _TENSOR_MODEL_PARALLEL_GROUP
assert _TENSOR_MODEL_PARALLEL_GROUP is None, 'tensor model parallel group is already initialized'
for i in range(num_tensor_model_parallel_groups):
ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)
group = torch.distributed.new_group(ranks)
print(ranks)
if rank in ranks:
_TENSOR_MODEL_PARALLEL_GROUP = group
[0, 1]
[2, 3]
[4, 5]
[6, 7]
[8, 9]
[10, 11]
[12, 13]
[14, 15]
# Build the pipeline model-parallel groups and embedding groups
#构建流水线并行groups和embedding groups
for i in range(num_pipeline_model_parallel_groups):
ranks = range(i, world_size, num_pipeline_model_parallel_groups)
print(ranks)
group = torch.distributed.new_group(ranks)
if rank in ranks:
_PIPELINE_MODEL_PARALLEL_GROUP = group
_PIPELINE_GLOBAL_RANKS = ranks
# Setup embedding group (to exchange gradients between
# first and last stages).
if len(ranks) > 1:
embedding_ranks = [ranks[0], ranks[-1]]
position_embedding_ranks = [ranks[0]]
print(embedding_ranks)
print(position_embedding_ranks)
if pipeline_model_parallel_split_rank is not None:
if ranks[pipeline_model_parallel_split_rank] not in embedding_ranks:
embedding_ranks = [ranks[0], ranks[pipeline_model_parallel_split_rank], ranks[-1]]
if ranks[pipeline_model_parallel_split_rank] not in position_embedding_ranks:
position_embedding_ranks = [ranks[0], ranks[pipeline_model_parallel_split_rank]]
else:
embedding_ranks = ranks
position_embedding_ranks = ranks
group = torch.distributed.new_group(embedding_ranks)
if rank in embedding_ranks:
_EMBEDDING_GROUP = group
if rank in ranks:
_EMBEDDING_GLOBAL_RANKS = embedding_ranks
group = torch.distributed.new_group(position_embedding_ranks)
if rank in position_embedding_ranks:
_POSITION_EMBEDDING_GROUP = group
if rank in ranks:
_POSITION_EMBEDDING_GLOBAL_RANKS = position_embedding_ranks
运行结果：
[0, 4, 8, 12]
[0, 12]
[0]
[1, 5, 9, 13]
[1, 13]
[1]
[2, 6, 10, 14]
[2, 14]
[2]
[3, 7, 11, 15]
[3, 15]
[3]
参考：
https://zhuanlan.zhihu.com/p/470279673
模型分布式环境初始化：
以两台分别有8个GPU服务器为例，训练具有12层的transformer layers，
标签：pipeline,group,ranks,概览,并行,groups,model,parallel,Yuan2.0 From： https://www.cnblogs.com/wangzhilun/p/17975268