资讯

历史

科技

环境与自然

成长

游戏

财经

文学与艺术

美食

健康

家居

文化

情感

汽车

三农

军事

旅行

运动

教育

生活

星座命理

Transformer模型训练与内存优化完全指南

创作时间:

作者:

@小白创作中心

Transformer模型训练与内存优化完全指南

引用

CSDN

https://blog.csdn.net/qq_40206371/article/details/139077249

本文深入探讨了Transformer模型在训练过程中的内存占用和优化策略。从模型权重、优化器状态到梯度和前向激活，详细分析了各个组件的内存需求，并提供了包括梯度累积、混合精度训练、优化器选择在内的多种优化方案。对于从事机器学习和深度学习的开发者和研究人员具有重要参考价值。

1 模型操作的解剖

Transformer架构主要包括三类主要操作：

张量收缩：线性层和多头注意力的组件都进行批量矩阵-矩阵乘法。这些操作是训练Transformer时计算最密集的部分。
统计归一化：Softmax和层归一化比张量收缩的计算密集度要低一些，涉及一个或多个归约操作，其结果通过映射应用。
元素级操作：包括偏差、dropout、激活和残差连接。这些是计算最不密集的操作。

2 模型内存刨析

2.1 模型内存

2.1.1 根据模型名称中的数字推断模型大小

默认情况下，Hugging Face 的类如 TextGenerationPipeline 或 AutoModelForCausalLM 会以 float32 精度加载模型。这意味着每个参数需要 4 字节（32 位），因此一个具有 80 亿参数的“8B”模型将需要大约 32GB 的内存。这可能有些浪费！大多数现代语言模型都是在“bfloat16”精度下训练的，这种精度只需要每个参数 2 字节。

模型内存计算公式如下：

int8 模型内存 = 1 * 参数量（字节）
fp16 和 bf16 模型内存 = 2 * 参数量（字节）
fp32 模型内存 = 4 * 参数量（字节）

2.1.2 推理总内存

除了用于存储模型权重的内存外，在实际的前向传播过程中还会产生一些额外的开销。根据经验，这些额外开销通常控制在总内存的20%以内，因此推理总内存≈1.2×模型内存。

2.2 训练模型内存

训练模型使用了比将模型放在GPU上多得多的内存。在训练过程中有许多组件使用GPU内存：

模型权重
优化器状态
梯度
一个参数对应一个梯度值
保存的前向激活用于梯度计算
临时缓冲区
特定功能的内存

2.1 模型权重

fp32训练：4字节 * 参数数量
混合精度训练：6字节 * 参数数量（在内存中保留一个fp32和一个fp16的模型）

2.2 优化器状态：

取决于优化器的具体类型
如果是SGD就不需要额外显存开销
如果是带一阶动量（momentum）的SGD就是1倍的参数量
如果是Adam的话就要在momentum的基础上加上二阶动量，优化器状态所占显存就是参数的2倍

2.3 梯度：

fp32或混合精度训练：4字节 * 参数数量（梯度始终保持在fp32）

2.4 前向激活

大小取决于许多因素，关键因素是序列长度、隐藏大小和批量大小。有输入和输出通过前向和后向函数传递，并为梯度计算保存前向激活。

3 高效训练总览

方法/工具	提高训练速度	优化内存使用
批量大小选择	是	是
梯度累积	否	是
梯度检查点	否	是
混合精度训练	是	否
优化器选择	是	是
数据预加载	是	否
DeepSpeed Zero	否	是
torch.compile	是	否
参数高效微调（PEFT）	否	是

3.1 选择批量大小

为了达到最佳性能，首先要确定适当的批量大小。推荐使用2^N大小的批量大小和输入/输出神经元数量。通常是8的倍数，但根据使用的硬件和模型的数据类型，这个数值可以更高。Tensor Core要求根据数据类型和硬件定义乘数。例如，对于fp16数据类型，推荐使用8的倍数，除非是A100 GPU，在这种情况下使用64的倍数。

3.2 梯度累积

梯度累积方法旨在以较小的增量计算梯度，而不是一次性为整个批次计算。这种方法涉及通过对模型进行前向和后向传播，并在此过程中累积梯度，以较小的批次迭代计算梯度。一旦累积了足够数量的梯度，就执行模型的优化步骤。通过使用梯度累积，可以增加有效的批量大小，超出GPU内存容量的限制。然而，重要的是要注意，梯度累积引入的额外前向和后向传播可能会减慢训练过程。

通过添加 gradient_accumulation_steps 参数到 TrainingArguments 来启用梯度累积：

training_args = TrainingArguments(per_device_train_batch_size=1, 
                                  gradient_accumulation_steps=4, 
                                  **default_args)
#——————>有效批量大小变为4

较高数量的梯度累积步骤可能导致更明显的训练减慢。假设，不使用梯度累积时 per_device_train_batch_size=4 已达到GPU的限制。如果希望以大小为64的批次进行训练，不要将 per_device_train_batch_size 设置为1并且 gradient_accumulation_steps 设置为64。相反，保持 per_device_train_batch_size=4 并设置 gradient_accumulation_steps=16。这样可以在更好地利用可用GPU资源的同时，获得相同的有效批量大小。或者，使用Accelerate完全控制训练循环。

3.2.1 使用pytorch的方式实现梯度累计

######初始化和设置
steps = len(loader)
#steps 代表训练数据加载器 loader 中的总步数（批次数量
validation_steps = int(validation_steps * gradient_accumulation_steps)
'''
validation_steps 被重新计算为原先的validation_steps 乘以 gradient_accumulation_steps
'''  

for step, batch in enumerate(loader, 1):
    #enumerate(loader, 1) 返回批次数据和当前步数 step（从1开始计数）。
    '''
    假设已经准备好将 batch 中的数据分为模型输入 inputs 和目标 targets。
    '''
    outputs = model(inputs)
    #模型将 inputs 作为输入，计算输出 outputs
    loss = loss_fn(outputs, targets)
    #使用损失函数 loss_fn 计算输出 outputs 和目标 targets 之间的损失 loss
    
    if gradient_accumulation_steps > 1:
        loss = loss / gradient_accumulation_steps
    '''
    如果 gradient_accumulation_steps 大于1，将损失 loss 除以 gradient_accumulation_steps
    这样每一步的梯度相对较小，有助于稳定训练过程。
    '''
    loss.backward()
    #执行反向传播计算梯度
    
    if step % gradient_accumulation_steps == 0 or step == steps:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        optimizer.step()
        model.zero_grad()
    '''
    每 gradient_accumulation_steps 步或在训练结束时（step == steps）
    进行梯度裁剪（防止梯度爆炸），然后进行一次优化步骤并清零梯度。
    '''
    if step % validation_steps == 0:
        validation_loop()
    #每 validation_steps 步，调用 validation_loop() 进行模型验证

3.2.2 使用Accelerator的方式实现

3.3 梯度检查点

即使将批量大小设置为1并使用梯度累积，一些大型模型仍可能面临内存问题。这是因为还有其他组件也需要内存存储。一种方法是在向后传递过程中计算梯度时保存前向传递的所有激活，会导致显著的内存开销。另一种方法是丢弃激活，并在向后传递时需要时重新计算它们，会引入相当大的计算开销并减慢训练过程。梯度检查点提供了这两种方法之间的折中方案，在计算图中策略性地保存选定的激活，只需要重新计算一小部分激活以获取梯度。

要在Trainer中启用梯度检查点，请向TrainingArguments传递相应的标志：

training_args = TrainingArguments(
    per_device_train_batch_size=1, 
    gradient_accumulation_steps=4, 
    gradient_checkpointing=True, 
    **default_args
)

3.3.1 pytorch普通实现

3.4 混合精度训练

混合精度训练是一种旨在通过使用较低精度的数值格式来优化模型训练的计算效率的技术。传统上，大多数模型使用32位浮点精度（fp32或float32）来表示和处理变量。然而，并非所有变量都需要这种高精度级别来实现准确的结果。通过将某些变量的精度降低到如16位浮点（fp16或float16）等较低的数值格式，我们可以加速计算。由于这种方法中的一些计算以半精度进行，而有些仍以全精度进行，因此这种方法被称为混合精度训练。

3.4.1 fp16

尽管梯度也以半精度计算，但它们在优化步骤中转换回全精度，因此在这里不会节省内存。虽然混合精度训练结果计算更快，但也可能导致使用更多GPU内存，尤其是在小批量大小的情况下。这是因为模型现在以16位和32位精度同时存在于GPU上（GPU上的原始模型大小的1.5倍）。

要启用混合精度训练，将fp16标志设置为True：

training_args = TrainingArguments(per_device_train_batch_size=4, 
                                  fp16=True, 
                                  **default_args)

3.5 优化器选择

对于Transformer，最常用的优化器是Adam或AdamW（带权重衰减的Adam）。Adam通过存储之前梯度的滚动平均来实现良好的收敛性，然而，它会增加与模型参数数量相当的额外内存占用。

例如，对于一个具有30亿参数的模型，如“google-t5/t5-3b”：

标准的AdamW优化器将需要24GB的GPU内存，因为它为每个参数使用8字节（8*3 => 24GB）
Adafactor优化器将需要超过12GB。它每个参数略多于4字节，所以是4*3再加上一些额外的。
8位BNB量化优化器只会使用6GB（2*3），如果所有优化器状态都被量化。

为了解决这个问题，可以使用其他优化器。

3.5.1 Adafactor

Adafactor不会为权重矩阵中的每个元素存储滚动平均，而是保留汇总信息（按行和按列的滚动平均的总和），显著减少了其占用空间。然而，与Adam相比，Adafactor在某些情况下可能收敛速度较慢。

training_args = TrainingArguments(per_device_train_batch_size=4, 
                                  optim="adafactor", 
                                  **default_args)

3.5.2 8位Adam

与Adafactor聚合优化器状态不同，8位Adam保留完整状态并对其进行量化。量化意味着它以较低的精度存储状态，并且只在优化时解量化。这类似于混合精度训练背后的思想。

training_args = TrainingArguments(per_device_train_batch_size=4, 
                                  optim="adamw_bnb_8bit", 
                                  **default_args)

3.6 数据预加载

为了达到极高的训练速度，能够以GPU能够处理的最大速度喂入数据是一个重要的要求。默认情况下，所有操作都发生在主进程中，这可能无法足够快速地从磁盘读取数据，从而创建瓶颈，导致GPU未充分利用。

配置以下参数以减少瓶颈：