大模型训练优化方法

写在前面

在训练模型尤其是大模型的时候，如何加快训练速度以及优化显存利用率是一个很关键的问题。本文主要参考HF上的一篇文章：https://huggingface.co/docs/transformers/perf_train_gpu_one，以及笔者在实际训练中的一些经验，给出一些比较实用的方法。

先看一个总览的表：

方法	加快训练速度	优化显存利用率
Batch size choice	Yes	Yes
Gradient accumulation	No	Yes
Gradient checkpointing	No	Yes
Mixed precision training	Yes	No
Optimizer choice	Yes	Yes
Data preloading	Yes	No
DeepSpeed Zero	No	Yes
torch.compile	Yes	No

其实加快训练速度本质就是提高数据吞吐量，每秒能够处理多少samples，在不爆显存的情况下，尽可能的提高Batch size，但在如今模型参数巨大的情况下，可能一张卡都不够放模型，例如llama-2-7B以FP16加载，都需要接近13G的显存，如果以FP32加载，那就需要25G的显存，在消费级显卡30、40系只有24G显存的卡上就难以训练或者微调。所以就有了一些列的优化显存利用率的方法。注：本篇先不讨论分布式训练的优化，将会在下篇文章讨论。

1.Batch size choice

增大Batch size既可以加快训练速度，又可以提高显存利用率。这个没啥原理，每秒能够处理的samples变多了，唯一需要祷告的就是不出现OOM。不过Batch size需要注意的是其值最好设置为2^N，为什么呢？这个和矩阵乘、模型维度、硬件都有关系，不同的Batch size是会影响到计算效率的，这是NVIDIA官方做的关于Batch size的实验[2],下面的图是NVIDIA推荐的Batch size：

2.Gradient Accumulation

梯度累积。这方法让俺想起了当年在学校实验室炼丹的过程，真是贫穷人士的灵丹妙药，没有足够的显卡搞不了大的Batch size，但是想要有大Batch size的效果，那就来试试梯度累积。

简单来说原理就是：每次来一个Batch size的数据，进行前向反向传播，保存梯度不清零，不更新模型参数，梯度积累到一定次数后，根据累积的梯度更新参数，然后梯度清零，进行下一次循环。所以该种方法加快不了训练速度，只是提高了显存的利用率，让你原来Batch size只能设置成2，现在可以设成8了！

先看一下不进行梯度累积的伪代码：

for batch in dataset:  # 遍历每个数据集的批次  
    optimizer.zero_grad()  # 梯度清零  
  
    outputs = model(inputs)  # 前向传播，计算输出  
    loss = compute_loss(outputs, labels)  # 计算损失  
  
    loss.backward()  # 反向传播，计算梯度  
    optimizer.step()  # 更新模型参数

看下进行梯度累积的伪代码：

for batch in dataset:  # 遍历每个数据集的批次  
    # 重置梯度  
    if batch_index % num_accumulation_steps == 0:  
        optimizer.zero_grad()  # 每 num_accumulation_steps 步重置一次梯度  
  
    # 前向传播  
    outputs = model(inputs)  # 计算输出  
    loss = compute_loss(outputs, labels)  # 计算损失  
  
    # 反向传播（累积梯度）  
    loss.backward()  # 反向传播，计算并累积梯度  
  
    # 累积了足够的梯度后更新参数  
    if (batch_index + 1) % num_accumulation_steps == 0:  
        optimizer.step()  # 更新模型参数

使用梯度累积时需要注意学习率的设置，毕竟小Batch size和大Batch size的学习率还是有区别的，经验的话就是尽量调大点。

3.Gradient checkpointing

梯度检查点。当你训练大一点的模型，Batch size已经设置成1了，但还是出现OOM了，此时如果换不了大显存的显卡，你还是想硬训，那可以试试梯度检查点。

原理还是比较简单的，因为在训练过程中，除了load模型占了一大部分显存外，还有一部分占显存的就是为了在反向传播过程中计算梯度，保存的前向传播的所有激活值。梯度检查点就是选择性地保存一些的激活值，这样在反向传播过程中计算梯度时重新计算没保存的这部分激活值，就是以时间换空间。所以这种方法不会加快训练训练速度，还会减慢训练速度，因为部分节点会进行两次前向传播，但极大地提高了显存的利用率。根据HF给出的数据使用梯度检查点，会使训练速度大约降低20%。

下面是一些动图简单的反映了梯度检查点的过程。源于该篇blog[3]：

正常的前向反向传播过程，需要保存的前向传播的所有激活值

选择第三个点作为检查点

这时，只需要保存第一个节点，选择的检查点以及最后一个节点，节省了第二个和第三个节点保存的显存，但反向传播的时候需要通过前向传播重新计算第二个和第三个节点的激活值

4.Mixed precision training

混合精度训练。要了解混合精度训练，首先要了解各个数值类型，正常模型训练的数值类型都是FP32，这个数值类型虽然精度和范围更大，但占用的显存和计算量也更大。所以混合精度训练就是在训练时一部分计算使用FP16，一部分计算使用FP32，加快训练速度同时也减少了显存占用（理论上应该是减少了显存占用，但实际上并没有特别明显）。但是注意：模型较小batch也较小，硬件也比较拉的情况下，混合精度训练可能会没有效果甚至会更慢。这时候训练速度主要消耗在IO（在GPU和GPU间传送数据），以及混合精度频繁进行FP16与FP32的转换。

那哪些计算使用FP16，哪些使用FP32呢？在NVIDIA的这篇Paper中[4]

首先会拷贝一份FP32权重的副本，然后整个计算过程为：权重从FP32转成FP16进行前向传播并计算loss，然后继续使用FP16反向传播计算梯度，最后转成FP32更新到FP32的权重上，这样循环计算。用FP32保存权重主要是为了避免溢出，保持一定的精度。Paper中实验表明，用FP16保存权重会造成80%的精度损失。细节可以看下Paper。

目前torch实现的是自动混合精度，他会自动的管理精度，并不是像上面paper中的前向反向都使用FP16，比如归一化（batch normalization），梯度累积时等需要高精度的也还是用的FP32，细节可以看[5]，还有如果显卡支持，还可以实现BF16，TF32等数据类型。下面是伪代码：

import torch  
import torch.nn as nn  
import torch.optim as optim  
from torch.cuda.amp import autocast, GradScaler  
  
# 定义模型、损失函数和优化器  
model = MyModel().cuda()  
criterion = nn.CrossEntropyLoss()  
optimizer = optim.Adam(model.parameters(), lr=0.001)  
  
# 初始化梯度缩放器  
scaler = GradScaler()  
  
for data, target in train_loader:  
    data, target = data.cuda(), target.cuda()  
  
    optimizer.zero_grad()  
  
    # 在autocast上下文中进行前向传播和损失计算  
    with autocast():  
        output = model(data)  
        loss = criterion(output, target)  
  
    # 使用梯度缩放器进行反向传播计算梯度  
    scaler.scale(loss).backward()  
  
    # 更新优化器  
    scaler.step(optimizer)  
  
    # 更新缩放器  
    scaler.update()

5.Optimizer choice

选择合适的优化器是既可以加快训练速度又可以提高显存利用率的。目前训练Transformer类模型最常用的优化器是Adam或AdamW（带权重衰减的Adam），Adam通过存储先前梯度的滚动平均值可以实现比较好的收敛，但这增加了与模型参数数量成正比的额外显存占用。如果安装了NVIDIA/apex，那么adamw_apex_fused是训练速度最快的优化器。

HF的Trainer目前集成了很多优化器，也就是指定一个参数的是事情。adamw_hf, adamw_torch, adamw_torch_fused, adamw_apex_fused, adamw_anyprecision, adafactor, or adamw_bnb_8bit这些等等。

简单看下两种AdamW优化器的替代选择的显存对比

对于一个3B参数的模型，“google-t5/t5-3b”

标准AdamW优化器需要24G左右的显存，它每个参数使用8字节存储(8*3 => 24GB)

Adafactor优化器需要12G左右的显存，它每个参数使用4字节存储(4*3 => 12GB)

adamw_bnb_8bit（量化后的AdamW）优化器需要6G左右的显存，它每个参数使用2字节存储(2*3 => 6GB)

但这这些显存占用明显改进的优化器，可能会带来收敛速度慢，精度下降等问题，还是要具体情况具体分析，俺一般使用adamw_apex_fused。

6.Data preloading

数据预加载。想要加快训练速度，很重要的一点就是要能够以GPU可以处理的最大速度来提供数据，使GPU的利用率达到100%。主要有以下两个方法加快数据供给。

DataLoader(pin_memory=True, …)，将数据预加载到CPU的锁页内存中，什么是锁页内存，简单来说就是可直接访问该页CPU内存；未锁页的内存，会被更换到磁盘上，如果下次需要该页了，则需要重新加载到内存里。这就会需要额外的时间了。相关的一些资料[6][7]原文中也给了了相应的速度:

锁页内存和GPU显存之间的拷贝速度大约是6GB/s
可分页内存和GPU显存间的拷贝速度大约是3GB/s

但要注意，如果数据量特别大，CPU的内存不是特别大的情况下，谨慎使用此参数，遇到过内存爆了的情况。。

第二个是DataLoader(num_workers=4, …)，给DataLoader使用更多的线程预加载数据。

7.DeepSpeed Zero

Deepspeed 核心就是四个字”显存优化“，其并不能加快训练速度，如果你显存够用是可以不使用Deepspeed 的，但训练大模型一般来说显存都不够用。

8.Using torch.compile

PyTorch 2.0 引入了一个新的编译函数，无需对现有的 PyTorch 代码进行任何修改，只需添加一行代码即可优化代码：model = torch.compile(model)。即可进行训练和推理加速。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述