内存优化概述

作者: Salman Mohammadi

torchtune 附带一系列即插即用的内存优化组件，为您提供极大的灵活性来 tune 我们的配方以适应您的硬件。本页简要介绍这些组件及其使用方法。为了方便起见，我们在下表中总结了这些组件

内存优化组件

组件	何时使用？
模型精度	您通常希望将其保留为默认值 bfloat16。与使用 float32 时的 4 字节相比，它每个模型参数使用 2 字节。
激活检查点	当您受到内存约束并希望使用更大的模型、批大小或上下文长度时使用。请注意，它会降低训练速度。
激活卸载	与激活检查点类似，当内存受限时可以使用它，但可能会降低训练速度。**应该** 与激活检查点一起使用。
梯度累积	在内存受限时有助于模拟更大的批大小。与反向传播中的优化器不兼容。当您已经可以容纳至少一个样本而不会 OOM，但数量不足时使用它。
低精度优化器	当您想要减小优化器状态的大小时使用。这在使用带有动量的优化器（如 Adam）训练大型模型时相关。请注意，较低精度的优化器可能会降低训练的稳定性/准确性。
将优化器步骤融合到反向传播中	当您有较大的梯度并且可以容纳足够大的批大小时使用它，因为它与 gradient_accumulation_steps 不兼容。
将优化器/梯度状态卸载到 CPU	将优化器状态和（可选）梯度卸载到 CPU，并在 CPU 上执行优化器步骤。这可以显著减少 GPU 内存使用量，但会牺牲 CPU RAM 和训练速度。仅当其他技术不足时才优先使用它。
低秩自适应 (LoRA)	当您想要显著减少可训练参数的数量，节省训练期间的梯度和优化器内存，并显著加快训练速度时使用。这可能会降低训练精度
量化低秩自适应 (QLoRA)	当您训练大型模型时使用，因为量化将节省 1.5 字节 * (模型参数数量)，但可能会牺牲一些训练速度和准确性。
权重分解低秩自适应 (DoRA)	LoRA 的变体，可能会提高模型性能，但会稍微增加内存消耗。

注意

目前，本教程侧重于单设备优化。请随时关注我们的页面更新，了解分布式微调的最新内存优化功能。

模型精度

这里发生了什么？

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

注意

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

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

• bfloat16，称为“半精度”，每个模型和优化器参数使用 2 字节 - 有效地减少了 fp32 一半的内存，并且还提高了训练速度。通常，如果您的硬件支持使用 bfloat16 进行训练，我们建议使用它 - 这是我们配方的默认设置。

注意

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

听起来很棒！我该如何使用它？

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

激活检查点

这里发生了什么？

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

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

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

听起来很棒！我该如何使用它？

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

激活卸载

这里发生了什么？

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

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

当您受到内存限制时，此设置对于较大的批大小或较长的上下文长度尤其有帮助。虽然当然将张量从 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 内存不足，分页优化器还将卸载到 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

这里发生了什么？

我们在上面提到了优化器状态的概念 - 有状态优化器用于维护梯度统计状态的内存，以及模型梯度 - 用于存储我们在执行模型反向传播时使用的梯度的张量。我们通过来自 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 上保留参数副本并预分配梯度内存来工作。因此，预计您的 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 的配置即可。

所有其余的 LoRA 参数对于 QLoRA 保持不变 - 查看上面关于 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 非常相似，您可以使用 DoRA 以及我们任何 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。