内存优化概述¶

torchtune 提供了许多即插即用的内存优化组件，这些组件为您提供了很大的灵活性，可以根据您的硬件调整我们的配方。本页简要介绍了这些组件的词汇表以及您可能如何使用它们。为了简化操作，我们在下表中总结了这些组件：

Memory optimization components¶
组件	何时使用？
模型精度	您通常希望将其保留为默认的`bfloat16`。与使用`float32`相比，它每个模型参数使用2字节而不是4字节。
激活检查点	在内存受限且希望使用更大的模型、批量大小或上下文长度时使用。请注意，这会减慢训练速度。
激活卸载	类似于激活检查点，这可以在内存受限时使用，但可能会降低训练速度。这应该与激活检查点一起使用。
梯度累积	在内存受限时有助于模拟更大的批量大小。与反向优化器不兼容。当您已经可以在不出现内存不足的情况下至少容纳一个样本，但不足以容纳更多样本时使用它。
低精度优化器	当您希望减少优化器状态的大小时使用。这在训练大型模型并使用带有动量的优化器（如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 VRAM。

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

此设置对于较大的批量大小或较长的上下文长度特别有帮助，尤其是在内存受限的情况下。当然，将张量从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中使用“steps”的其他组件，例如指标记录，或 learning rate schedulers，“step”被计为对模型参数的单个更新，而不是对数据的单个模型前向传递。假设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 命令行界面启用低精度优化器：

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¶

这里发生了什么？

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

此优化器可以包装任何基础优化器，并通过保持优化器状态并在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 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 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来通过加速训练来解决这个问题。

听起来很棒！我该怎么使用它？

您可以使用我们的任何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 是另一种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。