《多GPU大模型训练与微调手册》

在这里插入图片描述

多GPU微调预备知识

1. 参数数据类型 torch.dtype

在这里插入图片描述

1.1 半精度 half-precision

torch.float16：fp16 就是 float16，1个 sign（符号位），5个 exponent bits(指数位)，10个 mantissa bits(小数位)
torch.bfloat16：bf 16 就是 brain float16，1个 :符号位，8个exponent bits（指数位），7个mantissa bits（小数位）
区别：bf16 牺牲了精度（小数位），实现了比 fp16 更大的范围（多了三个指数位）。

1.2 全精度 single-precision

torch.float32：fp 32 就是 float32，1个 sign（符号位），8个 exponent bits(指数位)，23个 mantissa bits(小数位)

2. 显卡环境

2.1 参数量与显存换算

例如，实验室是单机多卡：8卡A6000（40G）服务器 320G显存

① CUDA_VISIBLE_DEVICES 控制显卡可见性

通过CUDA_VISIBLE_DEVICES环境变量控制哪些GPU可以被torch调用：

代码控制：

# 必须置于 import torch 之前，准确地说在 torch.cuda 的调用之前
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3,4,5,6,7'
import torch
torch.cuda.device_count()
# 8

命令行控制：

CUDA_VISIBLE_DEVICES=0,1 python train.py

② 推理换算

模型加载：
（1）目前模型的参数绝大多数都是float32类型, 每个参数占用 4 个字节。所以一个粗略的计算方法就是，每10亿个参数(1 billion＝10亿)，占用4G显存 (实际应该是10^9 * 4 / 1024 / 1024 / 1024 = 3.725G，为了方便可以记为4G)，即 1B Params= 4G VRAM。比如LLaMA的参数量为7000559616个Params，那么全精度加载这个模型参数需要的显存为：7000559616 * 4 /1024/1024/1024 = 26.08G。
（2）显存不够，可以用半精度的fp16/bf16来加载，这样每个参数只占2个字节，所需显存就降为一半，只需要13.04G。
（3）如果显存还不够，可以采用int8的精度，显存再降一半，仅需6.5G，但是模型效果会更差一些。
（4）如果显存还是不够，int4精度显存再降一半，仅需3.26G。int4就是最低精度了，再往下模型推理效果就很难保证了。
模型推理：注意上面只是加载模型到显存，模型运算时的一些临时变量也需要申请空间，比如你beam search的时候。所以真正做推理的时候记得留一些Buffer，不然就容易OOM。如果显存还不够，就只能采用Memery Offload的技术，把部分显存的内容给挪到内存，但是这样会显著降低推理速度。

③ 训练换算

模型训练的时候显存使用包括如下几部分：

模型权重，计算方法和推理一样。
优化器：（1）如果你采用AdamW，每个参数需要占用8个字节，因为需要维护两个状态。也就说优化器使用显存是全精度(float32)模型权重的2倍。（2）如果采用bitsandbytes优化的AdamW，每个参数需要占用2个字节，也就是全精度(float32)模型权重的一半。（3）如果采用SGD，则优化器占用显存和全精度模型权重一样。
梯度：梯度占用显存和全精度(float32)模型权重一样。
计算图内部变量：有时候也叫Forward Activations。

如果模型想要训练，只看前3部分，需要的显存是至少推理的3-4倍。7B的全精度模型加载需要78G ~ 104G。 然后计算图内部变量这一部分只能在运行时候观测了，可以两个不同的batch的占用显存的差值大概估算出来。

优化的思路也就有了，目前市面上主流的一些计算加速的框架如DeepSpeed, Megatron等都在降低显存方面做了很多优化工作，比如量化，模型切分，混合精度计算，Memory Offload等等。

2.2 分布式架构

在这里插入图片描述
3种并行方式：

数据并行Data Paralleism：模型复制到不同GPU上，将数据切分后，分配到不同的GPU上。
模型并行Model Paralleism：将模型切分后，分配到不同的GPU上。分为张量并行和流水线并行。
- 张量并行Tensor Paralleism：对模型参数 tensor 切分，分配到不同的GPU进行计算，在参数更新的时候再进行同步。
- 流水线并行Pipeline Paralleism：对模型按层layer切分，分配到不同的GPU上进行计算。
混合并行Hybrid Paralleism：同时进行数据并行、张量并行、流水线并行。

下面3个分布式框架都是基于 Pytorch 的并行框架:

DP(torch.nn.DataParallel)：单机-单进程多线程进行实现的，它使用一个进程来计算模型权重，在每个batch处理期间将数据分发到每个GPU，每个GPU 分发到 batch_size/N 个数据，各个GPU的forward结果汇聚到master GPU上计算loss，计算梯度更新master GPU参数，将参数复制给其他GPU。（数据并行）
DDP(torch.nn.DistributedDataParallel)：可以单机/多机-多进程进行实现的，每个GPU对应的进程都有独立的优化器，执行自己的更新过程。每个进程都执行相同的任务，并且每个进程都与所有其他进程通信。进程（GPU）之间只传递梯度，这样网络通信就不再是瓶颈。（数据并行）
FSDP(torch.distributed.fsdp.FullyShardedDataParallel)：Pytorch最新的数据并行方案，在1.11版本引入的新特性，目的主要是用于训练大模型。我们都知道Pytorch DDP用起来简单方便，但是要求整个模型加载到一个GPU上，维护模型参数、梯度和优化器状态的每个 GPU 副本。FSDP则可以在数据并行的基础上，将模型参数和优化器分片分配到 GPU，这使得大模型的训练权重得以加载。（数据并行+模型并行）

这些在前面的博客已经讲过：

分布式并行训练（DP、DDP、DeepSpeed）
pytorch单精度、半精度、混合精度、单卡、多卡（DP / DDP）、FSDP、DeepSpeed模型训练

2.3 分布式工具

前面的分布式框架使用起来较为麻烦，因此分布式工具在底层对torch的分布式框架进行封装，实现更加方便的分布式训练和微调：

DerepSpeed（微软开发）
Accelerate（Huggingface开发）

① DerepSpeed—Zero

DerepSpeed的原理是基于微软的研究：Zero（零冗余优化），研究哪些部分是占用存储空间的，并对这些占用存储的数据进行优化。
在这里插入图片描述

存储空间的消耗 Memory Consumption主要包含两部分：

Model States(主)：模型参数Parameters、梯度Gradients、优化器Optimizer_State
Residual States(次)：前向传播激活值Activations、临时缓存区Temporal Buffers、内存碎片Unusable Fragmented Memory

知道了什么东西会占存储，以及它们占了多大的存储之后，我们就可以来谈如何优化存储了。注意到，在整个训练中，有很多states并不会每时每刻都用到；因此提出了三种Zero优化方法：

Zero-DP（优化Model States）：作者采取三个方法优化内存，Pos、Pg、Pp。大体思路都是一样的，把每个模型的参数、梯度、优化器状态分别平均分给所有的gpu，当时计算需要用到其他gpu的内容时，通过GPU之间的通讯传输，以通讯换内存。其中前两个方法不增加通讯成本，第三个方法会增加GPU之间的通信成本。
Zero-R（优化Residual States）：（1）激活函数：在前向传播计算完成激活函数之后，对把激活值丢弃，由于计算图还在，等到反向传播的时候，再次计算激活值，算力换内存。或者采取一个与cpu执行一个换入换出的操作。（2）临时缓冲区：模型训练过程中经常会创建一些大小不等的临时缓冲区，比如对梯度进行All Reduce啥的，解决办法就是预先创建一个固定的缓冲区，训练过程中不再动态创建，如果要创建临时数据，在固定缓冲区创建就好。（3）内存碎片：显存出现碎片的一大原因是时候gradient checkpointing后，不断地创建和销毁那些不保存的激活值，解决方法是预先分配一块连续的显存，将常驻显存的模型状态和checkpointed activation存在里面，剩余显存用于动态创建和销毁discarded activation复用了操作系统对内存的优化，不断内存整理。
混合精度训练：对于模型，我们肯定希望其参数越精准越好，也即我们用fp32（单精度浮点数，存储占4byte）来表示参数W。但是在forward和backward的过程中，fp32的计算开销也是庞大的。那么能否在计算的过程中，引入fp16或bf16（半精度浮点数，存储占2byte），来减轻计算压力呢？于是，混合精度训练（float2hlaf）就产生了，它的步骤如下图：

get fp32：存储一份fp32的Model States：parameter，momentum和variance
fp32-to-fp16：在forward开始之前，额外开辟一块存储空间，将fp32 parameter减半到fp16 parameter。
fp16 computing：正常做forward和backward，在此之间产生的activation和gradients，都用fp16进行存储。
update fp32 model states：用fp16 gradients去更新fp32下的model states。

Zero-Offload：GPU显存不够，CPU内存来凑。如下图，左边是正常的计算图，右侧是Zero-Offload的计算图。（⭕️表示state，正方形表示计算图，箭头表示数据流向、M表示模型参数，float2half表示32位转16位）其实就是forward和backward在GPU上计算，参数更新在CPU上。因为CPU与GPU通信数据开销很大，所以CPU和GPU传播的是gradient16，这样保证传播数据量最小。
Zero-Infinity：GPU内存不够，SSD外存来凑。

②Accelerate—Huggingface

安装Accelerate之后，用accelerate config配置Accelerate训练模式（Yes/No），配置完成后会根据你回答的问题生成一个yaml文件，我的位于~/.cache/huggingface/accelerate，如果是单机多卡，num_processes指的就是GPU数量（多机多卡不了解）。

在这里插入图片描述
然后运行accelerate test来测试脚本能否正常工作。
一切都ok后，我们就能开始训练了：

accelerate launch path_to_script.py --args_for_the_script

使用：用Accelerate对象包装模型、优化器、dataloarder、scheduler
在这里插入图片描述