文章目录
多GPU微调预备知识
1. 参数数据类型 torch.dtype
1.1 半精度 half-precision, 全精度 single-precision, 量化quant
-
torch.float16(1+5+10)
:fp16 就是 float16,1个 sign(符号位),5个 exponent bits(指数位),10个 mantissa bits(小数位) -
torch.bfloat16(1+8+7)
:bf 16 就是 brain float16,1个 :符号位,8个exponent bits(指数位),7个mantissa bits(小数位) -
区别:bf16 牺牲了精度(小数位),实现了比 fp16 更大的范围(多了三个指数位)。因此bf16以更低的存储实现了逼近fp32的精度。
-
torch.float32(1+8+23)
:fp 32 就是 float32,1个 sign(符号位),8个 exponent bits(指数位),23个 mantissa bits(小数位) -
LLM.int8(1+7)
:int8 就是 纯int值 无小数部分,1个 sign(符号位),7个 exponent bits(指数位), [ − 2 7 + 1 , 2 7 − 1 ] [-2^7+1,2^7-1] [−27+1,27−1]。是vector-wise quantization,对每个矩阵的行row
/列colum
的normalization constant(最大值),对原fp16数值x进行quant过程: y = x x m a x 127 y=\frac{x}{x_{max}}127 y=xmaxx127,得到[-127, 127]范围的int8数值。将int8还原回fp16的dequant过程: x = y 127 x m a x x=\frac{y}{127}x_{max} x=127yxmax
1.2 加速推理 Inference
- 使用 TensorFloat-32:在 Ampere 和更高版本的 CUDA 设备上,矩阵乘法和卷积可以使用 TensorFloat-32 (TF32) 模式进行更快但精度稍低的计算。默认情况下,PyTorch 为卷积启用 TF32 模式,但不启用矩阵乘法。除非您的网络需要完整的 float32 精度,否则我们建议启用 TF32 进行矩阵乘法。它可以显著加快计算速度,而数值精度损失通常可以忽略不计。
import torch
torch.backends.cuda.matmul.allow_tf32 = True
- 半精度权重:为了节省 GPU 内存并获得更快的速度,请尝试直接以半精度或 float16 加载和运行模型权重:
import torch
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
)
pipe = pipe.to("cuda")
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
- int4/int8 量化权重:使用 4 比特量化的不同变体,例如 NF4 (NormalFloat4 (默认) ) 或纯 FP4 量化。从理论分析和实证结果来看,我们建议使用 NF4 量化以获得更好的性能。其他选项包括 bnb_4bit_use_double_quant ,它在第一轮量化之后会进行第二轮量化,为每个参数额外节省 0.4 比特。最后是计算类型,虽然 4 比特 bitsandbytes 以 4 比特存储权重,但计算仍然以 16 或 32 比特进行,这里可以选择任意组合 (float16、bfloat16、float32 等)。如果使用 16 比特计算数据类型 (默认 torch.float32),矩阵乘法和训练将会更快。用户应该利用 transformers 中最新的
BitsAndBytesConfig
来更改这些参数。下面是使用 NF4 量化加载 4 比特模型的示例,例子中使用了双量化以及 bfloat16 计算数据类型以加速训练。传入huggingface的from_pertrain()。
# BitsAndBytesConfig设置推理精度
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16
)
1.3 混合精度训练amp
混合精度训练torch.cuda.amp
:除了可以**显著降低显存(VRAM)外,对于训练加速(加快实验速度)**也是有很大作用的。
正常的fp32训练流程:
for epoch in tqdm(range(num_epochs)):
for batch_x, batch_y in tqdm(train_dataloader):
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)
logits = model(batch_x)
loss = loss_fn(logits, batch_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
使用torch.cuda.amp
的混合精度训练:
scaler = torch.cuda.amp.GradScaler()
for epoch in tqdm(range(num_epochs)):
for batch_x, batch_y in tqdm(train_dataloader):
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)
# forward in auto_mix_precision mode
with torch.cuda.amp.autocast():
logits = model(batch_x)
loss = loss_fn(logits, batch_y)
# scale loss's precision before backward
scaler.scale(loss).backward()
scaler.step(optimizer) # optimizer.step()
scaler.update()
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切分(layer内的tensor不切分),分配到不同的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,这使得大模型的训练权重得以加载。(数据并行+模型并行
)
这些在前面的博客已经讲过:
2.3 分布式工具
前面的分布式框架使用起来较为麻烦,因此分布式工具在底层对torch的分布式框架进行封装,实现更加方便的分布式训练和微调:
DerepSpeed
(微软开发)Accelerate
(Huggingface开发)
① DeepSpeed—Zero
DerepSpeed的原理是基于微软的研究:Zero(零冗余优化)
,研究哪些部分是占用存储空间的,并对这些占用存储的数据进行优化。
GPU存储空间的消耗 Memory Consumption主要包含两部分:
- Model States(主):
模型参数Parameters
、梯度Gradients
、优化器Optimizer_State
- Residual States(次):
前向传播激活值Activations
、临时缓存区Temporal Buffers
、内存碎片Unusable Fragmented Memory
知道了什么东西会占存储,以及它们占了多大的存储之后,我们就可以来谈如何优化存储了。注意到,在整个训练中,有很多states并不会每时每刻都用到;因此提出了三种Zero优化方法:
-
Zero-DP(
优化Model States
):以普通的Data Parallel数据并行为Baseline(模型参数、优化器参数、梯度在每个GPU上都复制一份,存在冗余),作者采取三个方法优化显存消耗,Pos、Pg、Pp。大体思路都是一样的,把每个模型的优化器 Zero-1
、梯度 Zero-2
、模型参数 Zero-3
状态分别平均分给所有的GPU,当时计算需要用到其他GPU的内容时,通过GPU之间的通讯传输,以通讯时间换内存。(其中前两个方法不增加通讯成本,第三个方法会增加GPU之间的通信成本)。因此Zero的3个阶段分别再原始的数据并行
基础上,实现优化器并行(每个GPU仅优化它有的部分参数)
、梯度并行
、模型并行
。
- 如下2个GPU使用
数据并行
进行混合精度
训练Transforemr:每个GPU都复制了一份ModelParameters、Gradients、Optimizer States、Activations。不同的Data分别会送入不同的GPU进行计算。 - 假如在
数据并行
基础上使用了Zero-1 优化器并行
:每个GPU只存储一半的Optimizer States,先进行前向传播计算loss,再反向传播计算梯度,每个GPU将自己的梯度发送给其他GPU(如下图的红色和绿色箭头),每个GPU将来自其他GPU的梯度进行平均,平均梯度
用于更新每个GPU的优化器中对应部分的FP32模型参数
,然后将FP32的模型参数复制回FP16模型参数中(如下图的蓝色箭头),最后可以进行all gather
(因为每个GPU只优化其对应部分的模型参数,最后要将自己优化好的FP16模型参数,传给其他GPU)。
- 如下2个GPU使用
-
Zero-R(
优化Residual States
):(1)激活函数:在前向传播计算完成激活函数之后,对把激活值丢弃,由于计算图还在,等到反向传播的时候,再次计算激活值
,算力换内存。或者采取一个与cpu执行一个换入换出的操作。(2)临时缓冲区:模型训练过程中经常会创建一些大小不等的临时缓冲区,比如对梯度进行All Reduce啥的,解决办法就是预先创建一个固定的缓冲区,训练过程中不再动态创建
,如果要创建临时数据,在固定缓冲区创建就好。(3)内存碎片:GPU出现碎片的一大原因是时候gradient checkpointing后,不断地创建和销毁那些不保存的激活值,解决方法是预先分配一块连续的显存,将常驻显存的模型状态和checkpointed activation存在里面,剩余显存用于动态创建和销毁discarded activation复用了操作系统对内存的优化,不断内存整理。 -
Zero-Offload:
GPU显存不够,CPU内存来凑
。如下图,左边是正常的计算图,右侧是Zero-Offload的计算图。(⭕️表示state,正方形表示计算图,箭头表示数据流向、M表示模型参数,float2half表示32位转16位)其实就是forward和backward在GPU上计算
,参数更新在CPU上
。因为CPU与GPU通信数据开销很大,所以CPU和GPU传播的是gradient16,这样保证传播数据量最小。
-
Zero-Infinity:
GPU内存不够,SSD外存来凑
。
-
混合精度训练:对于模型,我们肯定希望其参数越精准越好,也即我们用
fp32(单精度浮点数,存储占4byte)
来表示参数W。但是在forward和backward的过程中,fp32的计算开销也是庞大的。那么能否在计算的过程中,引入fp16或bf16(半精度浮点数,存储占2byte)
,来减轻计算压力呢?于是,混合精度训练(float2hlaf
)就产生了,它的步骤如下图:(只在优化器
中是FP32
,在其他地方
的都是FP16
)
get fp32
:存储一份fp32的Model States:parameter,momentum和variancefp32-to-fp16
:在forward开始之前,额外开辟一块存储空间,将fp32 parameter减半到fp16 parameter。fp16 computing
:正常做forward和backward,在此之间产生的activation和gradients,都用fp16进行存储。update fp32 model states
:用fp16 gradients去更新fp32下的model states。
-
DeepSpeed设置总结:
②Accelerate—Huggingface
Accelerate以一种简单的方式集成了Pytorch DDP
和DeepSpeed Zero
,为开发者提供了更加简洁的接口。
配置Accelerate:安装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
多GPU训练:用Accelerate对象包装模型、优化器、dataloarder、scheduler,loss的反向传播时也要改成accelerate的。
显示的使用多GPU推理StableDiffusion:使用 --num_processes
参数来指定GPU个数,调用accelerate launch运行py脚本:
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipeline(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")
accelerate launch run_distributed.py --num_processes=8
更详细可以查看:使用 Accelerate 进行🤗分布式推理
GPU内存优化技术集合:
import torch
torch.backends.cuda.matmul.allow_tf32 = True
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
pipeline.enable_vae_slicing()
pipeline.enable_sequential_cpu_offload()
pipeline.enable_model_cpu_offload()
result = pipeline(prompt).images[0]
上面的例子使用了多种GPU内存优化的技术:
- 使用tf32替代 fp32
- Half precision weights半精度权重
- Sliced VAE decode for larger batches分片VAE解码用于更大的batch
- pipe.enable_sequential_cpu_offload():使用accelerate将权重转到CPU以节省内存
- enable_model_cpu_offload():用模型卸载快速推理和内存节省