大模型训练优化方法

写在前面

在训练模型尤其是大模型的时候，如何加快训练速度以及优化显存利用率是一个很关键的问题。本文主要参考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大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

第一阶段（10天）：初阶应用

该阶段让大家对大模型 AI有一个最前沿的认识，对大模型 AI 的理解超过 95% 的人，可以在相关讨论时发表高级、不跟风、又接地气的见解，别人只会和 AI 聊天，而你能调教 AI，并能用代码将大模型和业务衔接。

• 大模型 AI 能干什么？
• 大模型是怎样获得「智能」的？
• 用好 AI 的核心心法
• 大模型应用业务架构
• 大模型应用技术架构
• 代码示例：向 GPT-3.5 灌入新知识
• 提示工程的意义和核心思想
• Prompt 典型构成
• 指令调优方法论
• 思维链和思维树
• Prompt 攻击和防范
• …

第二阶段（30天）：高阶应用

该阶段我们正式进入大模型 AI 进阶实战学习，学会构造私有知识库，扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架，抓住最新的技术进展，适合 Python 和 JavaScript 程序员。

• 为什么要做 RAG
• 搭建一个简单的 ChatPDF
• 检索的基础概念
• 什么是向量表示（Embeddings）
• 向量数据库与向量检索
• 基于向量检索的 RAG
• 搭建 RAG 系统的扩展知识
• 混合检索与 RAG-Fusion 简介
• 向量模型本地部署
• …

第三阶段（30天）：模型训练

恭喜你，如果学到这里，你基本可以找到一份大模型 AI相关的工作，自己也能训练 GPT 了！通过微调，训练自己的垂直大模型，能独立训练开源多模态大模型，掌握更多技术方案。

到此为止，大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗？

• 为什么要做 RAG
• 什么是模型
• 什么是模型训练
• 求解器 & 损失函数简介
• 小实验2：手写一个简单的神经网络并训练它
• 什么是训练/预训练/微调/轻量化微调
• Transformer结构简介
• 轻量化微调
• 实验数据集的构建
• …

第四阶段（20天）：商业闭环

对全球大模型从性能、吞吐量、成本等方面有一定的认知，可以在云端和本地等多种环境下部署大模型，找到适合自己的项目/创业方向，做一名被 AI 武装的产品经理。

• 硬件选型
• 带你了解全球大模型
• 使用国产大模型服务
• 搭建 OpenAI 代理
• 热身：基于阿里云 PAI 部署 Stable Diffusion
• 在本地计算机运行大模型
• 大模型的私有化部署
• 基于 vLLM 部署大模型
• 案例：如何优雅地在阿里云私有部署开源大模型
• 部署一套开源 LLM 项目
• 内容安全
• 互联网信息服务算法备案
• …

学习是一个过程，只要学习就会有挑战。天道酬勤，你越努力，就会成为越优秀的自己。

如果你能在15天内完成所有的任务，那你堪称天才。然而，如果你能完成 60-70% 的内容，你就已经开始具备成为一名大模型 AI 的正确特征了。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】