内存优化概述

作者： Salman Mohammadi

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

内存优化组件

元件	何时使用？
模型精度	您通常希望将其保留为默认值。它为每个模型参数使用 2 个字节，而不是使用 时为 4 个字节。bfloat16float32
激活检查点	当您内存受限并希望使用更大的模型、批处理大小或上下文长度时使用。请注意，它会减慢训练速度。
激活卸载	与激活检查点类似，这可以在内存受限时使用，但可能会降低训练速度。这应该与激活检查点一起使用。
梯度累积	在内存受限时非常有用，可以模拟较大的批处理大小。与 backward 中的 optimizer 不兼容。当您已经可以在没有 OOMing 的情况下容纳至少一个样本，但又不够时，请使用它。
低精度优化器	当您想要减小优化器状态的大小时使用。这在训练大型模型和使用具有动量的优化器（如 Adam）时是相关的。请注意，较低精度的优化器可能会降低训练稳定性/准确性。
将 Optimizer Step 融合到 Backward Pass 中	当您具有较大的渐变并且可以容纳足够大的批量大小时，请使用它，因为这与 不兼容。gradient_accumulation_steps
将 Optimizer/Gradient 状态卸载到 CPU	将优化器状态和（可选）梯度卸载到 CPU，并在 CPU 上执行优化器步骤。这可用于显著降低 GPU 内存使用量，但代价是 CPU、RAM 和训练速度。仅当其他技术不够时，才优先使用它。
低秩适应 （LoRA）	当您想显著减少可训练参数的数量，在训练期间节省梯度和优化器内存，并显著加快训练速度时。这可能会降低训练准确性
量化低秩适应 （QLoRA）	当您训练大型模型时，由于量化将节省 1.5 字节 * （# of model parameters），但可能会牺牲一些训练速度和准确性。
权重分解低秩适应 （DoRA）	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 期间重新计算它们。

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

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

要启用激活检查点，请使用 。enable_activation_checkpointing=True

激活卸载

这是怎么回事？

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

有关如何通过 .torch.autograd.graph.saved_tensors_hooks()

此设置对于较大的批处理大小或较长的上下文长度（当您受内存限制时）特别有用。 当然，将张量从 GPU 移动到 CPU 并返回需要运行时和资源，但 torchtune 使用多个 CUDA 流（如果可用）来重叠与计算的额外通信 以隐藏额外的运行时。由于通信工作负载是可变的，具体取决于 tensor 的数量和大小 offloaded 的 API API 的 API API 的 Zip 的 Zip 的 Beta，除非还启用了 Activation Checkpointing，否则我们不建议使用它，在这种情况下，只有 checkpoint Tensor 将被卸载。

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

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

梯度累积

这是怎么回事？

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

total_batch_size = batch_size * gradient_accumulation_steps

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

注意

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

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

total_batch_size = batch_size * gradient_accumulation_steps * num_devices

当 GPU 中至少可以容纳一个样本时，梯度累积特别有用。在这种情况下，人为地将批次增加 与使用其他以内存换取速度的内存优化技术（如激活检查点）相比，累积梯度可能会为您提供更快的训练速度。

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

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

注意

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

优化器

低精度优化器

这是怎么回事？

除了在训练期间降低模型和优化器的精度外，我们还可以进一步降低优化器状态的精度。 我们所有的 recipes 都支持 torchao 库中的低精度优化器。 对于单个设备配方，我们还支持 bitsandbytes。

一个好的起点可能是 和 优化器。 两者都通过量化优化器 state dict 来减少内存。如果没有足够的可用 GPU 内存，分页优化器也会卸载到 CPU。在实践中， 您可以预期 bnb 的 PagedAdamW8bit 会节省更高的内存，但 torchao 的 AdamW8bit 会带来更高的训练速度。torchao.prototype.low_bit_optim.AdamW8bitbitsandbytes.optim.PagedAdamW8bit

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

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

tune run <RECIPE> --config <CONFIG> \
optimizer=torchao.prototype.low_bit_optim.AdamW8bit

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

或者直接修改配置文件：

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

将 Optimizer Step 融合到 Backward Pass 中

这是怎么回事？

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

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

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

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

在 torchtune 中，您可以使用标志启用此功能。此功能在使用有状态优化器时效果最佳 对于具有大量参数的模型，并且您不需要使用 Gradient Accumulation。 在微调 LoRA 配方时，您不会看到有意义的影响，因为在这种情况下，要更新的参数数量很少。optimizer_in_bwd

将 Optimizer/Gradient 状态卸载到 CPU

这是怎么回事？

我们在上面提到了优化器状态的概念 - 有状态优化器用来维护梯度统计状态的内存，以及 模型梯度 - 当我们执行模型向后传递时用于存储梯度的张量。我们支持在单设备配方中使用 CPU 卸载 通过 CPUOffloadOptimizer 从 .torchao

这个优化器可以包装任何基本优化器，并通过保持优化器状态并在 CPU 上执行优化器步骤来工作，从而减少 GPU 内存使用情况（按优化器状态的大小）。此外，我们还可以使用 offload_gradients=True 将梯度卸载到 CPU。

如果在单个设备上进行微调，另一种选择是使用上面提到的 from bitsandbytes，它只会卸载到 CPU 当没有足够的 GPU 可用时。PagedAdamW8bit

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

要在配方中使用此优化器，请将配置中的键设置为 ，该 将使用 优化器 作为基本优化器。例如，要使用此优化器卸载 优化器状态和梯度到 CPU：optimizertorchao.prototype.low_bit_optim.CPUOffloadOptimizerfused=True

tune run <RECIPE> --config <CONFIG> \
optimizer=optimizer=torchao.prototype.low_bit_optim.CPUOffloadOptimizer \
optimizer.offload_gradients=True \
lr=4e-5

或者直接修改配置文件：

optimizer:
  _component_: torchao.prototype.low_bit_optim.CPUOffloadOptimizer
  offload_gradients: True
  # additional key-word arguments can be passed to torch.optim.AdamW
  lr: 4e-5

或者直接在代码中使用它，这样您就可以更改 Base Optimizer：

from torchao.prototype.low_bit_optim import CPUOffloadOptimizer
from torch.optim import Adam

optimizer = CPUOffloadOptimizer(
    model.parameters(), # your model here
    Adam,
    lr=1e-5,
    fused=True
)

CPUOffloadOptimizer 页面中的一些有用提示：torchao

• 使用 optimizer CPU 卸载时，CPU optimizer 步骤通常是瓶颈。为了最大限度地减少减速，建议 （1） 使用完全训练，以便参数、梯度和优化器状态处于 ;（2） 为每个优化器步骤提供 GPU 更多工作，以分摊卸载时间（例如，激活检查点、梯度累积的更大批量）。bf16bf16

• 当 时，渐变累积应始终设置为 1，因为每次向后传递都会在 GPU 上清除渐变。offload_gradients=True

• 此优化器的工作原理是保留参数的副本并在 CPU 上预先分配梯度内存。因此，预计您的 RAM 使用量将增加 4 倍的模型大小。

• 此优化器仅支持单设备配方。要在分布式配方中使用 CPU 卸载，请改用。有关更多详细信息，请参阅 FSDP1 与 FSDP2 以了解它们有何不同。fsdp_cpu_offload=True

参数有效微调 （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","output_proj"]

model:
  _component_: torchtune.models.llama3.lora_llama3_8b
  apply_lora_to_mlp: True
  model.lora_attn_modules: ["q_proj", "k_proj", "v_proj","output_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, 256]

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

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

如上所述，这些参数也在 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","output_proj"] \
model.lora_rank=32 \
model.lora_alpha=64

model:
  _component_: torchtune.models.llama3.lora_llama3_8b
  apply_lora_to_mlp: True
  lora_attn_modules: ["q_proj", "k_proj", "v_proj","output_proj"]
  lora_rank: 32
  lora_alpha: 64

注意

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

量化低秩适应 （QLoRA）

这是怎么回事？

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

在考虑使用 QLoRA 来减少内存使用量时，值得注意的是 QLoRA 比 LoRA 慢，如果 您正在微调的模型很小。从数字上看，QLoRA 大约节省 1.5 字节 * （# 模型参数）。此外，尽管 QLoRA 量化了模型， 它通过在模型前向传递期间将量化参数向上转换到原始的更高精度数据类型来最大限度地减少精度下降 - 这种向上转换可能会对训练速度造成不利影响。 我们的 QLoRA 教程中的相关部分演示了通过加快训练来解决此问题的用法。torch.compile

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

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

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

要从命令行进行配置，请执行以下操作：

tune run lora_finetune_single_device --config llama3/8B_qlora_single_device \
model.apply_lora_to_mlp=True \
model.lora_attn_modules=["q_proj","k_proj","v_proj"] \
model.lora_rank=32 \
model.lora_alpha=64

或者，通过修改配置：

model:
  _component_: torchtune.models.qlora_llama3_8b
  apply_lora_to_mlp: True
  lora_attn_modules: ["q_proj", "k_proj", "v_proj"]
  lora_rank: 32
  lora_alpha: 64

权重分解低秩适应 （DoRA）

这是怎么回事？

DoRA 是另一种建立在 LoRA 之上的 PEFT 技术，它通过以下方式构建 进一步将预训练的权重分解为两个部分：magnitude 和 direction。magnitude 分量 是调整比例的标量向量，而 direction 分量对应于原始 LoRA 分解， 更新权重的方向。

由于添加了 magnitude 参数，DoRA 为 LoRA 训练增加了少量开销，但已证明 提高 LoRA 的性能，尤其是在低等级时。

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

与 LoRA 和 QLoRA 非常相似，您可以使用 DoRA 对我们的任何 LoRA 配方进行微调。我们对 LoRA 使用相同的模型生成器 就像我们对 DoRA 所做的那样，因此您可以将任何模型生成器的版本与 .例如，要使用 DoRA 进行微调，您可以与 ：lora_use_dora=Trueuse_dora=True

tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
model.use_dora=True

model:
  _component_: torchtune.models.lora_llama3_8b
  use_dora: True

由于 DoRA 扩展了 LoRA，因此自定义 LoRA 的参数是相同的。您还可以像量化低秩适应 （QLoRA） 一样，通过使用来收割来量化基本模型权重 节省更多内存！quantize=True

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.lora_rank=16 \
model.lora_alpha=32 \
model.use_dora=True \
model.quantize_base=True

model:
  _component_: torchtune.models.lora_llama3_8b
  apply_lora_to_mlp: True
  lora_attn_modules: ["q_proj", "k_proj", "v_proj"]
  lora_rank: 16
  lora_alpha: 32
  use_dora: True
  quantize_base: True

注意

在后台，我们通过添加模块启用了 DoRA，我们交换了该模块 out for when .use_dora=True