[转]详解MegatronLM Tensor模型并行训练(Tensor Parallel)

发布网友发布时间：2024-10-23 03:22

共1个回答

热心网友时间：2024-10-24 00:05

详解MegatronLM Tensor模型并行训练(Tensor Parallel)的主要内容如下：

背景介绍：Megatron-LM于2020年发布，专门针对十亿参数级别的语言模型进行训练，如具有38亿参数的类GPT-2的transformer模型和39亿参数的BERT模型。模型并行训练有层间并行(inter-layer)和层内并行(intra-layer)两种方式，分别对应模型的竖切和横切。

Tensor并行计算方法介绍：Tensor并行通过合理方式将输入矩阵和参数矩阵分块，分别进行计算。在参数weight矩阵上，可以通过横切或纵切进行模型并行。横切时，输入矩阵按行切分，参数矩阵按列切分；纵切时，输入矩阵按列切分，参数矩阵按行切分。在反向传播时，需要通过all-gather和all-rece进行梯度聚合与合并。

GPT Transformer应用：GPT Transformer由Attention模块和MLP模块组成。在MLP模块中，第一个MLP层采用按列切分，第二个MLP层则采用按行切分。在Attention模块中，对Tensor并行切分方式与MLP模块类似，首先按列切分，然后对最终的weight矩阵按行切分。

输出Embedding层应用：在GPT计算中，Embedding层的权重按列切分，通过all-gather通信得到最终结果。为了降低通信量，可以将GEMM和cross entropy loss进行融合，通信量从【batch-size x sequence-length x vocabulary-size】降低到【batch-size x sequence-length】。

总结：Tensor并行训练是一种有效提升大模型训练效率的方法，通过合理分块和通信策略，能够实现模型并行，加速训练过程并优化计算资源的使用。