Read paper: ZeRO-Infinity

ZeRO-Infinity

Abstract

处理的是GPU内存墙的问题,是一个新的异构系统技术,利用了GPU,CPU,NVMe的内存,在有限的资源去达到史无前例的模型规模,而无需进行代码重构.同时,它能达到优异的训练吞吐量和可扩展性,不受CPU或者NVMe带宽的限制

在现在的GPU集群中,能达到万亿乃至十万亿的参数量

在单个DGX-2节点上能微调万亿参数量级别的模型

超线性的可扩展性

Introduction

How to increase model scale

• 训练模型的系统技术进步
• 并行技术如流水线并行,张量并行,ZeRO等的使用

Why GPU memory is the bottleneck

DeepSpeed实现的3D并行能在800张V100的卡上扩展到万亿级别的参数量,通过充分利用集群中聚合的GPU内存;对于A100的卡而言,则需要320张才可以做到万亿参数量;而若是想要训练百万亿参数量,即使GPU的内存可以翻5倍(指80GB -> 400GB),也需要6K张卡才可以训练到.So GPU memory is the bottleneck

What does GPU Memory Wall effect data scientist

GPU内存墙也限制了Data Scientist对模型进行微调.由于模型量大,微调它们虽然比预训练的工作需要更少的GPU资源,但是也让人举步维艰.即使可以访问特定的计算资源,但是在计算节点上被分配的内存也有限,这又反过来使得想微调的模型大小受限.这样就使得没资源(大规模的GPU集群)的人,无法做大模型微调的工作

原文有一句: For example fine-tuning GPT-3 would require over 8 DGX-2 nodes(128 GPUs) with 3D parallelism to just fit the model for training, even though a single DGX-2 node (16-GPUs) has enough compute to fine-tune it in a reasonable time.

不知道作者是想表达fine-tune比pre-train更cheap;

还是说即使更cheap,对大多数研究人员也内存墙;

还是指有些平台需要fine-tune的资源等价于pre-train的资源,即使可以有更合理更节省的方案

3D parallelism lack of usability and flexibility

3D并行性需要对代码重构,施加DP,PP,TP.但是PP即流水线并行它对于一些具有复杂依赖的模型无能为力,因为不能拆分为负载均衡的stage

Unprecedented Model Scale

• infinity offload engine

GPU资源是有限,我们可以offload数据到CPU和NVMe上,以此支持训练大模型

• memory-centric tiling

优化GPU内存,以支持超大参数的单个层放入GPU内存中

Excellent Training Efficiency

• bandwidth-centric partitioning

引入一种新的数据分区策略以利用所有设备的聚合内存带宽

• communication overlap-centric design

应该是遮掩通信的开销

• optimization for NVMe access

优化NVMe访问,应该是对offload数据,以及回传数据进行优化

Ease of Use

无需重构代码(个人感觉是说pretrain的代码finetune无需调动),无需使用3D并行(根据contribution的说法来看,应该是不使用TP),对任意的模型架构均会自动执行训练时所需的所有通信和数据分区

Related Work

Memory Requirements

用于深度学习训练,本文主要集中于分析基于transformer架构的模型,因为全部超过1billion参数的SOTA的模型都是基于transformer架构的.

本文的分析是基于混合精度训练以及采用Adam优化器(采用此optimizer已经变成了基于transformer模型的一个标准,因此用它)

将讨论的内存分类:

• model state: [optimizer state,gradient,model parameter]
• residual state: 主要指的是激活(activation)内存

GPU的工作内存是指GPU必须的可用于支持训练的最小内存,假设model and residual states可以从GPU内存中被成功的卸载

Memory for Model States

• parameter和gradient是用FP16来存储
• optimizer states包含了momentum,variance,parameters和gradients,都是用FP32来存储

因此我们可得知一个参数需要20个字节的内存

用$hd$来代表hidden dimension,用$nl$表示transformer层的数目.考虑到在一个transformer块中几乎所有参数是来自块内的四个线性层:其中两个在多头注意力机制中①input -> QKV ([$hd$,$3hd$]);②$W^O$ ([$hd$,$hd$]);剩下两个在FFN中③[$hd$,$4hd$];④[$4hd$,$hd$].因此总参数量可以表示为

$$12 * nl *{hd}^2$$

Model States所需的总内存则乘于每个参数所需的20个字节,即为:

$$240*nl*{hd}^2$$

Memory for Residual States

Residual States主要包含的是activation memory,它的参数主要取决于模型的架构:批处理大小$bsz$,序列长度$seq$

本文采用activation checkpoint的方式来对activation所需的内存进行减少,因此存储的是激活检查点的数据

$ci$用于表示两个activation checkpoints间的transformer block的个数,因此$\frac{nl}{ci}$表示的是激活检查点的个数

激活检查点所需的总内存为:

$$2 * bsz * seq * hd * \frac{nl}{ci}$$

Model State Working Memory(MSWM)

模型状态工作内存是在所有Model States卸载到CPU/NVMe上后,在最大的单个算子上执行forward/backward pass所需的最小的GPU内存

因此需要从backward pass的最大算子考虑,因为backward需要涉及到parameters和gradient

在基于transformer架构的模型中,最大的算子是[$hd,4hd$]的线性层,所需的内存大小为:

$$16 * {hd}^{2}$$

其中$16$是通过FP16(2bytes)以及参数和梯度所需的参数量计算而得

Activation Working Memory(AWM)

激活工作内存是在反向传播重计算时所需的内存(在执行真正的反向传播之前)

它的大小是取决于两个连续的activation checkpoints.($c_i$代表的是两个连续的激活检查点间的transformer block数目)

所需的内存大小为:

$$bsz*seq*c_i*(16 * hd + 2 * attn\_heads * seq)$$

$c_i$取1,即每个block都有个activation checkpoints

其中本人并没有计算出来这个数目,计算过程如下图(仅供参考):

Bandwidth Requirements

如果我们要将暂时用不上的数据先offload到CPU或NVMe对应的存储空间,为了确保训练效率不受影响,因此得确保带宽够用,满足带宽需求

定义以下的字符用来作为效率的度量指标:

• $peak_{tp}$: 峰值计算吞吐量
• $bw$: 数据移动带宽
• $ait$: 算术强度,是总计算量$total\_computation$与计算所需数据$total\_data\_movement$的比率,描述了每次数据移动的计算量,越高的$ait$则说明对$bw$要求越低,因为对每个加载的数据,加速器可以做更多的计算

效率计算公式:

$$efficiency = \frac{compute\_time}{compute\_time + comm\_time}$$

$$compute\_time = \frac{total\_computation}{peak_{tp}}$$

$$comm\_time = \frac{total\_data\_movement}{bw}$$

$$ait = \frac{total\_computation}{total\_data\_movement}$$

联立上述四个式子,可得:

$$efficiency = \frac{ait * bw}{ait * bw + peak_{tp}}$$

在讨论$bw$前,需要先量化$ait$.我们知道每一个transformer_block的操作都是相同的,因此可以从一个迭代中的总计算量和总数据移动量来考虑算术强度,其中model states和activation checkpoints的ait可以分开考虑(之所以分开考虑,是因为其相关的参数不同)

Total Computation per iter

AIT w.r.t. Parameters and Gradients

AIT w.r.t. Optimizer States

AIT w.r.t. Activation Checkpoints

ZeRO-Infinity design