大语言模型框架-Megatron-LM 源码分析

这个 linear 前向是标准的 torch matmul，除非 sequence 有并行设置才会进行一定的 allgather 通信。

之前异步都被 RowLinear 配置 false，所以 BP 核心是 fused kernel。且可以使用低精度 16bit 的内核。

gradient_accumulation_fusion (bool): If true, fuses weight gradient accumulation to GEMMs. Requires the custom CUDA extension fused_weight_gradient_mlp_cuda module. To use gradient_accumulation_fusion you must install APEX with --cpp_ext and --cuda_ext. For example: "pip install --global-option="--cpp_ext" --global-option="--cuda_ext" ". Note that the extension requires CUDA>=11. Otherwise, you must turn off gradient accumulation fusion. Defaults to False.

如果不选 fuse kernel 则执行执行矩阵乘完成 BP 反向传播计算。

有意思的是 Apex 库也用的这个 linear 实现，库之前二次开发逐渐成为当前常用方式

这个 fusekernel 最终还是调用的 cublas 的 gemm

FP16 使用的 BF16，且也是 cublas gemm，并可以利用 tensor core 的加速。也有版本可以选择 at::Half* A,类型的 bit 16 类型内核。

下面是上面的 kernel 优化目的 Gradient accumulation fusion 的介绍。

所以 C 矩阵在 kernel 的输出是 32 bit，输入是两个 16bit 矩阵保证了输出累计的数据精度。

接下来我们再看列切分配置，其类似行切分。好处是省去下一步的 GEMM 之前的 allreduce 通信，所以 attention 和 mlp 的第一层 gemm megatron 选择列切，之后再行切下一个阶段 gemm。

其中 gather_output 由模型设计者决定，进而和后面的层进行配合，看输出是当前层是否需要聚合还是不需要聚合收集。这点是比较有意思，当前属于人工硬编码选择配置好策略。

初始化决定列切切分维度，这个 world size 通过 tensor parallel 的输入 shell 进行配置。

一般配置是 8 路，也就是张量并行在一个 8 卡的 server，如果用户是 4 卡 server 则配置 4 即可。

和 Row 并行的区别是用户配置是否进行一次 all gather output

allgather 原因是列切完，我们产生了两块局部结果。如下文所示是两个 partition 1 和 2。

allgather [图来源 PyTorch 官网]

2 流水并行

其核心代码在 p2p_communication 中实现有 4 个重要参数。先沟通改下需要传递的 tensor shape，"""Communicate tensor shapes between stages. Used to communicate tensor shapes before the actual tensor communication happens. This is required when the sequence lengths across micro batches are not uniform.

下面是真正的传递核心，使用的 P2POp，isend 传递是异步传递，同时通过函数去确定相邻的 rank，这样写这个函数就不用管拓扑了。

然后将这些操作符做批量通信发送，一批完成和自己上游和下游的 send recv 异步通信

调用在 _communicate 进行

当前方式是相当于有了 microbatch 的发送方式，也就是既有下个 microbatch 的前向，也有当前 batch 的反向，一批次做异步通信。适合整体都已经运行起来了，已经不是第一个 batch 的场景。

这是在 backward 中调用的上面的 API

流水并行的切分靠的是这个函数获取自己的这个 rank 到底是哪个 model chunk，相当于静态切分好自己是属于哪块，类似张量并行。当前并行策略都是可以考虑这种分配方式，静态编译好 partition 和 rank 映射关系，启动后获取这个关系决定自己的通信方式。

前向传播通过以下函数进行 p2p 通信。

底层还是通过 _communicate 调用实现类似上面的 bp 过程。当前相当于只需要给下个 rank send 所以 tensor_send_prev 为空。如果多个 microbatch 配置是第一个 batch 前向，或者没配置多个 microbatch 场景使用。

3 数据并行

数据并行一般有几点优化：将梯度通信和 BP 计算进行 overlap，同时可以使用低精度做梯度聚合，将 grad 组成成小桶聚合。

图来自 PyTorch，桶聚合相当于 batching 通信张量

通过 hook 注册梯度更新事件

Overlap 的核心是当 bucket 的 paramer 累计到都有 gradient，触发 allreduce 同步。BP 计算该做自己的继续做，产生好的 gradient，这部分只要 ready 就同时触发 allreduce 与进行的 BP 就无关了，但是没产生的 gradient 还不能进行计算。

grad_buffer 中是核心逻辑。

allreduce，图片来源 PyTorch

核心逻辑通过 GradBuffer 进行聚合成连续 buffer，再拆解。

本质是 allreduce 并可以选择是否是异步。

在 finish 中不同等待异步 allreduce

此处聚合 wait 所有的 handler 同步通信。

config.finalize_model_grads_func = finalize_model_grads 在没有流水并行下，执行刚才的同步通信等待。