模型并行-Megatron-LM

1.原理

  针对transformer模型的模型并行方法,利用矩阵乘法的特性,将每个层切开,放入不同GPU中。

示例图片
对于注意力层,是按注意力头进行切分; 对于MLP层,是按矩阵列进行切分。整个执行流程中,每个层的需要进行数据通信,因为对于矩阵切分后的乘法要得到最终的结果需要每个分块的结果进行矩阵加。

2.缺点

  使用多头注意力的transformer的头数、MLP隐藏层大小需要能够被GPU数整除。每计算一个层就需要进行all reduce,导致不同层之间的计算是同步的,且有较大的通信开销。

posted @ 2024-12-10 12:52  博客员626  阅读(13)  评论(0编辑  收藏  举报