Read paper: ZeRO-Infinity
本文最后更新于:1 年前
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个字节的内存
用来代表hidden dimension,用表示transformer层的数目.考虑到在一个transformer块中几乎所有参数是来自块内的四个线性层:其中两个在多头注意力机制中①input -> QKV ([,]);② ([,]);剩下两个在FFN中③[,];④[,].因此总参数量可以表示为
Model States所需的总内存则乘于每个参数所需的20个字节,即为:
Memory for Residual States
Residual States主要包含的是activation memory,它的参数主要取决于模型的架构:批处理大小,序列长度
本文采用activation checkpoint的方式来对activation所需的内存进行减少,因此存储的是激活检查点的数据
用于表示两个activation checkpoints间的transformer block的个数,因此表示的是激活检查点的个数
激活检查点所需的总内存为:
Model State Working Memory(MSWM)
模型状态工作内存是在所有Model States卸载到CPU/NVMe上后,在最大的单个算子上执行forward/backward pass所需的最小的GPU内存
因此需要从backward pass的最大算子考虑,因为backward需要涉及到parameters和gradient
在基于transformer架构的模型中,最大的算子是[]的线性层,所需的内存大小为:
其中是通过FP16(2bytes)以及参数和梯度所需的参数量计算而得
Activation Working Memory(AWM)
激活工作内存是在反向传播重计算时所需的内存(在执行真正的反向传播之前)
它的大小是取决于两个连续的activation checkpoints.(代表的是两个连续的激活检查点间的transformer block数目)
所需的内存大小为:
取1,即每个block都有个activation checkpoints
其中本人并没有计算出来这个数目,计算过程如下图(仅供参考):
Bandwidth Requirements
如果我们要将暂时用不上的数据先offload到CPU或NVMe对应的存储空间,为了确保训练效率不受影响,因此得确保带宽够用,满足带宽需求
定义以下的字符用来作为效率的度量指标:
- : 峰值计算吞吐量
- : 数据移动带宽
- : 算术强度,是总计算量与计算所需数据的比率,描述了每次数据移动的计算量,越高的则说明对要求越低,因为对每个加载的数据,加速器可以做更多的计算
效率计算公式:
联立上述四个式子,可得:
在讨论前,需要先量化.我们知道每一个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
本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 协议 ,转载请注明出处!