如何基于PyTorch来优化大模型训练的内存（显存）使用：8种方法总结

大模型虽然效果很好，但是对资源的消耗却非常高。更麻烦的其实不是训练过程慢，而是峰值内存（显存）的消耗直接决定了我们的硬件是否可以来针对大模型进行训练。最近LightningAI官方总结了使用Fabric降低大模型训练内存的方法。但是，它也适用于其它场景。因此，本文总结一下相关的方法。

<center>![](https://www.datalearner.com/resources/blog_images/090f7076-620f-4f48-a477-728c797ac5c6.png)</center>
<center></center>

作者先实现了一个视觉模型ViT-L-16，并使用全精度微调方法做了一次微调。然后开始使用描述的8个方法一步一步优化，占用的峰值显存大小从最开始26.84GB到最低只需要2GB左右显存即可完成一样的微调。同时，精度损失非常小（3%左右）。这些方法对于做实践的童鞋来说十分重要。建议大家好好阅读理解~~

不过需要注意的是，这些方法不涉及最新的类似QLoRA这种才提出的优化方法。

另外，本文使用的是开源库Fabric作为示例，这是一种将PyTorch高级功能简化成几行代码的开源库，本质都是PyTorch，所以不影响大家学习。

[TOC]

#### 1、混合精度训练（Mixed-Precision Training）
混合精度训练是一种用于训练深度神经网络的技术，旨在提高训练速度和效率。在混合精度训练中，将不同的数值精度用于网络中的不同计算部分，以充分利用现代图形处理器（GPU）的计算能力。

传统上，神经网络中的参数和激活值使用单精度浮点数（32位）进行计算。然而，使用更低精度的浮点数，如半精度浮点数（16位），可以显著减少内存占用和计算需求。混合精度训练利用了这种观察结果，将网络中的一部分计算转换为半精度浮点数。

<center>![](https://www.datalearner.com/resources/blog_images/85eb2c0a-2c4d-43c7-af83-94bfb65d7886.png)</center>
<center></center>

混合精度训练一般分为以下几个步骤：

- **前向传播**：网络的输入通过一系列的计算层，生成预测输出。在混合精度训练中，一些计算层使用单精度浮点数进行计算，而另一些计算层使用半精度浮点数进行计算。

- **反向传播**：计算预测输出与实际输出之间的误差，并将误差通过网络进行反向传播以更新网络参数。在反向传播过程中，使用单精度浮点数计算梯度，以便准确地更新网络参数。

- **梯度缩放**：由于半精度浮点数的范围较小，梯度可能会溢出或丢失精度。为了解决这个问题，通常会对梯度进行缩放，使其适应半精度的范围。缩放后的梯度将用于更新网络参数。

混合精度训练的主要优势在于它可以显著提高训练速度和减少内存消耗。通过使用半精度浮点数进行计算，可以加快计算速度，并在存储和传输数据时减少内存需求。然而，混合精度训练也可能带来一些挑战，如数值不稳定性和精度损失。因此，适当的梯度缩放和数值稳定性技巧是保证混合精度训练有效性的关键。

总之，混合精度训练是一种利用不同精度的浮点数来加速和优化深度神经网络训练过程的技术。它结合了单精度和半精度浮点数的优点，提供了更高效的计算。

下图是使用Fabric自动混合精度（也就是说不需要你去选择哪些权重使用半精度）代码示例：
```python
# 正常精度
fabric = Fabric(accelerator="cuda", devices=1)

# 自动混合精度
fabric = Fabric(accelerator="cuda", devices=1, precision="16-mixed")
```

下图是测试效果：

<center>![](https://www.datalearner.com/resources/blog_images/3402023a-fe38-4603-a64f-cbdc4f50328e.png)</center>
<center></center>

可以看到，使用混合精度`16-mixed`之后，峰值显存从26.84GB降低到了18.21GB，但是准确性几乎没有变化。同时，训练时间也下降了6倍，从之前的17.88分钟降低到只需要3.45分钟！效果惊人。

#### 2、低精度训练（Lower-Precision Training）
低精度训练，也被称为降低精度训练或减少精度训练，是一种在深度学习中使用较低精度数据类型（如半精度16位甚至更低）来训练神经网络的技术。

混合精度训练是同时使用单精度浮点数和半精度浮点数，但是低精度训练通常只是用较低精度的浮点数。相比较混合精度训练，低精度训练可能会引入较大的精度损失，别是在计算梯度时。这可能会导致模型的准确性下降，需要使用技术来减轻精度损失带来的影响。但是，低精度训练可以比混合精度训练进一步减少内存消耗，并提高计算速度，

低精度训练的过程通常包括以下几个步骤：

- **参数转换**：在训练开始之前，将神经网络的参数从传统的单精度浮点数转换为较低精度（如半精度）。

- **前向和反向传播**：输入数据通过网络，使用降低精度的参数和激活值进行前向和反向传播。在前向和反向传播过程中，使用较低精度的数据类型进行计算。

- **梯度累积和更新**：在反向传播过程中计算的梯度被累积并用于更新降低精度的参数。更新步骤确保参数收敛到最优值。

低精度训练也面临一些挑战。降低精度可能导致数值精度损失，从而降低模型的准确性。为了解决这个问题，可以采用梯度缩放、精度特定调整和混合精度优化等技术。

总之，低精度训练是一种利用较低精度数据类型（如半精度）来训练神经网络的技术。它具有减少内存使用、加快计算速度和提高能源效率等优点。然而，需要认真考虑并采用适当的技术来应对数值精度损失带来的挑战。

基于Fabric的低精度训练代码如下：
```python
fabric = Fabric(accelerator="cuda", precision="16-true")
```

不过，低精度f16的数值范围是-65504到65504，那么原始的数值可能会丢失精度产生NaN值，此时一般建议使用bf16，主要区别在于指数位数和小数位数的差别，如下图：

<center>![](https://www.datalearner.com/resources/blog_images/651f1131-979f-409b-9517-a295978793c7.png)</center>
<center></center>
简单说就是在底层的数值表示中，浮点数都是16位二进制表示，第一位一般是表示正数还是负数，第二部分是指数位，第三部分是尾数位。bf16增加了一个指数位，降低了一个尾数位，进而可以表示更多范围的值，但是牺牲了一点精度。

全精度训练、自动混合精度训练和低精度训练结果对比如下：

<center>![](https://www.datalearner.com/resources/blog_images/4f6b3e70-2ac3-4213-b88d-24b0cabb5180.png)</center>
<center></center>

可以看到，低精度训练的峰值显存降到了13.82GB，准确率几乎没有变化。

#### 3、降低训练批处理大小（Reducing the Batchsize）
批处理大小（batch size）是指在训练神经网络模型时，同时输入到模型中进行前向传播和反向传播的样本数量。降低批处理大小可以显著降低内存和显存的使用。

不过，这也会带来一些坏处：

- **梯度估计的噪声增加**：较小的批处理大小会导致每个批次中样本数量的减少。这可能会引入更多的噪声，因为每个小批次的样本可能无法很好地代表整个训练数据集。噪声可能会导致训练过程中的不稳定性，并使模型难以收敛到最佳性能。

- **训练收敛速度变慢**：较小的批处理大小可能导致训练过程的收敛速度变慢。较大的批处理大小通常可以提供更稳定和准确的梯度估计，从而加快收敛速度。通过使用较小的批处理大小，可能需要更多的训练迭代才能达到相同的性能水平。

- **泛化能力可能受影响**：较小的批处理大小可能会影响模型的泛化能力。较大的批处理大小通常有助于模型从更全面和多样化的数据中学习，有助于提高模型的泛化能力。通过使用较小的批处理大小，模型可能更容易过度拟合训练数据，而在新的未见过的数据上表现较差。

但是，如果你的资源不够，这是不得不做的事情。那么在LightningAI的测试中，结果如下：

<center>![](https://www.datalearner.com/resources/blog_images/6e543143-c242-4484-b409-58f1b4c8c8a2.png)</center>
<center></center>

可以看到，批处理大小降低到原始数据1/3之后，峰值显存已经降到了5.69GB了，而准确性似乎影响不大。不过，这里可能在其它模型训练中不一样。大家还是要注意观察。

#### 4、使用梯度累积创建微批次（Using Gradient Accumulation to Create Microbatches）
使用梯度累积创建微批次是一种训练神经网络时常用的技术。当可用的GPU内存无法容纳所需的批次大小时，梯度累积可以帮助解决这个问题。

在传统的批次训练中，模型的权重**在每个批次之后都会被更新**一次。然而，当批次大小很大时，可能会导致GPU内存不足，无法一次性处理整个批次。这就限制了模型的训练能力。

**使用梯度累积，可以将大批次拆分为多个小批次，并在每个小批次上计算梯度。**然后，这些小批次的梯度会被累积（通常是求和或平均），而不是立即更新模型的权重。这样，在多个迭代中逐渐累积梯度，直到达到目标的“虚拟”批次大小。

一旦累积的梯度达到了目标的批次大小，就会使用这些累积的梯度来更新模型的权重。这种方式使得模型在计算资源有限的情况下，仍然可以使用较大的批次大小进行训练。

需要注意的是，梯度累积只会增加运行时间，而不会影响模型的性能。虽然每个小批次的梯度相对较小，但它们在累积过程中仍然能够提供有效的梯度信息，从而保持了模型的训练质量。

总结来说，使用梯度累积创建微批次的技术允许在有限的GPU内存下进行训练，通过将大批次分解为小批次并逐渐累积梯度，从而提高了模型的训练效率和能力。

这部分需要改动一下训练过程的代码：

<center>![](https://www.datalearner.com/resources/blog_images/2c65b703-d0d3-471d-8e4d-ebf033a4e637.png)</center>
<center></center>

结果如下：
<center>![](https://www.datalearner.com/resources/blog_images/3ca716eb-a2c6-429a-8c3f-afe26be15b41.png)</center>
<center></center>

作者批次大小为16，累积4次计算一次，也就是实际批处理大小为4，显存已经降低到了4GB左右，精度依然没变。不过缺点就是训练时间相比之前增加到12.91分钟了。

#### 5、使用Leaner Optimizer
"Leaner Optimizer"是指一种优化算法或技术，旨在改进神经网络训练过程的效率和效果。Leaner Optimizer的主要目标是减少计算和内存需求，同时保持或提升模型的性能。

传统的优化算法如随机梯度下降（Stochastic Gradient Descent，SGD）或其变种是基于从整个训练数据集或一个小批次计算得到的梯度来更新模型参数的。这些算法在大规模数据集或参数数量众多的复杂模型上计算代价昂贵。

Leaner Optimizer采用各种策略来减轻这种计算负担。一种常见的方法是使用动量或自适应学习率等技术，加速收敛并更高效地利用可用资源。例如，Adam、RMSprop或Adagrad等算法根据参数的历史梯度动态调整学习率，从而实现更快的收敛和改善的训练效率。

但是，Adam优化器本身还有额外的参数（为每个模型维护一个均值和方差），这些参数增加了存储和计算的开销。而SGD是一种无状态优化器，参数量更少。

结合这两个优点之后的Learner Optimizer技术则是这里需要的降低内存的好方法。

所以，这里的方法是将Adam与SGD进行替换，并引入余弦衰减学习率调度器来弥补这一点，从而实现更好的收敛性。

```python
# 此前的方法
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)

# 新方法
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

num_steps = NUM_EPOCHS * len(train_loader)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps)
```

这样的方法会进一步降低显存消耗：

<center>![](https://www.datalearner.com/resources/blog_images/66885de3-5118-472e-b0fd-526e5c00de54.png)</center>
<center></center>

此时，峰值显存已经下降2GB左右，精度损失3%。

#### 6、在目标设备上以所需精度创建模型
在PyTorch中，通常会在CPU设备上创建模型，然后将其转移到目标设备并转换为所需的精度。

当我们在CPU上创建模型时，通常使用默认的浮点数精度（例如32位浮点数）来表示模型参数和计算中间结果。这个完整精度表示可以提供最高的精确度，但它也会占用更多的内存和计算资源。

在模型从CPU设备转移到目标设备之前，需要将模型参数和计算结果从CPU内存传输到目标设备的内存中。这涉及到数据的复制和转换。如果我们直接在CPU上使用完整精度进行计算和存储，那么在模型传输过程中，中间模型表示会占用更多的内存和带宽。

在这种特定情况下，即模型加载过程中，前向传递期间的内存使用峰值大于模型在完整精度表示中的大小。这意味着在模型加载期间，可能会出现比完整精度模型更大的内存占用。显然，这也有很大的浪费。

#### 7、分布式训练和Tensor分片（Distributed Training and Tensor Sharding）
如果你有多个GPU，那么就可以使用分布式训练。不过，由于本文主要是针对降低显存使用，因此这里使用的是完全分片数据并行（Fully Sharded Data Parallelism，FSDP）。它利用数据并行和张量并行将大型权重矩阵切片分配到多个设备上。

完全分片数据并行（Fully Sharded Data Parallelism，FSDP）是一种高级的分布式多GPU策略，用于在多个设备上分片大型权重矩阵。它结合了数据并行和张量并行的概念，以实现更高效的训练过程和更大规模的模型。

在FSDP中，权重矩阵被分成多个较小的片段，每个片段分配给不同的GPU设备。这种分片操作可以减少每个设备上的内存占用，使得更大的模型可以在有限的设备内存上进行训练。

具体而言，FSDP将数据并行和张量并行相结合。数据并行将输入数据分布到不同的GPU上进行并行计算，而张量并行则将模型的权重矩阵分成多个片段，并在每个GPU上进行并行计算。通过将数据并行和张量并行结合起来，FSDP能够充分利用多个GPU设备的计算能力，并以更高效的方式进行训练。

总之，FSDP是一种高级的分布式多GPU训练策略，通过将大型权重矩阵切片分布到多个设备上，以实现更高效的训练和更大规模的模型。它结合了数据并行和张量并行的思想，提供了内存节省和训练加速的优势。

由于上面采用的数据已经很少，因此这里作者比较了FSDP策略与原始全精度训练的显存使用。

<center>![](https://www.datalearner.com/resources/blog_images/5f6deede-eade-436c-afe2-032c3bb5c7dc.png)</center>
<center></center>

可以看到，采用FSDP策略之后显存从26.84GB降低到了6GB左右。使用fabric去执行FSDP策略代码只需要一行：

```python
fabric = Fabric(accelerator="cuda", devices=4, strategy="fsdp")
```

#### 8、参数卸载（Parameter Offloading）
参数卸载（Parameter Offloading）是一种优化训练过程的技术，特别是在分布式深度学习中使用多个计算设备进行模型训练时。它的目标是减少数据传输和存储开销，提高训练效率。

在传统的分布式训练中，每个计算设备都需要拥有完整的模型参数副本。这意味着在每个设备上进行的计算需要将所有参数复制到该设备的内存中，并且在每次更新参数时，还需要将这些更新传输回主服务器或其他设备，这会引起大量的数据传输和存储开销。

参数卸载的思想是将一部分模型参数从计算设备卸载到主服务器或其他专用设备上。这样，每个计算设备只需要保存和处理部分参数，而不是全部参数。计算设备只需传输和存储部分参数，从而降低了通信和内存开销。

通常，参数卸载涉及到两个主要步骤：

初始参数分发：在训练开始之前，将完整的模型参数从主服务器或专用设备分发到各个计算设备。
参数更新收集：在每次参数更新之后，将更新后的参数从计算设备收集回主服务器或专用设备，以便进行全局参数的更新和同步。
通过参数卸载，每个计算设备只需要处理部分参数，从而减少了计算和通信开销，提高了训练效率。这种技术在大规模分布式训练中尤为有用，可以加速训练过程并减少资源的消耗。

原文完整内容参考：https://lightning.ai/pages/community/tutorial/pytorch-memory-vit-llm/

如何基于PyTorch来优化大模型训练的内存（显存）使用：8种方法总结

欢迎大家关注DataLearner官方微信，接受最新的AI技术推送

相关博客

最热博客