尽管 Transformers 在自然语言处理领域很成功了, 由于大量参数,即使在现代图形处理单元 (GPU) 上训练它们或在生产中部署它们仍然具有挑战性。训练或推断如此大的模型, 我们很可能会遇见 out of memory (OOM) 或者处理数据时间过长。
尽管如此,有很多方法可以避免此类问题,因此本文的主要贡献是描述和展示如何在训练和推理脚本中应用所提供的方法。首先,我们文章会介绍一些基本方法比如梯度累积 (Gradient Accumulation), 冻结(Freezing),自动混合精度训练( Automatic Mixed Precision,)8位优化器( 8-bit Optimizers)和 Gradient Checkpointing, 然后介绍NLp领域特殊的优化方法比如 Dynamic Padding, Uniform Dynamic Padding, and Fast Tokenizers.
我们开始吧!
!pip uninstall -q -y transformers
Gradient Accumulation 背后的想法非常简单——模拟更大的批量。有时为了更好地收敛或提高性能,需要使用大批量大小,但是,它通常需要大量内存。这种问题的一种可能的解决方案是使用较小的批大小,但是,一方面,小批大小会导致训练或推理时间增加,另一方面,梯度下降算法对批量大小,并可能导致不稳定的收敛和性能下降。相反,我们可以运行乘法步骤(累积步骤)并累积(计算平均)梯度一定数量的累积步骤,然后当我们有足够的计算梯度时执行优化步骤。

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)
# backward pass
loss.backward()
# perform optimization step
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
optimizer.step()
model.zero_grad()
steps=len(loader)
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 and step==steps:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
optimizer.step()
model.zero_grad()
冻结是通过切换模型某些层中的计算梯度来加速训练和降低内存利用率的有效方法,几乎不会损失最终质量。
深度学习中的一个众所周知的事实是,低层学习输入数据模式,同时顶层学习特定于目标任务的高级特征。当使用某种优化算法(例如 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)
自动混合精度(AMP)是另一种在不损失最终质量的情况下减少内存消耗和训练时间的非常简单的方法,在 文章 'Mixed Precision Training' 中进行了介绍NVIDIA 和百度研究人员在 2017 年发表的论文。该方法背后的关键思想是使用较低的精度来将模型的梯度和参数保存在内存中,即建议的方法不是使用全精度(例如 float32),而是使用半精度(例如float16) 用于将张量保存在内存中。然而,当以较低的精度计算梯度时,一些值可能太小以至于它们被视为零,这种现象称为overflow 。为了防止“溢出”,原论文的作者提出了一种梯度缩放方法。
PyTorch 为使用自动混合精度提供了一个具有必要功能(从降低精度到梯度缩放)的包,称为 torch.cuda.amp 。自动混合精度作为上下文管理器实现,因此可以轻松插入到训练和推理脚本中。

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)
# backward pass
loss.backward()
# perform optimization step
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
optimizer.step()
model.zero_grad()
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()
8 位优化器的想法类似于自动混合精度,其中模型的参数和梯度保持在较低的精度,但 8 位优化器还额外将优化器的状态也保持在较低的精度。作者(Meta Research)在原始论文 8-bit Optimizers via Block-wise Quantization 中详细描述了 8 位优化器,并表明8 位优化器会显着降低内存利用率并略微加快训练速度。此外,作者研究了不同超参数设置的影响,并表明 8 位优化器对学习率、β 和权重衰减参数的不同选择是稳定的,而不会损失性能或损害收敛性。因此,作者为 8 位优化器提供了一个舒适的高级库,称为bitsandbytes 。

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)
# selecting parameters, which requires gradients
model_parameters=filter(lambda parameter: parameter.requires_grad, model.parameters())
# initializing optimizer
optimizer=torch.optim.AdamW(params=model_parameters, lr=2e-5, weight_decay=0.0)
print(f"32-bit Optimizer:\
\
{optimizer}")
换个方法!
!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:\
\
{bnb_optimizer}")
有时甚至使用小批量和其他优化技术,例如梯度累积、冻结或自动精确训练,我们仍然会耗尽内存,尤其是在模型足够大的情况下。为解决这个问题而提出的强大解决方案之一是 Gradient Checkpointing,它首先在 2016 年的 Training Deep Nets With Sublinear Memory Cost 论文中引入。作者证明了梯度检查点可以显著降低内存利用率,从 O(n) 到 O(sqrt{n}),其中 n 是模型中的层数。这种方法允许在单个 GPU 上训练大型模型,或者提供更多内存来增加批量大小以实现更好更快的收敛。

Gradient Checkpoint 背后的想法是在小块中计算梯度,同时在前向和反向传播过程中从内存中删除不必要的梯度,从而降低内存利用率,尽管这种方法需要更多的计算步骤来重现整个反向传播图。
https://miro.medium.com/max/1082/0*s7U1QDfSXuVd1LrFPyTorch 框架通过 torch.utils.checkpoint.checkpoint 和 torch.utils.checkpoint.checkpoint_sequential 函数从盒子中提供梯度检查点。
“具体来说,在正向传递中,函数将以 torch.no_grad() 方式运行,即不存储中间激活。相反,正向传递保存输入元组和函数参数。在反向传递中,保存的输入和检索函数,并再次计算函数的前向传递,现在跟踪中间激活,然后使用这些激活值计算梯度。”
此外,HuggingFace Transformers 也支持 Gradient Checkpoint。梯度检查点可以通过 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}")
HuggingFace Transformers 提供了两种类型的 Tokenizer:Base 和 Fast。它们之间的主要区别在于 Fast Tokenizers 是在 Rust 上编写的,因为 Python 在循环中非常慢,这在tokenization过程中是必要的。这是一种让我们在tokenization过程中获得额外加速的重要方式。 Tokenizer 的类型可以通过 HuggingFace Transformers API 在 from_pretrained 方法中轻松更改; transformers.AutoTokenizer 实例通过设置 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:\
\
{tokenizer}", end="\
"*3)
# initializing Fast version of Tokenizer
fast_tokenizer=AutoTokenizer.from_pretrained(model_path, use_fast=True)
print(f"Fast version Tokenizer:\
\
{fast_tokenizer}")
通常,模型是用一批输入来训练的,并且批中的每个输入都必须具有固定的大小,即批必须是矩阵的表示。固定大小通常是根据数据集中的长度分布、特征数量或其他因素相对选择的。在 NLP 任务中,输入大小称为文本长度,称为最大长度。不幸的是,不同的文本有不同的长度,因此为了处理这种情况,研究人员提出了填充标记和截断。当最大长度小于输入文本的长度时,应用截断,因此删除了一些标记(通常是持续的)。填充标记是特殊标记,当输入文本的长度小于最大长度时添加到输入文本的末尾,另外值得注意的是,填充标记不应该包括在某些任务中计算损失(例如 Masked Language Modeling 或命名实体识别)。

- Performance and Scalability: How To Fit a Bigger Model and Train It Faster
- Speeding up Transformer w/ Optimization Strategies
- Things you can try to speed up training speed and preventing memory shortage if you are using transformers.
- 8-bit Adam and other memory optimizations
- Fitting larger networks into memory.
- Optimization approaches for Transformers