内存高效训练

🤗 PEFT 使微调在参数上更高效，但并不会自动提升内存效率。本概览汇总了降低训练内存用量的技巧，并提供了详细指南的链接。

始终考虑选择更小的基础模型、更小的批处理大小或更短的序列长度来降低内存使用量。

训练内存概览

让我们剖析训练内存的分布，以便推理潜在的应对措施。我们将以一个使用 Adam 优化器训练的大型语言模型为例。在进行全量微调时，以下部分将占用内存：

基础模型参数：内存消耗高度取决于所选的 dtype。每个参数的位数越少（例如 float16 为 16 位），占用的内存就越少。一个 1B 模型在 float16（每个参数 16 位/2 字节）下大约需要 1e9 × 2 byte = 1.863 GiB 的内存。
所有可训练的基础模型参数 ×3，用于梯度（1×）和 Adam 优化器状态（2×），因此需要 5.59GB 内存
层与层之间中间激活值的内存，这部分难以预测，主要取决于所使用的计算数据类型以及序列长度 / 批大小。

使用更小的基础模型或更小的计算数据类型会降低所有内存占用点，而使用更短的序列或更小的批次主要影响梯度和激活内存。采用 PEFT 方法会减少可训练参数的数量，从而显著降低梯度和优化器状态，节省大量内存。

并非所有 PEFT 方法都同等地易于构建，某些方法在内存效率方面更容易实现。如果内存预算有限，建议查看 PEFT 方法对比套件，并按最大加速器内存使用量进行筛选。各方法的平均加速器内存使用量可能相当，但并非所有方法都能随激活值和序列长度同等地扩展；有些方法比其他方法更容易出现内存峰值。

在针对大型层（如语言建模头或嵌入层）以微调特定 token 时，请考虑使用可训练 token。

量化是降低基础模型内存消耗的最佳方式之一，并且根据所采用的量化方法，还可以减少激活内存的占用。由于 PEFT 方法仅占模型总参数的一小部分，PEFT 默认使用比基础模型更高的精度。这还可能带来一种效果：适配器可以缓解量化方法带来的一部分质量损失。请阅读PEFT 量化指南。

我们训练的模型由矩阵乘法、求和和赋值等操作组成，每个操作都会产生新的结果，并因此需要占用内存。如果这些中间结果不再需要，我们可以将这些操作融合起来以节省内存。这只是 torch.compile 可以为您执行的众多优化之一，因此请查阅 PEFT torch.compile 指南。

你可以通过仅保存每 n 层之间的梯度并在运行时计算其余梯度，来用计算换取内存。有关更多信息，请参阅 Transformers 的梯度检查点文档。

如果未使用 Diffusers 或 Transformers，您可能需要根据所使用的训练框架自行实现梯度检查点逻辑。

在训练语言模型（或分类任务）时，使用 NLLLoss 非常常见。你需要分配一个大小为 batch × sequence × vocabulary 的矩阵。对于特别长的序列或词表，这会迅速变得开销巨大。

使用 TRL 时，您可以选择使用 Liger kernel integration 或 Chunked NLLLoss。后者会将序列拆分为大小为 256 的块，以保持最大内存占用恒定。

NLL vs. Chunked NLL comparison

如果默认的块大小不适合您的设置，请查阅原始 TRL PR以获取有关如何调整块大小的更多信息。