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

Abstract

主要工作内容

方法便捷性亮点

在模型参数扩展上的改进

精度效果

Introduction

扩展性

contributions

  1. 实现了tensor-model-parallel
    • We implement a simple and efficient model parallel approach by making only a few targeted modifications to an existing PyTorch transformer implementation.
    • 我们通过对现有的PyTorch转换器实现进行一些有针对性的修改，实现了一种简单而高效的模型并行方法。
  2. 扩展效率好
    • We perform an in-depth empirical analysis of ourmodel and data parallel technique and demonstrate up to 76% scaling efficiency using 512 GPUs.
    • 我们对我们的模型和数据并行技术进行了深入的实证分析，并使用了512个gpu显示了高达76%的扩展效率。
  3. LN的位置重要
    • We show that careful attention to the placement of layer normalization in BERT-like models is critical to achieving increased accuracies as the model grows.
    • 我们表明，在类似bert的模型中，LN的位置对于随着模型的增长而提高精度是至关重要的。
  4. 参数量增加可以使精度进一步提高
    • We demonstrate that scaling the model size results in improved accuracies for both GPT-2 (studied up to 8.3 billion parameters) and BERT (studied up to 3.9B parameters) models.
  5. SOTA
    • We showcase that our models achieve state of the art results on test sets: perplexity on WikiText103 (10.8 ppl), accuracy on LAMBADA (66.5%), and accuracy on RACE (90.9%).
  6. 开源
    • We open source our code along with the training and evaluation pipelines at https://github.com/NVIDIA/Megatron-LM

Background and Challenges

  1. Neural Language Model Pretraining
  2. Transformer Language Models and Multi-Head Attention
  3. Data and Model Parallelism in Deep Learning

Model Parallel Transformers

  1. 在MLP中常见的一个操作就是mm+gelu
  2. 一种方法是沿着其行(row)将weight矩阵A分割，并沿着其列(columns)输入X分割，来实现tensor-model-parallel，从下图中我们可以看出，在该方法下，需要通过同步来保障语义对等
  3. 另一种方法是沿着它的列(columns)分割A，这种方法的好处是保障了各自在独立运行时的语义对等，对于一个module的正向输出在各卡上的同步，就需要2种所描述的方法来保障了
  4. 使用3的方法我们可以看到，使用了同一个X，这样在反向的时候需要梯度同步来保障module的反向输出在各卡上的同步，在PyTorch中，我们可以轻松实现这个功能 
    class f(torch.autograd.Function):
def forward(ctx, x):
 return x
def backward(ctx, gradient):
 all_reduce(gradient)
 return gradient
  5. 图示，对于一个MLP，进来先进行方法3，然后使用方法2，实现MLP的整体的tensor-model-parallel且语义和原始MLP对等；该方法同样可以扩展到Self-Attention模块
  6. 仔细想了下为什么是上述图片这样的结构，主要原因还是在于residual path的add需要语义上的一致，所以对于每个类MLP的结构而言，一输入就走列并行，使效率最大化；输出时走行并行，使结果语义对齐，从而兼顾了效率和语义。transformer类模型中类MLP结构一般如下所示
  7. 下一步想到了另一个参数量很大的地方，embedding层，虽然torch将embedding视为将input作为index，在weight上做index_select操作，实际上embedding也可以被视为，对input进行one_hot然后和weight进行mm的操作，这样想来，将embedding层想做linear层使用模型并行操作就顺多了。这部分在代码上比较清楚，就直接上代码了 
    class VocabParallelEmbedding():
def __init__(self, num_embeddings, embedding_dim, init_method=init.xavier_normal_):
   super(VocabParallelEmbedding, self).__init__()
   ...
   # 通过获取当前tensor_model_parallel的rank来确定当前卡要embedding的category id，初始化权重
   self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
          self.num_embeddings, get_tensor_model_parallel_rank(),
          self.tensor_model_parallel_size)
   self.num_embeddings_per_partition = self.vocab_end_index - self.vocab_start_index
   self.weight = Parameter(torch.empty(
          self.num_embeddings_per_partition, self.embedding_dim,
          dtype=args.params_dtype))
   ...

def forward(self, input_):
   if self.tensor_model_parallel_size > 1:
      # 把当前卡上不要的category idx mask掉
      input_mask = (input_ < self.vocab_start_index) | \
                   (input_ >= self.vocab_end_index)
      # Mask the input.
      masked_input = input_.clone() - self.vocab_start_index
      masked_input[input_mask] = 0
   else:
     masked_input = input_
   # Get the embeddings.
   output_parallel = F.embedding(masked_input, self.weight,
                             self.padding_idx, self.max_norm,
                             self.norm_type, self.scale_grad_by_freq,
                             self.sparse)
   # 把mask掉的category idx对应的embedding处理成0.
   if self.tensor_model_parallel_size > 1:
     output_parallel[input_mask, :] = 0.0
   # Reduce across all the model parallel GPUs.
   output = reduce_from_tensor_model_parallel_region(output_parallel)
   return output
  8. 说明完主要内容后，作者提到了，dropout, layer normalization, or residual connections在每张卡上是独立的，具体来说就是每张卡上都各自保留了LN的参数，允许各自优化自己的部分。这里每张卡指的是tensor-model-parrallel-group中的每张卡。

Setup

  1. Training Dataset
  2. Training Optimization and Hyperparamet

Experiments

  1. Scaling Analysis
    1. MODEL AND DATA PARALLELISM
  2. Language Modeling Results Using GPT-2
  3. Bi-directional Transformer Results Using BERT
    1. 值得一提的是，由于原始的bert的结构中，residual模块不包含LN，如果参数量加大会导致无法收敛的情况，这里Megatron把LN放进了residual，避免了无法收敛的情况

Conclusion and Future Work