Hugging Face's logo

TRL documentation

减少内存使用

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

减少内存使用

训练工作流通常可以优化以降低内存消耗，TRL 提供了多项内置功能来帮助实现这一点。

下面，我们将概述这些技术，并建议尝试不同的组合，以找出最适合您特定环境的配置。

每种方法都包含受支持训练器的示例。如果您不确定某项技术是否与您的训练器兼容，请查看相应训练器的文档。

有关其他策略（例如所有训练器均支持的梯度检查点），请参阅transformers性能指南。

截断

数据集中的序列长度可能差异很大。在对数据进行批处理时，序列会被填充以匹配该批次中最长的序列，这可能导致内存占用过高，即使大多数序列相对较短也是如此。

Truncation prompt-completion

为了降低内存占用，将序列截断至合理长度至关重要。虽然 TRL 训练器默认会截断序列，但你可能需要根据具体使用场景调整默认的截断长度。

DPO
SFT

DPO 截断通过 max_length 进行控制，它会截断组合后的 prompt+completion 序列。

DPO truncation

要设置截断参数，请使用以下代码片段：

from trl import DPOConfig

training_args = DPOConfig(..., max_length=...)

旧版参数 max_prompt_length 和 max_completion_length 现已移除；请在训练前对数据集中过长的提示/补全进行过滤或预截断。

如何选择合适的 max_length 值？

如果 max_length 设置得过小，大量 token 将被丢弃，无法参与训练。如果设置得过大，内存占用可能会激增，从而引发内存溢出（OOM）错误。在未使用 packing 或 padding-free 的情况下，较大的 max_length 还可能导致训练效率低下，因为其中包含大量 padding token。

为了帮助您选择合适的值，我们提供了一个实用工具，用于可视化数据集中的序列长度分布。

打包

此技术仅适用于 SFT 训练以及使用 FlashAttention（或其变体）的配置。

截断存在以下几个缺点：

  1. 信息丢失：序列末尾的重要标记可能会被丢弃。
  2. 选择截断长度：过短会丢失数据；过长会降低效率。

Packing 通过将多个序列分组到同一训练行中，并将每行填充至 max_length，从而缓解这些问题。

Packing

TRL 使用 最佳适应递减（BFD） 装箱算法来实现打包，该算法在高效地分组序列的同时最大限度地减少填充。当序列长度超过 max_length 时，不同的策略决定了如何处理溢出的 token。

TRL 支持三种策略：

  • "bfd"（默认值）：使用 Best-Fit Decreasing 打包。如果序列长度超过 max_length，则丢弃溢出的 token。

  • "bfd_split"：使用 Best-Fit Decreasing packing，但在打包之前，长序列会被拆分为长度 ≤ max_length 的块。这保留了所有 token，并遵循了 Fewer Truncations Improve Language Modeling 中提出的方法。

  • "wrapped"：所有 token 被拼接成一个流，然后分割为固定长度的块。这种策略最大限度地减少了填充，但可能会混合不相关的样本。该策略对应于文献中描述的 concatenate-then-split 预处理方法（例如 Fewer Truncations Improve Language Modeling）。其缺点在于会破坏数据集中很大一部分序列的连续性，从而损害性能，这一点在 Qwen3-Coder-Next Technical Report 中有所讨论。

如果所有序列的长度都小于 max_length，bfd 和 bfd_split 的行为完全相同，因为无需进行截断或拆分。

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    packing=True,
    packing_strategy="bfd",
    max_length=512,
)

用于参数高效微调的 PEFT

诸如 LoRA 等参数高效微调（PEFT）方法，是降低训练期间内存占用最有效的技术之一。PEFT 方法不训练模型的所有参数，而是仅训练少量适配器参数，从而大幅降低内存需求，并使得在有限硬件上微调更大规模的模型成为可能。

有关使用 PEFT 与 TRL 的详细信息，包括各种适配器方法、量化选项和高级配置，请参阅 PEFT 集成。

要使用 PEFT 来减少内存占用：

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer

dataset = load_dataset("trl-lib/Capybara", split="train")

peft_config = LoraConfig()

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=peft_config,
)

PEFT 可以与其他内存缩减技术（例如 4 位或 8 位量化）结合使用，以实现更大的内存节省。有关量化的示例，请参阅 PEFT 集成。

无填充

无填充批处理（Padding-free batching）是降低内存占用的另一种方法。在此方法中，首先采样一个批次，然后将其展平为单个序列，从而避免填充。与打包（packing）不同，打包通过组合不同样本的部分可能导致序列不完整，而无填充批处理可确保所有序列保持完整且不被截断。

Padding-free

DPO
SFT

Padding-free 功能在 DPOTrainer 中暂时不可用：自 DPO 重构以来，设置 padding_free=True 会发出警告并回退到标准填充。该功能计划在后续更新中恢复。

from trl import DPOConfig

training_args = DPOConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

激活值卸载

激活卸载是一种内存效率技术，它在前向传播期间将激活张量临时移至 CPU 内存，并仅在反向传播需要时将其移回，从而降低 GPU 显存（VRAM）的使用量。这可以显著降低峰值内存使用量，代价是训练时间略有增加。

要在您的 SFT 训练配置中启用激活值卸载：

from trl import SFTConfig

training_args = SFTConfig(..., activation_offloading=True)

在底层，激活卸载通过实现 PyTorch 的 saved_tensors_hooks 来拦截前向传播过程中的激活值。它会根据张量的大小和上下文智能地管理哪些张量需要卸载，从而避免卸载那些效率低下的输出张量。为了优化性能，它可以通过一个标志位（默认为 true）使用 CUDA 流，将计算与 CPU-GPU 数据传输重叠执行。

将序列填充为倍数

目前，此技术仅支持 SFT 和 Reward 训练器。

启用此选项后，可确保所有序列都填充为指定值的倍数。
通过将序列长度对齐到对内存友好的边界，这可以在某些硬件上提高计算效率。

SFT
奖励
from trl import SFTConfig

training_args = SFTConfig(..., pad_to_multiple_of=2048)

PyTorch 缓存分配器

在长时间运行中，尤其是在线强化学习（GRPO、RLOO、Online DPO）场景下，GPU 内存可能会发生碎片化。设置 expandable_segments:True 可以让 PyTorch 的缓存分配器扩展现有内存段，从而减少已分配内存与已保留内存之间的差距：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTORCH_ALLOC_CONF=expandable_segments:True  # canonical name since PyTorch 2.10

在 vLLM 睡眠模式（vllm_enable_sleep_mode=True）下，vLLM 仅读取 PYTORCH_CUDA_ALLOC_CONF 以绕过此设置，因此请勿仅设置 PYTORCH_ALLOC_CONF。

在线方法中禁用生成时的模型收集

在使用 DeepSpeed ZeRO-3 时，模型权重会分片分布在多个 GPU 上。在线方法涉及在训练过程中从模型生成补全内容。在此步骤中，模型权重会临时汇集到单个 GPU 上进行生成。对于非常大的模型，这种汇集操作可能导致 OOM 错误，如该 issue 所述：#2250。

如果遇到此问题，可以通过设置以下参数来禁用生成过程中模型权重的收集：

GRPO
在线 DPO
RLOO
from trl import GRPOConfig

training_args = GRPOConfig(..., ds3_gather_for_generation=False)

此调整可防止模型权重被收集，从而避免 OOM 错误，但可能会导致生成速度变慢。

vLLM 睡眠模式

当使用 vLLM 作为在线训练方法的生成后端时，您可以启用 sleep mode，在优化步骤期间将 vLLM 的参数和缓存卸载到 CPU 内存中，并在需要执行权重同步和生成时重新加载回 GPU 显存。

GRPO
RLOO
from trl import GRPOConfig

training_args = GRPOConfig(..., vllm_enable_sleep_mode=True)

将 vLLM 的权重和缓存卸载有助于保持较低的 GPU 内存占用，这在训练大型模型或使用有限的 GPU 资源时尤为有益。然而，从睡眠模式唤醒 vLLM 引擎会引入一些主机-设备传输延迟，这可能会对训练速度产生轻微影响。

梯度检查点

梯度检查点通过不在前向传播期间存储所有中间激活值，而是在反向传播期间重新计算它们，从而以计算换取内存。

from trl import SFTConfig

training_args = SFTConfig(..., gradient_checkpointing=True)

所有训练器默认启用梯度检查点以优化内存使用。如有需要，可通过设置 gradient_checkpointing=False 将其禁用。

选择性激活检查点

使用 SFTTrainer，你可以在前向传播过程中保存注意力输出，而不是在反向传播中重新计算它。这可以恢复长上下文下大部分因检查点导致的性能下降，代价是每层额外增加一个隐藏状态大小的张量。它会强制使用非重入检查点。

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"selective": True},
)

有关更多内存优化技术，请参阅 Transformers 性能指南。

Update on GitHub

←自定义训练 加速训练→