【Kaggle】如何有效避免OOM(out of memory)和漫长的炼丹过程

本文介绍一些避免transformers的OOM以及训练等流程太漫长的方法，主要参考了kaggle notebook Optimization approaches for Transformers | Kaggle，其中梯度累积Gradient Accumulation，冻结Freezing已经在之前的博客中介绍过，本文会依次介绍混合精度训练Automatic Mixed Precision, 8-bit Optimizers, and 梯度检查点Gradient Checkpointing, 然后介绍一些NLP专用的方法，比如Dynamic Padding, Uniform Dynamic Padding, and Fast Tokenizers.

Automatic Mixed Precision

作用：不损失最终质量的情况下减少内存消耗和训练时间
关键思想：是使用较低的精度将模型的梯度和参数保持在memory中，即不是使用全精度 (例如float32)，而是使用半精度 (例如float16) 将张量保持在memory中。但是，当以较低的精度计算梯度时，某些值可能很小，以至于它们被视为零，这种现象称为 “overflow”。为了防止 “overflow溢出”，原始论文的作者提出了一种梯度缩放方法。

PyTorch提供了一个具有必要功能 (从降低精度到梯度缩放) 的软件包，用于使用自动混合精度，称为torch.cuda.amp。自动混合精度可以轻松地将其插入训练和推理代码中。

Vanilla training loop

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

Training loop with Automatic Mixed Precision

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-bit Optimizers

8位优化器的思想类似于自动混合精度，其中模型的参数和梯度保持在较低的精度，但8位优化器还将优化器的状态保持在较低的精度。https://arxiv.org/abs/2110.02861作者表明8位优化器显著降低了内存利用率，略微加快了训练速度。此外，作者研究了不同超参数设置的影响，并表明8位优化器对不同的学习速率、beta和权重衰减参数的选择是稳定的，不会损失性能或损害收敛性。因此，作者为8位优化器提供了一个高级库，称为bitsandbytes。

Initializing optimizer via PyTorch API

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:\n\n{optimizer}")

32-bit Optimizer:

AdamW (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    eps: 1e-08
    lr: 2e-05
    maximize: False
    weight_decay: 0.0
)

Initializing optimizer via bitsandbytes API

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

8-bit Optimizer:

AdamW (
Parameter Group 0
    betas: (0.9, 0.999)
    eps: 1e-08
    lr: 2e-05
    weight_decay: 0.0
)

Gradient Checkpointing

有时，即使使用小批量和其他优化技术，例如梯度累积、冻结或自动精度训练，我们仍然可能耗尽内存，尤其是在模型足够大的情况下。作者证明了梯度检查点可以显著地将内存利用率从\(O（n）\)降低到\(O(\sqrt{n})\)，其中n是模型中的层数。这种方法实现了在单个GPU上训练大型模型，或提供更多内存以增加批处理大小，从而更好更快地收敛。

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

pytorch提供了torch.utils.checkpoint.checkpoint 和 torch.utils.checkpoint.checkpoint_sequential 函数来实现梯度检查点。

"Specifically, in the forward pass, function will run in torch.no_grad() manner, i.e., not storing the intermediate activations. Instead, the forward pass saves the inputs tuple and the function parameter. In the backwards pass, the saved inputs and function is retrieved, and the forward pass is computed on function again, now tracking the intermediate activations, and then the gradients are calculated using these activation values."

另外，huggingface同样支持梯度检查点，可以对PreTrainedModel instance使用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}")

Gradient Checkpointing: True

Fast Tokenizers

base和fast tokenizer的区别：fast是在rust编写的，因为python在循环中非常慢，fast可以让我们在tokenize时获得额外的加速。下图是tokenize工作的原理示意，Tokenizer类型可以通过更改 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}")

Base version Tokenizer:

PreTrainedTokenizer(name_or_path='microsoft/deberta-v3-base', vocab_size=128000, model_max_len=1000000000000000019884624838656, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'bos_token': '[CLS]', 'eos_token': '[SEP]', 'unk_token': '[UNK]', 'sep_token': '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token': '[MASK]'})


Fast version Tokenizer:

PreTrainedTokenizerFast(name_or_path='microsoft/deberta-v3-base', vocab_size=128000, model_max_len=1000000000000000019884624838656, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'bos_token': '[CLS]', 'eos_token': '[SEP]', 'unk_token': '[UNK]', 'sep_token': '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token': '[MASK]'})

Dynamic Padding

即对输入的mini batch动态进行padding，将batch的输入填充到该batch的最大输入长度，可以将训练速度提高35%甚至50%，注意，pad token不应包括在某些任务（比如MLM和NER）的损失计算过程中。

Uniform Dynamic Padding

这是基于动态填充的方法，其思想是预先按文本的相应长度对文本进行排序，在训练或推理期间比动态填充需要更少的计算。但不建议在训练期间使用统一的动态填充，因为训练意味着输入的shuffle。