大模型训练框架Megatron原理系列转载_megatron框架

原文：[源码解析] 模型并行分布式训练Megatron (1) --- 论文 & 基础

https://www.cnblogs.com/rossiXYZ/p/15840803.html

目录

• [源码解析] 模型并行分布式训练Megatron (1) --- 论文 & 基础
  • 0x00 摘要
  • 0x01 Introduction
    • 1.1 问题
    • 1.2 数据并行
    • 1.3 模型并行
      • 1.3.1 通信
      • 1.3.2 张量并行
      • 1.3.3 流水线并行
    • 1.4 技术组合
    • 1.5 指导原则
  • 0x02 张量模型并行（Tensor Model Parallelism）
    • 2.1 原理
      • 2.1.1 行并行（Row Parallelism）
      • 2.1.2 列并行（Column Parallelism）
    • 2.2 Model Parallel Transformers
      • 2.2.1 Transformer
      • 2.2.2 切分 Transformer
      • 2.2.3 切分MLP
      • 2.2.4 切分self attention
      • 2.2.5 通信
      • 2.2.6 小结
  • 0x03 并行配置
    • 3.1 符号说明
    • 3.2 Tensor and Pipeline Model Parallelism
    • 3.3 Data and Model Parallelism
      • 3.3.1 Pipeline Model Parallelism.
      • 3.3.2 Data and Tensor Model Parallelism.
    • 3.4 Microbatch Size
    • 3.5 对比
      • 3.5.1 Tensor versus Pipeline Parallelism.
      • 3.5.2 Pipeline versus Data Parallelism.
      • 3.5.3 Tensor versus Data Parallelism.
  • 0x04 结论
  • 0xFF 参考

0x00 摘要

NVIDIA Megatron 是一个基于 PyTorch 的分布式训练框架，用来训练超大Transformer语言模型，其通过综合应用了数据并行，Tensor并行和Pipeline并行来复现 GPT3，值得我们深入分析其背后机理。

本系列大概有6～7篇文章，通过论文和源码和大家一起学习研究。

本文把 Megatron 的两篇论文/一篇官方PPT 选取部分内容，糅合在一起进行翻译分析，希望大家可以通过本文对 Megatron 思路有一个基本了解。

0x01 Introduction

1.1 问题

在NLP领域之中，大模型可以带来更精准强大的语义理解和推理能力，所以随着规模计算的普及和数据集的增大，使得模型的参数数量也以指数级的速度增长。训练这样大的模型非常具有挑战性，具体原因如下：

• (a) 对显存的挑战。即使是最大的GPU的主内存也不可能适合这些模型的参数，比如一个175B的GPT-3模型需要（175B * 4bytes）就是700GB模型参数空间，从而梯度也是700G，优化器状态是1400G，一共2.8TB。
• (b) 对计算的挑战。即使我们能够把模型放进单个GPU中（例如，通过在主机和设备内存之间交换参数），但是其所需的大量计算操作会导致漫长训练时间（例如，使用单个V100 NVIDIA GPU来训练1750亿个参数的GPT-3需要大约288年）。如何计算可以参见 2104.04473的附录 FLOATING-POINT OPERATIONS。
• (c) 对计算的挑战。不同并行策略对应的通信模式和通信量不同。
  • 数据并行：通信发生在后向传播的梯度规约all-reduce操作，通信量是每个GPU之上模型的大小。
  • 模型并行：我们在下面会详述。

这就需要采用并行化来加速。使用硬件加速器来横向扩展（scale out）深度神经网络训练主要有两种模式：数据并行，模型并行。

1.2 数据并行

数据并行模式会在每个worker之上复制一份模型，这样每个worker都有一个完整模型的副本。输入数据集是分片的，一个训练的小批量数据将在多个worker之间分割；worker定期汇总它们的梯度，以确保所有worker看到一个一致的权重版本。对于无法放进单个worker的大型模型，人们可以在模型之中较小的分片上使用数据并行。

数据并行扩展通常效果很好，但有两个限制：

• a）超过某一个点之后，每个GPU的batch size变得太小，这降低了GPU的利用率，增加了通信成本；
• b）可使用的最大设备数就是batch size，着限制了可用于训练的加速器数量。

1.3 模型并行

人们会使用一些内存管理技术，如激活检查点（activation checkpointing）来克服数据并行的这种限制，也会使用模型并行来对模型进行分区来解决这两个挑战，使得权重及其关联的优化器状态不需要同时驻留在处理器上。

模型并行模式会让一个模型的内存和计算分布在多个worker之间，以此来解决一个模型在一张卡上无法容纳的问题，其解决方法是把模型放到多个设备之上。

模型并行分为两种：流水线并行和张量并行，就是把模型切分的方式。

• 流水线并行（pipeline model parallel）是把模型不同的层放到不同设备之上，比如前面几层放到一个设备之上，中间几层放到另外一个设备上，最后几层放到第三个设备之上。
• 张量并行则是层内分割，把某一个层做切分，放置到不同设备之上，也可以理解为把矩阵运算分配到不同的设备之上，比如把某个矩阵乘法切分成为多个矩阵乘法放到不同设备之上。

具体如下图，上面是层间并行（流水线并行），纵向切一刀，前面三层给第一个GPU，后面三层给第二个GPU。下面是层内并行（tensor并行），横向切一刀，每个张量分成两块，分到不同GPU之上。

或者从另一个角度看看，两种切分同时存在，是正交和互补的（orthogonal and complimentary）。

图来自：GTC 2020: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

1.3.1 通信

我们接下来看看模型并行的通信状况。

• 张量并行：通信发生在每层的前向传播和后向传播过程之中，通信类型是all-reduce，不但单次通信数据量大，并且通信频繁。
• 流水线并行：通信在流水线阶段相邻的切分点之上，通信类型是P2P通信，单词通信数据量较少但是比较频繁，而且因为流水线的特点，会产生GPU空闲时间，这里称为流水线气泡（Bubble）。

比如下图之中，上方是原始流水线，下面是模型并行，中间给出了 Bubble 位置。

因为张量并行一般都在同一个机器之上，所以通过 NVLink 来进行加速，对于流水线并行，一般通过 Infiniband 交换机进行连接。

图来自 Megatron 论文。

1.3.2 张量并行

有些工作在张量（层内）模型并行化（ tensor (intra-layer) model parallelism）做出了一些尝试，即每个transformer 层内的矩阵乘法被分割到多个GPU上，虽然这种方法在NVIDIA DGX A100服务器（有8个80GB-A100 GPU）上对规模不超过200亿个参数的模型效果很好，但对更大的模型就会出现问题。因为较大的模型需要在多个多GPU服务器上分割，这导致了两个问题。

• (a) 张量并行所需的all-reduce通信需要通过服务器间的链接，这比多GPU服务器内的高带宽NVLink要慢；
• (b) 高度的模型并行会产生很多小矩阵乘法（GEMMs），这可能会降低GPU的利用率。

1.3.3 流水线并行

流水线模型并行化是另一项支持大型模型训练的技术。在流水线并行之中，一个模型的各层会在多个GPU上做切分。一个批次（batch）被分割成较小的微批（microbatches），并在这些微批上进行流水线式执行。

通过流水线并行，一个模型的层被分散到多个设备上。当用于具有相同transformer块重复的模型时，每个设备可以被分配相同数量的transformer层。Megatron不考虑更多的非对称模型架构，在这种架构下，层的分配到流水线阶段是比较困难的。在流水线模型并行中，训练会在一个设备上执行一组操作，然后将输出传递到流水线中下一个设备，下一个设备将执行另一组不同操作。

原生（naive）流水线会有这样的问题：一个输入在后向传递中看到的权重更新并不是其前向传递中所对应的。所以，流水线方案需要确保输入在前向和后向传播中看到一致的权重版本，以实现明确的同步权重更新语义。

模型的层可以用各种方式分配给worker，并且对于输入的前向计算和后向计算使用不同的schedule。层的分配策略和调度策略导致了不同的性能权衡。无论哪种调度策略，为了保持严格的优化器语义，优化器操作步骤（step）需要跨设备同步，这样，在每个批次结束时需要进行流水线刷新来完成微批执行操作（同时没有新的微批被注入）。Megatron引入了定期流水线刷新。

在每个批次的开始和结束时，设备是空闲的。我们把这个空闲时间称为流水线bubble，并希望它尽可能的小。根据注入流水线的微批数量，多达50%的时间可能被用于刷新流水线。微批数量与流水线深度（size）的比例越大，流水线刷新所花费的时间就越少。因此，为了实现高效率，通常需要较大的batch size。

一些方法将参数服务器与流水线并行使用。然而，这些都存在不一致的问题。TensorFlow的GPipe框架通过使用同步梯度下降克服了这种不一致性问题。然而，这种方法需要额外的逻辑来处理这些通信和计算操作流水线，并且会遇到降低效率的流水线气泡，或者对优化器本身的更改会影响准确性。

某些异步和bounded-staleness方法，如PipeMare、PipeDream和PipeDream-2BW完全取消了刷新，但这样会放松了权重更新语义。Megatron会在未来的工作中考虑这些方案。

1.4 技术组合

用户可以使用各种技术来训练他们的大型模型，每种技术都有不同的权衡。此外，这些技术也可以被结合起来使用。然而，结合这些技术会导致复杂的相互作用，对于系统拓扑是个极大的挑战，不仅要对模型做合理切割（依据算法特点），还需要做软硬件一体的系统架构设计，需要仔细推理以获得良好的性能。因此以下问题就特别重要：

应该如何组合并行技术，以便在保留严格的优化器语义的同时，在给定的batch size下最大限度地提高大型模型的训练吞吐量？

Megatron-LM 开发人员展示了一个如何结合流水线、张量和数据并行，名为PTD-P的技术，这项技术将以良好的计算性能（峰值设备吞吐量的52%）在1000个GPU上训练大型语言模型。PTD-P利用跨多GPU服务器的流水线并行、多GPU服务器内的张量并行和数据并行的组合，在同一服务器和跨服务器的GPU之间具有高带宽链接的优化集群环境中训练具有一万亿个参数的模型，并具有优雅的扩展性。

要实现这种规模化的吞吐量，需要在多个方面进行创新和精心设计：

• 高效的核（kernel）实现，这使大部分计算操作是计算绑定（compute-bound）而不是内存绑定（memory-bound。
• 在设备上对计算图进行智能分割，以减少通过网络发送的字节数，同时也限制设备的空闲时间。
• 实施特定领域的通信优化和使用高速硬件（比如最先进的GPU，并且同一服务器内和不同服务器GPU之间使用高带宽链接）。

1.5 指导原则

Megatron 开发者研究了各种组合之间如何影响吞吐量，基于这些研究得出来分布式训练的一些指导原则：

• 不同的并行模式以复杂的方式互相作用：并行化策略影响通信量、核的计算效率，以及worker因流水线刷新（流水线气泡）而等待的空闲时间。例如，张量模型并行在多GPU服务器中是有效的，但大模型必须采用流水线模型并行。

• 用于流水线并行的schdule对通信量、流水线气泡大小和用于存储激活的内存都有影响。Megatron 提出了一个新的交错schdule，与以前提出的schdule相比，它可以在稍微提高内存占用的基础上提高多达10%的吞吐量。

• 超参数的值，如microbatch size，对memory footprint、在worker上执行的核效果和流水线bubble大小有影响。

• 分布式训练是通信密集型的。使用较慢的节点间连接或更多的通信密集型分区会阻碍性能。

0x02 张量模型并行（Tensor Model Parallelism）

2.1 原理

我们用 GEMM 来看看如何进行模型并行，这里要进行的是 XA = Y，对于模型来说，X 是输入，A是权重，Y是输出。从数学原理上来看，对于linear层就是把矩阵分块进行计算，然后把结果合并，对于非linear层则不做额外设计。

2.1.1 行并行（Row Parallelism）

我们先看看Row Parallelism，就是把 A 按照行分割成两部分。为了保证运算，同时我们也把 X 按照列来分割为两部分，这里 声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/2023面试高手/article/detail/344127