Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

URL

• paper: https://arxiv.org/pdf/1909.08053
• code: https://github.com/NVIDIA/Megatron-LM

TL;DR

• 本文是英伟达 Megatron-LM 大模型分布式并行训练框架的一篇简介，大概阐述了 Megatron-LM 的机制和效果，对于实现细节着墨不多。
• 从代码实现角度看，Megatron-LM 是对 PyTorch 进行了二次封装。
• 基本思想是 Tensor Parallelism，拆分一层的权重到不同 GPU 上并行运算。

Algorithm

Tensor Parallelism

• 重点是这里的 f 和 g 函数，下面是其实现：

class f(torch.autograd.Function):
  def forward(ctx, x):
    return x
  def backward(ctx, gradient):
    all_reduce(gradient)
    return gradient
  
class g(torch.autograd.Function):
  def forward(ctx, x):
    all_reduce(x)
    return x
  def backward(ctx, gradient):
    return gradient

• 在本例子中 All-reduce 实际上是 ReduceSum，实际上聚合运算（包括 forward and backward）需要 GPU 间同步，存在通信代价
• Self-Attention 可以 tensor parallelism 的一个重要原因是 softmax 是逐行做的而不是全局做的
  
• 由图 3 可以得出结论，一个 transformer layer 只有 4 个同步算子（两个 LayerNorm 和两个 Dropout），其他算子都可以被分解为并行算子

Thought

• 给我的感觉有点类似 MLC 或者机器学习编译中的 Tiling（只是升级为训练及推理全流程 tiling）
• tensor parallelism 这种策略很考验对计算流程的抽象，除非有大量人力把框架搭建好且把所有算子实现写好，否则会很难用起来（当然对于 Nvidia 这些都不是事），突然想起了 Neuwizard，一声叹息…