内存优化概述

作者: 萨尔曼·穆罕默迪

torchtune 随附一系列即插即用的内存优化组件，这些组件为您提供了极大的灵活性，可以将我们的配方 tune 适配到您的硬件上。本页面提供了一个简要的术语表，介绍这些组件及其使用方法。为了方便起见，我们已在以下表格中总结了这些组件：

内存优化组件

组件	何时使用？
模型精度	您通常希望将其保留为默认值 bfloat16。如果您在精度方面遇到训练稳定性或准确性问题，fp32 可能会有所帮助，但会显著增加内存使用量并降低训练速度。
激活检查点	在内存受限且需要处理更大批量大小或更长上下文长度时使用。请注意，这可能会降低训练速度。
梯度累积	在内存受限的情况下，有助于模拟更大的批量大小。通常比激活检查点更受青睐，以获得更快的训练速度。
低精度优化器	当您需要通过将精度降低到 bf16 以下来进一步减少内存使用时。请注意，较低精度的优化器可能会降低训练的稳定性和准确性。
将优化器步骤融合到反向传播中	在使用状态优化器时，特别是在对大型模型进行全微调且梯度内存使用量较高时，有助于减少内存使用。这与 gradient_accumulation_steps 不兼容，因此由于模型吞吐量降低，训练可能会变慢。
低秩自适应（LoRA）	当您希望显著减少可训练参数的数量、在训练过程中节省梯度和优化器内存，并大幅加快训练速度时。
量化低秩自适应（QLoRA）	当您需要比 LoRA 更大幅度的内存节省，且可接受一定的训练速度损失时，此方案非常有用。它适用于超大规模模型或硬件资源受限的场景。

注意

目前，本教程专注于单设备优化。请稍后查看此页面，以获取分布式微调的最新内存优化功能。

模型精度

这里发生了什么？

我们使用“精度”一词来指代用于表示模型和优化器参数的底层数据类型。 torchtune 支持两种数据类型：

注意

我们建议深入阅读Sebastian Raschka的关于混合精度技术的博客文章，以更深入地了解精度和数据格式相关的概念。

• fp32，通常被称为“全精度”，每个模型和优化器参数使用4个字节。

• bfloat16，被称为“半精度”，每个模型和优化器参数使用2个字节——实际上是 fp32 的一半内存，并且还能提高训练速度。通常情况下，如果你的硬件支持使用 bfloat16 进行训练，我们建议使用它——这是我们配方的默认设置。

注意

另一种常见的范式是“混合精度”训练：模型权重为 bfloat16（或 fp16），优化器状态为 fp32。目前，torchtune 不支持混合精度训练。

太棒了！我该如何使用它？

只需在我们所有的配方中使用 dtype 标志或配置项即可！例如，在 bf16 中使用半精度训练时， 设置 dtype=bf16。

激活检查点

这里发生了什么？

PyTorch 文档中的相关部分很好地解释了这个概念。 引用如下：

Activation checkpointing is a technique that trades compute for memory. Instead of keeping tensors needed for backward alive until they are used in gradient computation during backward, forward computation in checkpointed regions omits saving tensors for backward and recomputes them during the backward pass.

这种设置在内存受限时很有帮助，尤其是由于更大的批次大小或更长的上下文长度。然而，这些内存节省是以训练速度（即每秒处理的标记数）为代价的，并且在这种激活重新计算的情况下，大多数情况下训练速度会显著变慢。

太棒了！我该如何使用它？

要启用激活检查点，可以在我们的任何配方中使用 enable_activation_checkpointing 配置条目或标志， 例如 enable_activation_checkpointing=True。

激活卸载

这里发生了什么？

你可能刚刚读到了关于激活检查点的内容！与检查点类似，卸载是一种内存效率技术，它通过暂时将激活值移动到 CPU 并在反向传播时需要时再将其移回 GPU 来节省 GPU VRAM。

有关如何通过 saved_tensors_hooks 实现此功能的更多详细信息，请参阅 PyTorch autograd 钩子教程。

这种设置在处理较大的批次大小或内存受限时较长的上下文长度时特别有用。 然而，这些内存节省可能会以训练速度（即每秒标记数）为代价，因为需要运行时间和资源将张量从 GPU 移动到 CPU 再移回。torchtune 中的实现提供了 offload_with_streams 选项，使用多个 CUDA 流来重叠额外的通信与计算，从而隐藏额外的运行时间。由于通信工作量取决于被卸载的张量的数量和大小，通常不会卸载每一个激活值。事实上，可以将卸载与激活检查点结合使用，这样所有的激活要么在反向传播中稍后重新计算，要么从 CPU 恢复回来。

太棒了！我该如何使用它？

要启用激活卸载功能，请在我们的LoRA微调单设备配方中使用enable_activation_offloading配置项或标志，例如enable_activation_offloading=True。要允许使用流，请确保您使用的PyTorch版本高于2.5.0.dev20240907，并指定offload_with_streams=True。

梯度累积

这里发生了什么？

梯度累积允许您通过在更新模型参数之前对多个批次的梯度进行累积，来模拟大批次大小。具体而言，使用梯度累积时用于梯度更新的样本总数为：

total_batch_size = batch_size * gradient_accumulation_steps

例如：使用 batch_size=1 和 gradient_accumulation_steps=32，我们得到的总批量大小为 32。

注意

对于torchtune中的其他使用“步骤”的组件，例如指标记录或 learning rate schedulers，“步骤”被计为模型参数的一次更新，而不是数据的一次模型前向传播。 假设gradient_accumulation_steps = 4和log_every_n_steps = 10。 每10个全局步骤记录一次指标，这相当于每40次模型前向传播记录一次。 因此，当使用梯度累积进行训练时，指标记录会显得更少，进度条的更新也会更慢。

如果您正在使用我们的分布式训练方案，只需乘以设备数量即可：

total_batch_size = batch_size * gradient_accumulation_steps * num_devices

梯度累积在内存受限的情况下特别有用。在这种情况下，累积梯度可能会比启用激活检查点提供更好的训练速度，因为激活检查点通过重复计算来减少内存消耗。

太棒了！我该如何使用它？

我们所有的微调配方都支持通过累积梯度来模拟更大的批量大小。只需设置 gradient_accumulation_steps 标志或配置项即可。

注意

在将优化器步骤与反向传播合并时，梯度累积应始终设置为1。

低精度优化器

这里发生了什么？

除了在训练过程中降低模型和优化器的精度外，我们还可以进一步降低优化器状态的精度。 我们所有的单设备微调配方都支持来自 bitsandbytes 库的低精度优化器——一个不错的起点可能是 AdamW8bit 和 PagedAdamW8bit 优化器，我们已经用这些优化器测试过我们的配方。

太棒了！我该如何使用它？

要在您的配方中使用此功能，请确保已安装 bitsandbytes (pip install bitsandbytes)。然后，使用 torchtune CLI 启用低精度优化器：

tune run <RECIPE> --config <CONFIG> \
optimizer=bitsandbytes.optim.PagedAdamW

或者直接通过 修改配置文件：

optimizer:
  _component_: bitsandbytes.optim.PagedAdamW
  lr: 2e-5

将优化器步骤与反向传播融合

这里发生了什么？

带有状态的优化器（例如使用动量的优化器）由于其稳定的收敛特性，已成为现代深度学习中的默认选择。 然而，维护梯度统计信息的状态会增加额外的内存使用。一个直接的替代方案可能是转向无状态优化器，例如不带动量的随机梯度下降， 这种优化器不需要任何额外的内存使用，但在训练过程中可能会导致较差的收敛效果。

我们能否在这里找到一个折中的方法呢？让我们考虑一种技术，它能够在不牺牲其理想的收敛特性的情况下，使用“状态化”的优化器（如 AdamW），同时避免梯度统计信息的内存开销。 你可能会问，这是如何实现的呢？通过 **完全移除优化器在 step() 期间存储的梯度缓冲区** 来实现。

要了解其工作原理，我们鼓励您阅读相关的 PyTorch 教程： 如何通过将优化器步骤与反向传播合并来节省内存。

太棒了！我该如何使用它？

在 torchtune 中，您可以使用 optimizer_in_bwd 标志启用此功能，该功能目前仅支持我们的 单设备全微调配方。当梯度内存特别大时，此功能效果最佳； 例如，当使用带有大量参数的模型的状态优化器时，以及当您不需要使用 梯度累积 时。

参数高效的微调 (PEFT)

低秩适应（LoRA）

这里发生了什么？

您可以阅读我们关于使用LoRA微调Llama2的教程，以了解LoRA的工作原理以及如何使用它。 简单来说，LoRA大大减少了可训练参数的数量，从而在训练过程中节省了大量的梯度和优化器内存。

太棒了！我该如何使用它？

您可以使用任何带有 lora_ 前缀的配方进行微调，例如 lora_finetune_single_device。这些配方利用了 LoRA 支持的模型构建器，我们支持所有模型，并且还使用了 lora_ 前缀，例如 torchtune.models.llama3.llama3() 模型有一个对应的 torchtune.models.llama3.lora_llama3()。 我们的目标是提供一套全面的配置，让您能够快速开始使用 LoRA 进行训练， 只需指定名称中包含 _lora 的任何配置即可，例如：

tune run lora_finetune_single_device --config llama3/8B_lora_single_device

有两组参数可以用来定制LoRA以满足你的需求。首先，控制LoRA应该应用于模型中的哪些线性层的参数：

• lora_attn_modules: List[str] 接受一个字符串列表，指定要应用 LoRA 的模型的哪些层：

  • q_proj 将 LoRA 应用于查询投影层。

  • k_proj 将 LoRA 应用于键投影层。

  • v_proj 将LoRA应用于值投影层。

  • output_proj 将LoRA应用于注意力输出投影层。

  虽然增加更多需要微调的层可能会提高模型的准确性，但这将导致内存使用量增加和训练速度降低。

• apply_lora_to_mlp: Bool 将 LoRA 应用于每个 Transformer 层中的 MLP。

• apply_lora_to_output: Bool 将 LoRA 应用于模型的最终输出投影。 这通常是到词汇空间的投影（例如在语言模型中），但其他建模任务可能有不同的投影——分类器模型会将投影到类的数量，例如

注意

使用绑定嵌入（例如Gemma和Qwen2 1.5B及0.5B）进行最终输出投影的模型不支持apply_lora_to_output。

这些都指定在model标志或配置条目下，例如：

tune run lora_finetune_single_device --config llama3/8B_lora_single_device  \
model.apply_lora_to_mlp=True \
model.lora_attn_modules=["q_proj","k_proj","v_proj"]

model:
  apply_lora_to_mlp: True
  model.lora_attn_modules: ["q_proj", "k_proj", "v_proj"]

其次，控制 LoRA 对模型影响程度的参数：

• lora_rank: int 影响 LoRA 分解的尺度，其中 lora_rank << in_dim 和 lora_rank << out_dim - 模型中任意线性层的维度。具体来说，lora_rank 以线性方式减少存储的梯度数量， 从 in_dim * out_dim 减少到 lora_rank * (in_dim + out_dim)。通常情况下，我们有 lora_rank in [8, 128]。

• lora_alpha: float 影响 LoRA 更新的幅度。较大的 alpha 值会导致基础模型权重更大的更新，这可能会以训练稳定性为代价；相反，较小的 alpha 可以稳定训练，但可能以学习速度较慢为代价。 我们为这些参数提供了默认设置，这些设置已在我们的所有模型上进行了测试，但我们鼓励您根据具体的使用场景进行调整。通常情况下，会同时调整 lora_rank 和 lora_alpha，其中 lora_alpha ~= 2*lora_rank。

• lora_dropout 在LoRA层中引入了Dropout以帮助正则化训练。我们所有模型的默认值为0.0。

如上所述，这些参数也在model标志或配置条目下指定。

注意

为了更深入地了解LoRA参数如何影响训练期间的内存使用， 请参阅我们在Llama2 LoRA教程中的相关部分。

量化低秩适配 (QLoRA)

这里发生了什么？

QLoRA 是在 LoRA 之上的一种增强，它将 LoRA 中的冻结模型参数保持在 4 位量化精度，从而减少内存使用。 这是通过作者提出的新型 4 位 NormalFloat (NF4) 数据类型实现的，该数据类型允许参数内存使用量减少 4-8 倍，同时保留模型准确性。你可以阅读我们关于 使用 QLoRA 微调 Llama2 的教程，以深入了解其工作原理。

在考虑使用QLoRA来减少内存使用时，值得注意的是，QLoRA通过在模型前向传播过程中将量化参数提升到原始的更高精度数据类型，从而防止量化过程中的精度下降——这种提升可能会对训练速度产生影响。我们QLoRA教程中的相关部分展示了如何使用torch.compile来解决这一问题，从而加快训练速度。

太棒了！我该如何使用它？

您可以使用任何带有 lora_ 前缀的 LoRA 配方（例如 lora_finetune_single_device）通过 QLoRA 进行微调。这些配方利用了支持我们所有模型的 QLoRA 模型构建器，并且也使用了 qlora_ 前缀，例如 torchtune.models.llama3.llama3_8b() 模型对应一个 torchtune.models.llama3.qlora_llama3_8b()。 我们的目标是提供一套全面的配置，让您能够快速开始使用 QLoRA 进行训练，只需指定名称中包含 _qlora 的任何配置即可，例如：

tune run lora_finetune_single_device --config llama3/8B_qlora_single_device

QLoRA 中其余所有 LoRA 参数保持不变 - 请查看上方关于 LoRA 的部分以了解如何配置。