文章目录
多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()：用模型卸载快速推理和内存节省
2998
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
