内存优化概述

作者: Salman Mohammadi

torchtune 提供了一系列即插即用的内存优化组件，让您可以非常灵活地根据硬件 调整（tune）我们的配方。本页简要概述了这些组件及其使用方法。为了方便起见，我们在下表中总结了这些组件

内存优化组件

组件	何时使用？
模型精度	通常建议保留默认的 bfloat16。使用 bfloat16 时，每个模型参数占用 2 字节，而不是使用 float32 时的 4 字节。
激活检查点	当内存受限且想使用更大的模型、批大小或上下文长度时使用。请注意，这会降低训练速度。
激活卸载	与激活检查点类似，当内存受限时可以使用此功能，但这可能会降低训练速度。这应该与激活检查点一起使用。
梯度累积	内存受限时有助于模拟更大的批大小。与反向传播中的优化器不兼容。当您至少可以放入一个样本而不会出现 OOM（内存不足）但又放不下足够多的样本时使用此功能。
低精度优化器	当您想减小优化器状态的大小时使用。这在训练大型模型和使用带有动量的优化器（如 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 文档中相关章节很好地解释了这一概念。引用如下

激活检查点是一种用计算换取内存的技术。它不是在反向传播期间将反向传播所需的张量一直保存在内存中直到用于梯度计算，而是在检查点区域的前向计算中省略保存反向传播所需的张量，并在反向传播过程中重新计算它们。

此设置对于内存受限的情况很有帮助，特别是由于批大小较大或上下文长度较长时。然而，节省内存的代价是训练速度（即每秒处理的 token 数），并且在大多数情况下，由于重新计算激活，训练速度会显著降低。

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

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

激活卸载

这是怎么回事？

您可能刚刚阅读了有关激活检查点的内容！与检查点类似，卸载是一种内存效率技术，它允许通过将激活暂时移动到 CPU 并在反向传播需要时将其带回，从而节省 GPU 显存。

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

当内存受限时，此设置对于更大的批大小或更长的上下文长度特别有帮助。虽然当然需要运行时和资源才能将张量从 GPU 移动到 CPU 再移回，但 torchtune 中的实现使用了多个 CUDA 流（如果可用），以便将额外的通信与计算重叠以隐藏额外的运行时。由于通信工作负载取决于被卸载的张量数量和大小而变化，因此我们不建议使用此功能，除非同时启用了激活检查点，在这种情况下，只会卸载被检查点的张量。

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

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

梯度累积

这是怎么回事？

梯度累积允许您通过在多个批次上累积梯度，然后在优化器更新模型参数之前，模拟更大的批次大小。具体来说，使用梯度累积时，用于梯度更新的样本总数是

total_batch_size = batch_size * gradient_accumulation_steps

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

注意

对于 torchtune 中使用“步”的其他组件，例如指标记录或学习 率 调度器，一个“步”被计为模型参数的一次更新，而不是数据的一次模型前向传播。假设 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 内存不足，Paged 优化器也会卸载到 CPU。实际上，您会发现 bnb 的 PagedAdamW8bit 能节省更多内存，而 torchao 的 AdamW8bit 训练速度更快。

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

要在配方中使用此功能，请确保已安装 torchao (pip install torchao) 或 bitsandbytes (pip install bitsandbytes)。然后，使用torchtune 命令行界面启用低精度优化器

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

这是怎么回事？

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

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

如果在单设备上进行微调，另一种选择是使用 bitsandbytes 中的 PagedAdamW8bit，如上文所述，它仅在 GPU 内存不足时才会卸载到 CPU。

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

要在您的配方中使用此优化器，请将配置中的 optimizer 键设置为 torchao.prototype.low_bit_optim.CPUOffloadOptimizer，它将使用 torch.optim.AdamW 优化器并设置 fused=True 作为基础优化器。例如，要使用此优化器将优化器状态和梯度都卸载到 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 上保留参数副本并预分配梯度内存。因此，预计您的内存使用量会增加到模型大小的 4 倍。

• 此优化器仅支持单设备配方。要在分布式配方中使用 CPU 卸载，请改用 fsdp_cpu_offload=True。有关更多详细信息，请参阅torch.distributed.fsdp.FullyShardedDataParallel，并参阅FSDP1 vs 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 在 [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 通过加速训练来解决这个问题。

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

您可以使用我们任何 LoRA 配方进行 QLoRA 微调，即带有 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 是另一种 PEFT 技术，它在 LoRA 的基础上进一步将预训练权重分解为两个分量：幅度和方向。幅度分量是一个标量向量，用于调整尺度，而方向分量对应于原始 LoRA 分解并更新权重的方向。

由于增加了幅度参数，DoRA 会给 LoRA 训练带来少量开销，但已证明它可以提高 LoRA 的性能，尤其是在低秩时。

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

与 LoRA 和 QLoRA 非常相似，您可以使用我们的任何 LoRA 配方进行 DoRA 微调。我们用于 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 替换为该模块。