梯度检测点和梯度累计

sunghosts

已于 2023-10-16 09:51:16 修改

阅读量457

点赞数 2

分类专栏： NLP 文章标签： python

于 2023-09-28 15:42:53 首次发布

NLP 专栏收录该内容

37 篇文章 3 订阅

订阅专栏

概述

随着模型变得越来越大，GPU 内存不足而获取内存的情况CUDA: out of memory (OOM) error变得越来越普遍。
在本文中，我们将讨论一些通过梯度技巧提高训练过程效率并优化使用 GPU 内存的方法。

梯度累积

梯度累积是一种用于在训练大型模型或处理大批量数据时克服内存限制的技术。通常，在反向传播过程中，模型参数的梯度会在每批数据之后计算和更新。然而，在梯度累积中，不是在每个批次之后更新参数，而是在执行参数更新之前在多个批次上累积梯度。通过减少存储梯度的内存需求，可以实现更加内存高效的训练过程。
例如，如果梯度累积设置为累积四个批次的梯度，则对这些批次的模型参数的梯度进行求和，并在累积后应用一次参数更新。这通过执行更少的更新来减少内存占用，并支持使用更大的批量或更复杂的模型进行训练。
梯度累积是深度学习中用于增加训练期间有效批量大小的技术。通常，神经网络的权重根据单批训练数据计算的梯度进行更新。然而，对于较大的模型或数据集，批量大小可能会受到 GPU 内存容量的限制，导致矢量化导致收敛时间显着延长。
如下图（源）所示，梯度累积将样本批次（用于训练神经网络）分成几个按顺序运行的小批次。简而言之，梯度累积背后的想法是在几个小批量中迭代累积梯度。
一旦我们通过上述过程积累了足够的梯度，我们就运行模型的优化步骤（通过通常的optimizer.step()）来增加总体批量大小。
下面的代码示例（来源）显示了模型如何受到梯度累积的积极影响。

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

trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()
print_summary(result)

> BEFORE
Time: 57.82
Samples/second: 8.86
GPU memory: 14949 MB

> AFTER
Time: 66.03
Samples/second: 7.75
GPU memory: 8681 MB

梯度累积可能会导致收敛速度变慢和训练时间更长，因为在进行更新之前梯度是在多个小批量中累积的。然而，在内存有限且需要更大的有效批量大小的情况下，它可能是一种有用的技术（特别是在对比学习中，由于大训练批量中增加了多样性，较大的批量大小会导致更好的学习）。
下面的代码有助于说明梯度累积背后的基本思想。在其中，我们训练一个num_iterations迭代循环，并在每次迭代中，accumulation_step在更新权重之前处理小批量。
在每次迭代期间，使用分别计算每个小批量的梯度compute_gradients()。然后将每个小批量的梯度累积在cumulative_gradients变量中。在处理了accumulation_steps个小批量之后，累积的梯度被用来更新权重update_weights()。

for i in range(num_iterations):
    accumulated_gradients = 0
    for j in range(accumulation_steps):
        batch = next(training_batch)
        gradients = compute_gradients(batch)
        accumulated_gradients += gradients
    update_weights(accumulated_gradients)

梯度累计和增大batch有区别吗

梯度累积（gradient accumulation）和增大批次大小（batch size）在某种程度上有相似的效果，但它们在实现和结果上有所不同。

梯度累积是指在多个小批次（mini-batch）上计算梯度，然后将这些梯度累积起来，在一定数量的小批次之后再更新模型参数。这样做的目的是在内存消耗较小的情况下，模拟使用较大批次大小的效果。梯度累积通常在显存受限的GPU上训练大型神经网络时使用。

增大批次大小意味着每次更新模型参数时使用更多的数据。这可能会加快模型的收敛速度，因为每次迭代都利用了更多的数据，梯度估计也更准确。但这会增加内存消耗，因为每个批次的数据都需要同时加载到内存中。此外，过大的批次大小可能导致模型过拟合，因为每次更新时梯度估计过于准确，模型可能会对噪声数据敏感。

总结一下，梯度累积和增大批次大小在一定程度上有相似的效果，都可以提高模型训练的稳定性和收敛速度。但梯度累积的主要优势在于减小内存消耗，而增大批次大小可能导致更快的收敛。具体选择哪种方法取决于你的硬件条件和模型需求。

梯度检查点

梯度检查点是一种用于在反向传播期间权衡内存使用量和计算时间的技术。在深度神经网络中，反向传播需要存储中间激活，以便在反向传播期间计算梯度。然而，对于具有大量层或有限内存的模型，存储所有中间激活可能会占用大量内存。
梯度检查点通过在反向传播期间有选择地重新计算中间激活的子集来解决此问题。不是存储所有激活，而是仅缓存其中的一个子集，通常是计算梯度所需的子集。剩余的中间激活在向后传递期间即时重新计算。通过重新计算而不是存储所有中间激活，可以减少内存使用，但代价是增加计算时间。
检查点的中间激活的具体选择取决于模型的内存要求和计算权衡。通过策略性地设置检查点激活，梯度检查点可以实现更高效的内存训练，从而可以使用更大的模型或减少深度学习任务中的内存瓶颈。
梯度检查点有助于减少训练反向传播阶段的内存需求，特别是在具有大量层或参数的模型中。
为了计算后向传递期间的梯度，通常会保存前向传递的所有激活。这会产生很大的内存开销。
梯度检查点仅存储其中的一个子集，而不是存储前向传递期间的所有中间激活。在向后传递过程中，缺失的中间激活会即时重新计算，从而减少训练期间所需的内存量。
或者，人们可以忘记前向传递期间的所有激活，并在后向传递期间根据需要重新计算它们。然而，这会增加大量的计算开销并减慢训练速度。
这种权衡允许使用更大的模型或批量大小，否则由于内存限制而无法实现。
您可以考虑通过两种方式进行梯度检查点：
总之，梯度累积通过在更新模型参数之前累积多个批次的梯度来解决内存限制，而梯度检查点有选择地重新计算中间激活的子集，以减少反向传播期间的内存使用。这两种技术都提供了在深度学习训练中优化内存和计算资源的方法。
下面的代码（源代码），添加了梯度检查点和梯度累积，我们可以看到节省了一些内存，但训练时间变慢了。正如HuggingFace提到的，一个好的经验法则是梯度检查点会使训练速度减慢 20%。

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

trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()
print_summary(result)

> BEFORE
Time: 66.03
Samples/second: 7.75
GPU memory: 8681 MB

> AFTER
Time: 85.47
Samples/second: 5.99
GPU memory occupied: 6775 MB.