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

多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]$。是vector-wise quantization，对每个矩阵的行row/列colum的normalization constant（最大值），对原fp16数值x进行quant过程：$y=\frac{x}{x_{max}}127$，得到[-127, 127]范围的int8数值。将int8还原回fp16的dequant过程：$x=\frac{y}{127}{x}_{max}$
  

1.2 加速推理 Inference

1. 使用 TensorFloat-32：在 Ampere 和更高版本的 CUDA 设备上，矩阵乘法和卷积可以使用 TensorFloat-32 （TF32） 模式进行更快但精度稍低的计算。默认情况下，PyTorch 为卷积启用 TF32 模式，但不启用矩阵乘法。除非您的网络需要完整的 float32 精度，否则我们建议启用 TF32 进行矩阵乘法。它可以显著加快计算速度，而数值精度损失通常可以忽略不计。

import torch
torch.backends.cuda.matmul.allow_tf32 = True
1
2

2. 半精度权重：为了节省 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]
1
2
3
4
5
6
7
8
9
10
11
12

3. 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
2
3
4
5
6
7

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()

1
2
3
4
5
6
7
8
9
10
11

使用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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14

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
1
2
3
4
5
6

• 命令行控制：

CUDA_VISIBLE_DEVICES=0,1 python train.py
1

② 推理换算

• 模型加载：
  （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. 模型权重，计算方法和推理一样。
2. 优化器：（1）如果你采用AdamW，每个参数需要占用8个字节，因为需要维护两个状态。也就说优化器使用显存是全精度(float32)模型权重的2倍。（2）如果采用bitsandbytes优化的AdamW，每个参数需要占用2个字节，也就是全精度(float32)模型权重的一半。（3）如果采用SGD，则优化器占用显存和全精度模型权重一样。
3. 梯度：梯度占用显存和全精度(float32)模型权重一样。
4. 计算图内部变量：有时候也叫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，这使得大模型的训练权重得以加载。（数据并行+模型并行）

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

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

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）。
    

• 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）
  

  1. get fp32：存储一份fp32的Model States：parameter，momentum和variance
  2. fp32-to-fp16：在forward开始之前，额外开辟一块存储空间，将fp32 parameter减半到fp16 parameter。
  3. fp16 computing：正常做forward和backward，在此之间产生的activation和gradients，都用fp16进行存储。
  4. 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
1

多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")
1
2
3
4
5
6
7
8
9
10
11
12
13

accelerate launch run_distributed.py --num_processes=8
1

更详细可以查看：使用 Accelerate 进行