一文了解预训练相关加速技巧

Posted 2021-12-28 一夜了

1. 简介

主要介绍deepspeed，Megatron，OneFlow等开源训练工具。

2. deepspeed介绍

deepspeed是微软开源的模型训练框架. 主要有流水线并行, 以及ZeRO内存优化技术等.

2.1 3D 并行性：扩展到万亿参数模型

数据并行训练、模型并行训练和流水线并行训练

2.1.1 数据并行

最常见的 DNN 并行化训练方法是数据并行化，这种方法把输入数据分散到各个 workers中（每个worker拥有全部模型）运行。不幸的是，尽管在加速数据并行的性能优化方面取得了一些进展，但若要在云基础设施上训练就会产生很大的通信开销。而且，随着GPU 计算速度的飞快提升，所有模型训练的耗时瓶颈会更进一步转向通信环节。

2.1.2 模型并行

模型并行化是另一种并行化训练形式，这种方法是把算子分散到各个 worker 上进行计算，这通常用于训练大型 DNN 模型。本文中，模型并行指的就是把模型的不同layer放在不同的机器上，不涉及将同一个layer切分到不同机器上的场景。
下图是模型并行化，显示了一个计算时间线，该示例有四台机器和一个管道。

• 在正向阶段，每个阶段对本阶段中的层的minibatch执行正向传递，并将结果发送到下一阶段。输出级在完成前向传递后，计算minibatch的损失。
• 在后向阶段，每个阶段形成后向通道，逐一将损失传播到前一阶段。
  

2.1.3 流水线并行

PipeDream combines traditional data parallelism with model parallelism enhanced with pipelining

a. 主要思想：

将模型分成多个stage，每个stage包括连续的多个layers。每个stage由独立的GPU执行forward和backward pass。

与数据并行相比，流水线并行降低了communication以及Pipelining overlaps computation and communication。

b. 1F1B

question: Given a model and a set of machines, PipeDream’s first challenge is to automatically partition layers of the model across available machines so as to minimize overall training time.

c. PipeDream’s Partitioning Algorithm：

假定模型有N layer以及M个机器。上面的问题就可以转化为动态规划问题。

• 首先会通过1000次的mini-batch，统计出每层的计算时间T
• PipeDream估计的通信时间是需要传输的数据量除以通信链路上的网络带宽
  
  Ｔ(i -> j, m)表示单个stage中，复制了ｍ份，从layer i到ｊ的时间。Ｔ表示单个layer前向和后向梯度运算的时间。Ｗ表示ｌ层想其他(m-1)个机器权重传输的时间。
  
  

其中$2\ast {Ｃ}_i$表示the time taken to communicate the activations and gradients between layers i and i +1

d. Weight Stashing：

Weight Stashing maintains multiple versions of the weights, one for each active minibatch. When performing the forward pass, each stage processes a minibatch using the latest version of weights available. After completing the forward pass, PipeDream stores the weights used as part of the intermediate state for that minibatch. When performing the minibatch’s backward pass, the same version of the weights is used to compute the weight gradient.

主要是为了维护多个版本的权重结果，权重存储确保在一个stage内，在给定的minibatch内，相同的模型参数版本用于前向和反向传递。
例如，在流水线并行图 中，minibatch 5 使用来自机器 1 上的minibatch 1 和机器 2 上的mibibatch 2。 但是不能保证在同一个mini-batch，在不同的stage中，权重参数保持一致。

e. Vertical Sync：

Vertical Sync eliminates the potential inconsistency across stages.

主要是为了解决同一个mini-batch在不同stage中，权重不一致的问题。
如下例所示:初始没做任何优化的参数迭代

经过weight stashing后的参数迭代：

经过vertical Sync后的参数迭代：

2.2 ZeRo and ZeRO-Offload （4） - 内存优化

2.2.1 GPU内存怎么被消耗？

a 1.5B parameter GPT-2 model requires 3GB of memory for its weights(or parameters) in 16-bit precision, yet, it cannot be trained on a single GPU with 32GB memory using Tensorflow or PyTorch.

大部分原因是GPU显存被模型训练中的model state消耗。例如保存optimizer states,gradients, and parameters信息的tensor，以及一些activations, temporary buffers and fragmented memory 。
举例：

Adam requires storing two optimizer states, i) the time averaged momentum and ii) variance of the gradients to compute the updates.Therefore, to train a model with ADAM, there has to be enough memory to hold a copy of both the momentum and variance of the gradients. In addition, there needs to be enough memory to store the gradients and the weights themselves.

总结： 假定一个模型参数为A，使用adam优化器，fp16混合精度训练，则parameters为2A，gradients为2A，optimizer states with fp32：parameters（4A），momentum(4A)，variance（4A）

通过优化后，数据并行处理后，理论上可以训练模型的大小：

实际测量的模型可以训练的大小：

2.2.2 ZeRO介绍

ZeRO leverages the aggregate computation and memory resources of data parallelism to reduce the memory and compute requirements of each device (GPU) used for model training. ZeRO reduces the memory consumption of each GPU by partitioning the various model training states (weights, gradients, and optimizer states) across the available devices (GPUs and CPUs) in the distributed training hardware. Concretely, ZeRO is being implemented as incremental stages of
optimizations, where optimizations in earlier stages are available in the later stages.

• ZeRO（The Zero Redundancy
  Optimizer）是一种用于大规模分布式深度学习的新型内存优化技术。ZeRO可以在当前的GPU集群上训练具有1000亿个参数的深度学习模型，其吞吐量是当前最佳系统的三到五倍。

• ZeRO通过跨数据并行进程划分模型状态（参数，梯度和优化器状态），而不是复制它们，从而消除了数据并行进程之间的内存冗余。它在训练期间使用动态通信方式，以在分布式设备之间共享必要的状态，以保持数据粒度的计算粒度和通信量。

• 另外, ZeRO是ZeRO-DP和ZeRO-R 的集合：

下面分别进行介绍:

ZeRO-DP：

• 优化器状态分区（Pos）– 内存减少4倍，通信量与数据并行性相同
• 添加渐变分区（Pos + g）– 内存减少8倍，与数据并行性相同的通信量
• 添加参数分区（Pos + g + p）– 内存减少与数据并行度Nd成线性关系。例如，拆分64个GPU（Nd = 64）将减少64倍的内存。通信量略有增加50％。
  ZeRO-DP具体优化过程：
  (1)
  (2)每一块GPU用来负责模型的一部分，梯度累加器经过4路并行进行累加
  
  (3) Mo是模型的最初阶段，模型参数从GPU0广播到GPU1/2/3
  
  (4) 各自的Mo模型使用各自的数据进行推理，仅仅记录激活值
  
  （5）在M0推理过后，GPU1/2/3可以删除M0部分的参数
  
  （6）在BP的过程中，M3被暂时保存，于此同时计算各自的梯度
  
  （7）将M3的梯度汇总到GPU3的M3网络
  
  （8）目前GPU0/1/2的M3参数以及激活层进行释放
  
  （9）依次类推到M0上的模型参数和梯度计算
  
  （10）所有temp的激活、参数都可以删除掉，只有模型的梯度
  （11）优化策略进行了并行的计算
  
  （12）优化计算的FP32被直接成FP16
  
  （13）将优化后的FP16参数更新至模型参数层
  
  ZeRO-R：
• 1）ZeRO-R optimizes activation memory by identifying and removing
  activation replication in existing MP approaches through activation
  partitioning. It also offloads activations to CPU when appropriate
  （ZeRO-R通过激活分区来识别和删除现有方法中的激活复制，从而优化激活内存）
• 2）ZeRO-R defines appropriate size for temporary buffers to strike for
  a balance of memory and computation
  efficiency.（ZeRO-R为临时缓冲区定义了适当的大小，以达到内存和计算效率的平衡）
• 3）ZeRO-R proactively manages memory based on the different lifetime
  of tensors, preventing memory
  fragmentation（ZeRO-R根据张量的不同生命周期主动管理内存，防止内存碎片）

2.2.3 ZeRO-offload介绍

The CPU memory bandwidth is at least an order of magnitude faster than the PCI-E bandwidth between CPU and GPU, while the GPU memory is another order of magnitude faster than even the CPU memory.

• ZeRO-offload的主要思想是利用cpu作为模型计算的工具之一，由于parameters，gradients，and the optimizer states主要是模型训练内存的主要瓶颈，所以可以将一部分model state放到cpu，有cpu进行计算，从而通过增加通信，从而降低内存的瓶颈。
• 模型训练的每次迭代的计算复杂度主要为 O(MB)，其中M表示model size，B表示batch size。所以为了避免cpu计算成为模型训练的瓶颈，需要将计算复杂度小于O(MB)的model state放到cpu上。
• GPU内存传输速度 > CPU内存传输速度 > CPU-GPU之间PCI-E传输速度。
• 将FWD+BWD放到GPU，norm calculations, weight updates放到CPU。

2.2.4 ZeRO-Infinity

ZeRO-infinity能够使用GPU，CPU和非易失性内存Express或简称NVMe

Bandwidth-Centric Partitioning

with the partitioned parameter and allgather based approach in ZeRO-Infinity, all PCIe links are active in parallel, each bringing in 1/𝑑𝑝𝑡ℎ portion of the parameter where 𝑑𝑝 is the data parallel degree. As a result, the effective communication bandwidth between NVMe or CPU to the GPU, increases linearly with the 𝑑𝑝 degree.

Overlap Centric Design

ZeRO-Infinity has an overlap engine that not only overlaps GPU-GPU communication with GPU computation, but also overlaps the NVMe to CPU, and CPU to GPU communication, all at the same time.

Infinity Offload Engine

DeepNVMe： a powerful C++ NVMe read/write library in the infinity offload engine that supports bulk read/write requests for asynchronous completion, and explicit synchronization requests to flush ongoing read/writes.Pinned memory management layer
Pinned memory management layer：To ensure high performance tensor reads (or writes) from (to) NVMe/CPU storage, the source (or destination) tensors must reside in pinned memory buffers.

3. Megatron

NVIDIA 基于 PyTorch 开发了 Megatron，本质上是一个专用于 GPT 的模型库，所有的代码都是 Python 脚本，NVIDIA 为 GPT 专门定制了分布式训练所需的算子、 流水并行调度器、模型并行所需的通信原语等功能。
它不仅支持传统分布式训练的数据并行，也支持模型并行，包括 Tensor 并行和 Pipeline 并行两种模型并行方式
具体的代码：https://github.com/nvidia/megatron-lm

tensor模型并行：

4. OneFlow

内部宣称是最快的深度学习框架，除了上面介绍的优化之外，还有：

4.1 优化方法：

1. 减少全局内存的访问

1.1. GPU的内存层次

• Global memory 可以被所有thread访问。位置是在device memory上，并配置了L1、L2缓存。Device
  memory访问延迟（access latency）高，访问带宽（bandwidth）低。
• Local memory 是每个thread内部私有的，但位置也是在device memory上，所以延迟差不多一样高，带宽差不多一样低。
• Shared memory 是每个block的所有thread共享同一块shared
  memory，对block内部的thread可见，当block执行结束时会被释放。由于shared
  memory存储位置是在片上内存（on-chip memory），所以不需要缓存，读取速度很快。

具体措施：减少global memory的访问可以明显提升GPU计算的性能。

• 将element-wise的Kernel跟上一个计算Kernel合并，通常可以减少global
  memory的访问需求，如将add合并到前面的conv_data_grad、bn、dropout、matmul等op上。
• 借助shared memory合并带有Reduce计算的Kernel
• 减少实际需要的访存大小，例如，以非常常见的relu为例，relu的后向op relu_grad在做计算时，消费前向的输入x或者输出y（一般是消费输出），但其实后向计算仅需要判断对应位置的元素是否大于0，所以可以将后向对y的消费替换为一个前向op产出的bitset，理论上可以省掉冗余访问y的操作，可以减少大约1/3的global memory访问。

2. 确保全局内存访问合并

3. 优化Kernel计算量

• 减少坐标变换的次数
• 根据Tensor的Shape优化计算量
• 尽量使用int32表示下标
• 表达式的等价变换

4. 延迟隐藏

4.2 缺点：

开源出来半年多，没有MoE框架，需要踩很多的坑。

5. PaddlePaddle

介绍：首先在单机内使用模型并行和分组参数切片组合的 2D 策略，这么选择的原因是这两个策略通信量较大，适合使用机器内的卡间通信；然后为了承载千亿规模模型，再叠加流水线并行策略，使用多台机器共同分担；最后为了做到高效，在外层又叠加了数据并行来增加并发数量，提升整体训练速度。
对比：百度提出了 “4D混合并行”，本质上是 Megatron 里的数据并行 + 模型并行 + 流水并行 + DeepSpeed 里的 ZeRO 优化

6. 其他的优化方法：

6.1 混合精度

6.2 梯度累加

Gradient Accumulation 就是把一个大 Batch 拆分成多个 micro-batch ， 每个 micro-batch 前后向计算后的梯度累加，在最后一个micro-batch累加结束后，统一更新模型。

6.3 后向重计算

Checkpointing 的核心思想 是在前向网络中标记少量的 Tensor （被 Checkpointing 的 Tensor ），前向计算就只会保留这些被标记的 Tensor， 其余的前向的 activation，会通过在反向传播中根据 Checkpointing 的 Tensor 临时重新计算一遍前向得到。这样就使得大量的 activation 不需要一直保存到后向计算，有效减少了大量 Tensor 的生命周期，使得内存复用效率大幅提升。

7. 其他的优化方法

8. 参考文献

1. https://juejin.cn/post/7002945085734141988
2. PipeDream: Fast and Efficient Pipeline Parallel DNN Training
3. PipeDream: Greneralized Pipeline Parallelism for DNN Training
4. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
5. ZeRO-Offload: Democratizing Billion-Scale Model Training
6. ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
7. https://mp.weixin.qq.com/s/ec20Z7kIp5Vyi1dAusXz1A
8. https://cloud.tencent.com/developer/article/1832354
9. https://mp.weixin.qq.com/s/cr-lYVvn1AQ7BN1VfzfuNg
10. Training Deep Nets with Sublinear Memory Cost
11. https://zhuanlan.zhihu.com/p/76638962
12. https://zhuanlan.zhihu.com/p/366906920
13. https://www.zhihu.com/question/498135156