Transformer 训练优化

前言

 

自 BERT 出现以来，NLP 领域已经进入了大模型的时代，大模型虽然效果好，但是毕竟不是人人都有着丰富的 GPU 资源，在训练时往往就捉襟见肘，出现显存 out of memory 的问题，或者训练时间非常非常的久，因此，这篇文章主要解决的问题就是如何在 GPU 资源受限的情况下训练 transformers 库上面的大模型。

尽管 Huggingface 开源的 Transformers 在自然语言处理（NLP）任务中取得了惊人的成功，但由于里面的模型参数数量庞大，即使是使用 GPU 进行训练或者部署，也仍具有非常大的挑战性，因为用如此大的模型进行训练或推理，会很容易发生显存不足（OOM）以及训练时间过长的问题。

然而，有很多方法可以避免显存不足以及训练时间过长的方法，这篇文章的主要贡献就是介绍了这些方法的原理以及如何实现，具体包括以下几种方法：

1. 梯度累积（Gradient Accumulation）

2. 冻结（Freezing）

3. 自动混合精度（Automatic Mixed Precision）

4. 8位优化器（8-bit Optimizers）

5. 梯度检查点（Gradient Checkpointing）

6. 快速分词器（Fast Tokenizers）

7. 动态填充（Dynamic Padding）

8. 均匀动态填充（Uniform Dynamic Padding）

其中 1-5 是神经网络通用的方法，可以用在任何网络的性能优化上，6-8 是针对 NLP 领域的性能优化方法。

1 梯度累积

梯度累积背后的想法非常简单，就是为了模拟更大的批量（batch）。有时，为了更好地收敛或提高性能，需要使用大批量进行训练，但是，这通常需要更大的显存。这个问题的一种可能的解决方案是使用较小的批量，但是，一方面，小批量训练会增加训练和推理时间，另一方面，梯度下降算法对批量大小的选择非常敏感，小批量可能会导致不稳定的收敛和性能降低。

所以，我们可以先执行几次前向传播和反向传播，使得梯度进行累积，当我们有足够的计算梯度时，再对参数进行优化，从而利用小显存，模拟大批量的效果，并且训练时间也不会大幅增加。

代码实现

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

2 冻结

冻结是一种非常有效的方法，通过取消计算模型某些层中的梯度计算（如 embedding 层，bert 的前几层），可以大大加快训练速度并且降低了显存占用，而且几乎不会损失模型的性能。

深度学习中的一个众所周知的事实是，网络的底层学习输入数据的通用特征，而网络顶层学习目标任务特定的高级特征，所以在对预训练模型进行微调时，一般网络底层的参数都不怎么需要变，这些都是通用的知识，需要学习的是顶层的那些参数，当使用某种优化算法（如 SGD、AdamW 或 RMSprop）执行优化步骤时，网络的底层的梯度就都很小，因此参数几乎保持不变，这也被称为梯度消失，因此，与其花费大量的时间和算力来计算底层这些“无用”梯度，并对此类梯度很小的参数进行优化，不如直接冻结它们，直接不计算梯度也不进行优化。

PyTorch 为关闭梯度计算提供了一个舒适的 API，可以通过 torch.Tensor 的属性 requires_ grad 设置。

代码实现

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)

3 自动混合精度

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

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

代码实现

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

4 8-bit Optimizers

8-bit Optimizers 的思想类似于自动混合精度（模型的参数和梯度使用较低的精度保存），但 8-bit Optimizers 还让优化器的状态使用低精度保存。作者（Meta Research）在最初的论文 “8-bit Optimizers via Block-wise Quantization” 中详细介绍了 8-bit Optimizers，表明 8-bit Optimizers 显著降低了显存占用，略微加快了训练速度。

此外，作者研究了不同超参数设置的影响，表明 8-bit Optimizers 对不同的学习率、beta 和权重衰减参数的效果是稳定的，不会降低性能或影响收敛性。因此，作者为 8 位优化器提供了一个高级库，叫做 bitsandbytes。

代码实现

!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}")

5 梯度检查点

有时候，即使用了上面的几种方法，显存可能还是不够，尤其是在模型足够大的情况下。那么梯度检查点（Gradient Checkpointing）就是压箱底的招数了，这个方法第一次在 “Training Deep Nets With Sublinear Memory Cost”，作者表明梯度检查点可以显著降低显存利用率，从 O(n) 降低到 ，其中 n 是模型的层数。这种方法允许在单个 GPU 上训练大型模型，或者提供更多内存以增加批量大小，从而更好更快地收敛。 

梯度检查点背后的思想是在小数据块中计算梯度，同时在正向和反向传播过程中从内存中移除不必要的梯度，从而降低内存利用率，但是这种方法需要更多的计算步骤来再现整个反向传播图，其实就是一种用时间来换空间的方法。

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

PyTorch框架里也有梯度检查点的实现，通过这两个函数：torch.utils.checkpoint.checkpoint  和 torch.utils.checkpoint.checkpoint_sequential 这边引用一段 torch 官网对梯度检查点的介绍。

梯度检查点通过用计算换取内存来工作。检查点部分不是存储整个计算图的所有中间激活以进行反向计算，而是不保存中间激活，而是在反向过程中重新计算它们。它可以应用于模型的任何部分。具体而言，在前向传播中，该函数将以 torch.no_grad() 的方式运行，即不存储中间激活。然而，前向传播保存了输入元组和函数参数。在反向传播时，检索保存的输入和函数，然后再次对函数进行前向传播，现在跟踪中间激活，然后使用这些激活值计算梯度。

此外，HuggingFace Transformers 也支持梯度检查点。梯度检查点可以通过 PreTrainedModel 实例的 gradient_checkpointing_enable 方法执行，一行代码直接搞定！

代码实现

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}")

6 快速分词器

HuggingFace Transformers 提供两种类型的分词器：基本分词器和快速分词器。它们之间的主要区别在于，快速分词器是在 Rust 上编写的，因为 Python 在循环中非常慢，但在分词的时候又要用到循环。快速分词器是一种非常简单的方法，允许我们在分词的时候获得额外的加速。要使用快速分词器也很简单，只要把 transformers.AutoTokenizer 里面的 from_pretrained 方法的 use_fast 的值修改为 True 就可以了。

▲ 分词器是如何工作的 

代码实现

from transformers import AutoTokenizer

# initializing Base version of Tokenizer
model_path = "microsoft/deberta-v3-base"
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
print(f"Base version Tokenizer:\n\n{tokenizer}", end="\n"*3)

# initializing Fast version of Tokenizer
fast_tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
print(f"Fast version Tokenizer:\n\n{fast_tokenizer}")

7 动态填充

通常来说，模型是用批量数据输入训练的，批中的每个输入必须具有固定大小，即一批量的数据必须是矩阵的表示，所有批量数据的尺寸都一样。固定尺寸通常是根据数据集中的长度分布、特征数量和其他因素来选择的。在 NLP 任务中，输入大小称为文本长度，或者最大长度（max length）。

然而，不同的文本具有不同的长度，为了处理这种情况，研究人员提出了填充标记和截断。当最大长度小于输入文本的长度时，会使用截断，因此会删除一些标记。当输入文本的长度小于最大长度时，会将填充标记，比如 [PAD]，添加到输入文本的末尾，值得注意的是，填充标记不应包含在某些任务的损失计算中（例如掩蔽语言建模或命名实体识别）

▲ 固定长度填充

然而，填充标记有明显的缺点。比如在输入文本相对于选定的最大长度非常短的情况下，效率就很低，需要更多的额外内存，比如我有一条文本长度 512，然后其他文本长度都在 10 左右，那么如果将 max seq 设置为 512，就会导致很多无效计算。

为了防止额外的计算操作，研究人员提出了一种非常有效的方法，就是将批量的输入填充到这一批量的最大输入长度，如下图所示，这种方法可以将训练速度提高 35% 甚至 50%，当然这种方法加速的效果取决于批量的大小以及文本长度的分布，批量越小，加速效果越明显，文本长度分布越不均，加速效果也越好。

▲ 动态填充

8 均匀动态填充

还有一种基于动态填充的方法，叫做均匀动态填充。其思想是在分 batch 时，先按文本的长度对文本进行排序，这样同一个 batch 里面的文本长度就都差不多。这种方法非常有效，在训练或推理期间的计算量都比动态填充要来的少。但是，不建议在训练期间使用均匀动态填充，因为训练时数据最好是要 shuffer 的，但是推理时如果一次性要推理很多文本的话可以考虑这么做

▲ 均匀动态填充

9 总结

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

参考文献

[1] Performance and Scalability: How To Fit a Bigger Model and Train It Faster

https://huggingface.co/docs/transformers/performance

[2] Speeding up Transformer w/ Optimization Strategies

https://www.kaggle.com/code/rhtsingh/speeding-up-transformer-w-optimization-strategies

[3] Things you can try to speed up training speed and preventing memory shortage if you are using transformers.

https://www.kaggle.com/competitions/AI4Code/discussion/327777

[4] 8-bit Adam and other memory optimizations

https://www.kaggle.com/competitions/feedback-prize-2021/discussion/303131

[5] Fitting larger networks into memory.

https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9