内存优化概述

作者: 萨尔曼·穆罕默迪

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

内存优化组件

组件	何时使用？
模型精度	您通常希望将其保留为默认值 bfloat16。当使用 float32 时，每个模型参数使用 2 字节而不是 4 字节。
激活检查点	在内存受限且希望使用更大模型、批量大小或上下文长度时使用。请注意，这将降低训练速度。
激活卸载	与激活检查点类似，这在内存受限时可以使用，但可能会降低训练速度。这应该与激活检查点结合使用。
梯度累积	在内存受限以模拟更大批量大小时非常有用。与反向传播中的优化器不兼容。当您至少能容纳一个样本而不发生内存溢出，但无法容纳足够多时使用它。
低精度优化器	在需要减小优化器状态大小时使用此选项。这在训练大型模型并使用带动量的优化器（如 Adam）时尤为相关。请注意，低精度优化器可能会降低训练的稳定性或准确性。
将优化器步骤融合到反向传播中	当您有较大的梯度并且可以容纳足够大的批次大小时使用，因为这与 gradient_accumulation_steps 不兼容。
将优化器/梯度状态卸载到 CPU	将优化器状态和（可选的）梯度卸载到 CPU，并在 CPU 上执行优化器步骤。这可以显著降低 GPU 内存占用，但会增加 CPU 内存消耗并降低训练速度。仅在其它技术不足时优先使用此方法。
低秩自适应（LoRA）	当您希望显著减少可训练参数的数量、在训练期间节省梯度和优化器内存，并大幅加快训练速度时。这可能会降低训练精度。
量化低秩自适应（QLoRA）	在训练大型模型时，量化将节省 1.5 字节 *（模型参数数量），但可能会以牺牲部分训练速度和精度为代价。
权重分解的低秩自适应（DoRA）	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.

此设置在内存受限的情况下非常有用，尤其是当受限于较大的批次大小或更长的上下文长度时。 然而，这些内存节省是以训练速度（即每秒 token 数）为代价的， 在大多数情况下，由于这种激活重计算，训练速度会显著变慢。

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

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

激活卸载

这里发生了什么？

您可能刚刚阅读了关于激活检查点的介绍！与检查点类似，卸载是一种内存效率优化技术，它通过将激活值临时移至 CPU 并在反向传播需要时将其取回，从而节省 GPU 显存。

请参阅 PyTorch autograd 钩子教程，了解更多关于如何通过 torch.autograd.graph.saved_tensors_hooks() 实现的详细信息。

此设置在处理较大的批次大小或较长的上下文长度时特别有用，尤其是在内存受限的情况下。 当然，将张量从 GPU 移动到 CPU 再移回需要时间和资源，但 torchtune 的实现使用了多个 CUDA 流（如果可用的话），以便将额外的通信与计算重叠，从而隐藏额外的运行时间。由于通信工作量会根据要卸载的张量数量和大小而变化，我们不建议在没有启用 激活检查点 的情况下使用此功能，因为只有被检查点化的张量会被卸载。

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

要启用激活卸载功能，请在我们的LoRA微调单设备配方中使用enable_activation_offloading配置条目或标志，例如enable_activation_offloading=True。要允许使用流，请确保您使用的PyTorch版本等于或晚于当前版本。

梯度累积

这里发生了什么？

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

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

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

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

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

注意

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

优化器

低精度优化器

这里发生了什么？

除了在训练过程中通过 降低模型和优化器的精度 来减少计算开销外，我们还可以进一步降低优化器状态的精度。 我们的所有配方都支持来自 torchao 库的低精度优化器。 对于单设备配方，我们还支持 bitsandbytes。

一个好的起点可能是 torchao.prototype.low_bit_optim.AdamW8bit 和 bitsandbytes.optim.PagedAdamW8bit 优化器。 两者通过量化优化器状态字典来减少内存。分页优化器在 GPU 内存不足时也会卸载到 CPU 上。实际上， 你可以从 bnb 的 PagedAdamW8bit 中获得更高的内存节省，但从 torchao 的 AdamW8bit 中可以获得更高的训练速度。

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

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

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

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

这里发生了什么？

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

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

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

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

在torchtune中，你可以通过使用optimizer_in_bwd标志启用此功能。当使用带有大量参数的模型和状态优化器时，该功能效果最佳，并且当你不需要使用梯度累积时效果更佳。在微调LoRA配方时，你不会看到显著的影响，因为在这种情况下被更新的参数数量较小。

将优化器/梯度状态卸载到CPU

这里发生了什么？

我们上面提到的优化器状态概念——由有状态优化器使用的内存，用于维护梯度统计信息的状态，以及模型梯度——是在执行模型反向传播时用于存储梯度的张量。在我们的单设备配方中，我们通过 CPUOffloadOptimizer 支持使用 CPU 卸载功能，该功能来自 torchao。

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

如果在单设备上进行微调，另一个选项是使用 bitsandbytes 中提到的 PagedAdamW8bit，如 上方 所述，当 GPU 不足时，它将仅卸载到 CPU。

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

要在配方中使用此优化器，请将配置中的 optimizer 键设置为 torchao.prototype.low_bit_optim.CPUOffloadOptimizer，这将使用带有 fused=True 作为基础优化器的 torch.optim.AdamW 优化器。例如，要使用此优化器将优化器状态和梯度卸载到 CPU：

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

或在代码中直接使用，从而允许您更改基础优化器：

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
)

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

• 当使用优化器CPU卸载时，CPU优化器步骤通常是瓶颈所在。为了最小化减速，建议（1）使用完整的 bf16 训练，以便参数、梯度和优化器状态都在 bf16 中；以及（2）在每个优化器步骤中给GPU更多的工作量以分摊卸载时间（例如，使用激活检查点的大批量大小，或梯度累积）。

• 当 offload_gradients=True 时，梯度累积应始终设置为 1，因为每次反向传播时 GPU 上的梯度都会被清除。

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

• 此优化器仅支持单设备配方。要在分布式配方中使用CPU卸载，请改用 fsdp_cpu_offload=True。有关更多详细信息，请参阅 torch.distributed.fsdp.FullyShardedDataParallel，并查看 FSDP1 与 FSDP2 以了解它们的区别。

参数高效的微调 (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","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_dim 和 lora_rank << out_dim - 模型中任意线性层的维度。具体来说，lora_rank 以线性方式减少存储的梯度数量， 从 in_dim * out_dim 减少到 lora_rank * (in_dim + out_dim)。通常情况下，我们有 lora_rank in [8, 256]。

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

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

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

在考虑使用QLoRA来减少内存使用时，值得注意的是，QLoRA比LoRA慢，并且如果你微调的模型较小，可能不值得这样做。具体来说，QLoRA可以节省大约 1.5 字节 *（模型参数的数量）*。此外，尽管 QLoRA 对模型进行了量化，但它通过在模型前向传播过程中将量化参数提升到原始更高精度的数据类型，从而最小化了精度下降——这种提升可能会对训练速度产生负面影响。 我们 QLoRA 教程中的 相关部分 展示了如何使用 torch.compile 来解决这一问题，从而加快训练速度。

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

您可以使用 QLoRA 与我们的任何 LoRA 配方进行微调，即带有 lora_ 前缀的配方，例如 lora_finetune_single_device。这些配方利用了 QLoRA 启用的模型构建器，我们支持所有模型，并且还使用了 qlora_ 前缀，例如： torchtune.models.llama3.llama3_8b() 模型有一个对应的 torchtune.models.llama3.qlora_llama3_8b()。 我们旨在提供一套全面的配置，以便您能够快速开始使用 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 技术，它通过进一步分解预训练权重为两个组件来构建：幅度和方向。幅度组件是一个标量向量，用于调整比例；方向组件对应于原始的 LoRA 分解，并更新权重的方向。

由于增加了幅度参数，DoRA 会给 LoRA 训练带来轻微开销，但已证明其能提升 LoRA 的性能，尤其是在低秩情况下。

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

与LoRA和QLoRA类似，您可以使用我们的任何LoRA配方进行微调。我们使用相同的模型构建器来构建LoRA模型 和DoRA模型，因此您可以将任何模型构建器的lora_版本与use_dora=True一起使用。例如，要使用DoRA微调 torchtune.models.llama3.llama3_8b()，您需要使用torchtune.models.llama3.lora_llama3_8b()并设置use_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

注意

在内部，我们通过添加 DoRALinear 模块启用了 DoRA， 当 use_dora=True 时，我们会将其替换为 LoRALinear。