显存优化技术→LLM 训练

文章源自@Vadim Irtlach大佬在kaggle的开源notebook

  1. 梯度累积（Gradient Accumulation）
  2. 冻结（Freezing）
  3. 自动混合精度（Automatic Mixed Precision）
  4. 8位优化器（8-bit Optimizers）
  5. 梯度检查点（Gradient Checkpointing）
  6. Parameter-Efficient Fine-Tuning（PEFT）
  7. Torch FSDP+CPU offload

0 估算模型所需的RAM

首先，我们需要了解如何根据参数量估计模型大致所需的 RAM，这在实践中有很重要的参考意义。我们需要通过估算设置 batch_size，设置模型精度，选择微调方法和参数分布方法等。

接下来，我们用 LLaMA-6B 模型为例估算其大致需要的内存。首先考虑精度对所需内存的影响：

  • fp32 精度，一个参数需要 32 bits, 4 bytes.
  • fp16 精度，一个参数需要 16 bits, 2 bytes.
  • int8 精度，一个参数需要 8 bits, 1 byte.

其次，考虑模型需要的 RAM 大致分三个部分：

  • 模型参数
  • 梯度
  • 优化器参数

模型参数：等于参数量*每个参数所需内存。

  • 对于 fp32，LLaMA-6B 需要 6B*4 bytes = 24GB内存
  • 对于 int8，LLaMA-6B 需要 6B*1 byte = 6GB

梯度：同上↑，等于参数量*每个梯度参数所需内存。

优化器参数：不同的优化器所储存的参数量不同。
对于常用的 AdamW 来说，需要储存两倍的模型参数（用来储存一阶和二阶momentum）。

  • fp32 的 LLaMA-6B，AdamW 需要 6B*8 bytes = 48 GB
  • int8 的 LLaMA-6B，AdamW 需要 6B*2 bytes = 12 GB

除此之外，CUDA kernel 也会占据一些 RAM，大概 1.3GB 左右，查看方式如下。

综上，int8 精度的 LLaMA-6B 模型大致需要 6GB+6GB+12GB+1.3GB = 25.3GB 左右。

再根据LLaMA的架构（hidden_size = 4096, intermediate_size =11008, num_hidden_layers = 32, context_length = 2048）计算中间变量内存。

每个 instance 需要：

所以一张 A100（80GB RAM）大概可以在 int8 精度；batch_size = 50 的设定下进行全参数训练。

查看消费级显卡的内存和算力：
2023 GPU Benchmark and Graphics Card Comparison Chart

梯度累积

梯度累积背后的想法非常简单，就是为了模拟更大的批量（batch）。有时，为了更好地收敛或提高性能，需要使用大批量进行训练，但是，这通常需要更大的显存。这个问题的一种可能的解决方案是使用较小的批量，但是，一方面，小批量训练会增加训练和推理时间，另一方面，梯度下降算法对批量大小的选择非常敏感，小批量可能会导致不稳定的收敛和性能降低。所以，我们可以先执行几次前向传播和反向传播，使得梯度进行累积，当我们有足够的计算梯度时，再对参数进行优化，从而利用小显存，模拟大批量的效果，并且训练时间也不会大幅增加。

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
steps = len(loader)

# perform validation loop each `validation_steps` training steps!
validation_steps = int(validation_steps * gradient_accumulation_steps)

for step, batch in enumerate(loader, 1):

    # prepare inputs and targets for the model and loss function respectively.

    # forward pass
    outputs = model(inputs)

    # computing loss
    loss = loss_fn(outputs, targets)

    # accumulating gradients over steps
    if gradient_accumulation_steps > 1:
        loss = loss / gradient_accumulation_steps

    # backward pass
    loss.backward()

        # perform optimization step after certain number of accumulating steps and at the end of epoch
    if step % gradient_accumulation_steps == 0 or step == steps:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        optimizer.step()
        model.zero_grad()

            # perform validation loop
    if step % validation_steps == 0:
        validation_loop()

冻结

冻结是一种非常有效的方法，通过取消计算模型某些层中的梯度计算（如embedding层，bert的前几层），可以大大加快训练速度并且降低了显存占用，而且几乎不会损失模型的性能。
深度学习中的一个众所周知的事实是，网络的底层学习输入数据的通用特征，而网络顶层学习目标任务特定的高级特征，所以在对预训练模型进行微调时，一般网络底层的参数都不怎么需要变，这些都是通用的知识，需要学习的是顶层的那些参数，当使用某种优化算法（如SGD、AdamW或RMSprop）执行优化步骤时，网络的底层的梯度就都很小，因此参数几乎保持不变，这也被称为梯度消失，因此，与其花费大量的时间和算力来计算底层这些“无用”梯度，并对此类梯度很小的参数进行优化，不如直接冻结它们，直接不计算梯度也不进行优化。
PyTorch为关闭梯度计算提供了一个舒适的API，可以通过torch.Tensor的属性requires_ grad设置。

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def freeze(module):
    """
    Freezes module's parameters.
    """
    for parameter in module.parameters():
        parameter.requires_grad = False

def get_freezed_parameters(module):
    """
    Returns names of freezed parameters of the given module.
    """

    freezed_parameters = []
    for name, parameter in module.named_parameters():
        if not parameter.requires_grad:
            freezed_parameters.append(name)

    return freezed_parameters
import torch
from transformers import AutoConfig, AutoModel


# initializing model
model_path = "microsoft/deberta-v3-base"
config = AutoConfig.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, config=config)


# freezing embeddings and first 2 layers of encoder
freeze(model.embeddings)
freeze(model.encoder.layer[:2])

freezed_parameters = get_freezed_parameters(model)
print(f"Freezed parameters: {freezed_parameters}")

# selecting parameters, which requires gradients and initializing optimizer
model_parameters = filter(lambda parameter: parameter.requires_grad, model.parameters())
optimizer = torch.optim.AdamW(params=model_parameters, lr=2e-5, weight_decay=0.0)

自动混合精度

自动混合精度（AMP）是另一种在不损失最终质量的情况下减少显存消耗和训练时间的方法，该方法由NVIDIA和百度研究人员在2017年的“Mixed Precision Training”论文中提出。该方法背后的关键思想是使用较低的精度将模型的梯度和参数保留在内存中，即不使用全精度（float32），而是使用半精度（例如float16）将张量保存在内存中。然而，当以较低精度计算梯度时，某些值可能太小，以至于被视为零，这种现象被称为“溢出”。为了防止“溢出”，原始论文的作者提出了一种梯度缩放方法。只要处理得当，这种方式训练得到的模型在性能上几乎不会有太大损失。

代码实现

PyTorch从1.6的版本开始提供了一个包：torch.cuda.amp，具有使用自动混合精度所需的功能（从降低精度到梯度缩放），自动混合精度作为上下文管理器实现，因此可以随时随地的插入到训练和推理脚本中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from torch.cuda.amp import autocast, GradScaler


scaler = GradScaler()

for step, batch in enumerate(loader, 1):

    # prepare inputs and targets for the model and loss function respectively.

    # forward pass with `autocast` context manager
    with autocast(enabled=True):
        outputs = model(inputs)

    # computing loss
    loss = loss_fn(outputs, targets)

    # scale gradint and perform backward pass
    scaler.scale(loss).backward()

    # before gradient clipping the optimizer parameters must be unscaled.
    scaler.unscale_(optimizer)

    # perform optimization step
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

    scaler.step(optimizer)
    scaler.update()

Memory Efficient Optimizer

Memory Efficient Optimizer 的基本思想类似于自动混合精度（模型的参数和梯度使用较低的精度保存），减少优化器所维护的运行状态的内存占用，如Adafactor。
作者（Meta Research）在最初的论文“8-bit Optimizers via Block-wise Quantization”中详细介绍了8-bit Optimizers，表明8-bit Optimizers显著降低了显存占用，略微加快了训练速度。此外，作者研究了不同超参数设置的影响，表明8-bit Optimizers对不同的学习率、beta和权重衰减参数的效果是稳定的，不会降低性能或影响收敛性。因此，作者为8位优化器提供了一个高级库，叫做bitsandbytes。

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
!pip install -q bitsandbytes-cuda110
def set_embedding_parameters_bits(embeddings_path, optim_bits=32):
    """
    https://github.com/huggingface/transformers/issues/14819#issuecomment-1003427930
    """

    embedding_types = ("word", "position", "token_type")
    for embedding_type in embedding_types:
        attr_name = f"{embedding_type}_embeddings"

        if hasattr(embeddings_path, attr_name): 
            bnb.optim.GlobalOptimManager.get_instance().register_module_override(
                getattr(embeddings_path, attr_name), 'weight', {'optim_bits': optim_bits}
            )

import bitsandbytes as bnb


# selecting parameters, which requires gradients
model_parameters = filter(lambda parameter: parameter.requires_grad, model.parameters())

# initializing optimizer 
bnb_optimizer = bnb.optim.AdamW(params=model_parameters, lr=2e-5, weight_decay=0.0, optim_bits=8)
# bnb_optimizer = bnb.optim.AdamW8bit(params=model_parameters, lr=2e-5, weight_decay=0.0) # equivalent to the above line

# setting embeddings parameters
set_embedding_parameters_bits(embeddings_path=model.embeddings)

print(f"8-bit Optimizer:\n\n{bnb_optimizer}")

梯度检查点

有时候，即使用了上面的几种方法，显存可能还是不够，尤其是在模型足够大的情况下。那么梯度检查点（Gradient Checkpointing）就是压箱底的招数了，这个方法第一次在 “Training Deep Nets With Sublinear Memory Cost” ，作者表明梯度检查点可以显著降低显存利用率，从 O(n) 降低到 O(n) ，其中n是模型的层数。这种方法允许在单个GPU上训练大型模型，或者提供更多内存以增加批量大小，从而更好更快地收敛。梯度检查点背后的思想是在小数据块中计算梯度，同时在正向和反向传播过程中从内存中移除不必要的梯度，从而降低内存利用率，但是这种方法需要更多的计算步骤来再现整个反向传播图，其实就是一种用时间来换空间的方法。

梯度检查点如何在正向和反向传播过程中工作

PyTorch框架里也有梯度检查点的实现，通过这两个函数：torch.utils.checkpoint.checkpoint和torch.utils.checkpoint.checkpoint_sequential
这边引用一段torch官网对梯度检查点的介绍：
梯度检查点通过用计算换取内存来工作。检查点部分不是存储整个计算图的所有中间激活以进行反向计算，而是不保存中间激活，而是在反向过程中重新计算它们。它可以应用于模型的任何部分。 具体而言，在前向传播中，该函数将以torch.no_grad()的方式运行，即不存储中间激活。然而，前向传播保存了输入元组和函数参数。在反向传播时，检索保存的输入和函数，然后再次对函数进行前向传播，现在跟踪中间激活，然后使用这些激活值计算梯度。
此外，HuggingFace Transformers也支持梯度检查点。梯度检查点可以通过PreTrainedModel实例的gradient_checkpointing_enable方法执行，一行代码直接搞定！

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
from transformers import AutoConfig, AutoModel
# https://github.com/huggingface/transformers/issues/9919
from torch.utils.checkpoint import checkpoint

# initializing model
model_path = "microsoft/deberta-v3-base"
config = AutoConfig.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, config=config)

# gradient checkpointing
model.gradient_checkpointing_enable()
print(f"Gradient Checkpointing: {model.is_gradient_checkpointing}")

参数高效微调 PEFT

参数高效微调 Parameter-Efficient Fine-Tuning（PEFT）技术，可以仅微调少量或者额外的模型参数，并固定住大部分预训练参数，从而大大降低 LLM 的训练成本。
内容较多，单独整理成一篇技术博客，欢迎阅读 《PEFT技术简介》

卸载（Offloading）

卸载（Offloading）是将未使用的数据暂时卸载到CPU或不同的设备中，然后在需要时再将其读回，参数动态地从 GPU -> CPU, CPU -> GPU 进行转移，从而节省 GPU 内存。
这个想法的一个具体实现是ZeRO，它将参数、梯度和优化器状态分割到所有可用的硬件上，并根据实际需要再将它们具体化。

Huggingface 这篇博文解释了 ZeRO 的大致实现方法：
_https://huggingface.co/blog/zero-deepspeed-fairscale_

借助 torch 实现 FSDP，只需要将 model 用 FSDPwarp 一下；同样，cpu_offload 也只需要一行代码：
_https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/_

总结

即使在现代GPU上，优化内存和时间也是开发模型的必要步骤，因此，本文介绍了加速训练和减少transformers等大型模型内存消耗的最强大、最流行的方法。

显存优化
#转载
显存优化技术→LLM 训练
http://example.com/2023/05/19/2023-05-19-显存优化技术→LLM训练/
作者
Ning Shixian
发布于
2023年5月19日
许可协议