内存优化概述

作者： Salman Mohammadi

Torchtune 附带了许多即插即用的内存优化组件，为您提供了很大的灵活性 到我们的配方到您的硬件。本页提供了这些组件的简要词汇表以及如何使用它们。 为了方便起见，我们在下表中总结了这些组件：tune

内存优化组件

元件	何时使用？
模型精度	您通常希望将其保留为默认值。如果您因精度而难以提高训练稳定性或准确性，fp32 可能会有所帮助，但会显著增加内存使用量并降低训练速度。bfloat16
激活检查点	当您内存受限并且需要处理更大的批处理大小或更长的上下文长度时使用。请注意，这可能会减慢训练速度。
梯度累积	在内存受限时非常有用，可以模拟较大的批处理大小。通常比激活检查点更可取，以获得更好的训练速度。
低精度优化器	当您需要通过降低优化器状态中的精度来进一步减少超出使用的内存使用量时。请注意，较低精度的优化器可能会降低训练稳定性/准确性。bf16
将 Optimizer Step 融合到 Backward Pass 中	有助于在使用有状态优化器时减少内存使用量，尤其是在对具有高梯度内存使用率的大型模型进行全面微调时。这与 不兼容，因此训练可能会因模型吞吐量降低而减慢。gradient_accumulation_steps
低秩适应 （LoRA）	当您想显著减少可训练参数的数量，在训练期间节省梯度和优化器内存，并显著加快训练速度时。
量化低秩适应 （QLoRA）	当您需要比 LoRA 更多的内存节省时，可能会牺牲一些训练速度。适用于非常大的模型或有限的硬件。

注意

在当前状态下，本教程侧重于单设备优化。更新此页面后请立即签到 了解分布式微调的最新内存优化功能。

模型精度

这是怎么回事？

我们使用术语 “precision” 来指代用于表示模型和优化器参数的基础数据类型。 我们在 torchtune 中支持两种数据类型：

注意

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

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

• bfloat16，称为“半精度”，每个模型和优化器参数使用 2 个字节 - 实际上是一半 的内存，还可以提高训练速度。通常，如果您的硬件支持使用 进行训练，则 我们建议使用它 - 这是我们配方的默认设置。fp32bfloat16

注意

另一种常见的范例是“混合精度”训练：其中模型权重在 （or ） 中，optimizer 状态为 。目前，我们不支持 torchtune 中的混合精度训练。bfloat16fp16fp32

听起来不错！我该如何使用它？

只需在我们所有的配方中使用 flag 或 config 条目即可！例如，要在 中使用半精度训练， 设置。dtypebf16dtype=bf16

激活检查点

这是怎么回事？

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

激活检查点是一种用计算换取内存的技术。 而不是保持 backward 所需的张量处于活动状态，直到它们在 Checkpointed 中 backward 和 forward 计算期间的梯度计算 regions 省略了为 backward 保存张量，并在 backward pass 期间重新计算它们。

当您受内存限制时，尤其是由于较大的批处理大小或较长的上下文长度时，此设置非常有用。 然而，这些内存节省是以训练速度（即每秒令牌数）为代价的， 在大多数情况下，由于这种激活重新计算，训练可能会减慢很多。

听起来不错！我该如何使用它？

要启用激活检查点，请使用 config 条目或标志 在我们的任何示例中，例如 .enable_activation_checkpointingenable_activation_checkpointing=True

激活卸载

这是怎么回事？

您可能刚刚阅读了有关激活检查点的信息！与 checkpointing 类似，卸载是一种内存 效率技术，允许通过临时将激活移动到 CPU 并引入 GPU VRAM 来保存 GPU VRAM 他们需要时在后传中返回。

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

此设置对于较大的批处理大小或较长的上下文长度（当您受内存限制时）特别有用。 然而，这些内存节省可能是以训练速度（即每秒令牌数）为代价的，因为它需要运行时 以及将 Tensor 从 GPU 移动到 CPU 并返回的资源。torchtune 中的实现可以选择使用多个 CUDA 流，以便将额外的通信与计算重叠，以隐藏额外的 运行。由于通信工作负载是可变的，具体取决于要卸载的张量的数量和大小，因此它是 通常不会卸载每个激活。事实上，once 可以将 offloading 与 activations 结合使用 检查点，其中所有激活将在稍后向后重新计算或从 CPU 返回。offload_with_streams

听起来不错！我该如何使用它？

要启用激活卸载，请使用 config 条目或标志 在我们的 LoRa 微调单个设备配方中，例如 .要允许 使用流，请确保您使用的是高于 PyTorch 2.5.0.dev20240907 的 Torch 版本，并且 指定。enable_activation_offloadingenable_activation_offloading=Trueoffload_with_streams=True

梯度累积

这是怎么回事？

梯度累积允许您通过在多个梯度上累积梯度来模拟大批量 batch中。具体来说，使用的样本总数 对于梯度更新，当使用梯度累积时：

total_batch_size = batch_size * gradient_accumulation_steps

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

注意

对于 torchtune 中使用 “steps” 的其他组件，例如 metric logging 或 ，“step” 计为 对模型参数的单个更新，而不是对数据进行单个模型前向传递。 假设 和 。 指标将每 10 个全局步骤记录一次，这相当于每 40 个模型向前传递一次。 因此，在使用梯度累积进行训练时，度量日志记录的出现频率会降低。 进度条的更新速度可能会更慢。gradient_accumulation_steps = 4log_every_n_steps = 10

如果您使用的是我们的分布式配方之一，只需乘以设备数量即可：

total_batch_size = batch_size * gradient_accumulation_steps * num_devices

梯度累积在内存受限时特别有用。在这种情况下， 累积梯度可能比启用激活提供更好的训练速度 checkpointing，因为激活 checkpointing 减少了内存消耗，但代价是重复 计算。

听起来不错！我该如何使用它？

我们所有的微调配方都支持通过累积梯度来模拟更大的批量大小。只需设置 flag 或 config 条目。gradient_accumulation_steps

注意

将优化器步骤融合到向后通道时，梯度累积应始终设置为 1。

低精度优化器

这是怎么回事？

除了在训练期间降低模型和优化器的精度外，我们还可以进一步降低优化器状态的精度。 我们所有的单设备微调配方都支持来自 bitsandbytes 库的低精度优化器 - 一个好的起点可能是 AND Optimizers，我们已经用它测试了我们的配方。AdamW8bitPagedAdamW8bit

听起来不错！我该如何使用它？

要在配方中使用它，请确保已安装 bitsandbytes （）。然后，启用 使用 torchtune CLI 的低精度优化器：pip install bitsandbytes

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

或者直接修改配置文件：

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

将 Optimizer Step 融合到 Backward Pass 中

这是怎么回事？

有状态优化器（例如使用 momentum 的优化器）由于其稳定的收敛属性而成为现代深度学习中的默认设置。 但是，维护梯度统计信息的状态是以额外的内存使用为代价的。直接的替代方案可能是 转向无状态优化器，例如 Stochastic Gradient Descent Without Momentum，它不需要任何额外的内存使用，但可能会导致训练期间收敛更差。

我们能在这里找到一个中间地带吗？让我们考虑一种技术，该技术允许使用像 AdamW 这样的 “有状态” 优化器，而不会产生梯度统计的内存开销，也不会牺牲它们理想的收敛特性。 你可能会问，这怎么可能？通过完全删除优化器在其 .step()

要了解其工作原理，我们鼓励您通读有关此概念的相关 PyTorch 教程：如何通过将优化器步骤融合到向后传递中来节省内存。

听起来不错！我该如何使用它？

在 torchtune 中，您可以使用标志启用此功能，目前仅在我们的 单设备完整 Finetune 配方。当梯度内存特别大时，此功能效果最佳; 例如，当将有状态优化器与具有大量参数的模型一起使用时，以及当您不需要使用梯度累积时。optimizer_in_bwd

参数有效微调 （PEFT）

低秩适应 （LoRA）

这是怎么回事？

您可以阅读我们关于使用 LoRA 微调 Llama2 的教程，以了解 LoRA 的工作原理和使用方法。 简而言之，LoRA 大大减少了可训练参数的数量，从而节省了大量的梯度和优化器 训练期间的记忆。

听起来不错！我该如何使用它？

您可以使用我们任何带有前缀的配方进行微调，例如 lora_finetune_single_device。这些示例利用 支持 LoRA 的模型构建器，我们所有的模型都支持它，并且还使用前缀，例如 该模型具有相应的 . 我们的目标是提供一套全面的配置，以便您快速开始使用 LoRA 进行训练。 只需指定名称中包含任何 的配置，例如：lora_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 应用于每个变压器层中的 MLP。

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

注意

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

这些都在 flag 或 config 条目下指定，即：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_dimlora_rank << out_dimlora_rankin_dim * out_dimlora_rank * (in_dim + out_dim)lora_rank in [8, 128]

• lora_alpha: float影响 LoRA 更新的量级。Alpha 越大，基本模型权重的更新就越大 ，可能会以牺牲训练稳定性为代价，相反，较小的 alpha 可以以较慢的学习速度为代价来稳定训练。 我们为这些参数提供了默认设置，我们已经在所有模型中测试了这些设置，但我们鼓励您调整它们 添加到您的特定用例中。通常，一个联合更改 和 一起，其中 .lora_ranklora_alphalora_alpha ~= 2*lora_rank

• lora_dropout在 LoRA 层中引入 dropout 以帮助规范训练。我们所有模型都默认为 0.0。

如上所述，这些参数也在 flag 或 config 条目下指定。model

注意

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

量化低秩适应 （QLoRA）

这是怎么回事？

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

在考虑使用 QLoRA 来减少内存使用量时，值得注意的是 QLoRA 可以防止量化过程中的精度下降 通过在模型前向传递期间将量化参数向上转换到原始的更高精度数据类型 - 此向上转换可以 对训练速度造成惩罚。我们的 QLoRA 教程中的相关部分演示了通过加快训练来解决此问题的用法。torch.compile

听起来不错！我该如何使用它？

您可以将 QLoRA 与我们的任何 LoRA 配方（即带有前缀的配方）进行微调，例如 lora_finetune_single_device。这些示例利用 启用了 QLoRA 的模型构建器，我们所有的模型都支持它，并且还使用前缀，例如 该模型具有相应的 . 我们的目标是提供一套全面的配置，以便您快速开始使用 QLoRA 进行培训。 只需指定名称中包含任何 的配置，例如：lora_qlora__qlora

tune run lora_finetune_single_device --config llama3/8B_qlora_single_device

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